[None][fix] Enable INT8 weight-only (W8A16) MoE for non-gated activations

cpp/tensorrt_llm/thop/moeOp.cpp

@@ -475,8 +475,19 @@ class FusedMoeRunner : public torch::CustomClassHolder
         {
             // Note: The weight shape for INT8 weight only quantization is different, e.g., fc2_expert_weights:
             // [num_experts, inter_size, hidden_size]
-            TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier * 2,
-                "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter 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.
+            if (isGatedActivation(base_activation_type))
+            {
+                TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier * 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,
+                    "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.");
+            }
         }
         else
         {
@@ -749,8 +760,6 @@ class FusedMoeRunner : public torch::CustomClassHolder
         }
         TORCH_CHECK(fc1_expert_weights.sizes()[0] == fc2_expert_weights.sizes()[0],
             "fc1_expert_weights and fc2_expert_weights must have the same number of experts.");
-        TORCH_CHECK(fc1_expert_weights.sizes()[1] == fc2_expert_weights.sizes()[2] * mInnerDimMultiplier * 2,
-            "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size.");
 
         TORCH_CHECK(!input_sf.has_value() || isWMxfp4AMxfp8Quant() || isNvfp4Quant(),
             "Block-scaling factors provided for non block-scaling quantization");
@@ -761,6 +770,13 @@ class FusedMoeRunner : public torch::CustomClassHolder
         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)
+        {
+            // Note: The weight shape for INT8 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];
+        }
         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
@@ -794,6 +810,38 @@ class FusedMoeRunner : public torch::CustomClassHolder
             reinterpret_cast<float const*>(swiglu_beta.has_value() ? swiglu_beta.value().const_data_ptr() : nullptr),
             reinterpret_cast<float const*>(swiglu_limit.has_value() ? swiglu_limit.value().const_data_ptr() : nullptr));
 
+        // Validate the fc1/fc2 inter-size relationship now that the activation type (gated vs
+        // non-gated) is finalized. INT8-woq uses a transposed weight layout, so its fc1/fc2 dim
+        // 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 (isGatedActivation(base_activation_type))
+            {
+                TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier * 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,
+                    "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.");
+            }
+        }
+        else
+        {
+            if (isGatedActivation(base_activation_type))
+            {
+                TORCH_CHECK(fc1_expert_weights.sizes()[1] == fc2_expert_weights.sizes()[2] * mInnerDimMultiplier * 2,
+                    "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size.");
+            }
+            else
+            {
+                TORCH_CHECK(fc1_expert_weights.sizes()[1] == fc2_expert_weights.sizes()[2] * mInnerDimMultiplier,
+                    "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.");
+            }
+        }
+
         setRunnerProfiles(profile_ids);
 
         auto stream = at::cuda::getCurrentCUDAStream(input.get_device());

tensorrt_llm/_torch/modules/fused_moe/quantization.py

@@ -1333,14 +1333,14 @@ def create_weights(self, module: torch.nn.Module):
         # since the quantized weights have their own layout
         w3_w1_weight_shape = (module.expert_size_per_partition,
                               module.hidden_size,
-                              module.intermediate_size_per_partition * 2)
+                              module.expand_intermediate_size_per_partition)
         w2_weight_shape = (module.expert_size_per_partition,
                            module.intermediate_size_per_partition,
                            module.hidden_size)
 
         fc31_weight_scale = nn.Parameter(torch.empty(
             module.expert_size_per_partition,
-            module.intermediate_size_per_partition * 2,
+            module.expand_intermediate_size_per_partition,
             dtype=module.dtype),
                                          requires_grad=False)
         module.register_parameter("fc31_weight_scale", fc31_weight_scale)
