Add KV-canary SWA + DeepSeek-V4 pool support (#26810)

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

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,
+    )

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,
 }
 
 

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",
+}

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

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)

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",
+    ),
 }

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

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

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)