Fix deterministic inference for Inkling (#33417)

python/sglang/kernels/ops/attention/flash_attn/cute/batch_invariance.py

@@ -0,0 +1,16 @@
+"""Process-wide so it need not thread through the autograd entry points. Set
+before the first kernel compile; it keys the forward compile cache.
+"""
+
+from __future__ import annotations
+
+_batch_invariant = False
+
+
+def set_batch_invariant(enabled: bool) -> None:
+    global _batch_invariant
+    _batch_invariant = bool(enabled)
+
+
+def is_batch_invariant() -> bool:
+    return _batch_invariant

python/sglang/kernels/ops/attention/flash_attn/cute/block_info.py

@@ -22,6 +22,7 @@ class BlockInfo:
     window_size_left: Optional[Int32] = None
     window_size_right: Optional[Int32] = None
     qhead_per_kvhead_packgqa: cutlass.Constexpr[int] = 1
+    batch_invariant: cutlass.Constexpr[bool] = False
 
     @cute.jit
     def get_n_idx_left_right(
@@ -159,6 +160,11 @@ class BlockInfo:
         n_block_min: Int32,
     ) -> Int32:
         """If we have separate iterations with causal or local masking at the start, where do we stop"""
+        # The boundary carries -seqlen_q and the two paths round differently, so a
+        # row's bits depend on how many queries share the pass. Masking every
+        # block is identity below the diagonal.
+        if const_expr(self.batch_invariant):
+            return n_block_min
         m_idx_min = m_block * self.tile_m
         if const_expr(self.qhead_per_kvhead_packgqa > 1):
             m_idx_min = m_idx_min // self.qhead_per_kvhead_packgqa

python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_sm100.py

@@ -242,6 +242,7 @@ class FlashAttentionForwardSm100:
         v_dequant: bool = False,
         q_sf_interleaved: bool = False,
         kv_sf_interleaved: bool = False,
+        batch_invariant: bool = False,
     ):
         # MXFP8 block-scaled attention (see interface._flash_attn_fwd):
         #   qk_blockscaled: Q/K fp8 e4m3 + per-32 UE8M0 scales; QK^T runs as
@@ -331,6 +332,7 @@ class FlashAttentionForwardSm100:
         self.pack_gqa = pack_gqa
         # relative (sheared) bias
         self.has_bias = has_bias
+        self.batch_invariant = batch_invariant
         self.rel_extent_padded = rel_extent_padded
         assert rel_extent_padded % n_block_size == 0
         self.bias_n_max = rel_extent_padded // n_block_size if has_bias else 0
@@ -2069,6 +2071,7 @@ class FlashAttentionForwardSm100:
             qhead_per_kvhead_packgqa=(
                 self.qhead_per_kvhead if const_expr(self.pack_gqa) else 1
             ),
+            batch_invariant=self.batch_invariant,
         )
         SeqlenInfoCls = partial(
             SeqlenInfoQK.create,

python/sglang/kernels/ops/attention/flash_attn/cute/interface.py

@@ -14,6 +14,9 @@ from cutlass import Float32, Int32
 from quack.compile_utils import make_fake_tensor as fake_tensor
 
 from sglang.kernels.jit.utils import is_arch_support_pdl
+from sglang.kernels.ops.attention.flash_attn.cute.batch_invariance import (
+    is_batch_invariant,
+)
 from sglang.kernels.ops.attention.flash_attn.cute.cache_utils import get_jit_cache
 from sglang.kernels.ops.attention.flash_attn.cute.testing import is_fake_mode
 
@@ -1071,6 +1074,7 @@ def _flash_attn_fwd(
             # Benchmark hook: time just the shear-prep kernels, skip the attention kernel.
             return out, lse
 
+    batch_invariant = is_batch_invariant()
     compile_key = (
         dtype,
         head_dim,
@@ -1130,6 +1134,7 @@ def _flash_attn_fwd(
         sfq.ndim if sfq is not None else None,
         sfk.ndim if sfk is not None else None,
         sfv.ndim if sfv is not None else None,
+        batch_invariant,
         fa_logging.get_fa_log_level(),
     )
 
@@ -1353,6 +1358,7 @@ def _flash_attn_fwd(
                             v_dequant=v_blockscaled,
                             q_sf_interleaved=q_sf_interleaved,
                             kv_sf_interleaved=kv_sf_interleaved,
+                            batch_invariant=batch_invariant,
                         )
                     ),
                 )

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -268,6 +268,13 @@ class FlashAttentionBackend(AttentionBackend):
             )
 
             self._get_scheduler_metadata = None
