From fca4ef9d698882cfaa84e65b7fdbb84ea9d5a5e6 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 8 Jun 2026 11:00:29 -0700 Subject: [PATCH] Fix SWA pool resolution for EAGLE draft workers (#27491) --- .../layers/attention/trtllm_mha_backend.py | 24 +++--- .../attention_methods/dense_attention.py | 8 +- .../unit/spec/test_resolve_swa_kv_pool.py | 78 +++++++++++++++++++ 3 files changed, 100 insertions(+), 10 deletions(-) create mode 100644 test/registered/unit/spec/test_resolve_swa_kv_pool.py diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 17f565ef1..e8e6634cf 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -125,9 +125,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): ) # SWA hybrid models split the KV cache into full and SWA pools with - # separate index spaces; SWA layers need a translated page_table. Resolve - # the pool from the allocator (stable at construction), not from - # token_to_kv_pool, which FROZEN_KV MTP swaps per forward call. + # separate index spaces; SWA layers need a translated page_table. self._swa_kv_pool: Optional[SWAKVPool] = self._resolve_swa_kv_pool(model_runner) # Forward metadata @@ -147,14 +145,22 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): def _resolve_swa_kv_pool(model_runner: ModelRunner) -> Optional[SWAKVPool]: """Return the SWAKVPool to translate against, or None for non-SWA models. - Read it from the allocator: in FROZEN_KV MTP the draft shares the - target's SWA allocator while its own token_to_kv_pool stays non-SWA - until swapped per call. The getattr only tolerates the minimal - allocator stub used by attention test fixtures. + EAGLE draft workers share the target allocator for token bookkeeping, + but own a separate draft KV pool. Do not use the target allocator's + SWA mapping for that draft pool. 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 SWA. """ + active_pool = model_runner.token_to_kv_pool + if isinstance(active_pool, SWAKVPool): + return active_pool + + if model_runner.is_draft_worker: + if not model_runner.spec_algorithm.is_frozen_kv_mtp(): + return None + allocator = model_runner.token_to_kv_pool_allocator - get_kvcache = getattr(allocator, "get_kvcache", None) - kvcache = get_kvcache() if get_kvcache is not None else None + kvcache = allocator.get_kvcache() return kvcache if isinstance(kvcache, SWAKVPool) else None def _maybe_translate_swa( diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index b26a3f4cd..84863d063 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py @@ -15,6 +15,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMo from sglang.srt.model_executor.forward_context import ForwardContext, forward_context from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.server_args import set_global_server_args_for_scheduler +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from ..mock_server_args import make_mock_server_args @@ -316,6 +317,8 @@ class MockModelRunner(ModelRunner): self.tp_size = 1 self.dp_size = 1 self.pp_size = 1 + self.is_draft_worker = False + self.spec_algorithm = SpeculativeAlgorithm.NONE speculative_num_draft_tokens = ( max(case.input_lens) if case.forward_mode.is_target_verify() @@ -367,7 +370,10 @@ class MockModelRunner(ModelRunner): enable_memory_saver=False, enable_alt_stream=False, ) - self.token_to_kv_pool_allocator = SimpleNamespace(page_size=case.page_size) + self.token_to_kv_pool_allocator = SimpleNamespace( + page_size=case.page_size, + get_kvcache=lambda: self.token_to_kv_pool, + ) self.attn_cp_size = 1 self.attention_chunk_size = None self.hisparse_coordinator = None diff --git a/test/registered/unit/spec/test_resolve_swa_kv_pool.py b/test/registered/unit/spec/test_resolve_swa_kv_pool.py new file mode 100644 index 000000000..d97b4fd14 --- /dev/null +++ b/test/registered/unit/spec/test_resolve_swa_kv_pool.py @@ -0,0 +1,78 @@ +"""Unit tests for TRTLLMHAAttnBackend._resolve_swa_kv_pool.""" + +import unittest +from unittest.mock import MagicMock + +from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-large") + +_resolve = TRTLLMHAAttnBackend._resolve_swa_kv_pool + + +def _mock_runner( + *, + active_pool=None, + is_draft_worker=False, + spec_algorithm=SpeculativeAlgorithm.NONE, + allocator_kvcache=None, +): + runner = MagicMock() + runner.token_to_kv_pool = active_pool + runner.is_draft_worker = is_draft_worker + runner.spec_algorithm = spec_algorithm + runner.token_to_kv_pool_allocator.get_kvcache.return_value = allocator_kvcache + return runner + + +class TestResolveSwaKvPool(CustomTestCase): + def test_active_pool_is_swa_returns_it(self): + swa = MagicMock(spec=SWAKVPool) + runner = _mock_runner(active_pool=swa) + self.assertIs(_resolve(runner), swa) + + def test_non_swa_active_pool_falls_through_to_allocator(self): + swa = MagicMock(spec=SWAKVPool) + runner = _mock_runner(active_pool=MagicMock(), allocator_kvcache=swa) + self.assertIs(_resolve(runner), swa) + + def test_allocator_kvcache_not_swa_returns_none(self): + runner = _mock_runner(active_pool=MagicMock(), allocator_kvcache=MagicMock()) + self.assertIsNone(_resolve(runner)) + + def test_draft_worker_non_frozen_kv_returns_none(self): + runner = _mock_runner( + active_pool=MagicMock(), + is_draft_worker=True, + spec_algorithm=SpeculativeAlgorithm.EAGLE, + allocator_kvcache=MagicMock(spec=SWAKVPool), + ) + self.assertIsNone(_resolve(runner)) + + def test_draft_worker_frozen_kv_mtp_returns_allocator_swa(self): + swa = MagicMock(spec=SWAKVPool) + runner = _mock_runner( + active_pool=MagicMock(), + is_draft_worker=True, + spec_algorithm=SpeculativeAlgorithm.FROZEN_KV_MTP, + allocator_kvcache=swa, + ) + self.assertIs(_resolve(runner), swa) + + def test_non_draft_worker_ignores_spec_algorithm(self): + swa = MagicMock(spec=SWAKVPool) + runner = _mock_runner( + active_pool=MagicMock(), + is_draft_worker=False, + spec_algorithm=SpeculativeAlgorithm.EAGLE, + allocator_kvcache=swa, + ) + self.assertIs(_resolve(runner), swa) + + +if __name__ == "__main__": + unittest.main()