@@ -1375,10 +1375,19 @@ def load_expert_w3_w1_weight(self, module: torch.nn.Module,
         w1_weight_shard = load_weight_shard(w1_weight, module.tp_size,
                                             module.tp_rank,
                                             TensorParallelMode.COLUMN)
-        w3_weight_shard = load_weight_shard(w3_weight, module.tp_size,
-                                            module.tp_rank,
-                                            TensorParallelMode.COLUMN)
-        w31_weight_shard = torch.cat([w3_weight_shard, w1_weight_shard], dim=0)
+
+        # w3_weight (gate_proj) is empty for non-gated MoE (e.g. Nemotron-H squared-ReLU).
+        # Only concatenate the gate projection when present; otherwise the single
+        # up-projection fills the (non-doubled) intermediate buffer. The unquantized
+        # fused-MoE path handles non-gated experts the same way.
+        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)
+            w31_weight_shard = torch.cat([w3_weight_shard, w1_weight_shard],
+                                         dim=0)
+        else:
+            w31_weight_shard = w1_weight_shard
 
         weight_dtype = torch.int8
 
@@ -1419,25 +1428,33 @@ def load_expert_w2_weight(self, module: torch.nn.Module,
                             non_blocking=True)
 
     def load_quant_scales(self, module: torch.nn.Module, weights: Dict):
-        # fc31 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
-        ]
+        # fc31 scales. w1 (up_proj) is always present; w3 (gate_proj) is absent
+        # for non-gated MoE (e.g. Nemotron-H squared-ReLU). Only concatenate the
+        # gate-projection scales when the gate weights are present; otherwise the
+        # up-projection scales alone fill the (non-doubled) fc31 scale buffer.
         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
         ]
-        w3_w1_scales = torch.cat(
-            [torch.stack(all_w3_scales),
-             torch.stack(all_w1_scales)], dim=-1)
+        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)
         w3_w1_scales = w3_w1_scales.to(module.dtype)
         module.fc31_weight_scale.data.copy_(w3_w1_scales.contiguous())