+            if model_runner.server_args.enable_deterministic_inference:
+                # Must precede the first kernel compile.
+                from sglang.kernels.ops.attention.flash_attn.cute.batch_invariance import (
+                    set_batch_invariant,
+                )
+
+                set_batch_invariant(True)
         else:
             raise ValueError(f"Invalid version: {self.fa_impl_ver=}")
 

python/sglang/srt/model_executor/runner/flashinfer_autotune.py

@@ -52,6 +52,10 @@ def should_run_flashinfer_autotune(
         return False
     if mr.server_args.disable_flashinfer_autotune:
         return False
+    if mr.server_args.enable_deterministic_inference:
+        # Tuned configs are per problem shape, so the reduction order would follow
+        # the batch shape.
+        return False
 
     server_args = mr.server_args
     if for_speculative_draft:

python/sglang/srt/models/inkling_common/attn.py

@@ -121,7 +121,7 @@ def _rel_proj_kernel_eligible(r: torch.Tensor) -> bool:
 
 
 class RelLogitsProj(nn.Module):
-    def __init__(self, d_rel: int, rel_extent: int):
+    def __init__(self, d_rel: int, rel_extent: int, *, deterministic: bool = False):
         super().__init__()
         self.d_rel = d_rel
         self.rel_extent = rel_extent
@@ -134,7 +134,11 @@ class RelLogitsProj(nn.Module):
         # territory. Rounding moves with the fold (r*tau rounds to bf16 before
         # the GEMM instead of after); flag-off keeps the exact legacy post-scale.
         self._prescale_tau = envs.SGLANG_OPT_USE_INKLING_FUSED_LOG_TAU.get()
-        self._proj_dispatch = envs.SGLANG_OPT_USE_INKLING_REL_PROJ_DISPATCH.get()
+        # The dispatch keys off the token count, so a row's kernel would follow
+        # the batch composition.
+        self._proj_dispatch = (
+            envs.SGLANG_OPT_USE_INKLING_REL_PROJ_DISPATCH.get() and not deterministic
+        )
 
     def _project(self, r: torch.Tensor) -> torch.Tensor:
         """``einsum("thd,de->the", r, proj)`` -- but dispatched: in production
@@ -309,7 +313,11 @@ class InklingAttention(nn.Module):
             self.rel_extent = rel_extent
             self.local_extent = None
 
-        self.rel_logits_proj = RelLogitsProj(self.d_rel, self.rel_extent)
+        self.rel_logits_proj = RelLogitsProj(
+            self.d_rel,
+            self.rel_extent,
+            deterministic=get_exec().deterministic.enable_deterministic_inference,
+        )
         # Fold the conditional log-scaling tau into the fused prologue's q
         # path (deletes the external scale kernel; bit-exact rounding).
         self._fused_log_tau = envs.SGLANG_OPT_USE_INKLING_FUSED_LOG_TAU.get()

python/sglang/srt/models/inkling_common/kernels/comm.py

@@ -455,9 +455,12 @@ def symm_mem_all_reduce(
     ):
         n = input.numel()
         num_tokens = input.shape[0] if input.dim() >= 2 else n
+        # select_ar_config() keys off the token count; plain multimem below
+        # reduces in a shape-independent order.
         res = (
             _get_inkling_ar_resources(comm)
             if envs.SGLANG_OPT_USE_INKLING_CUSTOM_AR.get()
+            and not get_exec().deterministic.enable_deterministic_inference
             else None
         )
         # Custom kernels need a 16B-vector-multiple size (validate() enforces it);

python/sglang/srt/server_args.py

@@ -7911,8 +7911,12 @@ class ServerArgs:
                     # CUDA: use NCCL tree algorithm
                     os.environ["NCCL_ALGO"] = "allreduce:tree"
                     self.disable_custom_all_reduce = True
+                    # should_torch_symm_mem_allreduce() takes the
+                    # symmetric-memory path only below a byte threshold, so
+                    # which reduce runs would follow the token count.
+                    self.enable_torch_symm_mem = False
                     logger.warning(
-                        "NCCL_ALGO is set to 'allreduce:tree' and custom all reduce is disabled for deterministic inference when TP size > 1."
+                        "NCCL_ALGO is set to 'allreduce:tree', and custom and symmetric-memory all reduce are disabled for deterministic inference when TP size > 1."
                     )
 
     def _handle_unified_memory_pool(self):

python/sglang/test/cache_consistency_jitter.py

@@ -0,0 +1,490 @@
+"""Cache/scheduling consistency test: score the same token contexts repeatedly
+and require bitwise-identical greedy output tokens and top-k logprobs.
+
+The workload generates a continuation for each independently sampled prefix,
+then submits shuffled requests starting from different cuts of that
+continuation. This scores each shared context multiple times after different
+batching and cache histories and compares the resulting tokens and logprobs, so
+a hybrid-cache (SWA + Mamba/sconv) or overlap-scheduler bug that corrupts a
+reused context shows up as a diverging token or logprob rather than as a small
+accuracy drop no eval would catch.
+
+The engine is put under additional pressure in three ways: probabilistic timing
+jitter is enqueued after CUDA event record and wait operations to amplify stream
+races, ``max_total_tokens`` constrains the KV cache, and periodic request
+retractions introduce further scheduling and cache disruption.
+
+This assumes a batch-invariant serving path: reducing the same context under a
+different batch shape must produce the same bits, or the differences reported are
+in reduction order rather than in cache and scheduling behaviour.
+"""
+
+from __future__ import annotations
+
+import gc
+import os
+import random
+import time
+from collections import defaultdict
+from collections.abc import Iterator
+from contextlib import contextmanager
+from functools import partial
+from typing import TYPE_CHECKING, Any, NamedTuple
+
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import globaltimer, smid
+
+if TYPE_CHECKING:
+    from sglang.srt.entrypoints.engine import Engine
+
+
+@triton.jit
+def _random_jitter_kernel(
+    enabled: tl.tensor,
+    PROBABILITY: tl.constexpr,
+    MAX_TIME_US: tl.constexpr,
+) -> None:
+    if tl.load(enabled) != 0:
+        start = globaltimer().to(tl.uint64)
+        sample = start ^ (smid().to(tl.uint64) << 32)
+        sample ^= sample << 13
+        sample ^= sample >> 7
+        sample ^= sample << 17
+        if sample.to(tl.uint32) < int(PROBABILITY * 2**32):
+            sample ^= sample << 13
+            sample ^= sample >> 7
+            sample ^= sample << 17
+            duration_ns = (
+                sample.to(tl.uint32).to(tl.uint64) * (MAX_TIME_US * 1_000 + 1)
+            ) >> 32
+            now = globaltimer().to(tl.uint64)
+            while now - start < duration_ns:
+                now = globaltimer().to(tl.uint64)
+
+
+def _random_jitter(enabled: torch.Tensor, probability: float, max_time_us: int) -> None:
+    kernel: Any = _random_jitter_kernel
+    kernel[(1,)](
+        enabled,
+        PROBABILITY=probability,
+        MAX_TIME_US=max_time_us,
+        num_warps=1,
+    )
+
+
+def _run_scheduler_process_with_jitter(
+    server_args: Any,
+    port_args: Any,
+    gpu_id: int,
+    *scheduler_args: Any,
+    jitter_probability: float,
+    jitter_max_time_us: int,
+    **scheduler_kwargs: Any,
+) -> None:
+    from sglang.srt.managers.scheduler import Scheduler, run_scheduler_process
+
+    original_record = torch.cuda.Event.record
+    original_wait = torch.cuda.Event.wait
+    original_capture_end = torch.cuda.CUDAGraph.capture_end
+    original_flush_cache = Scheduler.flush_cache
+    capture_streams: dict[tuple[torch.device, int], torch.cuda.Stream] = {}
+    probability = jitter_probability / server_args.tp_size
+    jitter_active = False
+
+    # Compile and load the kernel before an event hook can run during CUDA
+    # graph capture, where Triton compilation is not allowed. Captured kernels
+    # retain this device address and read its current value at every replay.
+    with torch.cuda.device(gpu_id):
+        enabled = torch.zeros(1, dtype=torch.uint8, device="cuda")
+        _random_jitter(enabled, probability, jitter_max_time_us)
+        torch.cuda.synchronize()
+
+    def wrap_event(original: Any) -> Any:
+        def with_jitter(
+            event: torch.cuda.Event,
+            stream: torch.cuda.Stream | None = None,
+        ) -> None:
+            original(event, stream)
+            target_stream = torch.cuda.current_stream() if stream is None else stream
+            with torch.cuda.stream(target_stream):
+                capturing = torch.cuda.is_current_stream_capturing()
+                if capturing or jitter_active:
+                    _random_jitter(enabled, probability, jitter_max_time_us)
+            if capturing:
+                capture_streams[(target_stream.device, target_stream.cuda_stream)] = (
+                    target_stream
+                )
+
+        return with_jitter
+
+    def flush_cache(self: Any, empty_cache: bool = True) -> bool:
+        nonlocal jitter_active
+        success = original_flush_cache(self, empty_cache)
+        if success and not jitter_active:
+            enabled.fill_(1)
+            torch.cuda.synchronize()
+            jitter_active = True
+        return success
+
+    torch.cuda.Event.record = wrap_event(original_record)
+    torch.cuda.Event.wait = wrap_event(original_wait)
+    Scheduler.flush_cache = flush_cache
+
+    # A post-record delay can extend a captured side stream past the
+    # event that originally joined it. Join those streams before ending
+    # capture; eager execution needs no corresponding special case.
+    def capture_end(self: torch.cuda.CUDAGraph) -> None:
+        origin_stream = torch.cuda.current_stream()
+        origin_key = (origin_stream.device, origin_stream.cuda_stream)
+        try:
+            for key, source_stream in capture_streams.items():
+                if key == origin_key:
+                    continue
+                done = torch.cuda.Event()
+                original_record(done, source_stream)
+                original_wait(done, origin_stream)
+            original_capture_end(self)
+        finally:
+            capture_streams.clear()
+
+    torch.cuda.CUDAGraph.capture_end = capture_end
+    print(
+        "Installed stream-sync jitter "
+        f"(probability={jitter_probability}, world_size={server_args.tp_size}, "
+        f"per_rank_probability={probability}, "
+        f"time_us=uniform[0,{jitter_max_time_us}], starts_after_cache_flush=True)",
+        flush=True,
+    )
+
+    run_scheduler_process(
+        server_args,
+        port_args,
+        gpu_id,
+        *scheduler_args,
+        **scheduler_kwargs,
+    )
+
+
+@contextmanager
+def get_jitter_engine(
+    *,
+    jitter_probability: float = 0.1,
+    jitter_max_time_us: int = 10_000,
+    retract_interval: int = 500,
+    **engine_kwargs: Any,
+) -> Iterator[Engine]:
+    from sglang.srt.entrypoints.engine import Engine
+
+    assert 0.0 <= jitter_probability <= 1.0, f"{jitter_probability=}"
+    assert jitter_max_time_us > 0, f"{jitter_max_time_us=}"
+    assert retract_interval > 0, f"{retract_interval=}"
+    assert "model_path" in engine_kwargs, "model_path is required"
+
+    # Force chunked prefill and optimistic admission into the deliberately
+    # undersized pool; metrics let the workload verify that retraction occurred.
+    engine_kwargs.setdefault("chunked_prefill_size", 512)
+    engine_kwargs.setdefault("enable_metrics", True)
+    engine_kwargs.setdefault("max_total_tokens", 20_480)
+    engine_kwargs.setdefault("schedule_conservativeness", 0.05)
+
+    if (
+        engine_kwargs.get("enable_dp_attention")
+        and "dist_init_addr" not in engine_kwargs
+    ):
+        import portpicker
+
+        from sglang.srt.server_args import DP_ATTENTION_HANDSHAKE_PORT_DELTA
+
+        port = portpicker.pick_unused_port_range(  # pyright: ignore[reportAttributeAccessIssue]
+            DP_ATTENTION_HANDSHAKE_PORT_DELTA + 1
+        )[
+            0
+        ]
+        engine_kwargs["dist_init_addr"] = f"127.0.0.1:{port}"
+
+    scheduler_process = (
+        partial(
+            _run_scheduler_process_with_jitter,
+            jitter_probability=jitter_probability,
+            jitter_max_time_us=jitter_max_time_us,
+        )
+        if jitter_probability > 0
+        else Engine.run_scheduler_process_func
+    )
+
+    class JitterEngine(Engine):
+        run_scheduler_process_func = staticmethod(scheduler_process)
+
+    engine = None
+    test_env = {
+        "SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS": "True",
+        "SGLANG_TEST_RETRACT": "True",
+        "SGLANG_TEST_RETRACT_INTERVAL": str(retract_interval),
+    }
+    original_test_env = {name: os.environ.get(name) for name in test_env}
+    os.environ.update(test_env)
+    # enable_metrics registers collectors on the process-global prometheus
+    # registry, so a second engine in the same process (a CustomTestCase retry)
+    # would die on re-registration and hide the original failure.
+    from prometheus_client import REGISTRY as _prom_registry
+
+    collectors_before = set(_prom_registry._collector_to_names)
+    try:
+        engine = JitterEngine(**engine_kwargs)
+        yield engine
+    finally:
+        if engine is not None:
+            engine.shutdown()
+            del engine
+            gc.collect()
+            torch.cuda.empty_cache()
+        for collector in set(_prom_registry._collector_to_names) - collectors_before:
+            _prom_registry.unregister(collector)
+        for name, value in original_test_env.items():
+            if value is None:
+                os.environ.pop(name)
+            else:
+                os.environ[name] = value
+
+
+# --- internals -------------------------------------------------------------
+
+
+class _Request(NamedTuple):
+    prefix: int
+    cut: int
+    input_ids: list[int]
+    max_new_tokens: int
+
+    @property
+    def label(self) -> str:
+        return f"prefix{self.prefix}/cut{self.cut}"
+
+
+# An observation key identifies a predicted continuation position within one
+# independently sampled prefix and its baseline continuation.
+_Key = tuple[int, int]
+# (token id at position, top-k token->logprob)
+_Obs = tuple[int, dict[int, float]]
+
+
+def _topk_equal(a: dict[int, float], b: dict[int, float]) -> bool:
+    """Compare top-k maps bitwise, tolerating rank-k boundary ties."""
+    if len(a) != len(b):
+        return False
+    if any(a[token] != b[token] for token in a.keys() & b.keys()):
+        return False
+    return all(
+        source[token] == min(other.values())
+        for only, source, other in (
+            (a.keys() - b.keys(), a, b),
+            (b.keys() - a.keys(), b, a),
+        )
+        for token in only
+    )
+
+
+def _record_response(
+    observations: dict[_Key, list[_Obs]],
+    req: _Request,
+    out: dict[str, Any],
+    baseline: list[int] | None,
+) -> list[int]:
+    """Record every scored position of one response; returns its output ids.
+
+    `baseline` is the prefix's expected greedy continuation (None while
+    recording the baseline itself). Positions past the first divergence from
+    the baseline are dropped: their context differs from every other request
+    at the same nominal position, so comparisons there are meaningless. The
+    divergence position itself still has identical context and is recorded —
+    its top-k diff is exactly the interesting signal.
+    """
+
+    meta = out["meta_info"]
+
+    out_ids = [int(t) for t in out["output_ids"]]
+    out_lps = meta["output_token_logprobs"]
+    out_top = meta["output_top_logprobs"]
+    assert (
+        len(out_ids) == req.max_new_tokens
+    ), f"{req.label}: got {len(out_ids)} output tokens, expected {req.max_new_tokens}"
+    expected = out_ids if baseline is None else baseline[req.cut :]
+    for m, (tid, lp_entry, top) in enumerate(
+        zip(out_ids, out_lps, out_top, strict=True)
+    ):
+        assert int(lp_entry[1]) == tid, f"{req.label}: output logprob misaligned at {m}"
+        pos = len(req.input_ids) + m
+        topk_map = {int(entry[1]): float(entry[0]) for entry in top}
+        assert len(topk_map) == len(top), f"duplicate token ids in top-k: {top}"
+        observations[(req.prefix, pos)].append((tid, topk_map))
+        if tid != expected[m]:
+            break
+    return out_ids
+
+
+def _generate_batch(
+    engine: Engine, requests: list[_Request], topk: int
+) -> list[dict[str, Any]]:
+    """One batched generate call; per-request fields go through the batch
+    API's list parameters, results come back in submission order."""
+    outs = engine.generate(
+        input_ids=[list(r.input_ids) for r in requests],
+        sampling_params=[
+            {"temperature": 0.0, "max_new_tokens": r.max_new_tokens, "ignore_eos": True}
+            for r in requests
+        ],
+        return_logprob=[True] * len(requests),
+        # -1 requests output logprobs only and permits full radix-cache reuse.
+        logprob_start_len=[-1] * len(requests),
+        top_logprobs_num=[topk] * len(requests),
+    )
+    results = list(outs)
+    assert all(isinstance(r, dict) for r in results)
+    return results
+
+
+def _total_retracted_requests() -> float:
+    from prometheus_client import CollectorRegistry
+    from prometheus_client import multiprocess as prom_multiprocess
+
+    assert "PROMETHEUS_MULTIPROC_DIR" in os.environ
+    registry = CollectorRegistry()
+    prom_multiprocess.MultiProcessCollector(registry)
+    return sum(
+        sample.value
+        for metric in registry.collect()
+        if "retracted_req" in metric.name
+        for sample in metric.samples
+        if sample.name.endswith("_total")
+    )
+
+
+def run_jitter_test(
+    engine: Engine,
+    *,
+    num_unique_prefixes: int = 10,  # independently sampled random-token prefixes
+    requests_per_prefix: int = 10,  # output-prefix cut requests per prefix
+    prefix_len_min: int = 2048,  # prefix length range, inclusive
+    prefix_len_max: int = 4096,
+    new_tokens: int = 2048,  # baseline generation length
+    topk: int = 8,  # how many logprob entries per position are compared
+    seed: int = 0,  # seeds every draw: token sequence, lengths, cuts, shuffle
+    min_retracted_requests: int = 10,
+    vocab_size: int = 199_998,  # excludes special tokens / padded embedding tail
+) -> None:
+    """Run the workload on a caller-constructed engine and assert bitwise
+    consistency across all overlapping observations."""
+    assert min_retracted_requests >= 0
+    prefix_rng = random.Random(seed + 1)
+    prefix_lens = sorted(
+        prefix_rng.randint(prefix_len_min, prefix_len_max)
+        for _ in range(num_unique_prefixes)
+    )
+    # Random output-cut lengths per (prefix, cut) pair, capped so every
+    # request still decodes at least 64 tokens; 0 stays in range (an exact
+    # duplicate of the baseline scheduled in a different batch mix =
+    # run-to-run check).
+    cut_rng = random.Random(seed + 2)
+    cut_lens = [
+        [cut_rng.randint(0, new_tokens - 64) for _ in range(requests_per_prefix)]
+        for _ in range(num_unique_prefixes)
+    ]
+    token_rng = random.Random(seed)
+    prefix_tokens = [
+        [token_rng.randrange(vocab_size) for _ in range(prefix_len)]
+        for prefix_len in prefix_lens
+    ]
+
+    observations: dict[_Key, list[_Obs]] = defaultdict(list)
+
+    # --- Phase 1: greedy baselines, one request per prefix; outputs define
+    # the prefixes' canonical continuations ---
+    t1 = time.monotonic()
+    baseline_reqs = [
+        _Request(prefix, 0, tokens, new_tokens)
+        for prefix, tokens in enumerate(prefix_tokens)
+    ]
+    outs = _generate_batch(engine, baseline_reqs, topk)
+    print(f"[timing] phase 1 (baselines): {time.monotonic() - t1:.1f}s")
+    baselines = {}
+    for req, out in zip(baseline_reqs, outs, strict=True):
+        baselines[req.prefix] = _record_response(observations, req, out, baseline=None)
+
+    # The scheduler only flushes when fully idle; the batch API can return
+    # before the last request has fully drained, so retry briefly.
+    for attempt in range(10):
+        flush_result = engine.flush_cache()
+        if getattr(flush_result, "success", True):
+            break
+        time.sleep(1.0)
+    else:
+        raise AssertionError(f"cache flush failed after retries: {flush_result}")
+    retractions_before_jitter = (
+        _total_retracted_requests() if min_retracted_requests else 0.0
+    )
+
+    # --- Phase 2: shuffled jitter batch — per prefix, requests from the
+    # prefix plus a random-length cut of the baseline output, decoding out to
+    # the baseline's total length ---
+    jitter_reqs = [
+        _Request(
+            prefix,
+            cut,
+            tokens + baselines[prefix][:cut],
+            new_tokens - cut,
+        )
+        for prefix, tokens in enumerate(prefix_tokens)
+        for cut in cut_lens[prefix]
+    ]
+    random.Random(seed).shuffle(jitter_reqs)
+    t2 = time.monotonic()
+    jitter_outs = _generate_batch(engine, jitter_reqs, topk)
+    print(f"[timing] phase 2 (jitter batch): {time.monotonic() - t2:.1f}s")
+    for req, out in zip(jitter_reqs, jitter_outs, strict=True):
+        _record_response(observations, req, out, baseline=baselines[req.prefix])
+
+    # --- Cross-check every observation sharing a key ---
+    num_cross_checked_positions = 0
+    num_mismatches = 0
+    mismatch_details: list[str] = []
+    for key in sorted(observations):
+        group = observations[key]
+        if len(group) >= 2:
+            num_cross_checked_positions += 1
+        ref_tid, ref_topk = group[0]
+        for tid, top in group[1:]:
+            if tid == ref_tid and _topk_equal(ref_topk, top):
+                continue
+            num_mismatches += 1
+            if len(mismatch_details) < 5:
+                shared = ref_topk.keys() & top.keys()
+                mismatch_details.append(
+                    f"prefix={key[0]} pos={key[1]} token {ref_tid} vs {tid} "
+                    f"max|dlogprob|="
+                    f"{max((abs(ref_topk[t] - top[t]) for t in shared), default=0.0):.3e}"
+                )
+    print(
+        f"checked {len(observations)} positions, "
+        f"{num_cross_checked_positions} observed by >=2 requests, "
+        f"{sum(len(g) for g in observations.values())} observations total, "
+        f"{num_mismatches} mismatches"
+    )
+    # Guard against the test silently becoming vacuous: most continuation
+    # positions must be observed by several requests.
+    assert num_cross_checked_positions >= len(prefix_lens) * new_tokens // 2, (
+        f"only {num_cross_checked_positions} positions were cross-checked; "
+        "the request construction no longer produces overlapping observations"
+    )
+    assert num_mismatches == 0, "\n".join(
+        [f"{num_mismatches} mismatching observations"] + mismatch_details
+    )
+    if min_retracted_requests:
+        retractions_during_jitter = (
+            _total_retracted_requests() - retractions_before_jitter
+        )
+        assert retractions_during_jitter >= min_retracted_requests, (
+            f"expected at least {min_retracted_requests} retractions during the jitter "
+            f"batch, got {retractions_during_jitter}"
+        )

