Add three check_tensor checkpoints to the NVFP4 GEMM path:
- fp4_input: activations before FP4 quantization
- fp4_act_scales: activation block scales after quantization
- fp4_gemm_output: GEMM output before bias/reshape
Registered per-layer in ModelOptNvFp4LinearMethod.process_weights_after_loading
so each linear layer gets its own named checkpoints (e.g.,
"model.layers.1.self_attn.o_proj.fp4_gemm_output").
Also removes the post-crash KV cache full scan (was already
removed in nan_detector.py, this syncs the state).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Remove _check_all_kv_cache() — scanning the KV cache after the
forward is circular (the forward just wrote NaN into it). Keep
the on-assignment check (check_kv_blocks) which catches stale NaN
in recycled blocks before they're used.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Add check_tensor() to NaNDetector for checking arbitrary tensors at
any point in the forward pass. Uses torch.isfinite() — all ops stay
on GPU, CUDA-graph compatible, writes to the same per-token flag
array as the RMSNorm kernel checks.
Any module can register checkpoints via register() and call
check_tensor(tensor, idx) in its forward. update_layer_names()
picks up _nan_detect_indices dicts for readable names.
Wire into DeepseekV2Attention to check attn_output before o_proj,
distinguishing "attention produced NaN" from "o_proj produced NaN".
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
When NaN/Inf is detected in real tokens, scan all KV cache blocks
and log which ones contain NaN before raising RuntimeError. This
helps distinguish "stale NaN in recycled cache block" from "compute
produced NaN" without needing to reproduce the issue.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Raise RuntimeError when NaN/Inf is detected in real tokens during
the forward pass. KV cache block checks remain log-only since stale
NaN in recycled blocks is diagnostic, not fatal.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
When VLLM_NAN_DETECT=1, check recycled KV cache blocks for stale
NaN/Inf before they are assigned to a new request. This catches the
"poisoned pool" scenario where one bad request leaves NaN in KV cache
blocks that then corrupt subsequent requests via attention.
Note: block zeroing (_zero_block_ids) only runs for Mamba/SSM models.
Standard attention models reuse blocks without zeroing, so stale NaN
from a previous request persists until overwritten by new KV writes.
Changes:
- Scheduler: also drain new_block_ids when VLLM_NAN_DETECT=1
- NaNDetector: add check_kv_blocks() for recycled block checking,
accept kv_caches in finalize(), handle uint8->fp8 viewing
- Model runner: call check_kv_blocks before zeroing, pass kv_caches
to finalize, add comment documenting the no-zero behavior
- Padding NaN logging: downgraded to debug level
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
torch.compile doesn't support logging methods in traced code. The
logger.warning_once() call in RMSNorm.forward_cuda() was causing
compilation failures when NaN detection is enabled.
This log message was informational only (warning about bypassing
Oink/batch-invariant paths), so removing it doesn't affect functionality.
Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
Padding rows in the swizzled scale tensor were uninitialized (torch.empty),
containing stale NaN from prior GPU allocations. The TRT-LLM mm_fp4 kernel
with use_8x4_sf_layout=True reads padding scales and applies them to real
rows, contaminating output with NaN.
Zero-filling ensures padding scales contribute 0 * data = 0.
Fixes: https://github.com/flashinfer-ai/flashinfer/issues/2861
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>
This reverts commit ef2c4f778d.
The zero-init workaround is unnecessary — the NaN was caused by a
different issue (int64 expert IDs in the routing simulator). Reverting
to restore the original torch.empty allocation which avoids the
overhead of pre-allocated zero-init buffers.
Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Reduce the test matrix by trimming:
- NUM_TOKENS_HIDDEN_SIZES: 19 → 4 combos (keep small, misaligned,
medium-aligned, large-misaligned)
- GROUP_SIZES: drop [1, 64] (redundant with [1, 128])
- Inline group_size/tma_alignment combos to 3 meaningful cases
instead of full Cartesian product (7)
Covers the same code paths (per-token, per-block, TMA alignment)
with fewer redundant combinations.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Reduce the test matrix from 864 to 216 cases (~4x speedup) by trimming
redundant hidden sizes (keep 8, 769, 8192 — covers small, misaligned,
large) and token counts (keep 7, 4096 — small, large), and quant scales
(keep 0.01, 10.0 — extreme ends).
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Add per-token NaN/Inf detection to all RMSNorm CUDA kernels by
piggybacking on the existing variance reduction. NaN/Inf propagates
through sum-of-squares naturally, so a single isnan(variance) ||
isinf(variance) check on thread 0 after the CUB reduction detects
it -- zero additional kernel launches, memory reads, or register
pressure. Each block writes to its own int8 flag slot (no atomics).
Controlled by VLLM_NAN_DETECT=1. When enabled:
- Bypasses Oink/batch-invariant paths (with warning)
- Reports per-token, per-layer NaN/Inf with layer names
- Distinguishes real-token errors from padding-token warnings
- CUDA-graph compatible (fixed flag buffer address)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Gate behind VLLM_NIXL_NAN_DETECT=1 env var. Checks KV cache blocks
for NaN on the decoder side after recv completes and on the prefiller
side after send is confirmed. Handles uint8-stored fp8 KV caches
(MLA cross-layer) by viewing as float8_e4m3fn before isnan check.
Uses a fast two-pass approach: single torch.isnan().any() across all
layers first, only doing per-layer breakdown if NaN is found.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
