minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

fix: per-sequence last-token embedding in EAGLE3/MTP draft for batched multimodal spec decoding (#27846)

Browse Source
This commit is contained in:
Jackey Hua
2026-06-11 13:33:16 -07:00
committed by GitHub
parent 10219bd9d6
commit 24c5d76f74
2 changed files with 9 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/models/llama_eagle3.py
+4 -3
View File
@@ -198,9 +198,10 @@ class LlamaModel(nn.Module):
and not forward_batch.forward_mode.is_draft_extend(include_v2=True) and not forward_batch.forward_mode.is_draft_extend(include_v2=True)
): ):
assert embeds is not None assert embeds is not None
embeds = torch.cat( last_indices = (
[embeds[:-1], self.embed_tokens(input_ids[-1].unsqueeze(0))] forward_batch.extend_start_loc + forward_batch.extend_seq_lens - 1
) ).long()
embeds[last_indices] = self.embed_tokens(input_ids[last_indices])
if embeds is None: if embeds is None:
embeds = self.embed_tokens(input_ids) embeds = self.embed_tokens(input_ids)
else: else:
python/sglang/srt/models/qwen3_5_mtp.py
+5 -5
View File
@@ -163,11 +163,11 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
and not forward_batch.forward_mode.is_draft_extend(include_v2=True) and not forward_batch.forward_mode.is_draft_extend(include_v2=True)
): ):
assert input_embeds is not None assert input_embeds is not None
input_embeds = torch.cat( last_indices = (
[ forward_batch.extend_start_loc + forward_batch.extend_seq_lens - 1
input_embeds[:-1], ).long()
self.model.embed_tokens(input_ids[-1].unsqueeze(0)), input_embeds[last_indices] = self.model.embed_tokens(
] input_ids[last_indices]
) )
if input_embeds is None: if input_embeds is None: