[Kernel] Add NaN/Inf checks to NVFP4 linear pipeline

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>

vllm/model_executor/layers/quantization/modelopt.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any
 import torch
 from torch.nn.parameter import Parameter
 
+import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
@@ -1167,6 +1168,22 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
         # Convert layer to NVFP4 linear kernel format
         convert_to_nvfp4_linear_kernel_format(self.backend, layer)
 
+        # Register NaN detection checkpoints for FP4 pipeline.
+        if envs.VLLM_NAN_DETECT:
+            from vllm.model_executor.layers.nan_detector import NaNDetector
+
+            detector = NaNDetector.get()
+            prefix = getattr(layer, "prefix", "nvfp4")
+            layer._nan_detect_indices = {
+                "fp4_input": detector.register(f"{prefix}.fp4_input"),
+                "fp4_act_scales": detector.register(
+                    f"{prefix}.fp4_act_scales"
+                ),
+                "fp4_gemm_output": detector.register(
+                    f"{prefix}.fp4_gemm_output"
+                ),
+            }
+
     def apply(
         self,
         layer: torch.nn.Module,

vllm/model_executor/layers/quantization/utils/nvfp4_utils.py

@@ -216,11 +216,26 @@ def apply_nvfp4_linear(
     output_dtype = x.dtype
     output_shape = [*x.shape[:-1], output_size]
 
+    # NaN/Inf checks: input to FP4 quantization
+    nan_check = envs.VLLM_NAN_DETECT and hasattr(layer, "_nan_detect_indices")
+    if nan_check:
+        from vllm.model_executor.layers.nan_detector import NaNDetector
+
+        detector = NaNDetector.get()
+        ndi = layer._nan_detect_indices
+        detector.check_tensor(x, ndi["fp4_input"])
+
     # Quantize BF16 or FP16 to (FP4 and interleaved block scale)
     x_fp4, x_blockscale = scaled_fp4_quant(
         x, input_global_scale_inv, is_sf_swizzled_layout=True, backend=backend.value
     )
 
+    # NaN/Inf checks: FP4 quantized activation scales
+    if nan_check:
+        detector.check_tensor(
+            x_blockscale.view(x.shape[0], -1), ndi["fp4_act_scales"]
+        )
+
     # Validate dtypes
     assert x_fp4.dtype == torch.uint8
     assert weight.dtype == torch.uint8
@@ -263,6 +278,10 @@ def apply_nvfp4_linear(
     # Slice output to remove N-dimension padding
     out = slice_nvfp4_output(out, output_size)
 
+    # NaN/Inf checks: GEMM output
+    if nan_check:
+        detector.check_tensor(out, ndi["fp4_gemm_output"])
+
     if bias is not None:
         out = out + bias