[QSA] Bound the RoPE position before it becomes an address

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>

vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py

@@ -23,6 +23,8 @@ def _norm_rope(
     pos_w,
     cos_sin_ptr,
     cos_sin_stride,
+    cos_sin_rows,  # bounds-guard instrumentation
+    guard_ctr,  # bounds-guard persistent int64[3], SITE_QSA_ROPE_POS
     norm_weight_ptr,
     eps,
     IS_MROPE: tl.constexpr,
@@ -30,6 +32,72 @@ def _norm_rope(
     MROPE_W: tl.constexpr,
 ):
     """Apply Gemma RMSNorm and selected-axis NeoX RoPE to register rows."""
+    # Bounds-guard instrumentation (ring-residue crash replay): the position is a
+    # ROW INDEX into the cos/sin table and nothing bounds it. The masks below
+    # bound which rotary PAIR is read, never which row, so
+    # `cos_sin_ptr + pos * cos_sin_stride` is an address formed from an
+    # unchecked value.
+    #
+    # On the pooled cross-chunk path that value is read straight out of a
+    # raw-key RING row (:305-317). The ring is excluded from zeroing
+    # (single_type_kv_cache_manager.py:96-101), so a re-claimed block keeps the
+    # previous occupant's bytes -- and nothing ties the row the reader selects
+    # (`first_position % STATE_SIZE`) to a row any occupant ever committed, so
+    # it can be allocator residue rather than a position at all. A crash replay drove
+    # exactly this and compute-sanitizer reported 11 invalid reads at
+    # +4.6e18 bytes, the filler bytes decoded as an int64.
+    #
+    # This is the ONLY guard in the series that covers that chain: the request guard
+    # bounds the pre-indexer's `request`, which sits upstream of the ring read
+    # and cannot see a bad position.
+    # WHICH AXIS. Under the production mrope section (11, 11, 10)
+    # all three components are live row indices: the IS_MROPE branch masks
+    # COLUMNS per axis while the row comes unmasked from pos_rows[axis], so a
+    # bad pos_h or pos_w faults exactly as a bad pos_t does. Recording pos_t
+    # unconditionally would name the wrong axis and, for a t-clean/h-bad case,
+    # record an in-range value next to a non-zero count.
+    #
+    # The filler under tl.where must be strongly negative, NOT 0: an offender
+    # that is entirely negative (i_t == -1 is the shape that started this
+    # series) would otherwise max to 0 and read back as position 0 -- the
+    # defect re-entering at the one site with a reproduced fault.
+    _FILL: tl.constexpr = -(1 << 62)
+    bad_t = (pos_t < 0) | (pos_t >= cos_sin_rows)
+    if IS_MROPE:
+        bad_h = (pos_h < 0) | (pos_h >= cos_sin_rows)
+        bad_w = (pos_w < 0) | (pos_w >= cos_sin_rows)
+    else:
+        # the non-mrope branch reads only pos_t (:63-66), and with NO mask at
+        # all, which is why this guard clamps rather than masks.
+        # always-false, but correctly shaped and typed for the max below.
+        bad_h = pos_h < pos_h
+        bad_w = pos_w < pos_w
+    guard_bad = bad_t | bad_h | bad_w
+    guard_n = tl.sum(guard_bad.to(tl.int32), axis=0)
+    if guard_n > 0:
+        worst_t = tl.max(tl.where(bad_t, pos_t, _FILL), axis=0).to(tl.int64)
+        worst_h = tl.max(tl.where(bad_h, pos_h, _FILL), axis=0).to(tl.int64)
+        worst_w = tl.max(tl.where(bad_w, pos_w, _FILL), axis=0).to(tl.int64)
+        worst = tl.maximum(worst_t, tl.maximum(worst_h, worst_w))
+        # slot 2 names the axis: 0 = t, 1 = h, 2 = w. Derived from the MASKS
+        # (deliberately), not by comparing the winner's VALUE against each
+        # axis's max. Value comparison mis-attributes whenever two axes share
+        # the worst value -- which is the common case here, because pos_h and
+        # pos_w are frequently equal to pos_t (the non-2D path sets them so,
+        # and a stale ring tail writes the same int64 into all three). Asking
+        # "did this axis have any offender at all", in t/h/w order, cannot.
+        any_t = tl.max(bad_t.to(tl.int32), axis=0) > 0
+        any_h = tl.max(bad_h.to(tl.int32), axis=0) > 0
+        axis = tl.where(any_t, 0, tl.where(any_h, 1, 2))
+        tl.atomic_add(guard_ctr, guard_n.to(tl.int64))
+        tl.atomic_max(guard_ctr + 1, worst + _GUARD_BIAS)
+        tl.atomic_max(guard_ctr + 2, axis.to(tl.int64) + _GUARD_BIAS)
+    # Clamp rather than skip: the rotation is then wrong for these rows, which
+    # is what a masked fault IS. The counter is the result, not the output.
+    pos_t = tl.minimum(tl.maximum(pos_t, 0), cos_sin_rows - 1)
+    pos_h = tl.minimum(tl.maximum(pos_h, 0), cos_sin_rows - 1)
+    pos_w = tl.minimum(tl.maximum(pos_w, 0), cos_sin_rows - 1)
+
     TILE_T: tl.constexpr = x.shape[0]
     TILE_H: tl.constexpr = x.shape[1]
     D: tl.constexpr = x.shape[2]
@@ -127,6 +195,8 @@ def _qsa_prepare_kernel(
     pos_stride_axis,
     pos_stride_token,
     cos_sin_ptr,
+    cos_sin_rows,  # bounds-guard instrumentation
+    rope_guard_ctr,  # bounds-guard persistent int64[3], SITE_QSA_ROPE_POS
     q_norm_weight_ptr,
     k_norm_weight_ptr,
     eps,
@@ -289,6 +359,8 @@ def _qsa_prepare_kernel(
             pos_w,
             cos_sin_ptr,
             D // 2,
+            cos_sin_rows,
+            rope_guard_ctr,
             q_norm_weight_ptr,
             eps,
             IS_2D_POSITIONS,
@@ -449,6 +521,8 @@ def _qsa_prepare_kernel(
                 pos_w + tl.arange(0, 1),
                 cos_sin_ptr,
                 D // 2,
+                cos_sin_rows,
+                rope_guard_ctr,
                 k_norm_weight_ptr,
                 eps,
                 IS_K_MROPE,
@@ -611,6 +685,8 @@ def qsa_prepare(
         pos_stride_axis,
         pos_stride_token,
         cos_sin_cache,
+        cos_sin_cache.shape[0],
+        guard_telemetry.counter(guard_telemetry.SITE_QSA_ROPE_POS, q.device),
         q_norm_weight,
         k_norm_weight,
         eps,