tests/unittest/_torch/modules/test_fused_moe.py (1)

481-497: Remove duplicate variable initialization in grouped_gemm function.

The d and m_indices variables are initialized twice in the nested function.

 def grouped_gemm(a: torch.Tensor, b: torch.Tensor, a_sf: torch.Tensor,
                  b_sf: torch.Tensor,
                  offset_array: torch.Tensor) -> torch.Tensor:
-    d = torch.empty((a.shape[0], b.shape[1]),
-                    device=b.device,
-                    dtype=torch.bfloat16)
-    m_indices = torch.empty(a.shape[0], device=b.device, dtype=torch.int32)
-    for idx in range(offset_array.numel() - 1):
-        m_indices[offset_array[idx]:offset_array[idx + 1]] = idx
-
     num_groups, n, k_ = b.shape
     d = torch.empty((a.shape[0], b.shape[1]),
                     device=b.device,
                     dtype=torch.bfloat16)
     m_indices = torch.empty(a.shape[0], device=b.device, dtype=torch.int32)
     for idx in range(offset_array.numel() - 1):
         m_indices[offset_array[idx]:offset_array[idx + 1]] = idx

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

368-374: Address workaround limitations and fix line length.

The code contains a workaround that restricts functionality to top-1 routing and has a line that exceeds the 120-character limit.

Apply this diff to fix the line length and document the limitation:

         if self.apply_router_weight_on_input:
-            assert self.routing_method.top_k == 1, "Current workaround only supports top-1 routing"
-            assert x.dtype != torch.float8_e4m3fn, "Current workaround for apply_router_weight_on_input does not support fp8 input"
+            assert self.routing_method.top_k == 1, "Current workaround only supports top-1 routing"
+            assert x.dtype != torch.float8_e4m3fn, (
+                "Current workaround for apply_router_weight_on_input does not support fp8 input"
+            )
             x = x * token_final_scales.to(x.dtype)
             # TODO: remove this once we have correct fusedmoe kernel ready
             token_final_scales = None

The TODO indicates incomplete implementation. Would you like me to open an issue to track the proper implementation of router weight application with the fused MoE kernel?

---

247-290: Add comprehensive documentation for the DeepGEMM function.

The function has extensive constraints and requirements that should be documented for maintainability.

Add documentation explaining the constraints:

 @nvtx_range("[DG]")
 def deepgemm_fp8_group_blockwise_gemm(
     a: torch.Tensor,
     b: torch.Tensor,
     sfa: torch.Tensor,
     sfb: torch.Tensor,
     masked_m: torch.Tensor,
     expected_m: int,
 ) -> torch.Tensor:
+    """
+    Perform FP8 group blockwise GEMM using DeepGEMM library.
+    
+    Args:
+        a: Input activation tensor [G, M, K] in FP8 format
+        b: Weight tensor [G, N, K] in FP8 format  
+        sfa: Scaling factors for tensor a
+        sfb: Scaling factors for tensor b
+        masked_m: Per-group M dimensions as int32 tensor
+        expected_m: Expected maximum M dimension
+        
+    Returns:
+        Output tensor [G, M, N] in bfloat16
+        
+    Constraints:
+        - Tensors must be contiguous with stride(-1) == 1
+        - All tensors must be FP8 e4m3fn format
+        - Output must be N-major for optimal performance
+    """
     d = torch.empty((a.shape[0], a.shape[1], b.shape[1]),
                     device=b.device,
                     dtype=torch.bfloat16)

tensorrt_llm/quantization/utils/fp8_utils.py (3)

32-52: Fix naming inconsistency between function name and return type.

The function name suggests E8M0 format but returns E4M3FN format, which is confusing.

Either rename the function or fix the return type:

 @nvtx_range("[DG] quantization")
 @torch.compile(dynamic=True)
