From 03940fc071d27c4a2e04c867c79ac9a6a9062855 Mon Sep 17 00:00:00 2001 From: Zijie Chen <300606707+zijiecode@users.noreply.github.com> Date: Wed, 16 Sep 2026 21:47:38 -0700 Subject: [PATCH 1/4] Reuse ASM context-prefill gather indices across attention layers --- .../srt/layers/attention/aiter_backend.py | 87 +++++++++-- ...test_asm_context_prefill_gather_indices.py | 138 ++++++++++++++++++ 2 files changed, 216 insertions(+), 9 deletions(-) create mode 100644 test/registered/amd/test_asm_context_prefill_gather_indices.py diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 31d93871ecd7..86cc37b22210 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/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,10 @@ 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). +_ASM_CTX_GATHER_CHECK = os.environ.get("SGLANG_ASM_CTX_GATHER_CHECK", "0") == "1" + + def _asm_context_prefill_gather_indices( kv_indptr: torch.Tensor, kv_indices: torch.Tensor, @@ -1059,6 +1069,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 _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: @@ -3492,15 +3566,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 +3600,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, diff --git a/test/registered/amd/test_asm_context_prefill_gather_indices.py b/test/registered/amd/test_asm_context_prefill_gather_indices.py new file mode 100644 index 000000000000..5bf652f39900 --- /dev/null +++ b/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() From 3c6350ccf863ddcc10eca20d0d969543177a921d Mon Sep 17 00:00:00 2001 From: Zijie Chen <300606707+zijiecode@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:54:02 -0700 Subject: [PATCH 2/4] Address ASM context prefill review comments --- python/sglang/srt/layers/attention/aiter_backend.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 86cc37b22210..c36c3692213f 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -185,7 +185,7 @@ def _aiter_fp8_asm_supports_gqa(num_q_heads: int, num_kv_heads: int) -> bool: # Cross-check the per-batch fast indices against the generic gather (syncs). -_ASM_CTX_GATHER_CHECK = os.environ.get("SGLANG_ASM_CTX_GATHER_CHECK", "0") == "1" +_GFX_ASM_CTX_GATHER_CHECK = os.environ.get("SGLANG_ASM_CTX_GATHER_CHECK", "0") == "1" def _asm_context_prefill_gather_indices( @@ -1098,7 +1098,7 @@ def _asm_context_prefill_indices( cu_k = fm.kv_indptr[: bs + 1] if cu_k.dtype != torch.int32: cu_k = cu_k.to(torch.int32) - if _ASM_CTX_GATHER_CHECK: + if _GFX_ASM_CTX_GATHER_CHECK: ref = _asm_context_prefill_gather_indices( fm.kv_indptr[: bs + 1], fm.kv_indices, @@ -3547,6 +3547,10 @@ 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 and any other model + # meeting the guards below: gfx950, 256-dim Q/K/V, supported GQA, + # non-vectorized FP8 KV, and no sliding window, sinks, or soft cap. + # Other configurations fall through to the attention paths below. if ( is_gfx95_supported() and forward_batch.forward_mode.is_extend() From 8a35b16b94486fea9c90ea245a8fa7fa31c8c5f2 Mon Sep 17 00:00:00 2001 From: Zijie Chen <300606707+zijiecode@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:01:43 -0700 Subject: [PATCH 3/4] Clarify current Qwen3.5 ASM prefill scope --- python/sglang/srt/layers/attention/aiter_backend.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index c36c3692213f..56ce77125c93 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -3547,9 +3547,7 @@ 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 and any other model - # meeting the guards below: gfx950, 256-dim Q/K/V, supported GQA, - # non-vectorized FP8 KV, and no sliding window, sinks, or soft cap. + # This applies to Qwen3.5 full-attention layers only currently, # Other configurations fall through to the attention paths below. if ( is_gfx95_supported() From 6294712014daf3811d167f73cadcc92dc2f4252b Mon Sep 17 00:00:00 2001 From: Zijie Chen <300606707+zijiecode@users.noreply.github.com> Date: Tue, 22 Sep 2026 08:56:19 -0700 Subject: [PATCH 4/4] fix(amd): rename ASM context gather check environment variable --- python/sglang/srt/layers/attention/aiter_backend.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 56ce77125c93..4809bbcb1342 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -185,7 +185,9 @@ def _aiter_fp8_asm_supports_gqa(num_q_heads: int, num_kv_heads: int) -> bool: # Cross-check the per-batch fast indices against the generic gather (syncs). -_GFX_ASM_CTX_GATHER_CHECK = os.environ.get("SGLANG_ASM_CTX_GATHER_CHECK", "0") == "1" +_GFX_ASM_CTX_GATHER_CHECK = ( + os.environ.get("SGLANG_GFX_ASM_CTX_GATHER_CHECK", "0") == "1" +) def _asm_context_prefill_gather_indices(