From 612e3d64642c5e3a0edba7d37d53f3b2ae9cd8da Mon Sep 17 00:00:00 2001 From: halt95 Date: Tue, 22 Sep 2026 04:13:08 -0400 Subject: [PATCH 012/103] qsa: bound the indexer page-table values --- .../qwen4_exp/nvidia/ops/qsa_indexer.py | 76 ++++++++++++++++++- 1 file changed, 72 insertions(+), 4 deletions(-) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py index 9ddd65d..a46f915 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa_indexer.py @@ -10,6 +10,14 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import ( ) from vllm.platforms import current_platform 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) + _TOPK_WORKSPACE_BYTES = 1024 * 1024 _DECODE_BLOCK_N = 64 @@ -28,6 +36,8 @@ def _qsa_mqa_paged_uniform_kernel( stride_cache_token, stride_table_req, stride_logits_row, + num_cache_blocks, # bounds-guard instrumentation + guard_ctr, # bounds-guard instrumentation: persistent int64[3] PAGE_SIZE: tl.constexpr, PAGE_TABLE_WIDTH: tl.constexpr, NUM_HEADS: tl.constexpr, @@ -72,6 +82,11 @@ def _qsa_mqa_paged_uniform_kernel( other=0.0, ) column_offsets = tl.arange(0, BLOCK_N) + # Bounds-guard instrumentation: accumulate in registers inside the pipelined loop + # and do the cross-lane reduction + atomics once after it, so an + # atomic-bearing conditional does not sit in the loop body. + guard_bad_n = tl.zeros([BLOCK_N], dtype=tl.int32) + guard_bad_col = tl.full([BLOCK_N], -1, tl.int32) for tile in tl.range(tile_start, tile_end, num_stages=STAGES): columns = tile * BLOCK_N + column_offsets live = columns < max_visible @@ -80,11 +95,22 @@ def _qsa_mqa_paged_uniform_kernel( physical_page = tl.load( page_table_ptr + request * stride_table_req + logical_page, mask=live, - other=0, + other=-1, ) + # Mirror of the QSA attention kernel (qsa.py:137): the VALUE read out of + # the page table is untrusted and must be bounded before it becomes an + # address. Folds into the existing live mask. + page_bad = live & ((physical_page < 0) | (physical_page >= num_cache_blocks)) + guard_bad_n += page_bad.to(tl.int32) + guard_bad_col = tl.maximum( + guard_bad_col, tl.where(page_bad, logical_page, -1).to(tl.int32) + ) + live = live & (physical_page >= 0) & (physical_page < num_cache_blocks) + # physical_page * block stride can overflow int32 for large caches. + safe_page = tl.maximum(physical_page, 0).to(tl.int64) keys = tl.load( k_cache_ptr - + physical_page[:, None].to(tl.int64) * stride_cache_block + + safe_page[:, None] * stride_cache_block + page_offset[:, None] * stride_cache_token + dims[None, :], mask=live[:, None] & (dims[None, :] < HEAD_DIM), @@ -106,6 +132,15 @@ def _qsa_mqa_paged_uniform_kernel( & (columns[:, None] < visible[None, :]), ) + # Bounds-guard instrumentation: one reduction and at most one atomic per program. + guard_n = tl.sum(guard_bad_n, axis=0) + if guard_n > 0: + tl.atomic_add(guard_ctr, guard_n.to(tl.int64)) + tl.atomic_max(guard_ctr + 1, (request).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max( + guard_ctr + 2, (tl.max(guard_bad_col, axis=0)).to(tl.int64) + _GUARD_BIAS + ) + @triton.jit(do_not_specialize=["num_rows", "query_offset"]) def _qsa_mqa_paged_prefill_kernel( @@ -124,6 +159,8 @@ def _qsa_mqa_paged_prefill_kernel( num_rows, query_offset, page_table_width, + num_cache_blocks, # bounds-guard instrumentation + guard_ctr, # bounds-guard instrumentation: persistent int64[3] PAGE_SIZE: tl.constexpr, NUM_HEADS: tl.constexpr, HEAD_DIM: tl.constexpr, @@ -181,6 +218,11 @@ def _qsa_mqa_paged_prefill_kernel( other=0.0, ) column_offsets = tl.arange(0, BLOCK_N) + # Bounds-guard instrumentation: accumulate in registers inside the pipelined loop + # and do the cross-lane reduction + atomics once after it, so an + # atomic-bearing conditional does not sit in the loop body. + guard_bad_n = tl.zeros([BLOCK_N], dtype=tl.int32) + guard_bad_col = tl.full([BLOCK_N], -1, tl.int32) for tile in tl.range(k_tile_start, k_tile_end, num_stages=STAGES): columns = tile * BLOCK_N + column_offsets live = columns < max_visible @@ -189,11 +231,22 @@ def _qsa_mqa_paged_prefill_kernel( physical_page = tl.load( page_table_ptr + request * stride_table_req + logical_page, mask=live, - other=0, + other=-1, + ) + # Mirror of the QSA attention kernel (qsa.py:137): the VALUE read out of + # the page table is untrusted and must be bounded before it becomes an + # address. Folds into the existing live mask. + page_bad = live & ((physical_page < 0) | (physical_page >= num_cache_blocks)) + guard_bad_n += page_bad.to(tl.int32) + guard_bad_col = tl.maximum( + guard_bad_col, tl.where(page_bad, logical_page, -1).to(tl.int32) ) + live = live & (physical_page >= 0) & (physical_page < num_cache_blocks) + # physical_page * block stride can overflow int32 for large caches. + safe_page = tl.maximum(physical_page, 0).to(tl.int64) keys = tl.load( k_cache_ptr - + physical_page[:, None].to(tl.int64) * stride_cache_block + + safe_page[:, None] * stride_cache_block + page_offset[:, None] * stride_cache_token + dims[None, :], mask=live[:, None] & (dims[None, :] < HEAD_DIM), @@ -217,6 +270,15 @@ def _qsa_mqa_paged_prefill_kernel( mask=store_mask, ) + # Bounds-guard instrumentation: one reduction and at most one atomic per program. + guard_n = tl.sum(guard_bad_n, axis=0) + if guard_n > 0: + tl.atomic_add(guard_ctr, guard_n.to(tl.int64)) + tl.atomic_max(guard_ctr + 1, (request).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max( + guard_ctr + 2, (tl.max(guard_bad_col, axis=0)).to(tl.int64) + _GUARD_BIAS + ) + @triton.jit def _expand_qsa_indices_kernel( @@ -360,6 +422,8 @@ def warmup_qsa_mqa_paged_decode( k_cache.stride(1), page_table.stride(0), columns, + k_cache.shape[0], + TritonWarmupTensor(torch.int64), PAGE_SIZE=page_size, PAGE_TABLE_WIDTH=page_table_width, NUM_HEADS=num_heads, @@ -423,6 +487,8 @@ def _prefill_logits( num_queries, query_offset, page_table.shape[1], + k_cache.shape[0], + guard_telemetry.counter(guard_telemetry.SITE_QSA_PREFILL, q.device), PAGE_SIZE=k_cache.shape[1], NUM_HEADS=q.shape[1], HEAD_DIM=q.shape[2], @@ -551,6 +617,8 @@ def qsa_select_paged_decode( *k_cache.stride()[:2], *page_table.stride()[:-1], *logits.stride()[:-1], + k_cache.shape[0], + guard_telemetry.counter(guard_telemetry.SITE_QSA_DECODE, q.device), PAGE_SIZE=k_cache.shape[1], PAGE_TABLE_WIDTH=page_table.shape[1], NUM_HEADS=q.shape[1], -- 2.43.0