[GDN] Bound the decode state index

Ports the check from upstream vLLM PR #50021.

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

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,