diff --git a/python/sglang/srt/models/llama_eagle3.py b/python/sglang/srt/models/llama_eagle3.py index dce33dcd4..a4022dba0 100644 --- a/python/sglang/srt/models/llama_eagle3.py +++ b/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 diff --git a/test/registered/8-gpu-models/test_kimi_k25.py b/test/registered/8-gpu-models/test_kimi_k25.py index 18e6b3c2a..5f903f102 100644 --- a/test/registered/8-gpu-models/test_kimi_k25.py +++ b/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(