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: