Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

benchmarks/bench_trtllm_gen_fused_moe_autotuner.py (1)

21-25: Guard FP8 quantization against all-zero inputs.

All-zero inputs make max zero, scale infinite, and the quantized tensor NaN (0 × ∞), which breaks the benchmark when buffers start cleared. Please handle the zero case before inverting the scale.

Apply this diff to fix the issue:

 def fp8_quantize(x):
-    max = x.float().abs().nan_to_num().max()
-    scale = FLOAT8_E4M3_MAX / max
-    x = (x * scale).to(torch.float8_e4m3fn)
-    return x, 1.0 / scale
+    max_val = x.float().abs().nan_to_num().max()
+    if max_val == 0:
+        return torch.zeros_like(x, dtype=torch.float8_e4m3fn), 1.0
+    scale = FLOAT8_E4M3_MAX / max_val
+    quantized = (x * scale).to(torch.float8_e4m3fn)
+    return quantized, 1.0 / scale

benchmarks/bench_trtllm_gen_fused_moe_autotuner.py (1)

67-89: Consider refactoring lambda to a named function.

The lambda expression is quite long (22 lines). While functional, a named function would improve readability and follow PEP 8 style guidelines.

Apply this diff to refactor:

+    def run_moe():
+        return trtllm_fp8_per_tensor_scale_moe(
+            routing_logits,
+            None,  # routing_bias
+            hidden_states,
+            w13,
+            output1_scale_scalar,
+            output1_scales_gate_scalar,
+            w2,
+            output2_scale_scalar,
+            num_experts,
+            top_k,
+            None,  # n_group
+            None,  # topk_group
+            intermediate_size,
+            0,  # local_expert_offset
+            num_experts,
+            1.0,  # routed_scaling_factor
+            False,  # use_routing_scales_on_input
+            None,
+            RoutingMethodType.TopK.value,
+            enable_pdl,
+            tune_max_num_tokens
+        )
+
-    fn = lambda: trtllm_fp8_per_tensor_scale_moe(
-        routing_logits,
-        None,  # routing_bias
-        hidden_states,
-        w13,
-        output1_scale_scalar,
-        output1_scales_gate_scalar,
-        w2,
-        output2_scale_scalar,
-        num_experts,
-        top_k,
-        None,  # n_group
-        None,  # topk_group
-        intermediate_size,
-        0,  # local_expert_offset
-        num_experts,
-        1.0,  # routed_scaling_factor
-        False,  # use_routing_scales_on_input
-        None,
-        RoutingMethodType.TopK.value,
-        enable_pdl,
-        tune_max_num_tokens
-    )
+    fn = run_moe

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

benchmarks/bench_trtllm_gen_fused_moe_autotuner.py (1)

23-27: Guard FP8 quantization against zero max.

Line 23 lets scale = FLOAT8_E4M3_MAX / 0 when the tensor is all zeros, so x * scale turns into NaNs and the benchmark fails immediately. Bail out with zeros and a unit inverse-scale before dividing.

Apply this diff to harden the quantizer:

 def fp8_quantize(x):
-    max = x.float().abs().nan_to_num().max()
-    scale = FLOAT8_E4M3_MAX / max
-    x = (x * scale).to(torch.float8_e4m3fn)
-    return x, 1.0 / scale
+    max_val = x.float().abs().nan_to_num().amax()
+    if max_val.item() == 0:
+        zeros = torch.zeros_like(x, dtype=torch.float8_e4m3fn)
+        return zeros, torch.tensor(1.0, device=x.device, dtype=torch.float32)
+    scale = FLOAT8_E4M3_MAX / max_val
+    quantized = (x * scale).to(torch.float8_e4m3fn)
+    return quantized, scale.reciprocal()

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

include/flashinfer/trtllm/fused_moe/DevKernel.h (1)

116-123: Critical: Fix the pow2 condition to include zero

This issue was previously flagged: the condition data.mPaddingLog2 > 0 excludes mPaddingLog2 == 0, causing it to fall into the tile-N branch (false). However, mPaddingLog2 == 0 represents tile size 1 (2^0 = 1), which is a power of 2 and should follow the pow2 path. When zero takes the false branch, mTileTokensDim is not populated for pow2 flows, leading to divide-by-zero runtime faults in the CUDA kernels.

Apply this diff:

