iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm][DSv4.1] Disambiguate ROCm attention class names (#58689)

Browse Source 
Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com>
Co-authored-by: Shanshan Shen <467638484@qq.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
This commit is contained in:
Tuukka Sarvi
2026-10-15 21:38:39 +08:00
committed by GitHub
co-authored by Shanshan Shen Kevin H. Luu
parent 1463501553
commit 0951fc183c
11 changed files with 86 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+2 -2
View File
@@ -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 (
tests/models/multimodal/generation/test_vit_cudagraph.py
+1 -1
View File
@@ -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"
),
tests/v1/attention/test_deepseek_v4_rocm_adaptive.py
+2 -2
View File
@@ -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())
tests/v1/attention/test_rocm_aiter_triton_sparse_mla.py
+6 -3
View File
@@ -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)
tests/v1/attention/test_rocm_attention_backends_selection.py
+42
View File
@@ -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,
vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py
+1
View File
@@ -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",
}
)
vllm/models/deepseek_v41/amd/model.py
+11 -3
View File
@@ -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
vllm/models/deepseek_v41/amd/mono_decode.py
+2 -2
View File
@@ -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:
vllm/models/deepseek_v41/amd/rocm.py
+15 -15
View File
@@ -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.
vllm/models/deepseek_v41/attention.py
+1 -1
View File
@@ -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,
vllm/v1/attention/backends/registry.py
+3
View File
@@ -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"
)