From 8ca58f171461e0628a78abfae7b6df337ffe05d6 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Tue, 22 Sep 2026 04:13:07 -0400 Subject: [PATCH] [QSA] Bound the pre-indexer request row Co-Authored-By: Hinata Signed-off-by: MOVZX --- .../qwen4_exp/nvidia/ops/qsa_prepare.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py index a26b699657..1f04526fab 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py @@ -5,6 +5,14 @@ import torch from vllm.triton_utils import tl, triton +from vllm.v1.worker import guard_telemetry + +# 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) + @triton.jit @@ -107,6 +115,7 @@ def _store_rotated(dst, y, o1, o2, scale): "num_state_blocks", "num_compressed_blocks", "num_k_work", + "num_requests", ] ) def _qsa_prepare_kernel( @@ -141,6 +150,8 @@ def _qsa_prepare_kernel( num_state_blocks, num_compressed_blocks, num_k_work, + num_requests, # bounds-guard instrumentation + guard_ctr, # bounds-guard instrumentation: persistent int64[3] HQ: tl.constexpr, D: tl.constexpr, TILE_T_Q: tl.constexpr, @@ -300,6 +311,17 @@ def _qsa_prepare_kernel( request, work_in_request = tl.split(work_metadata) if request < 0: return + # Second bounds guard: the early return above rejects only the -1 + # sentinel the builder writes for inactive work items. Any other + # corrupted value was loaded straight out of query_start_loc, which has + # num_requests + 1 entries -- 8 bytes for a single-request batch. The + # chunk_end load below derives its index from query_end, so a bad row + # also walked logical_positions; returning here stops both. + if request >= num_requests: + tl.atomic_add(guard_ctr, 1) + tl.atomic_max(guard_ctr + 1, (request).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max(guard_ctr + 2, (work_in_request).to(tl.int64) + _GUARD_BIAS) + return query_start = tl.load(query_start_loc_ptr + request) query_end = tl.load(query_start_loc_ptr + request + 1) @@ -612,6 +634,8 @@ def qsa_prepare( state_cache.shape[0], compressed_cache.shape[0], num_k_work, + min(query_start_loc.numel() - 1, state_block_table.shape[0]), + guard_telemetry.counter(guard_telemetry.SITE_QSA_PRE_INDEXER, q.device), HQ=num_q_heads, D=head_dim, TILE_T_Q=TILE_T_Q,