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

[Bugfix] Fix packed KV block zeroing stride (#50276)

Browse Source 
Signed-off-by: wangxian001 <120719093+wangxian001@users.noreply.github.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Co-authored-by: Nick Hill <nickhill123@gmail.com>
This commit is contained in:
wangxian001
2026-08-06 13:49:06 -07:00
committed by GitHub
co-authored by Nick Hill
parent adc3e03517
commit d6af803f43
2 changed files with 81 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/worker/test_kv_block_zeroer.py
+39
View File
@@ -21,6 +21,7 @@ def test_block_ids_are_not_overwritten_while_copy_is_in_flight():
zeroer._meta = (
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
torch.tensor([page_size_el], dtype=torch.int64, device=device),
torch.tensor([page_size_el], dtype=torch.int64, device=device),
page_size_el // page_size_el, # max_chunks = 1
page_size_el, # blk_size
1, # n_segs
@@ -71,6 +72,7 @@ def test_non_uniform_page_sizes():
device=device,
),
torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
torch.tensor(seg_page_sizes, dtype=torch.int64, device=device),
max_ps // blk_size,
blk_size,
2,
@@ -88,6 +90,41 @@ def test_non_uniform_page_sizes():
assert torch.all(storage[3] == 1)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_packed_segment_zeros_only_its_last_block_page():
"""A packed KV segment steps by block stride but clears only its page."""
device = torch.device("cuda")
num_blocks = 4
block_stride_el = 12
page_size_el = 4
page_offset_el = 3
backing = torch.ones(
(num_blocks, block_stride_el), dtype=torch.int32, device=device
)
zeroer = KVBlockZeroer.__new__(KVBlockZeroer)
zeroer.device = device
zeroer._meta = (
torch.tensor(
[backing.data_ptr() + page_offset_el * backing.element_size()],
dtype=torch.uint64,
device=device,
),
torch.tensor([block_stride_el], dtype=torch.int64, device=device),
torch.tensor([page_size_el], dtype=torch.int64, device=device),
1,
page_size_el,
1,
)
zeroer.zero_block_ids([num_blocks - 1])
torch.accelerator.synchronize()
expected = torch.ones_like(backing)
expected[-1, page_offset_el : page_offset_el + page_size_el] = 0
assert torch.equal(backing, expected)
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
def test_warmup_compiles_every_n_blocks_specialization():
"""After warmup, no launch should trigger a first-request JIT compile.
@@ -105,6 +142,7 @@ def test_warmup_compiles_every_n_blocks_specialization():
zeroer._meta = (
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
torch.tensor([page_size_el], dtype=torch.int64, device=device),
torch.tensor([page_size_el], dtype=torch.int64, device=device),
1, # max_chunks
page_size_el, # blk_size
1, # n_segs
@@ -141,6 +179,7 @@ def test_warmup_respects_available_block_count():
zeroer._meta = (
torch.tensor([storage.data_ptr()], dtype=torch.uint64, device=device),
torch.tensor([page_size_el], dtype=torch.int64, device=device),
torch.tensor([page_size_el], dtype=torch.int64, device=device),
1,
page_size_el,
1,
vllm/v1/worker/utils.py
+42 -17
View File
@@ -56,6 +56,7 @@ def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None:
@triton.jit(do_not_specialize=["n_blocks"])
def _zero_kv_blocks_kernel(
seg_addrs_ptr,
seg_block_strides_ptr,
seg_page_sizes_ptr,
block_ids_ptr,
n_blocks,
@@ -70,10 +71,10 @@ def _zero_kv_blocks_kernel(
buffer. For backends where K/V is outermost (block_dim=1) there are
two segments per buffer (one for K, one for V).
Segments may have different page sizes (e.g. models with multiple KV
cache groups like MLA + DSA indexer). Each segment's page size is
read from seg_page_sizes_ptr; programs whose chunk_index falls beyond
their segment's page size early-exit.
Segments may have different block strides and page sizes (e.g. packed
KV views or models with multiple KV cache groups like MLA + DSA
indexer). Each segment's block stride determines where a logical block
begins, while its page size determines how many elements are cleared.
seg_addrs_ptr holds absolute byte addresses (int64) for each segment,
allowing segments to live in different CUDA allocations.
@@ -88,6 +89,7 @@ def _zero_kv_blocks_kernel(
remainder = pid % work_per_block
seg_index = remainder // MAX_CHUNKS
chunk_index = remainder % MAX_CHUNKS
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:
return
@@ -95,7 +97,7 @@ def _zero_kv_blocks_kernel(
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) * page_size_el.to(tl.int64)
block_id.to(tl.int64) * block_stride_el.to(tl.int64)
+ chunk_index.to(tl.int64) * BLOCK_SIZE
)
cols = tl.arange(0, BLOCK_SIZE).to(tl.int64)
@@ -126,18 +128,21 @@ class KVBlockZeroer:
Block IDs from the scheduler reference logical blocks whose size
may differ from the kernel block size (virtual block splitting).
Each segment's page_size_el accounts for this ratio so that
``block_id * page_size_el`` lands at the correct offset.
Each virtual block is represented as an independent segment so its
physical block stride and zeroed page span remain independent.
Only AttentionSpec layers are processed; Mamba layers are skipped.
"""
self.device = device
self._meta: tuple[torch.Tensor, torch.Tensor, int, int, int] | None = None
self._meta: (
tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int] | None
) = None
if runner_only_attn_layers is None:
runner_only_attn_layers = set()
seen_ptrs: set[int] = set()
seg_addrs: list[int] = []
seg_block_strides: list[int] = []
seg_page_sizes: list[int] = []
for group in attn_groups_iter:
@@ -147,6 +152,7 @@ class KVBlockZeroer:
if group.kv_cache_group_id >= len(kernel_block_sizes):
continue
kernel_bs = kernel_block_sizes[group.kv_cache_group_id]
assert spec.block_size % kernel_bs == 0
ratio = spec.block_size // kernel_bs
block_dim = group.backend.get_kv_cache_block_dim(
kernel_bs,
@@ -167,22 +173,32 @@ class KVBlockZeroer:
seen_ptrs.add(dp)
el = kv.element_size()
cur_bytes = kv.stride(block_dim) * el
assert cur_bytes % 4 == 0
kernel_block_el = cur_bytes // 4
cur_page_el = kernel_block_el * ratio
block_stride_bytes = cur_bytes
block_stride_bytes = kv.stride(block_dim) * el
assert block_stride_bytes % 4 == 0
assert kv.shape[block_dim] % ratio == 0
outer_dims = [
d
for d in range(block_dim)
if kv.stride(d) * el > block_stride_bytes
]
outer_strides = [kv.stride(d) * el for d in outer_dims]
inner_dims = [
d for d in range(kv.ndim) if d != block_dim and d not in outer_dims
]
kernel_page_bytes = el + sum(
(kv.shape[d] - 1) * kv.stride(d) * el for d in inner_dims
)
assert kernel_page_bytes % 4 == 0
logical_block_stride_bytes = block_stride_bytes * ratio
for outer in iprod(*(range(kv.shape[d]) for d in outer_dims)):
off_bytes = sum(i * s for i, s in zip(outer, outer_strides))
seg_addrs.append(dp + off_bytes)
seg_page_sizes.append(cur_page_el)
assert (dp + off_bytes) % 4 == 0
for virtual_index in range(ratio):
seg_addrs.append(
dp + off_bytes + virtual_index * block_stride_bytes
)
seg_block_strides.append(logical_block_stride_bytes // 4)
seg_page_sizes.append(kernel_page_bytes // 4)
if not seg_addrs:
self._meta = None
@@ -195,6 +211,7 @@ class KVBlockZeroer:
)
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,
blk_size,
@@ -205,12 +222,20 @@ class KVBlockZeroer:
"""Zero the KV cache memory for the given block IDs."""
if not block_ids or self._meta is None:
return
seg_addrs, seg_page_sizes, max_chunks, blk_size, n_segs = self._meta
(
seg_addrs,
seg_block_strides,
seg_page_sizes,
max_chunks,
blk_size,
n_segs,
) = 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,)
_zero_kv_blocks_kernel[grid](
seg_addrs,
seg_block_strides,
seg_page_sizes,
idx,
n_blocks,