[ROCm][DSV4][Perf] Pack sparse-MLA ragged decode metadata into graph buffers

tests/kernels/attention/test_rocm_triton_attn_dsv4.py

@@ -1089,46 +1089,55 @@ def test_sparse_attn_decode_scrubs_untrusted_cache_by_default() -> None:
 def test_rocm_ragged_graph_buffer_view_tracks_source_width(
     monkeypatch, on_gfx950: bool
 ) -> None:
-    from vllm.models.deepseek_v4.amd import rocm as rocm_mod
+    from vllm.v1.attention.ops import rocm_aiter_mla_sparse as ops_mod
+
+    monkeypatch.setattr(ops_mod, "_ON_GFX950", on_gfx950)
 
-    monkeypatch.setattr(rocm_mod, "_ON_GFX950", on_gfx950)
+    num_rows = 2
+    dense_width = 4
+    max_entries_per_row = 8
+    indices_buffer = torch.full((16,), -1, dtype=torch.int32, device="cuda")
+    indptr_buffer = torch.full((3,), -1, dtype=torch.int32, device="cuda")
 
-    indices_buffer = torch.full((16,), -1, dtype=torch.int32)
-    indptr_buffer = torch.full((3,), -1, dtype=torch.int32)
-    first_indices = torch.tensor([3, 5, 7], dtype=torch.int32)
-    first_indptr = torch.tensor([0, 1, 3], dtype=torch.int32)
-    first_view, first_indptr_view = rocm_mod._copy_ragged_to_graph_buffers(
-        first_indices,
-        first_indptr,
+    first_dense = torch.tensor(
+        [[3, -1, -1, -1], [5, 7, -1, -1]], dtype=torch.int32, device="cuda"
+    )
+    first_lens = torch.tensor([1, 2], dtype=torch.int32, device="cuda")
+    first_view, first_indptr_view = ops_mod.build_ragged_into_graph_buffers(
+        first_dense,
+        first_lens,
         indices_buffer,
         indptr_buffer,
-        num_rows=2,
-        max_entries_per_row=8,
+        num_rows,
+        max_entries_per_row,
     )
 
-    second_indices = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32)
-    second_indptr = torch.tensor([0, 2, 6], dtype=torch.int32)
-    second_view, second_indptr_view = rocm_mod._copy_ragged_to_graph_buffers(
-        second_indices,
-        second_indptr,
+    second_dense = torch.tensor(
+        [[1, -1, 2, -1], [3, 4, 5, 6]], dtype=torch.int32, device="cuda"
+    )
+    second_lens = torch.tensor([3, 4], dtype=torch.int32, device="cuda")
+    second_view, second_indptr_view = ops_mod.build_ragged_into_graph_buffers(
+        second_dense,
+        second_lens,
         indices_buffer,
         indptr_buffer,
-        num_rows=2,
-        max_entries_per_row=8,
+        num_rows,
+        max_entries_per_row,
     )
 
-    expected_first_entries = (
-        first_indices.numel() if on_gfx950 else indices_buffer.numel()
-    )
-    expected_second_entries = (
-        second_indices.numel() if on_gfx950 else indices_buffer.numel()
+    expected_entries = (
+        num_rows * dense_width if on_gfx950 else num_rows * max_entries_per_row
     )
-    assert first_view.numel() == expected_first_entries
-    assert second_view.numel() == expected_second_entries
+    assert first_view.numel() == expected_entries
+    assert second_view.numel() == expected_entries
     assert first_view.data_ptr() == second_view.data_ptr() == indices_buffer.data_ptr()
     assert first_indptr_view.data_ptr() == second_indptr_view.data_ptr()
-    assert torch.equal(second_view[: second_indices.numel()], second_indices)
-    assert torch.equal(second_indptr_view, second_indptr)
+    assert torch.equal(
+        second_view[:6], torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.int32).cuda()
+    )
+    assert torch.equal(
+        second_indptr_view, torch.tensor([0, 2, 6], dtype=torch.int32).cuda()
+    )
 
 
 def test_rocm_capture_metadata_sets_adaptive_marker(monkeypatch) -> None:

tests/kernels/models/amd/test_dsv41_gfx942_decode_meta.py

@@ -6,8 +6,8 @@
 swa_decode_meta.launch replaces the decode part of the ROCm SWA metadata
 build: torch.ge for is_valid_token, the zero fill of the unused lengths,
 _compute_swa_indices_and_lens_kernel (or the DSpark draft block's
-ComputeDSparkNoncausalSWAIndicesKernel), build_ragged_indices_from_dense and
-_copy_ragged_to_graph_buffers. topk_pack.launch replaces
+ComputeDSparkNoncausalSWAIndicesKernel) and
+build_ragged_into_graph_buffers. topk_pack.launch replaces
 compute_global_topk_ragged_indices_and_indptr. For each random decode batch
 vLLM's ops run first. Then every output buffer is overwritten with -7
 (is_valid_token with its negation), the fused launch runs, and every output
@@ -90,13 +90,12 @@ def _swa_buffers(width, noncausal):
 def _swa_old_ops(c, b):
     """The decode part of vLLM's SWA metadata build and its ROCm ragged
     step."""
-    from vllm.models.deepseek_v41.amd.rocm import _copy_ragged_to_graph_buffers
     from vllm.v1.attention.backends.mla.sparse_swa import (
         _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL,
         _COMPUTE_SWA_INDICES_AND_LENS_KERNEL,
     )
     from vllm.v1.attention.ops.rocm_aiter_mla_sparse import (
-        build_ragged_indices_from_dense,
+        build_ragged_into_graph_buffers,
     )
 
     n, width = c["n"], b["width"]
@@ -131,10 +130,14 @@ def _swa_old_ops(c, b):
             num_tokens=n,
             token_offset=0,
         )
-    ragged, indptr = build_ragged_indices_from_dense(
-        b["indices"][:n].reshape(n, width), b["lens"][:n]
+    build_ragged_into_graph_buffers(
+        b["indices"][:n].reshape(n, width),
+        b["lens"][:n],
+        b["ragged"],
+        b["indptr"],
+        n,
+        width,
     )
-    _copy_ragged_to_graph_buffers(ragged, indptr, b["ragged"], b["indptr"], n, width)
 
 
 @pytest.mark.parametrize("noncausal", [False, True])