[Feature] Add leveled invariant-check primitive for nan/inf/oob validity checks (#32308)

python/sglang/srt/environ.py

@@ -210,6 +210,22 @@ class ToolStrictLevel(IntEnum):
     PARAMETER = 2
 
 
+class InvariantCheckLevel(IntEnum):
+    """Signal level for value/index validity checks (see invariants.py).
+
+    OFF: data layer only (sanitize/containment); no detection, no signal.
+    WARN: detect + throttled log/count; degrade, never crash (prod on-demand).
+    STRICT: detect + crash on GUARD/FATAL violations (CI default).
+
+    The data layer is unconditional and independent of this level; only the
+    detection + signal layer is gated here.
+    """
+
+    OFF = 0
+    WARN = 1
+    STRICT = 2
+
+
 class Envs:
 
     # Raise on bare server_args field assignments after resolution; mutation
@@ -786,6 +802,12 @@ class Envs:
     # page alignment). Off in prod; tests turn it on to fail-fast on
     # numerical / index violations instead of getting silent NaN cascades.
     SGLANG_ENABLE_ASYNC_ASSERT = EnvBool(False)
+    # Signal level for value/index validity checks (nan/inf/oob/...); see
+    # invariants.py. OFF (prod default) runs only the free data layer, WARN
+    # adds throttled logging, STRICT (CI default) crashes on violations.
+    # Supersedes SGLANG_ENABLE_ASYNC_ASSERT, which is bridged as STRICT until
+    # every callsite migrates.
+    SGLANG_INVARIANT_CHECK = EnvInt(InvariantCheckLevel.OFF)
     # Sanitize NaN logits before sampling kernels and log a throttled warning
     # (see sanitize_nan_logits).
     SGLANG_SANITIZE_NAN_LOGITS = EnvBool(False)

python/sglang/srt/models/deepseek_v4_dspark.py

@@ -49,12 +49,17 @@ from sglang.srt.speculative.ragged_verify import (
     read_ragged_verify_mode,
 )
 from sglang.srt.utils import add_prefix, is_blackwell_supported
-from sglang.srt.utils.async_probe import maybe_detect_in_closed_range
+from sglang.srt.utils.invariants import Bucket, InClosedRange, Invariant, expect
 
 logger = logging.getLogger(__name__)
 
 _PAD_NUM_HEADS = 64
 
+# DSpark confidence is a per-token score that must stay in [0, 1].
+_CONFIDENCE = Invariant(
+    "dspark.model.confidence", Bucket.GUARD, InClosedRange(0.0, 1.0)
+)
+
 
 def apply_rotary_emb(
     x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False
@@ -749,9 +754,7 @@ class DeepseekV4ForCausalLMDSpark(nn.Module):
             markov_embed_stack = None
         confidence_raw = confidence_head(x_post_hc, markov_embed_stack)
         confidence = confidence_head.apply_sts(confidence_raw)
-        maybe_detect_in_closed_range(
-            confidence, 0.0, 1.0, "DSpark confidence must lie in [0, 1]."
-        )
+        expect(_CONFIDENCE, confidence)
         return confidence
 
     def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None:

python/sglang/srt/speculative/dspark_components/dspark_draft.py

@@ -26,11 +26,28 @@ from sglang.srt.speculative.spec_info import (
     spec_scale_global_num_tokens,
 )
 from sglang.srt.speculative.spec_utils import draft_tp_context
-from sglang.srt.utils.async_probe import maybe_detect_nan
+from sglang.srt.utils.invariants import Bucket, Invariant, NotNaN, expect
 
 logger = logging.getLogger(__name__)
 
 
+def _one_hot_token0(probs: torch.Tensor) -> torch.Tensor:
+    degenerate = torch.isnan(probs[:, :1])
+    one_hot = torch.zeros_like(probs)
+    one_hot[:, 0] = 1.0
+    return torch.where(degenerate, one_hot, probs)
+
+
+# Draft step logits: NaN is a bug, but -inf is legitimate (masking). The data
+# layer lives downstream (probs one-hot below, or the fast kernel's clamp).
+_DRAFT_STEP_LOGITS = Invariant("dspark.draft.step_logits", Bucket.GUARD, NotNaN())
+# Draft sampling probs: SOFTEN (tolerate + count), matching the original
+# unconditional clamp; an all-NaN row would otherwise make multinomial raise.
+_DRAFT_PROBS = Invariant(
+    "dspark.draft.probs", Bucket.SOFTEN, NotNaN(), recover=_one_hot_token0
+)
+
+
 class DraftBlockResult(msgspec.Struct, frozen=True):
     draft_tokens: torch.Tensor
     corrected_logits: Optional[torch.Tensor]
@@ -181,13 +198,13 @@ def sample_draft_block(
     if not any_sampling:
 
         def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
-            maybe_detect_nan(step_logits, f"dspark draft step {step_idx}")
+            expect(_DRAFT_STEP_LOGITS, step_logits, msg=f"step {step_idx}")
             return torch.argmax(step_logits, dim=-1)
 
     else:
 
         def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
-            maybe_detect_nan(step_logits, f"dspark draft step {step_idx}")
+            expect(_DRAFT_STEP_LOGITS, step_logits, msg=f"step {step_idx}")
             if fast_sampling:
                 exp_noise = torch.empty(
                     step_logits.shape, dtype=torch.float32, device=step_logits.device
@@ -202,11 +219,7 @@ def sample_draft_block(
                 probs = torch.softmax(
                     step_logits.float() / temperatures[:, None], dim=-1
                 )
-                # All-NaN rows make multinomial raise; clamp to one-hot token 0.
-                degenerate_rows = torch.isnan(probs[:, :1])
-                one_hot_token0 = torch.zeros_like(probs)
-                one_hot_token0[:, 0] = 1.0
-                probs = torch.where(degenerate_rows, one_hot_token0, probs)
+                probs = expect(_DRAFT_PROBS, probs)
                 argmax_tokens = torch.argmax(step_logits, dim=-1)
                 sampled_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
                 return torch.where(greedy_mask, argmax_tokens, sampled_tokens)

python/sglang/srt/speculative/dspark_components/dspark_planner.py

@@ -12,7 +12,7 @@ from sglang.kernels.ops.speculative.dspark.dspark_schedule import (
     compute_sort_survival,
 )
 from sglang.srt.distributed import get_tp_group
-from sglang.srt.environ import envs
+from sglang.srt.environ import InvariantCheckLevel, envs
 from sglang.srt.layers.dp_attention import is_dp_attention_enabled
 from sglang.srt.managers.overlap_utils import (
     CONFIDENCE_RELAY_RING_LAG,
@@ -41,14 +41,25 @@ from sglang.srt.speculative.ragged_verify import (
     read_ragged_verify_mode,
     round_up_grid,
 )
-from sglang.srt.utils.async_probe import (
-    maybe_assert_async,
-    maybe_detect_in_closed_range,
-)
 from sglang.srt.utils.common import require_mlp_tp_gather
+from sglang.srt.utils.invariants import (
+    Bucket,
+    InClosedRange,
+    Invariant,
+    IsTrue,
+    expect,
+    resolve_level,
+)
 
 logger = logging.getLogger(__name__)
 
+# DSpark confidence is a per-token score that must stay in [0, 1].
+_CONFIDENCE = Invariant(
+    "dspark.planner.confidence", Bucket.GUARD, InClosedRange(0.0, 1.0)
+)
+# Scheduled verify lengths must not exceed the per-step token budget.
+_VERIFY_LEN_BUDGET = Invariant("dspark.verify_len_budget", Bucket.GUARD, IsTrue())
+
 
 class VerifyWindow(msgspec.Struct, frozen=True):
     positions_2d: torch.Tensor
@@ -580,12 +591,13 @@ class DSparkVerifyPlanner:
             cfg=self._schedule_cfg,
         ).to(device=device, dtype=torch.int32)
 
-        if envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        if resolve_level() >= InvariantCheckLevel.WARN:
             verify_lens_64 = verify_lens.to(torch.int64)
             effective_floor = max(self._schedule_cfg.min_verify_len, 1)
-            maybe_assert_async(
+            expect(
+                _VERIFY_LEN_BUDGET,
                 (verify_lens_64 - effective_floor).sum() <= budget,
-                f"DSpark verify-len budget violated (budget={budget})",
+                msg=f"budget={budget}",
             )
 
         if envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER.get():
@@ -907,7 +919,7 @@ def compute_confidence(
         markov_embed_stack = None
     confidence_raw = confidence_head(draft_hidden, markov_embed_stack)
     confidence = confidence_head.apply_sts(confidence_raw)
-    maybe_detect_in_closed_range(confidence, 0.0, 1.0, "DSpark confidence")
+    expect(_CONFIDENCE, confidence)
     return confidence
 
 

python/sglang/srt/speculative/dspark_components/dspark_verify.py

@@ -37,7 +37,11 @@ from sglang.srt.speculative.dspark_components.dspark_planner import (
     apply_logits_adjustments_strided,
 )
 from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
-from sglang.srt.utils.async_probe import maybe_detect_nan
+from sglang.srt.utils.invariants import Bucket, Invariant, NotNaN, expect
+
+# Draft proposal probs feeding rejection sampling; the data layer is the
+# in-kernel NaN-q guard in reject_sampling.py, so this is signal-only.
+_VERIFY_DRAFT_PROBS = Invariant("dspark.verify.draft_probs", Bucket.GUARD, NotNaN())
 
 
 def verify_logits_adjustments_are_noop(sampling_info) -> bool:
@@ -678,7 +682,7 @@ def accept_draft_tokens(
         temperatures=draft_block.temperatures,
         rows_per_request=gamma_rows,
     ).view(bs, gamma_rows, vocab)
-    maybe_detect_nan(draft_probs, "dspark verify: draft_probs")
+    expect(_VERIFY_DRAFT_PROBS, draft_probs)
     if not sampling_info.is_any_greedy:
         return AcceptSampling.execute(
             candidates=candidates,

python/sglang/srt/utils/invariants.py

@@ -0,0 +1,245 @@
+"""Value/index validity checks -- the invariant-check model.
+
+Two layers per check:
+  * data layer (sanitize / containment): unconditional, branchless, ~free.
+    May live in `recover` here or inside the kernel (then `recover=None`).
+  * signal layer (detect + log/crash): gated by SGLANG_INVARIANT_CHECK
+    (off / warn / strict).
+
+A `Bucket` (blast radius x recoverability) and the level decide whether a hit
+crashes, logs, or is silent. Detection is async (no GPU-CPU sync): crashes via
+torch's async assert, counts via a pinned-memory readback.
+"""
+
+from __future__ import annotations
+
+import enum
+import logging
+from typing import Callable, Optional
+
+import torch
+
+from sglang.srt.environ import InvariantCheckLevel, envs
+
+logger = logging.getLogger(__name__)
+
+
+class Bucket(enum.Enum):
+    """Invariant classification by blast radius and recoverability."""
+
+    SOFTEN = "soften"  # recoverable + legitimate event (never a bug)
+    GUARD = "guard"  # recoverable + is-a-bug / root cause unknown
+    FATAL_CONTAINABLE = "fatal_containable"  # no correct fallback; cheap containment
+    FATAL_UNCONTAINABLE = "fatal_uncontainable"  # no containment; global corruption
+
+
+def resolve_level() -> InvariantCheckLevel:
+    """Current signal level; an unset SGLANG_INVARIANT_CHECK falls back to the
+    legacy SGLANG_ENABLE_ASYNC_ASSERT=true as STRICT so CI keeps failing loud."""
+    if envs.SGLANG_INVARIANT_CHECK.is_set():
+        return InvariantCheckLevel(envs.SGLANG_INVARIANT_CHECK.get())
+    if envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return InvariantCheckLevel.STRICT
+    return InvariantCheckLevel.OFF
+
+
+class Property:
+    """A per-element validity predicate; returns an elementwise bool, no reduction."""
+
+    name: str
+
+    def ok(self, value: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class NotNaN(Property):
+    """Not NaN, tolerating +-Inf (e.g. legitimately masked -inf logits)."""
+
+    name = "not_nan"
+
+    def ok(self, value: torch.Tensor) -> torch.Tensor:
+        return ~torch.isnan(value)
+
+
+class NotInf(Property):
+    """Not +-Inf (e.g. fp16 overflow), tolerating NaN."""
+
+    name = "not_inf"
+
+    def ok(self, value: torch.Tensor) -> torch.Tensor:
+        return ~torch.isinf(value)
+
+
+class Finite(Property):
+    """Neither NaN nor Inf -- use when inf is also a bug for this tensor."""
+
+    name = "finite"
+
+    def ok(self, value: torch.Tensor) -> torch.Tensor:
+        return torch.isfinite(value)
+
+
+class InRange(Property):
+    """Half-open [lo, hi) -- unifies oob / range / index-domain checks."""
+
+    def __init__(self, lo, hi):
+        self.lo = lo
+        self.hi = hi
+        self.name = f"in_range[{lo},{hi})"
+
+    def ok(self, value: torch.Tensor) -> torch.Tensor:
+        return (value >= self.lo) & (value < self.hi)
+
+
+class InClosedRange(Property):
+    """Closed [lo, hi] -- for value sanity (e.g. a probability / confidence)."""
+
+    def __init__(self, lo, hi):
+        self.lo = lo
+        self.hi = hi
+        self.name = f"in_closed_range[{lo},{hi}]"
+
+    def ok(self, value: torch.Tensor) -> torch.Tensor:
+        return (value >= self.lo) & (value <= self.hi)
+
+
+class PageAligned(Property):
+    def __init__(self, page_size: int):
+        self.page_size = page_size
+        self.name = f"page_aligned[{page_size}]"
+
+    def ok(self, value: torch.Tensor) -> torch.Tensor:
+        return value % self.page_size == 0
+
+
+class IsTrue(Property):
+    """The value is itself the boolean condition (for derived / scalar asserts)."""
+
+    name = "is_true"
+
+    def ok(self, value: torch.Tensor) -> torch.Tensor:
+        return value
+
+
+# name -> Invariant; enumerated by the CI meta-test to enforce injection coverage.
+_REGISTRY: dict[str, Invariant] = {}
+
+
+class Invariant:
+    """A declared invariant; constructing one registers it. `recover` is the
+    optional python-side data layer (None = signal-only, in the kernel);
+    FATAL_UNCONTAINABLE forbids it."""
+
+    def __init__(
+        self,
+        name: str,
+        bucket: Bucket,
+        prop: Property,
+        *,
+        recover: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+    ):
+        if bucket is Bucket.FATAL_UNCONTAINABLE and recover is not None:
+            raise ValueError(f"uncontainable FATAL {name!r} cannot have a recover")
+        if name in _REGISTRY:
+            raise ValueError(f"duplicate invariant {name!r}")
+        self.name = name
+        self.bucket = bucket
+        self.prop = prop
+        self.recover = recover
+        _REGISTRY[name] = self
+
+
+def registered_invariants() -> dict[str, Invariant]:
+    return dict(_REGISTRY)
+
+
+def _get_rank() -> int:
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        return torch.distributed.get_rank()
+    return 0
+
+
+# Per-(rank, name) checks between aggregated log flushes.
+_FLUSH_EVERY = 512
+
+
+class _CheckReporter:
+    """Sync-free per-(rank, name) hit counter. Counts accumulate on-device and
+    mirror to pinned host via a non-blocking copy; the host reads the (stale)
+    count on a later call. First hit logs immediately, the rest on a cadence."""
+
+    def __init__(self):
+        self._dev: dict[str, torch.Tensor] = {}
+        self._host: dict[str, torch.Tensor] = {}
+        self._calls: dict[str, int] = {}
+        self._logged_total: dict[str, int] = {}
+
+    def record(self, key: str, bucket: Bucket, hit_count: torch.Tensor, msg: str):
+        dev = self._dev.get(key)
+        if dev is None:
+            dev = torch.zeros(1, dtype=torch.int64, device=hit_count.device)
+            self._dev[key] = dev
+            self._host[key] = torch.zeros(
+                1, dtype=torch.int64, pin_memory=hit_count.is_cuda
+            )
+            self._calls[key] = 0
+            self._logged_total[key] = 0
+
+        dev.add_(hit_count.reshape(1).to(torch.int64))
+        self._host[key].copy_(dev, non_blocking=True)  # async, no sync
+        self._calls[key] += 1
+
+        total = int(self._host[key][0])  # stale host read, no sync
+        last = self._logged_total[key]
+        first_hit = last == 0 and total > 0
+        cadence = self._calls[key] % _FLUSH_EVERY == 0
+        if total > last and (first_hit or cadence):
+            level = logging.INFO if bucket is Bucket.SOFTEN else logging.WARNING
+            logger.log(
+                level,
+                "invariant-check [%s]: +%d hit(s), %d total. %s",
+                key,
+                total - last,
+                total,
+                msg,
+            )
+            self._logged_total[key] = total
+
+
+_reporter = _CheckReporter()
+
+
+def _crashes(bucket: Bucket, level: InvariantCheckLevel) -> bool:
+    """The (bucket x level) crash decision, pure and total."""
+    if bucket is Bucket.SOFTEN:
+        return False
+    if bucket is Bucket.FATAL_UNCONTAINABLE:
+        return level >= InvariantCheckLevel.WARN
+    return level == InvariantCheckLevel.STRICT  # GUARD, FATAL_CONTAINABLE
+
+
+def _signal(ok: torch.Tensor, *, inv: Invariant, level: InvariantCheckLevel, msg: str):
+    if _crashes(inv.bucket, level):
+        # Loud: async assert surfaces at the next sync point (no CPU sync).
+        torch._assert_async(ok.all(), f"invariant-check FAILED [{inv.name}]: {msg}")
+        return
+    _reporter.record(f"{inv.name}@rank{_get_rank()}", inv.bucket, (~ok).sum(), msg)
+
+
+def expect(
+    inv: Invariant,
+    value: Optional[torch.Tensor],
+    *,
+    msg: str = "",
+) -> Optional[torch.Tensor]:
+    """Check `inv` over `value` per the (bucket x level) matrix. The data layer
+    (`inv.recover`) is applied unconditionally; only detection is gated. Returns
+    the recovered value."""
+    level = resolve_level()
+    if level >= InvariantCheckLevel.WARN and value is not None and value.numel() > 0:
+        detail = f"{inv.prop.name}: {msg}" if msg else inv.prop.name
+        _signal(inv.prop.ok(value), inv=inv, level=level, msg=detail)
+
+    if inv.recover is not None and value is not None:
+        value = inv.recover(value)
+    return value