iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
vLLM-4x-RTX-3090/patches/0012-qsa-bound-the-indexer-page-table-values.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

183 lines
7.7 KiB
Diff
Raw Permalink Blame History

From 612e3d64642c5e3a0edba7d37d53f3b2ae9cd8da 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 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
Reference in New Issue View Git Blame Copy Permalink