-def per_token_cast_to_fp8_e8m0(
+def per_token_cast_to_fp8_e4m3fn(
         x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Cast tensor to FP8 E4M3FN format with per-token quantization.
+    
+    Args:
+        x: Input tensor (2D or 3D), last dimension must be divisible by 128
+        
+    Returns:
+        Tuple of (quantized tensor in FP8 E4M3FN, scaling factors)
+    """

---

54-80: Fix naming inconsistency in per_block_cast_to_fp8_e8m0.

Similar to per_token_cast_to_fp8_e8m0, this function name suggests E8M0 format but actually returns E4M3FN format (lines 65 and 77).

Either rename the function or change the implementation:

-def per_block_cast_to_fp8_e8m0(
+def per_block_cast_to_fp8_e4m3fn(
         x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

---

470-480: Remove incorrect assertion.

The assertion at line 472 checks if the last dimension is divisible by 2, but unlike silu_and_mul_masked_post_quant_fwd, this function doesn't divide k by 2.

     assert input.is_contiguous()
     assert len(input.shape) == 2
-    assert input.shape[-1] % 2 == 0

     # FP8 quantization parameters
     finfo = torch.finfo(torch.float8_e4m3fn)

tensorrt_llm/quantization/utils/fp8_utils.py (2)

304-316: Fix docstring formatting issues.

The docstring needs proper formatting according to D205 and D415 style rules.

 def silu_and_mul_masked_post_quant_fwd(
     input: torch.Tensor,
     quant_group_size: int,
     masked_m: torch.Tensor,
     scale_ue8m0: bool = False,
 ):
     """
-    input shape [g, m, k]
-    output shape [g, m, k // 2], dtype fp8
-    output_scale [g, k // 4, m // 2 // 128], dtype int32
-    quant_group_size int
-    masked_m shape [g]
+    Perform fused SiLU activation and multiplication with masked post-quantization.
+
+    Args:
+        input: Input tensor of shape [g, m, k]
+        quant_group_size: Quantization group size (int)
+        masked_m: Mask tensor of shape [g]
+        scale_ue8m0: Whether to scale to UE8M0 format
+
+    Returns:
+        Tuple of:
+        - output: Tensor of shape [g, m, k // 2] in fp8 dtype
+        - output_scale: Tensor of shape [g, k // 4, m] in int32 dtype
     """

---

457-468: Fix docstring formatting issues.

The docstring needs proper formatting according to D205 and D415 style rules.

 def per_token_quant_and_transform(
     input: torch.Tensor,
     quant_group_size: int = 128,
     scale_ue8m0: bool = True,
 ):
     """
-    input shape [g, m, k]
-    output shape [g, m, k // 2], dtype fp8
-    output_scale [g, k // 4, m // 2 // 128], dtype int32
-    quant_group_size int
-    masked_m shape [g]
+    Perform per-token quantization and transformation to FP8.
+
+    Args:
+        input: Input tensor of shape [m, k]
+        quant_group_size: Quantization group size (default: 128)
+        scale_ue8m0: Whether to scale to UE8M0 format (default: True)
+
+    Returns:
+        Tuple of:
+        - output: Tensor of shape [m, k] in fp8 dtype
+        - output_scale: Tensor of shape [m, scale_k] in int32 dtype
     """

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

Learnt from: moraxu
PR: NVIDIA/TensorRT-LLM#6303
File: tests/integration/test_lists/qa/examples_test_list.txt:494-494
Timestamp: 2025-07-28T17:06:08.621Z
Learning: In TensorRT-LLM testing, it's common to have both CLI flow tests (test_cli_flow.py) and PyTorch API tests (test_llm_api_pytorch.py) for the same model. These serve different purposes: CLI flow tests validate the traditional command-line workflow, while PyTorch API tests validate the newer LLM API backend. Both are legitimate and should coexist.

Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#5616
File: tensorrt_llm/executor/worker.py:375-384
Timestamp: 2025-07-17T09:01:27.402Z
Learning: In tensorrt_llm/executor/worker.py, the LoRA adapter cache optimization logic that checks `is_adapter_in_cpu_cache()` and conditionally passes None for weights/config has a known race condition issue that cannot be solved with simple error handling or verification checks. This is a known limitation that requires a more comprehensive solution.