执行摘要

本PR修复了Kimi K2.5模型在启用分布式注意力和推测解码时，因多模态输入嵌入处理不当而导致的启动崩溃问题。通过调整llama_eagle3.py中的forward函数逻辑，并添加相应测试，确保了系统稳定性，但review中指出了代码可读性和assert使用的潜在风险。

功能与动机

该PR旨在解决issue #21336中报告的bug，即当存在多模态输入时，draft embedding无法正确处理，引发边界错误。作者在PR body中明确指出：'This issue is if there are mm input, draft embedding cannot handle it (will have out of boundary issue)'，并参考了qwen3_5_mtp.py中的现有修复模式，以确保一致性。

实现拆解

主要改动集中在两个文件：

• python/sglang/srt/models/llama_eagle3.py：在forward函数中添加了条件逻辑，处理多模态输入嵌入。关键代码如下：
  

  if input_embeds is None:
    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 # 存在风险的assert
      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：新增测试变体'TP8+DP8+MTP'，验证在分布式注意力配置下的模型启动和运行，确保修复覆盖相关场景。

评论区精华

review中，gemini-code-assist[bot]提出了关键反馈：

• 'The assert embeds is not None on this line is problematic... making this assert redundant and potentially dangerous.'（正确性风险）

• 'The logic for determining the embeddings is a bit nested and could be hard to follow. Refactoring it into an if/elif/else structure would make the different cases clearer.'（设计改进）

• 'Commented-out code reduces readability... please remove them.'（代码风格）
  这些讨论揭示了代码质量改进点，但PR已获批准，未在本次变更中完全解决。

风险与影响

风险：

• assert语句可能在不必要时触发运行时错误，影响系统稳定性。
• 嵌套逻辑降低了代码可读性，可能增加后续维护成本。
• 测试文件中的注释代码可能导致混淆，需清理以避免误解。

影响：

• 对用户：修复了崩溃问题，提升了Kimi K2.5模型在复杂配置下的可用性。
• 对系统：嵌入处理逻辑变更遵循现有模式，降低了回归风险，但需监控性能影响。
• 对团队：强调了代码review中关注正确性和可维护性的重要性。

关联脉络

与此PR相关的历史PR包括#21004（'[Fix] Add EPLB rebalance support for Kimi K2.5'），后者也针对Kimi K2.5模型进行了修复。这显示该模型在推测解码和多模态输入场景下存在多个稳定性问题，团队正在通过渐进式修复持续优化。整体上，这些变更反映了对高性能推理场景中边缘案例处理的重视。