[sgl] Fix kimi-k2.5 EAGLE3 MLA draft embeds for batched MM prefill (#27647)

2026-06-09 11:26:48 -07:00
parent 5babb902a9
commit 8ae328e5f0
1 changed files with 4 additions and 3 deletions
@@ -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: