From 8ae328e5f0425a3c80ae930fd761400d3cdb3d22 Mon Sep 17 00:00:00 2001 From: Bi Xue Date: Tue, 9 Jun 2026 11:26:48 -0700 Subject: [PATCH] [sgl] Fix kimi-k2.5 EAGLE3 MLA draft embeds for batched MM prefill (#27647) --- python/sglang/srt/models/kimi_k25_eagle3.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/models/kimi_k25_eagle3.py b/python/sglang/srt/models/kimi_k25_eagle3.py index 621f554b4..42f5db3e1 100644 --- a/python/sglang/srt/models/kimi_k25_eagle3.py +++ b/python/sglang/srt/models/kimi_k25_eagle3.py @@ -261,9 +261,10 @@ class Eagle3MLAModel(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: