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

[AMD] aiter: resolve SWA KV pool for draft workers + guard paged decode (#38756)

Browse Source
This commit is contained in:
Vignesh Sethuraman
2026-09-11 10:13:48 -07:00
committed by GitHub
parent 4309c7ce19
commit ddf02a4f58
2 changed files with 145 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/aiter_backend.py
+57 -15
View File
@@ -349,10 +349,14 @@ class AiterAttnBackend(AttentionBackend):
self.req_to_token_pool = model_runner.req_to_token_pool
self.token_to_kv_pool = model_runner.token_to_kv_pool
# sliding window attention
# sliding window attention. Resolve the SWA pool rather than reading it
# straight off the active pool: a frozen-KV MTP draft worker's active
# pool is its own draft pool, but its draft path reads target KV, so the
# SWA mapping must still come from the target allocator. Mirrors
# TRTLLMHAAttnBackend._resolve_swa_kv_pool.
self.swa_kv_pool = self._resolve_swa_kv_pool(model_runner)
self.use_sliding_window_kv_pool = (
isinstance(model_runner.token_to_kv_pool, SWAKVPool)
and model_runner.token_to_kv_pool.swa_layer_nums > 0
self.swa_kv_pool is not None and self.swa_kv_pool.swa_layer_nums > 0
)
# Detect SHUFFLE 5D ("vectorized") KV cache layout. When active
@@ -783,7 +787,7 @@ class AiterAttnBackend(AttentionBackend):
)
if self.use_sliding_window_kv_pool:
swa_slot_mapping = self.token_to_kv_pool.full_to_swa_index_mapping.long()
swa_slot_mapping = self.swa_kv_pool.full_to_swa_index_mapping.long()
if swa_dest_buf is not None:
swa_page_table = swa_dest_buf
@@ -845,7 +849,7 @@ class AiterAttnBackend(AttentionBackend):
page_table = torch.zeros(bs, max_blocks, dtype=torch.int32, device=device)
if self.use_sliding_window_kv_pool:
swa_slot_mapping = self.token_to_kv_pool.full_to_swa_index_mapping.long()
swa_slot_mapping = self.swa_kv_pool.full_to_swa_index_mapping.long()
if swa_page_table_dest is not None:
swa_page_table = swa_page_table_dest
@@ -1204,7 +1208,7 @@ class AiterAttnBackend(AttentionBackend):
self.cuda_graph_swa_out_cache_loc[:n].zero_()
else:
self.cuda_graph_swa_out_cache_loc[:n].copy_(
self.token_to_kv_pool.translate_loc_from_full_to_swa(
self.swa_kv_pool.translate_loc_from_full_to_swa(
forward_batch.out_cache_loc
)
)
@@ -1234,7 +1238,7 @@ class AiterAttnBackend(AttentionBackend):
swa_page_table = None
swa_out_cache_loc = None
if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
swa_out_cache_loc = self.token_to_kv_pool.translate_loc_from_full_to_swa(
swa_out_cache_loc = self.swa_kv_pool.translate_loc_from_full_to_swa(
forward_batch.out_cache_loc
)
max_kv_len = forward_batch.seq_lens_cpu.max().item()
@@ -1282,7 +1286,7 @@ class AiterAttnBackend(AttentionBackend):
# AITER attention kernels require int32 page indices;
# full_to_swa_index_mapping is stored as int64.
swa_page_table = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
self.swa_kv_pool.translate_loc_from_full_to_swa(
kv_indices
).to(torch.int32)
)
@@ -1693,11 +1697,9 @@ class AiterAttnBackend(AttentionBackend):
# AITER attention kernels (e.g. mha_batch_prefill_func)
# require int32 page indices; full_to_swa_index_mapping is
# stored as int64.
swa_page_table = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
self.indices_updater_prefill.kv_indices
).to(torch.int32)
)
swa_page_table = self.swa_kv_pool.translate_loc_from_full_to_swa(
self.indices_updater_prefill.kv_indices
).to(torch.int32)
self.forward_metadata = ForwardMetadata(
self.indices_updater_prefill.kv_indptr,
@@ -1908,7 +1910,7 @@ class AiterAttnBackend(AttentionBackend):
# AITER attention kernels require int32 page indices;
# full_to_swa_index_mapping is stored as int64.
swa_page_indices = (
self.token_to_kv_pool.translate_loc_from_full_to_swa(
self.swa_kv_pool.translate_loc_from_full_to_swa(
page_indices
).to(torch.int32)
)
@@ -2302,6 +2304,45 @@ class AiterAttnBackend(AttentionBackend):
"(speculative_eagle_topk=1 and SGLANG_AITER_UNIFIED_VERIFY=1)."
)
@staticmethod
def _resolve_swa_kv_pool(model_runner):
"""Return the SWAKVPool to translate against, or None for non-SWA models.
EAGLE draft workers share the target allocator for token bookkeeping but
own a separate draft KV pool, so the target allocator's SWA mapping must
not be used for them. FROZEN_KV MTP is the exception: its draft path reads
target KV directly, so it still needs the allocator pool when the active
pool is not itself an SWAKVPool. Mirrors
``TRTLLMHAAttnBackend._resolve_swa_kv_pool``.
"""
active_pool = model_runner.token_to_kv_pool
if isinstance(active_pool, SWAKVPool):
return active_pool
if getattr(model_runner, "is_draft_worker", False):
if not model_runner.spec_algorithm.is_frozen_kv_mtp():
return None
kvcache = model_runner.token_to_kv_pool_allocator.get_kvcache()
return kvcache if isinstance(kvcache, SWAKVPool) else None
@staticmethod
def _reject_paged_decode_sliding_window(layer):
"""Reject sliding-window layers on the aiter paged-decode path.
``paged_attention_ragged`` takes no sliding-window argument, so a
sliding-window layer routed to it would attend over the full context and
silently return wrong results. Raise instead of silently dropping the
window. Layers with ``sliding_window_size`` unset or -1 are unaffected.
"""
if layer.sliding_window_size is not None and layer.sliding_window_size > -1:
raise ValueError(
"aiter paged decode cannot honor sliding-window "
f"attention (layer {layer.layer_id} has "
f"sliding_window_size={layer.sliding_window_size}). "
"Enable the unified attention path "
"(SGLANG_USE_AITER_UNIFIED_ATTN=1) or select a "
"different attention backend."
)
def forward_extend(
self,
q: torch.Tensor,
@@ -2357,7 +2398,7 @@ class AiterAttnBackend(AttentionBackend):
k_cache, v_cache = self.token_to_kv_pool.get_kv_buffer(
layer.layer_id
)
slot_mapping_swa = token_to_kv_pool.full_to_swa_index_mapping
slot_mapping_swa = self.swa_kv_pool.full_to_swa_index_mapping
launch_reshape_and_cache_flash(
k.view(-1, layer.tp_k_head_num, layer.qk_head_dim),
@@ -3310,6 +3351,7 @@ class AiterAttnBackend(AttentionBackend):
sinks=sinks,
)
else:
self._reject_paged_decode_sliding_window(layer)
# Drop FP8 KV upcast: keep paged cache in native FP8 and use ``fp8_e4m3`` for
# in-kernel dequant in ``paged_attention_ragged``. (HIP maps CLI e5m2/e4m3 to
# ``fp8_dtype``; aiter has no ``fp8_e5m2`` string.)
test/srt/test_aiter_swa_kv_pool_resolver.py
+88
View File
@@ -0,0 +1,88 @@
"""PR3: aiter SWA KV-pool resolver + paged-decode sliding-window guard.
``_resolve_swa_kv_pool`` returns the active pool for every non-draft worker
(behaviour-preserving), skips the target SWA mapping for EAGLE draft workers
(which own a separate draft pool), and falls back to the allocator's KV cache
for FROZEN_KV MTP. ``_reject_paged_decode_sliding_window`` fails loudly when a
sliding-window layer would reach ``paged_attention_ragged`` (which has no window
argument) instead of silently returning full-context (wrong) results.
Both are pure branching logic, exercised here with mocks (no GPU kernel).
"""
import unittest
from types import SimpleNamespace
from unittest.mock import MagicMock
from sglang.srt.layers.attention.aiter_backend import AiterAttnBackend
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.test_utils import CustomTestCase
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd-mi35x")
def _model_runner(active_pool, *, is_draft=False, frozen_kv_mtp=False, alloc_pool=None):
return SimpleNamespace(
token_to_kv_pool=active_pool,
is_draft_worker=is_draft,
spec_algorithm=SimpleNamespace(is_frozen_kv_mtp=lambda: frozen_kv_mtp),
token_to_kv_pool_allocator=SimpleNamespace(get_kvcache=lambda: alloc_pool),
)
class TestResolveSwaKvPool(CustomTestCase):
def test_active_pool_is_swa_returns_it(self):
swa = MagicMock(spec=SWAKVPool)
mr = _model_runner(swa)
self.assertIs(AiterAttnBackend._resolve_swa_kv_pool(mr), swa)
def test_non_draft_non_swa_active_falls_back_to_allocator(self):
alloc_swa = MagicMock(spec=SWAKVPool)
mr = _model_runner(object(), alloc_pool=alloc_swa)
self.assertIs(AiterAttnBackend._resolve_swa_kv_pool(mr), alloc_swa)
def test_non_draft_no_swa_anywhere_returns_none(self):
mr = _model_runner(object(), alloc_pool=object())
self.assertIsNone(AiterAttnBackend._resolve_swa_kv_pool(mr))
def test_eagle_draft_worker_skips_target_mapping(self):
# Draft worker that is NOT frozen-KV MTP owns its own draft pool -> None,
# even if the allocator holds an SWA pool.
mr = _model_runner(
object(),
is_draft=True,
frozen_kv_mtp=False,
alloc_pool=MagicMock(spec=SWAKVPool),
)
self.assertIsNone(AiterAttnBackend._resolve_swa_kv_pool(mr))
def test_frozen_kv_mtp_draft_worker_uses_allocator_pool(self):
alloc_swa = MagicMock(spec=SWAKVPool)
mr = _model_runner(
object(), is_draft=True, frozen_kv_mtp=True, alloc_pool=alloc_swa
)
self.assertIs(AiterAttnBackend._resolve_swa_kv_pool(mr), alloc_swa)
class TestPagedDecodeSlidingWindowGuard(CustomTestCase):
def test_raises_for_sliding_window_layer(self):
layer = SimpleNamespace(sliding_window_size=1024, layer_id=7)
with self.assertRaises(ValueError) as ctx:
AiterAttnBackend._reject_paged_decode_sliding_window(layer)
msg = str(ctx.exception)
self.assertIn("sliding-window", msg)
self.assertIn("SGLANG_USE_AITER_UNIFIED_ATTN=1", msg)
self.assertIn("layer 7", msg)
def test_inert_when_window_unset(self):
layer = SimpleNamespace(sliding_window_size=None, layer_id=0)
AiterAttnBackend._reject_paged_decode_sliding_window(layer)
def test_inert_when_window_is_negative_one(self):
layer = SimpleNamespace(sliding_window_size=-1, layer_id=0)
AiterAttnBackend._reject_paged_decode_sliding_window(layer)
if __name__ == "__main__":
unittest.main()