diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index a45caac7c9e..cb69d85c43f 100644 --- a/vllm/_custom_ops.py +++ b/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, diff --git a/vllm/envs.py b/vllm/envs.py index d6240df3605..ced501b2052 100755 --- a/vllm/envs.py +++ b/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: diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 2500387debe..d2eb45fc780 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/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 diff --git a/vllm/model_executor/layers/fused_moe/cutlass_moe.py b/vllm/model_executor/layers/fused_moe/cutlass_moe.py index 75ee776646b..67fde4d7106 100644 --- a/vllm/model_executor/layers/fused_moe/cutlass_moe.py +++ b/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, ) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 75283b9bbe3..3c99b0600b1 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/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,