[AMD] Reuse KV gather indices across ASM context prefill layers

python/sglang/srt/layers/attention/aiter_backend.py

@@ -7,6 +7,7 @@
 """
 
 import logging
+import os
 from dataclasses import dataclass
 from enum import Enum, auto
 from typing import TYPE_CHECKING, Optional
@@ -158,6 +159,11 @@ class ForwardMetadata:
     swa_out_cache_loc: Optional[torch.Tensor] = None
     local_kv_lens: Optional[torch.Tensor] = None
     verify_token_table: Optional[torch.Tensor] = None
+    # ASM context-chunk prefill: KV slots to gather and cu_seqlens_k, computed
+    # once per batch by AiterAttnBackend._asm_context_prefill_indices.
+    asm_ctx_ready: bool = False
+    asm_ctx_tok_idx: Optional[torch.Tensor] = None
+    asm_ctx_cu_k: Optional[torch.Tensor] = None
 
 
 _AITER_PARTITION_SIZE_ROCM = 256
@@ -178,6 +184,12 @@ def _aiter_fp8_asm_supports_gqa(num_q_heads: int, num_kv_heads: int) -> bool:
     return (num_q_heads // num_kv_heads) in _AITER_FP8_ASM_GQA_RATIOS
 
 
+# Cross-check the per-batch fast indices against the generic gather (syncs).
+_GFX_ASM_CTX_GATHER_CHECK = (
+    os.environ.get("SGLANG_GFX_ASM_CTX_GATHER_CHECK", "0") == "1"
+)
+
+
 def _asm_context_prefill_gather_indices(
     kv_indptr: torch.Tensor,
     kv_indices: torch.Tensor,
@@ -1059,6 +1071,70 @@ def _get_kv_indices_scratch(
             )
         return self._kv_indices_scratch[:required_tokens]
 
+    def _asm_context_prefill_indices(
+        self, forward_batch: ForwardBatch, bs: int, num_kv_slots: int
+    ):
+        """KV slots and cu_seqlens_k for the ASM context-chunk prefill, computed
+        once per batch and shared by every full-attention layer.
+
+        For a plain extend batch AiterIndicesUpdaterPrefill lays kv_indices out
+        token by token with kv_indptr = cumsum(seq_lens), so the slots to gather
+        are the first sum(seq_lens) entries and cu_seqlens_k is kv_indptr itself;
+        both follow from host-side lengths without a device sync. Anything else
+        (spec batches, missing host lengths, or a short table) takes the generic
+        gather, which validates the metadata on the device.
+        """
+        fm = self.forward_metadata
+        if fm.asm_ctx_ready:
+            return fm.asm_ctx_tok_idx, fm.asm_ctx_cu_k
+        fm.asm_ctx_ready = True
+        total_k = 0
+        if (
+            forward_batch.spec_info is None
+            and forward_batch.forward_mode.is_extend()
+            and forward_batch.seq_lens_cpu is not None
+        ):
+            total_k = int(forward_batch.seq_lens_cpu[:bs].sum())
+        if 0 < total_k <= fm.kv_indices.numel():
+            tok_idx = fm.kv_indices[:total_k]
+            cu_k = fm.kv_indptr[: bs + 1]
+            if cu_k.dtype != torch.int32:
+                cu_k = cu_k.to(torch.int32)
+            if _GFX_ASM_CTX_GATHER_CHECK:
+                ref = _asm_context_prefill_gather_indices(
+                    fm.kv_indptr[: bs + 1],
+                    fm.kv_indices,
+                    forward_batch.seq_lens[:bs],
+                    num_kv_slots,
+                    forward_batch.forward_mode,
+                )
+                assert (
+                    ref is not None
+                    and torch.equal(ref[0], tok_idx.to(torch.long))
+                    and torch.equal(ref[1].to(torch.int32), cu_k)
+                ), (
+                    "asm context prefill: fast gather indices differ from the generic gather"
+                )
+                logger.info(
+                    "[asm-context-prefill] fast gather indices verified: bs=%d total_k=%d",
+                    bs,
+                    total_k,
+                )
+        else:
+            gathered = _asm_context_prefill_gather_indices(
+                fm.kv_indptr[: bs + 1],
+                fm.kv_indices,
+                forward_batch.seq_lens[:bs],
+                num_kv_slots,
+                forward_batch.forward_mode,
+            )
+            if gathered is None:
+                return None, None
+            tok_idx, cu_k = gathered
+            cu_k = cu_k.to(torch.int32)
+        fm.asm_ctx_tok_idx, fm.asm_ctx_cu_k = tok_idx, cu_k
+        return tok_idx, cu_k
+
     def _set_uniform_qo_indptr(
         self, bs: int, tokens_per_req: int, device: torch.device
     ) -> torch.Tensor:
@@ -3473,6 +3549,8 @@ def forward_extend(
             # faster; gathering the paged fp8 KV into a contiguous varlen
             # buffer costs only ~20 us per layer at 70k context. The no-prefix
             # first chunk already takes the ASM branch below.
+            # This applies to Qwen3.5 full-attention layers only currently,
+            # Other configurations fall through to the attention paths below.
             if (
                 is_gfx95_supported()
                 and forward_batch.forward_mode.is_extend()
@@ -3492,15 +3570,10 @@ def forward_extend(
             ):
                 bs = forward_batch.batch_size
                 k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id)
-                gathered = _asm_context_prefill_gather_indices(
-                    self.forward_metadata.kv_indptr[: bs + 1],
-                    self.forward_metadata.kv_indices,
-                    forward_batch.seq_lens[:bs],
-                    self.token_to_kv_pool.get_key_buffer(layer.layer_id).shape[0],
-                    forward_batch.forward_mode,
+                tok_idx, cu_k = self._asm_context_prefill_indices(
+                    forward_batch, bs, k_cache.shape[0]
                 )
-                if gathered is not None:
-                    tok_idx, cu_k = gathered
+                if tok_idx is not None:
                     hk = layer.tp_k_head_num * layer.qk_head_dim
                     hv = layer.tp_v_head_num * layer.v_head_dim
                     # uint8 view: index_select is not implemented for fp8.
@@ -3531,7 +3604,7 @@ def forward_extend(
                         k_descale.reshape(1),
                         v_descale.reshape(1),
                         self.qo_indptr[:bs0],
-                        cu_k.to(torch.int32),
+                        cu_k,
                         self.forward_metadata.max_q_len,
                         int(self.forward_metadata.max_kv_len),
                         softmax_scale=layer.scaling,

test/registered/amd/test_asm_context_prefill_gather_indices.py

@@ -0,0 +1,138 @@
+"""Per-batch gather indices for the ROCm ASM context-chunk prefill.
+
+For a plain extend batch the prefill indices updater lays kv_indices out token by
+token with kv_indptr = cumsum(seq_lens), so the ASM context prefill can take the
+first sum(seq_lens) entries and kv_indptr itself without any device sync. Guards:
+the fast path equals the generic (syncing) gather, is computed once per batch,
+and batches that break the invariant fall back to the generic gather.
+"""
+
+import types
+import unittest
+from unittest.mock import patch
+
+import torch
+
+from sglang.test.ci.ci_register import register_amd_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_amd_ci(est_time=30, suite="stage-b-test-1-gpu-small-amd")
+
+
+@unittest.skipUnless(torch.version.hip, "ROCm attention backend")
+class TestAsmContextPrefillGatherIndices(CustomTestCase):
+    @classmethod
+    def setUpClass(cls):
+        from sglang.srt.layers.attention import aiter_backend as ab
+
+        cls.ab = ab
+        cls.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def _batch(self, seq_lens, pad=256, spec_info=None, kv_indptr=None):
+        seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int32)
+        seq_lens_dev = seq_lens_cpu.to(self.device)
+        bs = len(seq_lens)
+        if kv_indptr is None:
+            kv_indptr = torch.zeros(bs + 1, dtype=torch.int32, device=self.device)
+            kv_indptr[1:] = torch.cumsum(seq_lens_dev, 0)
+        total = int(kv_indptr[bs].item())
+        # pool slots in a shuffled order, padded the way the updater pads
+        pool = torch.randperm(total + pad, device=self.device)[:total].to(torch.int32)
+        kv_indices = torch.cat([pool, pool[:1].expand(pad)])
+        fm = self.ab.ForwardMetadata(
+            kv_indptr=kv_indptr,
+            kv_indices=kv_indices,
+            qo_indptr=None,
+            kv_last_page_len=None,
+            max_q_len=1,
+            max_kv_len=max(seq_lens),
+        )
+        backend = types.SimpleNamespace(forward_metadata=fm)
+        fb = types.SimpleNamespace(
+            seq_lens=seq_lens_dev,
+            seq_lens_cpu=seq_lens_cpu,
+            spec_info=spec_info,
+            forward_mode=types.SimpleNamespace(is_extend=lambda: True),
+        )
+        return backend, fb, bs, total + pad
+
+    def test_fast_path_matches_generic_gather(self):
+        torch.manual_seed(0)
+        for seq_lens in ([5], [7, 3, 12], [1, 1, 1, 1], [1024, 17, 300]):
+            with self.subTest(seq_lens=seq_lens):
+                backend, fb, bs, num_slots = self._batch(seq_lens)
+                fm = backend.forward_metadata
+                tok_idx, cu_k = self.ab.AiterAttnBackend._asm_context_prefill_indices(
+                    backend, fb, bs, num_slots
+                )
+                ref_tok, ref_cu = self.ab._asm_context_prefill_gather_indices(
+                    fm.kv_indptr[: bs + 1], fm.kv_indices, fb.seq_lens, num_slots
+                )
+                self.assertTrue(torch.equal(tok_idx.to(torch.long), ref_tok))
+                self.assertTrue(torch.equal(cu_k, ref_cu.to(torch.int32)))
+                self.assertEqual(cu_k.dtype, torch.int32)
+                self.assertEqual(tok_idx.numel(), sum(seq_lens))
+                # cached: the second call for the same batch returns the same tensors
+                again = self.ab.AiterAttnBackend._asm_context_prefill_indices(
+                    backend, fb, bs, num_slots
+                )
+                self.assertIs(again[0], tok_idx)
+                self.assertIs(again[1], cu_k)
+
+    def test_spec_batch_uses_generic_gather(self):
+        backend, fb, bs, num_slots = self._batch([9, 4], spec_info=object())
+        fm = backend.forward_metadata
+        tok_idx, cu_k = self.ab.AiterAttnBackend._asm_context_prefill_indices(
+            backend, fb, bs, num_slots
+        )
+        ref_tok, ref_cu = self.ab._asm_context_prefill_gather_indices(
+            fm.kv_indptr[: bs + 1], fm.kv_indices, fb.seq_lens, num_slots
+        )
+        self.assertTrue(torch.equal(tok_idx.to(torch.long), ref_tok))
+        self.assertTrue(torch.equal(cu_k, ref_cu.to(torch.int32)))
+
+    def test_short_table_falls_back(self):
+        # Host lengths exceed the table; the generic helper clamps them.
+        backend, fb, bs, num_slots = self._batch([6, 6], pad=0)
+        fb.seq_lens_cpu = torch.tensor([6, 100], dtype=torch.int32)
+        fb.seq_lens = fb.seq_lens_cpu.to(self.device)
+        tok_idx, cu_k = self.ab.AiterAttnBackend._asm_context_prefill_indices(
+            backend, fb, bs, num_slots
+        )
+        # the generic gather clamps the lengths to the table, so it still returns
+        self.assertIsNotNone(tok_idx)
+        self.assertEqual(tok_idx.numel(), 12)
+
+    def test_missing_host_lengths_uses_and_caches_generic_gather(self):
+        for speculative in (False, True):
+            with self.subTest(speculative=speculative):
+                backend, fb, bs, num_slots = self._batch(
+                    [9, 4], spec_info=object() if speculative else None
+                )
+                fm = backend.forward_metadata
+                fb.seq_lens_cpu = None
+                ref_tok, ref_cu = self.ab._asm_context_prefill_gather_indices(
+                    fm.kv_indptr[: bs + 1], fm.kv_indices, fb.seq_lens, num_slots
+                )
+                with patch.object(
+                    self.ab,
+                    "_asm_context_prefill_gather_indices",
+                    wraps=self.ab._asm_context_prefill_gather_indices,
+                ) as generic:
+                    tok_idx, cu_k = (
+                        self.ab.AiterAttnBackend._asm_context_prefill_indices(
+                            backend, fb, bs, num_slots
+                        )
+                    )
+                    again = self.ab.AiterAttnBackend._asm_context_prefill_indices(
+                        backend, fb, bs, num_slots
+                    )
+                self.assertEqual(generic.call_count, 1)
+                self.assertTrue(torch.equal(tok_idx.to(torch.long), ref_tok))
+                self.assertTrue(torch.equal(cu_k, ref_cu.to(torch.int32)))
+                self.assertIs(again[0], tok_idx)
+                self.assertIs(again[1], cu_k)
+
+
+if __name__ == "__main__":
+    unittest.main()