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