chore: import upstream snapshot with attribution

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

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+import io
+import os
+import string
+import time
+from typing import ClassVar, Literal, Optional
+
+from sglang.srt.kv_canary.config import CanaryMode
+from sglang.srt.kv_canary.runner.swa_divergence import SwaDivergenceLog
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kv_canary.mode_config import _MODE_CONFIGS, _ModeConfig
+from sglang.test.kv_canary.utils import build_canary_server_args, post_parallel_generate
+from sglang.test.kv_canary.violation_assert_mixin import CanaryViolationAssertMixin
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+# Long prompt body shared by all canary e2e tests. The repetition count is chosen
+# 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
+
+
+class CapturedServerE2EBase(CanaryViolationAssertMixin, CustomTestCase):
+    process: ClassVar[Optional[object]] = None
+    base_url: ClassVar[str] = DEFAULT_URL_FOR_TEST
+    _stdout_buf: ClassVar[Optional[io.StringIO]] = None
+    _stderr_buf: ClassVar[Optional[io.StringIO]] = None
+
+    @classmethod
+    def tearDownClass(cls) -> None:
+        if cls.process is not None:
+            kill_process_tree(cls.process.pid)
+        for buf in (cls._stdout_buf, cls._stderr_buf):
+            if buf is not None:
+                buf.close()
+        cls._stdout_buf = None
+        cls._stderr_buf = None
+
+    def _captured_log_text(
+        self, side: Optional[Literal["prefill", "decode"]] = None
+    ) -> str:
+        stdout_text = (
+            self._stdout_buf.getvalue() if self._stdout_buf is not None else ""
+        )
+        stderr_text = (
+            self._stderr_buf.getvalue() if self._stderr_buf is not None else ""
+        )
+        return stdout_text + stderr_text
+
+    def assert_log_contains(self, substring: str) -> None:
+        log_text = self._captured_log_text()
+        if substring not in log_text:
+            raise AssertionError(
+                f"Expected substring {substring!r} not found in captured log. "
+                f"Log tail:\n{log_text[-2000:]}"
+            )
+
+
+class CanaryE2EBase(CapturedServerE2EBase):
+    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
+    # 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
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls._cfg = _MODE_CONFIGS[cls.model_mode]
+        server_env = os.environ.copy()
+        server_env.setdefault("SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT", "1")
+        server_env.update(cls.extra_env)
+        if cls.model_mode == "swa":
+            server_env.setdefault(
+                "SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL", "20"
+            )
+            # 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()
+
+        server_args = build_canary_server_args(
+            kv_canary_mode=cls.kv_canary_mode,
+            mode_cfg=cls._cfg,
+            extra_server_args=(
+                "--max-total-tokens",
+                "65536",
+                "--skip-server-warmup",
+                *cls.extra_server_args,
+            ),
+        )
+        cls.process = popen_launch_server(
+            cls._cfg.model_path,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=server_args,
+            env=server_env,
+            return_stdout_stderr=(cls._stdout_buf, cls._stderr_buf),
+        )
+
+    def make_prompts(self, n: int) -> list[str]:
+        if self.use_unique_prompts:
+            return _make_unique_prompts(n)
+        return [_LONG_PROMPT_BODY] * n
+
+    def send_parallel_requests(
+        self,
+        n: int = 8,
+        *,
+        assert_all_success: bool = True,
+        max_new_tokens: int = 2048,
+        timeout: float = 240.0,
+        ignore_eos: Optional[bool] = None,
+    ) -> list[dict]:
+        """Fan out n parallel /generate requests; return list of response dicts."""
+        if ignore_eos is None:
+            ignore_eos = self.model_mode == "swa"
+        results = post_parallel_generate(
+            url=self.base_url + "/generate",
+            prompts=self.make_prompts(n),
+            max_new_tokens=max_new_tokens,
+            timeout=timeout,
+            ignore_eos=ignore_eos,
+        )
+        if assert_all_success:
+            for result in results:
+                self.assertEqual(result.get("status_code"), 200, result)
+        return results
+
+    def maybe_assert_swa_divergence_observed(self) -> None:
+        if self.model_mode == "swa":
+            self.assert_swa_divergence_observed()
+
+    def assert_swa_divergence_observed(
+        self,
+        *,
+        min_swa_out_of_window_tokens: int = 1,
+        min_swa_full_idx_divergence: int = 1,
+        require_verify_lag: bool = True,
+        flush_wait_seconds: float = 3.0,
+        max_retries: int = 10,
+    ) -> None:
+        """Assert that the SWA path was genuinely exercised.
+
+        Three signals must all hold:
+          - ``swa_out_of_window_tokens >= 1``: at least one token has slid out of the
+            sliding window (its SWA mapping is 0). This only appears once a request decodes
+            past the window, so the window evicts — proves the SWA window slide actually ran.
+          - ``swa_full_idx_divergence >= 1``: SWA pool has actually remapped at least one
+            slot to a non-identity index (i.e. real slot reuse / eviction occurred). The
+            workload must drive SWA pool pressure for this to fire — required because the
+            "pool reuse" path is the one production hits under sustained long-context
+            traffic, and we must keep it covered.
+          - ``verify_swa < verify_full``: SWA verify kernel processed fewer tokens than
+            FULL — proves both kernel groups ran and the window short-circuited SWA.
+
+        The first two signals are checked as the *peak* across all sampled forwards, not
+        only the last sample. The divergence reporter snapshots one live forward batch per
+        interval; under PP it snapshots a single micro-batch, which may hold only in-window
+        requests even when another micro-batch diverged. "Was the SWA path ever exercised?"
+        is a max-over-samples question, so a trailing in-window sample must not mask an
+        earlier diverging one. ``verify_swa``/``verify_full`` are monotonic running totals,
+        so the lag check reads the last sample.
+        """
+        samples: list[tuple[SwaDivergenceLog, str]] = []
+        for _ in range(max_retries):
+            time.sleep(flush_wait_seconds)
+            samples = SwaDivergenceLog.find_all(self._captured_log_text())
+            if samples:
+                break
+
+        if not samples:
+            raise AssertionError(
+                "No kv_canary swa_divergence line found in server log after "
+                f"{max_retries} retries (wait={flush_wait_seconds}s each). "
+                f"Log tail:\n{self._captured_log_text()[-2000:]}"
+            )
+
+        peak_out_of_window = max(p.swa_out_of_window_tokens for p, _ in samples)
+        peak_full_idx_divergence = max(p.swa_full_idx_divergence for p, _ in samples)
+        last_parsed, last_line = samples[-1]
+
+        if peak_out_of_window < min_swa_out_of_window_tokens:
+            raise AssertionError(
+                f"SWA path not exercised: peak swa_out_of_window_tokens={peak_out_of_window} "
+                f"< min={min_swa_out_of_window_tokens} across {len(samples)} samples. "
+                f"Last line: {last_line}"
+            )
+        if peak_full_idx_divergence < min_swa_full_idx_divergence:
+            raise AssertionError(
+                f"SWA pool reuse not exercised: peak swa_full_idx_divergence={peak_full_idx_divergence} "
+                f"< min={min_swa_full_idx_divergence} across {len(samples)} samples. The workload "
+                f"did not drive enough SWA pool pressure to force slot remap. Last line: {last_line}"
+            )
+        if require_verify_lag and not (
+            last_parsed.verify_swa < last_parsed.verify_full
+        ):
+            raise AssertionError(
+                f"SWA path not exercised: verify_swa={last_parsed.verify_swa} "
+                f"not strictly less than verify_full={last_parsed.verify_full}. "
+                f"Line: {last_line}"
+            )
+
+
+def _make_unique_prompts(n: int) -> list[str]:
+    if n > len(_UNIQUE_PROMPT_FIRST_CHARS):
+        raise ValueError(
+            f"unique prompt count {n} exceeds supported count "
+            f"{len(_UNIQUE_PROMPT_FIRST_CHARS)}"
+        )
+
+    return [
+        (
+            f"{_UNIQUE_PROMPT_FIRST_CHARS[i]}"
+            f"{hex(i * 0x9E3779B1 & 0xFFFFFFFF)[2:]} "
+            f"{_LONG_PROMPT_BODY}"
+        )
+        for i in range(n)
+    ]

python/sglang/test/kv_canary/fixtures.py

@@ -0,0 +1,291 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import List, Optional
+
+import torch
+
+from sglang.jit_kernel.kv_canary import consts
+from sglang.jit_kernel.kv_canary.verify import CANARY_SLOT_BYTES, RealKvSource
+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
+
+DEFAULT_DEVICE: torch.device = torch.device("cuda")
+
+
+@dataclass
+class FakeMHAPool:
+    layer_num: int
+    k_buffer: List[torch.Tensor]
+    v_buffer: List[torch.Tensor]
+    page_size: int = 1
+
+    def get_contiguous_buf_infos(self):
+        ptrs = [b.data_ptr() for b in self.k_buffer] + [
+            b.data_ptr() for b in self.v_buffer
+        ]
+        lens = [b.nbytes for b in self.k_buffer] + [b.nbytes for b in self.v_buffer]
+        item_lens = [b[0].nbytes * self.page_size for b in self.k_buffer] + [
+            b[0].nbytes * self.page_size for b in self.v_buffer
+        ]
+        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,
+    *,
+    num_slots: int = 16,
+    dim: int = 8,
+    layer_num: int = 2,
+) -> FakeMHAPool:
+    k_layers = [
+        torch.zeros(num_slots, dim, dtype=torch.float16, device=device)
+        for _ in range(layer_num)
+    ]
+    v_layers = [
+        torch.zeros(num_slots, dim, dtype=torch.float16, device=device)
+        for _ in range(layer_num)
+    ]
+    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,
+        ring_capacity=1024,
+        sweep_interval=0,
+        real_kv_hash_mode=consts.RealKvHashMode.NONE,
+        enable_write_input_assert=False,
+        enable_verify_token_assert=True,
+        stats_print_every_n_steps=100,
+    )
+
+
+def make_req_to_token_pool(
+    device: torch.device = DEFAULT_DEVICE,
+    *,
+    max_reqs: int = 8,
+    max_seq_len: int = 32,
+) -> SimpleNamespace:
+    req_to_token = torch.zeros(max_reqs, max_seq_len, dtype=torch.int32, device=device)
+    return SimpleNamespace(
+        req_to_token=req_to_token, size=max_reqs, max_context_len=max_seq_len
+    )
+
+
+def make_forward_batch(
+    device: torch.device = DEFAULT_DEVICE,
+    *,
+    bs: int = 2,
+    seq_lens_list: tuple[int, ...] = (3, 4),
+    req_pool_indices: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+    seq_lens_sum: Optional[int] = None,
+    extend_prefix_lens: Optional[torch.Tensor] = None,
+    extend_prefix_lens_cpu: Optional[list] = None,
+    extend_seq_lens: Optional[torch.Tensor] = None,
+    extend_seq_lens_cpu: Optional[list] = None,
+    is_extend: bool = False,
+    is_target_verify: bool = False,
+    is_draft_extend_v2: bool = False,
+    spec_info: Optional[object] = None,
+    input_ids: Optional[torch.Tensor] = None,
+    positions: Optional[torch.Tensor] = None,
+    out_cache_loc: Optional[torch.Tensor] = None,
+    num_token_non_padded_cpu: Optional[int] = None,
+) -> SimpleNamespace:
+    seq_lens_default = list(seq_lens_list[:bs])
+    if req_pool_indices is None:
+        req_pool_indices = torch.tensor([1, 2][:bs], dtype=torch.int64, device=device)
+    if seq_lens is None:
+        seq_lens = torch.tensor(seq_lens_default, dtype=torch.int32, device=device)
+    if seq_lens_sum is None:
+        seq_lens_sum = int(sum(seq_lens_default))
+    if input_ids is None:
+        input_ids = torch.zeros(bs, dtype=torch.int32, device=device)
+    if positions is None:
+        # Default to decode-canonical: positions = seq_lens - 1 (one-token-per-req decode write
+        # at the post-bump tail). Plan input derives decode prefix_lens from positions
+        # directly, so the default must keep parity.
+        positions = (seq_lens.to(torch.int64) - 1).clamp(min=0).to(torch.int32)
+    if out_cache_loc is None:
+        out_cache_loc = torch.zeros(bs, dtype=torch.int32, device=device)
+
+    if is_extend:
+        forward_mode = ForwardMode.EXTEND
+    elif is_target_verify:
+        forward_mode = ForwardMode.TARGET_VERIFY
+    elif is_draft_extend_v2:
+        forward_mode = ForwardMode.DRAFT_EXTEND_V2
+    else:
+        forward_mode = ForwardMode.DECODE
+    return SimpleNamespace(
+        forward_mode=forward_mode,
+        batch_size=bs,
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens,
+        seq_lens_sum=seq_lens_sum,
+        extend_prefix_lens=extend_prefix_lens,
+        extend_prefix_lens_cpu=extend_prefix_lens_cpu,
+        extend_seq_lens=extend_seq_lens,
+        extend_seq_lens_cpu=extend_seq_lens_cpu,
+        spec_info=spec_info,
+        input_ids=input_ids,
+        positions=positions,
+        out_cache_loc=out_cache_loc,
+        num_token_non_padded_cpu=num_token_non_padded_cpu,
+        req_all_ids_flat=None,
+        req_all_ids_lens=None,
+    )
+
+
+def make_buffer_group(
+    *,
+    device: torch.device = DEFAULT_DEVICE,
+    kind: PoolKind = PoolKind.FULL,
+    has_v: bool = True,
+    has_real_kv: bool = False,
+    real_kv_source: Optional[RealKvSource] = None,
+    swa_index_lut: Optional[torch.Tensor] = None,
+    num_slots: int = 4,
+    kv_token_id_vs_position_offset: int = 0,
+) -> CanaryBufferGroup:
+    def _zero() -> torch.Tensor:
+        return torch.zeros(
+            num_slots, CANARY_SLOT_BYTES, dtype=torch.uint8, device=device
+        )
+
+    if has_real_kv:
+        source = real_kv_source or RealKvSource(
+            tensor=torch.zeros(num_slots, 16, dtype=torch.uint8, device=device),
+            page_size=1,
+            num_bytes_per_token=16,
+            read_bytes=16,
+        )
+        real_kv_sources = (source,)
+    else:
+        real_kv_sources = ()
+
+    return CanaryBufferGroup(
+        kind=kind,
+        k_head=_zero(),
+        k_tail=_zero(),
+        v_head=_zero() if has_v else None,
+        v_tail=_zero() if has_v else None,
+        real_kv_sources_k=real_kv_sources,
+        real_kv_sources_v=real_kv_sources if has_v else (),
+        swa_index_lut=swa_index_lut,
+        kv_token_id_vs_position_offset=kv_token_id_vs_position_offset,
+    )
+
+
+def make_radix_cache(
+    slot_lists: List[List[int]], device: torch.device = DEFAULT_DEVICE
+):
+    cache = RadixCache.__new__(RadixCache)
+    cache.device = device
+    cache.page_size = 1
+    cache.disable = False
+
+    root = TreeNode()
+    root.value = torch.tensor(
+        slot_lists[0] if slot_lists else [], dtype=torch.int32, device=device
+    )
+    cache.root_node = root
+
+    current = root
+    for child_slots in slot_lists[1:]:
+        child = TreeNode()
+        child.value = torch.tensor(child_slots, dtype=torch.int32, device=device)
+        child.parent = current
+        current.children[child.id] = child
+        current = child
+
+    return cache
+
+
+register_pool_attacher(FakeMHAPool, attach_mha)
+register_pool_attacher(FakeSWAPool, attach_swa)