-
         # fc2 scales
         all_w2_scales = [
             load_weight_shard(weights[f"{expert_id}.w2.weight_scale"],

tests/unittest/_torch/modules/moe/quantize_utils.py

@@ -2596,12 +2596,16 @@ def __init__(
         swiglu_beta: Optional[torch.Tensor] = None,
         swiglu_limit: Optional[torch.Tensor] = None,
     ):
-        assert activation_type == ActivationType.Swiglu, (
-            "Only Swiglu activation is supported for W8A16RefGatedMLPFusedMoE"
-        )
+        assert activation_type in (
+            ActivationType.Swiglu,
+            ActivationType.Relu2,
+            ActivationType.Silu,
+        ), f"Unsupported activation for W8A16RefGatedMLPFusedMoE: {activation_type}"
         # Store the original quant_config for assertion in load_weights
         self._original_quant_config = model_config.quant_config if model_config else None
-        # Create experts without quantization config since we'll dequantize weights
+        # Create experts without quantization config since we'll dequantize weights.
+        # Forward activation_type so the base builds gated (Swiglu) or non-gated
+        # (squared-ReLU / SiLU) experts as appropriate.
         super().__init__(
             num_experts=num_experts,
             routing_method=routing_method,
@@ -2610,6 +2614,7 @@ def __init__(
             dtype=dtype,
             model_config=ModelConfig(),  # No quant_config
             bias=bias,
+            activation_type=activation_type,
             swiglu_alpha=swiglu_alpha,
             swiglu_beta=swiglu_beta,
             swiglu_limit=swiglu_limit,
@@ -2628,24 +2633,33 @@ def load_weights(self, weights_list: List[Dict]):
             # Get quantized weights and scales
             w1 = weights[f"{expert}.w1.weight"]
             s1 = weights[f"{expert}.w1.weight_scale"]
-            w3 = weights[f"{expert}.w3.weight"]
-            s3 = weights[f"{expert}.w3.weight_scale"]
             w2 = weights[f"{expert}.w2.weight"]
             s2 = weights[f"{expert}.w2.weight_scale"]
 
             # Dequantize weights: w_dequant = (w.float() * scale).to(dtype)
             # Note: weights are (out_features, in_features), need transpose for matmul
             w1_dequant = (w1.T.contiguous().float() * s1).to(self.dtype).T.contiguous()
-            w3_dequant = (w3.T.contiguous().float() * s3).to(self.dtype).T.contiguous()
             w2_dequant = (w2.T.contiguous().float() * s2).to(self.dtype).T.contiguous()
 
             # Load as regular weights (no scales)
-            gate_up_proj_weights = [{}, {}]
             down_proj_weights = [{}]
-            gate_up_proj_weights[0]["weight"] = w1_dequant
-            gate_up_proj_weights[1]["weight"] = w3_dequant
             down_proj_weights[0]["weight"] = w2_dequant
-            self.experts[expert].gate_up_proj.load_weights(gate_up_proj_weights)
+
+            # Gated experts pack gate (w3) + up (w1) into gate_up_proj; non-gated
+            # experts (squared-ReLU) have only the up (w1) projection.
+            if self._is_gated:
+                w3 = weights[f"{expert}.w3.weight"]
+                s3 = weights[f"{expert}.w3.weight_scale"]
+                w3_dequant = (w3.T.contiguous().float() * s3).to(self.dtype).T.contiguous()
+                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, ref_output, weight_dtype=torch.int8):
@@ -2683,10 +2697,6 @@ def create_weights(self, **quant_kwargs) -> Dict[str, torch.Tensor]:
             w2_weight = torch.randint(
                 -128, 127, (self.hidden_size, self.intermediate_size), dtype=torch.int8
             ).cuda()
-            w3_weight = torch.randint(
-                -128, 127, (self.intermediate_size, self.hidden_size), dtype=torch.int8
-            ).cuda()
-
             # Per-channel scales
             w1_scale = (
                 torch.randn(self.intermediate_size, dtype=self.dtype, device="cuda")
@@ -2696,10 +2706,20 @@ def create_weights(self, **quant_kwargs) -> Dict[str, torch.Tensor]:
                 torch.randn(self.hidden_size, dtype=self.dtype, device="cuda")
                 / self.intermediate_size
             )
-            w3_scale = (
-                torch.randn(self.intermediate_size, dtype=self.dtype, device="cuda")
-                / self.hidden_size
-            )
+
+            # Non-gated experts (e.g. Nemotron-H squared-ReLU) have no gate (w3)
+            # projection, so emit empty w3 tensors. Mirrors NVFP4QuantizeUtil.
+            if self._is_gated:
+                w3_weight = torch.randint(
+                    -128, 127, (self.intermediate_size, self.hidden_size), dtype=torch.int8
+                ).cuda()
+                w3_scale = (
+                    torch.randn(self.intermediate_size, dtype=self.dtype, device="cuda")
+                    / self.hidden_size
+                )
+            else:
+                w3_weight = torch.empty(0, dtype=torch.int8, device="cuda")
+                w3_scale = torch.empty(0, dtype=self.dtype, device="cuda")
 
             weights[f"{expert_id}.w1.weight"] = w1_weight
             weights[f"{expert_id}.w2.weight"] = w2_weight

tests/unittest/_torch/modules/moe/test_moe_backend.py

@@ -791,14 +791,17 @@ def generate_element_wise_test_params() -> List:
             SEQ_LENS_TO_TEST,
             DTYPES_TO_TEST,
             [MoeBackendType.CUTLASS, MoeBackendType.TRTLLM],
-            [None, QuantAlgo.NVFP4],
+            [None, QuantAlgo.NVFP4, QuantAlgo.W8A16],
         ):
             if skip_reason:
                 continue
             if backend_type == MoeBackendType.CUTLASS and activation_type == ActivationType.Silu:
                 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:
+                continue
             test_id = f"act={activation_type.name}-{base_test_id}"
             param_values = (
                 dtype,