diff --git a/tests/test_config.py b/tests/test_config.py index 93bcdbc258..70c8728d25 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -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 diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 2b5e512cca..0a4c5cad88 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -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 ( diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 74a481e4b5..e2dc1bb45a 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -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