python/sglang/test/kv_canary/mode_config.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class _ModeConfig:
+    model_path: str
+    json_model_override_args: Optional[str] = None
+
+
+_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",
+    ),
+}

python/sglang/test/kv_canary/pd_fixture.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import uuid
+from typing import ClassVar, Literal, Optional
+
+from sglang.srt.kv_canary.config import CanaryMode
+from sglang.test.kv_canary.mode_config import _MODE_CONFIGS, _ModeConfig
+from sglang.test.kv_canary.utils import build_canary_server_args, post_parallel_generate
+from sglang.test.kv_canary.violation_assert_mixin import CanaryViolationAssertMixin
+from sglang.test.server_fixtures.disaggregation_fixture import (
+    PDDisaggregationServerBase,
+)
+
+_SHORT_PROMPT_BODY = ("The quick brown fox jumps over the lazy dog. " * 8).strip()
+
+
+class CanaryPDFixture(CanaryViolationAssertMixin, PDDisaggregationServerBase):
+    capture_per_side_logs = True
+
+    model_mode: ClassVar[Literal["mha", "swa", "dsv4"]]
+    kv_canary_mode: ClassVar[CanaryMode] = CanaryMode.LOG
+    extra_server_args: ClassVar[tuple[str, ...]] = (
+        "--kv-canary-real-data",
+        "partial",
+        "--skip-server-warmup",
+    )
+
+    _cfg: ClassVar[Optional[_ModeConfig]] = None
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        super().setUpClass()
+        cls._cfg = _MODE_CONFIGS[cls.model_mode]
+        cls.model = cls._cfg.model_path
+
+        canary_args = build_canary_server_args(
+            kv_canary_mode=cls.kv_canary_mode,
+            mode_cfg=cls._cfg,
+            extra_server_args=cls.extra_server_args,
+        )
+        cls.extra_prefill_args = list(canary_args)
+        cls.extra_decode_args = list(canary_args)
+        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.
+            cls.extra_prefill_env = {
+                **cls.extra_prefill_env,
+                "SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION": "1",
+            }
+            cls.extra_decode_env = {
+                **cls.extra_decode_env,
+                "SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION": "1",
+            }
+        cls.launch_all()
+
+    def send_parallel_short_requests(
+        self,
+        n: int,
+        *,
+        assert_all_success: bool = True,
+        max_new_tokens: int = 100,
+        timeout: float = 60.0,
+        distinct_prompts: bool = False,
+    ) -> list[dict]:
+        if distinct_prompts:
+            # Diverge at the very first token (request index before the per-call
+            # nonce) so requests share no radix-dedupable prefix beyond a
+            # tokenizer-added BOS, and retries never hit earlier attempts' cache.
+            nonce = uuid.uuid4().hex[:8]
+            prompts = [f"{i} {nonce} {_SHORT_PROMPT_BODY}" for i in range(n)]
+        else:
+            prompts = [_SHORT_PROMPT_BODY] * n
+        results = post_parallel_generate(
+            url=self.lb_url + "/generate",
+            prompts=prompts,
+            max_new_tokens=max_new_tokens,
+            timeout=timeout,
+        )
+        if assert_all_success:
+            for result in results:
+                self.assertEqual(result.get("status_code"), 200, result)
+        return results
+
+    def _captured_log_text(
+        self, side: Optional[Literal["prefill", "decode"]] = None
+    ) -> str:
+        if side == "prefill":
+            stdout_buf = type(self)._prefill_stdout_buf
+            stderr_buf = type(self)._prefill_stderr_buf
+        elif side == "decode":
+            stdout_buf = type(self)._decode_stdout_buf
+            stderr_buf = type(self)._decode_stderr_buf
+        else:
+            raise ValueError(f"Unsupported side={side!r}")
+        stdout_text = stdout_buf.getvalue() if stdout_buf is not None else ""
+        stderr_text = stderr_buf.getvalue() if stderr_buf is not None else ""
+        return stdout_text + stderr_text

