[Core] Add --enable-nan-fault-tolerance for NaN detection and request abort by elvircrn · Pull Request #50283 · vllm-project/vllm

elvircrn · 2026-07-29T14:45:15Z

Overview

Add --enable-nan-fault-tolerance for MRV2: abort on NaN logits and withhold the unchecked chunk from GPU prefix reuse and CPU offloading. Earlier validated prefixes stay reusable.

This consolidates fork PR #5 into this PR. The other open NaN fixes address individual kernels; this adds engine-level recovery.

Claims

--enable-detect-nans-in-logits exposes detection as an engine argument, retaining the legacy environment variable. Fault tolerance implies detection and requires MRV2.
Check intermediate prefill, final prefill, and decode logits before publishing new KV. Abort the failing request without evicting earlier validated chunks.
Delay OffloadingConnector and eager SimpleCPUOffloadConnector stores until validation. SimpleCPU lazy mode already selects published GPU blocks. No additional GPU-to-CPU status readbacks or KV scans.
Zero newly allocated GPU blocks on reuse and prevent async hybrid checkpoints from being published ahead of validation.

Validation

Rebased onto upstream main 0c2584929. Preserves upstream speculative-stop accounting and offload-manager locking; NaN metrics work with engine arguments even when the legacy environment flag is off.

CPU-only scheduler regressions cover sync/async prefill, first/later chunk failures, eager/lazy offloading, late copy acknowledgements, clean completion, and hybrid checkpoint publication. The new regressions fail on the pre-fix head (ab687e51e, 20 failures).

739 passed, 2 skipped across the affected scheduler, cache-manager, offload-worker, CPU-manager, and tier-manager suites. All applicable pre-commit hooks passed, including Ruff, mypy, and CI test tethering. Tests used a CPU-only harness to permit MRV2 configuration without CUDA; no model or GPU kernels ran.

CPU test coverage

tests/v1/simple_kv_offload/test_scheduler.py
tests/v1/core/test_single_type_kv_cache_manager.py
tests/v1/kv_connector/unit/offloading_connector/{test_scheduler,test_worker}.py
tests/v1/kv_offload/{cpu/test_manager,tiering/test_tiering_offloading}.py
tests/v1/metrics/test_stats.py
tests/v1/core/{test_scheduler,test_async_scheduler}.py and tests/v1/streaming_input/test_scheduler_streaming.py, filtered to NaN, stop, local-prefill, and throttle checks

The CUDA suite now includes eager/lazy SimpleCPU offloading. The updated CUDA end-to-end cases, VLLM_GPU_SYNC_CHECK=error, and model evaluations have not been run on this macOS host. Previous GPU results do not validate this update.

VLLM_USE_V2_MODEL_RUNNER=1 VLLM_GPU_SYNC_CHECK=error \
  CUDA_VISIBLE_DEVICES=0 .venv/bin/python -m pytest -sv \
  tests/v1/e2e/general/test_nan_fault_tolerance.py

Details

Scheduling advances computed-token counts before outputs arrive. Stores now stop at the completed-token boundary; GPU publication already waits for the logits check. A failed chunk stays unpublished, while earlier validated blocks can finish offloading and remain cache hits. Retries recompute the missing suffix normally.

Known limitations

Logits only: finite logits do not prove KV is uncorrupted; earlier cached prefixes are retained, not retrospectively scrubbed.
External KV: NIXL and other connectors are not covered. Existing corrupt CPU/disk entries are not invalidated. Disk stores downstream of these offload gates inherit the same admission boundary.
Custom samplers: must supply per-request NaN counts.
Cost: intermediate-prefill logits, NaN checks, block zeroing, and delayed cache availability add overhead. This update has not been benchmarked.

Pull Request Checklist

I used vLLM's /pr-checklist skill.
AI assistance was used: Claude Opus 4.8 and OpenAI Codex.
Design Fit: Reuses completion counters and cache publication; no new worker protocol or cache reset.
Testing and Validation: CPU regressions pass; updated GPU/model validation and human review remain required.
Code Quality and Style: Focused changes and existing test suites.
Pull Request Contents: Scope, evidence, and limitations updated.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

claude

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

fangyuchu · 2026-08-03T04:07:41Z

The overall design makes sense to me. I have two questions:

1. A (small) race window between schedule-time prefix-cache commit and output-time NaN detection

