[Bugfix] Zero new KV blocks for quantized + sliding-window hybrid caches by EdalatiAli · Pull Request #47574 · vllm-project/vllm

EdalatiAli · 2026-07-03T20:31:49Z

Purpose

A hybrid KV cache that quantizes full-attention layers but keeps sliding-window layers in higher precision (--kv-cache-dtype fp8 --kv-cache-dtype-skip-layers sliding_window) produces NaN / all-zero output once the input exceeds the sliding window. Uniform fp8 and bf16 are fine; only the hybrid combo fails (seen on Blackwell + FlashInfer).

Root cause

Attention groups share one block pool. The sliding-window group frees blocks mid-request; those pages get reallocated to the fp8 group and read back including slots not written this step. As bf16 that leftover is harmless
garbage, but as fp8 it decodes to NaN/Inf and corrupts attention. The existing zeroing guard (needs_kv_cache_zeroing) was gated only on Mamba (#35219), so this config was left exposed.

Fix

Extend needs_kv_cache_zeroing to also fire for block-dropping attention (sliding-window / chunked-local) + quantized (FP8/NVFP4) KV cache:

needs_kv_cache_zeroing = has_mamba_layers or (has_block_dropping_layers and has_quantized_kv_cache)

Uniform / unquantized caches stay False → zero overhead.

Test Plan

Run RULER (niah_single_1) evals on CohereLabs/command-a-plus-05-2026-w4a4 with various kv cache quantization options, before and after the fix.
Here are the repro steps on B200, FlashInfer backend, lm-eval-harness 0.4.x.

Serve the hybrid config:

vllm serve CohereLabs/command-a-plus-05-2026-w4a4 \
  --tool-call-parser cohere_command4 --reasoning-parser cohere_command4 \
  --enable-auto-tool-choice \
  --kv-cache-dtype fp8 --kv-cache-dtype-skip-layers sliding_window

Run the needle eval (repeat with --kv-cache-dtype fp8 alone, and with no
--kv-cache-dtype flag for the bf16 reference):

lm_eval --model local-completions \
  --model_args "base_url=http://127.0.0.1:8000/v1/completions,model=CohereLabs/command-a-plus-05-2026-w4a4,num_concurrent=32,tokenized_requests=False" \
  --tasks niah_single_1 \
  --metadata '{"max_seq_lengths":[4096,32768,131072],"tokenizer":"CohereLabs/command-a-plus-05-2026-w4a4"}' \
  --gen_kwargs "temperature=0.9,top_p=0.95,repetition_penalty=1.04" \
  --trust_remote_code

Test Result

KV cache config	4K	32K	132K
bf16 (reference)	0.994	0.388	0.556
fp8 uniform	0.994	0.376	0.462
fp8 + skip-SWA w/o fix	0.002	0.000	0.002
fp8 + skip-SWA with fix	0.990	0.376	0.520

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Hybrid configs mixing a block-dropping attention type (sliding-window / chunked-local) with a quantized (FP8/NVFP4) KV cache could read stale bytes from reused blocks, decoding to NaN. Extend needs_kv_cache_zeroing to cover them. Signed-off-by: EdalatiAli <aliedalati@cohere.com> Co-authored-by: Cursor <cursoragent@cursor.com>

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.

tlrmchlsmth · 2026-07-13T20:00:41Z

+    @property
+    def has_block_dropping_layers(self) -> bool:
+        """Any group uses an attention type that frees KV blocks mid-request
+        (sliding-window or chunked-local)."""
+        return any(
+            isinstance(g.kv_cache_spec, (SlidingWindowSpec, ChunkedLocalAttentionSpec))
+            for g in self.kv_cache_groups
+        )
+
+    @property
+    def has_quantized_kv_cache(self) -> bool:
+        """Any group stores its KV cache in a quantized dtype (FP8/NVFP4)."""
+        return any(
+            isinstance(g.kv_cache_spec, AttentionSpec)
+            and g.kv_cache_spec.kv_quant_mode != KVQuantMode.NONE
+            for g in self.kv_cache_groups
+        )


I feel the condition is a little bit off. The fundamental issue is that we have multiple KV cache groups stored at different precisions, right? Any way to capture that more directly?

Agreed. Switched to has_mixed_precision_kv_cache to directly capture the case. Thank you for the comment!

NickLucche

@EdalatiAli I might be reading it wrong, but aren't we also including models that are not using different dtypes per sw/fa layers (in fact any model that uses SW and kv fp8 which is quite common)?
I believe we should only target those with hybrid accuracy.

Signed-off-by: EdalatiAli <aliedalati@cohere.com>

NickLucche

I think this is addressing a more generic issue when mixing dtypes in kv cache, which is one may argue is also what happens with mamba.
Anyways I am happy with the revised check, although we might consider merging mamba/fa condition, I think the current check form might be clearer.
Thanks @EdalatiAli !

EdalatiAli · 2026-07-20T21:18:33Z

@NickLucche Thank you for the insightful comments!

…hes (vllm-project#47574) Signed-off-by: EdalatiAli <aliedalati@cohere.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

ZJY0516 · 2026-07-25T02:55:52Z

Has any test been done for PD? For Mamba, we found that GPU zero ops and PD transfer may have a data race. @NickLucche @tlrmchlsmth

NickLucche · 2026-07-27T13:10:34Z

Syncd offline with @ZJY0516 , I believe #48481 is addressing known data races in a generic way that would include the addition here.
Happy to help out and monitor accuracy regressions regardless

…hes (vllm-project#47574) Signed-off-by: EdalatiAli <aliedalati@cohere.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>

mergify Bot added v1 bug Something isn't working labels Jul 3, 2026

EdalatiAli force-pushed the fix-skipswa-kvquant branch from 14e0b47 to a7f997a Compare July 9, 2026 15:39

EdalatiAli marked this pull request as ready for review July 9, 2026 17:25

EdalatiAli requested review from heheda12345 and ivanium as code owners July 9, 2026 17:25

claude Bot reviewed Jul 9, 2026

View reviewed changes

EdalatiAli changed the title ~~[Bugfix] Zero freshly allocated KV blocks for quantized + sliding-window hybrid caches~~ [Bugfix] Zero new KV blocks for quantized + sliding-window hybrid caches Jul 9, 2026

tlrmchlsmth reviewed Jul 13, 2026

View reviewed changes

Merge branch 'main' into fix-skipswa-kvquant

6f7e7e2

NickLucche reviewed Jul 16, 2026

View reviewed changes

fix the condition for kv block zeroing

6def96d

Signed-off-by: EdalatiAli <aliedalati@cohere.com>

NickLucche approved these changes Jul 20, 2026

View reviewed changes

NickLucche added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 20, 2026

Merge branch 'main' into fix-skipswa-kvquant

8cf6fe9

tlrmchlsmth approved these changes Jul 20, 2026

View reviewed changes

tlrmchlsmth enabled auto-merge (squash) July 20, 2026 13:02

tlrmchlsmth merged commit 8ce53a6 into vllm-project:main Jul 20, 2026
81 checks passed

elvircrn mentioned this pull request Jul 24, 2026

KVBlockZeroer crashes with non-uniform page sizes on GLM-5.2 FP8 #49696

Closed

ZJY0516 mentioned this pull request Jul 24, 2026

[Bugfix] Zero non-uniform KV page sizes in KVBlockZeroer #49746

Closed

4 tasks

sashko-zakharchuk mentioned this pull request Aug 27, 2026

[Doc] Correct the Mamba rationale on needs_kv_cache_zeroing #54004

Open

benenzhu mentioned this pull request Oct 18, 2026

[Bugfix][KV Connector][ROCm] MoRIIO READ: skip zeroing READ destinations for mixed-precision caches #62313

Closed

5 of 6 tasks

Navigation Menu

Repository navigation

[Bugfix] Zero new KV blocks for quantized + sliding-window hybrid caches - #47574

[Bugfix] Zero new KV blocks for quantized + sliding-window hybrid caches#47574
tlrmchlsmth merged 4 commits into
vllm-project:mainfrom
EdalatiAli:fix-skipswa-kvquant

EdalatiAli commented Jul 3, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

tlrmchlsmth Jul 13, 2026

Uh oh!

EdalatiAli Jul 20, 2026

Uh oh!

NickLucche left a comment

Uh oh!

NickLucche left a comment •

edited

Loading

Uh oh!

Uh oh!

EdalatiAli commented Jul 20, 2026

Uh oh!

ZJY0516 commented Jul 25, 2026

Uh oh!

NickLucche commented Jul 27, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

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

Uh oh!

Purpose

Root cause

Fix

Test Plan

Test Result

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

tlrmchlsmth Jul 13, 2026

Choose a reason for hiding this comment

Uh oh!

EdalatiAli Jul 20, 2026

Choose a reason for hiding this comment

Uh oh!

NickLucche left a comment

Choose a reason for hiding this comment

Uh oh!

NickLucche left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

EdalatiAli commented Jul 20, 2026

Uh oh!

ZJY0516 commented Jul 25, 2026

Uh oh!

NickLucche commented Jul 27, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

EdalatiAli commented Jul 3, 2026 •

edited

Loading

NickLucche left a comment •

edited

Loading