From 36afd70f9c064de4bbf552100b0dc91ee882c80d Mon Sep 17 00:00:00 2001 From: Olga Miroshnichenko Date: Sat, 29 Aug 2026 03:17:18 +0300 Subject: [PATCH] [ROCm][Bugfix] Use token-level KV indices in the aiter ASM context-prefill gather (#36852) --- .../srt/layers/attention/aiter_backend.py | 104 ++++++++++------- .../test_aiter_asm_prefill_gather.py | 106 ++++++++++++++++++ 2 files changed, 170 insertions(+), 40 deletions(-) create mode 100644 test/registered/attention/test_aiter_asm_prefill_gather.py diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index e25583559..1f87e3b9d 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -149,6 +149,62 @@ class ForwardMetadata: _AITER_PARTITION_SIZE_ROCM = 256 +def _asm_context_prefill_gather_indices( + kv_indptr: torch.Tensor, + kv_indices: torch.Tensor, + seq_lens: torch.Tensor, + num_kv_slots: int, + forward_mode=None, +): + """KV-pool slots to gather for the ASM context-chunk prefill. + + kv_indptr/kv_indices are token-level for every page_size: + AiterIndicesUpdaterPrefill sets kv_indptr = cumsum(seq_lens) and writes one + kv_indices entry per token, so token t of sequence i lives in pool slot + kv_indices[kv_indptr[i] + t]. There is no page arithmetic to apply here. + + Returns (tok_idx, cu_seqlens_k), or None if the metadata disagrees with + seq_lens (mixed/spec batches) and the caller must use the paged kernel. + """ + bs = kv_indptr.numel() - 1 + device = kv_indices.device + kv_indptr = kv_indptr.to(torch.long) + seq_lens = seq_lens.to(device=device, dtype=torch.long) + # kvlen must not exceed the tokens this batch actually has in kv_indices, + # otherwise the gather runs off the end of the table. + seq_lens = torch.minimum(seq_lens, kv_indptr[1:] - kv_indptr[:bs]) + total_k = int(seq_lens.sum().item()) + cu_k = torch.zeros(bs + 1, dtype=torch.long, device=device) + torch.cumsum(seq_lens, 0, out=cu_k[1:]) + seq_ids = torch.repeat_interleave(torch.arange(bs, device=device), seq_lens) + pos_in_seq = torch.arange(total_k, device=device) - cu_k[seq_ids] + kv_slot = kv_indptr[seq_ids] + pos_in_seq + if total_k and int(kv_slot.max().item()) >= kv_indices.numel(): + logger.warning( + "[asm-context-prefill] metadata mismatch, falling back:" + " mode=%s bs=%s kv_slot_max=%s kv_indices=%s seq_lens=%s kv_indptr=%s", + forward_mode, + bs, + int(kv_slot.max().item()), + kv_indices.numel(), + seq_lens.tolist(), + kv_indptr.tolist(), + ) + return None + tok_idx = kv_indices[kv_slot].to(torch.long) + if total_k and int(tok_idx.max().item()) >= num_kv_slots: + logger.warning( + "[asm-context-prefill] gather index out of pool, falling back:" + " mode=%s bs=%s tok_idx_max=%s num_kv_slots=%s", + forward_mode, + bs, + int(tok_idx.max().item()), + num_kv_slots, + ) + return None + return tok_idx, cu_k + + class AiterAttnBackend(AttentionBackend): # kv_indptr/qo_indptr are preallocated at (req pool + 1); an extend batch @@ -2581,47 +2637,15 @@ class AiterAttnBackend(AttentionBackend): ): bs = forward_batch.batch_size k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(layer.layer_id) - page = self.page_size - kv_indptr = self.forward_metadata.kv_indptr[: bs + 1] - kv_pages = self.forward_metadata.kv_indices - seq_lens = forward_batch.seq_lens[:bs].to(torch.long) - # kvlen must not exceed the pages this batch actually has in - # kv_indices (metadata is page-granular for plain extend, but - # can disagree with seq_lens in mixed/spec batches -> OOB - # gather). Clamp per-seq kvlen to pages*page and fall back to - # the paged kernel on any inconsistency. - pages_per_seq = (kv_indptr[1 : bs + 1] - kv_indptr[:bs]).to(torch.long) - kvlen_cap = pages_per_seq * page - seq_lens = torch.minimum(seq_lens, kvlen_cap) - total_k = int(seq_lens.sum().item()) - cu_k = torch.zeros(bs + 1, dtype=torch.long, device=q.device) - torch.cumsum(seq_lens, 0, out=cu_k[1:]) - seq_ids = torch.repeat_interleave( - torch.arange(bs, device=q.device), seq_lens + 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, ) - pos_in_seq = torch.arange(total_k, device=q.device) - cu_k[seq_ids] - page_slot = kv_indptr[seq_ids].to(torch.long) + pos_in_seq // page - asm_cp_ok = bool(int(page_slot.max().item()) < kv_pages.numel()) - if not asm_cp_ok: - logger.warning( - "[asm-context-prefill] metadata mismatch, falling back:" - " mode=%s bs=%s page_slot_max=%s kv_pages=%s seq_lens=%s" - " kv_indptr=%s", - forward_batch.forward_mode, - bs, - int(page_slot.max().item()), - kv_pages.numel(), - seq_lens.tolist(), - kv_indptr.tolist(), - ) - if asm_cp_ok: - tok_idx = ( - kv_pages[page_slot].to(torch.long) * page + pos_in_seq % page - ) - asm_cp_ok = int(tok_idx.max().item()) < ( - self.token_to_kv_pool.get_key_buffer(layer.layer_id).shape[0] - ) - if asm_cp_ok: + if gathered is not None: + tok_idx, cu_k = gathered 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. diff --git a/test/registered/attention/test_aiter_asm_prefill_gather.py b/test/registered/attention/test_aiter_asm_prefill_gather.py new file mode 100644 index 000000000..1b4fdd81b --- /dev/null +++ b/test/registered/attention/test_aiter_asm_prefill_gather.py @@ -0,0 +1,106 @@ +"""Index arithmetic for the aiter ASM context-chunk prefill KV gather. + +kv_indptr/kv_indices are token-level for every page_size, so the gather must +resolve token t of sequence i to kv_indices[kv_indptr[i] + t]. Pure indexing, +so this runs on CPU. +""" + +import unittest + +import torch + +from sglang.srt.layers.attention.aiter_backend import ( + _asm_context_prefill_gather_indices, +) +from sglang.test.ci.ci_register import register_amd_ci + +register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd") + + +def _paged_pool(seq_lens, page_size, seed, headroom=2): + """Lay each sequence out over shuffled pages, as a page allocator would. + + Returns (req_to_token, num_kv_slots): req_to_token[i][t] is the pool slot + holding token t of sequence i. Pages are handed out in shuffled order, so a + correct gather cannot rely on sequences being contiguous in the pool. + """ + num_pages = headroom * sum((n + page_size - 1) // page_size for n in seq_lens) + g = torch.Generator().manual_seed(seed) + free_pages = torch.randperm(num_pages, generator=g).tolist() + req_to_token = [] + for n in seq_lens: + slots = [] + for _ in range((n + page_size - 1) // page_size): + base = free_pages.pop() * page_size + slots.extend(range(base, base + page_size)) + req_to_token.append(slots[:n]) + return req_to_token, num_pages * page_size + + +def _token_level_metadata(req_to_token): + """Build kv_indptr/kv_indices the way AiterIndicesUpdaterPrefill does.""" + seq_lens = torch.tensor([len(s) for s in req_to_token], dtype=torch.long) + kv_indptr = torch.zeros(len(req_to_token) + 1, dtype=torch.long) + torch.cumsum(seq_lens, 0, out=kv_indptr[1:]) + kv_indices = torch.tensor( + [slot for slots in req_to_token for slot in slots], dtype=torch.int32 + ) + return kv_indptr, kv_indices, seq_lens + + +class TestAsmPrefillGatherIndices(unittest.TestCase): + def _check(self, seq_lens, page_size, seed=0xA17E4): + req_to_token, num_kv_slots = _paged_pool(seq_lens, page_size, seed) + kv_indptr, kv_indices, lens = _token_level_metadata(req_to_token) + + gathered = _asm_context_prefill_gather_indices( + kv_indptr, kv_indices, lens, num_kv_slots + ) + self.assertIsNotNone(gathered, "gather rejected valid metadata") + tok_idx, cu_k = gathered + + expected = torch.tensor( + [slot for slots in req_to_token for slot in slots], dtype=torch.long + ) + self.assertEqual(tok_idx.tolist(), expected.tolist()) + self.assertEqual(cu_k.tolist(), kv_indptr.tolist()) + + def test_page_sizes(self): + # Chunked prefill shapes: several sequences with a long prefix. + for page_size in (1, 16, 64): + with self.subTest(page_size=page_size): + self._check([43616, 1024, 512], page_size) + + def test_single_sequence(self): + for page_size in (1, 16, 64): + with self.subTest(page_size=page_size): + self._check([2048], page_size) + + def test_unaligned_seq_lens(self): + for page_size in (16, 64): + with self.subTest(page_size=page_size): + self._check([1000, 33, 1], page_size) + + def test_falls_back_when_metadata_is_short(self): + # seq_lens longer than kv_indices has entries for (mixed/spec batches): + # clamped, not an out-of-bounds gather. + req_to_token, num_kv_slots = _paged_pool([512, 512], 64, seed=7) + kv_indptr, kv_indices, lens = _token_level_metadata(req_to_token) + gathered = _asm_context_prefill_gather_indices( + kv_indptr, kv_indices, lens + 128, num_kv_slots + ) + self.assertIsNotNone(gathered) + tok_idx, cu_k = gathered + self.assertEqual(cu_k.tolist(), kv_indptr.tolist()) + self.assertEqual(int(tok_idx.numel()), int(lens.sum())) + + def test_falls_back_when_slot_exceeds_pool(self): + req_to_token, _ = _paged_pool([512, 512], 64, seed=11) + kv_indptr, kv_indices, lens = _token_level_metadata(req_to_token) + self.assertIsNone( + _asm_context_prefill_gather_indices(kv_indptr, kv_indices, lens, 16) + ) + + +if __name__ == "__main__": + unittest.main()