Fix Kimi K2.5 dp attention+ spec decoding launch crash (#21391)

python/sglang/srt/models/llama_eagle3.py

@@ -150,7 +150,18 @@ class LlamaModel(nn.Module):
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> torch.Tensor:
         if input_embeds is None:
-            embeds = self.embed_tokens(input_ids)
+            embeds = forward_batch.mm_input_embeds
+            if (
+                forward_batch.forward_mode.is_extend()
+                and forward_batch.contains_mm_inputs()
+                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))]
+                )
+            if embeds is None:
+                embeds = self.embed_tokens(input_ids)
         else:
             embeds = input_embeds
 

test/registered/8-gpu-models/test_kimi_k25.py

@@ -38,11 +38,15 @@ class TestKimiK25(unittest.TestCase):
             "--speculative-num-steps=3",
             "--speculative-eagle-topk=1",
             "--speculative-num-draft-tokens=4",
-            "--mem-frac=0.85",
             "--model-loader-extra-config",
             '{"enable_multithread_load": true, "num_threads": 64}',
         ]
 
+        dp_attn_args = [
+            "--dp=8",
+            "--enable-dp-attention",
+        ]
+
         variants = [
             ModelLaunchSettings(
                 KIMI_K25_MODEL_PATH,
@@ -56,6 +60,12 @@ class TestKimiK25(unittest.TestCase):
                 extra_args=base_args + eagle3_args,
                 variant="TP8+MTP",
             ),
+            ModelLaunchSettings(
+                KIMI_K25_MODEL_PATH,
+                tp_size=8,
+                extra_args=base_args + dp_attn_args + eagle3_args,
+                variant="TP8+DP8+MTP",
+            ),
         ]
 
         run_combined_tests(