From 678e73a9ee89f0ad2c63c02bbc53c2be96e83d95 Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Sun, 31 May 2026 09:57:43 +0800 Subject: [PATCH] Add a deterministic token oracle and production write-input assertion (#26815) --- python/sglang/srt/environ.py | 5 +- python/sglang/srt/kv_canary/api.py | 4 + python/sglang/srt/kv_canary/config.py | 9 +- .../sglang/srt/kv_canary/perturb/__init__.py | 0 .../srt/kv_canary/perturb/next_token_swap.py | 83 ++++++ .../srt/kv_canary/runner/canary_manager.py | 5 + .../srt/kv_canary/runner/kernel_launcher.py | 4 +- .../single_forward_manager/manager.py | 44 +++- .../srt/kv_canary/token_oracle/__init__.py | 0 .../srt/kv_canary/token_oracle/install.py | 22 ++ .../srt/kv_canary/token_oracle/oracle.py | 50 ++++ .../kv_canary/token_oracle/oracle_manager.py | 110 ++++++++ .../srt/kv_canary/token_oracle/sampler.py | 59 +++++ .../sglang/srt/model_executor/model_runner.py | 6 + python/sglang/srt/server_args.py | 2 + python/sglang/test/kv_canary/fixtures.py | 1 + .../sglang/test/kv_canary/runner_test_base.py | 2 + python/sglang/test/mock_model/utils.py | 9 +- .../kv_canary/test_self_e2e_pr_25015.py | 86 +++++++ .../kv_canary/test_self_unit_token_oracle.py | 57 +++++ .../test_self_e2e_perturb_next_token_swap.py | 36 +++ .../test_self_unit_canary_mock_wiring.py | 238 ++++++++++++++++++ .../mock_model/test_self_unit_install.py | 50 ++++ .../mock_model/test_self_unit_oracle.py | 127 ++++++++++ .../test_self_unit_oracle_torch_vs_ref.py | 85 +++++++ .../test_self_unit_sampler_hookpoint.py | 49 ++++ .../unit/server_args/test_server_args.py | 56 +++++ 27 files changed, 1192 insertions(+), 7 deletions(-) create mode 100644 python/sglang/srt/kv_canary/perturb/__init__.py create mode 100644 python/sglang/srt/kv_canary/perturb/next_token_swap.py create mode 100644 python/sglang/srt/kv_canary/token_oracle/__init__.py create mode 100644 python/sglang/srt/kv_canary/token_oracle/install.py create mode 100644 python/sglang/srt/kv_canary/token_oracle/oracle.py create mode 100644 python/sglang/srt/kv_canary/token_oracle/oracle_manager.py create mode 100644 python/sglang/srt/kv_canary/token_oracle/sampler.py create mode 100644 test/registered/kv_canary/test_self_e2e_pr_25015.py create mode 100644 test/registered/kv_canary/test_self_unit_token_oracle.py create mode 100644 test/registered/mock_model/test_self_e2e_perturb_next_token_swap.py create mode 100644 test/registered/mock_model/test_self_unit_canary_mock_wiring.py create mode 100644 test/registered/mock_model/test_self_unit_install.py create mode 100644 test/registered/mock_model/test_self_unit_oracle.py create mode 100644 test/registered/mock_model/test_self_unit_oracle_torch_vs_ref.py create mode 100644 test/registered/mock_model/test_self_unit_sampler_hookpoint.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index e77eab7a3..547ed453c 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -750,9 +750,12 @@ class Envs: SGLANG_PLUGINS = EnvStr("") # =================================================================== - # KV-Canary (testing-only) + # KV-Canary / Token-Oracle (testing-only) # =================================================================== SGLANG_KV_CANARY_RING_CAPACITY = EnvInt(1024) + SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT = EnvBool(False) + SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS = EnvInt(50) + SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB = EnvFloat(0.0) SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE = EnvBool(False) SGLANG_KV_CANARY_ENABLE_MHA_V = EnvBool(False) diff --git a/python/sglang/srt/kv_canary/api.py b/python/sglang/srt/kv_canary/api.py index fd6708cb0..2d1eaab90 100644 --- a/python/sglang/srt/kv_canary/api.py +++ b/python/sglang/srt/kv_canary/api.py @@ -13,6 +13,7 @@ from sglang.srt.kv_canary.runner.canary_manager import CanaryManager from sglang.srt.model_executor.forward_batch_info import ForwardBatch if TYPE_CHECKING: + from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager from sglang.srt.model_executor.model_runner import ModelRunner from sglang.srt.server_args import ServerArgs @@ -23,6 +24,7 @@ def install_canary( *, server_args: "ServerArgs", model_runner: "ModelRunner", + token_oracle_manager: Optional["TokenOracleManager"] = None, ) -> Optional[CanaryManager]: config = CanaryConfig.from_env(server_args) if config.mode is CanaryMode.NONE: @@ -59,7 +61,9 @@ def install_canary( req_to_token_pool=model_runner.req_to_token_pool, launch_capacities=launch_capacities, swa_window_size=swa_window_size, + token_oracle_manager=token_oracle_manager, speculative_num_steps=speculative_num_steps, + is_eagle_draft_decode=model_runner.is_draft_worker, ) _patch_model_forward(model_runner=model_runner, manager=manager) diff --git a/python/sglang/srt/kv_canary/config.py b/python/sglang/srt/kv_canary/config.py index 729bffc17..0b7f6cd73 100644 --- a/python/sglang/srt/kv_canary/config.py +++ b/python/sglang/srt/kv_canary/config.py @@ -20,7 +20,7 @@ class CanaryMode(str, Enum): class CanaryConfig: """Top-level canary configuration. All knobs live here; nothing reads env vars deeper in the stack. - Constructed once inside install_canary(server_args, model_runner) via + Constructed once inside install_canary(server_args, model_runner, token_oracle_manager) via CanaryConfig.from_env(server_args), then frozen and threaded through the canary stack. Subsequent runtime never mutates it. @@ -33,11 +33,17 @@ class CanaryConfig: 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. + enable_write_input_assert: bool. True = launch_canary_write_kernel additionally compares + forward_batch.input_ids[i] / positions[i] against caller-supplied expected_input_tokens[i] / + expected_input_positions[i]; mismatch records a violation. Only useful when something else + (e.g. token_oracle.oracle_manager.fill_expected_inputs) is feeding the expected_* placeholders + per forward — canary itself knows no oracle. """ mode: CanaryMode ring_capacity: int sweep_interval: int + enable_write_input_assert: bool @classmethod def from_env(cls, server_args: "ServerArgs") -> "CanaryConfig": @@ -51,4 +57,5 @@ class CanaryConfig: mode=CanaryMode(mode_raw), ring_capacity=envs.SGLANG_KV_CANARY_RING_CAPACITY.get(), sweep_interval=server_args.kv_canary_sweep_interval, + enable_write_input_assert=envs.SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT.get(), ) diff --git a/python/sglang/srt/kv_canary/perturb/__init__.py b/python/sglang/srt/kv_canary/perturb/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/kv_canary/perturb/next_token_swap.py b/python/sglang/srt/kv_canary/perturb/next_token_swap.py new file mode 100644 index 000000000..d17da61d4 --- /dev/null +++ b/python/sglang/srt/kv_canary/perturb/next_token_swap.py @@ -0,0 +1,83 @@ +"""Swap two requests' sampled next tokens at the sampler exit. + +KV path is untouched, so kv_canary KV-side fail_reasons stay silent. The +token-oracle input check downstream MUST report fail_reason=write_token — this +validates that the input-check link is genuinely active. +""" + +from __future__ import annotations + +import logging +import random +from dataclasses import dataclass +from typing import Optional + +import torch + +from sglang.srt.environ import envs + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class NextTokenSwapConfig: + prob: float + warmup_steps: int + + @classmethod + def from_env(cls) -> "NextTokenSwapConfig": + return cls( + prob=envs.SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB.get(), + warmup_steps=envs.SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS.get(), + ) + + +_config: Optional[NextTokenSwapConfig] = None +_step_counter: int = 0 + + +def _get_config() -> NextTokenSwapConfig: + global _config + if _config is None: + _config = NextTokenSwapConfig.from_env() + return _config + + +def maybe_perturb_swap_next_tokens( + batch_next_token_ids: torch.Tensor, +) -> torch.Tensor: + global _step_counter + + config = _get_config() + step = _step_counter + _step_counter += 1 + + if config.prob <= 0.0: + return batch_next_token_ids + if step < config.warmup_steps: + return batch_next_token_ids + if batch_next_token_ids.shape[0] < 2: + return batch_next_token_ids + + if random.random() >= config.prob: + return batch_next_token_ids + + batch_size = batch_next_token_ids.shape[0] + i = random.randrange(batch_size) + j = random.randrange(batch_size) + while j == i: + j = random.randrange(batch_size) + + swapped = batch_next_token_ids.clone() + swapped[i], swapped[j] = ( + batch_next_token_ids[j].clone(), + batch_next_token_ids[i].clone(), + ) + + logger.info( + "kv_canary perturb next_token_swap: swapped i=%d j=%d step=%d", + i, + j, + step, + ) + return swapped diff --git a/python/sglang/srt/kv_canary/runner/canary_manager.py b/python/sglang/srt/kv_canary/runner/canary_manager.py index 27d0c3d83..bff3026be 100644 --- a/python/sglang/srt/kv_canary/runner/canary_manager.py +++ b/python/sglang/srt/kv_canary/runner/canary_manager.py @@ -23,6 +23,7 @@ from sglang.srt.kv_canary.single_forward_manager.manager import ( _PreOpsMaybeInsideGraphOutput, ) from sglang.srt.kv_canary.state import CanaryDeviceState +from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager if TYPE_CHECKING: from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache @@ -42,7 +43,9 @@ class CanaryManager: req_to_token_pool: "ReqToTokenPool", launch_capacities: CanaryLaunchCapacities, swa_window_size: int = 0, + token_oracle_manager: Optional[TokenOracleManager] = None, speculative_num_steps: int = 1, + is_eagle_draft_decode: bool = False, ) -> None: self.config = config self._req_to_token_pool = req_to_token_pool @@ -107,6 +110,8 @@ class CanaryManager: per_forward_write_req_capacity=launch_capacities.per_forward_write_req_capacity, per_forward_write_entry_capacity=launch_capacities.per_forward_write_entry_capacity, d2h_stream=self._d2h_stream, + token_oracle_manager=token_oracle_manager, + is_eagle_draft_decode=is_eagle_draft_decode, ) for _ in range(num_sfms) ) diff --git a/python/sglang/srt/kv_canary/runner/kernel_launcher.py b/python/sglang/srt/kv_canary/runner/kernel_launcher.py index e7a2c3684..63f9f97c4 100644 --- a/python/sglang/srt/kv_canary/runner/kernel_launcher.py +++ b/python/sglang/srt/kv_canary/runner/kernel_launcher.py @@ -64,8 +64,8 @@ def launch_endpoints_per_forward( forward_batch: "ForwardBatch", expected_inputs: ExpectedInputs, violation_log: ViolationLog, - enable_write_input_assert: bool = False, - enable_verify_token_assert: bool = False, + enable_write_input_assert: bool, + enable_verify_token_assert: bool, ) -> None: positions = _canonicalize_boundary_int64(forward_batch.positions, _POSITIONS) out_cache_loc = _canonicalize_boundary_int64(forward_batch.out_cache_loc, _OUT_LOC) diff --git a/python/sglang/srt/kv_canary/single_forward_manager/manager.py b/python/sglang/srt/kv_canary/single_forward_manager/manager.py index c9e0ff23f..c85f9395e 100644 --- a/python/sglang/srt/kv_canary/single_forward_manager/manager.py +++ b/python/sglang/srt/kv_canary/single_forward_manager/manager.py @@ -2,7 +2,7 @@ from __future__ import annotations from dataclasses import dataclass from enum import IntEnum -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional import torch @@ -22,6 +22,7 @@ from sglang.srt.kv_canary.single_forward_manager.data import ( PostOpsInsideGraphOutputBuffer, ) from sglang.srt.kv_canary.state import CanaryDeviceState +from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager from sglang.srt.utils.phase_checker import SimplePhaseChecker if TYPE_CHECKING: @@ -67,6 +68,8 @@ class SingleForwardManager: per_forward_write_req_capacity: int, per_forward_write_entry_capacity: int, d2h_stream: torch.cuda.Stream, + token_oracle_manager: Optional[TokenOracleManager], + is_eagle_draft_decode: bool, ) -> None: self._config = config self._device = device @@ -76,6 +79,8 @@ class SingleForwardManager: self._req_to_token_pool = req_to_token_pool self._swa_window_size = swa_window_size self._d2h_stream = d2h_stream + self._token_oracle_manager: Optional[TokenOracleManager] = token_oracle_manager + self._is_eagle_draft_decode: bool = is_eagle_draft_decode self._write_req_capacity = per_forward_write_req_capacity self._write_entry_capacity = per_forward_write_entry_capacity @@ -153,6 +158,22 @@ class SingleForwardManager: bs_capacity=self._write_req_capacity, device=self._device ) + enable_write_input_assert = self._should_enable_write_input_assert_for_launch( + forward_batch + ) + if enable_write_input_assert: + manager = self._token_oracle_manager + if manager is None: + raise RuntimeError( + "kv-canary: enable_write_input_assert=True requires a TokenOracleManager; pass " + "token_oracle_manager=install_oracle_sampler(oracle=...) into " + "install_canary(...)" + ) + manager.fill_expected_inputs( + forward_batch=forward_batch, + expected_inputs_out=expected_inputs, + ) + plan_input.fill_from_forward_batch(forward_batch=forward_batch) violation_log = self._device_state.violation_log @@ -180,6 +201,8 @@ class SingleForwardManager: forward_batch=forward_batch, expected_inputs=expected_inputs_slice, violation_log=violation_log, + enable_write_input_assert=enable_write_input_assert, + enable_verify_token_assert=False, ) return _PreOpsMaybeInsideGraphOutput( @@ -202,6 +225,9 @@ class SingleForwardManager: violation_log = self._device_state.violation_log num_tokens = int(forward_batch.positions.shape[0]) expected_inputs_slice = pre_ops_output.expected_inputs.slice(num_tokens) + enable_write_input_assert = self._should_enable_write_input_assert_for_launch( + forward_batch + ) for group_idx, group in enumerate(self._buffer_groups): launch_endpoints_per_forward( endpoints=self._endpoints, @@ -212,6 +238,8 @@ class SingleForwardManager: forward_batch=forward_batch, expected_inputs=expected_inputs_slice, violation_log=violation_log, + enable_write_input_assert=enable_write_input_assert, + enable_verify_token_assert=False, ) verify_plan_enable_combined = _torch_reduce_minimum( @@ -233,6 +261,20 @@ class SingleForwardManager: self._enable_warner.tick(self._output_buffer.verify_plan_enable) + def _should_enable_write_input_assert_for_launch( + self, forward_batch: "ForwardBatch" + ) -> bool: + if not self._config.enable_write_input_assert: + return False + forward_mode = forward_batch.forward_mode + if ( + self._is_eagle_draft_decode + and forward_mode is not None + and forward_mode.is_decode() + ): + return False + return True + def _is_head_tag(tag: CanaryLaunchTag) -> bool: return tag in ( diff --git a/python/sglang/srt/kv_canary/token_oracle/__init__.py b/python/sglang/srt/kv_canary/token_oracle/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/python/sglang/srt/kv_canary/token_oracle/install.py b/python/sglang/srt/kv_canary/token_oracle/install.py new file mode 100644 index 000000000..2a57a76ee --- /dev/null +++ b/python/sglang/srt/kv_canary/token_oracle/install.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +from sglang.srt.kv_canary.token_oracle.oracle import HashOracle +from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager +from sglang.srt.kv_canary.token_oracle.sampler import install_oracle_sampler + +if TYPE_CHECKING: + from sglang.srt.server_args import ServerArgs + + +def install_token_oracle_from_env( + *, server_args: "ServerArgs", vocab_size: int +) -> Optional[TokenOracleManager]: + # Must be called before create_sampler() so the factory is present when the + # Sampler is first constructed. + if server_args.sampling_backend != "token_oracle": + return None + + oracle = HashOracle(vocab_size=vocab_size) + return install_oracle_sampler(oracle=oracle) diff --git a/python/sglang/srt/kv_canary/token_oracle/oracle.py b/python/sglang/srt/kv_canary/token_oracle/oracle.py new file mode 100644 index 000000000..2fc70ac71 --- /dev/null +++ b/python/sglang/srt/kv_canary/token_oracle/oracle.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Protocol + +import torch + + +class TokenOracle(Protocol): + """Deterministic (generalized_req_id, position) -> token_id mapping.""" + + def expected_tokens( + self, *, generalized_req_ids: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: ... + + +@dataclass(frozen=True, slots=True, kw_only=True) +class HashOracle: + """token_id = splitmix64(generalized_req_id XOR position) % vocab_size.""" + + vocab_size: int + + def expected_tokens( + self, *, generalized_req_ids: torch.Tensor, positions: torch.Tensor + ) -> torch.Tensor: + x = generalized_req_ids.to(torch.int64) ^ positions.to(torch.int64) + x = _splitmix64_tensor(x) + return _uint64_mod(x, self.vocab_size).to(torch.int32) + + +_C1: int = -4658895280553007687 # 0xBF58476D1CE4E5B9 as signed int64 +_C2: int = -7723592293110705685 # 0x94D049BB133111EB as signed int64 + + +def _splitmix64_tensor(x: torch.Tensor) -> torch.Tensor: + x = (x ^ _logical_shr(x, 30)) * _C1 + x = (x ^ _logical_shr(x, 27)) * _C2 + x = x ^ _logical_shr(x, 31) + return x + + +def _logical_shr(x: torch.Tensor, n: int) -> torch.Tensor: + return (x >> n) & ((1 << (64 - n)) - 1) + + +def _uint64_mod(x: torch.Tensor, mod: int) -> torch.Tensor: + offset = (1 << 64) % mod + base = x % mod + correction = (x < 0).to(x.dtype) * offset + return (base + correction) % mod diff --git a/python/sglang/srt/kv_canary/token_oracle/oracle_manager.py b/python/sglang/srt/kv_canary/token_oracle/oracle_manager.py new file mode 100644 index 000000000..0935607c7 --- /dev/null +++ b/python/sglang/srt/kv_canary/token_oracle/oracle_manager.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.srt.kv_canary.expected_inputs import ExpectedInputs +from sglang.srt.kv_canary.token_oracle.oracle import TokenOracle + +if TYPE_CHECKING: + from sglang.srt.model_executor.forward_batch_info import ForwardBatch + + +class TokenOracleManager: + def __init__(self, *, oracle: TokenOracle) -> None: + self.oracle = oracle + + def fill_expected_inputs( + self, + *, + forward_batch: "ForwardBatch", + expected_inputs_out: ExpectedInputs, + ) -> None: + positions = forward_batch.positions + input_ids = forward_batch.input_ids + num_tokens = int(input_ids.shape[0]) + + if num_tokens == 0: + return + + generalized_req_ids = _build_generalized_req_id_per_token( + forward_batch=forward_batch, + num_tokens=num_tokens, + generalized_req_ids_per_row=select_generalized_req_ids( + vanilla_req_ids=forward_batch.rids_int, + bootstrap_room_ids_int=forward_batch.bootstrap_room_ids_int, + ), + ) + if forward_batch.forward_mode.is_extend(): + expected_tokens = input_ids + else: + expected_tokens = self.oracle.expected_tokens( + generalized_req_ids=generalized_req_ids, + positions=positions.to(torch.int64), + ) + expected_inputs_out.tokens[:num_tokens].copy_(expected_tokens.to(torch.int64)) + expected_inputs_out.positions[:num_tokens].copy_(positions.to(torch.int64)) + + def sample_next_tokens( + self, *, generalized_req_ids: torch.Tensor, logits_positions: torch.Tensor + ) -> torch.Tensor: + return self.oracle.expected_tokens( + generalized_req_ids=generalized_req_ids, + positions=logits_positions.to(torch.int64) + 1, + ) + + +def _build_generalized_req_id_per_token( + *, + forward_batch: "ForwardBatch", + num_tokens: int, + generalized_req_ids_per_row: torch.Tensor, +) -> torch.Tensor: + forward_mode = forward_batch.forward_mode + if forward_mode.is_target_verify(): + per_req = int(forward_batch.spec_info.draft_token_num) + result = _expand_uniform(generalized_req_ids_per_row, per_req) + elif forward_mode.is_draft_extend(include_v2=True): + per_req = int(forward_batch.spec_info.num_tokens_per_req) + result = _expand_uniform(generalized_req_ids_per_row, per_req) + elif forward_mode.is_extend(): + extend_seq_lens = forward_batch.extend_seq_lens + if extend_seq_lens is None: + raise RuntimeError( + "_build_generalized_req_id_per_token: extend_seq_lens is None in extend mode" + ) + lens = extend_seq_lens.to(torch.int64) + result = torch.repeat_interleave(generalized_req_ids_per_row, lens) + else: + result = generalized_req_ids_per_row + + if int(result.shape[0]) != num_tokens: + raise RuntimeError( + f"fill_expected_inputs: sum(lens)={int(result.shape[0])} != num_tokens={num_tokens}" + ) + return result + + +def _expand_uniform(values: torch.Tensor, per_row: int) -> torch.Tensor: + bs = int(values.shape[0]) + return values.unsqueeze(1).expand(bs, per_row).reshape(bs * per_row) + + +def select_generalized_req_ids( + *, + vanilla_req_ids: torch.Tensor, + bootstrap_room_ids_int: torch.Tensor | None, +) -> torch.Tensor: + if bootstrap_room_ids_int is None: + return vanilla_req_ids + + bootstrap_room_ids_int = bootstrap_room_ids_int.to( + device=vanilla_req_ids.device, + dtype=torch.int64, + ) + return torch.where( + bootstrap_room_ids_int >= 0, + bootstrap_room_ids_int, + vanilla_req_ids.to(torch.int64), + ) diff --git a/python/sglang/srt/kv_canary/token_oracle/sampler.py b/python/sglang/srt/kv_canary/token_oracle/sampler.py new file mode 100644 index 000000000..2b96f4882 --- /dev/null +++ b/python/sglang/srt/kv_canary/token_oracle/sampler.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, List + +import torch + +from sglang.srt.kv_canary.perturb.next_token_swap import maybe_perturb_swap_next_tokens +from sglang.srt.kv_canary.token_oracle.oracle import TokenOracle +from sglang.srt.kv_canary.token_oracle.oracle_manager import ( + TokenOracleManager, + select_generalized_req_ids, +) +from sglang.srt.layers.sampler import Sampler, register_sampler_backend + +if TYPE_CHECKING: + from sglang.srt.layers.logits_processor import LogitsProcessorOutput + from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo + + +def install_oracle_sampler(*, oracle: TokenOracle) -> TokenOracleManager: + manager = TokenOracleManager(oracle=oracle) + register_sampler_backend( + "token_oracle", + lambda: _OracleSampler(token_oracle_manager=manager), + ) + return manager + + +class _OracleSampler(Sampler): + def __init__(self, *, token_oracle_manager: TokenOracleManager) -> None: + super().__init__() + self._token_oracle_manager = token_oracle_manager + + def forward( + self, + logits_output: "LogitsProcessorOutput", + sampling_info: "SamplingBatchInfo", + return_logprob: bool, + top_logprobs_nums: List[int], + token_ids_logprobs: List[List[int]], + positions: torch.Tensor, + ) -> torch.Tensor: + vanilla_req_ids = sampling_info.rids_int + if vanilla_req_ids is None: + raise RuntimeError( + "_OracleSampler.forward: generalized_req_id source tensor is None; " + "token oracle requires a per-forward generalized_req_id source tensor " + "(set in ForwardBatch.init_new when SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE=1)" + ) + batch_next_token_ids = self._token_oracle_manager.sample_next_tokens( + generalized_req_ids=select_generalized_req_ids( + vanilla_req_ids=vanilla_req_ids, + bootstrap_room_ids_int=sampling_info.bootstrap_room_ids_int, + ), + logits_positions=positions, + ) + + batch_next_token_ids = maybe_perturb_swap_next_tokens(batch_next_token_ids) + return batch_next_token_ids diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index dc460cdf4..8d59c848d 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -107,6 +107,7 @@ from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner from sglang.srt.kv_canary.api import install_canary from sglang.srt.kv_canary.runner.canary_manager import context_tuple +from sglang.srt.kv_canary.token_oracle.install import install_token_oracle_from_env from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.attention_registry import ( ATTENTION_BACKENDS, @@ -650,6 +651,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): if self.server_args.elastic_ep_backend: ElasticEPStateManager.init(self.server_args) + self._token_oracle_manager = install_token_oracle_from_env( + server_args=server_args, + vocab_size=self.model_config.vocab_size, + ) # Load the model self.sampler = create_sampler() self.load_model() @@ -759,6 +764,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.canary_manager = install_canary( server_args=server_args, model_runner=self, + token_oracle_manager=self._token_oracle_manager, ) # Init ngram embedding token table diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 7a12d739a..920a78d33 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -100,6 +100,8 @@ LLAMA4_MODEL_ARCHS = ( ) SAMPLING_BACKEND_CHOICES = {"flashinfer", "pytorch", "ascend"} +if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get(): + SAMPLING_BACKEND_CHOICES.add("token_oracle") LOAD_FORMAT_CHOICES = [ "auto", diff --git a/python/sglang/test/kv_canary/fixtures.py b/python/sglang/test/kv_canary/fixtures.py index de92101d7..c6d934d36 100644 --- a/python/sglang/test/kv_canary/fixtures.py +++ b/python/sglang/test/kv_canary/fixtures.py @@ -136,6 +136,7 @@ def make_base_config() -> CanaryConfig: mode=CanaryMode.RAISE, ring_capacity=1024, sweep_interval=0, + enable_write_input_assert=False, ) diff --git a/python/sglang/test/kv_canary/runner_test_base.py b/python/sglang/test/kv_canary/runner_test_base.py index 09656f1d5..d2f3f48b8 100644 --- a/python/sglang/test/kv_canary/runner_test_base.py +++ b/python/sglang/test/kv_canary/runner_test_base.py @@ -25,11 +25,13 @@ def make_config( mode: CanaryMode = CanaryMode.RAISE, ring_capacity: int = 1024, sweep_interval: int = 0, + enable_write_input_assert: bool = False, ) -> CanaryConfig: return CanaryConfig( mode=mode, ring_capacity=ring_capacity, sweep_interval=sweep_interval, + enable_write_input_assert=enable_write_input_assert, ) diff --git a/python/sglang/test/mock_model/utils.py b/python/sglang/test/mock_model/utils.py index c8b220ee4..60888249a 100644 --- a/python/sglang/test/mock_model/utils.py +++ b/python/sglang/test/mock_model/utils.py @@ -20,7 +20,7 @@ _MOCK_MODEL_SERVER_ARGS_NO_CANARY: list[str] = [ "--load-format", "dummy", "--sampling-backend", - "pytorch", + "token_oracle", "--disable-piecewise-cuda-graph", ] @@ -51,7 +51,12 @@ def mock_model_server_args(*extra_args: str, canary_mode: str = "raise") -> list def mock_model_server_env(*, input_check_enabled: bool = True) -> dict[str, str]: """Return env overrides for popen_launch_server in mock-model + canary mode.""" - return {} + return { + "SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT": ( + "1" if input_check_enabled else "0" + ), + "SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE": "1", + } def run_mock_model_bench_serving( diff --git a/test/registered/kv_canary/test_self_e2e_pr_25015.py b/test/registered/kv_canary/test_self_e2e_pr_25015.py new file mode 100644 index 000000000..91be27c95 --- /dev/null +++ b/test/registered/kv_canary/test_self_e2e_pr_25015.py @@ -0,0 +1,86 @@ +"""Regression for PR #25015 EAGLE positions misalign: revert the fix and expect canary fire.""" + +from __future__ import annotations + +import unittest +from typing import ClassVar + +from sglang.srt.kv_canary.config import CanaryMode +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.e2e_base import CanaryE2EBase + +register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") + +_SPEC_EAGLE_TOKEN_ORACLE_ENV = { + "SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT": "0", + "SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE": "1", +} +_SPEC_EAGLE_REVERT_PR_ENV = { + **_SPEC_EAGLE_TOKEN_ORACLE_ENV, + "SGLANG_DEBUG_REVERT_PR": "25015", +} +_CUDA_GRAPH_MAX_BS = 1 +_EAGER_DRAFT_REQUEST_COUNT = 20 +assert _EAGER_DRAFT_REQUEST_COUNT > _CUDA_GRAPH_MAX_BS + +_SPEC_EAGLE_SERVER_ARGS = ( + "--sampling-backend", + "token_oracle", + "--speculative-algorithm", + "EAGLE", + "--cuda-graph-max-bs", + str(_CUDA_GRAPH_MAX_BS), + "--max-running-requests", + "32", +) + + +class _EaglePositionsBase(CanaryE2EBase): + model_mode = "mha" + # LOG mode keeps the server alive after the first violation so server warmup + this test's + # parallel requests both run; we then read the violation log to assert the position bit fired. + kv_canary_mode = CanaryMode.LOG + extra_server_args = _SPEC_EAGLE_SERVER_ARGS + revert_pr: ClassVar[bool] + + @classmethod + def setUpClass(cls) -> None: + if cls is _EaglePositionsBase: + raise unittest.SkipTest("abstract base; concrete subclasses set revert_pr") + cls.extra_env = ( + _SPEC_EAGLE_REVERT_PR_ENV if cls.revert_pr else _SPEC_EAGLE_TOKEN_ORACLE_ENV + ) + super().setUpClass() + + def test_pr_25015_eagle_positions(self) -> None: + self.send_parallel_requests( + n=_EAGER_DRAFT_REQUEST_COUNT, + assert_all_success=not self.revert_pr, + max_new_tokens=32, + timeout=60.0, + ) + + if self.revert_pr: + self.assert_violation_logged_any( + launch_tag_patterns=("*",), + fail_reason="verify_position", + flush_wait_seconds=0.0, + ) + else: + self.assert_no_violation(wait_seconds=2.0) + + +class TestEaglePositionsMisalignRegression(_EaglePositionsBase): + """Revert PR #25015 fix and expect canary to fire a position-mismatch violation.""" + + revert_pr = True + + +class TestEaglePositionsMatchWithFix(_EaglePositionsBase): + """With the PR #25015 fix in place, no canary fires.""" + + revert_pr = False + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kv_canary/test_self_unit_token_oracle.py b/test/registered/kv_canary/test_self_unit_token_oracle.py new file mode 100644 index 000000000..f8866473e --- /dev/null +++ b/test/registered/kv_canary/test_self_unit_token_oracle.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.kv_canary.expected_inputs import ExpectedInputs +from sglang.srt.kv_canary.token_oracle.oracle import HashOracle +from sglang.srt.kv_canary.token_oracle.oracle_manager import TokenOracleManager +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.kv_canary.fixtures import DEFAULT_DEVICE +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=1, stage="extra-a", runner_config="1-gpu-small") + + +class TestTokenOracleManager(CustomTestCase): + def setUp(self) -> None: + self.device = DEFAULT_DEVICE + + def test_fill_expected_inputs_expands_draft_extend_generalized_req_ids_per_token( + self, + ) -> None: + """Verify EAGLE draft extend maps one request row to every draft token.""" + forward_batch = SimpleNamespace( + forward_mode=ForwardMode.DRAFT_EXTEND, + spec_info=SimpleNamespace(num_tokens_per_req=4), + rids_int=torch.tensor([3, 7], dtype=torch.int64, device=self.device), + bootstrap_room_ids_int=None, + input_ids=torch.tensor( + [101, 102, 103, 104, 201, 202, 203, 204], + dtype=torch.int64, + device=self.device, + ), + positions=torch.arange(8, dtype=torch.int64, device=self.device), + extend_seq_lens=torch.tensor([1, 1], dtype=torch.int64, device=self.device), + ) + expected_inputs = ExpectedInputs.allocate(capacity=8, device=self.device) + manager = TokenOracleManager(oracle=HashOracle(vocab_size=32000)) + + manager.fill_expected_inputs( + forward_batch=forward_batch, + expected_inputs_out=expected_inputs, + ) + + self.assertTrue( + torch.equal(expected_inputs.tokens[:8], forward_batch.input_ids) + ) + self.assertTrue( + torch.equal(expected_inputs.positions[:8], forward_batch.positions) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/mock_model/test_self_e2e_perturb_next_token_swap.py b/test/registered/mock_model/test_self_e2e_perturb_next_token_swap.py new file mode 100644 index 000000000..40231926c --- /dev/null +++ b/test/registered/mock_model/test_self_e2e_perturb_next_token_swap.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +import unittest + +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.mock_model.perturb_e2e_base import MockModelPerturbE2EBase + +register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") + + +class TestPerturbNextTokenSwap(MockModelPerturbE2EBase): + """Mock-model self-test: swap two reqs' sampled next tokens at the sampler exit. + + KV path is untouched, so kv_canary KV-side fail_reasons stay silent. The + token-oracle input check downstream MUST report fail_reason=write_token. + Validates the input-check link is genuinely active. + """ + + extra_env = { + "SGLANG_KV_CANARY_PERTURB_NEXT_TOKEN_SWAP_PROB": "0.1", + "SGLANG_KV_CANARY_PERTURB_WARMUP_STEPS": "0", + } + extra_server_args = ("--skip-server-warmup",) + + def test_swap_triggers_input_check_violation_but_kv_paths_silent(self) -> None: + """Verify next_token swap fires write_token violation while KV reasons stay silent.""" + self.send_parallel_requests(n=4, timeout=30.0) + self.assert_log_contains("kv_canary perturb next_token_swap: swapped") + self.assert_any_launch_tag_violation_reported(fail_reason="write_token") + self.assert_any_launch_tag_violation_absent(fail_reason="verify_real_kv_hash") + self.assert_any_launch_tag_violation_absent(fail_reason="verify_position") + self.assert_any_launch_tag_violation_absent(fail_reason="verify_chain_hash") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/mock_model/test_self_unit_canary_mock_wiring.py b/test/registered/mock_model/test_self_unit_canary_mock_wiring.py new file mode 100644 index 000000000..7916867d7 --- /dev/null +++ b/test/registered/mock_model/test_self_unit_canary_mock_wiring.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +import dataclasses +import unittest + +import torch + +from sglang.srt.kv_canary.expected_inputs import ExpectedInputs +from sglang.srt.kv_canary.token_oracle.oracle import HashOracle +from sglang.srt.kv_canary.token_oracle.sampler import install_oracle_sampler +from sglang.srt.model_executor.forward_batch_info import ( + ForwardMode, + _stable_hash_str_to_i64, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.mock_model.utils import mock_model_server_args, mock_model_server_env +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") + + +@dataclasses.dataclass +class _StubForwardBatch: + input_ids: torch.Tensor + positions: torch.Tensor + req_pool_indices: torch.Tensor + forward_mode: ForwardMode + extend_seq_lens: object + rids_int: torch.Tensor + bootstrap_room_ids_int: torch.Tensor | None = None + spec_info: object | None = None + seq_lens: torch.Tensor | None = None + + +def _scalar_expected_token( + oracle: HashOracle, *, generalized_req_id: int, position: int +) -> int: + out = oracle.expected_tokens( + generalized_req_ids=torch.tensor([generalized_req_id], dtype=torch.int64), + positions=torch.tensor([position], dtype=torch.int64), + ) + return int(out.tolist()[0]) + + +class TestFillExpectedInputs(CustomTestCase): + def test_sample_next_tokens_uses_next_position(self) -> None: + oracle = HashOracle(vocab_size=32000) + hook = install_oracle_sampler(oracle=oracle) + + rid_a = "req-a" + hashed_a = _stable_hash_str_to_i64(rid_a) + out = hook.sample_next_tokens( + generalized_req_ids=torch.tensor([hashed_a], dtype=torch.int64), + logits_positions=torch.tensor([5], dtype=torch.int64), + ) + + self.assertEqual( + out.tolist(), + [_scalar_expected_token(oracle, generalized_req_id=hashed_a, position=6)], + ) + + def test_fill_expected_inputs_decode_one_token_per_req(self) -> None: + """Verify decode mode fills one expected token per request.""" + oracle = HashOracle(vocab_size=32000) + hook = install_oracle_sampler(oracle=oracle) + + rid_a = "req-a" + rid_b = "req-b" + fb = _StubForwardBatch( + input_ids=torch.tensor([0, 0], dtype=torch.int64), + positions=torch.tensor([10, 20], dtype=torch.int64), + req_pool_indices=torch.tensor([5, 7], dtype=torch.int64), + forward_mode=ForwardMode.DECODE, + extend_seq_lens=None, + rids_int=torch.tensor( + [_stable_hash_str_to_i64(rid_a), _stable_hash_str_to_i64(rid_b)], + dtype=torch.int64, + ), + ) + expected_inputs = ExpectedInputs.allocate( + capacity=8, device=torch.device("cpu") + ) + + hook.fill_expected_inputs( + forward_batch=fb, + expected_inputs_out=expected_inputs, + ) + + self.assertEqual( + expected_inputs.tokens[:2].tolist(), + [ + _scalar_expected_token( + oracle, + generalized_req_id=_stable_hash_str_to_i64(rid_a), + position=10, + ), + _scalar_expected_token( + oracle, + generalized_req_id=_stable_hash_str_to_i64(rid_b), + position=20, + ), + ], + ) + self.assertEqual(expected_inputs.positions[:2].tolist(), [10, 20]) + + def test_fill_expected_inputs_prefers_bootstrap_room_ids(self) -> None: + """Verify PD oracle checks can key by bootstrap room without rewriting rids_int.""" + oracle = HashOracle(vocab_size=32000) + hook = install_oracle_sampler(oracle=oracle) + + rid_a = "prefill-local-rid" + rid_b = "regular-rid" + hashed_a = _stable_hash_str_to_i64(rid_a) + hashed_b = _stable_hash_str_to_i64(rid_b) + fb = _StubForwardBatch( + input_ids=torch.tensor([0, 0], dtype=torch.int64), + positions=torch.tensor([10, 20], dtype=torch.int64), + req_pool_indices=torch.tensor([5, 7], dtype=torch.int64), + forward_mode=ForwardMode.DECODE, + extend_seq_lens=None, + rids_int=torch.tensor([hashed_a, hashed_b], dtype=torch.int64), + bootstrap_room_ids_int=torch.tensor([1234, -1], dtype=torch.int64), + ) + expected_inputs = ExpectedInputs.allocate( + capacity=8, device=torch.device("cpu") + ) + + hook.fill_expected_inputs( + forward_batch=fb, + expected_inputs_out=expected_inputs, + ) + + self.assertEqual(fb.rids_int.tolist(), [hashed_a, hashed_b]) + self.assertEqual( + expected_inputs.tokens[:2].tolist(), + [ + _scalar_expected_token(oracle, generalized_req_id=1234, position=10), + _scalar_expected_token( + oracle, generalized_req_id=hashed_b, position=20 + ), + ], + ) + self.assertEqual(expected_inputs.positions[:2].tolist(), [10, 20]) + + def test_fill_expected_inputs_extend_uses_forward_input_ids(self) -> None: + """Verify extend mode checks prompt tokens already present in the forward batch.""" + oracle = HashOracle(vocab_size=32000) + hook = install_oracle_sampler(oracle=oracle) + + rid_a = "req-a" + rid_b = "req-b" + hashed_a = _stable_hash_str_to_i64(rid_a) + hashed_b = _stable_hash_str_to_i64(rid_b) + fb = _StubForwardBatch( + input_ids=torch.tensor([101, 102, 103, 201], dtype=torch.int64), + positions=torch.tensor([0, 1, 2, 0], dtype=torch.int64), + req_pool_indices=torch.tensor([5, 7], dtype=torch.int64), + forward_mode=ForwardMode.EXTEND, + extend_seq_lens=torch.tensor([3, 1], dtype=torch.int64), + rids_int=torch.tensor([hashed_a, hashed_b], dtype=torch.int64), + ) + expected_inputs = ExpectedInputs.allocate( + capacity=8, device=torch.device("cpu") + ) + + hook.fill_expected_inputs( + forward_batch=fb, + expected_inputs_out=expected_inputs, + ) + + self.assertEqual( + expected_inputs.tokens[:4].tolist(), + [101, 102, 103, 201], + ) + self.assertEqual(expected_inputs.positions[:4].tolist(), [0, 1, 2, 0]) + + def test_fill_expected_inputs_zero_tokens_is_noop( + self, + ) -> None: + """Verify filling zero expected tokens leaves the output buffer unchanged.""" + hook = install_oracle_sampler(oracle=HashOracle(vocab_size=100)) + + rid_a = "req-a" + rid_b = "req-b" + fb = _StubForwardBatch( + input_ids=torch.empty(0, dtype=torch.int64), + positions=torch.empty(0, dtype=torch.int64), + req_pool_indices=torch.tensor([5, 7], dtype=torch.int64), + forward_mode=ForwardMode.DECODE, + extend_seq_lens=None, + rids_int=torch.tensor( + [_stable_hash_str_to_i64(rid_a), _stable_hash_str_to_i64(rid_b)], + dtype=torch.int64, + ), + ) + expected_inputs = ExpectedInputs.allocate( + capacity=4, device=torch.device("cpu") + ) + initial_tokens = expected_inputs.tokens.clone() + + hook.fill_expected_inputs( + forward_batch=fb, + expected_inputs_out=expected_inputs, + ) + + self.assertEqual(expected_inputs.tokens.tolist(), initial_tokens.tolist()) + + +class TestMockModelServerLaunchHelpers(CustomTestCase): + def test_mock_model_server_args_adds_canary_defaults(self) -> None: + """Verify mock model launch args include KV canary defaults before user args.""" + args = mock_model_server_args("--tp", "2") + + self.assertIn("--load-format", args) + self.assertIn("dummy", args) + self.assertIn("--sampling-backend", args) + self.assertIn("token_oracle", args) + self.assertIn("--kv-canary", args) + self.assertIn("raise", args) + self.assertEqual(args[-2:], ["--tp", "2"]) + + def test_mock_model_server_env_enables_input_check_by_default(self) -> None: + """Verify mock model launch env enables canary input checking by default.""" + env = mock_model_server_env() + + self.assertEqual(env["SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT"], "1") + self.assertEqual(env["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"], "1") + + def test_mock_model_server_env_can_disable_input_check(self) -> None: + """Verify mock model launch env can disable canary input checking.""" + env = mock_model_server_env(input_check_enabled=False) + + self.assertEqual(env["SGLANG_KV_CANARY_ENABLE_WRITE_INPUT_ASSERT"], "0") + self.assertEqual(env["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"], "1") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/mock_model/test_self_unit_install.py b/test/registered/mock_model/test_self_unit_install.py new file mode 100644 index 000000000..28f655b52 --- /dev/null +++ b/test/registered/mock_model/test_self_unit_install.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import os +import unittest +from types import SimpleNamespace + +os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = "1" + +from sglang.srt.kv_canary.token_oracle.install import install_token_oracle_from_env +from sglang.srt.kv_canary.token_oracle.oracle import HashOracle +from sglang.srt.layers.sampler import _CUSTOM_SAMPLER_FACTORIES +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") + + +def _make_server_args(*, sampling_backend: str) -> SimpleNamespace: + return SimpleNamespace(sampling_backend=sampling_backend) + + +class TestInstallTokenOracleFromEnv(CustomTestCase): + def test_install_token_oracle_from_env_disabled_returns_none(self) -> None: + """Verify server-arg-disabled token oracle installation (sampling_backend != 'token_oracle') returns no TokenOracleManager.""" + server_args = _make_server_args(sampling_backend="auto") + hook = install_token_oracle_from_env(server_args=server_args, vocab_size=1000) + self.assertIsNone(hook) + + def test_install_token_oracle_from_env_enabled_registers_oracle_backend( + self, + ) -> None: + """Verify token oracle installation via sampling_backend='token_oracle' registers the oracle backend.""" + server_args = _make_server_args(sampling_backend="token_oracle") + hook = install_token_oracle_from_env(server_args=server_args, vocab_size=512) + self.assertIsNotNone(hook) + self.assertIn("token_oracle", _CUSTOM_SAMPLER_FACTORIES) + + def test_install_token_oracle_from_env_enabled_returns_hook_with_hash_oracle( + self, + ) -> None: + """Verify token oracle installation via sampling_backend='token_oracle' returns a TokenOracleManager wrapping a HashOracle.""" + server_args = _make_server_args(sampling_backend="token_oracle") + hook = install_token_oracle_from_env(server_args=server_args, vocab_size=256) + self.assertIsNotNone(hook) + self.assertIsInstance(hook.oracle, HashOracle) + self.assertEqual(hook.oracle.vocab_size, 256) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/mock_model/test_self_unit_oracle.py b/test/registered/mock_model/test_self_unit_oracle.py new file mode 100644 index 000000000..fe7f02e6c --- /dev/null +++ b/test/registered/mock_model/test_self_unit_oracle.py @@ -0,0 +1,127 @@ +from __future__ import annotations + +import random +import unittest + +import torch + +from sglang.jit_kernel.kv_canary.consts import splitmix64 +from sglang.srt.kv_canary.token_oracle.oracle import ( + HashOracle, + _splitmix64_tensor, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") + + +_U64_MASK: int = (1 << 64) - 1 + + +def _signed_to_unsigned_i64(value: int) -> int: + return value & _U64_MASK + + +def _call(oracle: HashOracle, *, generalized_req_id: int, position: int) -> int: + out = oracle.expected_tokens( + generalized_req_ids=torch.tensor([generalized_req_id], dtype=torch.int64), + positions=torch.tensor([position], dtype=torch.int64), + ) + return int(out.tolist()[0]) + + +class TestHashOracle(CustomTestCase): + def test_hash_oracle_is_deterministic_for_same_inputs(self) -> None: + """Verify HashOracle returns the same token for identical inputs.""" + oracle = HashOracle(vocab_size=32000) + + first = _call(oracle, generalized_req_id=7, position=42) + second = _call(oracle, generalized_req_id=7, position=42) + + self.assertEqual(first, second) + + def test_hash_oracle_output_in_vocab_range(self) -> None: + """Verify HashOracle outputs stay within the configured vocabulary range.""" + vocab_size = 1024 + oracle = HashOracle(vocab_size=vocab_size) + + generalized_req_ids = torch.arange(0, 64, dtype=torch.int64).repeat_interleave( + 64 + ) + positions = torch.arange(0, 64, dtype=torch.int64).repeat(64) + tokens = oracle.expected_tokens( + generalized_req_ids=generalized_req_ids, positions=positions + ).tolist() + + for token in tokens: + self.assertTrue(0 <= token < vocab_size) + + +class TestSplitmix64Tensor(CustomTestCase): + def test_splitmix64_tensor_matches_scalar_ref_on_random_inputs(self) -> None: + """Verify tensor SplitMix64 matches the scalar reference on random inputs.""" + rng = random.Random(0) + num_cases = 1000 + unsigned_inputs: list[int] = [ + rng.randrange(0, 1 << 64) for _ in range(num_cases) + ] + + signed_inputs = [ + value if value < (1 << 63) else value - (1 << 64) + for value in unsigned_inputs + ] + actual = _splitmix64_tensor(torch.tensor(signed_inputs, dtype=torch.int64)) + + actual_unsigned = [_signed_to_unsigned_i64(v) for v in actual.tolist()] + expected_unsigned = [splitmix64(v) for v in unsigned_inputs] + + self.assertEqual(actual_unsigned, expected_unsigned) + + def test_splitmix64_tensor_known_vectors(self) -> None: + """Verify tensor SplitMix64 matches scalar reference values for known inputs.""" + inputs = [0, 1, -1, 1 << 32, (1 << 63) - 1, -(1 << 63)] + expected_unsigned = [splitmix64(_signed_to_unsigned_i64(v)) for v in inputs] + + actual = _splitmix64_tensor(torch.tensor(inputs, dtype=torch.int64)) + actual_unsigned = [_signed_to_unsigned_i64(v) for v in actual.tolist()] + + self.assertEqual(actual_unsigned, expected_unsigned) + + def test_splitmix64_tensor_preserves_shape_and_dtype(self) -> None: + """Verify tensor SplitMix64 preserves input shape and int64 dtype.""" + shape = (3, 4, 5) + rng = torch.Generator().manual_seed(42) + inputs = torch.randint( + low=-(1 << 62), + high=(1 << 62), + size=shape, + dtype=torch.int64, + generator=rng, + ) + + out = _splitmix64_tensor(inputs) + + self.assertEqual(out.shape, inputs.shape) + self.assertEqual(out.dtype, torch.int64) + + def test_splitmix64_tensor_is_deterministic(self) -> None: + """Verify tensor SplitMix64 returns stable values for repeated calls.""" + inputs = torch.tensor([0, 1, 2, 3, 1 << 40, -7], dtype=torch.int64) + + first = _splitmix64_tensor(inputs.clone()).tolist() + second = _splitmix64_tensor(inputs.clone()).tolist() + + self.assertEqual(first, second) + + def test_splitmix64_tensor_is_injective_on_distinct_inputs(self) -> None: + """Verify tensor SplitMix64 maps distinct sampled inputs to distinct outputs.""" + inputs = torch.arange(-1000, 1000, dtype=torch.int64) + + out = _splitmix64_tensor(inputs).tolist() + + self.assertEqual(len(set(out)), len(out)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/mock_model/test_self_unit_oracle_torch_vs_ref.py b/test/registered/mock_model/test_self_unit_oracle_torch_vs_ref.py new file mode 100644 index 000000000..3043e66ec --- /dev/null +++ b/test/registered/mock_model/test_self_unit_oracle_torch_vs_ref.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +import random +import unittest + +import torch + +from sglang.jit_kernel.kv_canary.verify_ref import splitmix64 +from sglang.srt.kv_canary.token_oracle.oracle import HashOracle +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-small") + + +class TestHashOracleTorchVsRef(CustomTestCase): + def test_hash_oracle_matches_scalar_splitmix64_ref(self) -> None: + """Verify single-item HashOracle calls match the scalar SplitMix64 reference.""" + rng = random.Random(0) + vocab_size = 32000 + num_cases = 1000 + + generalized_req_ids: list[int] = [] + positions: list[int] = [] + for _ in range(num_cases): + generalized_req_ids.append(rng.randrange(0, 1 << 60)) + positions.append(rng.randrange(0, 1 << 60)) + + ref_tokens: list[int] = [ + splitmix64(generalized_req_ids[i] ^ positions[i]) % vocab_size + for i in range(num_cases) + ] + + oracle = HashOracle(vocab_size=vocab_size) + torch_tokens: list[int] = [] + generalized_req_ids_tensor = torch.tensor( + generalized_req_ids, dtype=torch.int64 + ) + positions_tensor = torch.tensor(positions, dtype=torch.int64) + for i in range(num_cases): + out = oracle.expected_tokens( + generalized_req_ids=generalized_req_ids_tensor[i : i + 1], + positions=positions_tensor[i : i + 1], + ) + torch_tokens.append(int(out.tolist()[0])) + + for i in range(num_cases): + self.assertEqual( + torch_tokens[i], + ref_tokens[i], + f"mismatch at case {i}: generalized_req_id={generalized_req_ids[i]} " + f"position={positions[i]}: torch={torch_tokens[i]} ref={ref_tokens[i]}", + ) + + def test_hash_oracle_batched_matches_scalar_splitmix64_ref(self) -> None: + """Verify batched HashOracle calls match the scalar SplitMix64 reference.""" + rng = random.Random(1) + vocab_size = 32000 + num_cases = 1000 + + generalized_req_ids = [rng.randrange(0, 1 << 60) for _ in range(num_cases)] + positions = [rng.randrange(0, 1 << 60) for _ in range(num_cases)] + ref_tokens = [ + splitmix64(generalized_req_ids[i] ^ positions[i]) % vocab_size + for i in range(num_cases) + ] + + oracle = HashOracle(vocab_size=vocab_size) + out = oracle.expected_tokens( + generalized_req_ids=torch.tensor(generalized_req_ids, dtype=torch.int64), + positions=torch.tensor(positions, dtype=torch.int64), + ) + torch_tokens = out.tolist() + + for i in range(num_cases): + self.assertEqual( + torch_tokens[i], + ref_tokens[i], + f"batched mismatch at case {i}: generalized_req_id={generalized_req_ids[i]} " + f"position={positions[i]}: torch={torch_tokens[i]} ref={ref_tokens[i]}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/mock_model/test_self_unit_sampler_hookpoint.py b/test/registered/mock_model/test_self_unit_sampler_hookpoint.py new file mode 100644 index 000000000..2d4e2c44a --- /dev/null +++ b/test/registered/mock_model/test_self_unit_sampler_hookpoint.py @@ -0,0 +1,49 @@ +"""install_oracle_sampler registration into sglang's sampler-backend registry. + +Instantiating the registered _OracleSampler factory requires a live distributed (TP) group +plus a populated global ServerArgs, so the forward-path behavior of _OracleSampler is covered +by the e2e harness rather than this unit file. Here we only assert the registration-side +contract: the backend name shows up in the registry / choice set, and second install replaces +the factory with one bound to the new oracle. +""" + +from __future__ import annotations + +import os +import unittest + +os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = "1" + +from sglang.srt.kv_canary.token_oracle.oracle import HashOracle +from sglang.srt.kv_canary.token_oracle.sampler import install_oracle_sampler +from sglang.srt.layers.sampler import _CUSTOM_SAMPLER_FACTORIES +from sglang.srt.server_args import SAMPLING_BACKEND_CHOICES +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=60, stage="extra-a", runner_config="1-gpu-small") + + +class TestInstallOracleSampler(CustomTestCase): + def test_install_oracle_sampler_twice_returns_distinct_hooks_with_replaced_oracle( + self, + ) -> None: + """Verify reinstalling the oracle sampler replaces the registered factory.""" + oracle_a = HashOracle(vocab_size=100) + oracle_b = HashOracle(vocab_size=100) + + hook_a = install_oracle_sampler(oracle=oracle_a) + self.assertIn("token_oracle", _CUSTOM_SAMPLER_FACTORIES) + self.assertIn("token_oracle", SAMPLING_BACKEND_CHOICES) + factory_a = _CUSTOM_SAMPLER_FACTORIES["token_oracle"] + self.assertIs(hook_a.oracle, oracle_a) + + hook_b = install_oracle_sampler(oracle=oracle_b) + factory_b = _CUSTOM_SAMPLER_FACTORIES["token_oracle"] + self.assertIs(hook_b.oracle, oracle_b) + self.assertIsNot(hook_a, hook_b) + self.assertIsNot(factory_a, factory_b) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index b04faa8fb..6464d65e9 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -1,8 +1,11 @@ +import importlib import json +import os import tempfile import unittest from unittest.mock import MagicMock, patch +import sglang.srt.server_args as server_args_module from sglang.srt.arg_groups.speculative_hook import handle_speculative_decoding from sglang.srt.server_args import PortArgs, ServerArgs, prepare_server_args from sglang.test.ci.ci_register import register_cpu_ci @@ -659,5 +662,58 @@ class TestCutedslMoeMaxNumTokens(unittest.TestCase): self.assertEqual(args.cutedsl_moe_max_num_tokens(), 512) +class TestSamplingBackendTokenOracleEnvGate(CustomTestCase): + """The 'token_oracle' choice is gated on SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE. + + The choice set is built once at server_args.py import time, so each subtest + reloads the module with the env var set to the desired value. + """ + + def _reload_server_args_with_env(self, *, enabled: bool): + previous = os.environ.get("SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE") + os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = "1" if enabled else "0" + try: + return importlib.reload(server_args_module) + finally: + if previous is None: + os.environ.pop("SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE", None) + else: + os.environ["SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE"] = previous + + def test_token_oracle_rejected_when_env_disabled(self): + reloaded = self._reload_server_args_with_env(enabled=False) + self.assertNotIn("token_oracle", reloaded.SAMPLING_BACKEND_CHOICES) + + with self.assertRaises(SystemExit): + reloaded.prepare_server_args( + [ + "--model-path", + DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, + "--sampling-backend", + "token_oracle", + ] + ) + + def test_token_oracle_accepted_when_env_enabled(self): + reloaded = self._reload_server_args_with_env(enabled=True) + self.assertIn("token_oracle", reloaded.SAMPLING_BACKEND_CHOICES) + + parsed = reloaded.prepare_server_args( + [ + "--model-path", + DEFAULT_SMALL_MODEL_NAME_FOR_TEST_QWEN, + "--sampling-backend", + "token_oracle", + # Explicit device so ServerArgs.__post_init__ does not call + # get_device() (fails on CPU-only CI runners) and does not run + # _handle_cpu_backends (which would override sampling_backend + # to "pytorch", masking what we want to verify). + "--device", + "cuda", + ] + ) + self.assertEqual(parsed.sampling_backend, "token_oracle") + + if __name__ == "__main__": unittest.main()