[haregali]

Purpose

With int8_per_token_head (or fp8/int4 per-token-head) KV cache on KV-sharing models (Gemma-4, Gemma-3n), the first max_num_seqs requests after engine start generate corrupted leading tokens. Found while investigating #50702 and #50749; the reporter has since verified on their rig that this defect is real but distinct from the primary failures those issues report (see the issue threads). Supersedes #51061, closed while we finished verification.

TritonAttentionImpl._ensure_scale_caches fills the cache's inline scale bytes with 1.0 when an impl first builds its scale views. KV-sharing layers alias the target layer's cache view, and dummy runs skip their attention (attn_metadata is None early return), so their first touch happens during the first real forward pass, after the target layers have already written per-token scales for the prompt. The fill resets those scales to 1.0 while leaving the quantized data intact; K dequantizes ~200x too large and the first decode steps of the entire first batch go bad. Later batches are unaffected, so under multiprocess serving the failure looks stochastic and spuriously correlated with prefix caching, prompt length, and concurrency — matching the confusing observations in both issues.

The fix initializes each cache view's scale region exactly once, keyed by view data_ptr with a storage weakref for liveness. Keying by view rather than storage keeps packed layouts correct (multiple layers in one storage each own their region). Skipping the fill for kv-sharing layers alone would not cover cross-group tensor sharing, where the same wipe can occur.

Evidence

Deterministic repro: google/gemma-4-E2B-it, TRITON_ATTN, in-process engine, greedy, 48 number-sequence prompts whose correct continuation is checkable, 4 copies each.

• Before: exactly the first max_num_seqs requests fail, every run (mns=4 → requests 0–3, 6 → 0–5, 8 → 0–7). bf16 passes.
• Instrumentation: 20 late fills, all on the aliased storage, none elsewhere; wiped scales read back as exactly 1.0; quantized codes byte-identical to a healthy run.
• Independent confirmation: pre-creating all scale views at init on unfixed code also eliminates the failure.
• After: zero failures across prefix caching on/off, prompt-length residues, mns 4/6/8, and multiprocess mode; first-vs-later-batch KV bytes identical; bf16 unchanged.

Test Plan

pytest tests/v1/attention/test_triton_attn_scale_init.py   # new, 5 tests, red on main
pre-commit ruff-check / ruff-format / mypy-3.12            # pass
vllm bench latency (gemma-4-E2B, int8_per_token_head): 3.127s vs 3.138s on main

Notes

No open PR touches this path (searched 2026-08-05). The reporters of #50702/#50749 run v0.26.0, which has this code; their confirmation on a patched build would establish how much of their production failure this explains. AI assistance was used; every changed line was human-reviewed and all experiments ran locally on one RTX 4090.

[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.

[haregali]

gsm8k, 250 samples, 5-shot, gemma-4-E2B-it, TRITON_ATTN, RTX 4090:

config	strict-match	flexible-extract
this PR, int8_per_token_head	0.108 ± 0.020	0.132 ± 0.022
main, int8_per_token_head	0.108 ± 0.020	0.132 ± 0.022
main, bf16	0.124 ± 0.021	0.152 ± 0.023

No accuracy change. The bug only affects the first batch per engine start (~6 of 250 samples), below eval sensitivity; the deterministic repro in the description is the sensitive measurement. This confirms no regression.

[mergify]

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

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