[Bugfix] Align Qwen GDN gates with speculative tokens (#51812)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Jiangyun Zhu committed
Aug 11, 2026 at 23:35 UTC
5af7c8dad798bf899813f8f3c6b9eaf08a748e17
1 file changed
+6
-2
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+6
-2
@@ -1252,9 +1252,13 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1252
if spec_sequence_masks is not None:
1253
if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0:
1254
mixed_qkv_spec = mixed_qkv
1255
+ a_spec = a
1256
+ b_spec = b
1257
mixed_qkv_non_spec = None
1258
else:
1259
mixed_qkv_spec = mixed_qkv.index_select(0, spec_token_indx)
1260
+ a_spec = a.index_select(0, spec_token_indx)
1261
+ b_spec = b.index_select(0, spec_token_indx)
1262
mixed_qkv_non_spec = mixed_qkv.index_select(0, non_spec_token_indx)
1263
else:
1264
mixed_qkv_spec = None
@@ -1379,8 +1383,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1383
core_attn_out_spec, last_recurrent_state = (
1384
fused_sigmoid_gating_delta_rule_update(
1385
A_log=self.A_log,
1382
- a=a,
1383
- b=b,
1386
+ a=a_spec,
1387
+ b=b_spec,
1388
dt_bias=self.dt_bias,
1389
q=query_spec,
1390
k=key_spec,