Fix fp16 NaN flake in spec CI: bf16 eagle fixture; sanitize NaN logits in sampler (#27883)

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)

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(

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:

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."""

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