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 6, 19:15 UTC
[Bugfix] Fix FlashInfer GDN warmup ValueError on SM90 GPUs (#36876)
Thomas Parnell committed
Mar 13, 2026 at 07:09 UTC
f296a1966dca96cd69e5c1fa1264edbf611a1bd6
1 file changed
+8
-2
vllm/model_executor/models/qwen3_next.py
+8
-2
@@ -137,7 +137,7 @@ def fi_chunk_gated_delta_rule(
137
fi_state = initial_state.to(torch.float32)
138
fi_g = g.to(torch.float32)
139
fi_beta = beta.to(torch.float32)
140
- output, final_state = chunk_gated_delta_rule_fi(
140
+ result = chunk_gated_delta_rule_fi(
141
q=q,
142
k=k,
143
v=v,
@@ -147,8 +147,14 @@ def fi_chunk_gated_delta_rule(
147
output_final_state=output_final_state,
148
cu_seqlens=cu_seqlens,
149
)
150
+ # FlashInfer returns (output, state) when output_final_state=True,
151
+ # or just output when output_final_state=False.
152
# Unsqueeze back to 4D (1, L, H, D) to match fla output format
151
- return output.unsqueeze(0), final_state
153
+ if output_final_state:
154
+ output, final_state = result
155
+ return output.unsqueeze(0), final_state
156
+ else:
157
+ return result.unsqueeze(0), None
158
159
160
@CustomOp.register("chunk_gated_delta_rule")