python/sglang/test/kv_canary/pp_fixture.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from typing import ClassVar
+
+from sglang.test.kv_canary.consts import SWA_POOL_SERVER_ARGS
+from sglang.test.kv_canary.e2e_base import CanaryE2EBase
+
+PP_SIZE: int = 2
+
+
+class CanaryPPFixture(CanaryE2EBase):
+
+    model_mode: ClassVar[str] = "swa"
+    workload_n_batches: ClassVar[int] = 2
+
+    @classmethod
+    def setUpClass(cls) -> None:
+        cls.extra_server_args = (
+            "--pp-size",
+            str(PP_SIZE),
+            "--disable-cuda-graph",
+            *SWA_POOL_SERVER_ARGS,
+            *cls.extra_server_args,
+        )
+        super().setUpClass()

python/sglang/test/kv_canary/runner_test_base.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import torch
+
+from sglang.jit_kernel.kv_canary.consts import RealKvHashMode
+from sglang.jit_kernel.kv_canary.verify import CanaryLaunchTag
+from sglang.srt.kv_canary import endpoint as endpoint_module
+from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup
+from sglang.srt.kv_canary.capacities import CanaryLaunchCapacities
+from sglang.srt.kv_canary.config import CanaryConfig, CanaryMode
+from sglang.srt.kv_canary.perturb.config import PerturbConfig
+from sglang.srt.kv_canary.runner import kernel_launcher as kernel_launcher_module
+from sglang.srt.kv_canary.runner.canary_manager import CanaryManager
+from sglang.test.kv_canary.fixtures import (
+    DEFAULT_DEVICE,
+    make_buffer_group,
+    make_req_to_token_pool,
+)
+from sglang.test.test_utils import CustomTestCase
+
+
+def make_config(
+    *,
+    mode: CanaryMode = CanaryMode.RAISE,
+    ring_capacity: int = 1024,
+    sweep_interval: int = 0,
+    real_kv_hash_mode: RealKvHashMode = RealKvHashMode.NONE,
+    enable_write_input_assert: bool = False,
+    enable_verify_token_assert: bool = True,
+    stats_print_every_n_steps: int = 100,
+) -> CanaryConfig:
+    return CanaryConfig(
+        mode=mode,
+        ring_capacity=ring_capacity,
+        sweep_interval=sweep_interval,
+        real_kv_hash_mode=real_kv_hash_mode,
+        enable_write_input_assert=enable_write_input_assert,
+        enable_verify_token_assert=enable_verify_token_assert,
+        stats_print_every_n_steps=stats_print_every_n_steps,
+    )
+
+
+class RecordingEndpoint:
+    def __init__(self, *, kernel_kind: CanaryLaunchTag) -> None:
+        self.kernel_kind = kernel_kind
+        self.calls: list[dict[str, object]] = []
+
+    def launch_per_forward(self, **kwargs: object) -> None:
+        self.calls.append(kwargs)
+
+
+def make_perturb_config() -> PerturbConfig:
+    """Build a PerturbConfig with every probability pinned to 0 so the
+    perturb hooks do nothing during unit tests."""
+    return PerturbConfig(
+        req_to_token_prob=0.0,
+        real_kv_used_prob=0.0,
+        real_kv_unused_cache_prob=0.0,
+        real_kv_post_forward_prob=0.0,
+        target_group_kind=None,
+        warmup_steps=0,
+    )
+
+
+def make_manager(
+    *,
+    device: torch.device,
+    config: CanaryConfig | None = None,
+    perturb_config: PerturbConfig | None = None,
+    group: CanaryBufferGroup | None = None,
+    req_pool: SimpleNamespace | None = None,
+    per_forward_verify_capacity: int = 16,
+    speculative_num_steps: int = 1,
+) -> CanaryManager:
+    if config is None:
+        config = make_config()
+    if perturb_config is None:
+        perturb_config = make_perturb_config()
+    if group is None:
+        group = make_buffer_group(device=device)
+    if req_pool is None:
+        req_pool = make_req_to_token_pool(device=device, max_reqs=4, max_seq_len=8)
+    return CanaryManager(
+        config=config,
+        perturb_config=perturb_config,
+        buffer_groups=(group,),
+        device=device,
+        req_to_token_pool=req_pool,
+        launch_capacities=CanaryLaunchCapacities(
+            per_forward_verify_capacity=per_forward_verify_capacity,
+            per_forward_write_req_capacity=2,
+            per_forward_write_entry_capacity=8,
+        ),
+        speculative_num_steps=speculative_num_steps,
+    )
+
+
+class CanaryManagerTestCase(CustomTestCase):
+    def setUp(self) -> None:
+        self.device = DEFAULT_DEVICE
+        # Stub plan/verify/write kernels so CPU runs don't need CUDA JIT.
+        self._patchers = [
+            patch.object(
+                kernel_launcher_module,
+                "launch_canary_plan_kernels",
+                lambda **kwargs: None,
+            ),
+            patch.object(
+                endpoint_module, "launch_canary_verify_kernel", lambda **kwargs: None
+            ),
+            patch.object(
+                endpoint_module, "launch_canary_write_kernel", lambda **kwargs: None
+            ),
+        ]
+        for patcher in self._patchers:
+            patcher.start()
+            self.addCleanup(patcher.stop)

