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

[ROCm][DSv4.1] Enable DSpark adaptive verification (#58151)

Browse Source 
Signed-off-by: Tuukka Sarvi <tuukka.sarvi@amd.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>
This commit is contained in:
Tuukka Sarvi
2026-10-15 12:05:43 +08:00
committed by GitHub
co-authored by Kevin H. Luu
parent f46021a7fa
commit 10d2c6083d
5 changed files with 360 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.buildkite/test_areas/kernels.yaml
+5
View File
@@ -411,10 +411,14 @@ steps:
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/models/deepseek_v41/amd/rocm.py
- vllm/models/deepseek_v41/sparse_mla.py
- vllm/v1/attention/backends/mla/rocm_aiter_mla.py
- vllm/v1/attention/ops/rocm_aiter_mla_sparse.py
- vllm/v1/attention/selector.py
- vllm/platforms/rocm.py
- vllm/_aiter_ops.py
- tests/kernels/attention/test_rocm_triton_attn_dsv4.py
- tests/kernels/attention/test_rocm_aiter_mla_decode.py
- tests/kernels/attention/test_rocm_aiter_mla_decode_metadata.py
- tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
@@ -425,6 +429,7 @@ steps:
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_decode.py
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_fp8_support.py
- pytest -v -s tests/kernels/attention/test_rocm_aiter_mla_op_registration.py
- pytest -v -s tests/kernels/attention/test_rocm_triton_attn_dsv4.py::test_dsv41_adaptive_mla_swa_metadata_graph_replay
- label: ":nvidia: (L4) Quantization Kernels Shard %N"
key: kernels-quantization-test
tests/kernels/attention/test_rocm_triton_attn_dsv4.py
+266
View File
@@ -1983,6 +1983,272 @@ def test_dsv4_adaptive_mla_swa_metadata_graph_replay(monkeypatch) -> None:
torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
@requires_gfx950
@torch.inference_mode()
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,
compute_global_topk_ragged_indices_and_indptr,
)
from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
from vllm.v1.attention.ops import rocm_aiter_mla_sparse as mod
from vllm.v1.kv_cache_interface import MLAAttentionSpec, SlidingWindowMLASpec
device = torch.device("cuda")
num_reqs = 3
upper_query_len = 6
graph_tokens = num_reqs * upper_query_len
block_size = 128
compress_ratio = 2
compressed_block_size = block_size // compress_ratio
vllm_config = create_vllm_config(
model_name="facebook/opt-125m",
max_model_len=1024,
block_size=block_size,
max_num_seqs=num_reqs,
max_num_batched_tokens=graph_tokens,
hf_config_override={
"compress_ratios": [0, 1, 2],
"index_topk": 8,
"sliding_window": 32,
},
)
vllm_config.speculative_config = SimpleNamespace(
num_speculative_tokens=upper_query_len - 1,
parallel_drafting=False,
enable_adaptive_verification=True,
use_dspark=lambda: True,
)
mla_spec = MLAAttentionSpec(
block_size=block_size,
num_kv_heads=1,
head_size=HEAD_DIM,
dtype=torch.bfloat16,
tokens_per_state=compress_ratio,
cache_dtype_str="fp8_ds_mla",
model_version="deepseek_v4",
)
swa_spec = SlidingWindowMLASpec(
block_size=block_size,
num_kv_heads=1,
head_size=HEAD_DIM,
dtype=torch.uint8,
sliding_window=32,
cache_dtype_str="fp8_ds_mla",
model_version="deepseek_v4",
)
mla_builder = DeepseekV41ROCMAiterMLASparseMetadataBuilder(
mla_spec, ["c2a"], vllm_config, device
)
swa_builder = DeepseekV4ROCMAiterSparseSWAMetadataBuilder(
swa_spec, ["c2a"], vllm_config, device
)
assert (
mla_builder.get_cudagraph_support(vllm_config, mla_spec)
== AttentionCGSupport.ALWAYS
)
assert (
swa_builder.get_cudagraph_support(vllm_config, swa_spec)
== AttentionCGSupport.ALWAYS
)
seq_lens = torch.tensor([100, 110, 120], dtype=torch.int32, device=device)
seq_lens_cpu = seq_lens.cpu()
query_start_loc_cpu = torch.arange(
0, graph_tokens + 1, upper_query_len, dtype=torch.int32
)
block_table = torch.arange(num_reqs, dtype=torch.int32, device=device).view(
num_reqs, 1
)
def build_metadata(query_lens: list[int]):
query_lens_tensor = torch.tensor(query_lens, dtype=torch.int32, device=device)
query_start_loc = torch.cat(
[
torch.zeros(1, dtype=torch.int32, device=device),
query_lens_tensor.cumsum(0),
]
)
active_tokens = sum(query_lens)
slot_mapping = torch.full((graph_tokens,), -1, dtype=torch.int64, device=device)
slot_mapping[:active_tokens] = torch.arange(
active_tokens, dtype=torch.int64, device=device
)
common = CommonAttentionMetadata(
query_start_loc=query_start_loc,
query_start_loc_cpu=query_start_loc_cpu,
seq_lens=seq_lens,
seq_lens_cpu_upper_bound=seq_lens_cpu,
num_reqs=num_reqs,
num_actual_tokens=graph_tokens,
max_query_len=upper_query_len,
max_seq_len=int(seq_lens_cpu.max()),
block_table_tensor=block_table,
slot_mapping=slot_mapping,
causal=True,
)
return (
mla_builder.build_for_cudagraph_capture(common),
swa_builder.build_for_cudagraph_capture(common),
)
mla_metadata, swa_metadata = build_metadata([6, 6, 6])
metadata_ptrs = (
mla_metadata.req_id_per_token.data_ptr(),
mla_metadata.slot_mapping.data_ptr(),
swa_metadata.token_to_req_indices.data_ptr(),
swa_metadata.decode_swa_ragged_indices.data_ptr(),
swa_metadata.decode_swa_ragged_indptr.data_ptr(),
)
topk_indices = (
torch.arange(8, dtype=torch.int32, device=device)
.expand(graph_tokens, -1)
.contiguous()
)
torch.manual_seed(23)
num_heads = 16
q = (
torch.randn(
graph_tokens,
num_heads,
HEAD_DIM,
dtype=torch.bfloat16,
device=device,
)
* 0.125
)
swa_cache = _pack_fp8_ds_mla_cache(
torch.randn(
num_reqs * block_size,
HEAD_DIM,
dtype=torch.bfloat16,
device=device,
)
* 0.125,
block_size,
use_fnuz=False,
)
compressed_cache = _pack_fp8_ds_mla_cache(
torch.randn(
num_reqs * compressed_block_size,
HEAD_DIM,
dtype=torch.bfloat16,
device=device,
)
* 0.125,
compressed_block_size,
use_fnuz=False,
)
attn_sink = torch.linspace(-0.1, 0.1, num_heads, dtype=torch.float32, device=device)
out = torch.empty_like(q)
monkeypatch.setattr(mod, "_decode_gfx950_num_splits", lambda *args: 1)
def build_topk_ragged(mla_md, swa_md):
return compute_global_topk_ragged_indices_and_indptr(
topk_indices,
swa_md.token_to_req_indices,
mla_md.block_table,
compressed_block_size,
swa_md.is_valid_token,
)
def run_decode(mla_md, swa_md) -> None:
topk_ragged_indices, topk_ragged_indptr, topk_lens = build_topk_ragged(
mla_md, swa_md
)
mod.rocm_sparse_attn_decode(
q=q,
kv_cache=compressed_cache,
swa_k_cache=swa_cache,
swa_only=False,
topk_indices=None,
topk_lens=topk_lens,
swa_indices=swa_md.decode_swa_indices,
swa_lens=swa_md.decode_swa_lens,
swa_ragged_indices=swa_md.decode_swa_ragged_indices,
swa_ragged_indptr=swa_md.decode_swa_ragged_indptr,
topk_ragged_indices=topk_ragged_indices,
topk_ragged_indptr=topk_ragged_indptr,
attn_sink=attn_sink,
scale=HEAD_DIM**-0.5,
head_dim=HEAD_DIM,
nope_head_dim=NOPE_HEAD_DIM,
rope_head_dim=ROPE_HEAD_DIM,
output=out,
extra_cache_nan_free=True,
)
def reference(mla_md, swa_md) -> torch.Tensor:
topk_ragged_indices, topk_ragged_indptr, _ = build_topk_ragged(mla_md, swa_md)
return _ref_sparse_decode_ragged(
q=q,
main_cache=swa_cache,
main_rows=_rows_from_ragged(
swa_md.decode_swa_ragged_indices,
swa_md.decode_swa_ragged_indptr,
),
scale=HEAD_DIM**-0.5,
attn_sink=attn_sink,
block_size=block_size,
extra_cache=compressed_cache,
extra_rows=_rows_from_ragged(
topk_ragged_indices,
topk_ragged_indptr,
),
extra_block_size=compressed_block_size,
)
run_decode(mla_metadata, swa_metadata)
torch.accelerator.synchronize()
expected_full = reference(mla_metadata, swa_metadata)
torch.testing.assert_close(out, expected_full, atol=2e-2, rtol=2e-2)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
run_decode(mla_metadata, swa_metadata)
torch.accelerator.synchronize()
captured_full = out.clone()
reallocated_mla, reallocated_swa = build_metadata([2, 6, 4])
assert metadata_ptrs == (
reallocated_mla.req_id_per_token.data_ptr(),
reallocated_mla.slot_mapping.data_ptr(),
reallocated_swa.token_to_req_indices.data_ptr(),
reallocated_swa.decode_swa_ragged_indices.data_ptr(),
reallocated_swa.decode_swa_ragged_indptr.data_ptr(),
)
expected_owners = torch.tensor(
[0, 0, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2],
dtype=torch.int32,
device=device,
)
torch.testing.assert_close(reallocated_mla.req_id_per_token[:12], expected_owners)
torch.testing.assert_close(
reallocated_swa.token_to_req_indices[:12], expected_owners
)
assert reallocated_swa.is_valid_token[:12].all()
assert not reallocated_swa.is_valid_token[12:].any()
torch.testing.assert_close(
reallocated_swa.decode_swa_lens[12:],
torch.zeros(6, dtype=torch.int32, device=device),
)
expected_reallocated = reference(reallocated_mla, reallocated_swa)
run_decode(reallocated_mla, reallocated_swa)
torch.accelerator.synchronize()
torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
assert not torch.equal(captured_full, expected_reallocated)
graph.replay()
torch.accelerator.synchronize()
torch.testing.assert_close(out, expected_reallocated, atol=2e-2, rtol=2e-2)
# ---------------------------------------------------------------------------
# o-projection: fused inverse-RoPE + cached bf16 wo_a (rocm_inv_rope_einsum)
# ---------------------------------------------------------------------------
tests/v1/attention/test_deepseek_v4_rocm_adaptive.py
+52 -17
View File
@@ -10,6 +10,10 @@ from vllm.models.deepseek_v4.amd.rocm import (
DeepseekV4ROCMAiterMLASparseMetadataBuilder,
DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
)
from vllm.models.deepseek_v41.amd import rocm as v41_rocm
from vllm.models.deepseek_v41.amd.rocm import (
DeepseekV41ROCMAiterMLASparseMetadataBuilder,
)
from vllm.v1.attention.backend import AttentionCGSupport
from vllm.v1.attention.backends.mla import indexer
from vllm.v1.attention.backends.mla.indexer import (
@@ -51,7 +55,7 @@ def _make_indexer_builder(*, adaptive: bool, capacity: int = 12):
@pytest.mark.cpu_test
def test_deepseek_v4_rocm_adaptive_builders_support_varlen_full_graphs():
def test_deepseek_rocm_adaptive_builders_support_varlen_full_graphs():
adaptive_config = SimpleNamespace(
speculative_config=SimpleNamespace(enable_adaptive_verification=True)
)
@@ -62,6 +66,8 @@ def test_deepseek_v4_rocm_adaptive_builders_support_varlen_full_graphs():
for builder_cls in (
DeepseekV4ROCMAiterMLASparseMetadataBuilder,
DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
DeepseekV41ROCMAiterMLASparseMetadataBuilder,
v41_rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
):
assert (
builder_cls.get_cudagraph_support(adaptive_config, SimpleNamespace())
@@ -74,20 +80,25 @@ def test_deepseek_v4_rocm_adaptive_builders_support_varlen_full_graphs():
@pytest.mark.cpu_test
def test_deepseek_v4_rocm_adaptive_indexer_support(monkeypatch: pytest.MonkeyPatch):
@pytest.mark.parametrize(
("architecture", "backend_cls"),
[
("DeepseekV4ForCausalLM", DeepseekV4IndexerBackend),
("DeepseekV41ForCausalLM", DeepseekV41IndexerBackend),
],
)
def test_deepseek_rocm_adaptive_indexer_support(
monkeypatch: pytest.MonkeyPatch,
architecture: str,
backend_cls: type[DeepseekV4IndexerBackend],
):
_mock_rocm_platform(monkeypatch)
adaptive_config = _make_indexer_config(
architecture="DeepseekV4ForCausalLM", adaptive=True
)
fixed_config = _make_indexer_config(
architecture="DeepseekV4ForCausalLM", adaptive=False
)
adaptive_config = _make_indexer_config(architecture=architecture, adaptive=True)
fixed_config = _make_indexer_config(architecture=architecture, adaptive=False)
assert DeepseekV4IndexerBackend.supports_device_cpu_query_lens_mismatch()
assert backend_cls.supports_device_cpu_query_lens_mismatch()
assert indexer._rocm_supports_flattened_device_query_lens(adaptive_config)
assert (
DeepseekV4IndexerBackend.get_builder_cls() is DeepseekV32IndexerMetadataBuilder
)
assert backend_cls.get_builder_cls() is DeepseekV32IndexerMetadataBuilder
assert indexer._use_flattening(adaptive_config)
assert (
DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
@@ -104,12 +115,39 @@ def test_deepseek_v4_rocm_adaptive_indexer_support(monkeypatch: pytest.MonkeyPat
)
@pytest.mark.cpu_test
def test_deepseek_v41_rocm_fixed_k_path_is_unchanged(
monkeypatch: pytest.MonkeyPatch,
):
_mock_rocm_platform(monkeypatch)
fixed_config = _make_indexer_config(
architecture="DeepseekV41ForCausalLM", adaptive=False
)
for builder_cls in (
DeepseekV41ROCMAiterMLASparseMetadataBuilder,
v41_rocm.DeepseekV4ROCMAiterSparseSWAMetadataBuilder,
):
assert (
builder_cls.get_cudagraph_support(fixed_config, SimpleNamespace())
== AttentionCGSupport.UNIFORM_BATCH
)
assert not indexer._use_flattening(fixed_config)
assert (
DeepseekV32IndexerMetadataBuilder.get_cudagraph_support(
fixed_config, SimpleNamespace()
)
== AttentionCGSupport.UNIFORM_BATCH
)
@pytest.mark.cpu_test
@pytest.mark.parametrize(
"architecture",
["DeepseekV32ForCausalLM", "DeepseekV41ForCausalLM", "GlmMoeDsaForCausalLM"],
["DeepseekV32ForCausalLM", "GlmMoeDsaForCausalLM"],
)
def test_rocm_adaptive_indexer_flattening_is_scoped_to_deepseek_v4(
def test_rocm_adaptive_indexer_flattening_rejects_other_architectures(
monkeypatch: pytest.MonkeyPatch,
architecture: str,
):
@@ -125,9 +163,6 @@ def test_rocm_adaptive_indexer_flattening_is_scoped_to_deepseek_v4(
== AttentionCGSupport.UNIFORM_BATCH
)
if architecture == "DeepseekV41ForCausalLM":
assert not DeepseekV41IndexerBackend.supports_device_cpu_query_lens_mismatch()
@pytest.mark.cpu_test
def test_rocm_adaptive_indexer_preserves_single_request_uniform_path(
vllm/models/deepseek_v41/amd/rocm.py
+33 -2
View File
@@ -8,6 +8,7 @@ from typing import Any, cast
import torch
from vllm.config import VllmConfig
from vllm.distributed import (
get_tensor_model_parallel_world_size,
tensor_model_parallel_all_reduce,
@@ -41,6 +42,7 @@ from vllm.platforms.rocm import _ON_GFX950
from vllm.triton_utils import tl, triton
from vllm.utils.multi_stream_utils import execute_in_parallel
from vllm.v1.attention.backend import (
AttentionCGSupport,
CommonAttentionMetadata,
MultipleOf,
)
@@ -61,6 +63,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import (
rocm_paged_mxfp4_cache_layout,
)
from vllm.v1.kv_cache_interface import KVCacheSpec
from vllm.v1.worker.workspace import current_workspace_manager
logger = init_logger(__name__)
@@ -285,7 +288,35 @@ class DeepseekV4ROCMAiterSparseSWAMetadata(DeepseekSparseSWAMetadata):
prefill_swa_ragged_indptr: torch.Tensor | None = None
class DeepseekV41ROCMAiterMLASparseMetadataBuilder(DeepseekV4SparseMLAMetadataBuilder):
@classmethod
def get_cudagraph_support(
cls,
vllm_config: VllmConfig,
kv_cache_spec: KVCacheSpec,
) -> AttentionCGSupport:
spec_config = vllm_config.speculative_config
if spec_config is not None and spec_config.enable_adaptive_verification:
# Per-token request ownership and compressed slot mappings are built
# from device query boundaries into persistent buffers.
return AttentionCGSupport.ALWAYS
return super().get_cudagraph_support(vllm_config, kv_cache_spec)
class DeepseekV4ROCMAiterSparseSWAMetadataBuilder(DeepseekV41SparseSWAMetadataBuilder):
@classmethod
def get_cudagraph_support(
cls,
vllm_config: VllmConfig,
kv_cache_spec: KVCacheSpec,
) -> AttentionCGSupport:
spec_config = vllm_config.speculative_config
if spec_config is not None and spec_config.enable_adaptive_verification:
# SWA indices, lengths, and request ownership are built from device
# query boundaries and copied into persistent ragged buffers.
return AttentionCGSupport.ALWAYS
return super().get_cudagraph_support(vllm_config, kv_cache_spec)
# Keep fused multi-step decode disabled until update_draft_decode_metadata()
# also refreshes the ROCm-specific ragged SWA indices and indptrs.
supports_draft_decode_metadata_update = False
@@ -508,8 +539,8 @@ class DeepseekV4ROCMAiterMLASparseBackend(DeepseekV4SparseMLABackend):
return [64, 128]
@staticmethod
def get_builder_cls() -> type[DeepseekV4SparseMLAMetadataBuilder]:
return DeepseekV4SparseMLAMetadataBuilder
def get_builder_cls() -> type[DeepseekV41ROCMAiterMLASparseMetadataBuilder]:
return DeepseekV41ROCMAiterMLASparseMetadataBuilder
class DeepseekV41ROCMAiterSparseSWABackend(DeepseekSparseSWABackend):
vllm/v1/attention/backends/mla/indexer.py
+4 -7
View File
@@ -357,12 +357,6 @@ class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend):
class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend):
@classmethod
def supports_device_cpu_query_lens_mismatch(cls) -> bool:
# The ROCm flattened-query support above is validated for the
# DeepSeek-V4 adaptive DSpark path only.
return DeepseekV32IndexerBackend.supports_device_cpu_query_lens_mismatch()
@staticmethod
def get_name() -> str:
return "DEEPSEEK_V41_INDEXER"
@@ -953,7 +947,10 @@ def _rocm_supports_flattened_device_query_lens(vllm_config: VllmConfig) -> bool:
return (
current_platform.is_rocm()
and model_config is not None
and "DeepseekV4ForCausalLM" in model_config.architectures
and (
"DeepseekV4ForCausalLM" in model_config.architectures
or "DeepseekV41ForCausalLM" in model_config.architectures
)
)