82 lines
3.3 KiB
Diff
82 lines
3.3 KiB
Diff
From 85f78b1410981293e030f6256c0c16c84e1a12d7 Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
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
|
|
|