Full blocks are committed to the prefix cache at schedule time — before the forward writes their KV — while NaN detection only happens when that step's output is processed. With async scheduling, this interleaving is possible:

schedule(N): A will fill block X this step → X becomes matchable in the prefix cache (KV not yet written).
forward(N): A writes NaN KV into X.
schedule(N+1): B gets a prefix-cache hit on X — output(N) hasn't been processed yet, so nobody knows X is corrupt.
forward(N+1): B reads the NaN KV. Cache-hit blocks are never zeroed (they're not new allocations), so the zero-on-reallocation defense doesn't apply on this path.

My understanding is that this is largely self-containing: B's new tokens attend to every position of X, so the NaN almost certainly propagates into B's own logits. But could there be cases where B's KV cache gets corrupted while no NaN ever shows up in its logits?

2. Could the KV-zeroing logic be abstracted for reuse by fault-tolerance recovery?

We're working on another scenario with similar problem: retry-after-failure with KV reuse, so a request can resume generation from its existing KV after a fault instead of recomputing from scratch (). One problem we're thinking about is very similar in shape to yours: when a DP/EP rank fails mid-step, the KV written for the in-flight token (e.g., the last token of the last block) can be corrupt on the surviving ranks — NaN or plain garbage — and we need to wipe precisely that KV before resuming. The fault cause differs (numerical NaN vs. rank failure), but the cleanup primitive is the same.

Your zeroing logic is close to what we'd need, but it's currently tied to the allocation-time path (whole blocks, driven by new_block_ids_to_zero in SchedulerOutput). Do you think the zeroing mechanism could be factored into a reusable, on-demand function — e.g., zero_last_block_kv(self,request) — so that a fault-recovery path can call it to thoroughly zero the KV of the corrupted token(s) after a fault is detected, whatever the cause?

elvircrn · 2026-08-03T09:16:01Z

@fangyuchu Thanks for the feedback!

But could there be cases where B's KV cache gets corrupted while no NaN ever shows up in its logits?

In principle this should not be possible under normal circumstances but we did have this bug where, during a dummy run, we were accidentally writing NaNs into the KV$. See PR more info but specifically into slot 0. With slot 0 being the sentinel slot, later requests didn't have NaNs in their logits immediately.

As long as the request eventually picks up a NaN-containing block, NaNs should make their way into the logits.

Therefore I don't think this race should be a cause for concern unless there is a way to avoid it without paying a large performance penalty (is there?).

Ultimately, one could disable async scheduling in cases where the logic in the PR isn't able to recover model accuracy.

Your zeroing logic is close to what we'd need, but it's currently tied to the allocation-time path (whole blocks, driven by new_block_ids_to_zero in SchedulerOutput). Do you think the zeroing mechanism could be factored into a reusable, on-demand function — e.g., zero_last_block_kv(self,request) — so that a fault-recovery path can call it to thoroughly zero the KV of the corrupted token(s) after a fault is detected, whatever the cause?

Maybe we do this in a follow-up PR? I'd be happy to work on this.

tlrmchlsmth · 2026-08-03T19:48:32Z

+    enable_nan_fault_tolerance: bool = False
+    """When enabled, requests producing NaN logits are immediately aborted and
+    KV cache blocks are zeroed on reuse to prevent corruption from propagating
+    into prefix cache, CPU/disk offloading tiers, or disaggregated decode
+    instances. Implies --enable-detect-nans-in-logits."""


Move this to config/fault_tolerance.py?

njhill

Thanks @elvircrn it looks pretty clean overall

njhill · 2026-08-03T23:08:22Z

@@ -1184,7 +1183,7 @@ def record(

        if iteration_stats is None:
            return
-        if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
+        if self.vllm_config.observability_config.enable_detect_nans_in_logits:


nit: could make this flag a dedicated field

njhill · 2026-08-03T23:09:02Z

@@ -1360,6 +1360,7 @@ def get_kv_cache_config_from_groups(
            num_blocks=1,
            kv_cache_tensors=[],
            kv_cache_groups=kv_cache_groups,
+            force_zeroing=(vllm_config.observability_config.enable_nan_fault_tolerance),


nit: redundant parens

njhill · 2026-08-03T23:09:05Z

@@ -1419,6 +1420,7 @@ def get_kv_cache_config_from_groups(
        num_blocks=num_blocks,
        kv_cache_tensors=kv_cache_tensors,
        kv_cache_groups=kv_cache_groups,
+        force_zeroing=(vllm_config.observability_config.enable_nan_fault_tolerance),


nit: redundant parens

njhill · 2026-08-03T23:11:12Z

+                    # Collect this request's blocks so they can be evicted from
+                    # the prefix cache below. Prefill commits full blocks to the
+                    # prefix cache at schedule time (before the forward writes
+                    # KV), so a block holding NaN KV may already be matchable;
+                    # the plain free() path leaves it in the cache map.


Maybe reduce the size of the comment here

njhill · 2026-08-03T23:11:54Z

+        # Evict NaN-corrupted blocks from the prefix cache before freeing the
+        # aborted requests, so a later request with the same prefix cannot match
+        # (and read) the NaN KV as a cache hit. Eviction resets the block hash,
+        # which also makes free() prioritize these blocks for reallocation, where
+        # the existing zero-on-allocation path wipes their contents. This mirrors
+        # the KV-load-failure recovery path, which evicts corrupt blocks for the
+        # same reason.


Maybe trim this block of text down a bit

yea needs a human touch :)

njhill · 2026-08-03T23:12:30Z

@@ -1723,6 +1713,8 @@ def update_from_output(
        # to avoid expensive operations inside the loop.
        stopped_running_reqs: set[Request] = set()
        stopped_preempted_reqs: set[Request] = set()
+        nan_abort_req_ids: set[str] = set()
+        nan_block_ids: set[int] = set()


clearer?

Suggested change

nan_block_ids: set[int] = set()

nan_block_ids_to_evict: set[int] = set()

… abort Add two CLI flags for NaN corruption handling: - --enable-detect-nans-in-logits: enables NaN counting in logits per step (replaces VLLM_COMPUTE_NANS_IN_LOGITS env var with a proper CLI flag while keeping backward compat) - --enable-nan-fault-tolerance: implies the above, and additionally aborts requests when NaN logits are detected (during chunked prefill or decode) and forces KV cache block zeroing on reuse to prevent corruption from propagating into prefix cache, CPU/disk offloading tiers, or disaggregated decode instances via NIXL Signed-off-by: Elvir Crncevic <elvir@crncevic.dev> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

The NaN detection path was only wired into the legacy (V1) GPU model runner. The default V2 runner (vllm/v1/worker/gpu/model_runner.py) constructed its Sampler without detect_nans_in_logits, so num_nans was never computed and num_nans_in_logits never populated on the model runner output. The async output plumbing already forwards num_nans to ModelRunnerOutput; only the sampler flag was missing. Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

When a request is aborted for NaN in logits, its KV blocks are freed via the normal path, which returns them to the pool but leaves them in the prefix-cache hash map. Prefill commits full blocks to the prefix cache at schedule time (before the forward writes KV), so a block holding NaN KV can remain matchable and be reused as a prefix-cache hit by a later request with the same prefix -- bypassing the zero-on-reuse path, since cache hits are never re-zeroed. Evict the aborted request's blocks from the prefix cache before freeing them, mirroring the KV-load-failure recovery path. Eviction resets the block hash so free() also prioritizes the blocks for reallocation, where the existing zero-on-allocation logic wipes their contents. Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

…oleranceConfig - Move enable_nan_fault_tolerance from ObservabilityConfig to FaultToleranceConfig - Rename nan_block_ids -> nan_block_ids_to_evict in scheduler - Trim verbose comments in scheduler NaN abort/eviction paths - Cache detect_nans_in_logits as dedicated field in stat loggers - Remove redundant parens in kv_cache_utils force_zeroing - Update tests to use new config path Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Delay prefix-cache publication until asynchronous outputs have been validated, aggregate speculative NaN counts per request without host synchronization, and preserve corruption metrics on abort. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Cover sync and async scheduling across eager and CUDA graph execution, plus multi-row MTP verification, request abort, prefix-cache isolation, recovery, and corruption metrics on one small Qwen checkpoint. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Co-authored-by: Codex <codex@openai.com> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Assisted-by: Codex Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Initialize aggregation buffers on GPU and leave MRV1 unchanged. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Reduce per-row counts per request in one kernel and cover empty, long, and strided ranges. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Add short examples for regular and empty request row ranges. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Read fault tolerance settings from ParallelConfig. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Check replay size against the runner decode length, including speculative decoding. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Retain earlier validated GPU prefixes when a later chunk fails. Delay OffloadingConnector stores and SimpleCPU eager cursors until output validation, and keep async Mamba checkpoints behind the same boundary. Consolidate the offload fix from fork PR #5 into vllm-project#50283. Cover sync/async eager/lazy stores, late acknowledgements, and checkpoint publication. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Verify the upstream rebase keeps CLI-enabled NaN detection visible in request metrics and does not clear corruption on a later clean output. Co-authored-by: OpenAI Codex <noreply@openai.com> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

mergify · 2026-10-17T10:21:08Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @elvircrn.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

elvircrn requested review from ApostaC, ProExpertProg, WoosukKwon, alexm-redhat, heheda12345, hmellor, houseroad, ivanium, mgoin, njhill, orozery, robertgshaw2-redhat, tlrmchlsmth, yewentao256, youkaichao and ywang96 as code owners July 29, 2026 14:45

claude Bot reviewed Jul 29, 2026

View reviewed changes

mergify Bot added the v1 label Jul 29, 2026

elvircrn marked this pull request as draft July 29, 2026 14:46

elvircrn force-pushed the nan-fault-tolerance branch 2 times, most recently from 5576e4f to 63faf06 Compare July 30, 2026 09:10

mergify Bot added the mrv2 Model Runner V2 specific label Jul 31, 2026

elvircrn marked this pull request as ready for review July 31, 2026 17:52

claude Bot reviewed Jul 31, 2026

View reviewed changes

depthfirst-app Bot reviewed Jul 31, 2026

View reviewed changes

Comment thread vllm/v1/core/sched/scheduler.py Outdated

tlrmchlsmth mentioned this pull request Aug 3, 2026

[Core] Zero KV cache when NaN logits detected #50002

Closed

tlrmchlsmth reviewed Aug 3, 2026

View reviewed changes

njhill reviewed Aug 3, 2026

View reviewed changes

elvircrn and others added 26 commits October 16, 2026 01:49

Add tests for NaN fault tolerance and env var deprecation

9c0d56a

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Register NaN fault tolerance CLI args in argparser

1a5e43c

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Fix chunked prefill test: max_num_seqs must be <= max_num_batched_tokens

c9e6e96

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Fix structured output handling after rebase

584f328

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Force nan counting when nan fault tolerance is enabled

e027cb0

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Refactor NaN abort handling in scheduler

7143e21

Co-authored-by: Codex <codex@openai.com> Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Make _force_zeroing not exposed; Update docs for the NaN robustness flag

f0ecf1b

Assisted-by: Codex Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Comment update

c43d749

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Address NaN fault tolerance review comments

c2a77c6

Initialize aggregation buffers on GPU and leave MRV1 unchanged. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Use Triton to aggregate speculative NaN counts

77deca8

Reduce per-row counts per request in one kernel and cover empty, long, and strided ranges. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Document Triton NaN aggregation ranges

df4ccde

Add short examples for regular and empty request row ranges. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Clarify NaN aggregation inputs in kernel docstring

33c5f65

Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Fix NaN E2E config assertion after rebase

67e786e

Read fault tolerance settings from ParallelConfig. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Use request cache stats in NaN E2E test

e4ace85

Check replay size against the runner decode length, including speculative decoding. Assisted-by: Codex (OpenAI) Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Size NaN cache recovery prompt for hybrid KV blocks

61ad91a

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Cover MTP prefix cache block drop in NaN test

b728eee

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Allow two MTP cache blocks in NaN recovery test

5c92052

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

Require Model Runner V2 for NaN fault tolerance

172c8a6

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com>

elvircrn force-pushed the nan-fault-tolerance branch from e7d79bd to d51247d Compare October 15, 2026 23:56

mergify Bot removed the needs-rebase label Oct 15, 2026

mergify Bot added the needs-rebase label Oct 17, 2026

	nan_block_ids: set[int] = set()
	nan_block_ids_to_evict: set[int] = set()

Conversation

elvircrn commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Claims

Validation

Details

Known limitations

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

Uh oh!

fangyuchu commented Aug 3, 2026

1. A (small) race window between schedule-time prefix-cache commit and output-time NaN detection

2. Could the KV-zeroing logic be abstracted for reuse by fault-tolerance recovery?

Uh oh!

elvircrn commented Aug 3, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

njhill left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Oct 17, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

elvircrn commented Jul 29, 2026 •

edited

Loading