Add correctness-first Domino support to DFlash V2 by jianuo-huang · Pull Request #31328 · sgl-project/sglang

jianuo-huang · 2026-07-15T11:31:28Z

Summary

This PR adds Domino projector support to the current DFlash V2 worker for the
public Huang2020/Qwen3-8B-Domino-b16 checkpoint. It reuses SGLang's existing
target
verification, continuous-prefix acceptance, bonus-token handling, and KV-cache
commit paths unchanged.

The implementation is independent of the earlier integration in #28998 and is
based on the public Domino reference and the current DFlash V2 interfaces.
Addresses #28977 and #29511.

Design

Parse, validate, instantiate, and strictly load the checkpoint-native
bias-free prefix GRU and two-layer SiLU correction projector.
Generate 15 proposals for the public block-size-16 checkpoint. The first is
the full-vocabulary DFlash argmax; the GRU is initialized with
[verified_token, first_proposal], and every realized proposal is fed back
before the next correction step.
Keep the base LM-head projection full-vocabulary. Correction scoring uses one
per-request, block-shared candidate pool with K=2048 by default; set
--speculative-domino-candidate-pool-size 0 for full-vocabulary correction.
Use native nn.GRU for the two-token prefix and ATen GRUCell for the 13
one-token feedback updates.
Capture the complete rollout in the existing DFlash draft CUDA Graph, with an
eager fallback when the graph is disabled or unavailable.

The four performance changes are documented separately:

Supported and numerical scope

Tested scope: A100 80GB, CUDA, BF16, TP=1 for Qwen3-8B,
pure_draft_prefix_len=1, and block size 16.
Graph results cover exact capture batch shapes with
--disable-cuda-graph-padding. Padded replay is not claimed.
ATen GRUCell and cuDNN GRU use different BF16 operation orderings, so their
proposal tokens are not generally bit-exact.
K=2048 is an approximate search for the correction argmax. Target
verification is unchanged, but the observed K2048-versus-K0 output parity is
limited to the fixed workloads reported below; it is not a general
equivalence guarantee.

Validation

Registered eager GSM8K-200 gate with public models/data: score 0.950, average
speculative acceptance length 4.5544 (gates: >=0.90 and >4.0).
Exact-shape Perf 1 versus Perf 2 graph A/B: GSM8K-200 score 0.950 in both
arms and 0/208 fixed-512 ShareGPT text mismatches.
Perf 3 eager versus exact-shape graph replay: 0/21,900 proposal mismatches
over 20 random trials per batch. GSM8K-200 remained 0.950.
Perf 4 K2048 versus full-vocabulary K0: GSM8K-32 matched 32/32 output-ID
sequences with acceptance length 4.7975; ShareGPT C32 matched 128/128
texts with acceptance length 3.7980.
Focused config, weight-loading, runtime validation, full-chain, batch
independence, candidate-pool, sampler, and CUDA Graph replay tests:
22 passed, 23 subtests passed.

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python \
python3 -m pytest -q -s \
  test/registered/spec/dflash/test_dflash_domino.py

PYTHONPATH=python python3 -m pytest -q \
  test/registered/unit/spec/test_dflash_domino.py

Performance summary

The first table is a standalone Domino-rollout benchmark on one A100 80GB,
BF16/TP=1, 15 proposals, exact batch shapes, and full-vocabulary correction
(K=0). Target verification and the DFlash backbone are excluded. Values are
the median of six balanced measurements with 30 CUDA-event iterations each.

Batch	No optimization	Perf 1	Perf 2	Perf 3	No-opt to Perf 3
1	13.416 ms	5.888 ms	3.372 ms	2.363 ms (base: 0.674 ms)	5.68x
8	14.104 ms	6.023 ms	3.875 ms	2.881 ms (base: 0.775 ms)	4.90x
64	15.851 ms	8.821 ms	8.571 ms	7.808 ms (base: 4.555 ms)	2.03x

Parenthesized base-logit time is included in the total but belongs to DFlash; it
is not Domino-specific overhead.

Perf 4 uses a separate candidate-pool harness and aggregation, so it is reported
as its own K0 A/B rather than appended to the table above.

Batch	Full correction, K0	Candidate pool, K2048	Reduction
1	2.441 ms (base: 0.724 ms)	1.618 ms (base: 0.724 ms)	33.7%
8	2.879 ms (base: 0.817 ms)	2.122 ms (base: 0.817 ms)	26.3%
64	7.700 ms (base: 5.104 ms)	7.188 ms (base: 5.104 ms)	6.7%

The same parenthesized DFlash base-logit time is included in both Perf 4 arms.

Using SGLang's official serving benchmark runner:

Perf 2 versus Perf 1 improved output throughput by 11.2%, 3.4%, and
1.5% at concurrency 1/8/32.
Perf 3 versus Perf 2 improved output throughput by 5.8%, 3.5%, and
1.9%; C32 acceptance length changed by -0.16%.
Perf 4 K2048 versus K0 measured 1450.31 versus 1431.25 output tok/s at
C32 (+1.3%). This was one run per K and is treated as exploratory/noise-level,
not a serving-speedup claim.

