Fix SWA pool resolution for EAGLE draft workers (#27491)

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
+        # separate index spaces; SWA layers need a translated page_table.
-        # the pool from the allocator (stable at construction), not from
-        # token_to_kv_pool, which FROZEN_KV MTP swaps per forward call.
         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
+        EAGLE draft workers share the target allocator for token bookkeeping,
-        target's SWA allocator while its own token_to_kv_pool stays non-SWA
+        but own a separate draft KV pool. Do not use the target allocator's
-        until swapped per call. The getattr only tolerates the minimal
+        SWA mapping for that draft pool. FROZEN_KV MTP is the exception: its
-        allocator stub used by attention test fixtures.
+        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 = allocator.get_kvcache()
-        kvcache = get_kvcache() if get_kvcache is not None else None
         return kvcache if isinstance(kvcache, SWAKVPool) else None
 
     def _maybe_translate_swa(

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

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()