[Feature] Add leveled invariant-check primitive for nan/inf/oob validity checks (#32308)
This commit is contained in:
@@ -210,6 +210,22 @@ class ToolStrictLevel(IntEnum):
|
|||||||
PARAMETER = 2
|
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:
|
class Envs:
|
||||||
|
|
||||||
# Raise on bare server_args field assignments after resolution; mutation
|
# 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
|
# page alignment). Off in prod; tests turn it on to fail-fast on
|
||||||
# numerical / index violations instead of getting silent NaN cascades.
|
# numerical / index violations instead of getting silent NaN cascades.
|
||||||
SGLANG_ENABLE_ASYNC_ASSERT = EnvBool(False)
|
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
|
# Sanitize NaN logits before sampling kernels and log a throttled warning
|
||||||
# (see sanitize_nan_logits).
|
# (see sanitize_nan_logits).
|
||||||
SGLANG_SANITIZE_NAN_LOGITS = EnvBool(False)
|
SGLANG_SANITIZE_NAN_LOGITS = EnvBool(False)
|
||||||
|
|||||||
@@ -49,12 +49,17 @@ from sglang.srt.speculative.ragged_verify import (
|
|||||||
read_ragged_verify_mode,
|
read_ragged_verify_mode,
|
||||||
)
|
)
|
||||||
from sglang.srt.utils import add_prefix, is_blackwell_supported
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
_PAD_NUM_HEADS = 64
|
_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(
|
def apply_rotary_emb(
|
||||||
x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False
|
x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False
|
||||||
@@ -749,9 +754,7 @@ class DeepseekV4ForCausalLMDSpark(nn.Module):
|
|||||||
markov_embed_stack = None
|
markov_embed_stack = None
|
||||||
confidence_raw = confidence_head(x_post_hc, markov_embed_stack)
|
confidence_raw = confidence_head(x_post_hc, markov_embed_stack)
|
||||||
confidence = confidence_head.apply_sts(confidence_raw)
|
confidence = confidence_head.apply_sts(confidence_raw)
|
||||||
maybe_detect_in_closed_range(
|
expect(_CONFIDENCE, confidence)
|
||||||
confidence, 0.0, 1.0, "DSpark confidence must lie in [0, 1]."
|
|
||||||
)
|
|
||||||
return confidence
|
return confidence
|
||||||
|
|
||||||
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None:
|
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> None:
|
||||||
|
|||||||
@@ -26,11 +26,28 @@ from sglang.srt.speculative.spec_info import (
|
|||||||
spec_scale_global_num_tokens,
|
spec_scale_global_num_tokens,
|
||||||
)
|
)
|
||||||
from sglang.srt.speculative.spec_utils import draft_tp_context
|
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__)
|
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):
|
class DraftBlockResult(msgspec.Struct, frozen=True):
|
||||||
draft_tokens: torch.Tensor
|
draft_tokens: torch.Tensor
|
||||||
corrected_logits: Optional[torch.Tensor]
|
corrected_logits: Optional[torch.Tensor]
|
||||||
@@ -181,13 +198,13 @@ def sample_draft_block(
|
|||||||
if not any_sampling:
|
if not any_sampling:
|
||||||
|
|
||||||
def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
|
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)
|
return torch.argmax(step_logits, dim=-1)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
|
|
||||||
def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
|
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:
|
if fast_sampling:
|
||||||
exp_noise = torch.empty(
|
exp_noise = torch.empty(
|
||||||
step_logits.shape, dtype=torch.float32, device=step_logits.device
|
step_logits.shape, dtype=torch.float32, device=step_logits.device
|
||||||
@@ -202,11 +219,7 @@ def sample_draft_block(
|
|||||||
probs = torch.softmax(
|
probs = torch.softmax(
|
||||||
step_logits.float() / temperatures[:, None], dim=-1
|
step_logits.float() / temperatures[:, None], dim=-1
|
||||||
)
|
)
|
||||||
# All-NaN rows make multinomial raise; clamp to one-hot token 0.
|
probs = expect(_DRAFT_PROBS, probs)
|
||||||
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)
|
|
||||||
argmax_tokens = torch.argmax(step_logits, dim=-1)
|
argmax_tokens = torch.argmax(step_logits, dim=-1)
|
||||||
sampled_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
|
sampled_tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
|
||||||
return torch.where(greedy_mask, argmax_tokens, sampled_tokens)
|
return torch.where(greedy_mask, argmax_tokens, sampled_tokens)
|
||||||
|
|||||||
@@ -12,7 +12,7 @@ from sglang.kernels.ops.speculative.dspark.dspark_schedule import (
|
|||||||
compute_sort_survival,
|
compute_sort_survival,
|
||||||
)
|
)
|
||||||
from sglang.srt.distributed import get_tp_group
|
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.layers.dp_attention import is_dp_attention_enabled
|
||||||
from sglang.srt.managers.overlap_utils import (
|
from sglang.srt.managers.overlap_utils import (
|
||||||
CONFIDENCE_RELAY_RING_LAG,
|
CONFIDENCE_RELAY_RING_LAG,
|
||||||
@@ -41,14 +41,25 @@ from sglang.srt.speculative.ragged_verify import (
|
|||||||
read_ragged_verify_mode,
|
read_ragged_verify_mode,
|
||||||
round_up_grid,
|
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.common import require_mlp_tp_gather
|
||||||
|
from sglang.srt.utils.invariants import (
|
||||||
|
Bucket,
|
||||||
|
InClosedRange,
|
||||||
|
Invariant,
|
||||||
|
IsTrue,
|
||||||
|
expect,
|
||||||
|
resolve_level,
|
||||||
|
)
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
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):
|
class VerifyWindow(msgspec.Struct, frozen=True):
|
||||||
positions_2d: torch.Tensor
|
positions_2d: torch.Tensor
|
||||||
@@ -580,12 +591,13 @@ class DSparkVerifyPlanner:
|
|||||||
cfg=self._schedule_cfg,
|
cfg=self._schedule_cfg,
|
||||||
).to(device=device, dtype=torch.int32)
|
).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)
|
verify_lens_64 = verify_lens.to(torch.int64)
|
||||||
effective_floor = max(self._schedule_cfg.min_verify_len, 1)
|
effective_floor = max(self._schedule_cfg.min_verify_len, 1)
|
||||||
maybe_assert_async(
|
expect(
|
||||||
|
_VERIFY_LEN_BUDGET,
|
||||||
(verify_lens_64 - effective_floor).sum() <= 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():
|
if envs.SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER.get():
|
||||||
@@ -907,7 +919,7 @@ def compute_confidence(
|
|||||||
markov_embed_stack = None
|
markov_embed_stack = None
|
||||||
confidence_raw = confidence_head(draft_hidden, markov_embed_stack)
|
confidence_raw = confidence_head(draft_hidden, markov_embed_stack)
|
||||||
confidence = confidence_head.apply_sts(confidence_raw)
|
confidence = confidence_head.apply_sts(confidence_raw)
|
||||||
maybe_detect_in_closed_range(confidence, 0.0, 1.0, "DSpark confidence")
|
expect(_CONFIDENCE, confidence)
|
||||||
return confidence
|
return confidence
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -37,7 +37,11 @@ from sglang.srt.speculative.dspark_components.dspark_planner import (
|
|||||||
apply_logits_adjustments_strided,
|
apply_logits_adjustments_strided,
|
||||||
)
|
)
|
||||||
from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout
|
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:
|
def verify_logits_adjustments_are_noop(sampling_info) -> bool:
|
||||||
@@ -678,7 +682,7 @@ def accept_draft_tokens(
|
|||||||
temperatures=draft_block.temperatures,
|
temperatures=draft_block.temperatures,
|
||||||
rows_per_request=gamma_rows,
|
rows_per_request=gamma_rows,
|
||||||
).view(bs, gamma_rows, vocab)
|
).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:
|
if not sampling_info.is_any_greedy:
|
||||||
return AcceptSampling.execute(
|
return AcceptSampling.execute(
|
||||||
candidates=candidates,
|
candidates=candidates,
|
||||||
|
|||||||
@@ -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
|
||||||
@@ -0,0 +1,155 @@
|
|||||||
|
"""Unit tests for srt/utils/invariants.py -- CPU, no server.
|
||||||
|
|
||||||
|
Covers the (bucket x level) crash matrix, the unconditional data layer, the
|
||||||
|
throttled reporter, and self-registration / injection-coverage.
|
||||||
|
"""
|
||||||
|
|
||||||
|
import unittest
|
||||||
|
from unittest import mock
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from sglang.srt.environ import InvariantCheckLevel, envs
|
||||||
|
from sglang.srt.utils import invariants as ic
|
||||||
|
from sglang.srt.utils.invariants import (
|
||||||
|
Bucket,
|
||||||
|
Finite,
|
||||||
|
InRange,
|
||||||
|
Invariant,
|
||||||
|
expect,
|
||||||
|
registered_invariants,
|
||||||
|
)
|
||||||
|
from sglang.test.ci.ci_register import register_cpu_ci
|
||||||
|
from sglang.test.test_utils import CustomTestCase
|
||||||
|
|
||||||
|
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||||
|
|
||||||
|
|
||||||
|
# "test." namespace so the coverage meta-test isolates these from real invariants.
|
||||||
|
_SOFTEN = Invariant("test.soften", Bucket.SOFTEN, Finite(), recover=torch.nan_to_num)
|
||||||
|
_GUARD = Invariant("test.guard", Bucket.GUARD, Finite(), recover=torch.nan_to_num)
|
||||||
|
_FATAL_C = Invariant("test.fatal_containable", Bucket.FATAL_CONTAINABLE, InRange(0, 10))
|
||||||
|
_FATAL_U = Invariant(
|
||||||
|
"test.fatal_uncontainable", Bucket.FATAL_UNCONTAINABLE, InRange(0, 10)
|
||||||
|
)
|
||||||
|
|
||||||
|
# Injection manifest: every "test." invariant must have a test that triggers it.
|
||||||
|
_INJECTION_COVERAGE = {
|
||||||
|
"test.soften": "test_soften_never_crashes",
|
||||||
|
"test.guard": "test_guard_crashes_in_strict",
|
||||||
|
"test.fatal_containable": "test_fatal_containable_crashes_in_strict",
|
||||||
|
"test.fatal_uncontainable": "test_fatal_uncontainable_crashes_in_warn",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
class TestInvariants(CustomTestCase):
|
||||||
|
def _strict(self):
|
||||||
|
return envs.SGLANG_INVARIANT_CHECK.override(int(InvariantCheckLevel.STRICT))
|
||||||
|
|
||||||
|
def _warn(self):
|
||||||
|
return envs.SGLANG_INVARIANT_CHECK.override(int(InvariantCheckLevel.WARN))
|
||||||
|
|
||||||
|
def _off(self):
|
||||||
|
return envs.SGLANG_INVARIANT_CHECK.override(int(InvariantCheckLevel.OFF))
|
||||||
|
|
||||||
|
def test_crash_matrix(self):
|
||||||
|
L = InvariantCheckLevel
|
||||||
|
expected = {
|
||||||
|
(Bucket.SOFTEN, L.OFF): False,
|
||||||
|
(Bucket.SOFTEN, L.WARN): False,
|
||||||
|
(Bucket.SOFTEN, L.STRICT): False,
|
||||||
|
(Bucket.GUARD, L.OFF): False,
|
||||||
|
(Bucket.GUARD, L.WARN): False,
|
||||||
|
(Bucket.GUARD, L.STRICT): True,
|
||||||
|
(Bucket.FATAL_CONTAINABLE, L.OFF): False,
|
||||||
|
(Bucket.FATAL_CONTAINABLE, L.WARN): False,
|
||||||
|
(Bucket.FATAL_CONTAINABLE, L.STRICT): True,
|
||||||
|
(Bucket.FATAL_UNCONTAINABLE, L.OFF): False,
|
||||||
|
(Bucket.FATAL_UNCONTAINABLE, L.WARN): True,
|
||||||
|
(Bucket.FATAL_UNCONTAINABLE, L.STRICT): True,
|
||||||
|
}
|
||||||
|
for (bucket, level), want in expected.items():
|
||||||
|
self.assertEqual(ic._crashes(bucket, level), want, f"{bucket} @ {level}")
|
||||||
|
|
||||||
|
def test_recover_applied_even_when_off(self):
|
||||||
|
bad = torch.tensor([[float("nan"), 1.0]])
|
||||||
|
with self._off():
|
||||||
|
with mock.patch.object(torch, "_assert_async") as m:
|
||||||
|
out = expect(_GUARD, bad.clone())
|
||||||
|
self.assertTrue(torch.isfinite(out).all())
|
||||||
|
m.assert_not_called()
|
||||||
|
|
||||||
|
def test_soften_never_crashes(self):
|
||||||
|
bad = torch.tensor([[float("nan"), 1.0]])
|
||||||
|
with self._strict():
|
||||||
|
with mock.patch.object(torch, "_assert_async") as m:
|
||||||
|
out = expect(_SOFTEN, bad.clone())
|
||||||
|
m.assert_not_called()
|
||||||
|
self.assertTrue(torch.isfinite(out).all())
|
||||||
|
|
||||||
|
def test_guard_crashes_in_strict(self):
|
||||||
|
bad = torch.tensor([[float("nan"), 1.0]])
|
||||||
|
with self._strict():
|
||||||
|
with mock.patch.object(torch, "_assert_async") as m:
|
||||||
|
expect(_GUARD, bad.clone())
|
||||||
|
m.assert_called_once()
|
||||||
|
(cond, _msg), _ = m.call_args
|
||||||
|
self.assertFalse(bool(cond))
|
||||||
|
|
||||||
|
def test_fatal_containable_crashes_in_strict(self):
|
||||||
|
bad = torch.tensor([5, 20], dtype=torch.int64) # 20 is out of [0, 10)
|
||||||
|
with self._strict():
|
||||||
|
with mock.patch.object(torch, "_assert_async") as m:
|
||||||
|
expect(_FATAL_C, bad)
|
||||||
|
m.assert_called_once()
|
||||||
|
|
||||||
|
def test_fatal_uncontainable_crashes_in_warn(self):
|
||||||
|
bad = torch.tensor([5, 20], dtype=torch.int64) # 20 is out of [0, 10)
|
||||||
|
with self._warn():
|
||||||
|
with mock.patch.object(torch, "_assert_async") as m:
|
||||||
|
expect(_FATAL_U, bad)
|
||||||
|
m.assert_called_once()
|
||||||
|
|
||||||
|
def test_warn_counts_and_logs(self):
|
||||||
|
bad = torch.tensor([[float("nan"), 1.0]])
|
||||||
|
with self._warn():
|
||||||
|
with mock.patch.object(torch, "_assert_async") as m:
|
||||||
|
with self.assertLogs(ic.logger, level="WARNING") as cap:
|
||||||
|
expect(_GUARD, bad.clone())
|
||||||
|
m.assert_not_called()
|
||||||
|
self.assertTrue(any("test.guard" in line for line in cap.output))
|
||||||
|
|
||||||
|
def test_registry_and_validation(self):
|
||||||
|
reg = registered_invariants()
|
||||||
|
for name in _INJECTION_COVERAGE:
|
||||||
|
self.assertIn(name, reg)
|
||||||
|
|
||||||
|
with self.assertRaises(ValueError): # duplicate name
|
||||||
|
Invariant("test.guard", Bucket.GUARD, Finite(), recover=torch.nan_to_num)
|
||||||
|
|
||||||
|
with self.assertRaises(ValueError): # uncontainable FATAL forbids recover
|
||||||
|
Invariant(
|
||||||
|
"test.bad",
|
||||||
|
Bucket.FATAL_UNCONTAINABLE,
|
||||||
|
Finite(),
|
||||||
|
recover=torch.nan_to_num,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_injection_coverage_meta(self):
|
||||||
|
"""Every registered test-namespace invariant must have an injection test.
|
||||||
|
|
||||||
|
A later change broadens this filter to all namespaces so no real
|
||||||
|
GUARD/SOFTEN ships without a triggering test.
|
||||||
|
"""
|
||||||
|
uncovered = [
|
||||||
|
name
|
||||||
|
for name in registered_invariants()
|
||||||
|
if name.startswith("test.") and name not in _INJECTION_COVERAGE
|
||||||
|
]
|
||||||
|
self.assertEqual(
|
||||||
|
uncovered, [], f"invariants without an injection test: {uncovered}"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
unittest.main()
|
||||||
Reference in New Issue
Block a user