diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index cd41979bc72b..3e4b8eac03d5 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -211,7 +211,7 @@ class FusedMoeRunner : public torch::CustomClassHolder } FusedMoeRunner(c10::ScalarType activation_dtype, c10::ScalarType weight_dtype, c10::ScalarType output_dtype, - bool use_deepseek_fp8_block_scale, bool use_w4_group_scaling, bool use_int8_woq_per_channel, + bool use_deepseek_fp8_block_scale, bool use_w4_group_scaling, bool use_woq_per_channel, bool use_mxfp8_act_scaling, bool use_mxfp8_weight_scaling, bool use_fused_finalize) { mActivationDtype = activation_dtype; @@ -219,7 +219,7 @@ class FusedMoeRunner : public torch::CustomClassHolder mOutputDtype = output_dtype; mUseDeepSeekFP8BlockScaling = use_deepseek_fp8_block_scale; mUseW4GroupScaling = use_w4_group_scaling; - mUseINT8WoqPerChannel = use_int8_woq_per_channel; + mUseWoqPerChannel = use_woq_per_channel; mUseMxfp8ActScaling = use_mxfp8_act_scaling; mUseMxfp8WeightScaling = use_mxfp8_weight_scaling; mUseFusedFinalize = use_fused_finalize; @@ -234,6 +234,14 @@ class FusedMoeRunner : public torch::CustomClassHolder && mWeightDtype == c10::ScalarType::Float8_e4m3fn), "use_mxfp8_weight_scaling requires both activation and weight dtypes to be Float8_e4m3fn."); + // The per-channel weight-only path reinterprets fc2's dimensions as + // [num_experts, inter_size, hidden_size] and, for INT4, treats the trailing + // dim as packed two-per-byte. That is only meaningful for integer + // weight-only quantization, so reject other weight dtypes here rather + // than silently transposing them downstream. + TORCH_CHECK( + !mUseWoqPerChannel || isIntWeightOnlyQuant(), "use_woq_per_channel requires an INT8 or INT4 weight dtype."); + // keep consistent with cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp if (mActivationDtype == c10::ScalarType::Half && mWeightDtype == c10::ScalarType::Half) { @@ -471,21 +479,28 @@ class FusedMoeRunner : public torch::CustomClassHolder ActivationType base_activation_type = activation_type.has_value() ? static_cast(activation_type.value()) : ActivationType::Swiglu; - if (mUseINT8WoqPerChannel) + if (mUseWoqPerChannel) { - // Note: The weight shape for INT8 weight only quantization is different, e.g., fc2_expert_weights: - // [num_experts, inter_size, hidden_size] + // Note: The weight shape for per-channel weight-only quantization is dim-swapped, e.g., + // fc2_expert_weights: [num_experts, inter_size, hidden_size] // Mirror the non-woq else-branch below: gated activations (Swiglu/Geglu) require fc1's // intermediate dim to be 2x fc2's (one half each for gate and up), while non-gated // activations (Relu2/Identity/ReLU/SiLU/Gelu, e.g. Nemotron-H) require them to be equal. + // + // Under this dim-swapped layout the sub-byte packing sits on fc1's inter dim + // (sizes()[2]), while fc2's inter dim (sizes()[1]) is unpacked, so mInnerDimMultiplier + // multiplies the fc1 side here -- unlike the non-swapped branch below. For INT8 the + // multiplier is 1 and this is identical to the previous form; for INT4 it is 2 (see the + // isInt4Quant() branch in the constructor) and the previous form rejected every valid + // shape. if (isGatedActivation(base_activation_type)) { - TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier * 2, + TORCH_CHECK(fc1_expert_weights.sizes()[2] * mInnerDimMultiplier == fc2_expert_weights.sizes()[1] * 2, "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); } else { - TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier, + TORCH_CHECK(fc1_expert_weights.sizes()[2] * mInnerDimMultiplier == fc2_expert_weights.sizes()[1], "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size."); } } @@ -506,16 +521,19 @@ class FusedMoeRunner : public torch::CustomClassHolder int experts_per_token = token_selected_experts.sizes()[1]; int64_t num_rows = input.sizes()[0]; int64_t hidden_size = fc2_expert_weights.sizes()[1]; - int64_t unpadded_hidden_size_val - = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; int64_t inter_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; - if (mUseINT8WoqPerChannel) + if (mUseWoqPerChannel) { - // Note: The weight shape for INT8 weight only quantization is different, e.g., fc2_expert_weights: + // Note: The weight shape for per-channel weight-only quantization is different, e.g., fc2_expert_weights: // [num_experts, inter_size, hidden_size] hidden_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; inter_size = fc2_expert_weights.sizes()[1]; } + // Default the output width only after the per-channel weight-only layout is resolved: fc2's dims + // are transposed on that path (and packed two-per-byte for INT4), so hidden_size is not the + // logical width until the swap above. + int64_t unpadded_hidden_size_val + = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; if (isWMxfp4AMxfp8Quant() || isWMxfp4AFp8Quant()) { @@ -789,16 +807,19 @@ class FusedMoeRunner : public torch::CustomClassHolder int experts_per_token = token_selected_experts.sizes()[1]; int64_t num_rows = input.sizes()[0]; int64_t hidden_size = fc2_expert_weights.sizes()[1]; - int64_t unpadded_hidden_size_val - = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; int64_t inter_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; - if (mUseINT8WoqPerChannel) + if (mUseWoqPerChannel) { - // Note: The weight shape for INT8 weight only quantization is different, e.g., fc2_expert_weights: + // Note: The weight shape for per-channel weight-only quantization is different, e.g., fc2_expert_weights: // [num_experts, inter_size, hidden_size] hidden_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; inter_size = fc2_expert_weights.sizes()[1]; } + // Default the output width only after the per-channel weight-only layout is resolved: fc2's dims + // are transposed on that path (and packed two-per-byte for INT4), so hidden_size is not the + // logical width until the swap above. + int64_t unpadded_hidden_size_val + = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; int const num_experts_on_rank = fc2_expert_weights.sizes()[0]; auto const num_experts_total = static_cast(num_experts_on_rank * ep_size); auto parallelism_config @@ -851,16 +872,16 @@ class FusedMoeRunner : public torch::CustomClassHolder // ordering differs from the non-woq path; both mirror the gated/non-gated split used in // runMoe(). Gated activations (Swiglu/Geglu) require fc1's intermediate dim to be 2x fc2's; // non-gated (Relu2/Identity/ReLU/SiLU/Gelu, e.g. Nemotron-H) require them to be equal. - if (mUseINT8WoqPerChannel) + if (mUseWoqPerChannel) { if (isGatedActivation(base_activation_type)) { - TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier * 2, + TORCH_CHECK(fc1_expert_weights.sizes()[2] * mInnerDimMultiplier == fc2_expert_weights.sizes()[1] * 2, "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); } else { - TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier, + TORCH_CHECK(fc1_expert_weights.sizes()[2] * mInnerDimMultiplier == fc2_expert_weights.sizes()[1], "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size."); } } @@ -978,15 +999,20 @@ class FusedMoeRunner : public torch::CustomClassHolder int64_t const num_rows = input.sizes()[0]; int64_t hidden_size = fc2_expert_weights.sizes()[1]; int64_t inter_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; - if (mUseINT8WoqPerChannel) + if (mUseWoqPerChannel) { - // Note: The weight shape for INT8 weight only quantization is different, e.g., fc2_expert_weights: + // Note: The weight shape for per-channel weight-only quantization is different, e.g., fc2_expert_weights: // [num_experts, inter_size, hidden_size] hidden_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; inter_size = fc2_expert_weights.sizes()[1]; } - int64_t const group_size_ - = isInt4Quant() ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::int4_group_size : -1; + // Only group-scaled INT4 (W4A8_AWQ) carries a group size. Plain per-channel + // W4A16 must profile with -1 so the profiler builds QuantParams::Int, matching + // getQuantParams() and the runner chosen in the constructor; deriving this from + // isInt4Quant() alone would profile a groupwise configuration runMoe never uses. + int64_t const group_size_ = (isInt4Quant() && mUseW4GroupScaling) + ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::int4_group_size + : -1; int64_t const group_size = isWFP4A16Quant() ? TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::wfp4a16_group_size : group_size_; @@ -1072,7 +1098,7 @@ class FusedMoeRunner : public torch::CustomClassHolder bool mUseDeepSeekFP8BlockScaling = false; bool mUseW4GroupScaling = false; - bool mUseINT8WoqPerChannel = false; + bool mUseWoqPerChannel = false; bool mUseMxfp8ActScaling = false; bool mUseFusedFinalize = true; bool mUseMxfp8WeightScaling = false; @@ -2348,10 +2374,10 @@ class FusedMoeRunner : public torch::CustomClassHolder else if (isIntWeightOnlyQuant()) { TORCH_CHECK(quant_scales.has_value(), "Expecting quant scales for weight only quantization"); - if (mUseINT8WoqPerChannel) + if (mUseWoqPerChannel) { - TORCH_CHECK( - quant_scales.value().size() == 2, "Expecting 2 quant scales for INT8 weight only quantization"); + TORCH_CHECK(quant_scales.value().size() == 2, + "Expecting 2 quant scales for per-channel weight only quantization"); auto& fc1_weight_scales = quant_scales.value()[0]; auto& fc2_weight_scales = quant_scales.value()[1]; return kernels::QuantParams::Int(static_cast(fc1_weight_scales.data_ptr()), diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py index f7ad6d0fadc0..f6049aba3c55 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/fused_moe/trtllm_moe.py @@ -438,7 +438,7 @@ def _run_moe_with_alltoall( activation_type=activation_type, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale, use_w4_group_scaling=False, - use_int8_woq_per_channel=False, + use_woq_per_channel=False, use_mxfp8_act_scaling=False, min_latency_mode=False, use_fused_finalize=True, diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 8cee91add853..95be1790f232 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -101,7 +101,7 @@ def __init__( cluster_rank: int, use_deepseek_fp8_block_scale: bool, use_w4_group_scaling: bool, - use_int8_woq_per_channel: bool, + use_woq_per_channel: bool, use_mxfp8_act_scaling: bool, min_latency_mode: bool, use_fused_finalize: bool, @@ -123,7 +123,7 @@ def __init__( self.enable_alltoall = False self.use_deepseek_fp8_block_scale = use_deepseek_fp8_block_scale self.use_w4_group_scaling = use_w4_group_scaling - self.use_int8_woq_per_channel = use_int8_woq_per_channel + self.use_woq_per_channel = use_woq_per_channel self.use_mxfp8_act_scaling = use_mxfp8_act_scaling self.use_mxfp8_weight_scaling = use_mxfp8_weight_scaling self.min_latency_mode = min_latency_mode @@ -133,7 +133,7 @@ def __init__( instance_key = (x_dtype, weight_dtype, output_dtype, use_deepseek_fp8_block_scale, use_w4_group_scaling, - use_int8_woq_per_channel, use_mxfp8_act_scaling, + use_woq_per_channel, use_mxfp8_act_scaling, use_mxfp8_weight_scaling) if instance_key not in MoERunner.runner_dict: @@ -141,7 +141,7 @@ def __init__( instance_key] = torch.classes.trtllm.FusedMoeRunner( x_dtype, weight_dtype, output_dtype, use_deepseek_fp8_block_scale, use_w4_group_scaling, - use_int8_woq_per_channel, use_mxfp8_act_scaling, + use_woq_per_channel, use_mxfp8_act_scaling, use_mxfp8_weight_scaling, use_fused_finalize) self.fused_moe_runner = MoERunner.runner_dict[instance_key] @@ -161,7 +161,7 @@ def unique_id(self): self.enable_alltoall, self.use_deepseek_fp8_block_scale, self.use_w4_group_scaling, - self.use_int8_woq_per_channel, + self.use_woq_per_channel, self.use_mxfp8_act_scaling, self.min_latency_mode, self.use_fused_finalize, @@ -226,7 +226,7 @@ def fused_moe( enable_alltoall: bool = False, use_deepseek_fp8_block_scale: bool = False, use_w4_group_scaling: bool = False, - use_int8_woq_per_channel: bool = False, + use_woq_per_channel: bool = False, use_mxfp8_act_scaling: bool = False, min_latency_mode: bool = False, use_fused_finalize: bool = True, @@ -288,7 +288,7 @@ def fused_moe( cluster_rank=cluster_rank, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale, use_w4_group_scaling=use_w4_group_scaling, - use_int8_woq_per_channel=use_int8_woq_per_channel, + use_woq_per_channel=use_woq_per_channel, use_mxfp8_act_scaling=use_mxfp8_act_scaling, min_latency_mode=min_latency_mode, use_fused_finalize=use_fused_finalize, @@ -403,7 +403,7 @@ def _(input: torch.Tensor, enable_alltoall: bool = False, use_deepseek_fp8_block_scale: bool = False, use_w4_group_scaling: bool = False, - use_int8_woq_per_channel: bool = False, + use_woq_per_channel: bool = False, use_mxfp8_act_scaling: bool = False, min_latency_mode: bool = False, use_fused_finalize: bool = True, @@ -432,10 +432,18 @@ def _(input: torch.Tensor, gated_slot_lora_weight_ptrs: Optional[torch.Tensor] = None, token_to_slot: Optional[torch.Tensor] = None): seq_len = input.shape[0] - if use_int8_woq_per_channel: - # Note: The weight shape for INT8 weight only quantization is different, i.e., - # fc2_expert_weights: [num_experts, inter_size, hidden_size] - hidden_size = fc2_expert_weights.shape[2] + if use_woq_per_channel: + # Note: The weight shape for per-channel weight-only quantization is + # dim-swapped, i.e. fc2_expert_weights: [num_experts, inter_size, hidden_size]. + # + # Sub-byte weights are packed along that trailing hidden dim, so the stored + # extent is hidden_size / elements_per_byte and must be scaled back up to + # recover the logical hidden size. This mirrors the real op, which applies + # mInnerDimMultiplier; without it the fake kernel reports + # half the hidden size for INT4 and shape inference silently disagrees with + # the kernel under torch.compile. + inner_dim_multiplier = 2 if fc2_expert_weights.dtype == torch.quint4x2 else 1 + hidden_size = fc2_expert_weights.shape[2] * inner_dim_multiplier else: hidden_size = fc2_expert_weights.shape[1] diff --git a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py index 1f6dfaf93548..406445c8fce4 100755 --- a/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/moe/fused_moe/fused_moe_cutlass.py @@ -41,9 +41,9 @@ DeepSeekFP8BlockScalesFusedMoEMethod, FP8QDQFusedMoEMethod, MoEWeightLoadingMode, MXFP8CutlassFusedMoEMethod, NVFP4CutlassFusedMoEMethod, INT8WoqPerChannelFusedMoEMethod, - W4A16NVFP4CutlassFusedMoEMethod, W4A8MXFP4FP8CutlassFusedMoEMethod, - W4A8MXFP4MXFP8CutlassFusedMoEMethod, WFP4A16FusedMoEMethod, - WInt4AFP8FusedMoEMethod) + W4A16NVFP4CutlassFusedMoEMethod, W4A16WoqPerChannelFusedMoEMethod, + W4A8MXFP4FP8CutlassFusedMoEMethod, W4A8MXFP4MXFP8CutlassFusedMoEMethod, + WFP4A16FusedMoEMethod, WInt4AFP8FusedMoEMethod) # isort: on from .routing import BaseMoeRoutingMethod @@ -159,6 +159,15 @@ class CutlassFusedMoE(MoEImplBase): "sm_constraint": ("min", 80), "dtypes": {torch.float16, torch.bfloat16}, }, + # W4A16 (plain INT4 weight-only, per-channel scales): SM >= 80. + # Uses the legacy mixed-dtype weight-only path, which SM120/121 reach + # via the SM80 interleaved layout, because + # preprocess_weights_for_mixed_gemm remaps sm >= 120 to 80, so the same + # code covers SM80 through SM121. + QuantAlgo.W4A16: { + "sm_constraint": ("min", 80), + "dtypes": {torch.float16, torch.bfloat16}, + }, # W4A16_MXFP4: SM == 90 only QuantAlgo.W4A16_MXFP4: { "sm_constraint": ("exact", 90), @@ -509,7 +518,7 @@ def reserve_moe_lora_cuda_graph_workspace(self, max_num_tokens: int, cluster_rank=self.cluster_rank, use_deepseek_fp8_block_scale=False, use_w4_group_scaling=False, - use_int8_woq_per_channel=False, + use_woq_per_channel=False, use_mxfp8_act_scaling=False, min_latency_mode=False, use_fused_finalize=self.use_fused_finalize, @@ -743,6 +752,25 @@ def has_int8_woq_per_channel(self): return self.quant_config and self.quant_config.layer_quant_mode.is_int8_weight_only( ) and not self.quant_config.layer_quant_mode.has_per_group_scaling() + @property + def has_int4_woq_per_channel(self) -> bool: + """True for plain W4A16: INT4 weights with per-channel scales.""" + # Excluding per-group keeps W4A8_AWQ on WInt4AFP8FusedMoEMethod, which + # is selected by is_int4_weight_only_per_group(). + if not self.quant_config: + return False + layer_quant_mode = self.quant_config.layer_quant_mode + return layer_quant_mode.is_int4_weight_only( + ) and not layer_quant_mode.has_per_group_scaling() + + @property + def has_woq_per_channel(self) -> bool: + """True for either per-channel weight-only dtype; drives the C++ flag.""" + # Both share the dim-swapped weight layout and the 2-element scale list + # that the runner turns into QuantParams::Int. + return bool(self.has_int8_woq_per_channel + or self.has_int4_woq_per_channel) + def quantize_input( self, x: Union[torch.Tensor, Fp4QuantizedTensor], @@ -783,6 +811,10 @@ def quantize_input( elif self.has_int8_woq_per_channel: # No quantization needed here, handled in kernel pass + elif self.has_int4_woq_per_channel: + # Weight-only: activations stay in 16-bit, dequantization of the + # INT4 weights happens inside the mixed-dtype GEMM. + pass elif self.has_nvfp4: if hasattr( self, @@ -865,6 +897,11 @@ def _get_quant_method(self): return WInt4AFP8FusedMoEMethod() elif self.has_int8_woq_per_channel: return INT8WoqPerChannelFusedMoEMethod() + elif self.has_int4_woq_per_channel: + # Placed after is_int4_weight_only_per_group() so W4A8_AWQ keeps + # priority on WInt4AFP8FusedMoEMethod; this branch is reached + # only by plain per-channel INT4. + return W4A16WoqPerChannelFusedMoEMethod() elif self.quant_config.layer_quant_mode.has_w4a8_mxfp4_fp8(): return W4A8MXFP4FP8CutlassFusedMoEMethod() elif self.quant_config.layer_quant_mode.has_w4a16_mxfp4(): @@ -999,6 +1036,14 @@ def run_moe( if self.has_any_quant: if self.has_w4afp8: weight_dtype = torch.quint4x2 + elif self.has_int4_woq_per_channel: + # Packed INT4 is stored in an int8 parameter. Without this view + # the op sees Char, isInt8Quant() matches and + # create_weight_quant_runner builds CutlassMoeFCRunner -- a silent fallback that runs at the wrong element + # width and leaves mInnerDimMultiplier at 1. quint4x2 makes + # isInt4Quant() select the uint4b_t runner instead. + weight_dtype = torch.quint4x2 elif self.has_w4a16_mxfp4: weight_dtype = torch.uint8 @@ -1038,7 +1083,7 @@ def run_moe( enable_alltoall=enable_alltoall, use_deepseek_fp8_block_scale=self.has_deepseek_fp8_block_scales, use_w4_group_scaling=self.has_w4afp8 or self.has_w4a16_mxfp4, - use_int8_woq_per_channel=self.has_int8_woq_per_channel, + use_woq_per_channel=self.has_woq_per_channel, # use_mxfp8_act_scaling drives dynamic MXFP8 activation quantization # before the GEMM; required for both W4A8 MXFP4xMXFP8 and W8A8 # MXFP8xMXFP8 paths. @@ -1134,7 +1179,7 @@ def _run_moe_w4a16_nvfp4( enable_alltoall=enable_alltoall, use_deepseek_fp8_block_scale=False, use_w4_group_scaling=False, - use_int8_woq_per_channel=False, + use_woq_per_channel=False, use_mxfp8_act_scaling=False, min_latency_mode=False, use_fused_finalize=self.use_fused_finalize, diff --git a/tensorrt_llm/_torch/moe/fused_moe/quantization.py b/tensorrt_llm/_torch/moe/fused_moe/quantization.py index b7ee8e25e871..e7571c78ce11 100644 --- a/tensorrt_llm/_torch/moe/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/moe/fused_moe/quantization.py @@ -1610,6 +1610,204 @@ def load_quant_scales(self, module: torch.nn.Module, weights: Dict): module.fc2_weight_scale.data.copy_(w2_scales.contiguous()) +class W4A16WoqPerChannelFusedMoEMethod(FusedMoEMethodBase): + """Plain W4A16: INT4 weights, per-channel scales, 16-bit activations. + + Storage mirrors ``INT8WoqPerChannelFusedMoEMethod`` (the dim-swapped, + post-transpose layout the per-channel weight-only runner expects) with the + trailing output dimension halved, because two INT4 values are packed into + one int8 byte. That matches the dense W4A16 convention, where + ``WeightOnlyQuantLinearMethod.create_weights`` allocates + ``(in_features, out_features // 2)``, and the C++ ``FusedMoeRunner`` + constructor sets ``mInnerDimMultiplier = 2`` for INT4 weights. + """ + + eplb_support_status = EplbSupportStatus.NOT_SUPPORTED + + # 2 INT4 values per int8 byte, matching mInnerDimMultiplier in FusedMoeRunner. + PACKED_ELEMENTS_PER_BYTE = 2 + + def create_weights(self, module: torch.nn.Module) -> None: + """Allocate packed INT4 expert weights and full-width per-channel scales.""" + module.sm_version = get_sm_version() + module.sm_version = 80 if module.sm_version >= 90 else module.sm_version + module.preprocessor = preprocess_weights_for_mixed_gemm + + if not module.quant_config.layer_quant_mode.is_int4_weight_only(): + raise NotImplementedError( + f"W4A16 MoE requires INT4 weight-only quantization. Got: {module.quant_config.layer_quant_mode}." + ) + + # Storage container is int8; each byte holds two INT4 values. The + # logical widths are unpacked, so only the trailing (output) dim is + # halved. Sized from module.expand_intermediate_size_per_partition + # (twice the per-partition intermediate size for gated activations, once + # otherwise) so gated and non-gated both work without a hardcoded 2. + expand_inter = module.expand_intermediate_size_per_partition + w3_w1_weight_shape = (module.expert_size_per_partition, + module.hidden_size, + expand_inter // self.PACKED_ELEMENTS_PER_BYTE) + w2_weight_shape = (module.expert_size_per_partition, + module.intermediate_size_per_partition, + module.hidden_size // self.PACKED_ELEMENTS_PER_BYTE) + + # Scales stay at full logical width: one scale per output channel. + fc31_weight_scale = nn.Parameter(torch.empty( + module.expert_size_per_partition, expand_inter, dtype=module.dtype), + requires_grad=False) + module.register_parameter("fc31_weight_scale", fc31_weight_scale) + + fc2_weight_scale = nn.Parameter(torch.empty( + module.expert_size_per_partition, + module.hidden_size, + dtype=module.dtype), + requires_grad=False) + module.register_parameter("fc2_weight_scale", fc2_weight_scale) + + super().create_weights(module, torch.int8, w3_w1_weight_shape, + w2_weight_shape) + + self._online_eplb_not_supported(module) + + self.setup_quant_scales(module) + + def setup_quant_scales(self, module: torch.nn.Module) -> None: + """Publish the 2-element per-channel scale tuple the runner expects.""" + # Reuses the per-channel 2-scale tuple; FusedMoeRunner's quant-scale + # handling turns this into QuantParams::Int(fc1_scale, fc2_scale). + module.quant_scales = FusedMoEQuantScalesINT8WoqPerChannel( + fc31_weight_scale=module.fc31_weight_scale, + fc2_weight_scale=module.fc2_weight_scale, + ) + + @staticmethod + def _validate_alignment(num_rows: int, name: str, + module: torch.nn.Module) -> None: + """Fail with a diagnostic before preprocess_weights_for_mixed_gemm's bare asserts. + + ``preprocess_weights_for_mixed_gemm`` asserts + ``num_rows % rows_per_tile == 0`` with + ``rows_per_tile = 128 * 8 // BITS_PER_ELT_A``. That is activation-driven, + so it is 64 for INT4 exactly as for INT8: W4A16 inherits W8A16's TP + restriction, no worse. Raised here so the message names the offending + tensor and tp_size instead of surfacing as an opaque AssertionError + inside the preprocessor. + """ + if num_rows % 64 != 0: + raise ValueError( + f"W4A16 MoE requires the pre-transpose row count of {name} to be a " + f"multiple of 64, got {num_rows} (tp_size={module.tp_size}). " + "preprocess_weights_for_mixed_gemm interleaves 64-row tiles. " + "For w2_weight this dimension is the per-partition intermediate " + "size, so a tensor-parallel size that keeps it 64-aligned is " + "required; for w3_w1_weight it is the hidden size.") + + def load_expert_w3_w1_weight(self, module: torch.nn.Module, + w1_weight: torch.Tensor, + w3_weight: Optional[torch.Tensor], + dst_w3_w1_weight: torch.Tensor) -> None: + """Load the w1 (and, when gated, w3) weights for one expert.""" + w1_weight_shard = load_weight_shard(w1_weight, module.tp_size, + module.tp_rank, + TensorParallelMode.COLUMN) + if w3_weight is not None and w3_weight.numel() > 0: + w3_weight_shard = load_weight_shard(w3_weight, module.tp_size, + module.tp_rank, + TensorParallelMode.COLUMN) + # Gated: w3 first, matching the fc31 scale concatenation order. + w31_weight_shard = torch.cat([w3_weight_shard, w1_weight_shard], + dim=0) + else: + # Non-gated activations (e.g. Nemotron-H squared-ReLU) have no w3. + w31_weight_shard = w1_weight_shard + + if module.dtype not in (torch.float16, torch.bfloat16): + raise ValueError( + "activation dtype should be float16 or bfloat16, got " + f"{module.dtype}") + + # shape[-1] is the pre-transpose column count == num_rows the + # preprocessor sees after .T (it reads shape[1] of the 3-D view). + self._validate_alignment(w31_weight_shard.shape[-1], "w3_w1_weight", + module) + + # Checkpoint entries are already packed two INT4 per byte along the + # OUTPUT dim, i.e. w1/w3 arrive as (inter/2, hidden), so the COLUMN + # shard and the dim-0 concat both operate in packed coordinates. The + # transpose then yields (hidden, expand_inter/2), which is exactly the + # destination parameter. preprocess_weights_for_mixed_gemm consumes and + # returns packed bytes -- it preserves the byte count and only permutes + # (its subbyte_transpose step) -- so no dimension is + # halved here. Packing along the output dim is required because a packed + # tensor cannot be transposed to move the packing axis. + w31_weight_shard = module.preprocessor(w31_weight_shard.T.contiguous(), + torch.quint4x2, module.dtype, + module.sm_version).contiguous() + dst_w3_w1_weight.copy_(w31_weight_shard.view(dst_w3_w1_weight.dtype), + non_blocking=True) + + def load_expert_w2_weight(self, module: torch.nn.Module, + w2_weight: torch.Tensor, + dst_w2_weight: torch.Tensor) -> None: + """Load the w2 weight for one expert.""" + # ROW shard: the split is on w2's input dim, so sharding and the + # last-dim packing do not interact. + w2_weight_shard = load_weight_shard(w2_weight, module.tp_size, + module.tp_rank, + TensorParallelMode.ROW) + + self._validate_alignment(w2_weight_shard.shape[-1], "w2_weight", module) + + w2_weight_shard = module.preprocessor(w2_weight_shard.T.contiguous(), + torch.quint4x2, module.dtype, + module.sm_version).contiguous() + dst_w2_weight.copy_(w2_weight_shard.view(dst_w2_weight.dtype), + non_blocking=True) + + def load_quant_scales(self, module: torch.nn.Module, weights: Dict) -> None: + """Load per-output-channel scales, concatenating w3 only when gated.""" + # The keys below are the per-expert VANILLA layout. Fused checkpoints + # store gate_up_proj_weight_scale / down_proj_weight_scale instead, so + # reject that mode here rather than failing on an opaque KeyError. + if (module.weight_loading_mode == + MoEWeightLoadingMode.FUSED_GATE_UP_PROJ): + raise ValueError( + "W4A16 per-channel MoE does not support loading scales from " + "MoEWeightLoadingMode.FUSED_GATE_UP_PROJ checkpoints.") + all_w1_scales = [ + load_weight_shard(weights[f"{expert_id}.w1.weight_scale"], + module.tp_size, module.tp_rank, + TensorParallelMode.COLUMN) + for expert_id in module.initial_local_expert_ids + ] + has_w3_scales = all(f"{expert_id}.w3.weight_scale" in weights + for expert_id in module.initial_local_expert_ids) + if module.is_gated_activation and has_w3_scales: + all_w3_scales = [ + load_weight_shard(weights[f"{expert_id}.w3.weight_scale"], + module.tp_size, module.tp_rank, + TensorParallelMode.COLUMN) + for expert_id in module.initial_local_expert_ids + ] + w3_w1_scales = torch.cat( + [torch.stack(all_w3_scales), + torch.stack(all_w1_scales)], + dim=-1) + else: + w3_w1_scales = torch.stack(all_w1_scales) + module.fc31_weight_scale.data.copy_( + w3_w1_scales.to(module.dtype).contiguous()) + + all_w2_scales = [ + load_weight_shard(weights[f"{expert_id}.w2.weight_scale"], + module.tp_size, module.tp_rank, + TensorParallelMode.ROW) + for expert_id in module.initial_local_expert_ids + ] + w2_scales = torch.stack(all_w2_scales).to(module.dtype) + module.fc2_weight_scale.data.copy_(w2_scales.contiguous()) + + class WInt4AFP8FusedMoEMethod(FusedMoEMethodBase): eplb_support_status = EplbSupportStatus.NOT_SUPPORTED diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 908331efffb3..217223638818 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -179,6 +179,9 @@ l0_b200: - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_selects_padded_quant_method - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_nvfp4_situ_fc31_scale_c_drops_dequant_scale - unittest/_torch/moe/test_moe_backend.py::test_trtllm_gen_situ_rejects_quant_algos_without_fused_cubins + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_w4a16_weight_shapes_gated_and_nongated + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_w4a16_unaligned_rows_raise_diagnostic + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_w4a16_aligned_rows_accepted - unittest/_torch/moe/test_moe_backend.py::test_megamoe_bakes_situ_softcaps_as_uniform_scalars - unittest/_torch/moe/test_moe_backend.py::test_megamoe_plain_swiglu_carries_no_constants - unittest/_torch/moe/test_moe_backend.py::test_create_moe_forwards_situ_activation_as_one_carrier diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index 4de2ed6b5794..9ec4f8fa7fec 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml @@ -88,6 +88,7 @@ l0_dgx_b200: - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and NVFP4" - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A8_MXFP4_MXFP8" - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W8A16" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A16 and not MXFP4" - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and MXFP8 and not W4A8" # --- TRTLLM --- - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "TRTLLM and NVFP4 and not W4A8" diff --git a/tests/integration/test_lists/test-db/l0_dgx_h100.yml b/tests/integration/test_lists/test-db/l0_dgx_h100.yml index 7c671ca75747..9706c2e47f58 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_h100.yml +++ b/tests/integration/test_lists/test-db/l0_dgx_h100.yml @@ -165,6 +165,7 @@ l0_dgx_h100: # --- CUTLASS --- - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8_BLOCK_SCALES" - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W8A16" + - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A16 and not MXFP4" - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A16_MXFP4" - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and W4A8_AWQ" # --- MARLIN (SM90-only; focused DEP + ALLGATHER x NVFP4 matrix) --- diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 76f392254576..026373510002 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -41,6 +41,9 @@ l0_h100: # ------------- MoE: test_moe_backend (by backend) --------------- # ------------- MoE: test_single_gpu (by backend) --------------- - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_single_gpu -k "CUTLASS and not None" + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_w4a16_weight_shapes_gated_and_nongated + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_w4a16_unaligned_rows_raise_diagnostic + - unittest/_torch/moe/test_moe_backend.py::test_cutlass_w4a16_aligned_rows_accepted # ------------- Kimi K3 MoE unit tests --------------- # Any-CUDA (Triton trtllm::situ_and_mul op + torch reference, requires_cuda # only; not SM100/SM103-gated), so they run on Hopper rather than consuming diff --git a/tests/unittest/_torch/moe/moe_test_utils.py b/tests/unittest/_torch/moe/moe_test_utils.py index 7fbd836a2f85..283b518135c5 100644 --- a/tests/unittest/_torch/moe/moe_test_utils.py +++ b/tests/unittest/_torch/moe/moe_test_utils.py @@ -731,7 +731,7 @@ def should_skip_cutlass( "intermediate_size)" ) - # TP per-shard alignment: W8A16, NVFP4, W4A8_AWQ, and MXFP8 require + # TP per-shard alignment: W8A16, W4A16, NVFP4, W4A8_AWQ, and MXFP8 require # 128-aligned per-shard intermediate_size. W8A16 fails in # preprocess_weights_for_mixed_gemm (num_rows % rows_per_tile != 0). NVFP4 # pads to 128-alignment (NVFP4_ROW_ALIGNMENT in quantization.py:2312) but @@ -743,9 +743,20 @@ def should_skip_cutlass( # int32 UE8M0 SF packing, so a non-128-aligned per-shard intermediate # raises AssertionError. # W4A8_MXFP4_MXFP8 uses MXFP4 auto-padding that handles this correctly. + # + # W4A16 (W4A16WoqPerChannelFusedMoEMethod) shares W8A16's constraint exactly, + # and it is no worse: rows_per_tile = 128*8//BITS_PER_ELT_A in + # preprocess_weights_for_mixed_gemm is activation-driven, so it is 64 for + # both INT4 and INT8 weights. INT4's other row constraints there are weaker + # divisors of 64 (B_ROWS_PER_MMA = 32, elts_in_int32 = 8), so 64 binds in + # both cases. The listed quants + # use the stricter 128 here rather than 64 because these multi-GPU configs + # also feed NVFP4-style paths; W4A16 is grouped with W8A16 for the same + # reason it shares the loader shape. if moe_tp_size > 1 and model_config is not None: tp_alignment_quants = { QuantAlgo.W8A16, + QuantAlgo.W4A16, QuantAlgo.NVFP4, QuantAlgo.W4A8_AWQ, QuantAlgo.MXFP8, diff --git a/tests/unittest/_torch/moe/quantize_utils.py b/tests/unittest/_torch/moe/quantize_utils.py index b6d7441ef2f4..d85d6ef14410 100644 --- a/tests/unittest/_torch/moe/quantize_utils.py +++ b/tests/unittest/_torch/moe/quantize_utils.py @@ -209,6 +209,14 @@ def get_test_quant_params(quant_algo, x, backend_type=None): elif quant_algo == QuantAlgo.W8A16: quantize_util_cls = W8A16QuantizeUtil quant_config = QuantConfig(quant_algo=QuantAlgo.W8A16) + elif quant_algo == QuantAlgo.W4A16: + # Plain W4A16: INT4 weight-only, per-channel scales. QuantConfig -> + # quant_mode maps W4A16 to use_weight_only(use_int4_weights=True) with no + # PER_GROUP bit, so is_int4_weight_only() is true and + # has_per_group_scaling() is false -- which is exactly what + # has_int4_woq_per_channel keys off in fused_moe_cutlass.py. + quantize_util_cls = W4A16QuantizeUtil + quant_config = QuantConfig(quant_algo=QuantAlgo.W4A16) elif quant_algo == QuantAlgo.W4A8_AWQ: quantize_util_cls = W4A8AWQQuantizeUtil quant_config = QuantConfig(quant_algo=QuantAlgo.W4A8_AWQ) @@ -2738,6 +2746,245 @@ def create_ref_module( return super().create_ref_module(routing_method, ref_cls) +# int4_woq_per_channel (plain W4A16) +# +# Nibble order below is dictated by the two consumers of these bytes, both read +# rather than assumed: +# * preprocess_weights_for_mixed_gemm's subbyte_transpose builds +# cat([low, high]) with low = (t << 4) >> 4 and high = t >> 4, i.e. the LOW +# nibble is the lower-indexed logical element +# (tensorrt_llm/quantization/functional.py). +# * unpack_int4_packed_tensor_to_int8 writes elt_0 = (packed << 4) >> 4 to the +# even index and elt_1 = packed >> 4 to the odd index +# (cpp/tensorrt_llm/thop/weightOnlyQuantOp.cpp). +# Both agree: even logical index -> low nibble, odd -> high nibble. +# +# Packing runs along dim 0 (the OUTPUT channel dim) because that is the axis the +# MoE loader needs. A packed tensor cannot be transposed to move its packing +# axis, and torch.ops.trtllm._symmetric_quantize_last_axis_of_batched_matrix +# packs along the LAST axis, so it cannot be reused +# here; the arithmetic is reproduced explicitly instead. + + +def pack_int4_along_dim0(unpacked: torch.Tensor) -> torch.Tensor: + """Pack an (out, in) int8 tensor of INT4 values into (out // 2, in) bytes. + + Row ``2r`` of the input becomes the low nibble of output row ``r`` and row + ``2r + 1`` becomes the high nibble. + """ + assert unpacked.shape[0] % 2 == 0, ( + f"output dim must be even to pack two INT4 per byte, got {unpacked.shape[0]}" + ) + low = (unpacked[0::2] & 0x0F).to(torch.uint8) + high = (unpacked[1::2] & 0x0F).to(torch.uint8) + return (low | (high << 4)).view(torch.int8) + + +def unpack_int4_along_dim0(packed: torch.Tensor) -> torch.Tensor: + """Inverse of :func:`pack_int4_along_dim0`, sign-extending each nibble.""" + as_u8 = packed.view(torch.uint8) + # Double shift sign-extends the 4-bit value, as unpack_int4_packed_tensor_to_int8 does. + low = (as_u8 << 4).view(torch.int8) >> 4 + high = as_u8.view(torch.int8) >> 4 + stacked = torch.stack([low, high], dim=1) + return stacked.reshape(packed.shape[0] * 2, packed.shape[1]) + + +class W4A16RefGatedMLPFusedMoE(RefMLPFusedMoE): + """ + A derived class of RefMLPFusedMoE serving as a reference implementation of + plain W4A16 (INT4 weight-only, per-channel scales) for correctness testing. + + Mirrors W8A16RefGatedMLPFusedMoE: GatedMLP has no W4A16 path, so the weights + are unpacked and dequantized in load_weights and the non-quantized forward is + used. + """ + + def __init__( + self, + num_experts: int, + routing_method: BaseMoeRoutingMethod, + hidden_size: int, + intermediate_size: int, + dtype: Optional[torch.dtype] = None, + model_config: Optional[ModelConfig] = None, + bias=False, + activation_type: ActivationType = ActivationType.Swiglu, + swiglu_alpha: Optional[torch.Tensor] = None, + swiglu_beta: Optional[torch.Tensor] = None, + swiglu_limit: Optional[torch.Tensor] = None, + ) -> None: + if activation_type not in ( + ActivationType.Swiglu, + ActivationType.Relu2, + ActivationType.Silu, + ): + raise ValueError( + f"Unsupported activation for W4A16RefGatedMLPFusedMoE: {activation_type}" + ) + self._original_quant_config = model_config.quant_config if model_config else None + # Forward activation_type so the base builds gated (Swiglu) or non-gated + # (squared-ReLU / SiLU) experts, as W8A16RefGatedMLPFusedMoE does. + super().__init__( + num_experts=num_experts, + routing_method=routing_method, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + model_config=ModelConfig(), # No quant_config; weights are dequantized. + bias=bias, + activation_type=activation_type, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + swiglu_limit=swiglu_limit, + ) + + def _dequantize(self, packed: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: + """Unpack along the output dim, then apply the per-output-channel scale.""" + unpacked = unpack_int4_along_dim0(packed.cpu()).to(packed.device) + # unpacked is (out, in); scale is (out,) -> broadcast over the input dim. + return (unpacked.T.contiguous().float() * scale).to(self.dtype).T.contiguous() + + def load_weights(self, weights_list: List[Dict]) -> None: + """Unpack and dequantize the INT4 expert weights into the fp reference.""" + assert len(weights_list) == 1 + weights = weights_list[0] + + assert ( + self._original_quant_config + and self._original_quant_config.quant_algo == QuantAlgo.W4A16 + ), "expect quant_algo to be W4A16" + + for expert in range(self.num_experts): + w1_dequant = self._dequantize( + weights[f"{expert}.w1.weight"], weights[f"{expert}.w1.weight_scale"] + ) + w2_dequant = self._dequantize( + weights[f"{expert}.w2.weight"], weights[f"{expert}.w2.weight_scale"] + ) + + down_proj_weights = [{}] + down_proj_weights[0]["weight"] = w2_dequant + + # Gated experts pack gate (w3) + up (w1) into gate_up_proj; non-gated + # experts (squared-ReLU) have only the up (w1) projection and their + # w3 entries are empty tensors, which must not be unpacked. + if self._is_gated: + w3_dequant = self._dequantize( + weights[f"{expert}.w3.weight"], weights[f"{expert}.w3.weight_scale"] + ) + gate_up_proj_weights = [{}, {}] + gate_up_proj_weights[0]["weight"] = w1_dequant + gate_up_proj_weights[1]["weight"] = w3_dequant + self.experts[expert].gate_up_proj.load_weights(gate_up_proj_weights) + else: + up_proj_weights = [{}] + up_proj_weights[0]["weight"] = w1_dequant + self.experts[expert].up_proj.load_weights(up_proj_weights) + + self.experts[expert].down_proj.load_weights(down_proj_weights) + + def check_accuracy( + self, + output: torch.Tensor, + ref_output: torch.Tensor, + weight_dtype: torch.dtype = torch.quint4x2, + ) -> None: + """Compare against the reference at the same thresholds as the W8A16 arm.""" + # Same helper and the same percent thresholds as the W8A16 arm; only the + # dtype differs, which widens atol via calc_woq_tolerence's + # bits_in_type = 4. Deliberately NOT loosened + # beyond that: per-channel INT4 accuracy is the gate that decides whether + # this feature ships, so a relaxed threshold would hide the signal. + atol = calc_woq_tolerence(ref_output, weight_dtype) + moe_tp_size = getattr(self, "moe_tp_size", 1) + if moe_tp_size > 1: + check_accuracy(output, ref_output, rtol=1e-1, atol=atol, percent=0.96) + else: + check_accuracy(output, ref_output, rtol=1e-7, atol=atol, percent=0.99) + + +class W4A16QuantizeUtil(BaseQuantizeUtil): + """ + W4A16QuantizeUtil inherits from BaseQuantizeUtil to support correctness + testing for plain W4A16 quantized MoE modules (INT4 weight-only, + per-channel scales). + """ + + def create_weights(self, **quant_kwargs) -> Dict[str, torch.Tensor]: + """ + Create quantized weights for MoE experts using plain W4A16. + + Values are drawn directly in the INT4 range rather than quantized from a + BF16 tensor, mirroring W8A16QuantizeUtil which draws int8 values with + torch.randint. The reference dequantizes with the same scales, so this + exercises the loader, packing and kernel path exactly without + introducing a second quantizer whose rounding convention could drift. + + On-disk layout is packed along the OUTPUT dim -- w1/w3 are + (intermediate_size // 2, hidden_size) and w2 is + (hidden_size // 2, intermediate_size) -- which is what + W4A16WoqPerChannelFusedMoEMethod's loader expects: the COLUMN shard and + dim-0 concat then operate in packed coordinates and the subsequent + transpose yields the destination parameter shape directly. + """ + assert self.quant_config is not None and self.quant_config.quant_algo == QuantAlgo.W4A16, ( + "expect quant_algo to be W4A16" + ) + weights = {} + for expert_id in range(self.num_experts): + # INT4 signed range is [-8, 7]; randint's upper bound is exclusive. + w1_unpacked = torch.randint( + -8, 8, (self.intermediate_size, self.hidden_size), dtype=torch.int8 + ).cuda() + w2_unpacked = torch.randint( + -8, 8, (self.hidden_size, self.intermediate_size), dtype=torch.int8 + ).cuda() + + # Per-output-channel scales, at full (unpacked) logical width. + w1_scale = ( + torch.randn(self.intermediate_size, dtype=self.dtype, device="cuda") + / self.hidden_size + ) + w2_scale = ( + torch.randn(self.hidden_size, dtype=self.dtype, device="cuda") + / self.intermediate_size + ) + + # Non-gated experts (e.g. Nemotron-H squared-ReLU) have no gate (w3) + # projection, so emit empty w3 tensors. Mirrors W8A16QuantizeUtil. + if self._is_gated: + w3_unpacked = torch.randint( + -8, 8, (self.intermediate_size, self.hidden_size), dtype=torch.int8 + ).cuda() + w3_packed = pack_int4_along_dim0(w3_unpacked) + w3_scale = ( + torch.randn(self.intermediate_size, dtype=self.dtype, device="cuda") + / self.hidden_size + ) + else: + w3_packed = torch.empty(0, dtype=torch.int8, device="cuda") + w3_scale = torch.empty(0, dtype=self.dtype, device="cuda") + + weights[f"{expert_id}.w1.weight"] = pack_int4_along_dim0(w1_unpacked) + weights[f"{expert_id}.w2.weight"] = pack_int4_along_dim0(w2_unpacked) + weights[f"{expert_id}.w3.weight"] = w3_packed + weights[f"{expert_id}.w1.weight_scale"] = w1_scale + weights[f"{expert_id}.w2.weight_scale"] = w2_scale + weights[f"{expert_id}.w3.weight_scale"] = w3_scale + return weights + + def create_ref_module( + self, + routing_method: BaseMoeRoutingMethod, + ref_cls: type = W4A16RefGatedMLPFusedMoE, + ) -> torch.nn.Module: + """ + Create a reference module for correctness testing. + """ + return super().create_ref_module(routing_method, ref_cls) + + class W4A8AWQRefGatedMLPFusedMoE(nn.Module): """ A reference implementation of W4A8_AWQ quantization for MoE correctness testing. diff --git a/tests/unittest/_torch/moe/test_moe_backend.py b/tests/unittest/_torch/moe/test_moe_backend.py index 3f4d41f500a8..f498f0c30e52 100644 --- a/tests/unittest/_torch/moe/test_moe_backend.py +++ b/tests/unittest/_torch/moe/test_moe_backend.py @@ -1060,6 +1060,158 @@ def test_enumerate_megamoe_candidate_tactics_curated_space() -> None: megamoe_op.validate_megamoe_tactic(tuple(invalid_tactic)) +@pytest.mark.parametrize( + "activation_type", + [ActivationType.Swiglu, ActivationType.Relu2], + ids=["gated_swiglu", "nongated_relu2"], +) +def test_cutlass_w4a16_weight_shapes_gated_and_nongated( + activation_type: ActivationType, +) -> None: + """W4A16 MoE must size its weights from expand_intermediate_size_per_partition. + + Nemotron-H uses squared-ReLU and is NON-gated, so a class that assumes + gating (a literal ``* 2``) allocates twice the fc1 rows it should and the + loader then writes past the logical extent. This pins the W4A16 class to + the same ``expand_intermediate_size_per_partition`` sizing that + ``INT8WoqPerChannelFusedMoEMethod`` uses. + + Also pins the INT4 packing axis: two INT4 values share one int8 byte along + the trailing (output) dim, matching the dense W4A16 convention + ``(in_features, out_features // 2)`` in ``WeightOnlyQuantLinearMethod`` and + the runner's ``mInnerDimMultiplier = 2``. + """ + if not torch.cuda.is_available(): + pytest.skip("CUDA required to construct a Cutlass MoE backend") + + from tensorrt_llm._utils import get_sm_version + + sm_version = get_sm_version() + if sm_version < 80: + pytest.skip(f"CutlassFusedMoE W4A16 requires SM >= 80, got SM{sm_version}") + + num_experts, top_k = 4, 2 + # 64-aligned so preprocess_weights_for_mixed_gemm's row-tile constraint is + # satisfied and is not what this test measures. + hidden_size, intermediate_size = 128, 256 + dtype = torch.bfloat16 + + backend = create_test_backend( + backend_type=MoeBackendType.CUTLASS, + routing_method=RenormalizeMoeRoutingMethod(top_k=top_k), + num_experts=num_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + dtype=dtype, + quant_config=QuantConfig(quant_algo=QuantAlgo.W4A16), + mapping=Mapping(world_size=1, rank=0, tp_size=1), + activation_type=activation_type, + ) + backend.create_weights() + + is_gated = is_gated_activation(activation_type) + expected_expand = intermediate_size * (2 if is_gated else 1) + assert backend.expand_intermediate_size_per_partition == expected_expand + + # Storage container stays int8; the packing is expressed in the shape. + assert backend.w3_w1_weight.dtype == torch.int8 + assert backend.w2_weight.dtype == torch.int8 + + # Trailing (output) dim halved by INT4 packing; leading dims unpacked. + assert tuple(backend.w3_w1_weight.shape) == ( + num_experts, + hidden_size, + expected_expand // 2, + ) + assert tuple(backend.w2_weight.shape) == ( + num_experts, + intermediate_size, + hidden_size // 2, + ) + + # Scales are per output channel, so they stay at full logical width. + assert tuple(backend.fc31_weight_scale.shape) == (num_experts, expected_expand) + assert tuple(backend.fc2_weight_scale.shape) == (num_experts, hidden_size) + + # The whole point of 4-bit: exactly half the weight bytes of W8A16, whose + # params are these shapes with the trailing dim unpacked. + w4a16_bytes = ( + backend.w3_w1_weight.numel() * backend.w3_w1_weight.element_size() + + backend.w2_weight.numel() * backend.w2_weight.element_size() + ) + w8a16_bytes = num_experts * (hidden_size * expected_expand + intermediate_size * hidden_size) + assert w4a16_bytes * 2 == w8a16_bytes + + +@pytest.mark.parametrize( + "intermediate_size,tp_size", + [(1856, 2), (1856, 4), (928, 2), (2688, 4)], +) +def test_cutlass_w4a16_unaligned_rows_raise_diagnostic( + intermediate_size: int, tp_size: int +) -> None: + """A non-64-aligned row count must fail with a message that names the cause. + + ``preprocess_weights_for_mixed_gemm`` enforces the row-tile constraint as a + BARE assert -- ``assert num_rows % rows_per_tile == 0`` -- inside a helper + several frames below the loader. Left alone it surfaces as a bare + ``AssertionError`` with no tensor name, no dimension and no tp_size. + + The shapes chosen here are the ones that break W8A16 today: + intermediate_size 1856 at TP 2/4 and 2688 at TP 4. W4A16 inherits the same + restriction and is no worse, because ``rows_per_tile = + 128*8//BITS_PER_ELT_A`` is activation-driven and therefore 64 for both INT4 + and INT8; INT4's other row constraints there are weaker divisors of 64 + (B_ROWS_PER_MMA = 32, elts_in_int32 = 8). + """ + from tensorrt_llm._torch.modules.fused_moe.quantization import W4A16WoqPerChannelFusedMoEMethod + + class _DummyModule: + pass + + module = _DummyModule() + module.tp_size = tp_size + + # The 64-row constraint applies to the PER-SHARD row count, which is what + # the loader passes: _validate_alignment(w2_weight_shard.shape[-1], ...). + num_rows = intermediate_size // tp_size + + # Precondition: these really are the unaligned cases, so the test is not + # vacuously passing on an aligned shape. + assert num_rows % 64 != 0 + + with pytest.raises(ValueError) as excinfo: + W4A16WoqPerChannelFusedMoEMethod._validate_alignment(num_rows, "w2_weight", module) + + message = str(excinfo.value) + # The diagnostic must identify the tensor, the offending count and the + # tp_size -- the three facts a bare assert throws away. + assert "w2_weight" in message + assert str(num_rows) in message + assert f"tp_size={tp_size}" in message + assert "64" in message + + +@pytest.mark.parametrize("num_rows", [64, 128, 1024, 2048, 2816]) +def test_cutlass_w4a16_aligned_rows_accepted(num_rows: int) -> None: + """The aligned TP sizes the method claims to support must pass cleanly. + + Counterpart to the rejection test: pinning only the failure would leave the + validator free to reject every shape, so the accepted side is pinned too. + """ + from tensorrt_llm._torch.modules.fused_moe.quantization import W4A16WoqPerChannelFusedMoEMethod + + class _DummyModule: + pass + + module = _DummyModule() + module.tp_size = 1 + + assert num_rows % 64 == 0 + # Must not raise. + W4A16WoqPerChannelFusedMoEMethod._validate_alignment(num_rows, "w2_weight", module) + + def run_backend_moe( backend: MoE, backend_type: MoeBackendType, @@ -1147,6 +1299,7 @@ def run_backend_moe( QuantAlgo.W4A8_MXFP4_MXFP8, QuantAlgo.MXFP8, QuantAlgo.W8A16, + QuantAlgo.W4A16, QuantAlgo.W4A8_AWQ, ] @@ -1295,7 +1448,7 @@ def generate_element_wise_test_params() -> List: SEQ_LENS_TO_TEST, DTYPES_TO_TEST, [MoeBackendType.CUTLASS, MoeBackendType.TRTLLM], - [None, QuantAlgo.NVFP4, QuantAlgo.W8A16], + [None, QuantAlgo.NVFP4, QuantAlgo.W8A16, QuantAlgo.W4A16], ): if skip_reason: continue @@ -1303,8 +1456,11 @@ def generate_element_wise_test_params() -> List: continue if backend_type == MoeBackendType.TRTLLM and quant_algo is None: continue - # INT8 weight-only per-channel non-gated MoE is CUTLASS-path only. - if quant_algo == QuantAlgo.W8A16 and backend_type != MoeBackendType.CUTLASS: + # Per-channel weight-only (INT8 and INT4) non-gated MoE is CUTLASS-path only. + if ( + quant_algo in (QuantAlgo.W8A16, QuantAlgo.W4A16) + and backend_type != MoeBackendType.CUTLASS + ): continue test_id = f"act={activation_type.name}-{base_test_id}" param_values = ( diff --git a/tests/unittest/_torch/moe/test_moe_module.py b/tests/unittest/_torch/moe/test_moe_module.py index 489c42b8ab01..eadcba56d14b 100644 --- a/tests/unittest/_torch/moe/test_moe_module.py +++ b/tests/unittest/_torch/moe/test_moe_module.py @@ -107,6 +107,7 @@ W4A8MXFP4MXFP8TRTLLMGenFusedMoEMethod, W4A8NVFP4FP8TRTLLMGenFusedMoEMethod, W4A16MXFP4TRTLLMGenFusedMoEMethod, + W4A16WoqPerChannelFusedMoEMethod, WFP4A16FusedMoEMethod, WInt4AFP8FusedMoEMethod, ) @@ -944,6 +945,7 @@ def _test_moe_multi_gpu( QuantAlgo.W4A8_MXFP4_MXFP8, QuantAlgo.MXFP8, QuantAlgo.W8A16, + QuantAlgo.W4A16, QuantAlgo.W4A8_AWQ, ] @@ -1872,6 +1874,12 @@ def _get_fused_moe_method_class(quant_algo, backend_type): # W4A8_AWQ uses is_int4_weight_only_per_group() -> WInt4AFP8FusedMoEMethod QuantAlgo.W4A8_AWQ: WInt4AFP8FusedMoEMethod, QuantAlgo.W8A16: INT8WoqPerChannelFusedMoEMethod, + # Plain W4A16 (INT4 weight-only, per-channel) resolves via + # has_int4_woq_per_channel, which is checked AFTER + # is_int4_weight_only_per_group so W4A8_AWQ keeps priority. Asserting + # the concrete class here is what catches a silent fall-through to + # INT8WoqPerChannelFusedMoEMethod. + QuantAlgo.W4A16: W4A16WoqPerChannelFusedMoEMethod, QuantAlgo.W4A16_MXFP4: WFP4A16FusedMoEMethod, QuantAlgo.W4A8_MXFP4_FP8: W4A8MXFP4FP8CutlassFusedMoEMethod, QuantAlgo.W4A8_MXFP4_MXFP8: W4A8MXFP4MXFP8CutlassFusedMoEMethod, diff --git a/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py b/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py index 99c0b7bff096..06063e5d143b 100644 --- a/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py +++ b/tests/unittest/_torch/peft/test_moe_lora_grouped_gemm.py @@ -361,7 +361,7 @@ def test_reserve_prevents_growth_across_captures(): cluster_rank=0, use_deepseek_fp8_block_scale=False, use_w4_group_scaling=False, - use_int8_woq_per_channel=False, + use_woq_per_channel=False, use_mxfp8_act_scaling=False, min_latency_mode=False, use_fused_finalize=True,