-  if (data.mPaddingLog2 > 0) {
+  if (data.mPaddingLog2 >= 0) {

csrc/trtllm_fused_moe_kernel_launcher.cu (1)

66-81: Fix UB when tile not present; snap to nearest supported entry before walking neighbors.

std::find may return end() (e.g., nextPowerOfTwo=32 but 32 not in list). Using std::next(it) on end() is UB. Replace with lower_bound, clamp iterator, and snap tile_tokens_dim to an actual element; also guard neighbors.

-  // assume supported_tile_nums is sorted
-  int32_t tile_tokens_dim = std::clamp(nextPowerOfTwo(avg_tokens_per_expert),
-                                       supported_tile_nums.front(), supported_tile_nums.back());
-  auto it = std::find(supported_tile_nums.begin(), supported_tile_nums.end(), tile_tokens_dim);
+  // assume supported_tile_nums is sorted
+  int32_t tile_tokens_dim = std::clamp(nextPowerOfTwo(avg_tokens_per_expert),
+                                       supported_tile_nums.front(), supported_tile_nums.back());
+  auto it = std::lower_bound(supported_tile_nums.begin(), supported_tile_nums.end(), tile_tokens_dim);
+  if (it == supported_tile_nums.end()) { it = std::prev(supported_tile_nums.end()); }
+  // Optionally pick the closer neighbor if lower_bound overshoots
+  if (it != supported_tile_nums.begin()) {
+    auto prev = std::prev(it);
+    if (std::abs(*prev - tile_tokens_dim) <= std::abs(*it - tile_tokens_dim)) it = prev;
+  }
+  tile_tokens_dim = *it;
 
   std::set<int32_t> selected_tile_nums;
   selected_tile_nums.insert(tile_tokens_dim);
-  if (std::next(it) != supported_tile_nums.end()) {
+  if (std::next(it) != supported_tile_nums.end()) {
     selected_tile_nums.insert(*std::next(it));
-    if (std::next(std::next(it)) != supported_tile_nums.end()) {
+    if (std::next(std::next(it)) != supported_tile_nums.end()) {
       selected_tile_nums.insert(*std::next(std::next(it)));
     }
   }
-  if (it != supported_tile_nums.begin()) {
+  if (it != supported_tile_nums.begin()) {
     selected_tile_nums.insert(*std::prev(it));
   }

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h (1)

48-54: Remove nested trtllm::gen inside batchedGemm; it shadows real ::trtllm::gen.

This creates batchedGemm::trtllm::gen and breaks lookups for the actual ::trtllm::gen (dtype helpers, enums). Drop the nested block and reference global with ::.

-namespace trtllm {
-namespace gen {
-class CudaRunner;
-class GenCfg;
-}  // namespace gen
-}  // namespace trtllm
+// Forward declarations, when needed, should reference global namespace or use fully qualified ::trtllm::gen.
+// (Remove nested shadow to keep tg = ::trtllm::gen valid.)

And make the alias explicit:

-namespace tg = trtllm::gen;
+namespace tg = ::trtllm::gen;

Also applies to: 59-60

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h (2)

58-64: Avoid nested trtllm::gen in batchedGemm; use global to prevent shadowing.

Same shadowing issue as in GemmGatedActOptions.h. Remove the nested block and use ::trtllm::gen in aliases.

-namespace trtllm {
-namespace gen {
-class CudaRunner;
-class GenCfg;
-}  // namespace gen
-}  // namespace trtllm
+// Use ::trtllm::gen forward decls if needed; avoid introducing batchedGemm::trtllm::gen.

-namespace tg = trtllm::gen;
+namespace tg = ::trtllm::gen;

Also applies to: 69-71

---

388-415: Restore printing of mBatchedM/mBatchedN in dumpOptions.

Placeholders {} drop essential tuning logs. Serialize the actual vectors when dumpRuntimeParams is true.

-  if (dumpRuntimeParams) {
-    ss << "mBatchedM={}," << std::endl;
-    ss << "mBatchedN={}," << std::endl;
-  }
+  if (dumpRuntimeParams) {
+    ss << "mBatchedM={";
+    for (size_t i = 0; i < options.mBatchedM.size(); ++i) {
+      ss << (i ? ", " : "") << options.mBatchedM[i];
+    }
+    ss << "}," << std::endl;
+    ss << "mBatchedN={";
+    for (size_t i = 0; i < options.mBatchedN.size(); ++i) {
+      ss << (i ? ", " : "") << options.mBatchedN[i];
+    }
+    ss << "}," << std::endl;
+  }

benchmarks/bench_trtllm_gen_fused_moe_autotuner.py (3)

30-40: Remove or document the unused quant_mode parameter.

The quant_mode parameter at line 32 is not used within the function body.

If this parameter is reserved for future use, add a comment explaining its purpose. Otherwise, remove it:

 def bench_trtllm_gen_fused_moe_autotuner_fp8(
     tune_max_num_tokens: Optional[int],
-    quant_mode: Literal["Fp8-Per-Tensor"],
     num_tokens: int,
     num_experts: int,
     hidden_size: int,
     intermediate_size: int,
     top_k: int,
     warmups: int,
     iterations: int,
 ):

---

23-27: Guard FP8 quantization against all-zero inputs.

When the input tensor is all zeros, max becomes zero, scale becomes infinite, and the quantized result is NaN (0 × ∞). This will break benchmarking when buffers start cleared.

Apply this diff to handle the zero case:

 def fp8_quantize(x):
-    max = x.float().abs().nan_to_num().max()
-    scale = FLOAT8_E4M3_MAX / max
-    x = (x * scale).to(torch.float8_e4m3fn)
-    return x, 1.0 / scale
+    max_val = x.float().abs().nan_to_num().max()
+    if max_val == 0:
+        return torch.zeros_like(x, dtype=torch.float8_e4m3fn), 1.0
+    scale = FLOAT8_E4M3_MAX / max_val
+    quantized = (x * scale).to(torch.float8_e4m3fn)
+    return quantized, 1.0 / scale

---

60-68: Construct FP8 scale vectors without CPU conversion errors.

Lines 60-68 call torch.tensor([...], device=device) on CUDA scalars, which raises TypeError: can't convert CUDA tensor to numpy(). Build the vectors on device without Python lists.

Apply this diff to keep scales on CUDA:

-    output1_scale_scalar = torch.tensor(
-        [hidden_states_scale * w13_scale] * num_experts, device=device
-    )
+    scale_prod_1 = (hidden_states_scale * w13_scale).item()
+    output1_scale_scalar = torch.full(
+        (num_experts,),
+        scale_prod_1,
+        device=device,
+        dtype=torch.float32,
+    )
@@
-    output2_scale_scalar = torch.tensor(
-        [hidden_states_scale * w2_scale] * num_experts, device=device
-    )
+    scale_prod_2 = (hidden_states_scale * w2_scale).item()
+    output2_scale_scalar = torch.full(
+        (num_experts,),
+        scale_prod_2,
+        device=device,
+        dtype=torch.float32,
+    )

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h (1)

745-787: Fix variable shadowing that loses runtime dimensions in workspace calculation.

Line 767 creates a new reference auto& options = config.mOptions that shadows the options variable from line 750. This causes:

• Lines 752-764 to use complete options (with runtime M/N/K/valid fields from data)
• Lines 769-779 to use only config.mOptions (missing runtime dimensions)

The shadowing loses the populated mM, mN, mK, mValidM, mValidN, mValidK fields, potentially leading to incorrect workspace size calculations.

Apply this diff to remove the shadowing:

-    // Get options from config.
-    auto& options = config.mOptions;
+    // Use the same options variable created earlier (no shadowing)

Then use the existing options variable for all subsequent calculations.

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h (1)

652-665: Fix warning macro that prevents valid dimension clamping.

When options.mValidK > options.mK, the TLLM_LOG_WARNING at line 655 fires. Under TLLM_GEN_EXPORT_INTERFACE, this macro immediately returns false, so the subsequent clamp operation at lines 659-661 never executes. This breaks callers relying on updateOptions == true to sanitize oversized valid dimensions.

Replace the warning with a conditional log that doesn't trigger early return:

-  if (options.mValidM > options.mM || options.mValidN > options.mN ||
-      options.mValidK > options.mK) {
-    TLLM_LOG_WARNING(
-        options.mValidK <= options.mK,
-        "ValidM, ValidN, and ValidK must be less than or equal to M, N, and K respectively.");
-    if (updateOptions) {
-      options.mValidM = std::min(options.mValidM, options.mM);
-      options.mValidN = std::min(options.mValidN, options.mN);
-      options.mValidK = std::min(options.mValidK, options.mK);
-    } else {
-      return false;
-    }
-  }
+  if (options.mValidM > options.mM || options.mValidN > options.mN ||
+      options.mValidK > options.mK) {
+#ifdef TLLM_GEN_DEBUG
+    printArgs("WARNING: ValidM, ValidN, and ValidK must be less than or equal to M, N, and K respectively.\n");
+#endif
+    if (updateOptions) {
+      options.mValidM = std::min(options.mValidM, options.mM);
+      options.mValidN = std::min(options.mValidN, options.mN);
+      options.mValidK = std::min(options.mValidK, options.mK);
+    } else {
+      return false;
+    }
+  }

csrc/trtllm_fused_moe_kernel_launcher.cu (1)

381-392: Tile sets look good; ensure sorted+unique before selection to avoid duplicates.

When conditionally pushing 128/192/256, duplicates can appear across branches; keep the vector sorted and unique before passing to computeSelectedTileN.

-    std::vector<int32_t> mSupportedTileN = {8, 16, 32, 64, 128, 192, 256};
+    std::vector<int32_t> mSupportedTileN = {8, 16, 32, 64, 128, 192, 256};
+    mSupportedTileN.erase(std::unique(mSupportedTileN.begin(), mSupportedTileN.end()), mSupportedTileN.end());

-    std::vector<int32_t> mSupportedTileN = {8, 16, 32, 64, 128};
+    std::vector<int32_t> mSupportedTileN = {8, 16, 32, 64, 128};
+    mSupportedTileN.erase(std::unique(mSupportedTileN.begin(), mSupportedTileN.end()), mSupportedTileN.end());

-  std::vector<int32_t> supported_tile_nums = {8, 16, 32, 64};
+  std::vector<int32_t> supported_tile_nums = {8, 16, 32, 64};
   ...
   // after conditional push_backs
+  std::sort(supported_tile_nums.begin(), supported_tile_nums.end());
+  supported_tile_nums.erase(std::unique(supported_tile_nums.begin(), supported_tile_nums.end()),
+                            supported_tile_nums.end());

Also applies to: 730-740, 1322-1336

flashinfer/fused_moe/core.py (1)

125-133: Minor nit: duplicate enum in list.

MxE4m3 appears twice in trtllm_gen_dtype_has_scale. Harmless; remove duplicate for clarity.

-        DtypeTrtllmGen.MxE4m3,
         DtypeTrtllmGen.E2m1,
         DtypeTrtllmGen.MxE2m1,
-        DtypeTrtllmGen.MxE4m3,

include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h (1)

154-169: New validation is good; message tweaks optional.

DeepSeek FP8 and shuffled-MatrixA constraints on hidden/valid sizes are correct. Consider clarifying error text to include hiddenSizeStr for both fields.

include/flashinfer/trtllm/fused_moe/RoutingKernel.h (1)

53-57: Template split (isPow2_, UsePdl_) and new mTileTokensDim are wired correctly.

Defaults (mPaddingLog2=-1) are safe under non-pow2 path; Data->Params copy includes mTileTokensDim. Comment on mPtrPermutedIdxToTokenIdx shape still mentions a derived formula and may confuse readers.

-  // dim: [mTileTokensDim * mTopK + (mNumExperts × mTileTokensDim) - mNumExperts]
+  // dim: [total_permuted_tokens]; actual size equals padded tokens count reported via mPtrPermutedIdxSize[0]

Also applies to: 95-103, 104-157, 232-254, 274-301

benchmarks/bench_trtllm_gen_fused_moe_autotuner.py (1)

70-92: Prefer def over lambda assignment for clarity.

Assigning a lambda expression to a variable reduces readability and makes debugging harder (stack traces show <lambda> instead of a meaningful function name).

Refactor to use a proper function definition:

-    fn = lambda: trtllm_fp8_per_tensor_scale_moe(
-        routing_logits,
-        None,  # routing_bias
-        hidden_states,
-        w13,
-        output1_scale_scalar,
-        output1_scales_gate_scalar,
-        w2,
-        output2_scale_scalar,
-        num_experts,
-        top_k,
-        None,  # n_group
-        None,  # topk_group
-        intermediate_size,
-        0,  # local_expert_offset
-        num_experts,
-        1.0,  # routed_scaling_factor
-        False,  # use_routing_scales_on_input
-        None,
-        RoutingMethodType.TopK.value,
-        enable_pdl,
-        num_tokens if tune_max_num_tokens is None else tune_max_num_tokens,
-    )
+    def fn():
+        return trtllm_fp8_per_tensor_scale_moe(
+            routing_logits,
+            None,  # routing_bias
+            hidden_states,
+            w13,
+            output1_scale_scalar,
+            output1_scales_gate_scalar,
+            w2,
+            output2_scale_scalar,
+            num_experts,
+            top_k,
+            None,  # n_group
+            None,  # topk_group
+            intermediate_size,
+            0,  # local_expert_offset
+            num_experts,
+            1.0,  # routed_scaling_factor
+            False,  # use_routing_scales_on_input
+            None,
+            RoutingMethodType.TopK.value,
+            enable_pdl,
+            num_tokens if tune_max_num_tokens is None else tune_max_num_tokens,
+        )