From 0951fc183cc7a86b54be45320dd42679c7065995 Mon Sep 17 00:00:00 2001 From: Tuukka Sarvi Date: Thu, 15 Oct 2026 16:38:39 +0300 Subject: [PATCH] [ROCm][DSv4.1] Disambiguate ROCm attention class names (#58689) Signed-off-by: Tuukka Sarvi Co-authored-by: Shanshan Shen <467638484@qq.com> Co-authored-by: Kevin H. Luu --- .../attention/test_rocm_triton_attn_dsv4.py | 4 +- .../generation/test_vit_cudagraph.py | 2 +- .../test_deepseek_v4_rocm_adaptive.py | 4 +- .../test_rocm_aiter_triton_sparse_mla.py | 9 ++-- .../test_rocm_attention_backends_selection.py | 42 +++++++++++++++++++ .../warmup/flashinfer_sparse_mla_warmup.py | 1 + vllm/models/deepseek_v41/amd/model.py | 14 +++++-- vllm/models/deepseek_v41/amd/mono_decode.py | 4 +- vllm/models/deepseek_v41/amd/rocm.py | 30 ++++++------- vllm/models/deepseek_v41/attention.py | 2 +- vllm/v1/attention/backends/registry.py | 3 ++ 11 files changed, 86 insertions(+), 29 deletions(-) diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 97d8cbaa3e..66414f723e 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -1989,8 +1989,8 @@ def test_dsv41_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None: """Replay V4.1 ROCm decode after device-only query reallocation.""" from tests.v1.attention.utils import create_vllm_config from vllm.models.deepseek_v41.amd.rocm import ( - DeepseekV4ROCMAiterSparseSWAMetadataBuilder, DeepseekV41ROCMAiterMLASparseMetadataBuilder, + DeepseekV41ROCMAiterSparseSWAMetadataBuilder, compute_global_topk_ragged_indices_and_indptr, ) from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata @@ -2044,7 +2044,7 @@ def test_dsv41_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None: mla_builder = DeepseekV41ROCMAiterMLASparseMetadataBuilder( mla_spec, ["c2a"], vllm_config, device ) - swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder( + swa_builder = DeepseekV41ROCMAiterSparseSWAMetadataBuilder( swa_spec, ["c2a"], vllm_config, device ) assert ( diff --git a/tests/models/multimodal/generation/test_vit_cudagraph.py b/tests/models/multimodal/generation/test_vit_cudagraph.py index 8c7cf1a68e..b7df7aa055 100644 --- a/tests/models/multimodal/generation/test_vit_cudagraph.py +++ b/tests/models/multimodal/generation/test_vit_cudagraph.py @@ -255,7 +255,7 @@ MODEL_CONFIGS: dict[str, VitCudagraphTestConfig] = { vllm_runner_kwargs={ "load_format": "dummy", "attention_backend": ( - "ROCM_FLASHMLA_SPARSE_DSV4" + "ROCM_FLASHMLA_SPARSE_DSV41" if current_platform.is_rocm() else "FLASHMLA_SPARSE_DSV41" ), diff --git a/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py index 7c18949881..8157bd6397 100644 --- a/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py +++ b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py @@ -67,7 +67,7 @@ def test_deepseek_rocm_adaptive_builders_support_varlen_full_graphs(): DeepseekV4ROCMAiterMLASparseMetadataBuilder, DeepseekV4ROCMAiterSparseSWAMetadataBuilder, DeepseekV41ROCMAiterMLASparseMetadataBuilder, - v41_rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder, + v41_rocm.DeepseekV41ROCMAiterSparseSWAMetadataBuilder, ): assert ( builder_cls.get_cudagraph_support(adaptive_config, SimpleNamespace()) @@ -126,7 +126,7 @@ def test_deepseek_v41_rocm_fixed_k_path_is_unchanged( for builder_cls in ( DeepseekV41ROCMAiterMLASparseMetadataBuilder, - v41_rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder, + v41_rocm.DeepseekV41ROCMAiterSparseSWAMetadataBuilder, ): assert ( builder_cls.get_cudagraph_support(fixed_config, SimpleNamespace()) diff --git a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py index 6849c4e3eb..e132a45c36 100644 --- a/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py +++ b/tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py @@ -171,6 +171,11 @@ def test_paged_prefill_indices_match_gathered_prefill( from vllm.v1.kv_cache_interface import SlidingWindowMLASpec rocm = importlib.import_module(f"vllm.models.{model}.amd.rocm") + builder_cls = ( + rocm.DeepseekV41ROCMAiterSparseSWAMetadataBuilder + if model == "deepseek_v41" + else rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder + ) device = torch.device("cuda") gen = torch.Generator().manual_seed(compress_ratio) window, width, block_size, max_model_len = 16, 32, 16, 128 @@ -233,9 +238,7 @@ def test_paged_prefill_indices_match_gathered_prefill( monkeypatch.setattr( rocm_aiter_ops, "is_triton_sparse_mla_enabled", lambda: enabled ) - return rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder( - swa_spec, ["swa"], vllm_config, device - ).build(0, common) + return builder_cls(swa_spec, ["swa"], vllm_config, device).build(0, common) assert build(enabled=False).prefill_swa_ragged_indices is None md = build(enabled=True) diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index f103b43650..2d264a7421 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch import pytest import torch +from vllm.config.attention import AttentionConfig from vllm.platforms import current_platform from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.selector import AttentionSelectorConfig @@ -44,6 +45,47 @@ def cleared_attention_selector_cache(): _cached_get_attn_backend.cache_clear() +def test_deepseek_rocm_backend_registry_paths_are_model_specific(): + assert AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4.get_path() == ( + "vllm.models.deepseek_v4.amd.rocm.DeepseekV4ROCMAiterMLASparseBackend" + ) + assert AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV41.get_path() == ( + "vllm.models.deepseek_v41.amd.rocm.DeepseekV41ROCMAiterMLASparseBackend" + ) + + +@pytest.mark.parametrize( + "backend_name, expect_warning", + [ + ("ROCM_FLASHMLA_SPARSE_DSV4", True), + ("ROCM_FLASHMLA_SPARSE_DSV41", False), + ], + ids=["legacy", "canonical"], +) +def test_deepseek_v41_rocm_accepts_legacy_backend_alias(backend_name, expect_warning): + from vllm.models.deepseek_v41.amd import model as dsv41_model + from vllm.models.deepseek_v41.amd.rocm import ( + DeepseekV41ROCMAiterMLAAttention, + DeepseekV41ROCMAiterMLASparseBackend, + ) + + config = SimpleNamespace( + attention_config=AttentionConfig(backend=backend_name), + ) + + with patch.object(dsv41_model.logger, "warning_once") as warning_once: + selected = dsv41_model._select_dsv4_attn_cls(config) + + assert selected is DeepseekV41ROCMAiterMLAAttention + if expect_warning: + warning_once.assert_called_once() + else: + warning_once.assert_not_called() + assert ( + DeepseekV41ROCMAiterMLASparseBackend.get_name() == "ROCM_FLASHMLA_SPARSE_DSV41" + ) + + def test_aiter_unified_attention_uses_dedicated_metadata_builder(): from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( RocmAiterUnifiedAttentionBackend, diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 91ae7c44c2..5c01f4ca2e 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -30,6 +30,7 @@ _DEEPSEEK_V4_SPARSE_MLA_BACKENDS = frozenset( "ROCM_FLASHMLA_SPARSE_DSV4", "FLASHMLA_SPARSE_DSV41", "FLASHINFER_MLA_SPARSE_DSV41", + "ROCM_FLASHMLA_SPARSE_DSV41", "DEEPSEEK_SPARSE_SWA", } ) diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 639bea27ad..ee4ef76c68 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -141,12 +141,20 @@ class DeepseekV4MoE(DeepseekV4MoEBase): def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]: """Pick the ROCm sparse-MLA implementation for DeepSeek V4.1.""" backend = vllm_config.attention_config.backend - if backend is not None and ( - backend is not AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4 + if backend is not None and backend not in ( + AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4, + AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV41, ): raise ValueError( f"{backend.name} is not supported for DeepSeek V4.1 on ROCm; " - "use ROCM_FLASHMLA_SPARSE_DSV4." + "use ROCM_FLASHMLA_SPARSE_DSV41." + ) + if backend is AttentionBackendEnum.ROCM_FLASHMLA_SPARSE_DSV4: + # TODO: Remove DSV4 compatibility after a deprecation period. + logger.warning_once( + "ROCM_FLASHMLA_SPARSE_DSV4 is deprecated for DeepSeek V4.1 " + "and will be removed in a future release. " + "Use ROCM_FLASHMLA_SPARSE_DSV41 instead." ) return DeepseekV41ROCMAiterMLAAttention diff --git a/vllm/models/deepseek_v41/amd/mono_decode.py b/vllm/models/deepseek_v41/amd/mono_decode.py index ddeb6b2631..2dd17cc6d2 100644 --- a/vllm/models/deepseek_v41/amd/mono_decode.py +++ b/vllm/models/deepseek_v41/amd/mono_decode.py @@ -39,7 +39,7 @@ from vllm.platforms import current_platform if TYPE_CHECKING: from vllm.models.deepseek_v41.amd.model import DeepseekV4DecoderLayer from vllm.models.deepseek_v41.amd.mono.runner import MonoLayerWeights - from vllm.models.deepseek_v41.amd.rocm import DeepseekV4ROCMAiterSparseSWAMetadata + from vllm.models.deepseek_v41.amd.rocm import DeepseekV41ROCMAiterSparseSWAMetadata from vllm.models.deepseek_v41.sparse_mla import DeepseekV4FlashMLAMetadata logger = init_logger(__name__) @@ -301,7 +301,7 @@ class MonoDecodeLayer: swa, comp = None, None if isinstance(md, dict): swa = cast( - "DeepseekV4ROCMAiterSparseSWAMetadata | None", + "DeepseekV41ROCMAiterSparseSWAMetadata | None", md.get(attn.swa_cache_layer.prefix), ) if attn.compressed_cache_prefix is not None: diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 5802a54b3f..44b9336632 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -280,7 +280,7 @@ _TopkRagged = tuple[torch.Tensor, torch.Tensor, torch.Tensor] @dataclass -class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): +class DeepseekV41ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata): decode_swa_ragged_indices: torch.Tensor | None = None decode_swa_ragged_indptr: torch.Tensor | None = None # Built only for the aiter prefill, which reads the SWA cache in place. @@ -303,7 +303,7 @@ class DeepseekV41ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBu return super().get_cudagraph_support(vllm_config, kv_cache_spec) -class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder): +class DeepseekV41ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder): @classmethod def get_cudagraph_support( cls, @@ -348,7 +348,7 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, replay_start: torch.Tensor | None = None, - ) -> DeepseekV4ROCMAiterSparseSWAMetadata: + ) -> DeepseekV41ROCMAiterSparseSWAMetadata: base = super().build( common_prefix_len=common_prefix_len, common_attn_metadata=common_attn_metadata, @@ -394,7 +394,7 @@ class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBu base.decode_swa_width, ) - return DeepseekV4ROCMAiterSparseSWAMetadata( + return DeepseekV41ROCMAiterSparseSWAMetadata( **vars(base), decode_swa_ragged_indices=ragged_indices, decode_swa_ragged_indptr=ragged_indptr, @@ -529,10 +529,10 @@ class DeepseekV41RocmMxfp4Indexer(DeepseekV4Indexer): return q, q_scale, weights -class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend): +class DeepseekV41ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend): @staticmethod def get_name() -> str: - return "ROCM_FLASHMLA_SPARSE_DSV4" + return "ROCM_FLASHMLA_SPARSE_DSV41" @staticmethod def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: @@ -545,14 +545,14 @@ class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend): class DeepseekV41ROCMAiterSparseSWABackend(DeepseekSparseSWABackend): @staticmethod - def get_builder_cls() -> type["DeepseekV4ROCMAiterSparseSWAMetadataBuilder"]: - return DeepseekV4ROCMAiterSparseSWAMetadataBuilder + def get_builder_cls() -> type["DeepseekV41ROCMAiterSparseSWAMetadataBuilder"]: + return DeepseekV41ROCMAiterSparseSWAMetadataBuilder class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): """ROCm sparse MLA attention layer for DeepSeek V4.1.""" - backend_cls = DeepseekV4ROCMAiterMLASparseBackend + backend_cls = DeepseekV41ROCMAiterMLASparseBackend swa_backend_cls = DeepseekV41ROCMAiterSparseSWABackend _use_aiter_sparse_mla = False @@ -1021,7 +1021,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): else None, ) swa_metadata = cast( - DeepseekV4ROCMAiterSparseSWAMetadata | None, + DeepseekV41ROCMAiterSparseSWAMetadata | None, attn_metadata.get(self.swa_cache_layer.prefix), ) assert swa_metadata is not None @@ -1100,7 +1100,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): def _decode_topk_ragged( self, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, attn_metadata: DeepseekV4FlashMLAMetadata | None, num_decodes: int, num_decode_tokens: int, @@ -1139,7 +1139,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): q: torch.Tensor, positions: torch.Tensor, kv_cache: torch.Tensor | None, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, attn_metadata: DeepseekV4FlashMLAMetadata | None, swa_only: bool, output: torch.Tensor | None, @@ -1233,7 +1233,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): swa_k_cache: torch.Tensor, output: torch.Tensor | None, attn_metadata: DeepseekV4FlashMLAMetadata | None, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> None: """Sparse prefill into bf16 ``output``, or into ``mxfp8_out``. @@ -1399,7 +1399,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): def _prefill_topk_ragged( self, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, attn_metadata: DeepseekV4FlashMLAMetadata, compressed_k_cache: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1449,7 +1449,7 @@ class DeepseekV41ROCMAiterMLAAttention(DeepseekV4Attention): swa_k_cache: torch.Tensor, output: torch.Tensor | None, attn_metadata: DeepseekV4FlashMLAMetadata | None, - swa_metadata: DeepseekV4ROCMAiterSparseSWAMetadata, + swa_metadata: DeepseekV41ROCMAiterSparseSWAMetadata, mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> None: """Prefill straight off the paged caches, with no bf16 gather. diff --git a/vllm/models/deepseek_v41/attention.py b/vllm/models/deepseek_v41/attention.py index 7009b466a2..b68879c5c8 100644 --- a/vllm/models/deepseek_v41/attention.py +++ b/vllm/models/deepseek_v41/attention.py @@ -646,7 +646,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): if self.backend_cls.get_name() in ( "FLASHMLA_SPARSE_DSV41", "FLASHMLA_MEGA_ATTN_DSV41", - "ROCM_FLASHMLA_SPARSE_DSV4", + "ROCM_FLASHMLA_SPARSE_DSV41", ): from vllm.models.deepseek_v41.common.ops.cache_utils import ( _COMBINE_TOPK_SWA_INDICES_KERNEL, diff --git a/vllm/v1/attention/backends/registry.py b/vllm/v1/attention/backends/registry.py index e1c4c8d067..d01f5b09a2 100644 --- a/vllm/v1/attention/backends/registry.py +++ b/vllm/v1/attention/backends/registry.py @@ -117,6 +117,9 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta): "vllm.models.deepseek_v41.nvidia.flashinfer_sparse." "DeepseekV4FlashInferMLASparseBackend" ) + ROCM_FLASHMLA_SPARSE_DSV41 = ( + "vllm.models.deepseek_v41.amd.rocm.DeepseekV41ROCMAiterMLASparseBackend" + ) FLASHMLA_MEGA_ATTN_DSV41 = ( "vllm.models.deepseek_v41.sparse_mla.FlashMLAMegaAttnBackend" )