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