From b1fdcd6e2757401653c10e57e01b01e047a2bf60 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Tue, 22 Sep 2026 04:13:07 -0400 Subject: [PATCH] [GDN] Bound the decode state index Ports the check from upstream vLLM PR #50021. Co-Authored-By: Hinata Signed-off-by: MOVZX --- .../ops/fused_sigmoid_gating.py | 42 ++++++++++++++++--- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py index 7e0c7e05ca..c5e75ec4ee 100644 --- a/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py +++ b/vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py @@ -10,6 +10,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.heuristics( @@ -37,6 +45,7 @@ def fused_sigmoid_gating_delta_rule_update_kernel( cu_seqlens, ssm_state_indices, num_accepted_tokens, + guard_ctr, # bounds-guard instrumentation: persistent int64[3] scale, N: tl.int64, # num of sequences T: tl.int64, # num of tokens @@ -106,12 +115,32 @@ def fused_sigmoid_gating_delta_rule_update_kernel( i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1 else: i_t = 0 - # Load state index and check for invalid entries - state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to( - tl.int64 - ) - # Skip if state index is invalid (NULL_BLOCK_ID=0) + # PR #50021: i_t is num_accepted_tokens - 1 and is otherwise + # unbounded against a stride_indices_seq-column tensor. A zero + # accepted count gives i_t == -1, a read before this request's row + # (and before the tensor for i_n == 0); a stale or too-large count + # reads past the row. The state_idx <= 0 guard below rejects only + # non-positive values, so an out-of-range read returning a garbage + # positive int flowed into h0 + state_idx * stride_init_state_token + # and faulted the SM. + idx_in_row = (i_t >= 0) & (i_t < stride_indices_seq) + if (i_t < 0) or (i_t >= stride_indices_seq): + tl.atomic_add(guard_ctr, 1) + tl.atomic_max(guard_ctr + 1, (i_n).to(tl.int64) + _GUARD_BIAS) + tl.atomic_max(guard_ctr + 2, (i_t).to(tl.int64) + _GUARD_BIAS) + state_idx = tl.load( + ssm_state_indices + i_n * stride_indices_seq + i_t, + mask=idx_in_row, + other=0, + ).to(tl.int64) + # Skip invalid state indices without exposing uninitialized output. + # o is q.new_empty in the entry below, so the previous early return + # left this request's rows holding whatever the allocator returned. if state_idx <= 0: + zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty) + for _ in range(0, T): + tl.store(p_o, zero, mask=mask_v) + p_o += HV * V return p_h0 = h0 + state_idx * stride_init_state_token else: @@ -255,6 +284,9 @@ def fused_sigmoid_gating_delta_rule_update( cu_seqlens=cu_seqlens, ssm_state_indices=ssm_state_indices, num_accepted_tokens=num_accepted_tokens, + guard_ctr=guard_telemetry.counter( + guard_telemetry.SITE_GDN_DECODE_STATE, q.device + ), scale=scale, N=N, T=T,