MLA Based Eagle3

[IzzyPutterman]

Purpose

Test Plan

Test Result

---

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.

[chatgpt-codex-connector]

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

[gemini-code-assist]

Code Review

This pull request adds support for MLA-based Eagle3 speculative decoding for DeepseekV2/V3 models. The changes include adding a new deepseek_eagle3.py model implementation, updating configurations, and modifying the DeepseekV2 model to support auxiliary hidden state extraction for Eagle3.

I've found a critical issue in vllm/model_executor/models/deepseek_v2.py which includes a syntax error in an enumerate call and a likely logic error where an argument to a layer call was omitted. I've provided a code suggestion to fix this. The rest of the changes look good and correctly integrate the new model.

[gemini-code-assist]

There are a couple of issues in this loop:

1. There's a syntax error in the enumerate call. The start argument is misplaced and there's a missing comma. It should be enumerate(islice(...), start=self.start_layer).
2. The llama_4_scaling argument is no longer passed to the layer call, but the logic to compute it is still present. This seems like an accidental omission and could lead to incorrect behavior.

I've provided a suggestion to fix both issues.

Suggested change

	for idx, layer in enumerate(
	islice(self.layers, self.start_layer, self.end_layer)
	start=self.start_layer,
	):
	if idx in self.aux_hidden_state_layers:
	aux_hidden_states.append(hidden_states + residual)
	hidden_states, residual = layer(positions, hidden_states, residual)
	for idx, layer in enumerate(
	islice(self.layers, self.start_layer, self.end_layer),
	start=self.start_layer,
	):
	if idx in self.aux_hidden_state_layers:
	aux_hidden_states.append(hidden_states + residual)
	hidden_states, residual = layer(positions, hidden_states, residual,
	llama_4_scaling)

[mergify]

Hi @IzzyPutterman, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[mgoin]

Hi @IzzyPutterman do you have a model or usecase where this is useful? It would be helpful for review if you filled out the description and testing you did

[IzzyPutterman]

Hi @IzzyPutterman do you have a model or usecase where this is useful? It would be helpful for review if you filled out the description and testing you did

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. We have a Kimi K2 thinking Eagle3 internally, which might get released (not sure here).

[mergify]

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

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