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 5, 07:15 UTC
[Bugfix] Fix GDN FLA kernel crashes with NULL_BLOCK_ID=0 CUDA graph padding (#39064)
Signed-off-by: Vibhav Agarwal <vibhavagarwal5@gmail.com> Co-authored-by: vibhav-agarwal <vibhav.agarwal@glance.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Vibhav Agarwal committed
Apr 11, 2026 at 14:05 UTC
d4cb783c10ffc091af7f09a3b052dceadc06d075
2 files changed
+14
-13
vllm/model_executor/layers/fla/ops/fused_recurrent.py
+8
-7
@@ -106,12 +106,12 @@ 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
- # Load state index and check for PAD_SLOT_ID (-1)
109
+ # Load state index and check for invalid entries
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:
113
+ # Skip if state index is invalid (NULL_BLOCK_ID=0)
114
+ if state_idx <= 0:
115
return
116
p_h0 = h0 + state_idx * stride_init_state_token
117
else:
@@ -150,12 +150,12 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
150
151
# keep the states for multi-query tokens
152
if INPLACE_FINAL_STATE:
153
- # Load state index and check for PAD_SLOT_ID (-1)
153
+ # Load state index and check for invalid entries
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:
157
+ # Only store if state index is valid (not NULL_BLOCK_ID=0)
158
+ if final_state_idx > 0:
159
p_ht = ht + final_state_idx * stride_final_state_token
160
p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
161
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
@@ -292,7 +292,8 @@ def fused_recurrent_gated_delta_rule_packed_decode_kernel(
292
state_idx = tl.load(ssm_state_indices + i_n * stride_indices_seq).to(tl.int64)
293
p_o = o + (i_n * HV + i_hv) * V + o_v
294
295
- if state_idx < 0:
295
+ # Skip if state index is invalid (NULL_BLOCK_ID=0)
296
+ if state_idx <= 0:
297
zero = tl.zeros([BV], dtype=tl.float32).to(p_o.dtype.element_ty)
298
tl.store(p_o, zero, mask=mask_v)
299
return
vllm/model_executor/layers/fla/ops/fused_sigmoid_gating.py
+6
-6
@@ -106,12 +106,12 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
106
i_t = tl.load(num_accepted_tokens + i_n).to(tl.int64) - 1
107
else:
108
i_t = 0
109
- # Load state index and check for PAD_SLOT_ID (-1)
109
+ # Load state index and check for invalid entries
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:
113
+ # Skip if state index is invalid (NULL_BLOCK_ID=0)
114
+ if state_idx <= 0:
115
return
116
p_h0 = h0 + state_idx * stride_init_state_token
117
else:
@@ -155,12 +155,12 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
155
156
# keep the states for multi-query tokens
157
if INPLACE_FINAL_STATE:
158
- # Load state index and check for PAD_SLOT_ID (-1)
158
+ # Load state index and check for invalid entries
159
final_state_idx = tl.load(
160
ssm_state_indices + i_n * stride_indices_seq + i_t
161
).to(tl.int64)
162
- # Only store if state index is valid (not PAD_SLOT_ID)
163
- if final_state_idx >= 0:
162
+ # Only store if state index is valid (not NULL_BLOCK_ID=0)
163
+ if final_state_idx > 0:
164
p_ht = ht + final_state_idx * stride_final_state_token
165
p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
166
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)