python/sglang/test/kv_canary/utils.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from concurrent.futures import ThreadPoolExecutor
+
+import requests
+
+from sglang.srt.kv_canary.config import CanaryMode
+from sglang.test.kv_canary.mode_config import _ModeConfig
+
+
+def build_canary_server_args(
+    *,
+    kv_canary_mode: CanaryMode,
+    mode_cfg: _ModeConfig,
+    extra_server_args: tuple[str, ...] = (),
+) -> list[str]:
+    args = [
+        "--kv-canary",
+        kv_canary_mode.value,
+        "--disable-piecewise-cuda-graph",
+        "--context-length",
+        "16384",
+        *extra_server_args,
+    ]
+    if mode_cfg.json_model_override_args is not None:
+        args.extend(["--json-model-override-args", mode_cfg.json_model_override_args])
+    return args
+
+
+def post_parallel_generate(
+    *,
+    url: str,
+    prompts: list[str],
+    max_new_tokens: int,
+    timeout: float,
+    ignore_eos: bool = False,
+) -> list[dict]:
+    def _send(prompt: str) -> dict:
+        try:
+            resp = requests.post(
+                url,
+                json={
+                    "text": prompt,
+                    "sampling_params": {
+                        "max_new_tokens": max_new_tokens,
+                        "temperature": 0.0,
+                        "ignore_eos": ignore_eos,
+                    },
+                },
+                timeout=timeout,
+            )
+            return {"status_code": resp.status_code, "body": resp.text}
+        except requests.RequestException as exc:
+            return {"status_code": -1, "error": repr(exc)}
+
+    with ThreadPoolExecutor(max_workers=max(1, len(prompts))) as pool:
+        return list(pool.map(_send, prompts))

