[Bugfix] Pick the DeepSeek V4 eager cudagraph region per model runner (#52401)
Nick Hill committed
Aug 15, 2026 at 21:04 UTC
8efa13b700f1836657699cae2503dc2feab27fa0
3 files changed
+96
-66
tests/test_config.py
+24
-37
@@ -66,41 +66,18 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
66
assert envs.VLLM_USE_V2_MODEL_RUNNER is expected
67
68
69
-@pytest.mark.parametrize(
70
- "cudagraph_mode",
71
- [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE],
72
-)
73
-def test_deepseek_v4_rejects_mrv1_piecewise_cudagraph(cudagraph_mode):
74
- config = SimpleNamespace(
75
- use_v2_model_runner=False,
76
- model_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]),
77
- compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode),
78
- )
79
-
80
- with pytest.raises(ValueError, match="DeepSeek V4 does not support PIECEWISE"):
81
- VllmConfig._validate_mrv1_piecewise_cudagraph(config)
82
-
83
-
84
-@pytest.mark.parametrize(
85
- ("use_v2_model_runner", "architecture", "cudagraph_mode"),
86
- [
87
- (True, "DeepseekV4ForCausalLM", CUDAGraphMode.PIECEWISE),
88
- (False, "DeepseekV4ForCausalLM", CUDAGraphMode.NONE),
89
- (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL),
90
- (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL_DECODE_ONLY),
91
- (False, "LlamaForCausalLM", CUDAGraphMode.PIECEWISE),
92
- ],
93
-)
94
-def test_mrv1_piecewise_cudagraph_allowed(
95
- use_v2_model_runner, architecture, cudagraph_mode
96
-):
97
- config = SimpleNamespace(
98
- use_v2_model_runner=use_v2_model_runner,
99
- model_config=SimpleNamespace(architectures=[architecture]),
100
- compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode),
101
- )
102
-
103
- VllmConfig._validate_mrv1_piecewise_cudagraph(config)
69
+def test_rocm_defaults_deepseek_v4_to_mrv1(monkeypatch):
70
+ """ROCm keeps DeepSeek V4 on MRV1, which is still faster there."""
71
+ from vllm.config.vllm import default_v2_model_runner_architectures
72
+ from vllm.platforms import current_platform
73
+
74
+ monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
75
+ # The lookup is lru_cached against a fixed platform.
76
+ default_v2_model_runner_architectures.cache_clear()
77
+ try:
78
+ assert "DeepseekV4ForCausalLM" not in default_v2_model_runner_architectures()
79
+ finally:
80
+ default_v2_model_runner_architectures.cache_clear()
81
82
83
@pytest.mark.parametrize(
@@ -330,10 +307,20 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1(
307
),
308
],
309
)
333
-def test_is_default_v2_model_runner_model(model_config, expected):
310
+def test_is_default_v2_model_runner_model(model_config, expected, monkeypatch):
311
+ from vllm.config.vllm import default_v2_model_runner_architectures
312
+ from vllm.platforms import current_platform
313
+
314
+ # The expectations below are the platform-independent defaults; ROCm's
315
+ # DeepSeek V4 carve-out is covered by test_rocm_defaults_deepseek_v4_to_mrv1.
316
+ monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
317
+ default_v2_model_runner_architectures.cache_clear()
318
config = SimpleNamespace(model_config=model_config)
319
336
- assert VllmConfig._is_default_v2_model_runner_model(config) is expected
320
+ try:
321
+ assert VllmConfig._is_default_v2_model_runner_model(config) is expected
322
+ finally:
323
+ default_v2_model_runner_architectures.cache_clear()
324
325
326
@pytest.mark.skip_global_cleanup
vllm/config/vllm.py
+7
-25
@@ -66,10 +66,6 @@ else:
66
67
logger = init_logger(__name__)
68
69
-MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES = frozenset(
70
- {"DeepseekV4ForCausalLM"}
71
-)
72
-
69
DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
70
{
71
"DeepseekV2ForCausalLM",
@@ -87,6 +83,13 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
83
@lru_cache
84
def default_v2_model_runner_architectures() -> frozenset[str]:
85
"""Architectures defaulting to the V2 model runner on this platform."""
86
+ from vllm.platforms import current_platform
87
+
88
+ if current_platform.is_rocm():
89
+ # TODO(rocm): DeepSeek V4 is still faster on MRV1 on ROCm. The
90
+ # attention layer picks the eager cudagraph region MRV1 needs, so
91
+ # this is a perf default only; drop it once MRV2 catches up.
92
+ return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - {"DeepseekV4ForCausalLM"}
93
return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES
94
95
@@ -691,25 +694,6 @@ class VllmConfig:
694
return False
695
return is_default_v2_architecture or not model_config.is_moe
696
694
- def _validate_mrv1_piecewise_cudagraph(self) -> None:
695
- if self.use_v2_model_runner:
696
- return
697
- model_config = self.model_config
698
- if model_config is None:
699
- return
700
- if not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs():
701
- return
702
- architectures = getattr(model_config, "architectures", [])
703
- if any(
704
- arch in MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES
705
- for arch in architectures
706
- ):
707
- raise ValueError(
708
- "DeepSeek V4 does not support PIECEWISE CUDA graphs with "
709
- "Model Runner V1. Use Model Runner V2 or disable PIECEWISE "
710
- "CUDA graphs."
711
- )
712
-
697
@property
698
def needs_dp_coordinator(self) -> bool:
699
"""
@@ -1644,8 +1628,6 @@ class VllmConfig:
1628
"pipeline parallelism",
1629
)
1630
1647
- self._validate_mrv1_piecewise_cudagraph()
1648
-
1631
# final check of cudagraph mode after all possible updates
1632
if current_platform.is_cuda_alike():
1633
if (
vllm/models/deepseek_v4/attention.py
+65
-4
@@ -298,6 +298,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
298
eager_scratch_pool=eager_scratch_pool,
299
)
300
301
+ self._prepare_and_attn_fn = self._prepare_and_attn
302
+ if not vllm_config.use_v2_model_runner:
303
+ # MRV1's piecewise capture only tolerates the wide eager region: with
304
+ # the narrow one the attention input preparation stays in the captured
305
+ # graph and MRV1 produces garbage (#51430).
306
+ self._prepare_and_attn_fn = self._prepare_and_attn_eager
307
+
308
# Will be None on ROCm for now.
309
self.aux_stream_list = aux_stream_list
310
# [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
@@ -379,6 +386,64 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
386
self.eps,
387
)
388
389
+ self._prepare_and_attn_fn(
390
+ hidden_states,
391
+ qr,
392
+ kv,
393
+ kv_score,
394
+ indexer_kv_score,
395
+ indexer_weights,
396
+ positions,
397
+ o_padded,
398
+ )
399
+ o = o_padded[:, : self.n_local_heads, :]
400
+
401
+ # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
402
+ return self._o_proj(o, positions)
403
+
404
+ @eager_break_during_capture
405
+ def _prepare_and_attn_eager(
406
+ self,
407
+ hidden_states: torch.Tensor,
408
+ qr: torch.Tensor,
409
+ kv: torch.Tensor,
410
+ kv_score: torch.Tensor,
411
+ indexer_kv_score: torch.Tensor,
412
+ indexer_weights: torch.Tensor,
413
+ positions: torch.Tensor,
414
+ o_padded: torch.Tensor,
415
+ ) -> None:
416
+ """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly.
417
+
418
+ The nested ``_sparse_indexer_and_attn`` break runs inline, since
419
+ ``add_eager`` clears ``_capturing`` before invoking this.
420
+ """
421
+ self._prepare_and_attn(
422
+ hidden_states,
423
+ qr,
424
+ kv,
425
+ kv_score,
426
+ indexer_kv_score,
427
+ indexer_weights,
428
+ positions,
429
+ o_padded,
430
+ )
431
+
432
+ def _prepare_and_attn(
433
+ self,
434
+ hidden_states: torch.Tensor,
435
+ qr: torch.Tensor,
436
+ kv: torch.Tensor,
437
+ kv_score: torch.Tensor,
438
+ indexer_kv_score: torch.Tensor,
439
+ indexer_weights: torch.Tensor,
440
+ positions: torch.Tensor,
441
+ o_padded: torch.Tensor,
442
+ ) -> None:
443
+ """Attention input preparation followed by the sparse indexer and MLA.
444
+
445
+ Only the latter runs in the eager break.
446
+ """
447
attn_metadata = get_forward_context().attn_metadata
448
indexer = self.indexer
449
compressor = self.compressor
@@ -438,10 +503,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
503
positions,
504
o_padded,
505
)
441
- o = o_padded[:, : self.n_local_heads, :]
442
-
443
- # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
444
- return self._o_proj(o, positions)
506
507
def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
508
# Override point: the ROCm layer preshuffles this weight in place, so