fix: Domino inference patches for speculators-format checkpoints by orestis-z · Pull Request #48241 · vllm-project/vllm

orestis-z · 2026-07-10T11:08:51Z

Purpose

Follow-up to #47627 (eros483/feature/domino-head-for-dflash). Two fixes needed for correct Domino inference:

Fix 1 — `speculator.py`: step gap between prefix and correction start (bug, affects all Domino checkpoints)

With correction_start = prefix_len + 1, the original loop range(correction_start, num_steps) leaves positions [prefix_len, correction_start) unfilled — draft_tokens stays at token_id 0. For Huang2020/Qwen3-8B-Domino-b16 (pure_draft_prefix_len=1), step 1 is always token 0, collapsing acceptance length to ~1.0 across all benchmarks.

Fix: loop from prefix_len instead, using pure base logits for positions below correction_start:

# Before
for step in range(correction_start, self.num_speculative_steps):
    logits = self.model.compute_domino_logits(...)

# After
for step in range(prefix_len, self.num_speculative_steps):
    if step < correction_start:
        logits = base_step_logits          # pure base, no correction
    else:
        logits = self.model.compute_domino_logits(...)

Fix 2 — `algos.py`: forward Domino fields for speculators-format checkpoints (feature)

Huang2020/Qwen3-8B-Domino-b16 embeds Domino fields directly in dflash_config. Speculators-format checkpoints (from vllm-project/speculators) put them at the top level. Without this forwarding, dflash_config.get("projector_type") returns None and the Domino head is never activated.

Validation

Huang2020/Qwen3-8B-Domino-b16 on 4×H100 (Qwen/Qwen3-8B verifier, TP=4, spec_tokens=15):

Subset	DFlash v5 baseline	Domino (Huang2020)	Delta
HumanEval	2.228	3.420	+53.5%
math_reasoning	2.307	3.873	+67.9%
qa	1.931	2.994	+55.0%
question	2.021	3.087	+52.7%
rag	1.943	3.270	+68.3%
summarization	1.816	2.579	+42.0%
tool_call	2.085	2.960	+42.0%
translation	1.943	3.029	+55.9%
writing	2.026	3.129	+54.4%
Average	2.033	3.149	+54.9%

Without Fix 1, acceptance collapses to ~1.0 across all subsets.

Checklist

Purpose filled in (follow-up to Feature/domino head for dflash #47627)
Test results provided
Linting passes

🤖 Generated with Claude Code
Co-authored-by: Claude Sonnet 4.6 noreply@anthropic.com

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.

mergify · 2026-07-10T11:09:26Z

Documentation preview: https://vllm--48241.org.readthedocs.build/en/48241/

mergify · 2026-07-14T05:55:20Z

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

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

mergify · 2026-07-15T20:34:09Z

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

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

…the serve rationale DOMINO is not an engine-level speculative algorithm: engines expose Domino as "dflash" and enable the causal correction head (prefix_gru + embed_proj) from the checkpoint dflash_config.projector_type="domino". A speculative_algorithm=DOMINO would leak into SGLang ServerArgs and fail cryptically, so mirror the vLLM guardrail and raise at the SGLang ServerArgs builder. Also correct both guardrail messages. The previous wording claimed Domino cannot be served and that its correction head is inert, which is wrong: serving with DFLASH keeps the Domino head active on engines that support it (vllm-project/vllm#48241, sgl-project/sglang#31328). Signed-off-by: khazic <khazzz1c@gmail.com>

* feat(backends): add Domino drafter training backend Port the Domino training path from NeMo AutoModel (dflash/domino_core.py) into the SpeCo overlay as a DFlash variant, mirroring how DSpark extends DFlash. Domino adds a causal correction head on top of the DFlash parallel block backbone: a single-layer GRU encodes a causal state from each block's previous tokens, and a low-rank embed_proj over [backbone hidden | GRU state] emits a full-vocab logit delta added to the parallel base logits. Training jointly supervises the refined (final) and backbone-only (base) logits with a base-anchor curriculum loss = (1-lambda)*final + lambda*base, lambda decaying to 0. The shifted-label alignment (target x[a+1:a+1+block], prev [x[a], labels[:-1]], every position supervised) reuses the DSpark alignment, which equals AutoModel's shift_label Domino path. DominoTrainerBackend subclasses DFlashTrainerBackend; only build_model and the training forward differ. Wired through worker dispatch, base_trainer block-drafter gates, auto config routing, oldlogprob aux layers, and config keys. Domino is training-only: its GRU correction has no stock vLLM proposer, so the vLLM config builder raises and directs serving to DFLASH (the trained backbone). AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com> * test(domino): add GPU hardware smoke for the Domino backend AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com> * fix(domino): compute dual-logit CE in fp32 and avoid full final-logits clone AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com> * style(domino): drop unused torch import AI assistance was used for this change. Signed-off-by: khazic <khazzz1c@gmail.com> * fix(domino): guard None drafter model_path in build_model os.path.join crashes with a TypeError when rollout.drafter.model_path is None (training a Domino drafter from scratch with no pre-existing weights). Guard config_path so it falls back to None and the existing checks route to the from-scratch fallback config path instead of crashing. Signed-off-by: khazic <khazzz1c@gmail.com> * test(domino): skip lambda-base test without torch to fix CPU CI test_domino_lambda_base_schedule imported get_lambda_base from domino_trainer_backend, whose module subclasses the torch-based DFlash backend at import time, so the torch-free CPU unit-test job hit ModuleNotFoundError: No module named 'torch'. Guard with pytest.importorskip like every other test in the file. Signed-off-by: khazic <khazzz1c@gmail.com> * feat(domino): guard the SGLang serve path against DOMINO and correct the serve rationale DOMINO is not an engine-level speculative algorithm: engines expose Domino as "dflash" and enable the causal correction head (prefix_gru + embed_proj) from the checkpoint dflash_config.projector_type="domino". A speculative_algorithm=DOMINO would leak into SGLang ServerArgs and fail cryptically, so mirror the vLLM guardrail and raise at the SGLang ServerArgs builder. Also correct both guardrail messages. The previous wording claimed Domino cannot be served and that its correction head is inert, which is wrong: serving with DFLASH keeps the Domino head active on engines that support it (vllm-project/vllm#48241, sgl-project/sglang#31328). Signed-off-by: khazic <khazzz1c@gmail.com> * fix(domino): compute top5_correct so top5_acc is not always zero top5_correct was initialized to zero and never reduced, so the top5_correct_count diagnostic (which base_trainer turns into top5_acc) stayed pinned at 0. DFlash and DSpark both compute it; Domino did not. Mirror the existing top1 idiom in this forward: take topk over the base logits and overwrite the suffix rows from the Domino-corrected logits, so no second [num_active, vocab] tensor is materialized. Guard topk with min(5, vocab) like DSpark does, for small-vocab configs. Adds a CPU regression test that fails without the fix (top5=0 vs top1=2). Signed-off-by: khazic <khazzz1c@gmail.com> * docs(domino): correct the serve rationale in the module docstring The docstring still claimed Domino is training-only and that its correction head has no engine proposer. Domino is a projector_type sub-mode of DFlash: the serve method stays dflash and the head is enabled from the checkpoint, so align this with the guardrails in vllm_runtime and sglang_runtime. Signed-off-by: khazic <khazzz1c@gmail.com> --------- Signed-off-by: khazic <khazzz1c@gmail.com>

Signed-off-by: swylyu001 <swylyu001@gmail.com> Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

…Head Bug 1: Correction applied one position too early. The Domino head was never trained to correct the anchor position. The paper (x4.1.1) and training code (SpecForge specforge/core/domino.py) both leave the leading 1 + pure_draft_prefix_len positions as pure base. Fix: start the correction loop at prefix_len + 1 instead of prefix_len. Bug 2: DominoHead.embed_proj final projection output target config.vocab_size instead of draft_vocab_size. Training code in SpecForge uses nn.Linear(emb_dim, draft_vocab_size). When draft_vocab != target_vocab, corrections are applied at wrong token indices and checkpoint loading fails. Fix: use config.draft_vocab_size. Refs: - Paper: Domino (arXiv:2605.29707) x4.1.1, x4.1.2 - Training code: sgl-project/SpecForge specforge/core/domino.py - Inference code: jianuo-huang/Domino code/dflash.py spec_generate() Co-authored-by: Cursor Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

Two fixes on top of Eros483's feature/domino-head-for-dflash (99c7879): 1. algos.py: forward Domino fields (projector_type, shift_label, pure_draft_prefix_len, gru_hidden_dim, emb_dim) from speculators top-level config into dflash_config. Required for speculators-format checkpoints (Huang2020 embeds these directly in dflash_config; speculators puts them at the top level and needs this forwarding). 2. speculator.py: fix step-0 gap in _generate_domino_draft. With correction_start = prefix_len + 1, positions [prefix_len, correction_start) were never written and stayed as token_id 0, causing ~0% acceptance. Now these positions use pure base logits before GRU correction begins. Affects all Domino checkpoints regardless of format. Verified: Huang2020/Qwen3-8B-Domino-b16 scores 3.149 avg acceptance (+54.9% vs DFlash v5 baseline of 2.033) with these patches applied. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>

…le speculator Two fixes: 1. qwen3_dflash.py: also route already-prefixed domino_head.* weights to the Domino loader. Our training saves with domino_head.* prefix; the previous code only handled bare prefix_gru.* / embed_proj.* names. 2. speculator.py: restore v0.24.0-compatible base with Domino additions applied on top. The previous version copied from Eros483 branch called _prepare_eplb_forward which does not exist in v0.24.0 base class. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>

Training applies Domino correction from position `pure_draft_prefix_len` onward, but inference was using `prefix_len + 1`, skipping one extra position. This mismatch reduced acceptance rates for all Domino checkpoints. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>

…ger compat Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>

Split compute_logits into compute_draft_logits (raw draft-space logits) and scatter_logits_to_target (d2t mapping to verifier space). The Domino draft generation now operates in draft space: base logits + correction both have shape [N, draft_vocab_size], and scattering to target space happens after the addition rather than before. This fixes a shape mismatch crash when serving Domino checkpoints trained with draft_vocab_size != target_vocab_size (e.g. 32000 vs 151936). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com> Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>

benchislett · 2026-07-21T13:56:40Z

Would like to see an nsys profile so we can make sure the overhead is minimal. Kernel fusion may be necessary here

jianuo-huang · 2026-07-21T15:51:28Z

Related SGLang Domino profiling and optimization results are available in sgl-project/sglang#31328 (see the four Perf links in the PR description). They cover base-logit precompute, full-rollout CUDA Graph, GRU feedback, and a K=2048 candidate pool, with component and end-to-end serving measurements.

mergify · 2026-07-30T16:16:19Z

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

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

orestis-z requested review from WoosukKwon, njhill, sighingnow, vadiklyutiy and yewentao256 as code owners July 10, 2026 11:08

claude Bot reviewed Jul 10, 2026

View reviewed changes

mergify Bot added documentation Improvements or additions to documentation qwen Related to Qwen models v1 labels Jul 10, 2026

orestis-z marked this pull request as draft July 13, 2026 11:10

orestis-z mentioned this pull request Jul 13, 2026

feat(dflash): add Domino causal correction head (projector_type=domino) vllm-project/speculators#773

Draft

4 tasks

mergify Bot added the needs-rebase label Jul 14, 2026

orestis-z force-pushed the domino-speculators-patches branch from 3f75ad7 to e44ae8b Compare July 14, 2026 14:38

mergify Bot removed the needs-rebase label Jul 14, 2026

mergify Bot added the needs-rebase label Jul 15, 2026

khazic mentioned this pull request Jul 16, 2026

feat: add Domino drafter training backend verl-project/verl-SpeCo#17

Merged

orestis-z mentioned this pull request Jul 17, 2026

Feature/domino head for dflash #47627

Closed

8 tasks

mergify Bot added the speculative-decoding label Jul 20, 2026

orestis-z mentioned this pull request Jul 20, 2026

RFC: Ablation Study — Domino vs DSpark Performance Comparison vllm-project/speculators#757

Open

swylyu001 and others added 6 commits July 20, 2026 22:12

[SpecDecode] Support Domino projector in DFlash

d3d36fe

Signed-off-by: swylyu001 <swylyu001@gmail.com> Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

[Doc] Add DFlash/Domino speculative decoding docs

331c781

Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

[Doc] DFlash/Domino gains not yet validated in vLLM

ea14345

Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

chore: fixing inference code, pending update to main inference blocks

61b8af4

Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

chore: fixing inference code, pending update to main inference blocks

a7714aa

Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

arnabwithab and others added 6 commits July 20, 2026 22:12

chore: cleaning up

db5f6f5

Signed-off-by: Eros483 <arnabmandal2912@gmail.com>

fix: pass causal to capture() instead of __init__() for CudaGraphMana…

f006976

…ger compat Signed-off-by: Orestis Zambounis <orestis.zambounis@gmail.com>

orestis-z force-pushed the domino-speculators-patches branch from 5489323 to 490855b Compare July 20, 2026 22:16

mergify Bot removed the needs-rebase label Jul 20, 2026

Merge branch 'main' into domino-speculators-patches

57c8abf

Merge branch 'main' into domino-speculators-patches

4417156

mergify Bot added needs-rebase mrv2 Model Runner V2 specific labels Jul 30, 2026

mergify Bot added the dflash label Aug 20, 2026

github-project-automation Bot added this to Sprint - DFlash Aug 20, 2026

github-project-automation Bot moved this to Backlog in Sprint - DFlash Aug 20, 2026

Navigation Menu

Uh oh!

Uh oh!

fix: Domino inference patches for speculators-format checkpoints - #48241

fix: Domino inference patches for speculators-format checkpoints#48241
orestis-z wants to merge 14 commits into
vllm-project:mainfrom
orestis-z:domino-speculators-patches

orestis-z commented Jul 10, 2026

Uh oh!

claude Bot left a comment

Uh oh!

mergify Bot commented Jul 10, 2026

Uh oh!

mergify Bot commented Jul 14, 2026

Uh oh!

mergify Bot commented Jul 15, 2026

Uh oh!

benchislett commented Jul 21, 2026

Uh oh!

jianuo-huang commented Jul 21, 2026

Uh oh!

mergify Bot commented Jul 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Uh oh!

Conversation

orestis-z commented Jul 10, 2026

Purpose

Fix 1 — speculator.py: step gap between prefix and correction start (bug, affects all Domino checkpoints)

Fix 2 — algos.py: forward Domino fields for speculators-format checkpoints (feature)

Validation

Checklist

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Jul 10, 2026

Uh oh!

mergify Bot commented Jul 14, 2026

Uh oh!

mergify Bot commented Jul 15, 2026

Uh oh!

benchislett commented Jul 21, 2026

Uh oh!

jianuo-huang commented Jul 21, 2026

Uh oh!

mergify Bot commented Jul 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Fix 1 — `speculator.py`: step gap between prefix and correction start (bug, affects all Domino checkpoints)

Fix 2 — `algos.py`: forward Domino fields for speculators-format checkpoints (feature)