Final end-to-end serving benchmark

A final 135-cell TP1 serving sweep compared target-only decoding, the official
Qwen3-8B DFlash b16 draft, and Domino b16 K2048 on one A100 80GB.

Full settings and all numerical results: benchmark comment.

Rollout microbenchmark speedups are not end-to-end serving claims. Exact setup,
correctness data, caveats, figures, and benchmark commands are in the four
linked performance comments.

Checklist

Format changed files with the repository pre-commit hooks.
Add focused CPU/CUDA unit tests and a registered CUDA E2E gate.
Provide accuracy, acceptance-length, rollout, and official serving results.
Add user documentation after the supported runtime contract is agreed.

CI States

Latest PR Test (Base): ❌ Run #33199855245
Latest PR Test (Extra): ❌ Run #33199855055
Latest PR Test (AMD ROCm 7.2): ❌ Run #33199855228

gemini-code-assist

Code Review

This pull request introduces support for the Domino speculative decoding projector within the DFlash framework. It adds configuration parsing, runtime validation, weight loading, and sequential greedy rollout logic using a prefix GRU and embedding projection, along with corresponding unit and integration tests. The review feedback suggests optimizing performance by precomputing base logits in a single batched operation instead of sequentially inside the rollout loop, and improving robustness when checking for required projector weights by defensively handling potential parameter prefixes.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

…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>

gemini-code-assist · 2026-07-17T03:59:19Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist · 2026-07-21T14:25:01Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist · 2026-07-21T15:10:42Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist · 2026-07-21T15:16:32Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

gemini-code-assist Bot reviewed Jul 15, 2026

View reviewed changes

Comment thread python/sglang/srt/speculative/domino_utils.py

Comment thread python/sglang/srt/models/dflash.py

jianuo-huang force-pushed the feat/domino-correctness-first-upstream branch from 987e10f to a7bb8af Compare July 15, 2026 16:37

khazic mentioned this pull request Jul 16, 2026

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

Merged

amd-bot mentioned this pull request Jul 16, 2026

[CI Monitor] Daily Report - 2026-07-16 bingxche/sglang-ci-bot#137

Open

jianuo-huang marked this pull request as ready for review July 17, 2026 03:59

jianuo-huang requested review from Qiaolin-Yu, Ying1123, hnyls2002 and merrymercy as code owners July 17, 2026 03:59

jianuo-huang force-pushed the feat/domino-correctness-first-upstream branch 4 times, most recently from 8a43e4c to d6222a4 Compare July 21, 2026 14:19

jianuo-huang marked this pull request as draft July 21, 2026 14:23

jianuo-huang marked this pull request as ready for review July 21, 2026 14:24

jianuo-huang marked this pull request as draft July 21, 2026 14:27

jianuo-huang marked this pull request as ready for review July 21, 2026 15:10

jianuo-huang marked this pull request as draft July 21, 2026 15:14

jianuo-huang marked this pull request as ready for review July 21, 2026 15:16

jianuo-huang mentioned this pull request Jul 21, 2026

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

Draft

3 tasks

jianuo-huang force-pushed the feat/domino-correctness-first-upstream branch 2 times, most recently from ec28069 to 19ed8f1 Compare July 22, 2026 05:22

jianuo-huang mentioned this pull request Jul 22, 2026

feat: support tensor-parallel Domino rollout #32018

Closed

Qiaolin-Yu self-assigned this Jul 23, 2026

jianuo-huang mentioned this pull request Jul 26, 2026

Domino now runs in SGLang (DFlash V2 worker) — interested in DominoTree under real serving? slin-zhq/Domino-Tree#1

Open

This was referenced Aug 28, 2026

feat: add optimized Domino rollout to DFlash V2 jianuo-huang/sglang#3

Closed

feat: add optimized Domino rollout to DFlash V2 #36860

Closed

jianuo-huang added 4 commits August 29, 2026 02:32

feat: add correctness-first Domino support

32f3dd4

perf: precompute Domino base logits

f50594a

feat: capture Domino rollout in draft CUDA graph

1489e15

perf: add Domino candidate vocabulary pool

dba1120

jianuo-huang force-pushed the feat/domino-correctness-first-upstream branch from 19ed8f1 to dba1120 Compare August 28, 2026 18:33

jianuo-huang mentioned this pull request Aug 28, 2026

feat: add optimized Domino rollout to DFlash V2 #36899

Merged

Qiaolin-Yu closed this Sep 8, 2026

Navigation Menu

Uh oh!

Add correctness-first Domino support to DFlash V2 - #31328

Add correctness-first Domino support to DFlash V2#31328
jianuo-huang wants to merge 4 commits into
sgl-project:mainfrom
jianuo-huang:feat/domino-correctness-first-upstream

jianuo-huang commented Jul 15, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

jianuo-huang commented Jul 15, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Design

Supported and numerical scope

Validation

Performance summary

Final end-to-end serving benchmark

Checklist

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot commented Jul 17, 2026

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

gemini-code-assist Bot commented Jul 21, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

jianuo-huang commented Jul 15, 2026 •

edited by github-actions Bot

Loading