sgl-project/sglang · error · ValueError

Input probs contains NaN.

Error message

Input probs contains NaN.

What it means

The MUSA top-p sampling kernel wrapper optionally checks the probability tensor for NaN values before sampling. When check_nan=True and any element of probs is NaN, it raises immediately, because sampling from NaN probabilities yields garbage tokens.

Source

Thrown at python/sglang/kernels/aot/python/sgl_kernel/musa.py:248

        indices,
        maybe_top_p_arr,
        top_p_val,
        deterministic,
        generator,
    )
    return samples


def top_p_sampling_from_probs(
    probs: torch.Tensor,
    top_p: Union[torch.Tensor, float],
    indices: Optional[torch.Tensor] = None,
    deterministic: bool = True,
    generator: Optional[torch.Generator] = None,
    check_nan: bool = False,
) -> torch.Tensor:
    if check_nan and torch.any(torch.isnan(probs)):
        raise ValueError("Input probs contains NaN.")
    return _top_p_sampling_from_probs_internal(
        probs, indices, *_to_tensor_scalar_tuple(top_p), deterministic, generator
    )


def _top_k_top_p_sampling_from_probs_internal(
    probs: torch.Tensor,
    indices: Optional[torch.Tensor],
    maybe_top_k_arr: Optional[torch.Tensor],
    top_k_val: int,
    maybe_top_p_arr: Optional[torch.Tensor],
    top_p_val: float,
    deterministic: bool,
    generator: Optional[torch.Generator],
) -> torch.Tensor:
    device = probs.device
    probs = probs.float()
    maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None

View on GitHub (pinned to 0132848349)

Solutions

  1. Inspect logits before softmax for inf/NaN and fix masking so each row has at least one finite logit
  2. Keep check_nan=False only after verifying upstream numerics are sound; otherwise leave it on in tests/debug
  3. Clamp or normalize logits (e.g. subtract max, use softmax with dtype float32) before converting to probs

Example fix

# before
sampled = top_p_sampling_from_probs(probs, top_p, check_nan=True)
# after
assert not torch.isnan(logits).any()
probs = torch.softmax(logits.float(), dim=-1)
sampled = top_p_sampling_from_probs(probs, top_p, check_nan=True)
Defensive patterns

Strategy: validation

Validate before calling

if check_nan and torch.isnan(probs).any(): raise RuntimeError('NaN probs from upstream logits')

Type guard

def probs_finite(p): return bool(torch.isfinite(p).all())

Try / catch

except ValueError as e: if 'NaN' in str(e): log offending rows and fix logits

Prevention

When it happens

Trigger: Calling top_p_sampling_from_probs(probs, top_p, check_nan=True) where torch.any(torch.isnan(probs)) is true; also hit indirectly via top_k_top_p_sampling_from_probs with top_k_first order.

Common situations: Upstream numerics producing NaN logits (inf - inf, overflow in bf16 softmax, masked-out rows of all -inf), which the check converts into a loud failure instead of silent garbage.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/6e8983290751ab1c. Report an issue: GitHub.