[Quantization] [Eagle] Add complete quantization support to the draft model in Eagle by shreyas269 · Pull Request #28435 · vllm-project/vllm

shreyas269 · 2025-11-11T01:47:45Z

Purpose

This PR adds comprehensive quantization support for Eagle and Eagle3 draft models in speculative decoding, including full KV cache quantization support. Previously, Eagle draft models could not use quantized weights in fully connected layers, or quantized KV caches.

Recently, #26590 was merged to properly obtain the draft model's quantization config but it doesn't address the case where the entire draft model is quantized and we want to read input and weight scales of fc layer along with KV cache quantization scales.

This PR addresses the following:

Define get_draft_quant_config in utils to avoid duplication of code in llama_eagle.py and llama_eagle3.py.
Add ReplicatedLinear class to make fc layer in drafters quantizable and handle input and weight quantization scales (in llama with eagle/eagle3).
Handle and load KV cache quantization scales. Additionally, attempts to remap it to the expected name format in the model.

This PR is a duplicate to #27434 with an additional smoke test.

Test Plan

Tested with a base llama3 instruct model with a quantized Eagle draft model (one decoder layer + one FC layer) with static fp8 quantization. The quantization of the base/verifier and Eagle draft model was performed using ModelOpt.

The non-quantized models work exactly the same as before (no changes to behavior).

Added smoke tests
- New test file test_eagle_quantization.py with unit tests for draft model quantization
- Tests cover get_draft_quant_config() with and without draft models
- Tests verify FC layer behavior with quantization configs
- Tests check KV cache scale name handling and remapping

Test Result

Before:
KeyError: 'fc.input_scale'

After:

(Worker_TP1 pid=1060248) WARNING 10-21 15:19:08 [modelopt.py:103] Detected ModelOpt fp8 checkpoint. Please note that the format is experimental and could change.
(Worker_TP1 pid=1060248) INFO 10-21 15:19:08 [default_loader.py:309] Loading weights took 0.03 seconds
(Worker_TP1 pid=1060248) INFO 10-21 15:19:08 [eagle.py:973] Assuming the EAGLE head shares the same vocab embedding with the target model.
(Worker_TP1 pid=1060248) INFO 10-21 15:19:08 [eagle.py:995] Loading EAGLE LM head weights from the target model.
(Worker_TP0 pid=1060247) WARNING 10-21 15:19:08 [modelopt.py:103] Detected ModelOpt fp8 checkpoint. Please note that the format is experimental and could change.
(Worker_TP0 pid=1060247) WARNING 10-21 15:19:08 [modelopt.py:103] Detected ModelOpt fp8 checkpoint. Please note that the format is experimental and could change.
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00, 36.17it/s]

Acceptance rate of drafter:

Per-position acceptance rate: 0.830, 0.675, 0.452, 0.309, Avg Draft acceptance rate: 56.6%

The test_eagle_quantization and all pre-commit hooks pass successfully.

tests/model_executor/test_eagle_quantization.py::test_get_draft_quant_config_with_draft_model PASSED
tests/model_executor/test_eagle_quantization.py::test_get_draft_quant_config_without_draft_model PASSED
tests/model_executor/test_eagle_quantization.py::test_fc_layer_quant_config_usage[cuda:0] WARNING 10-30 16:20:51 [vllm.py:1005] Current vLLM config is not set.
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
WARNING 10-30 16:20:54 [vllm.py:1005] Current vLLM config is not set.
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
INFO 10-30 16:20:54 [parallel_state.py:1325] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0, EP rank 0
WARNING 10-30 16:20:54 [vllm.py:1005] Current vLLM config is not set.
PASSED
tests/model_executor/test_eagle_quantization.py::test_fc_layer_quant_config_usage[cuda:1] WARNING 10-30 16:20:55 [vllm.py:1005] Current vLLM config is not set.
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
WARNING 10-30 16:20:55 [vllm.py:1005] Current vLLM config is not set.
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0
PASSED
tests/model_executor/test_eagle_quantization.py::test_kv_cache_scale_name_handling PASSED
tests/model_executor/test_eagle_quantization.py::test_kv_cache_scale_name_no_scale PASSED
tests/model_executor/test_eagle_quantization.py::test_maybe_remap_kv_scale_name PASSED
tests/model_executor/test_eagle_quantization.py::test_load_weights_kv_scale_handling PASSED
========================================== 8 passed, 2 warnings in 7.70s ==========================================

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.

… model in Eagle Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

gemini-code-assist

Code Review

This pull request introduces comprehensive quantization support for Eagle and Eagle3 draft models, which is a significant enhancement. The changes are well-structured, including the refactoring of get_draft_quant_config into a shared utility to reduce code duplication. The replacement of torch.nn.Linear with ReplicatedLinear is correctly implemented to enable quantization for the fully connected layers. The addition of a dedicated test file with smoke tests for the new quantization features is also a great contribution that improves the robustness of the codebase.

I have one suggestion regarding code duplication in the load_weights methods of llama_eagle.py and llama_eagle3.py. Extracting the logic for handling KV cache scales into a shared function would further improve maintainability.

Overall, this is a well-executed PR that successfully adds an important feature.

gemini-code-assist · 2025-11-11T01:49:09Z

vllm/model_executor/models/llama_eagle.py

