[Spec Decode][MRV2] Support heterogeneous-vocab draft models (TLI)

[zmy1260222041]

Overview

Draft models with a different vocabulary than the target, use_heterogeneous_vocab=True from #38174, now run on Model Runner V2. Since #59541 they fall back to V1. This is the TLI item listed as TBD in #47172.

Claims

• method="draft_model" with use_heterogeneous_vocab=True runs on V2, and the V1 fallback is removed.
• On both directions of vocab mismatch, acceptance length on V2 matches V1 and GSM8K accuracy stays within noise of the target alone.
• V2 throughput with TLI is the same as or slightly above V1.

Validation

Tests I ran on main at 624a5dd with this change. The GPU tests ran on 1x L4.

pytest tests/test_config.py -k v2_model_runner -q
# 32 passed
pytest tests/v1/spec_decode/test_vocab_mapping.py -q
# 17 passed
pytest tests/v1/e2e/spec_decode/draft_model/test_draft_model.py -k heterogeneous_vocab -v -s
# Using V2 Model Runner
# acceptance_len=1.920
# exact text matches=96/100 (required=61)
# 1 passed
pytest tests/v1/e2e/spec_decode/draft_model/test_draft_model.py -k "correctness or realistic" -v
# 5 passed

Without the change, the new e2e test fails with Model Runner V2 does not yet support: heterogeneous-vocabulary draft models.

GSM8K, 500 questions, greedy, k=3, all prompts submitted at once, 1x L4. AL is the mean acceptance length from the spec decode metrics. Throughput is output tokens per second over the whole batch.

Target	Draft	Runner	GSM8K	AL	tok/s
Qwen2.5-1.5B-Instruct	none	V1	0.532		3666
Qwen2.5-1.5B-Instruct	none	V2	0.532		3724
Qwen2.5-1.5B-Instruct	SmolLM2-360M-Instruct	V1	0.542	2.11	1771
Qwen2.5-1.5B-Instruct	SmolLM2-360M-Instruct	V2	0.524	2.11	1832
SmolLM2-1.7B-Instruct	none	V1	0.458		1859
SmolLM2-1.7B-Instruct	none	V2	0.458		1816
SmolLM2-1.7B-Instruct	Qwen2.5-0.5B-Instruct	V1	0.448	3.13	2281
SmolLM2-1.7B-Instruct	Qwen2.5-0.5B-Instruct	V2	0.454	3.12	2311

The first pair has a target vocab larger than the draft's (151,936 vs 49,152), and the second pair is the reverse. These are the two pairs from #59541. The SmolLM2 target ends right away on raw few-shot prompts, so its rows use chat prompts. The standard error at 500 questions is about 2 points. With the 1.5B target at this batch size the draft is extra work, so throughput with drafting is lower than the target alone. That is the same on V1.

Greedy outputs vs the target alone

Greedy outputs with drafting do not match the target alone token for token on either runner, for any draft. Verify computes k+1 rows at once and plain decode computes one, so near-tied logits can flip, see #54506. To check that TLI adds nothing on top of that, the table compares each run with the target alone on the same runner, next to a same-vocab draft as a control, and looks at the first diverging token of every mismatch.

Target	Draft	Runner	Same as target	First divergence at a top-2 tie or one bf16 step apart
Qwen2.5-1.5B-Instruct	SmolLM2-360M-Instruct	V2	247/500	85%
Qwen2.5-1.5B-Instruct	Qwen2.5-0.5B-Instruct	V2	242/500	80%
SmolLM2-1.7B-Instruct	Qwen2.5-0.5B-Instruct	V2	289/500	95%
SmolLM2-1.7B-Instruct	SmolLM2-360M-Instruct	V2	296/500	96%

Reproduce

# Run once with VLLM_USE_V2_MODEL_RUNNER=0 and once with =1, from the repo root.
from vllm import LLM
from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
from tests.v1.e2e.spec_decode.utils import compute_acceptance_len

llm = LLM(
    "Qwen/Qwen2.5-1.5B-Instruct",
    max_model_len=4096,
    seed=0,
    disable_log_stats=False,
    speculative_config={
        "method": "draft_model",
        "model": "HuggingFaceTB/SmolLM2-360M-Instruct",
        "num_speculative_tokens": 3,
        "use_heterogeneous_vocab": True,
    },
)
print(evaluate_gsm8k_offline(llm, num_questions=500))
print(compute_acceptance_len(llm.get_metrics()))

For the SmolLM2 target, pass use_chat_completions=True to evaluate_gsm8k_offline.

Details

The change keeps the draft vocab inside StandaloneARSpeculator. Ids cross between vocabularies at two points only:

1. Target ids become draft ids once per step, in the expanded step-0 input buffer. This is done in place so the draft CUDA graph work in [MRV2][Spec Decode]Adapt cudagraph for draft model #60989 can keep reading that buffer. The dummy-run path reads the target's own input buffer, so it maps into a new tensor instead.
2. propose() maps the drafted ids to target ids, into a preallocated buffer, right before returning them.

compute_draft_logits masks logits outside the shared vocab, so the drafter only proposes tokens the target also has. Between draft steps the drafter keeps its own ids. V1 maps back and forth every step, but ids in the shared vocab map one to one, so the drafter sees the same inputs on both runners.

Tokenizers are loaded with each model's own path, mode and revision. The draft config inherits the target's tokenizer mode and revision, which #54780 fixes on the V1 side.

use_local_argmax_reduction bypasses compute_draft_logits, so the mask would be skipped. #51564 already rejects that combination in the config, so it is not repeated here.

Not in this PR: CUDA graphs for the standalone drafter, which #60989 is adding, and probabilistic drafting with a heterogeneous vocab, which needs the draft distribution mapped into the target vocab and is also unsupported on V1.

Duplicate check: #47172 and the open PRs for TLI and heterogeneous vocab on V2. #54780, #51564 and #59012 fix V1 or VocabMapping, and #61152 fixes probabilistic rejection sampling for narrower draft vocabularies. None of them ports TLI to V2.

This PR was written with AI assistance from Claude Code. I reviewed every changed line and ran the tests above myself.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

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.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