feat(speculative): add Domino online training path on top of DFlash - #2819
Merged
HuiyingLi merged 3 commits intoJun 28, 2026
Merged
Conversation
Domino (sgl-project/SpecForge#571) extends the parallel DFlash draft backbone with a lightweight causal correction head: a GRU encodes a causal state from each block's previous tokens, and a low-rank projection of [backbone hidden | GRU state] adds a correction to the parallel base logits. Training jointly supervises the Domino-refined and backbone-only logits with a base-anchor curriculum (loss = (1-lambda_base)*final_loss + lambda_base*base_loss, lambda_base decaying to 0). The DFlash path is unchanged: the draft model gains an optional Domino head (projector_type='domino'), a new DominoTrainerModule reuses DFlash anchor sampling and the block mask, and TrainDominoRecipe subclasses the DFlash recipe via four behavior-preserving seams. Adds an example YAML and unit tests (domino core + recipe seams). Signed-off-by: khazic <khazzz1c@gmail.com>
model.to(bfloat16) rounded the rotary inv_freq buffer to bf16, but HF from_pretrained and the SGLang serving runtime keep it in fp32. That train/inference RoPE mismatch grows with absolute position and erodes draft acceptance. Recompute fresh fp32 frequencies from the rotary config after any low-precision cast (a bf16 round-trip cannot be undone by upcasting), covering the training build and reload paths. Signed-off-by: khazic <khazzz1c@gmail.com>
…nt mask clone /simplify cleanups (behavior unchanged): _compute_weighted_losses returns only the three losses instead of also re-returning flattened views (the caller recomputes them as views where needed), and the pre-decay eval mask references weight_mask directly since the decay step rebinds it out-of-place, avoiding a per-step clone. Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Contributor
|
/ok to test 0c29737 |
HuiyingLi
approved these changes
Jun 28, 2026
24 tasks
akoumpa
added a commit
that referenced
this pull request
Aug 30, 2026
* build: bump FA3/FA4 pin to pick up CuTe compile-key fixes The pinned flash-attention ref 002cce0a1 (2026-07-03) predates both CuTe compile-key fixes, so every forward/backward with a tensor max_seqlen -- i.e. every varlen/packed-sequence step -- rebuilt the compile key and recompiled the kernel. Reported against a Qwen3.5-VL packed-sequence finetune on B300. Bump to ce088ab9ce0f (main), which adds over the old pin: - #2507 [CuTe, Bwd] fix backward compile key churn (max_seqlen tensor) - #2762 [CuTe, Fwd] stabilize tensor max_seqlen compile key - #2745 [CuTe] fix forward dynamic-shape correctness - #2819 [CuTe] speed up scalar SM100 mask compilation #2798 raised the cute requirement to nvidia-cutlass-dsl>=4.6.2, so pin the FA4 CUTLASS DSL install to 4.6.2 to match. Kept exact rather than floating so the image stays reproducible. Note this ref also feeds the FA3 (Hopper) wheel, which is built by default, so the bump affects the default x86 image and not just INSTALL_FA4=true builds. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: pin full flash-attention commit Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: bump cutlass-dsl to 4.6.2, ffpa-attn to 0.2.3, quack-kernels to 0.6.4 flash_attn.cute (FA4) at the pinned FLASH_ATTN_REF builds against nvidia-cutlass-dsl 4.6.2. quack-kernels and ffpa-attn both import the same `cutlass` package, so all three have to agree on one version. - quack-kernels 0.6.1 -> 0.6.4 (pins cutlass-dsl 4.6.2; 0.6.1 pinned 4.6.0) - ffpa-attn 0.2.2 -> 0.2.3 (first release pinning cutlass-dsl 4.6.2 and quack-kernels 0.6.4; 0.2.2 hard-pinned 4.6.0/0.6.1) - ffpa extra's explicit cutlass-dsl pin 4.6.0 -> 4.6.2 Because ffpa-attn 0.2.3 already declares the versions FA4 needs, the resolution is conflict-free -- no override-dependencies entries required. Corrects the apache-tvm-ffi note as well: FA4 declares apache-tvm-ffi>=0.1.12 but never imports tvm-ffi, so the <=0.1.11 cap that keeps tilelang working is not mutually exclusive with FA4. Both lock files regenerated with uv 0.8.22 (the version CI pins); the diff is confined to the six bumped packages. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: enable FA4 (INSTALL_FA4=true) by default Flips the default in both stages that declare the ARG: wheel_builder (gates building the flash_attn/cute wheel) and automodel_final (gates the nvidia-cutlass-dsl[cu13] install and the flash_attn/cute symlink). Docker does not inherit ARGs across stages, so flipping only the first would build the wheel while skipping the DSL install and the symlink -- a half-installed FA4. Both FA4 wheels are installed --no-deps, so flash-attn-4's declared apache-tvm-ffi>=0.1.12 bound is never resolved and the <=0.1.11 cap that keeps the tilelang kernels working still holds. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> * build: correct the apache-tvm-ffi comment The note claimed flash-attn-4 never imports tvm-ffi. That is wrong: flash_attn/cute/cache_utils.py has a bare top-level `import tvm_ffi`, reached eagerly through flash_attn.cute.__init__ -> interface.py, so importing flash_attn.cute hard-requires it. The cap itself is still right, but for a different reason. FA4 is installed --no-deps, so its declared >=0.1.12 floor is never resolved, and the tvm_ffi that is present comes via quack-kernels (>=0.1.6,<0.2) pinned here to 0.1.11. FA4 therefore runs against an ffi older than it declares; the two symbols it references, tvm_ffi.Function and tvm_ffi.__version__, both exist in 0.1.11. Note this is unexercised: nothing selects attn_implementation="flash_attention_4" until the FA4 backend lands, so flash_attn.cute is never imported by any test here. Worth knowing given docker/Dockerfile sets FLASH_ATTENTION_CUTE_DSL_CACHE_ENABLED=1, which turns on the AOT cache path in that same cache_utils.py. Verified against flash-attention ce088ab9, the pinned FLASH_ATTN_REF. Comment only; no dependency or lock change. Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com> --------- Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Objective
Add a Domino online training path on top of the existing DFlash speculative-decoding stack, ported from SpecForge (sgl-project/SpecForge#571). The existing DFlash path is left unchanged.
Domino extends the parallel DFlash draft backbone with a lightweight causal correction head. DFlash drafts a whole block in a single non-causal forward, so each predicted position is blind to the (drafted) tokens earlier in its own block. The Domino head fixes that: a GRU encodes a causal state from the block's previous tokens, and a low-rank projection of
[backbone hidden | GRU state]produces a logit delta that is added to the parallel base logits.Design
Training jointly supervises two logits with a base-anchor curriculum:
final_lossis the CE on the Domino-refined logits,base_lossthe CE on the backbone-only base logits.lambda_basedecays linearly to 0 over the firstlambda_base_decay_ratiofraction of training, so early steps keep the parallel backbone strong and later steps let the correction head take over.The integration reuses DFlash end to end:
DominoTrainerModulesubclassesDFlashTrainerModule, reusing anchor sampling, the noise-block construction, and the block attention mask. Only the head, the dual-logit loss, and the metrics are Domino-specific.TrainDominoRecipesubclassesTrainDFlashRecipethrough four behavior-preserving seam methods added to the DFlash recipe (_build_dflash_config,_build_trainer_module,_run_trainer_step,_log_extra_train_metrics). The DFlash recipe behavior is unchanged.prefix_gru,embed_proj) lives on the DFlash draft model, gated bydflash_config.projector_type='domino', so a pure-DFlash draft has no extra parameters.RoPE precision fix
While validating, I found that
model.to(bfloat16)rounds the Qwen3 rotaryinv_freqbuffer to bf16, whereas HFfrom_pretrainedand the SGLang serving runtime keep it in fp32. That train/inference RoPE mismatch grows with absolute position and erodes draft acceptance. The draft now recomputes fresh fp32 frequencies from the rotary config after any low-precision cast (a bf16 round-trip cannot be undone by upcasting), covering the training build and reload paths. This benefits DFlash and Domino alike.Files
nemo_automodel/components/speculative/dflash/domino_core.pyDominoTrainerModule, dual-logit loss, metrics,get_lambda_basenemo_automodel/components/speculative/dflash/draft_qwen3.pynemo_automodel/recipes/llm/train_dflash.pynemo_automodel/recipes/llm/train_domino.pyTrainDominoRecipeexamples/speculative/dflash/qwen3_domino.yamltests/unit_tests/...test_domino_core.py,test_train_domino.py, RoPE-fp32 regression testsValidation
Unit tests (full DFlash + Domino + speculative suite): 393 passed, 1 skipped on an A800 (including new RoPE-fp32 regression tests).
End-to-end smoke (Qwen3-0.6B target, ShareGPT-style data,
sdpabackend): a full run to the final checkpoint. Loss trends down,final_lossseparates belowbase_loss,lambda_basedecays on schedule, and the saved consolidated safetensors carries the Domino head (prefix_gru.*,embed_proj.*) withprojector_type: dominoin the config.Usage
The example expects the same target-model and dataset setup as the DFlash online-training examples, with the Domino architecture and curriculum parameters read from the YAML.
References