Kimi k2.5 MLA based eagle3

[jhaotingc]

Purpose

@IzzyPutterman is original author.

This allows for Eagles that share MLA instead of GQA for attention, so one can train Eagle3s for Kimi and Deepseek and use them across TRTLLM, SGL, and vLLM.

Test Plan

Acc benchmark:

lm_eval \
  --model local-completions \
  --model_args base_url=http://my_server:8001/v1/completions,model=/trt_llm_ci/data/llm-models/Kimi-K2.5-NVFP4,num_concurrent=16,tokenized_requests=False,trust_remote_code=True \
  --tasks gsm8k \
  --batch_size 16

Test Result

without Eagle3

local-completions ({'base_url': 'http://my_server:8001/v1/completions', 'model': '/trt_llm_ci/data/llm-models/Kimi-K2.5-NVFP4', 'num_concurrent': 16, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: None, batch_size: 16
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9295|±  |0.0071|
|     |       |strict-match    |     5|exact_match|↑  |0.9295|±  |0.0071|

with Eagle3

local-completions ({'base_url': 'http://my_server:8001/v1/completions', 'model': '/trt_llm_ci/data/llm-models/Kimi-K2.5-NVFP4', 'num_concurrent': 16, 'tokenized_requests': False}), gen_kwargs: ({}), limit: None, num_fewshot: None, batch_size: 16
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9265|±  |0.0072|
|     |       |strict-match    |     5|exact_match|↑  |0.9257|±  |0.0072|

Acceptance:

  --- Weighted averages (by accepted tokens) ---
  Mean acceptance length:   2.785
  Avg draft accept rate:    59.5%
  Per-position accept rate: [0.826, 0.594, 0.365]

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jhaotingc.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request adds support for Eagle3 speculative decoding for Deepseek and Kimi models. The changes include a new deepseek_eagle3.py model implementation, and modifications to deepseek_v2.py to support auxiliary hidden state extraction. The configuration and model registry are also updated. My review focuses on a logical issue in the new deepseek_eagle3.py file where a condition is always false, leading to dead code and a bypassed assertion. I've provided a suggestion to fix this.

_{Note: Security Review did not run due to the size of the PR.}

[benchislett]

See #35966: why does this PR need a new model implementation and the other one doesn't?

[IzzyPutterman]

See #35966: why does this PR need a new model implementation and the other one doesn't?

This one implements the MLA based eagle3 not GQA like the other PR
This is a rebase of my ancient PR: #30574

[benchislett]

I was able to get this working locally. Seems fine but https://huggingface.co/lightseekorg/kimi-k2.5-eagle3 does not run. Two other EAGLE heads, one internal and one public, work fine. Assuming there's an unrelated config issue with the broken one.

[jhaotingc]

The failed tests seem unrelated to our change @benchislett

[benchislett]

will rerun failed tests a couple times and update from main once again if that doesn't work. failing that, we'll force-merge tomorrow.

[leihuang-sketch]

@jhaotingc In my tests, I found that there was no improvement in performance, and the throughput remained almost the same without additional parameters for eagle3. Could it be because the output of the gsm8k dataset is too short?

[leihuang-sketch]

LGTM