Replace VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE env var with max_num_batched_tokens

The NVFP4 MoE blockscale buffer was sized using a hardcoded env var
defaulting to 163840, which is unnecessarily large. Replace it with
max_num_batched_tokens from the scheduler config, using the formula:
  buffer = mnbt * topk + 127 * num_experts
to account for CUTLASS 128-row alignment padding per expert.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Signed-off-by: mgoin <mgoin64@gmail.com>

vllm/_custom_ops.py

@@ -1758,6 +1758,7 @@ def scaled_fp4_experts_quant(
     expert_offsets: torch.Tensor,
     blockscale_offsets: torch.Tensor,
     topk: int,
+    max_num_batched_tokens: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Quantize input tensor to NVFP4 and return quantized tensor and scale, for
@@ -1767,6 +1768,9 @@ def scaled_fp4_experts_quant(
         input_global_scale: A scalar scaling factor for the entire tensor.
         expert_offsets: The expert offsets tensor
         blockscale_offsets: The blockscale offsets tensor
+        topk: Number of top-k experts selected
+        max_num_batched_tokens: Maximum tokens per scheduler iteration,
+            used to size the blockscale output buffer.
     Outputs:
         output: The quantized tensor in NVFP4
         output_scales: The blockscale tensor in FP8-E4M3
@@ -1776,18 +1780,19 @@ def scaled_fp4_experts_quant(
         f"input.ndim needs to be == 2, but got {input_tensor.ndim}."
     )
 
-    # Control the maximum number of tokens per expert supported by the
-    # NVFP4 MoE Expert Quantization. This is used to prevent the kernel
-    # from running out of memory. This value can also be increased to support
-    # larger models.
-    MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
     m_numtopk, k = input_tensor.shape
+    num_experts = expert_offsets.shape[0] - 1
+    # Worst-case padding: each expert's blockscale rows are rounded up to a
+    # multiple of 128 (CUTLASS swizzle constraint). If every expert has a
+    # remainder of 1 token, that's 127 padding tokens per expert.
+    max_buffer_rows = max_num_batched_tokens * topk + 127 * num_experts
 
-    assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
-        f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
-        f"{MAX_TOKENS_PER_EXPERT})"
-        f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use"
-        f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
+    assert m_numtopk <= max_num_batched_tokens * topk, (
+        f"m_numtopk must be less than max_num_batched_tokens * topk ("
+        f"{max_num_batched_tokens} * {topk} = "
+        f"{max_num_batched_tokens * topk})"
+        f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}."
+        f" Use --max-num-batched-tokens to increase this value."
     )
     scales_k = k // 16
     padded_k = (scales_k + (4 - 1)) // 4
@@ -1797,7 +1802,7 @@ def scaled_fp4_experts_quant(
         m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
     )
     output_scales = torch.empty(
-        MAX_TOKENS_PER_EXPERT * topk,
+        max_buffer_rows,
         padded_k,
         dtype=torch.int32,
         device=input_tensor.device,
@@ -1820,6 +1825,7 @@ def silu_and_mul_scaled_fp4_experts_quant(
     expert_offsets: torch.Tensor,
     blockscale_offsets: torch.Tensor,
     topk: int,
+    max_num_batched_tokens: int,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations.
@@ -1830,6 +1836,8 @@ def silu_and_mul_scaled_fp4_experts_quant(
         expert_offsets: The expert offsets tensor [n_experts+1]
         blockscale_offsets: The blockscale offsets tensor [n_experts+1]
         topk: Number of top-k experts selected
+        max_num_batched_tokens: Maximum tokens per scheduler iteration,
+            used to size the blockscale output buffer.
     Outputs:
         output: The quantized tensor in NVFP4 [m_topk, k/2]
         output_scales: The blockscale tensor in FP8-E4M3
@@ -1839,20 +1847,21 @@ def silu_and_mul_scaled_fp4_experts_quant(
         f"input.ndim needs to be == 2, but got {input_tensor.ndim}."
     )
 
-    # Control the maximum number of tokens per expert supported by the
-    # NVFP4 MoE Expert Quantization. This is used to prevent the kernel
-    # from running out of memory. This value can also be increased to support
-    # larger models.
-    MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
     m_numtopk, k_times_2 = input_tensor.shape
     assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
     k = k_times_2 // 2
+    num_experts = expert_offsets.shape[0] - 1
+    # Worst-case padding: each expert's blockscale rows are rounded up to a
+    # multiple of 128 (CUTLASS swizzle constraint). If every expert has a
+    # remainder of 1 token, that's 127 padding tokens per expert.
+    max_buffer_rows = max_num_batched_tokens * topk + 127 * num_experts
 
-    assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
-        f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
-        f"{MAX_TOKENS_PER_EXPERT})"
-        f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use"
-        f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
+    assert m_numtopk <= max_num_batched_tokens * topk, (
+        f"m_numtopk must be less than max_num_batched_tokens * topk ("
+        f"{max_num_batched_tokens} * {topk} = "
+        f"{max_num_batched_tokens * topk})"
+        f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}."
+        f" Use --max-num-batched-tokens to increase this value."
     )
     scales_k = k // 16
     padded_k = (scales_k + (4 - 1)) // 4
@@ -1862,7 +1871,7 @@ def silu_and_mul_scaled_fp4_experts_quant(
         m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
     )
     output_scales = torch.empty(
-        MAX_TOKENS_PER_EXPERT * topk,
+        max_buffer_rows,
         padded_k,
         dtype=torch.int32,
         device=input_tensor.device,

vllm/envs.py

@@ -177,7 +177,7 @@ if TYPE_CHECKING:
     VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
     VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600
     VLLM_MOONCAKE_BOOTSTRAP_PORT: int = 8998
-    VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840
+
     VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1
     VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16
     VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300
@@ -1316,13 +1316,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE": lambda: int(
         os.getenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(394 * 1024 * 1024))
     ),
-    # Control the maximum number of tokens per expert supported by the
-    # NVFP4 MoE CUTLASS Kernel. This value is used to create a buffer for
-    # the blockscale tensor of activations NVFP4 Quantization.
-    # This is used to prevent the kernel from running out of memory.
-    "VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE": lambda: int(
-        os.getenv("VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE", "163840")
-    ),
     # Specifies the thresholds of the communicated tensor sizes under which
     # vllm should use flashinfer fused allreduce. The variable should be a
     # JSON with the following format:

vllm/model_executor/layers/fused_moe/config.py

@@ -1170,6 +1170,9 @@ class FusedMoEConfig:
 
     moe_backend: str = "auto"
     max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE
+    # Maximum number of tokens in a single scheduler iteration.
+    # Used to size NVFP4 MoE blockscale buffers.
+    max_num_batched_tokens: int = 2048
     has_bias: bool = False
     is_act_and_mul: bool = True
     is_lora_enabled: bool = False

vllm/model_executor/layers/fused_moe/cutlass_moe.py

@@ -499,6 +499,7 @@ def run_cutlass_moe_fp4(
     k: int,
     e: int,
     device: torch.device,
+    max_num_batched_tokens: int,
     apply_router_weight_on_input: bool = False,
 ) -> None:
     """
@@ -605,6 +606,7 @@ def run_cutlass_moe_fp4(
         expert_offsets,
         blockscale_offsets,
         num_topk,
+        max_num_batched_tokens,
     )
     c1 = _resize_cache(workspace13, (m * topk, w1_n))
     c2 = _resize_cache(workspace2, (m * topk, n))
@@ -626,12 +628,22 @@ def run_cutlass_moe_fp4(
         # Note: c2 workspace is no longer needed since SiLU is fused with quantization.
         # c3 reuses workspace13 after c1 is consumed.
         int_fp4, int_blockscale = ops.silu_and_mul_scaled_fp4_experts_quant(
-            c1, a2_gscale, expert_offsets, blockscale_offsets, num_topk
+            c1,
+            a2_gscale,
+            expert_offsets,
+            blockscale_offsets,
+            num_topk,
+            max_num_batched_tokens,
         )
     else:
         apply_moe_activation(activation, c2, c1)
         int_fp4, int_blockscale = ops.scaled_fp4_experts_quant(
-            c2, a2_gscale, expert_offsets, blockscale_offsets, num_topk
+            c2,
+            a2_gscale,
+            expert_offsets,
+            blockscale_offsets,
+            num_topk,
+            max_num_batched_tokens,
         )
 
     ops.cutlass_fp4_moe_mm(
@@ -791,6 +803,7 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
             k=k,
             e=e,
             device=hidden_states.device,
+            max_num_batched_tokens=self.moe_config.max_num_batched_tokens,
             apply_router_weight_on_input=apply_router_weight_on_input,
         )
 

vllm/model_executor/layers/fused_moe/layer.py

@@ -556,6 +556,7 @@ class FusedMoE(CustomOp):
             moe_backend=vllm_config.kernel_config.moe_backend,
             router_logits_dtype=router_logits_dtype,
             max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
+            max_num_batched_tokens=vllm_config.scheduler_config.max_num_batched_tokens,
             has_bias=has_bias,
             is_act_and_mul=is_act_and_mul,
             is_lora_enabled=vllm_config.lora_config is not None,