[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:
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:
|
||||
|
||||
Reference in New Issue
Block a user