fix(moe): bound routing_replay_out dim0 from below in both validators

csrc/trtllm_fused_moe_kernel_launcher.cu

@@ -949,8 +949,12 @@ void cast_fp32_to_bf16(void* output, void const* input, int64_t num_elements, cu
 }  // namespace
 
 // Validate routing_replay_out tensor properties.
-// NOTE: dim0 >= num_tokens is intentionally NOT checked — with CUDA graphs the buffer
-// is pre-allocated at maximum batch size and reused across steps with varying num_tokens.
+// NOTE: dim0 is only bounded from below. The routing kernels write one replay row per
+// token unconditionally (DeepSeek launches numBlocks == num_tokens and writes row
+// blockIdx.x; the custom and llama4 kernels write row tokenIdx), so a buffer with fewer
+// rows than tokens is written past its end. Oversized buffers stay legal: with CUDA
+// graphs the buffer is pre-allocated at maximum batch size and reused across steps with
+// varying num_tokens.
 static void validate_routing_replay_out(TensorView const& replay, TensorView const& hidden_states,
                                         int64_t top_k) {
   TVM_FFI_ICHECK(replay.device().device_type == kDLCUDA)
@@ -959,6 +963,9 @@ static void validate_routing_replay_out(TensorView const& replay, TensorView con
       << "routing_replay_out must be on the same device as hidden_states";
   TVM_FFI_ICHECK(replay.ndim() == 2) << "routing_replay_out must be 2D [num_tokens, top_k]";
   TVM_FFI_ICHECK(replay.size(1) == top_k) << "routing_replay_out dim1 must equal top_k";
+  TVM_FFI_ICHECK(replay.size(0) >= hidden_states.size(0))
+      << "routing_replay_out dim0 must be >= num_tokens (" << hidden_states.size(0) << "), got "
+      << replay.size(0) << "; the routing kernel writes one replay row per token";
   TVM_FFI_ICHECK((replay.dtype() == DLDataType{kDLInt, 16, 1}))
       << "routing_replay_out must be int16 dtype";
   TVM_FFI_ICHECK(replay.IsContiguous())

flashinfer/fused_moe/core.py

@@ -5446,9 +5446,16 @@ def _validate_bf16_gemm1_activation_params(
 def _validate_routing_replay_out(
     routing_replay_out: Optional[torch.Tensor],
     top_k: int,
+    num_tokens: int,
     num_fused_shared_experts: int = 0,
 ) -> None:
-    """Validate routing_replay_out tensor properties before passing to C++ kernels."""
+    """Validate routing_replay_out tensor properties before passing to C++ kernels.
+
+    ``num_tokens`` bounds dim0 from below: the routing kernels write one replay row per
+    token unconditionally, so a shorter buffer is written past its end. Oversized buffers
+    stay legal for CUDA-graph capture at a fixed maximum batch size. It is required rather
+    than defaulted so that a new entry point cannot silently opt out of the bound.
+    """
     if routing_replay_out is None:
         return
     if num_fused_shared_experts > 0:
@@ -5468,6 +5475,12 @@ def _validate_routing_replay_out(
         raise ValueError(
             f"routing_replay_out dim1 must equal top_k={top_k}, got {routing_replay_out.shape[1]}"
         )
+    if routing_replay_out.shape[0] < num_tokens:
+        raise ValueError(
+            f"routing_replay_out dim0 must be >= num_tokens={num_tokens}, "
+            f"got {routing_replay_out.shape[0]}; the routing kernel writes one replay "
+            "row per token"
+        )
     if not routing_replay_out.is_contiguous():
         raise ValueError("routing_replay_out must be contiguous (packed row-major)")
 
@@ -5650,7 +5663,9 @@ def trtllm_bf16_moe(
         scalar return; will become ``[output]`` in v0.8.0).  Otherwise returns
         ``[gemm2_output, expert_weights, expanded_idx_to_permuted_idx]``.
     """
-    _validate_routing_replay_out(routing_replay_out, top_k)
+    _validate_routing_replay_out(
+        routing_replay_out, top_k, num_tokens=hidden_states.shape[0]
+    )
     _validate_bf16_gemm1_activation_params(
         activation_type,
         gemm1_alpha,
@@ -5867,7 +5882,9 @@ def trtllm_bf16_routed_moe(
         ``False``      ``Tensor``          ``[gemm2_output, expert_weights, expanded_idx_to_permuted_idx, gemm1_activation_output]``
         =============  ==================  =========================================================================
     """
-    _validate_routing_replay_out(routing_replay_out, top_k)
+    _validate_routing_replay_out(
+        routing_replay_out, top_k, num_tokens=hidden_states.shape[0]
+    )
     _validate_bf16_gemm1_activation_params(
         activation_type,
         gemm1_alpha,
@@ -6040,7 +6057,9 @@ def trtllm_fp8_per_tensor_scale_moe(
         Final MoE output when ``do_finalize`` is ``True``, otherwise
         ``[gemm2_output, expert_weights, expanded_idx_to_permuted_idx]``.
     """
-    _validate_routing_replay_out(routing_replay_out, top_k)
+    _validate_routing_replay_out(
+        routing_replay_out, top_k, num_tokens=hidden_states.shape[0]
+    )
     result = get_trtllm_moe_sm100_module().trtllm_fp8_per_tensor_scale_moe(
         routing_logits,
         routing_bias,
@@ -6179,7 +6198,9 @@ def trtllm_fp8_per_tensor_scale_routed_moe(
         Final MoE output when ``do_finalize`` is ``True``, otherwise
         ``[gemm2_output, expert_weights, expanded_idx_to_permuted_idx]``.
     """
-    _validate_routing_replay_out(routing_replay_out, top_k)
+    _validate_routing_replay_out(
+        routing_replay_out, top_k, num_tokens=hidden_states.shape[0]
+    )
     topk_ids_tensor, topk_weights, routing_mode = _split_precomputed_routing(topk_ids)
     result = get_trtllm_moe_sm100_module().trtllm_fp8_per_tensor_scale_routed_moe(
         routing_mode,
@@ -6663,7 +6684,12 @@ def trtllm_fp8_block_scale_moe(
             "Fused shared experts (num_fused_shared_experts > 0) are only supported "
             f"with DeepSeekV3 routing; got routing_method_type={routing_method_type}."
         )
-    _validate_routing_replay_out(routing_replay_out, top_k, nfse)
+    _validate_routing_replay_out(
+        routing_replay_out,
+        top_k,
+        num_tokens=hidden_states.shape[0],
+        num_fused_shared_experts=nfse,
+    )
     _validate_fp8_block_scale_gemm1_activation_params(
         fp8_quantization_type,
         activation_type,
@@ -7176,7 +7202,12 @@ def trtllm_fp4_block_scale_moe(
             "Fused shared experts (num_fused_shared_experts > 0) are only supported "
             f"with DeepSeekV3 routing; got routing_method_type={routing_method_type}."
         )
-    _validate_routing_replay_out(routing_replay_out, top_k, nsfe)
+    _validate_routing_replay_out(
+        routing_replay_out,
+        top_k,
+        num_tokens=hidden_states.shape[0],
+        num_fused_shared_experts=nsfe,
+    )
     return get_trtllm_moe_sm100_module().trtllm_fp4_block_scale_moe(
         RoutingInputMode.FromLogits,
         routing_logits,
@@ -7602,7 +7633,9 @@ def trtllm_mxint4_block_scale_moe(
         ``[output]`` when ``do_finalize`` is ``True``, otherwise
         ``[gemm2_output, expert_weights, expanded_idx_to_permuted_idx]``.
     """
-    _validate_routing_replay_out(routing_replay_out, top_k)
+    _validate_routing_replay_out(
+        routing_replay_out, top_k, num_tokens=hidden_states.shape[0]
+    )
     return get_trtllm_moe_sm100_module().trtllm_mxint4_block_scale_moe(
         routing_logits,
         routing_bias,

tests/moe/test_trtllm_gen_fused_moe.py

@@ -2471,6 +2471,71 @@ def test_fused_shared_experts_reject_replay_and_non_deepseek_routing():
                 )
 
 
+def test_routing_replay_out_rejects_undersized_dim0():
+    """An undersized replay buffer must be rejected, not written out of bounds.
+
+    Routing launches one block per token and writes row ``blockIdx.x``
+    unconditionally, so ``dim0 < num_tokens`` is a device-side buffer overflow
+    rather than a truncated result. Oversized buffers stay legal for CUDA-graph
+    pre-allocation, hence the ``>=`` bound. Host-side check, so no GPU needed.
+    """
+    num_experts = 4
+    num_tokens = 4
+    top_k = 1
+    fp8_kwargs = {
+        "routing_logits": torch.empty((num_tokens, num_experts), dtype=torch.bfloat16),
+        "routing_bias": None,
+        "hidden_states": torch.empty((num_tokens, 1), dtype=torch.bfloat16),
+        "hidden_states_scale": torch.empty((1, 1), dtype=torch.float32),
+        "gemm1_weights": torch.empty((1, 2, 1), dtype=torch.bfloat16),
+        "gemm1_weights_scale": torch.empty((1, 1, 1), dtype=torch.float32),
+        "gemm2_weights": torch.empty((1, 1, 1), dtype=torch.bfloat16),
+        "gemm2_weights_scale": torch.empty((1, 1, 1), dtype=torch.float32),
+    }
+    fp4_kwargs = {
+        "routing_logits": torch.empty((num_tokens, num_experts), dtype=torch.bfloat16),
+        "routing_bias": None,
+        "hidden_states": torch.empty((num_tokens, 2), dtype=torch.bfloat16),
+        "hidden_states_scale": None,
+        "gemm1_weights": torch.empty((1, 2, 1), dtype=torch.uint8),
+        "gemm1_weights_scale": torch.empty((1, 1, 1), dtype=torch.float8_e4m3fn),
+        "gemm1_bias": None,
+        "gemm1_alpha": None,
+        "gemm1_beta": None,
+        "gemm1_clamp_limit": None,
+        "gemm2_weights": torch.empty((1, 1, 1), dtype=torch.uint8),
+        "gemm2_weights_scale": torch.empty((1, 1, 1), dtype=torch.float8_e4m3fn),
+        "gemm2_bias": None,
+        "output1_scale_scalar": None,
+        "output1_scale_gate_scalar": None,
+        "output2_scale_scalar": None,
+    }
+    common_kwargs = {
+        "num_experts": num_experts,
+        "top_k": top_k,
+        "n_group": None,
+        "topk_group": None,
+        "intermediate_size": 1,
+        "local_expert_offset": 0,
+        "local_num_experts": num_experts,
+        "routed_scaling_factor": None,
+        "routing_method_type": RoutingMethodType.DeepSeekV3.value,
+    }
+
+    for op, op_kwargs in (
+        (trtllm_fp8_block_scale_moe, fp8_kwargs),
+        (trtllm_fp4_block_scale_moe, fp4_kwargs),
+    ):
+        with pytest.raises(ValueError, match=r"dim0 must be >= num_tokens"):
+            op(
+                **op_kwargs,
+                **common_kwargs,
+                routing_replay_out=torch.empty(
+                    (num_tokens - 1, top_k), dtype=torch.int16
+                ),
+            )
+
+
 def test_fp4_block_scale_moe_fused_shared_experts_reject_routed_only_tensors():
     """Routed-only expert-major tensors must fail host-side, not OOB on the GPU.