From 10d2c6083d3b7508f2409a127a1d03dd489e427c Mon Sep 17 00:00:00 2001 From: Tuukka Sarvi Date: Thu, 15 Oct 2026 07:05:43 +0300 Subject: [PATCH] [ROCm][DSv4.1] Enable DSpark adaptive verification (#58151) Signed-off-by: Tuukka Sarvi Co-authored-by: Kevin H. Luu --- .buildkite/test_areas/kernels.yaml | 5 + .../attention/test_rocm_triton_attn_dsv4.py | 266 ++++++++++++++++++ .../test_deepseek_v4_rocm_adaptive.py | 69 +++-- vllm/models/deepseek_v41/amd/rocm.py | 35 ++- vllm/v1/attention/backends/mla/indexer.py | 11 +- 5 files changed, 360 insertions(+), 26 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index a1224f6231..126fff9bdc 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -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 diff --git a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py index 51b77ce172..97d8cbaa3e 100644 --- a/tests/kernels/attention/test_rocm_triton_attn_dsv4.py +++ b/tests/kernels/attention/test_rocm_triton_attn_dsv4.py @@ -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) # --------------------------------------------------------------------------- diff --git a/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py index 914cd249fe..7c18949881 100644 --- a/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py +++ b/tests/v1/attention/test_deepseek_v4_rocm_adaptive.py @@ -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( diff --git a/vllm/models/deepseek_v41/amd/rocm.py b/vllm/models/deepseek_v41/amd/rocm.py index 3ba6eebcbd..5802a54b3f 100644 --- a/vllm/models/deepseek_v41/amd/rocm.py +++ b/vllm/models/deepseek_v41/amd/rocm.py @@ -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): diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py index 0382a81fe0..04287ca1f6 100644 --- a/vllm/v1/attention/backends/mla/indexer.py +++ b/vllm/v1/attention/backends/mla/indexer.py @@ -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 + ) )