python/sglang/test/kv_canary/violation_assert_mixin.py

@@ -0,0 +1,145 @@
+from __future__ import annotations
+
+import time
+from typing import Literal, Optional
+
+from sglang.srt.kv_canary.perturb.config import TargetGroupKind
+from sglang.test.kv_canary.violation_log_utils import (
+    assert_no_violation_in_log,
+    find_violation_in_log,
+)
+
+_Side = Optional[Literal["prefill", "decode"]]
+
+
+class CanaryViolationAssertMixin:
+    def _captured_log_text(self, side: _Side = None) -> str:
+        raise NotImplementedError
+
+    def assert_per_forward_violation_reported(
+        self,
+        *,
+        fail_reason: str,
+        target_group: Optional[TargetGroupKind] = None,
+        side: _Side = None,
+        flush_wait_seconds: float = 2.0,
+    ) -> None:
+        suffix = "" if target_group is None else f"_{target_group.name}"
+        self.assert_violation_logged_any(
+            launch_tag_patterns=(f"HEAD_*{suffix}", f"TAIL_*{suffix}"),
+            fail_reason=fail_reason,
+            side=side,
+            flush_wait_seconds=flush_wait_seconds,
+        )
+
+    def assert_sweep_violation_reported(
+        self,
+        *,
+        fail_reason: str,
+        target_group: TargetGroupKind,
+        side: _Side = None,
+        flush_wait_seconds: float = 2.0,
+        max_retries: int = 4,
+    ) -> None:
+        self.assert_violation_logged_any(
+            launch_tag_patterns=(f"SWEEP_*_{target_group.name}",),
+            fail_reason=fail_reason,
+            side=side,
+            flush_wait_seconds=flush_wait_seconds,
+            max_retries=max_retries,
+        )
+
+    def assert_any_launch_tag_violation_reported(
+        self,
+        *,
+        fail_reason: str,
+        side: _Side = None,
+        flush_wait_seconds: float = 3.0,
+        max_retries: int = 10,
+    ) -> None:
+        self.assert_violation_logged_any(
+            launch_tag_patterns=("*",),
+            fail_reason=fail_reason,
+            side=side,
+            flush_wait_seconds=flush_wait_seconds,
+            max_retries=max_retries,
+        )
+
+    def assert_any_launch_tag_violation_absent(
+        self, *, fail_reason: str, side: _Side = None
+    ) -> None:
+        self.assert_no_violation_matching(
+            launch_tag_patterns=("*",), fail_reason=fail_reason, side=side
+        )
+
+    def assert_violation_logged_any(
+        self,
+        *,
+        launch_tag_patterns: tuple[str, ...],
+        fail_reason: str,
+        side: _Side = None,
+        flush_wait_seconds: float = 2.0,
+        max_retries: int = 1,
+    ) -> None:
+        log_text = ""
+        for _ in range(max_retries):
+            time.sleep(flush_wait_seconds)
+            log_text = self._captured_log_text(side)
+            if find_violation_in_log(
+                log_text,
+                launch_tag_patterns=launch_tag_patterns,
+                fail_reason=fail_reason,
+            ):
+                return
+        side_label = "" if side is None else f" on side={side}"
+        other_side_diag = ""
+        if side in ("prefill", "decode"):
+            other = "decode" if side == "prefill" else "prefill"
+            try:
+                other_text = self._captured_log_text(other)
+                other_match = find_violation_in_log(
+                    other_text,
+                    launch_tag_patterns=launch_tag_patterns,
+                    fail_reason=fail_reason,
+                )
+                other_side_diag = (
+                    f"\n[diag] other side ({other}) buf len={len(other_text)} "
+                    f"contains_match={other_match}"
+                )
+            except (NotImplementedError, ValueError):
+                pass
+        raise AssertionError(
+            f"No canary violation matching launch_tag_patterns={launch_tag_patterns!r} "
+            f"fail_reason={fail_reason!r}{side_label} after max_retries={max_retries} "
+            f"(wait={flush_wait_seconds}s each). "
+            f"log_text len={len(log_text)}.{other_side_diag} Log tail:\n"
+            f"{log_text[-2000:]}"
+        )
+
+    def assert_no_violation_matching(
+        self,
+        *,
+        launch_tag_patterns: tuple[str, ...],
+        fail_reason: str,
+        side: _Side = None,
+    ) -> None:
+        log_text = self._captured_log_text(side)
+        if find_violation_in_log(
+            log_text,
+            launch_tag_patterns=launch_tag_patterns,
+            fail_reason=fail_reason,
+        ):
+            raise AssertionError(
+                f"Unexpected canary violation matching "
+                f"launch_tag_patterns={launch_tag_patterns!r} "
+                f"fail_reason={fail_reason!r}. Log tail:\n{log_text[-2000:]}"
+            )
+
+    def assert_no_violation(
+        self,
+        *,
+        side: _Side = None,
+        wait_seconds: float = 2.0,
+    ) -> None:
+        time.sleep(wait_seconds)
+        assert_no_violation_in_log(self._captured_log_text(side))

python/sglang/test/kv_canary/violation_log_utils.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+import fnmatch
+import re
+
+_VIOLATION_LINE_RE = re.compile(
+    r"kv_canary violation: launch_tag=(\S+) fail_reason=(\S+)"
+)
+
+
+def find_violation_in_log(
+    log_text: str,
+    *,
+    launch_tag_patterns: tuple[str, ...],
+    fail_reason: str,
+) -> bool:
+    for match in _VIOLATION_LINE_RE.finditer(log_text):
+        tag = match.group(1)
+        reason_field = match.group(2)
+        if fail_reason not in reason_field.split("+"):
+            continue
+        if any(fnmatch.fnmatchcase(tag, pattern) for pattern in launch_tag_patterns):
+            return True
+    return False
+
+
+def assert_no_violation_in_log(log_text: str) -> None:
+    if "kv_canary violation:" in log_text:
+        raise AssertionError(
+            f"Unexpected canary violation found. Log tail:\n{log_text[-2000:]}"
+        )