Repository navigation
Conversation
With int8/fp8_per_token_head KV cache on KV-sharing models (Gemma-4), the first scheduled batch after engine start emits corrupted leading tokens. TritonAttentionImpl._ensure_scale_caches builds its scale views lazily per impl and unconditionally fills the storage's inline scale regions with 1.0 on first touch. KV-sharing layers alias their target layer's cache tensor and, when warmup did not exercise them, first touch it mid-way through the first real forward pass - after the target layer already wrote real per-token scales for the prompt. The fill wipes those scales (leaving the int8 codes intact), so the prompt K/V dequantize against 1.0 instead of ~0.005, inflating K by ~200x and corrupting the first decode steps of every request in the first batch. Subsequent batches are clean because the views exist afterwards, which made the bug look stochastic under multiprocess serving and dependent on batch geometry, prefix caching, and prompt-length residue (vllm-project#50702, vllm-project#50749). Fix: track initialized storages on the class and fill each storage's scale regions exactly once, keyed by (data_ptr, nbytes) so re-allocated caches (sleep/wake_up) are re-initialized. Verified on a deterministic single-GPU reproduction (gemma-4-E2B-it, TRITON_ATTN, int8_per_token_head, in-process engine, 48 arithmetic continuation prompts x 4 copies, greedy): before the fix the entire first batch (exactly max_num_seqs requests) produces objectively wrong continuations in every run; after the fix zero requests are damaged across prefix-caching on/off, prompt-length residues, max_num_seqs 4/6/8, multiprocess mode, and first-vs-later-batch KV bytes become identical. bf16 outputs are unaffected. Signed-off-by: Hareesh <hgsata@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Rsg3Q3x8FCVCZd1fUFED5U
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
With
int8_per_token_head/fp8_per_token_headKV cache on KV-sharing models (Gemma-4 family), the first scheduled batch after engine start emits corrupted leading tokens. Likely the root cause behind #50702 and #50749.TritonAttentionImpl._ensure_scale_cachesbuilds its inline-scale views lazily per impl and unconditionallyfill_(1.0)s the storage's scale regions on first touch. KV-sharing layers alias their target layer's cache tensor; when warmup didn't exercise them, their first touch happens mid-way through the first real forward pass — after the target layers already wrote real per-token scales for the prompt. The fill wipes the scales (int8 codes stay intact), so prompt K/V dequantize against 1.0 instead of ~0.005: K inflated ~200×, first decode steps of every request in the first batch corrupted. Later batches are clean (views exist), which made the bug look stochastic under multiprocess serving and spuriously dependent on batch geometry / prefix caching / prompt-length residue — matching the confusing conditionals reported in both issues.Fix: initialize each storage's scale regions exactly once (class-level registry keyed by
(data_ptr, nbytes), so caches re-allocated by sleep/wake_up are re-initialized).Evidence
Deterministic single-GPU repro:
google/gemma-4-E2B-it, TRITON_ATTN, MRV1, in-process engine, 48 arithmetic-continuation prompts × 4 copies, greedy — output objectively checkable against the true sequence.max_num_seqsrequests damaged, every run (mns=4 → requests 0–3, mns=6 → 0–5, mns=8 → 0–7); int8-only (bf16 clean)._ensure_scale_caches: 33 fills wipe nothing (pre-write), 2 late fills wipe 22,384 / 13,394 written non-default scales on an aliased storage mid-first-batch.221, V ×24 = exactly1.0 / true_scale.Test Plan
Repro harness (~100 s/run on one RTX 4090) available on request; happy to contribute it as an integration test if desired.
Notes
_ensure_scale_cachesor this failure (searched 2026-08-04); [Bugfix] Fix int4_per_token_head for SWA models and gemma 4 #48842 fixes an unrelated int4 startup crash.