A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 10:00 UTC
[Bugfix][Kernel] Fix negative memory offset in GDN Triton kernel (#33326)
Signed-off-by: CarstyYou <186021327+CarstyYou@users.noreply.github.com>
CarstyYou committed
Jan 30, 2026 at 02:40 UTC
23591e631e1982610731d3de1ed09cad16094d32
1 file changed
+18
-15
vllm/model_executor/layers/fla/ops/fused_recurrent.py
+18
-15
@@ -106,13 +106,14 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
106
i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
107
else:
108
i_t = 0
109
- p_h0 = (
110
- h0
111
- + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(
112
- tl.int64
113
- )
114
- * stride_init_state_token
109
+ # Load state index and check for PAD_SLOT_ID (-1)
110
+ state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(
111
+ tl.int64
112
)
113
+ # Skip if state index is invalid (PAD_SLOT_ID = -1)
114
+ if state_idx < 0:
115
+ return
116
+ p_h0 = h0 + state_idx * stride_init_state_token
117
else:
118
p_h0 = h0 + bos * HV * K * V
119
p_h0 = p_h0 + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
@@ -149,17 +150,19 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
150
151
# keep the states for multi-query tokens
152
if INPLACE_FINAL_STATE:
152
- p_ht = (
153
- ht
154
- + tl.load(ssm_state_indices + i_n * stride_indices_seq + i_t).to(
155
- tl.int64
156
- )
157
- * stride_final_state_token
158
- )
153
+ # Load state index and check for PAD_SLOT_ID (-1)
154
+ final_state_idx = tl.load(
155
+ ssm_state_indices + i_n * stride_indices_seq + i_t
156
+ ).to(tl.int64)
157
+ # Only store if state index is valid (not PAD_SLOT_ID)
158
+ if final_state_idx >= 0:
159
+ p_ht = ht + final_state_idx * stride_final_state_token
160
+ p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
161
+ tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
162
else:
163
p_ht = ht + (bos + i_t) * stride_final_state_token
161
- p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
162
- tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
164
+ p_ht = p_ht + i_hv * K * V + o_k[:, None] * V + o_v[None, :]
165
+ tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
166
167
p_q += H * K
168
p_k += H * K