From ba42d161f3e1b686ef6fe5cea86efbbb363fb0f1 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Sat, 28 Mar 2026 19:37:10 -0400 Subject: [PATCH] [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) Signed-off-by: Tyler Michael Smith --- .../layers/quantization/modelopt.py | 17 +++++++++++++++++ .../layers/quantization/utils/nvfp4_utils.py | 19 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 78644f74d28..83053e0d8e1 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/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, diff --git a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py b/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py index bcb4769e4c9..95767e8c23d 100644 --- a/vllm/model_executor/layers/quantization/utils/nvfp4_utils.py +++ b/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