From 24c5d76f74004d16e697b23e4bf970e55d4182ae Mon Sep 17 00:00:00 2001 From: Jackey Hua <107608053+zhendonghua@users.noreply.github.com> Date: Thu, 11 Jun 2026 13:33:16 -0700 Subject: [PATCH] fix: per-sequence last-token embedding in EAGLE3/MTP draft for batched multimodal spec decoding (#27846) --- python/sglang/srt/models/llama_eagle3.py | 7 ++++--- python/sglang/srt/models/qwen3_5_mtp.py | 10 +++++----- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/python/sglang/srt/models/llama_eagle3.py b/python/sglang/srt/models/llama_eagle3.py index 3ce581943..7cd61389f 100644 --- a/python/sglang/srt/models/llama_eagle3.py +++ b/python/sglang/srt/models/llama_eagle3.py @@ -198,9 +198,10 @@ class LlamaModel(nn.Module): and not forward_batch.forward_mode.is_draft_extend(include_v2=True) ): assert embeds is not None - embeds = torch.cat( - [embeds[:-1], self.embed_tokens(input_ids[-1].unsqueeze(0))] - ) + last_indices = ( + 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: embeds = self.embed_tokens(input_ids) else: diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index 6b39a590c..c699a79aa 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -163,11 +163,11 @@ class Qwen3_5ForCausalLMMTP(nn.Module): and not forward_batch.forward_mode.is_draft_extend(include_v2=True) ): assert input_embeds is not None - input_embeds = torch.cat( - [ - input_embeds[:-1], - self.model.embed_tokens(input_ids[-1].unsqueeze(0)), - ] + last_indices = ( + forward_batch.extend_start_loc + forward_batch.extend_seq_lens - 1 + ).long() + input_embeds[last_indices] = self.model.embed_tokens( + input_ids[last_indices] ) if input_embeds is None: