From 103ea6246a6d3d82553d807a604407c50eb640bc Mon Sep 17 00:00:00 2001 From: Hareesh Date: Tue, 4 Aug 2026 10:37:22 -0700 Subject: [PATCH] [Bugfix] Initialize per-token-head KV scale regions once per storage 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 (#50702, #50749). Fix: track initialized storages in a class-level registry and fill each storage's scale regions exactly once, keyed by data_ptr with a weakref to the storage so freed-and-reallocated 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 Co-Authored-By: Claude Fable 5 --- .../attention/test_triton_attn_scale_init.py | 109 ++++++++++++++++++ vllm/v1/attention/backends/triton_attn.py | 32 ++++- 2 files changed, 139 insertions(+), 2 deletions(-) create mode 100644 tests/v1/attention/test_triton_attn_scale_init.py diff --git a/tests/v1/attention/test_triton_attn_scale_init.py b/tests/v1/attention/test_triton_attn_scale_init.py new file mode 100644 index 000000000000..155982634564 --- /dev/null +++ b/tests/v1/attention/test_triton_attn_scale_init.py @@ -0,0 +1,109 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Inline per-token-head KV scale regions must be initialized exactly once +per storage. + +KV-sharing layers alias their target layer's cache tensor and build their +scale views lazily on first forward, which can happen after the target layer +has already written real per-token scales in the same pass; re-initializing +on view creation wipes those scales and corrupts the first batch +(#50702, #50749). +""" + +import pytest +import torch + +from vllm.v1.attention.backends.triton_attn import TritonAttentionImpl + + +def _make_impl() -> TritonAttentionImpl: + impl = object.__new__(TritonAttentionImpl) + impl._k_scale_cache = None + impl._v_scale_cache = None + return impl + + +@pytest.fixture(autouse=True) +def _reset_registry(): + registry = getattr(TritonAttentionImpl, "_scale_initialized_storages", None) + saved = dict(registry) if registry is not None else None + if registry is not None: + registry.clear() + yield + if registry is not None: + registry.clear() + registry.update(saved) + + +def _kv_cache(num_blocks=4, nkv=1, block_size=16, head_size=256): + padded = head_size + 4 + return torch.zeros(num_blocks, nkv, block_size, 2 * padded, dtype=torch.int8) + + +def test_first_impl_initializes_scales_to_one(): + impl = _make_impl() + impl._ensure_scale_caches(_kv_cache()) + assert torch.all(impl._k_scale_cache == 1.0) + assert torch.all(impl._v_scale_cache == 1.0) + + +def test_aliasing_impl_preserves_written_scales(): + kv = _kv_cache() + writer = _make_impl() + writer._ensure_scale_caches(kv) + writer._k_scale_cache[1, 2, 0] = 0.123 + writer._v_scale_cache[3, 5, 0] = 0.456 + + alias = _make_impl() + alias._ensure_scale_caches(kv) + assert alias._k_scale_cache[1, 2, 0].item() == pytest.approx(0.123) + assert alias._v_scale_cache[3, 5, 0].item() == pytest.approx(0.456) + + +def test_new_storage_is_initialized(): + first = _make_impl() + kv1 = _kv_cache() + first._ensure_scale_caches(kv1) + + second = _make_impl() + kv2 = _kv_cache() + second._ensure_scale_caches(kv2) + assert torch.all(second._k_scale_cache == 1.0) + assert torch.all(second._v_scale_cache == 1.0) + + +def test_dead_registry_entry_is_reinitialized(): + """A registry entry whose storage died (freed-and-reallocated cache at a + reused address) must not suppress initialization of the new cache.""" + import weakref + + doomed = torch.zeros(8, dtype=torch.int8) + dead_ref = weakref.ref(doomed.untyped_storage()) + del doomed + + kv2 = _kv_cache() + kv2[0, 0, 0, 256:260] = -5 # dirty the first K-scale byte region + TritonAttentionImpl._scale_initialized_storages[kv2.data_ptr()] = dead_ref + assert dead_ref() is None + + impl = _make_impl() + impl._ensure_scale_caches(kv2) + assert torch.all(impl._k_scale_cache == 1.0) + assert torch.all(impl._v_scale_cache == 1.0) + + +def test_packed_layers_with_distinct_offsets_both_initialize(): + """Two layers packed into one storage at different offsets must each + initialize their own scale region (registry keys on view identity).""" + padded = 260 + backing = torch.zeros(2 * 4 * 1 * 16 * 2 * padded, dtype=torch.int8) + kv_a = backing[: backing.numel() // 2].view(4, 1, 16, 2 * padded) + kv_b = backing[backing.numel() // 2 :].view(4, 1, 16, 2 * padded) + + a = _make_impl() + a._ensure_scale_caches(kv_a) + b = _make_impl() + b._ensure_scale_caches(kv_b) + assert torch.all(a._k_scale_cache == 1.0) + assert torch.all(b._k_scale_cache == 1.0) + assert torch.all(b._v_scale_cache == 1.0) diff --git a/vllm/v1/attention/backends/triton_attn.py b/vllm/v1/attention/backends/triton_attn.py index 3f9b9ed0c144..1c9706ef7197 100644 --- a/vllm/v1/attention/backends/triton_attn.py +++ b/vllm/v1/attention/backends/triton_attn.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """High-Performance Triton-only Attention layer.""" +import weakref from dataclasses import dataclass from typing import ClassVar @@ -417,6 +418,21 @@ class TritonAttentionImpl(AttentionImpl): # Per-token-head quant: scale views carved from inline head padding. _k_scale_cache: torch.Tensor | None = None _v_scale_cache: torch.Tensor | None = None + # KV-cache views whose inline scale regions have been initialized; maps + # the view's data_ptr (base + storage_offset, so distinct layers packed + # into one storage each initialize their own region) to a weakref of the + # underlying storage, honored only while that exact storage is alive so a + # freed-and-reallocated cache is re-initialized. Relies on torch's + # storage PyObject preservation (torch >= 2.1) for weakref stability. + # Shared across impls because KV-sharing layers alias their target + # layer's cache view. + # NOTE: cumem sleep/wake keeps the storage object alive and remaps + # physical pages under it, so entries survive wake; post-wake scale + # sanity comes from the runner zeroing the cache on wake, and every + # written slot's scale is rewritten before it can be read. + _scale_initialized_storages: ClassVar[ + dict[int, "weakref.ReferenceType[torch.UntypedStorage]"] + ] = {} def _ensure_scale_caches(self, kv_cache: torch.Tensor) -> None: """Extract per-head scale views from the padded content dimension. @@ -470,7 +486,6 @@ def to_f32_units(elements: int) -> int: stride=(block_f32, slot_f32, head_f32), storage_offset=k_scale_off_f32, ) - self._k_scale_cache.fill_(1.0) # V scales (second content half) self._v_scale_cache = torch.as_strided( @@ -479,7 +494,20 @@ def to_f32_units(elements: int) -> int: stride=(block_f32, slot_f32, head_f32), storage_offset=v_scale_off_f32, ) - self._v_scale_cache.fill_(1.0) + + # Initialize the scale regions exactly once per underlying storage. + # KV-sharing layers alias their target layer's cache tensor and build + # their own views lazily on their first forward, which can happen + # AFTER the target layer has already written real per-token scales in + # the same pass; an unconditional fill here would wipe those scales + # and corrupt every request in the first batch (issues #50702/#50749). + registry = TritonAttentionImpl._scale_initialized_storages + view_ptr = kv_cache.data_ptr() + ref = registry.get(view_ptr) + if ref is None or ref() is None: + self._k_scale_cache.fill_(1.0) + self._v_scale_cache.fill_(1.0) + registry[view_ptr] = weakref.ref(kv_cache.untyped_storage()) def fused_output_quant_supported(self, quant_key: QuantKey): return quant_key == kFp8StaticTensorSym