tensorrt_llm/_torch/modules/fused_moe/quantization.py (2)

1481-1487: Guard against invalid scaling_vector_size; add asserts.

Hidden and intermediate sizes must be divisible by scaling_vector_size to avoid mis-shaped block scales and kernel mismatches.

Apply this diff:

 def create_weights(self,
                    module: torch.nn.Module,
                    weight_dtype,
                    weight_vec_size,
                    block_scales_dtype,
                    block_scales_vec_size,
-                   scaling_vector_size=16):
+                   scaling_vector_size=16):
 
-        module.scaling_vector_size = scaling_vector_size
+        module.scaling_vector_size = scaling_vector_size
+        assert module.hidden_size % module.scaling_vector_size == 0, (
+            f"hidden_size {module.hidden_size} must be divisible by scaling_vector_size "
+            f"{module.scaling_vector_size}"
+        )
+        assert module.intermediate_size_per_partition % module.scaling_vector_size == 0, (
+            f"intermediate_size_per_partition {module.intermediate_size_per_partition} must be divisible by "
+            f"scaling_vector_size {module.scaling_vector_size}"
+        )

---

2010-2058: Avoid hard-coded 32; use module.scaling_vector_size.

Reduces drift if the scaling width changes and keeps overrides consistent.

Apply this diff:

 class W4A8NVFP4FP8TRTLLMGenFusedMoEMethod(NVFP4TRTLLMGenFusedMoEMethod):

     def create_weights(self, module: torch.nn.Module):
         weight_vec_size = torch.iinfo(self.weight_dtype).bits // 4
         block_scales_vec_size = 1

-        NVFP4FusedMoEMethod.create_weights(self, module, self.weight_dtype,
-                                           weight_vec_size,
-                                           self.block_scales_dtype,
-                                           block_scales_vec_size, 32)
+        NVFP4FusedMoEMethod.create_weights(self, module, self.weight_dtype,
+                                           weight_vec_size,
+                                           self.block_scales_dtype,
+                                           block_scales_vec_size,
+                                           32)  # sets module.scaling_vector_size

@@
-    def load_expert_w3_w1_weight_scale_nvfp4(
+    def load_expert_w3_w1_weight_scale_nvfp4(
             self, module: torch.nn.Module, w1_weight_scale: torch.Tensor,
             w3_weight_scale: torch.Tensor,
             dst_w3_w1_weight_scale: torch.Tensor):
         return super().load_expert_w3_w1_weight_scale_nvfp4(
-            module, w1_weight_scale, w3_weight_scale, dst_w3_w1_weight_scale,
-            32)
+            module, w1_weight_scale, w3_weight_scale, dst_w3_w1_weight_scale,
+            module.scaling_vector_size)

@@
-    def load_expert_w2_weight_scale_nvfp4(self, module: torch.nn.Module,
+    def load_expert_w2_weight_scale_nvfp4(self, module: torch.nn.Module,
                                           w2_weight_scale: torch.Tensor,
                                           dst_w2_weight_scale: torch.Tensor):
         return super().load_expert_w2_weight_scale_nvfp4(
-            module, w2_weight_scale, dst_w2_weight_scale, 32)
+            module, w2_weight_scale, dst_w2_weight_scale,
+            module.scaling_vector_size)

tensorrt_llm/_torch/modules/fused_moe/create_moe.py (1)

49-53: Update warning message to include the new quant mode.

Keeps UX consistent with the actual supported set.

Apply this diff:

-            logger.warning(
-                "TRTLLMGenFusedMoE only supports fp8_block_scales, nvfp4, w4a16_mxfp4, w4a8_mxfp4_fp8 and w4a8_mxfp4_mxfp8. "
+            logger.warning(
+                "TRTLLMGenFusedMoE only supports fp8_block_scales, nvfp4, w4a16_mxfp4, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8 and w4a8_mxfp4_mxfp8. "
                 f"Check out details in quant_config: {quant_config}"
                 "Using CutlassFusedMoE instead.")

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (3)

115-116: Include new quant mode in the assertion message.

The condition allows it; the error text should reflect it.

Apply this diff:

-        assert self.has_deepseek_fp8_block_scales \
-            or self.has_nvfp4 or self.has_w4a16_mxfp4 or self.has_w4a8_nvfp4_fp8 \
-            or self.has_w4a8_mxfp4_fp8 or self.has_w4a8_mxfp4_mxfp8, "TRTLLMGenFusedMoE only supports fp8_block_scaling, nvfp4, w4a16_mxfp4, w4a8_mxfp4_fp8 and w4a8_mxfp4_mxfp8 dtypes."
+        assert self.has_deepseek_fp8_block_scales \
+            or self.has_nvfp4 or self.has_w4a16_mxfp4 or self.has_w4a8_nvfp4_fp8 \
+            or self.has_w4a8_mxfp4_fp8 or self.has_w4a8_mxfp4_mxfp8, \
+            "TRTLLMGenFusedMoE only supports fp8_block_scaling, nvfp4, w4a16_mxfp4, w4a8_nvfp4_fp8, w4a8_mxfp4_fp8 and w4a8_mxfp4_mxfp8 dtypes."

---

206-209: Consider enabling post‑quant allgather for W4A8 NVFP4 FP8.

If fp8_fp4_block_scale_moe_runner supports PQ allgather like nvfp4, include it here for performance parity.

Apply this diff if supported:

-        is_post_quant_allgather_supported = self.has_nvfp4 or self.has_w4a8_mxfp4_fp8 or self.has_w4a8_mxfp4_mxfp8
+        is_post_quant_allgather_supported = (
+            self.has_nvfp4
+            or self.has_w4a8_nvfp4_fp8
+            or self.has_w4a8_mxfp4_fp8
+            or self.has_w4a8_mxfp4_mxfp8
+        )

Please confirm the kernel supports PQ allgather; if not, ignore.

---

379-410: W4A8 NVFP4 FP8 forward: minor robustness and consistency tweaks.

• Pass router logits/bias as None under PQ allgather (matches other branches).
• Optional: verify whether padding to kernel alignment is ever required for this path.

Apply this diff:

-            outputs = torch.ops.trtllm.fp8_fp4_block_scale_moe_runner(
-                router_logits,
-                routing_bias,
+            outputs = torch.ops.trtllm.fp8_fp4_block_scale_moe_runner(
+                router_logits if not run_post_quant_allgather else None,
+                routing_bias if not run_post_quant_allgather else None,
                 hidden_states_fp8,
                 self.w3_w1_weight,
                 self.w3_w1_weight_scale.view(torch.float8_e4m3fn),
                 self.w2_weight,
                 self.w2_weight_scale.view(torch.float8_e4m3fn),
                 self.fc31_scale_c.data,
                 self.fc31_alpha.data,
                 self.fc2_alpha.data,

Also please validate on a model where hidden_size is not a multiple of common tiles (e.g., 6144) that no input padding is required for this kernel. If padding is required, mirror the mxfp4-fp8 path’s padding pattern.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Learnt from: djns99
PR: NVIDIA/TensorRT-LLM#6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

Learnt from: djns99
PR: NVIDIA/TensorRT-LLM#7104
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1475-1480
Timestamp: 2025-08-21T02:39:12.009Z
Learning: The min latency mode functionality in TensorRT-LLM MOE kernels (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu) is deprecated and no longer being maintained/updated, as confirmed by djns99. Bug reports and optimization suggestions for the computeStridesTmaWarpSpecializedLowLatencyKernel and related min latency code paths should be deprioritized.

Learnt from: djns99
PR: NVIDIA/TensorRT-LLM#6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.