diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/batch_invariance.py b/python/sglang/kernels/ops/attention/flash_attn/cute/batch_invariance.py new file mode 100644 index 000000000..3aae353d2 --- /dev/null +++ b/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 diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/block_info.py b/python/sglang/kernels/ops/attention/flash_attn/cute/block_info.py index f49253eef..962c08744 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/block_info.py +++ b/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 diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_sm100.py b/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_sm100.py index d9c12dfa5..07e8192cb 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_sm100.py +++ b/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, diff --git a/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py b/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py index 318cc2b28..2e4e00ddc 100644 --- a/python/sglang/kernels/ops/attention/flash_attn/cute/interface.py +++ b/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, ) ), ) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index fddf9f99d..ee2337f7e 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/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=}") diff --git a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py b/python/sglang/srt/model_executor/runner/flashinfer_autotune.py index 6790070a6..d7d1fda99 100644 --- a/python/sglang/srt/model_executor/runner/flashinfer_autotune.py +++ b/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: diff --git a/python/sglang/srt/models/inkling_common/attn.py b/python/sglang/srt/models/inkling_common/attn.py index fd659c719..7df207b06 100644 --- a/python/sglang/srt/models/inkling_common/attn.py +++ b/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() diff --git a/python/sglang/srt/models/inkling_common/kernels/comm.py b/python/sglang/srt/models/inkling_common/kernels/comm.py index 876baca84..fa0fc14c6 100644 --- a/python/sglang/srt/models/inkling_common/kernels/comm.py +++ b/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); diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 50900c44e..b754b5a45 100644 --- a/python/sglang/srt/server_args.py +++ b/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): diff --git a/python/sglang/test/cache_consistency_jitter.py b/python/sglang/test/cache_consistency_jitter.py new file mode 100644 index 000000000..ee5d1e1fc --- /dev/null +++ b/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}" + ) diff --git a/test/registered/models_e2e/test_inkling.py b/test/registered/models_e2e/test_inkling.py index a84441a2a..5ca5ef4ca 100644 --- a/test/registered/models_e2e/test_inkling.py +++ b/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()