Add a deterministic token oracle and production write-input assertion (#26815)

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)
 

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)

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

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

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

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_write_input_assert: bool,
-    enable_verify_token_assert: bool = False,
+    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)

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 (

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)

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

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

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

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

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

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

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

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(

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

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

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

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

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

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

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

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

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