[ROCm][Perf] Skip redundant sparse index remap on non-indexer layers (#51309)

Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com>
2026-10-13 05:25:56 +00:00
parent 2ce8012fa0
commit 0d2689e0df
5 changed files with 349 additions and 16 deletions
@@ -476,6 +476,8 @@ def test_sparse_mla_sink_forward_mqa_preserves_split_query(dtype):
        (num_tokens, 128), -1, dtype=torch.int32, device="cuda"
    )
    impl.topk_indices_buffer[:, : selected.shape[1]] = selected
+    impl.owns_indexer = True
+    impl.index_epoch = 0
    impl.q_concat_buffer = torch.empty_like(q)
    metadata = SimpleNamespace(
        attn_out_dtype=dtype,
@@ -0,0 +1,295 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if not current_platform.is_rocm():
+    pytest.skip(
+        "ROCm AITER sparse MLA index remap test requires ROCm.",
+        allow_module_level=True,
+    )
+
+from vllm._aiter_ops import is_aiter_found_and_supported
+
+if not is_aiter_found_and_supported():
+    pytest.skip(
+        "ROCm AITER sparse MLA index remap test requires a supported AITER "
+        "installation.",
+        allow_module_level=True,
+    )
+
+from vllm.model_executor.layers.attention import sparse_mla_attention as sparse_mla_mod
+from vllm.v1.attention.backends.mla import rocm_aiter_mla_sparse as sparse_mod
+
+NUM_HEADS = 16
+KV_LORA_RANK = 512
+QK_NOPE_HEAD_DIM = 128
+QK_ROPE_HEAD_DIM = 64
+HEAD_SIZE = 576
+NUM_TOKENS = 2
+TOPK_TOKENS = 4
+
+
+@pytest.fixture(autouse=True)
+def _reset_index_epochs():
+    # The counter is keyed by id(), and a freed test buffer can hand its id to
+    # the next one, so a leftover count would leak across tests.
+    sparse_mod._INDEX_EPOCHS.clear()
+    yield
+    sparse_mod._INDEX_EPOCHS.clear()
+
+
+def _make_metadata():
+    return SimpleNamespace(
+        num_actual_tokens=NUM_TOKENS,
+        req_id_per_token=torch.zeros(NUM_TOKENS, dtype=torch.int32, device="cpu"),
+        block_table=torch.arange(16, dtype=torch.int32, device="cpu").view(2, 8),
+        paged_kv_indptr=torch.zeros(NUM_TOKENS + 1, dtype=torch.int32, device="cpu"),
+        paged_kv_indices=torch.zeros(
+            NUM_TOKENS * TOPK_TOKENS, dtype=torch.int32, device="cpu"
+        ),
+        block_size=1,
+        topk_tokens=TOPK_TOKENS,
+        remapped_buf=None,
+        remapped_epoch=-1,
+    )
+
+
+def _patch_init_deps(monkeypatch, *, speculative: bool = False):
+    monkeypatch.setattr(
+        sparse_mod.AiterMLAHelper,
+        "check_num_heads_validity",
+        staticmethod(lambda num_heads: None),
+    )
+    vllm_config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=NUM_TOKENS),
+        model_config=SimpleNamespace(dtype=torch.bfloat16),
+        speculative_config=SimpleNamespace() if speculative else None,
+        attention_config=SimpleNamespace(hisparse_config=None),
+    )
+    monkeypatch.setattr(sparse_mod, "get_current_vllm_config", lambda: vllm_config)
+    monkeypatch.setattr(sparse_mla_mod, "get_current_vllm_config", lambda: vllm_config)
+    monkeypatch.setattr(
+        sparse_mla_mod, "get_tensor_model_parallel_world_size", lambda: 1
+    )
+    monkeypatch.setattr(
+        sparse_mod,
+        "current_workspace_manager",
+        lambda: SimpleNamespace(
+            get_simultaneous=lambda *specs: tuple(
+                torch.zeros(shape, dtype=dtype, device="cpu") for shape, dtype in specs
+            )
+        ),
+    )
+
+
+def _build_impl(indexer, topk_indices_buffer):
+    return sparse_mod.ROCMAiterMLASparseImpl(
+        num_heads=NUM_HEADS,
+        head_size=HEAD_SIZE,
+        scale=1.0,
+        num_kv_heads=1,
+        alibi_slopes=None,
+        sliding_window=None,
+        kv_cache_dtype="auto",
+        logits_soft_cap=None,
+        attn_type="decoder",
+        kv_sharing_target_layer_name=None,
+        topk_indices_buffer=topk_indices_buffer,
+        indexer=indexer,
+        q_lora_rank=None,
+        kv_lora_rank=KV_LORA_RANK,
+        qk_nope_head_dim=QK_NOPE_HEAD_DIM,
+        qk_rope_head_dim=QK_ROPE_HEAD_DIM,
+        qk_head_dim=QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM,
+        v_head_dim=QK_NOPE_HEAD_DIM,
+        kv_b_proj=None,
+    )
+
+
+def _new_buffer():
+    return torch.zeros(NUM_TOKENS, TOPK_TOKENS, dtype=torch.int32, device="cpu")
+
+
+def _patch_forward_deps(monkeypatch, impl, remap_calls):
+    monkeypatch.setattr(
+        sparse_mod,
+        "triton_convert_req_index_to_global_index",
+        lambda *args, **kwargs: remap_calls.append(args),
+    )
+    monkeypatch.setattr(
+        sparse_mod.AiterMLAHelper,
+        "get_mla_padded_q",
+        staticmethod(lambda num_heads, q: q),
+    )
+    attn_out = torch.zeros(NUM_TOKENS, NUM_HEADS, KV_LORA_RANK, device="cpu")
+    monkeypatch.setattr(
+        type(impl),
+        "_forward_mla",
+        lambda self, layer, q, kv_cache, attn_metadata: (attn_out, None),
+    )
+    return attn_out
+
+
+def _run_forward(impl, metadata):
+    return impl.forward_mqa(
+        torch.zeros(NUM_TOKENS, NUM_HEADS, HEAD_SIZE, device="cpu"),
+        torch.zeros(4, 1, HEAD_SIZE, device="cpu"),
+        metadata,
+        SimpleNamespace(_q_scale=None, _k_scale=None),
+    )
+
+
+@pytest.mark.parametrize("has_indexer", [True, False])
+def test_owns_indexer_is_structural(monkeypatch, has_indexer):
+    """owns_indexer reflects the layer, not the serving config."""
+    _patch_init_deps(monkeypatch)
+    buffer = _new_buffer()
+    indexer = SimpleNamespace(topk_indices_buffer=buffer) if has_indexer else None
+
+    impl = _build_impl(indexer, buffer)
+
+    assert impl.owns_indexer is has_indexer
+    assert impl.topk_indices_buffer is buffer
+
+
+def test_forward_mqa_remaps_once_per_metadata(monkeypatch):
+    """Skip layers sharing a metadata share its remap; the first one pays.
+
+    All layers read one top-k buffer, so the reuse has to key off the metadata
+    that owns ``paged_kv_indices``, not off the buffer being distinct.
+    """
+    _patch_init_deps(monkeypatch)
+    buffer = _new_buffer()
+    first = _build_impl(None, buffer)
+    second = _build_impl(None, buffer)
+
+    remap_calls: list[tuple] = []
+    attn_out = _patch_forward_deps(monkeypatch, first, remap_calls)
+    metadata = _make_metadata()
+
+    out, lse = _run_forward(first, metadata)
+    assert out is attn_out
+    assert lse is None
+    assert len(remap_calls) == 1
+    assert metadata.remapped_buf is buffer
+
+    _run_forward(second, metadata)
+    assert len(remap_calls) == 1
+
+
+def test_forward_mqa_remaps_for_each_metadata(monkeypatch):
+    """Skip layers on a second metadata must remap its own paged_kv_indices.
+
+    Regression: gating on ``indexer is not None`` skipped every layer in the
+    non-indexer metadata group, leaving its ``paged_kv_indices`` unwritten.
+    """
+    _patch_init_deps(monkeypatch)
+    buffer = _new_buffer()
+    indexer_layer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer)
+    skip_layer = _build_impl(None, buffer)
+
+    remap_calls: list[tuple] = []
+    _patch_forward_deps(monkeypatch, indexer_layer, remap_calls)
+
+    indexer_metadata = _make_metadata()
+    _run_forward(indexer_layer, indexer_metadata)
+    assert len(remap_calls) == 1
+    assert indexer_metadata.remapped_buf is buffer
+
+    # A distinct metadata instance carries its own, still-unwritten indices.
+    skip_metadata = _make_metadata()
+    _run_forward(skip_layer, skip_metadata)
+    assert len(remap_calls) == 2
+    assert skip_metadata.remapped_buf is buffer
+
+
+def test_forward_mqa_remaps_after_a_later_indexer_write(monkeypatch):
+    """A skip metadata must be refreshed once per indexer write, not once ever.
+
+    Regression for cross-layer index sharing (``index_topk_freq``). With
+    freq=4 on a 78-layer model the 57 skip layers share one metadata while the
+    21 indexer layers share another, so the skip metadata has to pick up every
+    new selection. Recording only the buffer left layers 4 onwards attending
+    to the first selection of the pass, which collapsed RULER niah_single_2 to
+    0.026 on GLM-5.3 and took GPQA from 0.884 to 0.672.
+    """
+    _patch_init_deps(monkeypatch)
+    buffer = _new_buffer()
+    # Construction order mirrors the layer stack: indexer, three skips, indexer.
+    _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer)
+    first_skip = _build_impl(None, buffer)
+    reuse_skip = _build_impl(None, buffer)
+    _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer)
+    next_skip = _build_impl(None, buffer)
+
+    remap_calls: list[tuple] = []
+    _patch_forward_deps(monkeypatch, first_skip, remap_calls)
+    skip_metadata = _make_metadata()
+
+    _run_forward(first_skip, skip_metadata)
+    assert len(remap_calls) == 1
+
+    # Same selection, same metadata: reuse.
+    _run_forward(reuse_skip, skip_metadata)
+    assert len(remap_calls) == 1
+
+    # A later indexer layer has since rewritten the buffer in place. The object
+    # is unchanged, so only the epoch can reveal it.
+    _run_forward(next_skip, skip_metadata)
+    assert len(remap_calls) == 2
+    assert skip_metadata.remapped_epoch == next_skip.index_epoch
+
+
+def test_index_epoch_tracks_the_preceding_indexer_write(monkeypatch):
+    """Layers between two indexer layers share the earlier one's epoch."""
+    _patch_init_deps(monkeypatch)
+    buffer = _new_buffer()
+    first_indexer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer)
+    skip_a = _build_impl(None, buffer)
+    skip_b = _build_impl(None, buffer)
+    second_indexer = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer)
+    skip_c = _build_impl(None, buffer)
+
+    assert skip_a.index_epoch == first_indexer.index_epoch
+    assert skip_b.index_epoch == first_indexer.index_epoch
+    assert second_indexer.index_epoch == first_indexer.index_epoch + 1
+    assert skip_c.index_epoch == second_indexer.index_epoch
+
+
+def test_forward_mqa_remaps_when_buffer_differs(monkeypatch):
+    """A layer reading a different buffer cannot reuse the metadata's remap."""
+    _patch_init_deps(monkeypatch)
+    remapped_buffer = _new_buffer()
+    other_buffer = _new_buffer()
+    layer = _build_impl(None, other_buffer)
+
+    remap_calls: list[tuple] = []
+    _patch_forward_deps(monkeypatch, layer, remap_calls)
+    metadata = _make_metadata()
+    metadata.remapped_buf = remapped_buffer
+
+    _run_forward(layer, metadata)
+    assert len(remap_calls) == 1
+    assert metadata.remapped_buf is other_buffer
+
+
+def test_forward_mqa_indexer_layer_always_remaps(monkeypatch):
+    """The indexer rewrites its buffer in place, so its remap is never stale."""
+    _patch_init_deps(monkeypatch)
+    buffer = _new_buffer()
+    impl = _build_impl(SimpleNamespace(topk_indices_buffer=buffer), buffer)
+
+    remap_calls: list[tuple] = []
+    _patch_forward_deps(monkeypatch, impl, remap_calls)
+    metadata = _make_metadata()
+    metadata.remapped_buf = buffer
+    metadata.remapped_epoch = impl.index_epoch
+
+    _run_forward(impl, metadata)
+    assert len(remap_calls) == 1
@@ -523,6 +523,8 @@ def test_forward_mqa_slices_full_batch_metadata_to_query_tokens(monkeypatch):
    impl.kv_cache_dtype = "auto"
    impl.num_heads = 4
    impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4)
