A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 20:00 UTC
[Bugfix][MoE] Snapshot max_cudagraph_capture_size into FusedMoEConfig (#44613)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aoshen02 committed
Jun 6, 2026 at 14:14 UTC
f87df1df9e93042fcfbd939d5fab0b7b0f47abc8
5 files changed
+7
-16
vllm/model_executor/layers/fused_moe/config.py
+2
@@ -1266,6 +1266,8 @@ class FusedMoEConfig:
1266
# cannot silently select one and drop the clamp.
1267
swiglu_limit: float | None = None
1268
1269
+ max_capture_size: int = 0
1270
+
1271
def __post_init__(self):
1272
if self.dp_size > 1:
1273
logger.debug_once(
vllm/model_executor/layers/fused_moe/experts/flashinfer_cutlass_moe.py
+1
-4
@@ -4,7 +4,6 @@
4
import torch
5
6
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
-from vllm.config import get_current_vllm_config
7
from vllm.logger import init_logger
8
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
9
from vllm.model_executor.layers.fused_moe.config import (
@@ -94,9 +93,7 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
93
# - pass per-block weight scales to the kernel
94
# - skip input activation quantization (kernel applies scaling)
95
self.use_deepseek_fp8_block_scale = quant_config.is_block_quantized
97
- self.max_capture_size = (
98
- get_current_vllm_config().compilation_config.max_cudagraph_capture_size
99
- )
96
+ self.max_capture_size = moe_config.max_capture_size
97
self.gemm1_clamp_limit: torch.Tensor | None = None
98
if quant_config.gemm1_clamp_limit is not None:
99
self.gemm1_clamp_limit = torch.tensor(
vllm/model_executor/layers/fused_moe/experts/trtllm_mxfp4_moe.py
+1
-5
@@ -79,11 +79,7 @@ class TrtLlmMxfp4ExpertsBase:
79
else:
80
self.gemm1_clamp_limit = None
81
82
- from vllm.config import get_current_vllm_config
83
-
84
- self.max_capture_size = (
85
- get_current_vllm_config().compilation_config.max_cudagraph_capture_size
86
- )
82
+ self.max_capture_size = moe_config.max_capture_size
83
84
@staticmethod
85
def _supports_current_device() -> bool:
vllm/model_executor/layers/fused_moe/layer.py
+1
@@ -326,6 +326,7 @@ class FusedMoE(PluggableLayer):
326
moe_backend=vllm_config.kernel_config.moe_backend,
327
router_logits_dtype=router_logits_dtype,
328
max_num_tokens=max_num_batched_tokens,
329
+ max_capture_size=compilation_config.max_cudagraph_capture_size,
330
has_bias=has_bias,
331
is_act_and_mul=is_act_and_mul,
332
is_lora_enabled=vllm_config.lora_config is not None,
vllm/model_executor/layers/quantization/mxfp4.py
+2
-7
@@ -3,7 +3,6 @@
3
4
import torch
5
6
-from vllm.config import get_current_vllm_config
6
from vllm.logger import init_logger
7
from vllm.model_executor.layers.attention import Attention
8
from vllm.model_executor.layers.fused_moe import (
@@ -141,9 +140,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):
140
self.weight_dtype = "gpt_oss_mxfp4"
141
self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
142
144
- self.max_capture_size = (
145
- get_current_vllm_config().compilation_config.max_cudagraph_capture_size
146
- )
143
+ self.max_capture_size = moe.max_capture_size
144
145
self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
146
self.moe_kernel: mk.FusedMoEKernel | None = None
@@ -475,9 +472,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
472
self.weight_dtype = "mxfp4"
473
self.mxfp4_backend, self.experts_cls = select_deepseek_v4_mxfp4_moe_backend(moe)
474
478
- self.max_capture_size = (
479
- get_current_vllm_config().compilation_config.max_cudagraph_capture_size
480
- )
475
+ self.max_capture_size = moe.max_capture_size
476
477
self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
478
self.moe_kernel: mk.FusedMoEKernel | None = None