From 975dca5bb5db302077674cfa9afe851ee700ad73 Mon Sep 17 00:00:00 2001 From: Arpan Tripathi <42506819+tripathiarpan20@users.noreply.github.com> Date: Wed, 16 Sep 2026 14:03:57 +0100 Subject: [PATCH] [Perf][GDN] Scatter mixed speculative outputs into the caller buffer (#57140) Signed-off-by: Arpan Tripathi Signed-off-by: zjy0516 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 Co-authored-by: Codex --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index bbbb1c8bd5..a57c984db4 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -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: