diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py b/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py index 1f04526fab..1b26fd8145 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa_prepare.py +++ b/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,