From f408ad2b734cc62e0782a02ecf2202ca9d9aa4f2 Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Sat, 28 Mar 2026 19:54:50 -0400 Subject: [PATCH] [Bugfix] Make check_tensor fullgraph-safe and support FP8 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove @torch.compiler.disable — fullgraph=True rejects it. Instead, inline the check in check_tensor() directly. All ops (view, to, isfinite, any, bitwise_or_) are traceable by dynamo. FP8 tensors are cast to float16 before torch.isfinite since isfinite doesn't support Float8 dtypes. Co-Authored-By: Claude Opus 4.6 (1M context) Signed-off-by: Tyler Michael Smith --- vllm/model_executor/layers/nan_detector.py | 35 ++++++++-------------- 1 file changed, 13 insertions(+), 22 deletions(-) diff --git a/vllm/model_executor/layers/nan_detector.py b/vllm/model_executor/layers/nan_detector.py index af51d910f3f..aebea3e115b 100644 --- a/vllm/model_executor/layers/nan_detector.py +++ b/vllm/model_executor/layers/nan_detector.py @@ -161,36 +161,27 @@ class NaNDetector: """Check *tensor* for NaN/Inf, writing per-token flags. Uses ``torch.isfinite`` -- all ops stay on GPU, no D2H sync. - CUDA-graph compatible (fixed output address). + CUDA-graph compatible and fullgraph=True safe. + + FP8 tensors are cast to float16 before checking since + ``torch.isfinite`` doesn't support FP8 dtypes. Args: - tensor: 2-D ``[num_tokens, hidden_size]`` tensor to check. + tensor: ``[num_tokens, ...]`` tensor to check. checkpoint_idx: index returned by :meth:`register`. """ if self._nan_flags is None: return - _check_tensor_impl( - tensor, self._nan_flags, checkpoint_idx + num_tokens = tensor.shape[0] + t = tensor.view(num_tokens, -1) + if t.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz, + torch.float8_e5m2, torch.float8_e5m2fnuz): + t = t.to(torch.float16) + has_bad = (~torch.isfinite(t)).any(dim=1) + self._nan_flags[checkpoint_idx, :num_tokens].bitwise_or_( + has_bad.to(torch.int8) ) - -@torch.compiler.disable -def _check_tensor_impl( - tensor: torch.Tensor, - nan_flags: torch.Tensor, - checkpoint_idx: int, -) -> None: - num_tokens = tensor.shape[0] - t = tensor.view(num_tokens, -1) - # torch.isfinite doesn't support FP8 — cast to float first. - if t.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz, - torch.float8_e5m2, torch.float8_e5m2fnuz): - t = t.to(torch.float16) - has_bad = (~torch.isfinite(t)).any(dim=1) - nan_flags[checkpoint_idx, :num_tokens].bitwise_or_( - has_bad.to(torch.int8) - ) - # ------------------------------------------------------------------ # Post-forward checking # ------------------------------------------------------------------