+    impl.owns_indexer = True
+    impl.index_epoch = 0
    expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16)
    forward_mla = Mock(return_value=(expected_output, None))
    monkeypatch.setattr(impl, "_forward_mla", forward_mla)
@@ -570,6 +572,8 @@ def test_forward_mqa_slices_decode_subset_for_bf16_and_fp8(monkeypatch, kv_cache
    impl.kv_cache_dtype = kv_cache_dtype
    impl.num_heads = 4
    impl.topk_indices_buffer = torch.arange(48, dtype=torch.int32).view(12, 4)
+    impl.owns_indexer = True
+    impl.index_epoch = 0
    expected_output = torch.zeros(2, 4, 512, dtype=torch.bfloat16)
    forward_mla = Mock(return_value=(expected_output, None))
    monkeypatch.setattr(impl, "_forward_mla", forward_mla)
@@ -96,6 +96,8 @@ def test_forward_mqa_prepares_triton_sparse_mla_inputs(
    impl.sinks = torch.randn(num_heads, generator=gen).to(device)
    impl.kv_cache_dtype = kv_cache_dtype
    impl.topk_indices_buffer = topk_indices.to(device)
+    impl.owns_indexer = True
+    impl.index_epoch = 0
    impl.use_aiter_sparse_mla = True
    layer = SimpleNamespace(
        _q_scale=torch.tensor([0.02], device=device),
@@ -401,6 +401,18 @@ class ROCMAiterMLASparseBackend(AttentionBackend):
        return on_mi3xx()


+# Indexer writes per shared top-k buffer, in layer construction order.
+_INDEX_EPOCHS: dict[int, int] = {}
+
+
+def _index_epoch(buf: object, owns_indexer: bool) -> int:
+    """Number of indexer writes up to and including this layer."""
+    # id(): tensors are unhashable, and the buffer outlives every reader.
+    epoch = _INDEX_EPOCHS.get(id(buf), 0) + (1 if owns_indexer else 0)
+    _INDEX_EPOCHS[id(buf)] = epoch
+    return epoch
+
+
@dataclass(kw_only=True)
 class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata):
    qo_indptr: torch.Tensor | None = None
@@ -409,6 +421,10 @@ class ROCMAiterMLASparseMetadata(SparseMLACommonMetadata):
    paged_kv_indptr: torch.Tensor | None = None
    attn_out_dtype: torch.dtype | None = None

+    # Source buffer and index_epoch of paged_kv_indices; -1 before any remap.
+    remapped_buf: object = None
+    remapped_epoch: int = -1
+
    # Persistent MLA metadata (only populated when persistent mode is enabled,
    # i.e. when the aiter sparse decode kernel supports work-stealing splits).
    work_meta_data: torch.Tensor | None = None
@@ -818,6 +834,10 @@ class ROCMAiterMLASparseImpl(SparseMLACommonImpl[ROCMAiterMLASparseMetadata]):
        self.masked_mha_available = False

        vllm_config = get_current_vllm_config()
+        self.owns_indexer: bool = indexer is not None
+        self.index_epoch: int = _index_epoch(
+            self.topk_indices_buffer, self.owns_indexer
+        )
        max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
        q_concat_shape = (max_tokens, num_heads, head_size)
        (self.q_concat_buffer,) = current_workspace_manager().get_simultaneous(
@@ -1162,25 +1182,35 @@ class ROCMAiterMLASparseImpl(SparseMLACommonImpl[ROCMAiterMLASparseMetadata]):

        num_actual_toks = q.shape[0]

-        # Get topk indices
-        assert self.topk_indices_buffer is not None
-        topk_indices = fit_kpool_indices_to_aiter(
-            self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens
-        )
-
        kv_c_and_k_pe_cache, block_stride_rows = flat_kv_row_view(
            kv_c_and_k_pe_cache, attn_metadata.block_size
        )
-        triton_convert_req_index_to_global_index(
-            attn_metadata.req_id_per_token[:num_actual_toks],
-            attn_metadata.block_table,
-            topk_indices,
-            attn_metadata.paged_kv_indptr,
-            attn_metadata.paged_kv_indices,
-            BLOCK_SIZE=attn_metadata.block_size,
-            BLOCK_STRIDE_ROWS=block_stride_rows,
-            NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
-        )
+
+        # Indexer layers may rewrite their buffer, so always remap. Non-indexer
+        # layers have their own KV cache group: remap once per indexer write.
+        if (
+            self.owns_indexer
+            or attn_metadata.remapped_epoch != self.index_epoch
+            or attn_metadata.remapped_buf is not self.topk_indices_buffer
+        ):
+            # Get topk indices
+            assert self.topk_indices_buffer is not None
+            topk_indices = fit_kpool_indices_to_aiter(
+                self.topk_indices_buffer[:num_actual_toks], attn_metadata.topk_tokens
+            )
+
+            triton_convert_req_index_to_global_index(
+                attn_metadata.req_id_per_token[:num_actual_toks],
+                attn_metadata.block_table,
+                topk_indices,
+                attn_metadata.paged_kv_indptr,
+                attn_metadata.paged_kv_indices,
+                BLOCK_SIZE=attn_metadata.block_size,
+                BLOCK_STRIDE_ROWS=block_stride_rows,
+                NUM_TOPK_TOKENS=attn_metadata.topk_tokens,
+            )
+            attn_metadata.remapped_buf = self.topk_indices_buffer
+            attn_metadata.remapped_epoch = self.index_epoch

        # write the latent and rope to kv cache
        if fp8_attention: