From 79f3183f86b89c3bda05d467041bf3ef9ef60426 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Wed, 12 Aug 2026 23:42:26 -0400 Subject: [PATCH] [Bugfix] Bound KV block zeroing launch geometry (#52058) Signed-off-by: Lucas Wilkinson Co-authored-by: OpenAI Codex --- tests/v1/worker/test_kv_block_zeroer.py | 77 ++++++++++++++++++++++--- vllm/v1/worker/utils.py | 46 ++++++--------- 2 files changed, 86 insertions(+), 37 deletions(-) diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index 58d1a2d3e7..3502898872 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -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 diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index f495579ddb..46a73935df 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -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, )