[Perf][GDN] Scatter mixed speculative outputs into the caller buffer (#57140)
Signed-off-by: Arpan Tripathi <tripathiarpan20@gmail.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Xavier Lyu <9938360+xavierlyu@users.noreply.github.com> Co-authored-by: Bohdan Podziubanchuk <114604629+Danbog32@users.noreply.github.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Codex <noreply@openai.com>
Arpan Tripathi committed
Sep 16, 2026 at 14:03 UTC
975dca5bb5db302077674cfa9afe851ee700ad73
1 file changed
+3
-7
vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py
+3
-7
@@ -1575,14 +1575,10 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
1575
1576
# 3. Merge core attention output
1577
if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
1578
- merged_out = torch.empty(
1579
- (1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
1580
- dtype=core_attn_out_non_spec.dtype,
1581
- device=core_attn_out_non_spec.device,
1578
+ core_attn_out.index_copy_(0, spec_token_indx, core_attn_out_spec.squeeze(0))
1579
+ core_attn_out.index_copy_(
1580
+ 0, non_spec_token_indx, core_attn_out_non_spec.squeeze(0)
1581
)
1583
- merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
1584
- merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
1585
- core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)
1582
elif spec_sequence_masks is not None:
1583
core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
1584
else: