[Bugfix] Pick the DeepSeek V4 eager cudagraph region per model runner (#52401)

2026-08-15 21:04:08 -07:00
parent 6593754e61
commit 8efa13b700
3 changed files with 95 additions and 65 deletions
@@ -66,41 +66,18 @@ def test_v2_model_runner_env_tri_state(monkeypatch, env_value, expected):
    assert envs.VLLM_USE_V2_MODEL_RUNNER is expected


-@pytest.mark.parametrize(
-    "cudagraph_mode",
-    [CUDAGraphMode.PIECEWISE, CUDAGraphMode.FULL_AND_PIECEWISE],
-)
-def test_deepseek_v4_rejects_mrv1_piecewise_cudagraph(cudagraph_mode):
-    config = SimpleNamespace(
-        use_v2_model_runner=False,
-        model_config=SimpleNamespace(architectures=["DeepseekV4ForCausalLM"]),
-        compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode),
-    )
+def test_rocm_defaults_deepseek_v4_to_mrv1(monkeypatch):
+    """ROCm keeps DeepSeek V4 on MRV1, which is still faster there."""
+    from vllm.config.vllm import default_v2_model_runner_architectures
+    from vllm.platforms import current_platform

-    with pytest.raises(ValueError, match="DeepSeek V4 does not support PIECEWISE"):
-        VllmConfig._validate_mrv1_piecewise_cudagraph(config)
-
-
-@pytest.mark.parametrize(
-    ("use_v2_model_runner", "architecture", "cudagraph_mode"),
-    [
-        (True, "DeepseekV4ForCausalLM", CUDAGraphMode.PIECEWISE),
-        (False, "DeepseekV4ForCausalLM", CUDAGraphMode.NONE),
-        (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL),
-        (False, "DeepseekV4ForCausalLM", CUDAGraphMode.FULL_DECODE_ONLY),
-        (False, "LlamaForCausalLM", CUDAGraphMode.PIECEWISE),
-    ],
-)
-def test_mrv1_piecewise_cudagraph_allowed(
-    use_v2_model_runner, architecture, cudagraph_mode
-):
-    config = SimpleNamespace(
-        use_v2_model_runner=use_v2_model_runner,
-        model_config=SimpleNamespace(architectures=[architecture]),
-        compilation_config=SimpleNamespace(cudagraph_mode=cudagraph_mode),
-    )
-
-    VllmConfig._validate_mrv1_piecewise_cudagraph(config)
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
+    # The lookup is lru_cached against a fixed platform.
+    default_v2_model_runner_architectures.cache_clear()
+    try:
+        assert "DeepseekV4ForCausalLM" not in default_v2_model_runner_architectures()
+    finally:
+        default_v2_model_runner_architectures.cache_clear()


@pytest.mark.parametrize(
@@ -330,10 +307,20 @@ def test_resolve_cudagraph_mode_adjusts_spec_decode_sizes_only_for_v1(
        ),
    ],
 )
-def test_is_default_v2_model_runner_model(model_config, expected):
+def test_is_default_v2_model_runner_model(model_config, expected, monkeypatch):
+    from vllm.config.vllm import default_v2_model_runner_architectures
+    from vllm.platforms import current_platform
+
+    # The expectations below are the platform-independent defaults; ROCm's
+    # DeepSeek V4 carve-out is covered by test_rocm_defaults_deepseek_v4_to_mrv1.
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: False)
+    default_v2_model_runner_architectures.cache_clear()
    config = SimpleNamespace(model_config=model_config)

-    assert VllmConfig._is_default_v2_model_runner_model(config) is expected
+    try:
+        assert VllmConfig._is_default_v2_model_runner_model(config) is expected
+    finally:
+        default_v2_model_runner_architectures.cache_clear()


@pytest.mark.skip_global_cleanup
@@ -66,10 +66,6 @@ else:

 logger = init_logger(__name__)

-MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES = frozenset(
-    {"DeepseekV4ForCausalLM"}
-)
-
 DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
    {
        "DeepseekV2ForCausalLM",
@@ -87,6 +83,13 @@ DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES = frozenset(
@lru_cache
 def default_v2_model_runner_architectures() -> frozenset[str]:
    """Architectures defaulting to the V2 model runner on this platform."""
+    from vllm.platforms import current_platform
+
+    if current_platform.is_rocm():
+        # TODO(rocm): DeepSeek V4 is still faster on MRV1 on ROCm. The
+        # attention layer picks the eager cudagraph region MRV1 needs, so
+        # this is a perf default only; drop it once MRV2 catches up.
+        return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES - {"DeepseekV4ForCausalLM"}
    return DEFAULT_V2_MODEL_RUNNER_ARCHITECTURES


@@ -691,25 +694,6 @@ class VllmConfig:
            return False
        return is_default_v2_architecture or not model_config.is_moe

-    def _validate_mrv1_piecewise_cudagraph(self) -> None:
-        if self.use_v2_model_runner:
-            return
-        model_config = self.model_config
-        if model_config is None:
-            return
-        if not self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs():
-            return
-        architectures = getattr(model_config, "architectures", [])
-        if any(
-            arch in MRV1_UNSUPPORTED_PIECEWISE_CUDAGRAPH_ARCHITECTURES
-            for arch in architectures
-        ):
-            raise ValueError(
-                "DeepSeek V4 does not support PIECEWISE CUDA graphs with "
-                "Model Runner V1. Use Model Runner V2 or disable PIECEWISE "
-                "CUDA graphs."
-            )
-
    @property
    def needs_dp_coordinator(self) -> bool:
        """
@@ -1644,8 +1628,6 @@ class VllmConfig:
                        "pipeline parallelism",
                    )

-        self._validate_mrv1_piecewise_cudagraph()
-
        # final check of cudagraph mode after all possible updates
        if current_platform.is_cuda_alike():
            if (
@@ -298,6 +298,13 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
                eager_scratch_pool=eager_scratch_pool,
            )

+        self._prepare_and_attn_fn = self._prepare_and_attn
+        if not vllm_config.use_v2_model_runner:
+            # MRV1's piecewise capture only tolerates the wide eager region: with
+            # the narrow one the attention input preparation stays in the captured
+            # graph and MRV1 produces garbage (#51430).
+            self._prepare_and_attn_fn = self._prepare_and_attn_eager
+
        # Will be None on ROCm for now.
        self.aux_stream_list = aux_stream_list
        # [0]: GEMM start / post-GEMM event0. [1..3]: GEMM done events;
@@ -379,6 +386,64 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
            self.eps,
        )

+        self._prepare_and_attn_fn(
+            hidden_states,
+            qr,
+            kv,
+            kv_score,
+            indexer_kv_score,
+            indexer_weights,
+            positions,
+            o_padded,
+        )
+        o = o_padded[:, : self.n_local_heads, :]
+
+        # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
+        return self._o_proj(o, positions)
+
+    @eager_break_during_capture
+    def _prepare_and_attn_eager(
+        self,
+        hidden_states: torch.Tensor,
+        qr: torch.Tensor,
+        kv: torch.Tensor,
+        kv_score: torch.Tensor,
+        indexer_kv_score: torch.Tensor,
+        indexer_weights: torch.Tensor,
+        positions: torch.Tensor,
+        o_padded: torch.Tensor,
+    ) -> None:
+        """Wide eager region: the whole of ``_prepare_and_attn`` runs eagerly.
+
+        The nested ``_sparse_indexer_and_attn`` break runs inline, since
+        ``add_eager`` clears ``_capturing`` before invoking this.
+        """
+        self._prepare_and_attn(
+            hidden_states,
+            qr,
+            kv,
+            kv_score,
+            indexer_kv_score,
+            indexer_weights,
+            positions,
+            o_padded,
+        )
+
+    def _prepare_and_attn(
+        self,
+        hidden_states: torch.Tensor,
+        qr: torch.Tensor,
+        kv: torch.Tensor,
+        kv_score: torch.Tensor,
+        indexer_kv_score: torch.Tensor,
+        indexer_weights: torch.Tensor,
+        positions: torch.Tensor,
+        o_padded: torch.Tensor,
+    ) -> None:
+        """Attention input preparation followed by the sparse indexer and MLA.
+
+        Only the latter runs in the eager break.
+        """
        attn_metadata = get_forward_context().attn_metadata
        indexer = self.indexer
        compressor = self.compressor
@@ -438,10 +503,6 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
            positions,
            o_padded,
        )
-        o = o_padded[:, : self.n_local_heads, :]
-
-        # Inverse-RoPE + wo_a + wo_b output projection (platform-specific).
-        return self._o_proj(o, positions)

    def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # Override point: the ROCm layer preshuffles this weight in place, so