89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
"""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()
|