From b225aefba7f828fbab2a48f8722c64fdb23ccfb0 Mon Sep 17 00:00:00 2001 From: Dorijan10 Date: Mon, 22 Jun 2026 14:56:21 +0000 Subject: [PATCH 1/5] Enable INT8 weight-only quantization for non-gated MoE experts The INT8 weight-only per-channel MoE path assumed gated activations (Swiglu/Geglu) in three places, rejecting or mis-handling non-gated experts (squared-ReLU, e.g. Nemotron-H) that the underlying CUTLASS kernels already support: - moeOp.cpp: the woq validation hardcoded fc1.inter == 2 * fc2.inter; now conditioned on isGatedActivation(), mirroring the existing non-woq branch. - INT8WoqPerChannelFusedMoEMethod: buffer sizing, weight loading, and scale loading assumed the doubled (gate+up) layout; now handle the single up-projection when the gate weight is absent, mirroring the unquantized fused-MoE path's existing non-gated handling. Gated models are unaffected (they retain the original code path). Signed-off-by: Dorijan10 --- cpp/tensorrt_llm/thop/moeOp.cpp | 49 ++++++++++++++++-- .../_torch/modules/fused_moe/quantization.py | 51 ++++++++++++------- 2 files changed, 79 insertions(+), 21 deletions(-) diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index feede19bf71a..7634faaf8a33 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -480,8 +480,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 { @@ -771,8 +782,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"); @@ -815,6 +824,38 @@ class FusedMoeRunner : public torch::CustomClassHolder reinterpret_cast(swiglu_alpha.has_value() ? swiglu_alpha.value().const_data_ptr() : nullptr), reinterpret_cast(swiglu_beta.has_value() ? swiglu_beta.value().const_data_ptr() : nullptr), reinterpret_cast(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); diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 42c9c84e2408..1314d5cde0ae 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -1312,14 +1312,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) @@ -1354,10 +1354,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 @@ -1398,25 +1407,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"], From 3573e0ed38f27491372c9c67b167220b01ce0769 Mon Sep 17 00:00:00 2001 From: Dorijan10 Date: Wed, 24 Jun 2026 16:36:17 +0100 Subject: [PATCH 2/5] [None][test] Add non-gated INT8 weight-only per-channel MoE test Covers the non-gated (squared-ReLU) path of INT8 weight-only per-channel fused MoE: the gate projection is absent, so the intermediate buffer is single-width (expand ratio 1). Complements the existing gated test and exercises the non-gated handling enabled by the preceding commit. Signed-off-by: Dorijan10 --- .../unittest/_torch/modules/test_fused_moe.py | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) diff --git a/tests/unittest/_torch/modules/test_fused_moe.py b/tests/unittest/_torch/modules/test_fused_moe.py index c27ca12f4647..f711e89d827f 100644 --- a/tests/unittest/_torch/modules/test_fused_moe.py +++ b/tests/unittest/_torch/modules/test_fused_moe.py @@ -41,6 +41,7 @@ NVFP4CutlassFusedMoEMethod # isort: on from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm._utils import get_sm_version, mpi_rank from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -2977,6 +2978,77 @@ def __init__(self, hidden_size, intermediate_size, expand_ratio, self.block_scales_vec_size = 4 # 4 fp8 values packed into int32 +@pytest.mark.skipif(torch.cuda.device_count() < 1, + reason="needs 1 GPU to run this test") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_fused_moe_int8_woq_per_channel_non_gated(dtype): + """Non-gated (squared-ReLU) INT8 weight-only per-channel MoE. + + The gate projection (w3) is absent for non-gated experts, so the + intermediate buffer is single-width (expand ratio 1, not 2). This + exercises the non-gated path in the INT8 weight-only CUTLASS fused-MoE. + """ + mapping = Mapping() + mapping.rank = mpi_rank() + + with torch.device(f'cuda:{mapping.rank}'): + SEQ_LEN = 4 + HIDDEN_SIZE = 768 + INTERMEDIATE_SIZE = 640 + NUM_EXPERTS = 3 + TOP_K = 2 + routing_method = RenormalizeMoeRoutingMethod(top_k=TOP_K) + torch.manual_seed(0) + torch.cuda.manual_seed(0) + x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") + router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), + dtype=dtype, + device="cuda") + + quant_config = QuantConfig(quant_algo=QuantAlgo.W8A16) + weights = {} + for expert_id in range(NUM_EXPERTS): + # Non-gated: only w1 (up) and w2 (down); no w3 (gate). + w1_weight = torch.randint(-128, + 127, (INTERMEDIATE_SIZE, HIDDEN_SIZE), + dtype=torch.int8).cuda() + w2_weight = torch.randint(-128, + 127, (HIDDEN_SIZE, INTERMEDIATE_SIZE), + dtype=torch.int8).cuda() + w1_scale = torch.randn( + (INTERMEDIATE_SIZE), dtype=dtype, device="cuda") / HIDDEN_SIZE + w2_scale = torch.randn( + (HIDDEN_SIZE), dtype=dtype, device="cuda") / INTERMEDIATE_SIZE + + weights[f"{expert_id}.w1.weight"] = w1_weight + weights[f"{expert_id}.w2.weight"] = w2_weight + weights[f"{expert_id}.w1.weight_scale"] = w1_scale + weights[f"{expert_id}.w2.weight_scale"] = w2_scale + + fused_moe = CutlassFusedMoE( + num_experts=NUM_EXPERTS, + routing_method=routing_method, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + dtype=dtype, + reduce_results=False, + model_config=ModelConfig(quant_config=quant_config), + activation_type=ActivationType.Relu2) + + # Non-gated buffer is single-width (expand ratio 1, not 2). + assert fused_moe.intermediate_size_expand_ratio == 1 + assert not fused_moe.is_gated_activation + + fused_moe.load_weights([weights]) + fused_moe.cuda() + + with torch.inference_mode(), autotune(): + output = fused_moe.forward(x, router_logits) + + assert output.shape == (SEQ_LEN, HIDDEN_SIZE) + assert torch.isfinite(output).all() + + def test_nvfp4_cutlass_get_weights_shapes_error_cases(): """Test NVFP4CutlassFusedMoEMethod.get_weights_shapes for error cases.""" method = NVFP4CutlassFusedMoEMethod() From dca287d8fe5f2a5605311bd3a8003f4c13f8a086 Mon Sep 17 00:00:00 2001 From: Dorijan10 Date: Wed, 24 Jun 2026 18:44:58 +0100 Subject: [PATCH 3/5] [None][fix] Apply INT8 woq size swap in runMoeMinLantency runMoeMinLantency derived hidden_size/inter_size from the non-woq weight layout and fed them to getWorkspaceInfo/runMoe without the INT8 weight-only per-channel swap that runMoe and runGemmProfile already apply. Add the same swap so all three paths use consistent dimensions for INT8 woq per-channel. Signed-off-by: Dorijan10 --- cpp/tensorrt_llm/thop/moeOp.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index 7634faaf8a33..e3440eaa3ae2 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -792,6 +792,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(num_experts_on_rank * ep_size); auto parallelism_config From 2d71065b495a6974fc0209b84bac32b4d83fc3df Mon Sep 17 00:00:00 2001 From: Dorijan10 Date: Fri, 26 Jun 2026 16:02:36 +0100 Subject: [PATCH 4/5] [None][test] Move non-gated INT8 W8A16 MoE coverage into the unified framework Port the non-gated (squared-ReLU) INT8 weight-only per-channel MoE coverage from the deprecated test_fused_moe.py smoke test into the unified moe/test_moe_backend.py framework, which validates against a dequantized reference rather than only output shape and finiteness. - quantize_utils.py: W8A16QuantizeUtil.create_weights gains a non-gated branch (empty w3); W8A16RefGatedMLPFusedMoE forwards activation_type to the base reference and loads the single up-projection when non-gated, mirroring the existing NVFP4 element-wise support. - test_moe_backend.py: add W8A16 to the element-wise parameter set, guarded to the CUTLASS backend (the path this PR fixes). - Remove the superseded non-gated smoke test from test_fused_moe.py. Validated on A100: 20 non-gated + 20 gated W8A16 cases pass (float16 + bfloat16, CUTLASS). Signed-off-by: Dorijan10 --- .../_torch/modules/moe/quantize_utils.py | 58 ++++++++++----- .../_torch/modules/moe/test_moe_backend.py | 5 +- .../unittest/_torch/modules/test_fused_moe.py | 72 ------------------- 3 files changed, 43 insertions(+), 92 deletions(-) diff --git a/tests/unittest/_torch/modules/moe/quantize_utils.py b/tests/unittest/_torch/modules/moe/quantize_utils.py index 0cb3fd613beb..2afbe19584ed 100644 --- a/tests/unittest/_torch/modules/moe/quantize_utils.py +++ b/tests/unittest/_torch/modules/moe/quantize_utils.py @@ -2280,12 +2280,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, @@ -2294,6 +2298,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, @@ -2312,24 +2317,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): @@ -2367,10 +2381,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") @@ -2380,10 +2390,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 diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 8006a8c890c8..7264380c7764 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -450,7 +450,7 @@ 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 @@ -458,6 +458,9 @@ 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: + continue test_id = f"act={activation_type.name}-{base_test_id}" param_values = ( dtype, diff --git a/tests/unittest/_torch/modules/test_fused_moe.py b/tests/unittest/_torch/modules/test_fused_moe.py index f711e89d827f..c27ca12f4647 100644 --- a/tests/unittest/_torch/modules/test_fused_moe.py +++ b/tests/unittest/_torch/modules/test_fused_moe.py @@ -41,7 +41,6 @@ NVFP4CutlassFusedMoEMethod # isort: on from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._torch.utils import ActivationType from tensorrt_llm._utils import get_sm_version, mpi_rank from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -2978,77 +2977,6 @@ def __init__(self, hidden_size, intermediate_size, expand_ratio, self.block_scales_vec_size = 4 # 4 fp8 values packed into int32 -@pytest.mark.skipif(torch.cuda.device_count() < 1, - reason="needs 1 GPU to run this test") -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_fused_moe_int8_woq_per_channel_non_gated(dtype): - """Non-gated (squared-ReLU) INT8 weight-only per-channel MoE. - - The gate projection (w3) is absent for non-gated experts, so the - intermediate buffer is single-width (expand ratio 1, not 2). This - exercises the non-gated path in the INT8 weight-only CUTLASS fused-MoE. - """ - mapping = Mapping() - mapping.rank = mpi_rank() - - with torch.device(f'cuda:{mapping.rank}'): - SEQ_LEN = 4 - HIDDEN_SIZE = 768 - INTERMEDIATE_SIZE = 640 - NUM_EXPERTS = 3 - TOP_K = 2 - routing_method = RenormalizeMoeRoutingMethod(top_k=TOP_K) - torch.manual_seed(0) - torch.cuda.manual_seed(0) - x = torch.randn((SEQ_LEN, HIDDEN_SIZE), dtype=dtype, device="cuda") - router_logits = torch.randn((SEQ_LEN, NUM_EXPERTS), - dtype=dtype, - device="cuda") - - quant_config = QuantConfig(quant_algo=QuantAlgo.W8A16) - weights = {} - for expert_id in range(NUM_EXPERTS): - # Non-gated: only w1 (up) and w2 (down); no w3 (gate). - w1_weight = torch.randint(-128, - 127, (INTERMEDIATE_SIZE, HIDDEN_SIZE), - dtype=torch.int8).cuda() - w2_weight = torch.randint(-128, - 127, (HIDDEN_SIZE, INTERMEDIATE_SIZE), - dtype=torch.int8).cuda() - w1_scale = torch.randn( - (INTERMEDIATE_SIZE), dtype=dtype, device="cuda") / HIDDEN_SIZE - w2_scale = torch.randn( - (HIDDEN_SIZE), dtype=dtype, device="cuda") / INTERMEDIATE_SIZE - - weights[f"{expert_id}.w1.weight"] = w1_weight - weights[f"{expert_id}.w2.weight"] = w2_weight - weights[f"{expert_id}.w1.weight_scale"] = w1_scale - weights[f"{expert_id}.w2.weight_scale"] = w2_scale - - fused_moe = CutlassFusedMoE( - num_experts=NUM_EXPERTS, - routing_method=routing_method, - hidden_size=HIDDEN_SIZE, - intermediate_size=INTERMEDIATE_SIZE, - dtype=dtype, - reduce_results=False, - model_config=ModelConfig(quant_config=quant_config), - activation_type=ActivationType.Relu2) - - # Non-gated buffer is single-width (expand ratio 1, not 2). - assert fused_moe.intermediate_size_expand_ratio == 1 - assert not fused_moe.is_gated_activation - - fused_moe.load_weights([weights]) - fused_moe.cuda() - - with torch.inference_mode(), autotune(): - output = fused_moe.forward(x, router_logits) - - assert output.shape == (SEQ_LEN, HIDDEN_SIZE) - assert torch.isfinite(output).all() - - def test_nvfp4_cutlass_get_weights_shapes_error_cases(): """Test NVFP4CutlassFusedMoEMethod.get_weights_shapes for error cases.""" method = NVFP4CutlassFusedMoEMethod() From 801e9813ef3ec7fbac04afb8c3c08af057b7d881 Mon Sep 17 00:00:00 2001 From: Dorijan10 Date: Fri, 24 Jul 2026 12:10:30 +0100 Subject: [PATCH 5/5] [None][chore] Fix pre-commit formatting (clang-format tabs, yapf, trailing whitespace) Formatting only, no functional change: two stray tabs in moeOp.cpp and two yapf line-wraps in load_quant_scales. 'git diff -w' on the moeOp.cpp hunk is empty. Signed-off-by: Dorijan10 --- cpp/tensorrt_llm/thop/moeOp.cpp | 4 ++-- tensorrt_llm/_torch/modules/fused_moe/quantization.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index e3440eaa3ae2..dcec3e7a20fa 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -480,7 +480,7 @@ 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] - // Mirror the non-woq else-branch below: gated activations (Swiglu/Geglu) require fc1's + // 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)) @@ -831,7 +831,7 @@ class FusedMoeRunner : public torch::CustomClassHolder reinterpret_cast(swiglu_alpha.has_value() ? swiglu_alpha.value().const_data_ptr() : nullptr), reinterpret_cast(swiglu_beta.has_value() ? swiglu_beta.value().const_data_ptr() : nullptr), reinterpret_cast(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 diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 1314d5cde0ae..77b5875fe18c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -1417,9 +1417,8 @@ def load_quant_scales(self, module: torch.nn.Module, weights: Dict): 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) + 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"], @@ -1429,7 +1428,8 @@ def load_quant_scales(self, module: torch.nn.Module, weights: Dict): ] w3_w1_scales = torch.cat( [torch.stack(all_w3_scales), - torch.stack(all_w1_scales)], dim=-1) + 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)