+            # Handle kv cache quantization scales
+            if self.quant_config is not None and (
+                scale_name := self.quant_config.get_cache_scale(name)
+            ):
+                # Loading kv cache quantization scales
+                param = params_dict[scale_name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                loaded_weight = (
+                    loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
+                )
+                weight_loader(param, loaded_weight)
+                loaded_params.add(scale_name)
+                continue
+            # Remapping the name FP8 kv-scale
+            if "scale" in name:
+                name = maybe_remap_kv_scale_name(name, params_dict)
+                if name is None:
+                    continue


This block of code for handling KV cache quantization scales and remapping FP8 kv-scale names is duplicated in vllm/model_executor/models/llama_eagle3.py (lines 220-237). To improve maintainability and avoid potential inconsistencies in the future, consider refactoring this logic into a shared utility function in vllm/model_executor/models/utils.py. You've already done this for get_draft_quant_config, and a similar approach would be beneficial here.

heheda12345 · 2025-11-12T08:25:21Z

CC @mgoin @Isotr0py @yewentao256 for quantization and @luccafong @benchislett for eagle

Isotr0py

The quantization part looks reasonable to me. See if others are fine with quantzaion for eagle. (I don't have much knowledge about spec decode)

Isotr0py · 2025-11-12T08:50:45Z

tests/model_executor/test_eagle_quantization.py

Seems that we should put this test under tests/models/quantization or tests/quantization?

@Isotr0py, Thanks for the review! I placed the test in tests/model_executor because this is testing the model executor's quantization integration for the Eagle model architecture, rather than end-to-end inference with quantized models.

That said, I'm happy to move it if you think tests/models/quantization or tests/quantization would be better for discoverability.

shreyas269 · 2025-11-14T15:38:13Z

Hey folks, just a quick follow-up on this PR. The reviewers were tagged previously (tagging them again below), and I’m happy to incorporate any changes needed. Whenever someone has time to take a look, I’d appreciate it! :)

CC @mgoin @Isotr0py @yewentao256 @luccafong @benchislett @rahul-tuli

mgoin

Looks good to me, although I think the test can be improved to use a real model

mgoin · 2025-11-14T17:28:18Z

tests/model_executor/test_eagle_quantization.py

I don't love all this Mocking and working with low-level classes and configs. I'd like to rather replace this with a single model e2e smoke test so we see that the config parsing, weight loading, and such all work together with a known model

I agree that using a single real model would be cleaner. The only reason I didn’t do that here is because our internal model isn’t publicly shareable, so I couldn’t include it in the test. Although we could get a public drafter model released but could take time to go through approval since I'm contributing via capital one).

Given that constraint, this felt like the most straightforward solution. Happy to update this if you have a suggestion for a better approach.

Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

mergify · 2025-11-15T17:12:26Z

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

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

Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

rahul-tuli · 2025-11-17T15:42:43Z

The PR looks good considering the constraints for using a real model, a follow up PR with a real model would be nice. Could you also fix the pre-commit checks?

mgoin · 2025-11-17T16:59:50Z

@shreyas269 please fix the pre-commit as I cannot do it for you due to branch permissions, thank you!

Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

shreyas269 · 2025-11-17T19:53:04Z

@mgoin, @rahul-tuli, fixed the pre-commit.

… model in Eagle (vllm-project#28435) Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

[Eagle] [Quantization] Add complete quantization support to the draft…

625bf1b

… model in Eagle Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

mergify bot added llama Related to Llama models speculative-decoding labels Nov 11, 2025

gemini-code-assist bot reviewed Nov 11, 2025

View reviewed changes

shreyas269 mentioned this pull request Nov 11, 2025

[Eagle] [Quantization] Add complete quantization support to the draft model in Eagle #27434

Closed

5 tasks

Isotr0py reviewed Nov 12, 2025

View reviewed changes

benchislett requested review from benchislett and luccafong November 14, 2025 17:18

mgoin added quantization ready ONLY add when PR is ready to merge/full CI is needed labels Nov 14, 2025

mgoin approved these changes Nov 14, 2025

View reviewed changes

Fix mkdocs - explicity specify return type

93c6208

Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

mergify bot added the needs-rebase label Nov 15, 2025

Merge branch 'main' into vllm-eagle-quantize-shreyas269

4425797

Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

mergify bot removed the needs-rebase label Nov 15, 2025

fix pre-commit

2f16dd5

Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

vllm-bot merged commit 95ae50b into vllm-project:main Nov 17, 2025
48 of 50 checks passed

devpatelio pushed a commit to SumanthRH/vllm that referenced this pull request Nov 29, 2025

[Quantization] [Eagle] Add complete quantization support to the draft…

c308ef3

… model in Eagle (vllm-project#28435) Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

kitaekatt pushed a commit to kitaekatt/vllm that referenced this pull request Dec 1, 2025

[Quantization] [Eagle] Add complete quantization support to the draft…

6b88eb7

… model in Eagle (vllm-project#28435) Signed-off-by: Shreyas Kulkarni <shreyas.gp269@gmail.com>

Uh oh!

Conversation

shreyas269 commented Nov 11, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Test Plan

Test Result

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Nov 11, 2025

Choose a reason for hiding this comment

Uh oh!

heheda12345 commented Nov 12, 2025

Uh oh!

Isotr0py left a comment

Choose a reason for hiding this comment

Uh oh!

Isotr0py Nov 12, 2025

Choose a reason for hiding this comment

Uh oh!

shreyas269 Nov 12, 2025

Choose a reason for hiding this comment

Uh oh!

shreyas269 commented Nov 14, 2025

Uh oh!

mgoin left a comment

Choose a reason for hiding this comment

Uh oh!

mgoin Nov 14, 2025

Choose a reason for hiding this comment

Uh oh!

shreyas269 Nov 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Nov 15, 2025

Uh oh!

rahul-tuli commented Nov 17, 2025

Uh oh!

mgoin commented Nov 17, 2025

Uh oh!

shreyas269 commented Nov 17, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

shreyas269 commented Nov 11, 2025 •

edited by github-actions bot

Loading

shreyas269 Nov 14, 2025 •

edited

Loading