test/registered/models_e2e/test_inkling.py

@@ -18,6 +18,7 @@ import unittest
 import requests
 
 from sglang.srt.utils import kill_process_tree
+from sglang.test.cache_consistency_jitter import get_jitter_engine, run_jitter_test
 from sglang.test.ci.ci_register import register_cuda_ci
 
 # Aliased so pytest does not collect the imported `test_`-prefixed helper as a test.
@@ -31,7 +32,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=250, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=450, stage="base-b", runner_config="1-gpu-large")
 
 # Defaults to the HF `test` revision; override MODEL/REVISION to point at a
 # local checkpoint. Empty REVISION drops the flag (for local paths).
@@ -187,5 +188,42 @@ class TestInklingServer(CustomTestCase):
         )
 
 
+class TestInklingCacheConsistency(CustomTestCase):
+    """Bitwise version of the KL check above: the same context is scored under
+    different batch shapes, cache histories and retraction timing, and every
+    overlapping observation must agree. Boots its own in-process engine because
+    the harness patches the scheduler process to inject stream-sync jitter."""
+
+    def test_scored_contexts_are_bitwise_identical(self):
+        engine_kwargs = {
+            "model_path": _MODEL_PATH,
+            "trust_remote_code": True,
+            "attention_backend": "fa4",
+            "page_size": 128,
+            "mamba_radix_cache_strategy": "extra_buffer",
+            "swa_full_tokens_ratio": 0.1,
+            "mamba_full_memory_ratio": 0.1,
+            "mem_fraction_static": 0.5,
+            "enable_deterministic_inference": True,
+        }
+        if _MODEL_REVISION:
+            engine_kwargs["revision"] = _MODEL_REVISION
+
+        with get_jitter_engine(**engine_kwargs) as engine:
+            # Sized to the 20480-token pool the harness pins: large enough to
+            # force retraction, small enough that the batch still admits.
+            run_jitter_test(
+                engine,
+                num_unique_prefixes=4,
+                requests_per_prefix=5,
+                prefix_len_min=384,
+                prefix_len_max=512,
+                new_tokens=256,
+                # How many requests get retracted tracks pool size, hence GPU
+                # memory, so only assert the path ran at all.
+                min_retracted_requests=1,
+            )
+
+
 if __name__ == "__main__":
     unittest.main()