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

flashinfer swa kv pool fix (dflash gemma 4) (#27737)

Browse Source
This commit is contained in:
David Wang
2026-06-12 11:35:05 -07:00
committed by GitHub
parent fa4273d2db
commit bb33594c1a
3 changed files with 96 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/spec/test_resolve_swa_kv_pool.py
+52 -33
View File
@@ -1,9 +1,11 @@
"""Unit tests for TRTLLMHAAttnBackend._resolve_swa_kv_pool."""
"""Unit tests for attention-backend SWA KV pool resolution."""
import unittest
from unittest.mock import MagicMock
from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend
from sglang.srt.layers.attention.trtllm_mha_backend import TRTLLMHAAttnBackend
from sglang.srt.mem_cache.base_swa_memory_pool import BaseSWAKVPool
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
@@ -11,7 +13,10 @@ 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
_RESOLVERS = (
("trtllm_mha", TRTLLMHAAttnBackend._resolve_swa_kv_pool, SWAKVPool),
("flashinfer", FlashInferAttnBackend._resolve_swa_kv_pool, BaseSWAKVPool),
)
def _mock_runner(
@@ -31,47 +36,61 @@ def _mock_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)
for name, resolve, pool_type in _RESOLVERS:
with self.subTest(backend=name):
swa = MagicMock(spec=pool_type)
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)
for name, resolve, pool_type in _RESOLVERS:
with self.subTest(backend=name):
swa = MagicMock(spec=pool_type)
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))
for name, resolve, _ in _RESOLVERS:
with self.subTest(backend=name):
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))
for name, resolve, pool_type in _RESOLVERS:
with self.subTest(backend=name):
runner = _mock_runner(
active_pool=MagicMock(),
is_draft_worker=True,
spec_algorithm=SpeculativeAlgorithm.EAGLE,
allocator_kvcache=MagicMock(spec=pool_type),
)
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)
for name, resolve, pool_type in _RESOLVERS:
with self.subTest(backend=name):
swa = MagicMock(spec=pool_type)
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)
for name, resolve, pool_type in _RESOLVERS:
with self.subTest(backend=name):
swa = MagicMock(spec=pool_type)
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__":