[None][feat] Add W4A16 per-channel weight-only support for Cutlass fused MoE

cpp/tensorrt_llm/thop/moeOp.cpp

@@ -211,15 +211,15 @@ 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;
         mWeightDtype = weight_dtype;
         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<ActivationType>(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<int>(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<float const*>(fc1_weight_scales.data_ptr()),

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,

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,15 +133,15 @@ 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:
             MoERunner.runner_dict[
                 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]