minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[ROCm][Bugfix] Use token-level KV indices in the aiter ASM context-prefill gather (#36852)

Browse Source
This commit is contained in:
Olga Miroshnichenko
2026-08-28 17:17:18 -07:00
committed by GitHub
parent 4944e50e2c
commit 36afd70f9c
2 changed files with 170 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/aiter_backend.py
+64 -40
View File
@@ -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.
test/registered/attention/test_aiter_asm_prefill_gather.py
+106
View File
@@ -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()