From 85f78b1410981293e030f6256c0c16c84e1a12d7 Mon Sep 17 00:00:00 2001 From: halt95 Date: Tue, 22 Sep 2026 04:13:08 -0400 Subject: [PATCH 014/103] worker: pool-bound the KV block zeroer --- vllm/v1/worker/utils.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index df2ebc7..71c1a2d 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -39,10 +39,17 @@ from vllm.v1.kv_cache_interface import ( UniformTypeKVCacheSpecs, create_kv_cache_views, ) +from vllm.v1.worker import guard_telemetry from vllm.v1.worker.block_table import get_block_table_width logger = init_logger(__name__) +# Bias so a negative row/index survives atomic_max on a zeroed counter +# -- it MUST be a tl.constexpr: a plain module-level int is not readable +# from a @triton.jit body -- an earlier build shipped one and every guarded kernel died +# with "NameError: Cannot access global variable _GUARD_BIAS" on the smoke. +_GUARD_BIAS = tl.constexpr(guard_telemetry.GUARD_BIAS) + def raise_if_nan_logits(num_nans_in_logits: Mapping[str, int]) -> None: if not any(num_nans_in_logits.values()): @@ -62,6 +69,8 @@ def _zero_kv_blocks_kernel( seg_block_strides_ptr, seg_page_sizes_ptr, block_ids_ptr, + num_pool_blocks, # bounds-guard instrumentation + guard_ctr, # bounds-guard instrumentation: persistent int64[3] BLOCK_SIZE: tl.constexpr, ): """Zero KV cache blocks across all segments in a single launch. @@ -91,6 +100,17 @@ def _zero_kv_blocks_kernel( if chunk_offset >= page_size_el: return block_id = tl.load(block_ids_ptr + block_index) + # Bounds-guard instrumentation (guard 5): block_id goes straight into + # block_id * block_stride_el with no pool bound, and this kernel masks only + # WITHIN a segment page, so a corrupt id writes zeros elsewhere in the KV + # allocation instead of faulting visibly. A valid pool id is always in + # range (every group's tensor is asserted to be an exact multiple of + # kv_cache_config.num_blocks), so this can only fire on corruption. + if (block_id < 0) or (block_id >= num_pool_blocks): + tl.atomic_add(guard_ctr, 1) + tl.atomic_max(guard_ctr + 1, (block_index).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max(guard_ctr + 2, (block_id).to(tl.int64) + _GUARD_BIAS) + return seg_addr = tl.load(seg_addrs_ptr + seg_index) ptr = tl.cast(seg_addr, tl.pointer_type(tl.int32)) block_offset = block_id.to(tl.int64) * block_stride_el.to(tl.int64) @@ -138,6 +158,8 @@ class KVBlockZeroer: Only AttentionSpec layers are processed; Mamba layers are skipped. """ self.device = device + # Bounds-guard instrumentation: pool bound for the guard in _zero_kv_blocks_kernel. + self._num_pool_blocks = int(num_blocks) self._meta: ( tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, int, int] | None ) = None @@ -246,6 +268,10 @@ class KVBlockZeroer: seg_block_strides, seg_page_sizes, idx, + self._num_pool_blocks, + guard_telemetry.counter( + guard_telemetry.SITE_KV_ZEROER, self.device + ), BLOCK_SIZE=blk_size, ) -- 2.43.0