Skip to content

[Bugfix][Spec Decode] Honour the draft's load_config on Model Runner V2 - #55337

Open
stecasta wants to merge 2 commits into
vllm-project:mainfrom
stecasta:fix/v2-draft-load-config
Open

stecasta wants to merge 2 commits into
vllm-project:mainfrom
stecasta:fix/v2-draft-load-config

Conversation

@stecasta

@stecasta stecasta commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

Purpose

load_eagle_model() builds the draft with get_model() but never passes load_config:

eagle_model = get_model(vllm_config=vllm_config, model_config=draft_model_config)

get_model() resolves that to load_config or vllm_config.load_config, so the draft silently inherits the target's load config and SpeculativeConfig.draft_load_config ("Load config for the draft model. If not specified, will use the load config from the target model", vllm/config/speculative.py) has no effect on Model Runner V2.

V1 passes it in LLMBaseProposer._get_model(), and V2's own Gemma4 speculator passes it in vllm/v1/worker/gpu/spec_decode/gemma4/speculator.py, so the eagle path is the outlier rather than a deliberate choice.

Fix

Forward speculative_config.draft_load_config. It defaults to None, which get_model() already resolves to the target's load config, so nothing changes when it is unset.

Validation

pytest tests/v1/spec_decode/test_draft_load_config.py -v   # 2 passed
pre-commit run --files <changed files>                     # all hooks passed

Run inside vllm/vllm-openai:nightly-8a728663c1 (0.28.1rc1.dev388+g8a728663c), unpatched then patched against the same installed package. Without the fix the first test fails with the bug's signature:

assert None is <object object at 0xfbde13c30f70>
 +  where None = _capture_draft_load_config(<object object at 0xfbde13c30f70>)

The second test pins the unset case so the fallback to the target's load config stays intact.

Why this is not a duplicate

I ran the checks in AGENTS.md. Nothing open touches draft_load_config on the V2 path. #54788 and #54826 also edit load_eagle_model for moe_backend and attention_backend respectively; this is a third, independent field and the three will need a trivial rebase against each other depending on merge order.

Model evaluation

Not applicable. The draft's load config selects how its weights are read (load format, download dir), not what they are. With draft_load_config unset, which is the default, behaviour is byte-identical, and that case is pinned by the second test.

AI assistance

AI assistance was used to author this change. I have reviewed every changed line, ran the tests and linters above myself, and can defend the design.

load_eagle_model() builds the draft with get_model() but never passes
load_config, so get_model() falls back to the target's and
SpeculativeConfig.draft_load_config is silently ignored. V1 passes it in
LLMBaseProposer._get_model(), and V2's own gemma4 speculator passes it,
so the eagle path is the outlier.

Forward speculative_config.draft_load_config. It defaults to None, which
get_model() already resolves to the target's load config, so behaviour is
unchanged when it is unset.

Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

@mergify mergify Bot added speculative-decoding mrv2 Model Runner V2 specific bug Something isn't working labels Sep 4, 2026
@coderabbitai

coderabbitai Bot commented Sep 4, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: cc3480aa-ad25-496e-bdb7-67c5cb5d24a7

📥 Commits

Reviewing files that changed from the base of the PR and between 619e85b and ceb4338.

📒 Files selected for processing (1)
  • tests/v1/spec_decode/test_draft_load_config.py

Included review availability: Your plan provides up to 10 included reviews per hour; 8 remain after this review.


📝 Summary

Summary by CodeRabbit

  • Bug Fixes

    • Draft models now use their dedicated loading configuration when speculative decoding is enabled.
    • Configurations without a draft loading setting continue to defer to the target model’s defaults.
  • Tests

    • Added coverage verifying that draft loading settings are correctly passed through, including when no setting is provided.

Walkthrough

load_eagle_model now passes speculative_config.draft_load_config to get_model. New tests verify forwarding for configured and unset values.

Changes

Draft model load configuration

Layer / File(s) Summary
Forward draft load configuration
vllm/v1/worker/gpu/spec_decode/eagle/utils.py, tests/v1/spec_decode/test_draft_load_config.py
load_eagle_model passes draft_load_config to get_model. Tests verify explicit values and None.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to ceb43

Draft models now honor their configured load settings while unset settings retain existing target-config fallback behavior. The covered forwarding paths present no remaining merge-blocking risk.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 2 files. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the bugfix and the affected Model Runner V2 speculative decoding behavior.
Description check ✅ Passed The description accurately explains the defect, the fix, compatibility behavior, validation, and scope.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
  • Fix all pre-merge checks with AI

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@sylvesterkaczmarek sylvesterkaczmarek left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct propagation boundary. get_model() already defines the None -> target load_config fallback, so forwarding draft_load_config here restores the explicit override without changing default behaviour. The sentinel/None pair pins both sides of the contract.

load_eagle_model reads speculative_config.attention_backend once this
branch picks up main, so the stub needs the field to keep working.

Signed-off-by: Stefano Castagnetta <scastagnetta@nvidia.com>

@sylvesterkaczmarek sylvesterkaczmarek left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Re-reviewed current ceb4338. The latest change only extends the draft load-config test fixture with attention_backend; the core handoff still passes the draft-specific load_config through to the loader, while None correctly defers to the target/default resolution. No blocker from me on this head.

@stecasta

stecasta commented Sep 8, 2026

Copy link
Copy Markdown
Contributor Author

Hi @njhill, this should be the last one in the series :)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working mrv2 Model Runner V2 specific speculative-decoding

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants