[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>
This commit is contained in:
Arpan Tripathi
2026-09-16 13:03:57 +00:00
committed by GitHub
co-authored by Xavier Lyu Bohdan Podziubanchuk zjy0516 Codex
parent 8b1d188046
commit 975dca5bb5
@@ -1575,14 +1575,10 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
# 3. Merge core attention output
if spec_sequence_masks is not None and core_attn_out_non_spec is not None:
merged_out = torch.empty(
(1, num_actual_tokens, *core_attn_out_spec.shape[2:]),
dtype=core_attn_out_non_spec.dtype,
device=core_attn_out_non_spec.device,
core_attn_out.index_copy_(0, spec_token_indx, core_attn_out_spec.squeeze(0))
core_attn_out.index_copy_(
0, non_spec_token_indx, core_attn_out_non_spec.squeeze(0)
)
merged_out.index_copy_(1, spec_token_indx, core_attn_out_spec)
merged_out.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec)
core_attn_out[:num_actual_tokens] = merged_out.squeeze(0)
elif spec_sequence_masks is not None:
core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
else: