[Spec] Add Domino projector support to DFlash speculative decoding (V2)

[heiheiha798]

Motivation

Add Domino projector support to DFlash speculative decoding so the public Domino draft checkpoints (e.g. Qwen3-8B-Domino-b16) can run on SGLang. Closes the request in #28977.

Domino augments the ordinary DFlash draft model with a GRU prefix encoder (prefix_gru) and an MLP (embed_proj) that emits a per-step bias on top of the target lm_head logits. Draft tokens for the block are then produced by a short sequential rollout (each selected token feeds the next GRU step). The verify / accept / bonus-token / draft-KV paths are the existing DFlash V2 paths, unchanged.

Scope of this PR: CUDA + TP=1 only

This first PR intentionally supports single-GPU (TP=1) on CUDA only. The Domino rollout is sequential — the token selected at step k is the input to the GRU at step k+1 — so TP>1 would require a cross-rank synchronization after every drafted token. That is non-trivial and is deliberately deferred to a follow-up commit/PR.

Behavior for unsupported configurations (clear errors, no silent fallback):

• TP>1: rejected early at DFlashWorkerV2.__init__ as soon as a Domino projector is detected (--tp-size 1 required).
• Non-CUDA device: the rollout raises NotImplementedError before touching CUDA graphs / Triton kernels.
• Added-vocab lm_head shards: explicit NotImplementedError.

Other DFlash limitations are unchanged and still apply (DP attention, PP>1, grammar-constrained decoding, return_logprob, return_hidden_states with overlap).

Implementation

• dflash_utils.py: recognize projector_type in {"domino", "causal_v5"}; add Domino fields to DFlashDraftConfig (gru_hidden_dim, emb_dim, pure_draft_prefix_len, shift_label) with validation. DFlashDraftConfig is migrated from @dataclass to msgspec.Struct to follow the repo container convention.
• models/dflash.py: instantiate prefix_gru + embed_proj only for Domino projectors; flatten_parameters() after weight load.
• domino_helper.py / domino_kernels.py / domino_rollout.py (new): Domino rollout state, Triton kernels (fused SiLU(z+s) @ fc2.T + bias + base_logits full-vocab argmax, fused GRU cell), and the sequential draft-block rollout with a CUDA-graph-captured loop.
• dflash_worker_v2.py: create the Domino helper/rollout only for Domino draft models, reject TP>1 early, and branch to the Domino rollout at the existing draft-token selection hook. Ordinary DFlash draft selection is untouched.

Scoring is full-vocab: each step computes the exact argmax of base_logits + domino_bias over the entire vocabulary (no candidate-pool approximation), so the result matches a dense reference exactly. No new environment variables are introduced.

Tests

• test/registered/unit/spec/test_dflash_domino_config.py (CPU, base-a-test-cpu): config parsing + validation (Domino fields, unsupported projector, missing required fields, pure_draft_prefix_len != 1, non-bool shift_label).
• test/registered/unit/spec/test_dflash_domino_kernel.py (CUDA, base-b 1-gpu-small): the fused full-vocab scoring kernel matches a dense torch reference across bf16/fp16/fp32 and multiple vocab sizes, plus a mutation check.

Local validation (RTX A6000, single GPU)

Setup	accept_length	tok/s
Ordinary DFlash (Qwen3-8B-DFlash-b16), baseline	3.066	98.1
Ordinary DFlash, after this change	3.066	98.1
Domino (Qwen3-8B-Domino-b16), V2 overlap	3.737	111.7
Domino, non-overlap (--disable-overlap-schedule)	3.66	106.9

Target model: Qwen3-8B. Prompt: one GSM8K sample, temperature=0, max_new_tokens=512. Ordinary DFlash numbers are unchanged by this PR; Domino output text is identical between overlap and non-overlap.

Follow-ups

• TP>1 support (per-step cross-rank sync for the GRU rollout).
• Optional perf work (e.g. candidate-pool scoring) only if it can be shown not to change the selected tokens.

Closes #28977

---

CI States

Latest PR Test (Base): ❌ Run #28005093292
Latest PR Test (Extra): ❌ Run #28005093197

[gemini-code-assist]

Warning

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

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.