diff --git a/python/sglang/srt/kv_canary/pool_patcher/adapters/dsv4.py b/python/sglang/srt/kv_canary/pool_patcher/adapters/dsv4.py new file mode 100644 index 000000000..4204e2dc6 --- /dev/null +++ b/python/sglang/srt/kv_canary/pool_patcher/adapters/dsv4.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +import torch + +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind +from sglang.srt.kv_canary.pool_patcher.buf_info_splice import patch_buf_info_method +from sglang.srt.kv_canary.pool_patcher.buffer_alloc import alloc_canary_buf + + +def attach_dsv4( + *, + pool: object, + device: torch.device, + kv_token_id_vs_position_offset: int, +) -> tuple[CanaryBufferGroup, ...]: + """Attach canary buffers to a DeepSeekV4TokenToKVPool. + + TODO: only the swa_kv_pool sub-pool is wired; c4_kv_pool / c128_kv_pool / + c4_indexer_kv_pool / compress state pools are left uncovered. + """ + sub_pool = pool.swa_kv_pool + num_slots = int(sub_pool.size) + + k_head = alloc_canary_buf(num_slots=num_slots, device=device) + k_tail = alloc_canary_buf(num_slots=num_slots, device=device) + + group = CanaryBufferGroup( + kind=PoolKind.SWA, + k_head=k_head, + k_tail=k_tail, + v_head=None, + v_tail=None, + swa_index_lut=pool.full_to_swa_index_mapping, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + + patch_buf_info_method( + pool, + method_name="get_state_buf_infos", + group=group, + has_v_half=False, + page_size=sub_pool.page_size, + ) + + return (group,) diff --git a/python/sglang/srt/kv_canary/pool_patcher/adapters/swa.py b/python/sglang/srt/kv_canary/pool_patcher/adapters/swa.py new file mode 100644 index 000000000..d9012b0c3 --- /dev/null +++ b/python/sglang/srt/kv_canary/pool_patcher/adapters/swa.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from typing import Optional + +import torch + +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind +from sglang.srt.kv_canary.pool_patcher.buf_info_splice import patch_buf_info_method +from sglang.srt.kv_canary.pool_patcher.buffer_alloc import alloc_canary_buf + + +def attach_swa( + *, + pool: object, + device: torch.device, + kv_token_id_vs_position_offset: int, +) -> tuple[CanaryBufferGroup, ...]: + full_group = _build_subpool_group( + sub_pool=pool.full_kv_pool, + kind=PoolKind.FULL, + device=device, + swa_lut=None, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + swa_group = _build_subpool_group( + sub_pool=pool.swa_kv_pool, + kind=PoolKind.SWA, + device=device, + swa_lut=pool.full_to_swa_index_mapping, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) + + patch_buf_info_method( + pool, + method_name="get_contiguous_buf_infos", + group=full_group, + has_v_half=True, + page_size=pool.page_size, + ) + patch_buf_info_method( + pool, + method_name="get_state_buf_infos", + group=swa_group, + has_v_half=True, + page_size=pool.page_size, + ) + return (full_group, swa_group) + + +def _build_subpool_group( + *, + sub_pool: object, + kind: PoolKind, + device: torch.device, + swa_lut: Optional[torch.Tensor], + kv_token_id_vs_position_offset: int, +) -> CanaryBufferGroup: + num_slots = int(sub_pool.k_buffer[0].shape[0]) + k_head = alloc_canary_buf(num_slots=num_slots, device=device) + k_tail = alloc_canary_buf(num_slots=num_slots, device=device) + v_head = alloc_canary_buf(num_slots=num_slots, device=device) + v_tail = alloc_canary_buf(num_slots=num_slots, device=device) + return CanaryBufferGroup( + kind=kind, + k_head=k_head, + k_tail=k_tail, + v_head=v_head, + v_tail=v_tail, + swa_index_lut=swa_lut, + kv_token_id_vs_position_offset=kv_token_id_vs_position_offset, + ) diff --git a/python/sglang/srt/kv_canary/pool_patcher/api.py b/python/sglang/srt/kv_canary/pool_patcher/api.py index 90a9a217d..6c1438412 100644 --- a/python/sglang/srt/kv_canary/pool_patcher/api.py +++ b/python/sglang/srt/kv_canary/pool_patcher/api.py @@ -7,12 +7,16 @@ import torch from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup from sglang.srt.kv_canary.config import CanaryConfig +from sglang.srt.kv_canary.pool_patcher.adapters.dsv4 import attach_dsv4 from sglang.srt.kv_canary.pool_patcher.adapters.mha import attach_mha +from sglang.srt.kv_canary.pool_patcher.adapters.swa import attach_swa +from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.mem_cache.memory_pool import ( KVCache, MHATokenToKVPool, MHATokenToKVPoolFP4, ) +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool logger = logging.getLogger(__name__) @@ -21,6 +25,8 @@ PoolAttacher = Callable[..., tuple[CanaryBufferGroup, ...]] _POOL_ATTACHERS: Dict[Type, PoolAttacher] = { MHATokenToKVPool: attach_mha, MHATokenToKVPoolFP4: attach_mha, + SWAKVPool: attach_swa, + DeepSeekV4TokenToKVPool: attach_dsv4, } diff --git a/python/sglang/test/kv_canary/consts.py b/python/sglang/test/kv_canary/consts.py new file mode 100644 index 000000000..152c9c178 --- /dev/null +++ b/python/sglang/test/kv_canary/consts.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from typing import Final + +# SWA e2e pool sizing for 8 reqs × ~7K prompt + 2K decode, SWA window 1024. +# FULL pool = max-total-tokens; must fit 8 × (7000 + 2048) = 72_384 to avoid preempt. +# SWA pool = max-total-tokens × ratio; +# ≥ 8 × 1024 = 8192 (else deadlock — in-flight footprint exceeds capacity); +# < 8 × 7000 = 56_000 (else allocator never recycles → swa_full_idx_divergence stays 0). +# Pick FULL=81920 (72_384 + headroom), SWA=16384 (2× in-flight floor, well under 56K). +SWA_POOL_SERVER_ARGS: Final[tuple[str, ...]] = ( + "--max-total-tokens", + "81920", + "--swa-full-tokens-ratio", + "0.2", +) + +DSV4_DEEPEP_CONFIG: Final[str] = ( + '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}' +) + +DSV4_POOL_SERVER_ARGS: Final[tuple[str, ...]] = ( + "--trust-remote-code", + "--tp", + "4", + "--dp", + "4", + "--enable-dp-attention", + "--moe-a2a-backend", + "deepep", + "--cuda-graph-max-bs", + "128", + "--max-running-requests", + "256", + "--deepep-config", + DSV4_DEEPEP_CONFIG, + "--mem-fraction-static", + "0.7", + "--speculative-algorithm", + "EAGLE", + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", +) + +DSV4_POOL_SERVER_ENV: Final[dict[str, str]] = { + "SGLANG_DSV4_FP4_EXPERTS": "0", + "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024", +} diff --git a/python/sglang/test/kv_canary/e2e_base.py b/python/sglang/test/kv_canary/e2e_base.py index 21d244e87..0c6cf059b 100644 --- a/python/sglang/test/kv_canary/e2e_base.py +++ b/python/sglang/test/kv_canary/e2e_base.py @@ -18,7 +18,9 @@ from sglang.test.test_utils import ( ) # Long prompt body shared by all canary e2e tests. The repetition count is chosen -# so the tokenised prompt is comfortably long; token count is roughly 7k after BPE. +# so the tokenised prompt comfortably exceeds the SWA sliding window of swa-mode +# fixtures (gemma-4-E2B); short prompts would never exercise the SWA-windowed +# verify path. Token count is roughly 7k after BPE. _LONG_PROMPT_BODY = ("The quick brown fox jumps over the lazy dog. " * 700).strip() _UNIQUE_PROMPT_FIRST_CHARS = string.ascii_letters + string.digits @@ -60,12 +62,14 @@ class CapturedServerE2EBase(CanaryViolationAssertMixin, CustomTestCase): class CanaryE2EBase(CapturedServerE2EBase): - model_mode: ClassVar[Literal["mha"]] + model_mode: ClassVar[Literal["mha", "swa", "dsv4"]] kv_canary_mode: ClassVar[CanaryMode] extra_env: ClassVar[dict[str, str]] = {} extra_server_args: ClassVar[tuple[str, ...]] = () use_unique_prompts: ClassVar[bool] = False - # Number of sequential request batches each test method sends. Default 1 keeps tests fast. + # SWA divergence assertions need slot recycling across batches; setting > 1 makes the + # test methods send N sequential batches so the SWA allocator's full→swa index mapping + # diverges from identity. Default 1 keeps MHA tests fast. workload_n_batches: ClassVar[int] = 1 _cfg: ClassVar[Optional[_ModeConfig]] = None @@ -75,6 +79,12 @@ class CanaryE2EBase(CapturedServerE2EBase): cls._cfg = _MODE_CONFIGS[cls.model_mode] server_env = os.environ.copy() server_env.update(cls.extra_env) + if cls.model_mode == "swa": + # SWA mode uses google/gemma-4-E2B-it, whose forward does a + # ``positions += 1`` in-place. canary's WRITE/VERIFY require + # forward_batch.positions to stay 0-indexed, so flip the gemma + # path to out-of-place shift for these tests. + server_env.setdefault("SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION", "1") cls._stdout_buf = io.StringIO() cls._stderr_buf = io.StringIO() diff --git a/python/sglang/test/kv_canary/fixtures.py b/python/sglang/test/kv_canary/fixtures.py index 5ade24525..44a18ae69 100644 --- a/python/sglang/test/kv_canary/fixtures.py +++ b/python/sglang/test/kv_canary/fixtures.py @@ -10,6 +10,7 @@ from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode from sglang.srt.kv_canary.pool_patcher.adapters.mha import attach_mha +from sglang.srt.kv_canary.pool_patcher.adapters.swa import attach_swa from sglang.srt.kv_canary.pool_patcher.api import register_pool_attacher from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode from sglang.srt.model_executor.forward_batch_info import ForwardMode @@ -35,6 +36,48 @@ class FakeMHAPool: return ptrs, lens, item_lens +@dataclass +class FakeSwaSubPool: + k_buffer: List[torch.Tensor] + v_buffer: List[torch.Tensor] + + +@dataclass +class FakeSWAPool: + full_kv_pool: object + swa_kv_pool: object + full_to_swa_index_mapping: torch.Tensor + page_size: int = 1 + + def get_contiguous_buf_infos(self): + return _kv_buf_infos( + k_buffer=self.full_kv_pool.k_buffer, + v_buffer=self.full_kv_pool.v_buffer, + page_size=self.page_size, + ) + + def get_state_buf_infos(self): + return _kv_buf_infos( + k_buffer=self.swa_kv_pool.k_buffer, + v_buffer=self.swa_kv_pool.v_buffer, + page_size=self.page_size, + ) + + +def _kv_buf_infos( + *, + k_buffer: List[torch.Tensor], + v_buffer: List[torch.Tensor], + page_size: int, +) -> tuple: + ptrs = [b.data_ptr() for b in k_buffer] + [b.data_ptr() for b in v_buffer] + lens = [b.nbytes for b in k_buffer] + [b.nbytes for b in v_buffer] + item_lens = [b[0].nbytes * page_size for b in k_buffer] + [ + b[0].nbytes * page_size for b in v_buffer + ] + return ptrs, lens, item_lens + + def make_mha_pool( device: torch.device = DEFAULT_DEVICE, *, @@ -53,6 +96,41 @@ def make_mha_pool( return FakeMHAPool(layer_num=layer_num, k_buffer=k_layers, v_buffer=v_layers) +def make_swa_pool( + device: torch.device = DEFAULT_DEVICE, + *, + full_slots: int = 16, + swa_slots: int = 8, + dim: int = 8, + layer_num: int = 1, +) -> FakeSWAPool: + full = FakeSwaSubPool( + k_buffer=[ + torch.zeros(full_slots, dim, dtype=torch.float16, device=device) + for _ in range(layer_num) + ], + v_buffer=[ + torch.zeros(full_slots, dim, dtype=torch.float16, device=device) + for _ in range(layer_num) + ], + ) + swa = FakeSwaSubPool( + k_buffer=[ + torch.zeros(swa_slots, dim, dtype=torch.float16, device=device) + for _ in range(layer_num) + ], + v_buffer=[ + torch.zeros(swa_slots, dim, dtype=torch.float16, device=device) + for _ in range(layer_num) + ], + ) + lut = torch.full((full_slots + 1,), -1, dtype=torch.int64, device=device) + lut[:swa_slots] = torch.arange(swa_slots, dtype=torch.int64, device=device) + return FakeSWAPool( + full_kv_pool=full, swa_kv_pool=swa, full_to_swa_index_mapping=lut + ) + + def make_base_config() -> CanaryConfig: return CanaryConfig( mode=CanaryMode.RAISE, @@ -189,3 +267,4 @@ def make_radix_cache( register_pool_attacher(FakeMHAPool, attach_mha) +register_pool_attacher(FakeSWAPool, attach_swa) diff --git a/python/sglang/test/kv_canary/mode_config.py b/python/sglang/test/kv_canary/mode_config.py index 1e60ad492..4743e9eaa 100644 --- a/python/sglang/test/kv_canary/mode_config.py +++ b/python/sglang/test/kv_canary/mode_config.py @@ -14,4 +14,10 @@ _MODE_CONFIGS: dict[str, _ModeConfig] = { "mha": _ModeConfig( model_path="Qwen/Qwen3-0.6B", ), + "swa": _ModeConfig( + model_path="google/gemma-4-E2B-it", + ), + "dsv4": _ModeConfig( + model_path="deepseek-ai/DeepSeek-V4-Flash", + ), } diff --git a/test/manual/kv_canary/test_self_e2e_baseline_dsv4.py b/test/manual/kv_canary/test_self_e2e_baseline_dsv4.py new file mode 100644 index 000000000..94b23a02a --- /dev/null +++ b/test/manual/kv_canary/test_self_e2e_baseline_dsv4.py @@ -0,0 +1,21 @@ +from __future__ import annotations + +import unittest +from test.registered.kv_canary.test_self_e2e_baseline import _BaselineBase + +from sglang.test.kv_canary.consts import ( + DSV4_POOL_SERVER_ARGS, + DSV4_POOL_SERVER_ENV, +) + + +class TestBaselineDsv4(_BaselineBase): + __test__ = True + + model_mode = "dsv4" + extra_server_args = DSV4_POOL_SERVER_ARGS + extra_env = DSV4_POOL_SERVER_ENV + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_e2e_baseline.py b/test/registered/kv_canary/test_self_e2e_baseline.py index 562025302..6974806a3 100644 --- a/test/registered/kv_canary/test_self_e2e_baseline.py +++ b/test/registered/kv_canary/test_self_e2e_baseline.py @@ -4,6 +4,7 @@ import unittest from sglang.srt.kv_canary.config import CanaryMode from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.consts import SWA_POOL_SERVER_ARGS from sglang.test.kv_canary.e2e_base import CanaryE2EBase register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") @@ -33,5 +34,10 @@ class TestBaselineMha(_BaselineBase): model_mode = "mha" +class TestBaselineSwa(_BaselineBase): + model_mode = "swa" + extra_server_args = SWA_POOL_SERVER_ARGS + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_pool_patcher.py b/test/registered/kv_canary/test_self_unit_pool_patcher.py index 5f6a35408..7760c845e 100644 --- a/test/registered/kv_canary/test_self_unit_pool_patcher.py +++ b/test/registered/kv_canary/test_self_unit_pool_patcher.py @@ -10,6 +10,7 @@ from sglang.test.kv_canary.fixtures import ( DEFAULT_DEVICE, make_base_config, make_mha_pool, + make_swa_pool, ) from sglang.test.test_utils import CustomTestCase @@ -41,6 +42,22 @@ class TestAttachCanaryBuffers(PoolPatcherHelper, CustomTestCase): self.assertIsNotNone(group.v_tail) self.assertEqual(group.v_head.shape, (16, CANARY_SLOT_BYTES)) + def test_canary_buffer_group_allocate_full_and_swa(self): + """Verify SWA pools allocate full and SWA canary buffers.""" + pool = make_swa_pool(self.device, full_slots=16, swa_slots=8) + groups_tuple = attach_canary_buffers( + pool=pool, + config=self.config, + device=self.device, + kv_token_id_vs_position_offset=0, + ) + groups = {g.kind: g for g in groups_tuple} + self.assertEqual(set(groups.keys()), {PoolKind.FULL, PoolKind.SWA}) + self.assertEqual(groups[PoolKind.FULL].k_head.shape[0], 16) + self.assertEqual(groups[PoolKind.SWA].k_head.shape[0], 8) + self.assertIsNotNone(groups[PoolKind.SWA].swa_index_lut) + self.assertIsNone(groups[PoolKind.FULL].swa_index_lut) + class TestPoolPatcherBufferInfos(PoolPatcherHelper, CustomTestCase): def test_get_contiguous_buf_infos_inserts_canary_entries(self): @@ -64,6 +81,24 @@ class TestPoolPatcherBufferInfos(PoolPatcherHelper, CustomTestCase): ptrs_after, _, _ = pool.get_contiguous_buf_infos() self.assertEqual(ptrs_after, ptrs_before) + def test_swa_attach_splices_full_into_contiguous_and_swa_into_state(self): + """Verify SWA patching splices canary buffers into both buffer lists.""" + pool = make_swa_pool(self.device, full_slots=16, swa_slots=8) + contiguous_before, _, _ = pool.get_contiguous_buf_infos() + state_before, _, _ = pool.get_state_buf_infos() + + attach_canary_buffers( + pool=pool, + config=self.config, + device=self.device, + kv_token_id_vs_position_offset=0, + ) + + contiguous_after, _, _ = pool.get_contiguous_buf_infos() + state_after, _, _ = pool.get_state_buf_infos() + self.assertEqual(len(contiguous_after), len(contiguous_before) + 4) + self.assertEqual(len(state_after), len(state_before) + 4) + def test_pd_layout_canary_inserted_correctly(self): """Verify PD (prefill-decode disaggregation) canary buffers are inserted in layout order.""" pool = make_mha_pool(self.device, num_slots=16, dim=8, layer_num=2)