From ddf02a4f5851edc02c9d5387ba4dedf3240e3bee Mon Sep 17 00:00:00 2001 From: Vignesh Sethuraman Date: Fri, 11 Sep 2026 10:13:48 -0700 Subject: [PATCH] [AMD] aiter: resolve SWA KV pool for draft workers + guard paged decode (#38756) --- .../srt/layers/attention/aiter_backend.py | 72 +++++++++++---- test/srt/test_aiter_swa_kv_pool_resolver.py | 88 +++++++++++++++++++ 2 files changed, 145 insertions(+), 15 deletions(-) create mode 100644 test/srt/test_aiter_swa_kv_pool_resolver.py diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 09b720e2f..7abaa12ed 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -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.) diff --git a/test/srt/test_aiter_swa_kv_pool_resolver.py b/test/srt/test_aiter_swa_kv_pool_resolver.py new file mode 100644 index 000000000..d9ca4f06a --- /dev/null +++ b/test/srt/test_aiter_swa_kv_pool_resolver.py @@ -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()