diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 6e7a26dce..b02400e1d 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -629,6 +629,9 @@ 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) + # Sanitize NaN logits before sampling kernels and log a throttled warning + # (see sanitize_nan_logits). + SGLANG_SANITIZE_NAN_LOGITS = EnvBool(True) # VLM SGLANG_VLM_CACHE_SIZE_MB = EnvInt(100) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index 84d57b8e2..941cdb353 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -16,6 +16,7 @@ from sglang.srt.layers.utils.logprob import get_token_ids_logprobs, get_top_logp from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo from sglang.srt.sampling.sampling_params import TOP_K_ALL from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils.async_probe import sanitize_nan_logits from sglang.srt.utils.common import ( get_bool_env_var, is_cuda, @@ -83,9 +84,10 @@ class Sampler(nn.Module): def _preprocess_logits( self, logits: torch.Tensor, sampling_info: SamplingBatchInfo ) -> torch.Tensor: - """Apply custom logit processors.""" + """Apply custom logit processors and sanitize non-finite logits.""" if sampling_info.has_custom_logit_processor: apply_custom_logit_processor(logits, sampling_info) + sanitize_nan_logits(logits, "sampler: next_token_logits") return logits def forward( diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index fb410f76e..abcc63fad 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -45,7 +45,11 @@ from sglang.srt.speculative.triton_ops.cache_locs import ( from sglang.srt.speculative.triton_ops.eagle import ( fill_bonus_tokens as fill_bonus_tokens, ) -from sglang.srt.utils.async_probe import maybe_detect_nan, maybe_detect_oob +from sglang.srt.utils.async_probe import ( + maybe_detect_nan, + maybe_detect_oob, + sanitize_nan_logits, +) from sglang.srt.utils.common import is_cuda, is_hip, is_musa, is_npu _is_cuda = is_cuda() @@ -474,6 +478,8 @@ class EagleVerifyInputV2Mixin: sampling_info = batch.sampling_info next_token_logits = logits_output.next_token_logits + sanitize_nan_logits(next_token_logits, "verify: target model logits") + # Apply penalty # This is a relaxed version of penalties for speculative decoding. if sampling_info.acc_additive_penalties is not None: diff --git a/python/sglang/srt/utils/async_probe.py b/python/sglang/srt/utils/async_probe.py index 84f510e45..b511476d5 100644 --- a/python/sglang/srt/utils/async_probe.py +++ b/python/sglang/srt/utils/async_probe.py @@ -5,12 +5,76 @@ When the gate is on, a violation surfaces as an assertion at the next CUDA sync point instead of as a silent NaN cascade or illegal-address crash. """ +import logging from typing import Optional import torch from sglang.srt.environ import envs +logger = logging.getLogger(__name__) + + +class _AsyncNanWarner: + """One-shot NaN monitor: device-side detection lands in pinned host + memory without any stream sync; the host reads the (slightly stale) flag + on a later call, warns once, and stops detecting.""" + + def __init__(self): + self._dev = None + self._host = None + self._warned = False + + def check(self, tensor: torch.Tensor, msg: str): + if self._warned or not tensor.is_cuda: + return + if self._dev is None: + self._dev = torch.zeros(1, dtype=torch.int32, device=tensor.device) + self._host = torch.zeros(1, dtype=torch.int32, pin_memory=True) + + # Report a hit enqueued on an earlier step (pinned read, no sync). + if int(self._host[0]): + logger.warning( + "NaN detected in %s; values were sanitized before sampling. " + "This usually indicates numerical overflow (e.g. fp16 " + "activations) or an upstream bug producing NaN. " + "Logged once; further occurrences are silent.", + msg, + ) + self._warned = True + return + + # Enqueue this step's detection (async, no sync). + self._dev.add_(torch.isnan(tensor).any().to(torch.int32)) + self._host.copy_(self._dev, non_blocking=True) + + +_nan_warner = _AsyncNanWarner() + + +def maybe_warn_nan(tensor: Optional[torch.Tensor], msg: str = ""): + """Non-fatal counterpart of maybe_detect_nan: throttled sync-free warning + instead of crashing. Callers sanitize the tensor themselves.""" + if envs.SGLANG_ENABLE_ASYNC_ASSERT.get(): + # The hard assert path already covers detection. + return + if tensor is None: + return + _nan_warner.check(tensor, msg) + + +def sanitize_nan_logits(logits: torch.Tensor, msg: str = ""): + """Detect NaN (assert in CI, throttled warning in prod), then sanitize in + place: NaN logits (e.g. fp16 activation overflow) are undefined behavior + in sampling kernels and can come back as out-of-vocab token ids. +-1e30 + rather than dtype min/max because callers divide logits by temperature, + which would overflow dtype min/max to +-Inf and softmax back to NaN.""" + maybe_detect_nan(logits, msg) + if not envs.SGLANG_SANITIZE_NAN_LOGITS.get(): + return + maybe_warn_nan(logits, msg) + torch.nan_to_num_(logits, nan=-1e30, posinf=1e30, neginf=-1e30) + def maybe_detect_nan(tensor: Optional[torch.Tensor], msg: str = ""): """Async NaN check — no GPU-CPU sync, error surfaces at next sync point.""" diff --git a/python/sglang/test/server_fixtures/spec_eagle_fixture.py b/python/sglang/test/server_fixtures/spec_eagle_fixture.py index f4041f761..fe3131753 100644 --- a/python/sglang/test/server_fixtures/spec_eagle_fixture.py +++ b/python/sglang/test/server_fixtures/spec_eagle_fixture.py @@ -61,7 +61,9 @@ class SpecEagleServerBase(CustomTestCase): mem_fraction_static = 0.75 max_running_requests = 8 chunked_prefill_size = 128 - dtype = "float16" + # bf16 rather than fp16: fp16 activations can overflow (-> Inf -> NaN) on + # degenerate draft branches in verify and trip the CI NaN asserts. + dtype = "bfloat16" cuda_graph_max_bs = None trust_remote_code = True