iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] Bound KV block zeroing launch geometry (#52058)

Browse Source 
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Lucas Wilkinson
2026-08-13 03:42:26 +00:00
committed by GitHub
co-authored by OpenAI Codex
parent f9538af28c
commit 79f3183f86
2 changed files with 86 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/worker/test_kv_block_zeroer.py
+69 -8
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
from types import SimpleNamespace
import pytest
@@ -10,6 +11,7 @@ from vllm.v1.kv_cache_interface import (
ChunkedLocalAttentionSpec,
SlidingWindowSpec,
)
from vllm.v1.worker import utils as worker_utils
from vllm.v1.worker.utils import (
AttentionGroup,
KVBlockZeroer,
@@ -121,10 +123,7 @@ def test_non_uniform_page_sizes():
seg_page_sizes = [page_size_a, page_size_b]
max_ps = max(seg_page_sizes)
def largest_power_of_2_divisor(n):
return n & -n
blk_size = min(min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes), 1024)
blk_size = min(1 << (max_ps - 1).bit_length(), 1024)
zeroer._meta = (
torch.tensor(
@@ -134,7 +133,7 @@ def test_non_uniform_page_sizes():
),
torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
max_ps // blk_size,
(max_ps + blk_size - 1) // blk_size,
blk_size,
2,
)
@@ -186,12 +185,74 @@ def test_packed_segment_zeros_only_its_last_block_page():
assert torch.equal(backing, expected)
def test_large_dsv4_launch_geometry(monkeypatch):
"""Keep the failing DSV4 shape efficient and within launch limits."""
device = torch.device("cpu")
n_blocks, n_segs = 6870, 181
layer_names = [f"layer.{i}" for i in range(n_segs)]
page_sizes = [9344 if i % 2 == 0 else 292 for i in range(n_segs)]
spec = SlidingWindowSpec(
block_size=1,
num_kv_heads=1,
head_size=1,
dtype=torch.int32,
sliding_window=1,
)
storages = {
name: torch.ones((1, page_size), dtype=torch.int32)
for name, page_size in zip(layer_names, page_sizes)
}
zeroer = KVBlockZeroer(
device,
attn_groups_iter=[
AttentionGroup(
_BlockFirstBackend, # type: ignore[arg-type]
[name],
spec,
group_id,
)
for group_id, name in enumerate(layer_names)
],
kernel_block_sizes=[1] * n_segs,
cache_dtype="auto",
static_forward_context={
name: SimpleNamespace(kv_cache=storage)
for name, storage in storages.items()
},
)
assert zeroer._meta is not None
_, _, seg_page_sizes, max_chunks, blk_size, n_segs = zeroer._meta
assert seg_page_sizes.tolist() == page_sizes
assert (max_chunks, blk_size, n_segs) == (10, 1024, 181)
captured_grids = []
class FakeKernel:
def __getitem__(self, grid):
captured_grids.append(grid)
return lambda *args, **kwargs: None
monkeypatch.setattr(worker_utils, "_zero_kv_blocks_kernel", FakeKernel())
monkeypatch.setattr(
worker_utils,
"async_tensor_h2d",
lambda values, **kwargs: torch.tensor(values, dtype=torch.int64),
)
zeroer.zero_block_ids(list(range(n_blocks)))
old_max_chunks = max(page_sizes) // 4
assert math.prod((n_blocks, n_segs, old_max_chunks)) > 2**31 - 1
assert captured_grids == [(n_blocks, n_segs, max_chunks)]
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_warmup_compiles_every_n_blocks_specialization():
def test_warmup_compiles_for_all_block_counts():
"""After warmup, no launch should trigger a first-request JIT compile.
``n_blocks`` is ``do_not_specialize``, so a single warmup launch must
cover every block count.
The block count is carried by the launch grid, so changing it must reuse
the warmup's compiled kernel.
"""
device = torch.device("cuda")
num_blocks = 64
vllm/v1/worker/utils.py
+17 -29
View File
@@ -17,7 +17,6 @@ from vllm.model_executor.models.interfaces import MultiModalEmbeddings
from vllm.model_executor.models.utils import extract_layer_index
from vllm.platforms import current_platform
from vllm.triton_utils import tl, triton
from vllm.utils.math_utils import largest_power_of_2_divisor
from vllm.utils.mem_utils import MemorySnapshot, format_gib
from vllm.utils.torch_utils import async_tensor_h2d
from vllm.v1.attention.backend import (
@@ -52,15 +51,12 @@ def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None:
raise RuntimeError(f"NaNs detected in logits: {corrupted_requests}")
@triton.jit(do_not_specialize=["n_blocks"])
@triton.jit
def _zero_kv_blocks_kernel(
seg_addrs_ptr,
seg_block_strides_ptr,
seg_page_sizes_ptr,
block_ids_ptr,
n_blocks,
N_SEGS: tl.constexpr,
MAX_CHUNKS: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
):
"""Zero KV cache blocks across all segments in a single launch.
@@ -78,29 +74,27 @@ def _zero_kv_blocks_kernel(
seg_addrs_ptr holds absolute byte addresses (int64) for each segment,
allowing segments to live in different CUDA allocations.
Programs are mapped as (block_index, seg_index, chunk_index).
Programs are mapped directly onto a 3-D grid as
(block_index, seg_index, chunk_index).
"""
pid = tl.program_id(0)
work_per_block = N_SEGS * MAX_CHUNKS
block_index = pid // work_per_block
if block_index >= n_blocks:
return
remainder = pid % work_per_block
seg_index = remainder // MAX_CHUNKS
chunk_index = remainder % MAX_CHUNKS
block_index = tl.program_id(0)
seg_index = tl.program_id(1)
chunk_index = tl.program_id(2)
block_stride_el = tl.load(seg_block_strides_ptr + seg_index)
page_size_el = tl.load(seg_page_sizes_ptr + seg_index)
if chunk_index >= page_size_el // BLOCK_SIZE:
chunk_offset = chunk_index.to(tl.int64) * BLOCK_SIZE
if chunk_offset >= page_size_el:
return
block_id = tl.load(block_ids_ptr + block_index)
seg_addr = tl.load(seg_addrs_ptr + seg_index)
ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32))
offset = (
block_id.to(tl.int64) * block_stride_el.to(tl.int64)
+ chunk_index.to(tl.int64) * BLOCK_SIZE
block_offset = block_id.to(tl.int64) * block_stride_el.to(tl.int64)
cols = chunk_offset + tl.arange(0, BLOCK_SIZE).to(tl.int64)
tl.store(
ptr + block_offset + cols,
tl.zeros([BLOCK_SIZE], dtype=tl.int32),
mask=cols < page_size_el,
)
cols = tl.arange(0, BLOCK_SIZE).to(tl.int64)
tl.store(ptr + offset + cols, tl.zeros([BLOCK_SIZE], dtype=tl.int32))
class KVBlockZeroer:
@@ -204,15 +198,12 @@ class KVBlockZeroer:
return
max_page_size_el = max(seg_page_sizes)
blk_size = min(
min(largest_power_of_2_divisor(ps) for ps in seg_page_sizes),
1024,
)
blk_size = min(1 << (max_page_size_el - 1).bit_length(), 1024)
self._meta = (
torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device),
torch.tensor(seg_block_strides, dtype=torch.int64, device=self.device),
torch.tensor(seg_page_sizes, dtype=torch.int64, device=self.device),
max_page_size_el // blk_size,
(max_page_size_el + blk_size - 1) // blk_size,
blk_size,
len(seg_addrs),
)
@@ -231,15 +222,12 @@ class KVBlockZeroer:
) = self._meta
n_blocks = len(block_ids)
idx = async_tensor_h2d(block_ids, device=self.device, dtype=torch.int64)
grid = (n_blocks * n_segs * max_chunks,)
grid = (n_blocks, n_segs, max_chunks)
_zero_kv_blocks_kernel[grid](
seg_addrs,
seg_block_strides,
seg_page_sizes,
idx,
n_blocks,
N_SEGS=n_segs,
MAX_CHUNKS=max_chunks,
BLOCK_SIZE=blk_size,
)