From f87df1df9e93042fcfbd939d5fab0b7b0f47abc8 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sat, 6 Jun 2026 14:14:41 +0800 Subject: [PATCH] [Bugfix][MoE] Snapshot max_cudagraph_capture_size into FusedMoEConfig (#44613) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 4.8 (1M context) --- vllm/model_executor/layers/fused_moe/config.py | 2 ++ .../layers/fused_moe/experts/flashinfer_cutlass_moe.py | 5 +---- .../layers/fused_moe/experts/trtllm_mxfp4_moe.py | 6 +----- vllm/model_executor/layers/fused_moe/layer.py | 1 + vllm/model_executor/layers/quantization/mxfp4.py | 9 ++------- 5 files changed, 7 insertions(+), 16 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 430947235e..0a070f77f8 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1266,6 +1266,8 @@ class FusedMoEConfig: # cannot silently select one and drop the clamp. swiglu_limit: float | None = None + max_capture_size: int = 0 + def __post_init__(self): if self.dp_size > 1: logger.debug_once( diff --git a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py index fd9446c2a2..ff259c828f 100644 --- a/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py @@ -4,7 +4,6 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -94,9 +93,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular): # - pass per-block weight scales to the kernel # - skip input activation quantization (kernel applies scaling) self.use_deepseek_fp8_block_scale = quant_config.is_block_quantized - self.max_capture_size = ( - get_current_vllm_config().compilation_config.max_cudagraph_capture_size - ) + self.max_capture_size = moe_config.max_capture_size self.gemm1_clamp_limit: torch.Tensor | None = None if quant_config.gemm1_clamp_limit is not None: self.gemm1_clamp_limit = torch.tensor( diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py index 43f800343c..2b8a5529ab 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py @@ -79,11 +79,7 @@ class TrtLlmMxfp4ExpertsBase: else: self.gemm1_clamp_limit = None - from vllm.config import get_current_vllm_config - - self.max_capture_size = ( - get_current_vllm_config().compilation_config.max_cudagraph_capture_size - ) + self.max_capture_size = moe_config.max_capture_size @staticmethod def _supports_current_device() -> bool: diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 2f4401563b..d1d6a4f03c 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -326,6 +326,7 @@ class FusedMoE(PluggableLayer): moe_backend=vllm_config.kernel_config.moe_backend, router_logits_dtype=router_logits_dtype, max_num_tokens=max_num_batched_tokens, + max_capture_size=compilation_config.max_cudagraph_capture_size, has_bias=has_bias, is_act_and_mul=is_act_and_mul, is_lora_enabled=vllm_config.lora_config is not None, diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 2c69fc7453..65714b4c9a 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -3,7 +3,6 @@ import torch -from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.fused_moe import ( @@ -141,9 +140,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase): self.weight_dtype = "gpt_oss_mxfp4" self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe) - self.max_capture_size = ( - get_current_vllm_config().compilation_config.max_cudagraph_capture_size - ) + self.max_capture_size = moe.max_capture_size self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {} self.moe_kernel: mk.FusedMoEKernel | None = None @@ -475,9 +472,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self.weight_dtype = "mxfp4" self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe) - self.max_capture_size = ( - get_current_vllm_config().compilation_config.max_cudagraph_capture_size - ) + self.max_capture_size = moe.max_capture_size self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {} self.moe_kernel: mk.FusedMoEKernel | None = None