From 0fb65afed023c1f421f2086591ab16d0dcd03a24 Mon Sep 17 00:00:00 2001 From: potoior <2986485901@qq.com> Date: Thu, 13 Aug 2026 18:36:25 +0800 Subject: [PATCH] fix(qwen3_moe): export W4A16 NVFP4 checkpoints via NVFP4 MoE path QUANT_NVFP4_A16 (nvfp4_a16) was introduced in 0.10.0 to distinguish W4A16 weight-only NVFP4 from W4A4 NVFP4, but Qwen3SparseMoeBlock._prepare_moe_weights only gated on QUANT_NVFP4. Mixed-precision W4A16_NVFP4 checkpoints such as Qwen3.6-35B-A3B-NVFP4 therefore fell through to the INT4 GPTQ Marlin repack and crashed with: AttributeError: 'NVFP4A16MarlinLinear' object has no attribute 'qweight' Two changes: - Treat QUANT_NVFP4_A16 as NVFP4 when gating the NVFP4 MoE plugin path. - Accept NVFP4A16MarlinLinear experts in repack_nvfp4_gated_moe_experts (their raw weight/weight_scale/weight_scale_2 buffers are the same ModelOpt NVFP4 format, so decode_modelopt_nvfp4 reuses them as-is). Verified: tensorrt-edgellm-export completes for Qwen3.6-35B-A3B-NVFP4 (llm). --- tensorrt_edgellm/checkpoint/repacking.py | 10 +++++++--- .../models/qwen3_moe/modeling_qwen3_moe.py | 5 +++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/tensorrt_edgellm/checkpoint/repacking.py b/tensorrt_edgellm/checkpoint/repacking.py index 5f2a34e10..0b2af2533 100644 --- a/tensorrt_edgellm/checkpoint/repacking.py +++ b/tensorrt_edgellm/checkpoint/repacking.py @@ -1320,7 +1320,8 @@ def repack_nvfp4_gated_moe_experts( :data:`NVFP4_MOE_INTERMEDIATE_SIZE_ALIGNMENT` for ``"concat"``. """ from ..models.linear import \ - is_nvfp4_linear # local import to avoid circular dep + NVFP4A16MarlinLinear # local import to avoid circular dep + from ..models.linear import is_nvfp4_linear if fc1_layout == "interleave": build_fc1_dense = _interleave_gated_moe_fc1 @@ -1364,8 +1365,11 @@ def repack_nvfp4_gated_moe_experts( gate = expert.gate_proj up = expert.up_proj down = expert.down_proj - if not (is_nvfp4_linear(gate) and is_nvfp4_linear(up) - and is_nvfp4_linear(down)): + + def _is_nvfp4(m): + return is_nvfp4_linear(m) or isinstance(m, NVFP4A16MarlinLinear) + + if not (_is_nvfp4(gate) and _is_nvfp4(up) and _is_nvfp4(down)): raise TypeError("Gated NVFP4 MoE experts must use NVFP4 quant") gate_dense = decode_modelopt_nvfp4(gate.weight, gate.weight_scale, diff --git a/tensorrt_edgellm/models/qwen3_moe/modeling_qwen3_moe.py b/tensorrt_edgellm/models/qwen3_moe/modeling_qwen3_moe.py index 42190b411..f3de6fcfb 100644 --- a/tensorrt_edgellm/models/qwen3_moe/modeling_qwen3_moe.py +++ b/tensorrt_edgellm/models/qwen3_moe/modeling_qwen3_moe.py @@ -64,7 +64,7 @@ import torch import torch.nn as nn -from ...config import QUANT_FP16, QUANT_NVFP4, ModelConfig +from ...config import QUANT_FP16, QUANT_NVFP4, QUANT_NVFP4_A16, ModelConfig from ..default.modeling_default import (MLP, Attention, OnnxSpec, RMSNorm, _make_flat_wrapper) from ..linear import FP16Linear, make_linear @@ -223,7 +223,8 @@ def __init__(self, config: ModelConfig, module_prefix: str = "") -> None: self.hidden_size = config.hidden_size self.group_size = config.quant.group_size self.zero_point_offset = config.quant.gptq_zero_point_offset - self._use_nvfp4_moe = config.quant.quant_type == QUANT_NVFP4 + self._use_nvfp4_moe = config.quant.quant_type in (QUANT_NVFP4, + QUANT_NVFP4_A16) self._use_fp16_moe = config.quant.quant_type == QUANT_FP16 # All paths compute the same SwiGLU expert FFN; the integer is just # each plugin's own enum for it. Int4MoePlugin names the elementwise