diff --git a/python/sglang/jit_kernel/kv_canary/plan/api.py b/python/sglang/jit_kernel/kv_canary/plan/api.py index ac38ed8bc..daf5a7f7f 100644 --- a/python/sglang/jit_kernel/kv_canary/plan/api.py +++ b/python/sglang/jit_kernel/kv_canary/plan/api.py @@ -38,7 +38,7 @@ def launch_canary_plan_kernels( prefix_lens[r] - swa_window_size) if SWA else 0. slot_idx = req_to_token[req_pool_indices[r], pos] (SWA-translated via full_to_swa_index_mapping if non-None); prev_slot_idx = req_to_token[req_pool_indices[r], pos-1] for pos > 0, else -1. (SWA windows do NOT reset the chain — - the writer chains across the entire prefix; verify within an SWA window dereferences the real + the writer chains across the entire prefix; sweep verify within an SWA window dereferences the real predecessor for chain-link reconstruction.) Expected-token gather: when ``req_to_verify_expected_tokens`` is supplied, ``expected_input_id = req_to_verify_expected_tokens[rp, pos + kv_token_id_vs_position_offset]`` when ``0 <= pos + diff --git a/python/sglang/jit_kernel/kv_canary/verify.py b/python/sglang/jit_kernel/kv_canary/verify.py index 978665440..e0121bd40 100644 --- a/python/sglang/jit_kernel/kv_canary/verify.py +++ b/python/sglang/jit_kernel/kv_canary/verify.py @@ -17,16 +17,20 @@ CANARY_SLOT_BYTES: Final[int] = consts.CANARY_FIELDS_PER_SLOT * 8 class CanaryLaunchTag(IntEnum): - """Unique tag per (head | tail) × (K | V) × (FULL | SWA) launch.""" + """Unique tag per (head | tail | sweep) × (K | V) × (FULL | SWA) launch.""" HEAD_K_FULL = 0 HEAD_V_FULL = 1 TAIL_K_FULL = 2 TAIL_V_FULL = 3 - HEAD_K_SWA = 4 - HEAD_V_SWA = 5 - TAIL_K_SWA = 6 - TAIL_V_SWA = 7 + SWEEP_K_FULL = 4 + SWEEP_V_FULL = 5 + HEAD_K_SWA = 6 + HEAD_V_SWA = 7 + TAIL_K_SWA = 8 + TAIL_V_SWA = 9 + SWEEP_K_SWA = 10 + SWEEP_V_SWA = 11 def _assert_contiguous(tensor: torch.Tensor, name: str) -> None: @@ -68,7 +72,8 @@ class VerifyPlan: """Flat verify entries consumed by launch_canary_verify_kernel. Each row is a self-contained (slot_idx, position, prev_slot_idx) triple, so the verify kernel makes no - assumption about the entry's source. prev_slot_idx == -1 flags a chain-seed entry (kernel + assumption about the entry's source — per-forward derivation, sweep over running reqs, and sweep over + radix-cache orphan slots all populate the same schema. prev_slot_idx == -1 flags a chain-seed entry (kernel anchors on the hardcoded CANARY_CHAIN_ANCHOR constant instead of reading a predecessor). Sized to a cuda-graph-captured capacity; active prefix is verify_num_valid[0]. Padding tail entries are diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py b/python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py index 5590997c7..56478806a 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py +++ b/python/sglang/jit_kernel/tests/kv_canary/test_pipeline_e2e.py @@ -635,7 +635,7 @@ def test_pipeline_ring_overflow_via_real_plan() -> None: @pytest.mark.parametrize( - "kernel_kind", [CanaryLaunchTag.HEAD_K_FULL, CanaryLaunchTag.TAIL_V_SWA] + "kernel_kind", [CanaryLaunchTag.HEAD_K_FULL, CanaryLaunchTag.SWEEP_V_SWA] ) def test_pipeline_kernel_kind_propagates(kernel_kind: CanaryLaunchTag) -> None: """Different CanaryLaunchTag values: violation ring's kernel_kind field matches on both sides.""" diff --git a/python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py b/python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py index 9db088701..5fded210c 100644 --- a/python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py +++ b/python/sglang/jit_kernel/tests/kv_canary/test_verify_hand.py @@ -966,7 +966,7 @@ class TestViolationRing: buf_pair = _buf_pair() _stamp_head(buf_pair, slot_idx=1, token=1) - for tag in (CanaryLaunchTag.HEAD_K_FULL, CanaryLaunchTag.TAIL_V_SWA): + for tag in (CanaryLaunchTag.HEAD_K_FULL, CanaryLaunchTag.SWEEP_V_SWA): plan_pair = _plan_pair_single(slot_idx=1, position=99) cuda_log, _ = run_verify_diff( buf_pair=buf_pair, plan_pair=plan_pair, kernel_kind=tag diff --git a/python/sglang/srt/kv_canary/config.py b/python/sglang/srt/kv_canary/config.py index 8a21139ad..729bffc17 100644 --- a/python/sglang/srt/kv_canary/config.py +++ b/python/sglang/srt/kv_canary/config.py @@ -30,10 +30,14 @@ class CanaryConfig: violations propagate to host as RuntimeError after the next D2H pump. ring_capacity: Violation ring capacity (rows in ViolationLog.violation_ring). Sized generously; overflow only drops detail beyond row N, the monotonic counter still grows. + sweep_interval: 0 disables sweep entirely; positive N means every N-th forward step the runner + additionally walks all radix-tree-held slots (overlap with per-forward HEAD/TAIL is harmless + redundancy) and verifies them. """ mode: CanaryMode ring_capacity: int + sweep_interval: int @classmethod def from_env(cls, server_args: "ServerArgs") -> "CanaryConfig": @@ -46,4 +50,5 @@ class CanaryConfig: return cls( mode=CanaryMode(mode_raw), ring_capacity=envs.SGLANG_KV_CANARY_RING_CAPACITY.get(), + sweep_interval=server_args.kv_canary_sweep_interval, ) diff --git a/python/sglang/srt/kv_canary/endpoint.py b/python/sglang/srt/kv_canary/endpoint.py index 1917879be..cd0ce0403 100644 --- a/python/sglang/srt/kv_canary/endpoint.py +++ b/python/sglang/srt/kv_canary/endpoint.py @@ -45,6 +45,11 @@ class CanaryEndpoint: expected_inputs: ExpectedInputs, violation_log: ViolationLog, ) -> None: + if _is_sweep_tag(self.kernel_kind): + raise NotImplementedError( + f"kv-canary: launch_per_forward not supported on sweep endpoint {self.kernel_kind.name}" + ) + context = self._make_verify_or_write_context( violation_log=violation_log, ) @@ -77,6 +82,25 @@ class CanaryEndpoint: expected_input_positions=expected_input_positions, ) + def launch_sweep( + self, + *, + verify_plan: VerifyPlan, + violation_log: ViolationLog, + ) -> None: + if not _is_sweep_tag(self.kernel_kind): + raise NotImplementedError( + f"kv-canary: launch_sweep not supported on non-sweep endpoint {self.kernel_kind.name}" + ) + + launch_canary_verify_kernel( + context=self._make_verify_or_write_context( + violation_log=violation_log, + ), + plan=verify_plan, + check_verify_expected_token=False, + ) + def _make_verify_or_write_context( self, *, @@ -93,6 +117,15 @@ class CanaryEndpoint: ) +def _is_sweep_tag(tag: CanaryLaunchTag) -> bool: + return tag in ( + CanaryLaunchTag.SWEEP_K_FULL, + CanaryLaunchTag.SWEEP_V_FULL, + CanaryLaunchTag.SWEEP_K_SWA, + CanaryLaunchTag.SWEEP_V_SWA, + ) + + def _resolve_canary_buf( *, slot: str, @@ -113,6 +146,8 @@ _FULL_LAYOUT: tuple[tuple[CanaryLaunchTag, str, str], ...] = ( (CanaryLaunchTag.HEAD_V_FULL, "HEAD", "V"), (CanaryLaunchTag.TAIL_K_FULL, "TAIL", "K"), (CanaryLaunchTag.TAIL_V_FULL, "TAIL", "V"), + (CanaryLaunchTag.SWEEP_K_FULL, "SWEEP", "K"), + (CanaryLaunchTag.SWEEP_V_FULL, "SWEEP", "V"), ) @@ -121,6 +156,8 @@ _SWA_LAYOUT: tuple[tuple[CanaryLaunchTag, str, str], ...] = ( (CanaryLaunchTag.HEAD_V_SWA, "HEAD", "V"), (CanaryLaunchTag.TAIL_K_SWA, "TAIL", "K"), (CanaryLaunchTag.TAIL_V_SWA, "TAIL", "V"), + (CanaryLaunchTag.SWEEP_K_SWA, "SWEEP", "K"), + (CanaryLaunchTag.SWEEP_V_SWA, "SWEEP", "V"), ) @@ -138,7 +175,8 @@ def build_endpoints_from_group( if half == "V" and not group.has_v_half: continue - canary_buf = _resolve_canary_buf(slot=slot, half=half, group=group) + buf_slot = "TAIL" if slot == "SWEEP" else slot + canary_buf = _resolve_canary_buf(slot=buf_slot, half=half, group=group) lut = group.swa_index_lut if pool_kind is PoolKind.SWA else None slot_view = device_state.slot_run_counters[tag.value : tag.value + 1] kernel_view = device_state.kernel_run_counters[tag.value : tag.value + 1] diff --git a/python/sglang/srt/kv_canary/radix_cache_walker.py b/python/sglang/srt/kv_canary/radix_cache_walker.py new file mode 100644 index 000000000..36f444d94 --- /dev/null +++ b/python/sglang/srt/kv_canary/radix_cache_walker.py @@ -0,0 +1,148 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.mem_cache.radix_cache import RadixCache +from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache + +if TYPE_CHECKING: + from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache + from sglang.srt.mem_cache.radix_cache import TreeNode + + +@dataclass(frozen=True, slots=True, kw_only=True) +class RadixCacheWalkResult: + slot_indices: torch.Tensor + positions: torch.Tensor + prev_slot_indices: torch.Tensor + + +def walk_radix_cache_for_canary( + *, + radix_cache: "BasePrefixCache", + unlocked_only: bool = False, + swa_resident_only: bool = False, +) -> RadixCacheWalkResult: + """Walk the radix tree and emit flat (slot_indices, positions, prev_slot_indices) tensors. + + With both flags False (default), emits every slot held by the radix cache (including slots + also referenced by a currently-running req — that overlap is harmless redundancy with the + per-forward HEAD/TAIL path). ``unlocked_only=True`` skips nodes still locked by a running + req. ``swa_resident_only=True`` skips SWA-tombstoned nodes (slots evicted from the SWA + window).""" + cache_type = type(radix_cache) + if cache_type is not RadixCache and cache_type is not SWARadixCache: + raise NotImplementedError( + f"walk_radix_cache_for_canary does not support {cache_type.__name__}" + ) + + slot_buf: list[int] = [] + position_buf: list[int] = [] + prev_slot_buf: list[int] = [] + + _walk_radix_subtree( + node=radix_cache.root_node, + radix_cache=radix_cache, + depth=0, + parent_last_slot=-1, + slot_buf=slot_buf, + position_buf=position_buf, + prev_slot_buf=prev_slot_buf, + is_root=True, + unlocked_only=unlocked_only, + swa_resident_only=swa_resident_only, + ) + + slot_tensor = torch.tensor(slot_buf, dtype=torch.int64) + position_tensor = torch.tensor(position_buf, dtype=torch.int64) + prev_slot_tensor = torch.tensor(prev_slot_buf, dtype=torch.int64) + return RadixCacheWalkResult( + slot_indices=slot_tensor, + positions=position_tensor, + prev_slot_indices=prev_slot_tensor, + ) + + +def _walk_radix_subtree( + *, + node: "TreeNode", + radix_cache: "BasePrefixCache", + depth: int, + parent_last_slot: int, + slot_buf: list[int], + position_buf: list[int], + prev_slot_buf: list[int], + is_root: bool, + unlocked_only: bool, + swa_resident_only: bool, +) -> None: + if isinstance(node.value, torch.Tensor): + node_slots = [int(s) for s in node.value.tolist()] + else: + node_slots = [] + + if unlocked_only: + emit_slots = not is_root and _node_is_unlocked_for_canary( + node=node, radix_cache=radix_cache + ) + else: + emit_slots = not is_root + if swa_resident_only: + emit_slots = emit_slots and _node_is_swa_resident_for_canary( + node=node, + radix_cache=radix_cache, + ) + + chain_last_slot = parent_last_slot + for j, slot in enumerate(node_slots): + prev = parent_last_slot if j == 0 else node_slots[j - 1] + if emit_slots: + slot_buf.append(slot) + position_buf.append(depth + j) + prev_slot_buf.append(prev) + chain_last_slot = slot + + child_depth = depth + len(node_slots) + for child in node.children.values(): + _walk_radix_subtree( + node=child, + radix_cache=radix_cache, + depth=child_depth, + parent_last_slot=chain_last_slot, + slot_buf=slot_buf, + position_buf=position_buf, + prev_slot_buf=prev_slot_buf, + is_root=False, + unlocked_only=unlocked_only, + swa_resident_only=swa_resident_only, + ) + + +def _node_is_unlocked_for_canary( + *, + node: "TreeNode", + radix_cache: "BasePrefixCache", +) -> bool: + if type(radix_cache) is RadixCache: + return node.lock_ref == 0 + + if type(radix_cache) is SWARadixCache: + return node.full_lock_ref == 0 + + raise NotImplementedError( + f"walk_radix_cache_for_canary does not support {type(radix_cache).__name__}" + ) + + +def _node_is_swa_resident_for_canary( + *, + node: "TreeNode", + radix_cache: "BasePrefixCache", +) -> bool: + if type(radix_cache) is not SWARadixCache: + return True + + return not node.swa_tombstone diff --git a/python/sglang/srt/kv_canary/runner/canary_manager.py b/python/sglang/srt/kv_canary/runner/canary_manager.py index 98f25be53..af33a44de 100644 --- a/python/sglang/srt/kv_canary/runner/canary_manager.py +++ b/python/sglang/srt/kv_canary/runner/canary_manager.py @@ -16,6 +16,7 @@ from sglang.srt.kv_canary.endpoint import ( CanaryEndpoint, build_endpoints_from_group, ) +from sglang.srt.kv_canary.runner.sweep import SweepOrchestrator from sglang.srt.kv_canary.runner.violation_manager import ViolationManager from sglang.srt.kv_canary.single_forward_manager.manager import ( SingleForwardManager, @@ -24,6 +25,7 @@ from sglang.srt.kv_canary.single_forward_manager.manager import ( from sglang.srt.kv_canary.state import CanaryDeviceState if TYPE_CHECKING: + from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -82,6 +84,14 @@ class CanaryManager: d2h_stream=self._d2h_stream, outer_step_counter_getter=self._get_outer_step_counter, ) + self._sweep_orchestrator = SweepOrchestrator( + config=config, + device_state=self._device_state, + buffer_groups=self._buffer_groups, + endpoints=self._endpoints, + swa_window_size=self._swa_window_size, + outer_step_counter_getter=self._get_outer_step_counter, + ) self._single_forward_managers: tuple[SingleForwardManager, ...] = ( SingleForwardManager( config=config, @@ -174,6 +184,7 @@ class CanaryManager: ) -> None: for idx in single_forward_indices: self._single_forward_managers[idx].post_ops_outside_graph() + self._sweep_orchestrator.maybe_run_sweep() self._outer_step_counter += 1 self._violation_manager.step() @@ -182,6 +193,9 @@ class CanaryManager: single_forward_manager.phase_checker.enable_assert() self._device_state.enable_chain_position_assert.fill_(1) + def attach_radix_cache(self, radix_cache: "BasePrefixCache") -> None: + self._sweep_orchestrator.attach_radix_cache(radix_cache) + def _get_outer_step_counter(self) -> int: return self._outer_step_counter diff --git a/python/sglang/srt/kv_canary/runner/kernel_launcher.py b/python/sglang/srt/kv_canary/runner/kernel_launcher.py index 418727308..e7a2c3684 100644 --- a/python/sglang/srt/kv_canary/runner/kernel_launcher.py +++ b/python/sglang/srt/kv_canary/runner/kernel_launcher.py @@ -88,6 +88,7 @@ def launch_endpoints_per_forward( for endpoint in endpoints if _endpoint_belongs_to_group(endpoint, group) and tag_filter(endpoint.kernel_kind) + and not _is_sweep_tag(endpoint.kernel_kind) and passes_v_half_gate(endpoint.kernel_kind) ] assert len(active_endpoints) > 0 @@ -106,12 +107,46 @@ def launch_endpoints_per_forward( ) +def launch_endpoints_sweep( + *, + endpoints: tuple[CanaryEndpoint, ...], + group: CanaryBufferGroup, + verify_plan: VerifyPlan, + violation_log: ViolationLog, +) -> None: + active_endpoints = [ + endpoint + for endpoint in endpoints + if _endpoint_belongs_to_group(endpoint, group) + and _is_sweep_tag(endpoint.kernel_kind) + and passes_v_half_gate(endpoint.kernel_kind) + ] + assert len(active_endpoints) > 0 + + for endpoint in active_endpoints: + endpoint.launch_sweep( + verify_plan=verify_plan, + violation_log=violation_log, + ) + + +def _is_sweep_tag(tag: CanaryLaunchTag) -> bool: + return tag in ( + CanaryLaunchTag.SWEEP_K_FULL, + CanaryLaunchTag.SWEEP_V_FULL, + CanaryLaunchTag.SWEEP_K_SWA, + CanaryLaunchTag.SWEEP_V_SWA, + ) + + def _is_v_half_tag(tag: CanaryLaunchTag) -> bool: return tag in ( CanaryLaunchTag.HEAD_V_FULL, CanaryLaunchTag.TAIL_V_FULL, + CanaryLaunchTag.SWEEP_V_FULL, CanaryLaunchTag.HEAD_V_SWA, CanaryLaunchTag.TAIL_V_SWA, + CanaryLaunchTag.SWEEP_V_SWA, ) diff --git a/python/sglang/srt/kv_canary/runner/sweep.py b/python/sglang/srt/kv_canary/runner/sweep.py new file mode 100644 index 000000000..32d343f28 --- /dev/null +++ b/python/sglang/srt/kv_canary/runner/sweep.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import logging +from collections.abc import Callable +from typing import TYPE_CHECKING, Optional + +from sglang.srt.kv_canary.buffer_group import CanaryBufferGroup, PoolKind +from sglang.srt.kv_canary.config import CanaryConfig +from sglang.srt.kv_canary.endpoint import CanaryEndpoint +from sglang.srt.kv_canary.runner.kernel_launcher import launch_endpoints_sweep +from sglang.srt.kv_canary.state import CanaryDeviceState +from sglang.srt.kv_canary.sweep_plan_builder import build_verify_plan_radix_sweep + +if TYPE_CHECKING: + from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache + +logger = logging.getLogger(__name__) + + +class SweepOrchestrator: + def __init__( + self, + *, + config: CanaryConfig, + device_state: CanaryDeviceState, + buffer_groups: tuple[CanaryBufferGroup, ...], + endpoints: tuple[CanaryEndpoint, ...], + swa_window_size: int, + outer_step_counter_getter: Callable[[], int], + ) -> None: + self._config = config + self._device_state = device_state + self._buffer_groups = buffer_groups + self._endpoints = endpoints + self._swa_window_size = swa_window_size + self._outer_step_counter_getter = outer_step_counter_getter + self._radix_cache: Optional["BasePrefixCache"] = None + + self._last_sweep_step: int = -1 + self._sweep_passes: int = 0 + + @property + def sweep_passes(self) -> int: + return self._sweep_passes + + def attach_radix_cache(self, radix_cache: "BasePrefixCache") -> None: + self._radix_cache = radix_cache + + def maybe_run_sweep(self) -> None: + if self._config.sweep_interval == 0: + return + outer_step_counter = self._outer_step_counter_getter() + if ( + self._last_sweep_step >= 0 + and outer_step_counter - self._last_sweep_step < self._config.sweep_interval + ): + return + self._last_sweep_step = outer_step_counter + + if self._radix_cache is None: + return + + violation_log = self._device_state.violation_log + for group in self._buffer_groups: + window = self._swa_window_size if group.kind is PoolKind.SWA else 0 + verify_plan = build_verify_plan_radix_sweep( + radix_cache=self._radix_cache, + swa_window_size=window, + full_to_swa_index_mapping=group.swa_index_lut, + ) + launch_endpoints_sweep( + endpoints=self._endpoints, + group=group, + verify_plan=verify_plan, + violation_log=violation_log, + ) + + self._sweep_passes += 1 + logger.info( + "[canary] sweep succeeded %d times (last_step=%d)", + self._sweep_passes, + outer_step_counter, + ) diff --git a/python/sglang/srt/kv_canary/runner/violation_reporter.py b/python/sglang/srt/kv_canary/runner/violation_reporter.py index ed81b667a..abc5826d8 100644 --- a/python/sglang/srt/kv_canary/runner/violation_reporter.py +++ b/python/sglang/srt/kv_canary/runner/violation_reporter.py @@ -69,6 +69,13 @@ class ViolationReporter: def _canary_kind_label(tag: CanaryLaunchTag) -> str: name_lower = tag.name.lower() + if tag in ( + CanaryLaunchTag.SWEEP_K_FULL, + CanaryLaunchTag.SWEEP_V_FULL, + CanaryLaunchTag.SWEEP_K_SWA, + CanaryLaunchTag.SWEEP_V_SWA, + ): + return name_lower return f"per_forward_{name_lower}" diff --git a/python/sglang/srt/kv_canary/state.py b/python/sglang/srt/kv_canary/state.py index 7d19c7216..8b9a5a35c 100644 --- a/python/sglang/srt/kv_canary/state.py +++ b/python/sglang/srt/kv_canary/state.py @@ -13,10 +13,10 @@ from sglang.srt.kv_canary.config import CanaryConfig class ViolationLog: """Global violation sink shared across all canary launches. - One instance per canary runner — every launch (head / tail, K / V half, FULL / SWA group) writes + One instance per canary runner — every launch (head / tail / sweep, K / V half, FULL / SWA group) writes into the same ring. The kernel_kind field stamped into each violation row identifies which launch fired (kernel_kind is a static IntEnum tag — :class:`CanaryLaunchTag` in - ``sglang.jit_kernel.kv_canary.verify`` — with a unique value per (head|tail, K|V, FULL|SWA) tuple). + ``sglang.jit_kernel.kv_canary.verify`` — with a unique value per (head|tail|sweep, K|V, FULL|SWA) tuple). Ring capacity is sized generously (≥ 1024) so overflow is a non-concern in practice — violations are cold-path and the host raises at the first one anyway (or just logs it in mode="log"). atomicAdd @@ -63,7 +63,7 @@ class CanaryDeviceState: no per-step allocation. Fields: - violation_log: The single ViolationLog shared by every launch (head / tail × K / V × + violation_log: The single ViolationLog shared by every launch (head / tail / sweep × K / V × FULL / SWA). All kernels atomicAdd into violation_log.violation_write_index and stamp their CanaryLaunchTag into each violation row. kernel_run_counters: Per-CanaryLaunchTag int64 counter array, shape [num_tags], device. The diff --git a/python/sglang/srt/kv_canary/sweep_plan_builder.py b/python/sglang/srt/kv_canary/sweep_plan_builder.py new file mode 100644 index 000000000..3e9f43168 --- /dev/null +++ b/python/sglang/srt/kv_canary/sweep_plan_builder.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.jit_kernel.kv_canary.verify import VerifyPlan +from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary + +if TYPE_CHECKING: + from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache + + +def build_verify_plan_radix_sweep( + *, + radix_cache: "BasePrefixCache", + swa_window_size: int, + full_to_swa_index_mapping: Optional[torch.Tensor], + unlocked_only: bool = False, +) -> VerifyPlan: + """Build a sweep VerifyPlan directly from the radix-cache walker. + + The walker covers every slot held by the radix tree. There is no exclusion for slots also owned + by running requests because the overlap with per-forward HEAD/TAIL coverage is harmless + redundancy. The helper applies the SWA LUT before writing the plan. + """ + device = radix_cache.req_to_token_pool.req_to_token.device + + walk_result = walk_radix_cache_for_canary( + radix_cache=radix_cache, + unlocked_only=unlocked_only, + ) + slot_indices = walk_result.slot_indices.to(device) + positions = walk_result.positions.to(device) + prev_slot_indices = walk_result.prev_slot_indices.to(device) + + if swa_window_size > 0: + assert ( + full_to_swa_index_mapping is not None + ), "full_to_swa_index_mapping is required when SWA is enabled" + slot_indices = _swa_translate( + indices=slot_indices, lut=full_to_swa_index_mapping + ) + prev_slot_indices = _swa_translate( + indices=prev_slot_indices, lut=full_to_swa_index_mapping + ) + + num_valid = int(slot_indices.shape[0]) + verify_plan = VerifyPlan.allocate(verify_capacity=max(1, num_valid), device=device) + + verify_plan.verify_slot_indices[:num_valid].copy_(slot_indices) + verify_plan.verify_expected_positions[:num_valid].copy_(positions) + verify_plan.verify_prev_slot_indices[:num_valid].copy_(prev_slot_indices) + verify_plan.verify_num_valid.fill_(num_valid) + verify_plan.enable.fill_(1) + + return verify_plan + + +def _swa_translate( + *, + indices: torch.Tensor, + lut: torch.Tensor, +) -> torch.Tensor: + # 0 is both SWAKVPool's evicted-sentinel and the kernel's kTokenToKvSlotPadding, + # so evicted indices propagate as the canonical "no real slot" value and the + # verify kernel handles them. + if indices.numel() == 0: + return indices + lut_dev = lut.to(indices.device).to(torch.int64) + anchor_mask = indices < 0 + safe = torch.where(anchor_mask, torch.zeros_like(indices), indices).to(torch.int64) + looked_up = lut_dev[safe] + return torch.where(anchor_mask, indices.to(torch.int64), looked_up) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index e6d176bed..a21ebc94b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -451,6 +451,9 @@ class Scheduler( self.disable_radix_cache = result.disable_radix_cache self.tree_cache = result.tree_cache + if (c := self.tp_worker.model_runner.canary_manager) is not None: + c.attach_radix_cache(self.tree_cache) + if self.enable_hisparse: # Coordinator was created inside ModelRunner.initialize() before CUDA graph capture self.hisparse_coordinator = self.tp_worker.model_runner.hisparse_coordinator diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index efd65e23f..7a12d739a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -778,6 +778,7 @@ class ServerArgs: disable_attn_tp_gather: bool = False gc_threshold: Optional[List[int]] = None kv_canary: str = "none" + kv_canary_sweep_interval: int = 0 # Context parallelism used in the long sequence prefill phase of DeepSeek v3.2 enable_dsa_prefill_context_parallel: bool = False dsa_prefill_cp_mode: str = "round-robin-split" @@ -6318,6 +6319,12 @@ class ServerArgs: "'raise' fails the server on the first detected mismatch (CI lane)." ), ) + parser.add_argument( + "--kv-canary-sweep-interval", + type=int, + default=ServerArgs.kv_canary_sweep_interval, + help="Every N forward steps, run a full-pool sweep.", + ) parser.add_argument( "--cuda-graph-max-bs", type=int, @@ -7367,6 +7374,11 @@ class ServerArgs: "When setting gc_threshold, it must contain 1 to 3 integers." ) + if self.kv_canary_sweep_interval > 0 and self.kv_canary == "none": + raise ValueError( + "--kv-canary-sweep-interval requires --kv-canary in {log, raise}" + ) + def check_lora_server_args(self): assert self.max_loras_per_batch > 0, "max_loras_per_batch must be positive" diff --git a/python/sglang/test/kv_canary/fixtures.py b/python/sglang/test/kv_canary/fixtures.py index 44a18ae69..de92101d7 100644 --- a/python/sglang/test/kv_canary/fixtures.py +++ b/python/sglang/test/kv_canary/fixtures.py @@ -135,6 +135,7 @@ def make_base_config() -> CanaryConfig: return CanaryConfig( mode=CanaryMode.RAISE, ring_capacity=1024, + sweep_interval=0, ) diff --git a/python/sglang/test/kv_canary/runner_test_base.py b/python/sglang/test/kv_canary/runner_test_base.py index a55c9805b..ec3cc687b 100644 --- a/python/sglang/test/kv_canary/runner_test_base.py +++ b/python/sglang/test/kv_canary/runner_test_base.py @@ -24,10 +24,12 @@ def make_config( *, mode: CanaryMode = CanaryMode.RAISE, ring_capacity: int = 1024, + sweep_interval: int = 0, ) -> CanaryConfig: return CanaryConfig( mode=mode, ring_capacity=ring_capacity, + sweep_interval=sweep_interval, ) diff --git a/test/registered/kv_canary/test_self_e2e_baseline.py b/test/registered/kv_canary/test_self_e2e_baseline.py index 6974806a3..5d9b11701 100644 --- a/test/registered/kv_canary/test_self_e2e_baseline.py +++ b/test/registered/kv_canary/test_self_e2e_baseline.py @@ -11,7 +11,7 @@ register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") class _BaselineBase(CanaryE2EBase): - """No perturb, kv-canary=log. Server should run clean with no canary + """No perturb, kv-canary=log, sweep off. Server should run clean with no canary violations and every request must come back 200.""" kv_canary_mode = CanaryMode.LOG diff --git a/test/registered/kv_canary/test_self_unit_endpoint.py b/test/registered/kv_canary/test_self_unit_endpoint.py index 521d7a478..8c406d2dd 100644 --- a/test/registered/kv_canary/test_self_unit_endpoint.py +++ b/test/registered/kv_canary/test_self_unit_endpoint.py @@ -63,6 +63,28 @@ class TestSelfUnitEndpoint(CustomTestCase): def setUp(self): self.device = DEFAULT_DEVICE + def test_launch_sweep_only_calls_verify(self): + """Verify sweep launch invokes only the verify kernel.""" + calls: list[str] = [] + with patch.object( + endpoint_module, + "launch_canary_verify_kernel", + lambda **kwargs: calls.append("verify"), + ), patch.object( + endpoint_module, + "launch_canary_write_kernel", + lambda **kwargs: calls.append("write"), + ): + ep = _make_endpoint( + device=self.device, kernel_kind=CanaryLaunchTag.SWEEP_K_FULL + ) + args = _make_kernel_args(self.device) + ep.launch_sweep( + verify_plan=args.verify_plan, + violation_log=args.violation_log, + ) + self.assertEqual(calls, ["verify"]) + def test_launch_per_forward_passes_kernel_kind(self): """Verify per-forward launch passes the endpoint kernel kind.""" captured: list[tuple[str, CanaryLaunchTag]] = [] @@ -109,33 +131,19 @@ class TestSelfUnitEndpoint(CustomTestCase): ): shared_log = ViolationLog.allocate(ring_capacity=2, device=self.device) ep_a = _make_endpoint( - device=self.device, kernel_kind=CanaryLaunchTag.HEAD_K_FULL + device=self.device, kernel_kind=CanaryLaunchTag.SWEEP_K_FULL ) ep_b = _make_endpoint( - device=self.device, kernel_kind=CanaryLaunchTag.HEAD_V_FULL + device=self.device, kernel_kind=CanaryLaunchTag.SWEEP_V_FULL ) - args = _make_kernel_args(self.device) - ep_a.launch_per_forward( - verify_plan=args.verify_plan, - write_plan=args.write_plan, - input_ids=args.input_ids, - positions=args.positions, - out_cache_loc=args.out_cache_loc, - enable_write_input_assert=args.enable_write_input_assert, - enable_verify_token_assert=args.enable_verify_token_assert, - expected_inputs=args.expected_inputs, + plan = VerifyPlan.allocate(verify_capacity=1, device=self.device) + ep_a.launch_sweep( + verify_plan=plan, violation_log=shared_log, ) - ep_b.launch_per_forward( - verify_plan=args.verify_plan, - write_plan=args.write_plan, - input_ids=args.input_ids, - positions=args.positions, - out_cache_loc=args.out_cache_loc, - enable_write_input_assert=args.enable_write_input_assert, - enable_verify_token_assert=args.enable_verify_token_assert, - expected_inputs=args.expected_inputs, + ep_b.launch_sweep( + verify_plan=plan, violation_log=shared_log, ) self.assertEqual(captured_rings[0], captured_rings[1]) diff --git a/test/registered/kv_canary/test_self_unit_radix_walker.py b/test/registered/kv_canary/test_self_unit_radix_walker.py new file mode 100644 index 000000000..6ef7dce93 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_radix_walker.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import unittest + +import torch + +from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary +from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache, TreeNode +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE, make_radix_cache +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-small") + + +class TestSelfUnitRadixWalker(CustomTestCase): + def setUp(self): + self.device = DEFAULT_DEVICE + + def test_single_node_chain_positions_increase(self): + """Verify a single radix chain emits increasing positions.""" + chain = [10, 20, 30, 40] + cache = make_radix_cache([[], chain], device=self.device) + result = walk_radix_cache_for_canary(radix_cache=cache) + self.assertEqual(result.slot_indices.tolist(), chain) + self.assertEqual(result.positions.tolist(), [0, 1, 2, 3]) + self.assertEqual(result.prev_slot_indices.tolist(), [-1, 10, 20, 30]) + + def test_child_node_first_slot_prev_is_parent_last(self): + """Verify child chains link their first slot to the parent tail.""" + parent = [7, 8] + child = [9, 10] + cache = make_radix_cache([[], parent, child], device=self.device) + result = walk_radix_cache_for_canary(radix_cache=cache) + self.assertEqual(result.slot_indices.tolist(), parent + child) + self.assertEqual(result.prev_slot_indices.tolist()[len(parent)], parent[-1]) + + def test_root_child_first_slot_prev_minus_one(self): + """Verify root child chains use -1 as the initial previous slot.""" + cache = make_radix_cache([[], [42, 43]], device=self.device) + result = walk_radix_cache_for_canary(radix_cache=cache) + self.assertEqual(int(result.prev_slot_indices[0]), -1) + + def test_position_equals_depth_from_root(self): + """Verify emitted positions match depth from the radix root.""" + cache = make_radix_cache([[], [1, 2], [3], [4, 5]], device=self.device) + result = walk_radix_cache_for_canary(radix_cache=cache) + self.assertEqual(result.positions.tolist(), [0, 1, 2, 3, 4]) + self.assertEqual(result.slot_indices.tolist(), [1, 2, 3, 4, 5]) + + def test_walk_includes_locked_nodes_by_default(self): + """Verify radix walking includes locked nodes by default.""" + cache = make_radix_cache([[], [1, 2], [3, 4]], device=self.device) + locked_node = next(iter(cache.root_node.children.values())) + locked_node.lock_ref = 1 + result = walk_radix_cache_for_canary(radix_cache=cache) + self.assertEqual(result.slot_indices.tolist(), [1, 2, 3, 4]) + + def test_walk_unlocked_only_skips_locked(self): + """Verify unlocked-only radix walking skips locked nodes.""" + cache = make_radix_cache([[], [1, 2], [3, 4]], device=self.device) + locked_node = next(iter(cache.root_node.children.values())) + locked_node.lock_ref = 1 + result = walk_radix_cache_for_canary(radix_cache=cache, unlocked_only=True) + self.assertEqual(result.slot_indices.tolist(), [3, 4]) + + def test_walk_unlocked_only_uses_swa_full_lock_ref(self): + """Verify SWA radix walking honors full-pool lock references.""" + cache = SWARadixCache.__new__(SWARadixCache) + cache.device = self.device + cache.page_size = 1 + cache.disable = False + + root = TreeNode() + root.value = torch.tensor([], dtype=torch.int32, device=self.device) + cache.root_node = root + + locked_child = TreeNode() + locked_child.value = torch.tensor([1, 2], dtype=torch.int32, device=self.device) + locked_child.parent = root + locked_child.full_lock_ref = 1 + root.children[locked_child.id] = locked_child + + unlocked_child = TreeNode() + unlocked_child.value = torch.tensor( + [3, 4], dtype=torch.int32, device=self.device + ) + unlocked_child.parent = root + root.children[unlocked_child.id] = unlocked_child + + result = walk_radix_cache_for_canary(radix_cache=cache, unlocked_only=True) + self.assertEqual(result.slot_indices.tolist(), [3, 4]) + + def test_swa_resident_only_skips_tombstoned_nodes(self): + """Verify SWA radix walking skips nodes whose SWA storage was evicted.""" + cache = SWARadixCache.__new__(SWARadixCache) + cache.device = self.device + cache.page_size = 1 + cache.disable = False + + root = TreeNode() + root.value = torch.tensor([], dtype=torch.int32, device=self.device) + cache.root_node = root + + tombstoned_child = TreeNode() + tombstoned_child.value = torch.tensor( + [1, 2], dtype=torch.int32, device=self.device + ) + tombstoned_child.parent = root + tombstoned_child.swa_tombstone = True + root.children[tombstoned_child.id] = tombstoned_child + + resident_child = TreeNode() + resident_child.value = torch.tensor( + [3, 4], dtype=torch.int32, device=self.device + ) + resident_child.parent = root + root.children[resident_child.id] = resident_child + + result = walk_radix_cache_for_canary( + radix_cache=cache, + swa_resident_only=True, + ) + self.assertEqual(result.slot_indices.tolist(), [3, 4]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_runner_sweep.py b/test/registered/kv_canary/test_self_unit_runner_sweep.py new file mode 100644 index 000000000..a2d02e1c6 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_runner_sweep.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import unittest +from unittest.mock import patch + +from sglang.srt.kv_canary import endpoint as endpoint_module +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.fixtures import ( + make_forward_batch, + make_radix_cache, + make_req_to_token_pool, +) +from sglang.test.kv_canary.runner_test_base import ( + CanaryManagerTestCase, + make_config, + make_manager, +) + +register_cuda_ci(est_time=45, stage="extra-a", runner_config="1-gpu-small") + + +def _run_one_cycle(manager, forward_batch) -> None: + with manager.with_ops_outside_graph( + single_forward_indices=[0], + maybe_inaccurate_forward_batch=forward_batch, + ): + with manager.with_active_single_forward_manager(0): + pre_ops_output = manager.pre_ops_maybe_inside_graph(forward_batch) + manager.post_ops_maybe_inside_graph(forward_batch, pre_ops_output) + + +class TestSelfUnitManagerSweep(CanaryManagerTestCase): + def test_sweep_every_n_cadence(self) -> None: + """Verify sweep execution follows the configured step cadence.""" + config = make_config(sweep_interval=4) + manager = make_manager(device=self.device, config=config) + forward_batch = make_forward_batch(self.device) + + sweep_calls: list[int] = [] + real_maybe = manager._sweep_orchestrator.maybe_run_sweep + + def _spy() -> None: + before = manager._sweep_orchestrator._last_sweep_step + real_maybe() + if manager._sweep_orchestrator._last_sweep_step != before: + sweep_calls.append(manager._outer_step_counter) + + with patch.object(manager._sweep_orchestrator, "maybe_run_sweep", _spy): + for _ in range(12): + _run_one_cycle(manager, forward_batch) + self.assertEqual(sweep_calls, [0, 4, 8]) + + def test_sweep_path_launches_sweep_kernels(self) -> None: + """Verify sweep paths launch sweep verify kernels.""" + config = make_config(sweep_interval=1) + manager = make_manager(device=self.device, config=config) + forward_batch = make_forward_batch(self.device) + manager._single_forward_managers[0].pre_ops_outside_graph( + maybe_inaccurate_forward_batch=forward_batch + ) + with manager.with_active_single_forward_manager(0): + manager.pre_ops_maybe_inside_graph(forward_batch) + + cache = make_radix_cache([[], [10, 11, 12]], device=self.device) + cache.req_to_token_pool = make_req_to_token_pool(self.device) + manager.attach_radix_cache(cache) + + sweep_kernel_kinds: list[str] = [] + with patch.object( + endpoint_module, + "launch_canary_verify_kernel", + lambda **kwargs: sweep_kernel_kinds.append( + kwargs["context"].kernel_kind.name + ), + ): + manager._sweep_orchestrator.maybe_run_sweep() + self.assertTrue(any("SWEEP" in kind for kind in sweep_kernel_kinds)) + + def test_sweep_allocates_verify_plan_from_walker_output(self) -> None: + """Verify sweep planning sizes the verify plan from walker output.""" + manager = make_manager(device=self.device) + cache = make_radix_cache([[], [10, 11], [12, 13, 14]], device=self.device) + cache.req_to_token_pool = make_req_to_token_pool(self.device) + manager.attach_radix_cache(cache) + + valid_counts: list[int] = [] + with patch.object( + endpoint_module, + "launch_canary_verify_kernel", + lambda **kwargs: valid_counts.append( + int(kwargs["plan"].verify_num_valid.item()) + ), + ): + manager._sweep_orchestrator.maybe_run_sweep() + self.assertTrue(all(count == 5 for count in valid_counts)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_sweep_plan_builder.py b/test/registered/kv_canary/test_self_unit_sweep_plan_builder.py new file mode 100644 index 000000000..b9d5c0731 --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_sweep_plan_builder.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import unittest + +import torch + +from sglang.srt.kv_canary.sweep_plan_builder import build_verify_plan_radix_sweep +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.fixtures import ( + DEFAULT_DEVICE, + make_radix_cache, + make_req_to_token_pool, +) +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-small") + + +class TestSelfUnitSweepPlanBuilder(CustomTestCase): + def setUp(self) -> None: + self.device = DEFAULT_DEVICE + + def test_build_verify_plan_radix_sweep(self) -> None: + """Verify radix sweep verify plans include cached slot chains.""" + empty_cache = make_radix_cache([[]], device=self.device) + empty_cache.req_to_token_pool = make_req_to_token_pool(self.device) + empty_out = build_verify_plan_radix_sweep( + radix_cache=empty_cache, + swa_window_size=0, + full_to_swa_index_mapping=None, + ) + self.assertEqual(int(empty_out.verify_num_valid.item()), 0) + + cache = make_radix_cache([[], [100, 101, 102]], device=self.device) + cache.req_to_token_pool = make_req_to_token_pool(self.device) + out = build_verify_plan_radix_sweep( + radix_cache=cache, + swa_window_size=0, + full_to_swa_index_mapping=None, + ) + self.assertEqual(int(out.verify_num_valid.item()), 3) + self.assertEqual(out.verify_slot_indices.dtype, torch.int64) + self.assertEqual(out.verify_expected_positions.dtype, torch.int64) + self.assertEqual(out.verify_prev_slot_indices.dtype, torch.int64) + self.assertEqual(out.verify_slot_indices[:3].tolist(), [100, 101, 102]) + self.assertEqual(out.verify_expected_positions[:3].tolist(), [0, 1, 2]) + self.assertEqual(out.verify_prev_slot_indices[:3].tolist(), [-1, 100, 101]) + + def test_radix_held_slot_still_swept(self) -> None: + """Verify held radix slots are still included in sweep plans.""" + cache = make_radix_cache([[], [42, 43, 44]], device=self.device) + cache.req_to_token_pool = make_req_to_token_pool(self.device) + out = build_verify_plan_radix_sweep( + radix_cache=cache, + swa_window_size=0, + full_to_swa_index_mapping=None, + ) + num_valid = int(out.verify_num_valid.item()) + self.assertEqual(num_valid, 3) + self.assertEqual( + set(out.verify_slot_indices[:num_valid].tolist()), {42, 43, 44} + ) + + def test_truly_free_slot_not_swept(self) -> None: + """Verify free radix slots are excluded from sweep plans.""" + empty_cache = make_radix_cache([[]], device=self.device) + empty_cache.req_to_token_pool = make_req_to_token_pool(self.device) + out = build_verify_plan_radix_sweep( + radix_cache=empty_cache, + swa_window_size=0, + full_to_swa_index_mapping=None, + ) + self.assertEqual(int(out.verify_num_valid.item()), 0) + + def test_swa_translate_preserves_evicted_as_padding_sentinel(self) -> None: + """Evicted (LUT=0) slots stay in the plan as the padding sentinel; the kernel does the skipping.""" + cache = make_radix_cache([[], [100, 101, 102]], device=self.device) + cache.req_to_token_pool = make_req_to_token_pool(self.device) + + lut = torch.zeros(200, dtype=torch.int64, device=self.device) + lut[100] = 500 + lut[101] = 0 + lut[102] = 502 + + out = build_verify_plan_radix_sweep( + radix_cache=cache, + swa_window_size=128, + full_to_swa_index_mapping=lut, + ) + num_valid = int(out.verify_num_valid.item()) + self.assertEqual(num_valid, 3) + self.assertEqual(out.verify_slot_indices[:num_valid].tolist(), [500, 0, 502]) + self.assertEqual( + out.verify_prev_slot_indices[:num_valid].tolist(), [-1, 500, 0] + ) + self.assertEqual(out.verify_expected_positions[:num_valid].tolist(), [0, 1, 2]) + + +if __name__ == "__main__": + unittest.main()