Repository navigation
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
Contributor
|
Important Draft PR not reviewedDraft PRs are not automatically reviewed by default.
To automatically review draft PRs, update your CodeRabbit configuration: reviews:
auto_review:
drafts: true
Comment |
h-guo18
added this pull request to stack #2758
October 12, 2026 09:10
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #2757 +/- ##
==========================================
- Coverage 79.82% 79.45% -0.37%
==========================================
Files 661 661
Lines 74928 75283 +355
==========================================
+ Hits 59810 59819 +9
- Misses 15118 15464 +346
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
|
h-guo18
added a commit
that referenced
this pull request
Oct 23, 2026
### What does this PR do? Type of change: new feature (performance), bug fix Speeds up DFlash-family drafter training. On Gemma-4-E4B DSpark streaming (PDX B300, 3 nodes = 1 vLLM serve + 2 trainers, 16 DP ranks, batch 4 × 4096), step time goes **1.690 → 0.484 s (3.49×)**, with the loss unchanged within run-to-run noise. Every round below was measured on the same window (steps 21-40, profiler off). | Change | Step time | |---|---| | FlexAttention over a block-sparse draft mask (`dflash_use_flex_attention`, opt-in). The dense `[B, 1, Q, KV]` mask had forced SDPA onto its sm80-only memory-efficient kernel. | 1.690 → 0.711 s | | Offline/streaming and DSpark online: project only the base-teacher rows a loss reads, not full `[B, seq, vocab]` logits; one device sync for DSpark's metrics instead of five | 0.711 → 0.629 s | | DSpark TVD chunked with `Tensor.split` instead of slicing, whose backward zero-filled a full `[N, vocab]` tensor per chunk | 0.629 → 0.541 s | | `torch.compile` of the draft stack and the DSpark TVD chunk (`dflash_use_torch_compile`, opt-in). With compile or FlexAttention on, the anchor count is fixed by the config and `seq_len`, so the compiled shapes stay static: recompiles in 40 steps drop from 4 to 1. Eager runs still size it to the batch. | 0.541 → 0.484 s | The same rounds, by phase and by GPU kernel: <img width="2040" height="2040" alt="image" src="https://github.com/user-attachments/assets/13ae02e5-a171-4f48-8739-c6828d018413" /> Also: - Offline/streaming LiLiCorr (#2589) reads its candidate target logits the same per-row way, instead of rebuilding full-sequence logits. - FlexAttention's conflicts with the attention sink and with dropout are rejected at convert time. - No-behavior-change cleanups: the attention visibility rule is written once (`_draft_mask_mod`) for SDPA and FlexAttention, and KD is switched by `dflash_self_logit_distillation` rather than by a nullable argument. **Why not HF's `flex_attention` dispatch (the path EAGLE takes)?** The BlockMask is built here either way; the difference is the kernel call. HF's `flex_attention_forward` passes `enable_gqa=True` whenever the query-head count is a power of two, and that backward took 568 ms on the Gemma-4-E4B drafter: 10× the path that repeats K/V first, and 2.5× slower than the SDPA baseline. At `head_dim` > 256, FlexAttention's default tiles fail on sm103 (misaligned address, launch failure), so the kernel tiles are pinned to 32×32 and the BlockMask block size has to divide them; `dflash_flex_attention.py` keeps both rules together. ### Usage ```yaml dflash: dflash_use_flex_attention: true dflash_use_torch_compile: true ``` On multi-node runs, point `TRITON_CACHE_DIR` and `TORCHINDUCTOR_CACHE_DIR` at node-local paths: on a shared filesystem, concurrent ranks race on the compile cache. ### Testing - `tests/unit/torch/speculative`: 387 passed on transformers 5.12. `tests/gpu/torch/speculative/plugins`: 14 passed, 1 skipped. - FlexAttention on torch 2.8.0 (manual, RTX PRO 6000): against an fp64 dense reference on padded-row shapes (seq 128, Q_LEN 144 and 244), forward and backward agree to ≤ 7e-7 relative in fp32, and compute-sanitizer memcheck is clean. Before the `q_block` clamp, memcheck caught the backward reading past `anchors`. - New tests cover FlexAttention vs. dense (GPU), the mask rule against a hand-written mask and on padded FlexAttention rows, TVD chunk-size invariance and `split` in the backward graph, anchor sampling vs. the legacy sampler in eager and static modes, the per-row teacher (lm_head never sees the full sequence offline), and FlexAttention conflicts at convert time and through a `from_pretrained` resume. - The cleanups were checked old vs. new: losses, gradients, dense masks and GPU BlockMasks are bitwise equal, except DSpark's rnn head, which differs at bf16 rounding. - Against main, offline DFlash KD, LiLiCorr and DSpark get bitwise-equal losses and gradients in bf16 from the per-row teacher. The DFlash regression recipe (eager) trains in the same time as on main with an identical loss curve. - Step times are medians over 16 ranks × steps 21-40 (profiler off) on the setup above. ### Before your PR is "*Ready for review*" - Is this change backward compatible?: ✅ Both speedups are opt-in: `dflash_use_flex_attention` and `dflash_use_torch_compile` default off. The latter defaulted to True on main but did nothing; the shipped recipes set it false. - If you copied code from any other sources or added a new PIP dependency, did you follow guidance in `CONTRIBUTING.md`: N/A - Did you write any new necessary tests?: ✅ - Did you update [Changelog](https://github.com/NVIDIA/Model-Optimizer/blob/main/CHANGELOG.rst)?: ❌ - Did you get Claude approval on this PR?: ❌ ### Additional Information Based on main. The step times were measured on Gemma-4-E4B, which needs #2186 to run; nothing in this PR is Gemma-specific. The streaming `/done` fix that was part of this PR is now #2756. Follow-up refactors are in draft #2757, stacked on this PR. 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Added an optional FlexAttention path for DFlash draft attention, with block-sparse masking. * Added optional training-time compilation for draft model components; compilation is disabled by default. * Teacher logits can now be obtained from supplied logits or captured hidden states during distillation. * **Bug Fixes** * Improved draft training behavior for batches with no valid anchors, ensuring trainable parameters receive gradients. * Added validation for FlexAttention configurations that do not support attention sinks or dropout. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Hao Guo <haoguo@nvidia.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
h-guo18
force-pushed
the
haoguo/dflash-family-cleanup
branch
from
October 23, 2026 07:36
28eb032 to
9377aa7
Compare
…norm core The next-token shift -- the base logits that predict the token at position i sit at i - 1 -- was spelled out at each of _teacher_logits' three callers: DFlash KD and DSpark as (safe_label_indices - 1).clamp(min=0), LiLiCorr as anchor + s for the token at anchor + 1 + s. A wrong shift trains without an error, just worse, so it now lives once: _teacher_logits takes the label positions and applies the shift itself, and the callers pass what they supervise. DFlash2's predecessor gather uses the same offset for a different quantity (the previous token, not the teacher) and keeps its own. _maybe_apply_base_final_norm took the producer's whole outputs dict, so _teacher_logits wrapped its one flag in a single-key dict to call it. The bool now has its own function, _apply_base_final_norm, which the dict reader delegates to; EAGLE keeps calling the dict reader, unchanged. Behavior is unchanged: against this branch's parent, offline DFlash KD, LiLiCorr (penalty and calibration) and DSpark give bitwise-equal losses and gradients in fp32 and bf16, with pre- and post-norm captures stored in float64. The teacher tests now pass labels and expect row label - 1, and fail without the shift; a new test pins that label 0 reads row 0 rather than wrapping to the last row. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
…he DFlash family Two pieces of the family's forwards were written out once per model. - The step with nothing to supervise returns a zero loss that every trainable draft parameter feeds, so DDP's reduction still sees each of them. There were four copies in two forms: DFlash's indexed one element of each trainable parameter, while Domino's forward, DSpark's forward and DSpark's degenerate-loss branch summed every whole parameter, reading all draft weights to produce a zero. HFDFlashModel._draft_zero_loss now has the cheap form once. - The concatenation of the base layers the draft's fc reads, hidden_states[lid + 1] for each target layer, appeared in the DFlash, DSpark and Domino forwards. It is now HFDFlashModel._target_hidden, which also carries the "+ 1" (hidden_states[0] is the embedding output) that modify()'s check on negative ids refers to. Losses and gradients are unchanged; DSpark's DDP-coverage tests, which require a gradient on every draft parameter for a batch with no valid anchor and for one with zero total weight, pass on the shared helper. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
_build_generate_swa_mask hand-wrote the draft's visibility rule a second time -- context before the block, the window counted from each query's real position, block-causal under dflash_draft_attention="causal" -- and its docstring asked readers to keep it in step with training. It is the training rule, _draft_mask_mod, for the one block being drafted, anchored at ctx_len, so it now evaluates that. The grid evaluation and the additive conversion that the training mask's dense path already did become _dense_additive_mask, shared by both. A verbatim copy of the old builder stays in the tests as the reference: over contexts of 1, 5 and 16 tokens, no window, a window of one block and of two, and bidirectional and causal blocks, the masks are equal element for element (and None where the old one returned None). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
Which token each block position is trained on, which positions count, and how they are weighted were written out five times: DFlash _compute_loss, DFlash2 _compute_loss (for the selector), LiLiCorr _block_targets, DSpark _compute_dspark_loss and Domino _compute_domino_loss. The copies differed only in the label offset (DFlash, DFlash2 and LiLiCorr train position k on the token at anchor + k; DSpark and Domino on anchor + k + 1) and in whether slot 0 is skipped, so a fix to one never reached the others. - _block_targets returns the label token, its sequence position and the supervision mask, given the class constant _label_shift (0, or 1 for DSpark and Domino). The teacher reads its rows from those positions, including LiLiCorr's candidate logits. - _block_tokens is the teacher-forced token gather that DSpark's and Domino's heads and DFlash2's predecessor ids each wrote by hand. - _block_position_weights holds the D-PACE and decay weighting. DSpark and Domino pass "decay" explicitly, which is what they do today: they never read dflash_loss_objective. - _masked_top1 replaces five copies of the masked top-1 accuracy. The in-loss zero branches keep routing through the draft's output (flat_logits.sum() * 0.0) rather than _draft_zero_loss: with DFlashModule as its own FSDP2 unit (the MiniMax-M2.7 launcher layout), reading the resharded parameters directly skips the unit's backward collectives on that rank and the job aborts with a collective mismatch. Behavior is unchanged. Against the previous head, a train step of DFlash (KD and CE), DFlash2 (ce, tv and lambda, with the selector), LiLiCorr, DSpark (vanilla and rnn heads, confidence head on) and Domino, offline and online, with D-PACE and decay, and steps that take the zero-loss branch, gives bitwise-equal loss, gradients and metrics in fp32 and bf16. DFlash2's predecessor ids now come from the shared gather; they differ from the old formula only at positions whose label lies past the sequence end, which carry no selector weight. A new test pins each variant's labels and supervision on a hand-worked example; it fails if DSpark's shift is set to 0 or if the anchor slot stops being excluded. No existing test did. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
…Spark DFlash's and DSpark's pseudo_speculative_generate repeated the same first half: run the base model on the prefix, take its greedy next token, gather the target layers onto the draft's device, build the block (anchor then mask tokens), its position ids and mask, and run the draft. Only the decoding after that differs (DFlash takes the argmax of positions 1..bs-1; DSpark walks its Markov head over the block). - _generation_base returns the base model's next token and hidden states, and carries the offline check both copies made. - _draft_block drafts the block that token anchors. Its position ids now come from _build_position_ids at anchor = ctx_len -- the training layout -- the way the generation mask already comes from the training mask rule. The comment it replaces described a uniform-block layout training has not used since the anchors were introduced. - _target_hidden takes an optional device, so the last two hand-written hidden_states[lid + 1] picks go through it. The docstring's note that generation runs with attention_mask=None was stale since the generation mask was added and now says what is passed. Generated tokens are unchanged: DFlash (full attention, sliding window, causal blocks), DFlash2, LiLiCorr, DSpark and Domino online give bitwise-equal draft tokens in fp32 and bf16 against the previous commit. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
… a second default Thirteen variant fields were read as getattr(config, "dflash_...", <default>) although DFlashConfig declares every one of them with a default, so each default lived twice. One pair had already drifted: dflash_fp32_master_weights defaults to True in the config and to False in HFDFlashModel.modify's fallback, which a reader would take for the default. The fallbacks never run, because modify() always receives a DFlashConfig, so the fields are now read directly. convert_to_dflash_model chose the registry for dflash_architecture_config.projector_type with an if/elif chain and listed the accepted names again by hand in its error. A table maps each projector type to its registry, and the error lists the table's keys. No behavior change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
…e skill tables The shared generation step raises one message for the whole family, "An offline DFlash-family model cannot run AR validation / pseudo_speculative_generate". The known-failure rows in the DFlash and DSpark skill references still quoted the two per-variant messages it replaced, so a search for the actual error found no row. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
h-guo18
force-pushed
the
haoguo/dflash-family-cleanup
branch
from
October 24, 2026 13:47
9377aa7 to
025e18a
Compare
h-guo18
added a commit
that referenced
this pull request
Oct 24, 2026
…evice DSpark's pseudo_speculative_generate adds its Markov bias, computed on the draft's device, to the backbone logits that the base model's lm_head returns. scripts/ar_validate.py loads with device_map="auto", and accelerate keeps a tied lm_head with embed_tokens on the first GPU while the draft can land on a later one. The logits then stay on the first GPU, so the sum raises "Expected all tensors to be on the same device" for every sample and AR validation reports nothing. This has been the case since DSpark was added. The backbone logits now move to the draft's device before the loop. On a single device the move is a no-op. Verified: a new GPU test dispatches the model with accelerate's dispatch_model, embed_tokens and lm_head on the CPU and everything else on cuda:0, and generates with DFlash and DSpark. DSpark fails without the fix and passes with it. The bitwise A/B against #2757 shows DSpark's drafted tokens unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
h-guo18
added a commit
that referenced
this pull request
Oct 24, 2026
…nt and head HFDominoModel did not override pseudo_speculative_generate, so AR validation ran DFlash's, which decodes block positions 1..bs-1 with the backbone alone. DFlash trains position k on the token at anchor + k; Domino trains it on anchor + k + 1 (_label_shift = 1). Every drafted token was therefore compared with the target one position too early, so drafts were almost never accepted, and the correction head that training optimizes and the exporter ships was never applied. domino.yaml keeps AR validation off, so shipped training did not hit this. DFlash's default decode now starts at block position 1 - _label_shift and returns at most bs - 1 + _label_shift tokens. For DFlash, DFlash2 and LiLiCorr (shift 0) that is the same slice as before. Domino overrides it the way DSpark does for its Markov head: it decodes the block left to right, stepping the GRU over the anchor and the tokens drafted so far, and from pure_draft_prefix_len on adds embed_proj([hidden_k; GRU state]) to the backbone logits, as _apply_domino_head does in training. As in DSpark's decode, the backbone logits and the drafted tokens' embeddings are moved to the draft's device: under device_map="auto" the embedding table, and a tied lm_head with it, can sit on another GPU. The eval forward warned that Domino's acceptance rates were backbone-only, but that method computes no acceptance rate, and AR validation now runs the head. The warning goes, and so do the other copies of the claim: the comments in domino.yaml, dspark.yaml and the Nemotron-3.5 DSpark warm-start recipe, the note in the Domino launcher example, and the speculative-decoding skill's Domino, DSpark and validation sheets. DSpark's generation has applied its Markov head since DSpark was added. The validation sheet keeps its rule against reporting such an AR as the model's result, now for the heads that AR validation still skips (LiLiCorr's reranker, DFlash2's selector), and adds the LiLiCorr and online DFlash2 examples to the ones it lists as shipping no benchmark task. The Domino sheet and example now say how to add the HF AR eval: DFlash's task_2 with --steps equal to the block size, since Domino drafts the whole block. The flag values do not change. Verified: a new test checks that Domino's greedy drafts equal the argmax of _apply_domino_head teacher-forced on them, with head weights large enough to overturn the backbone. It fails on the parent commit, and also when the head is skipped, its GRU state lags a step, or the prefix boundary is off by one. The GPU test that puts lm_head on another device now covers Domino as well, and fails without the move. Bitwise A/B against #2757: DFlash and DSpark loss, gradients and drafted tokens are unchanged; only Domino's drafted tokens change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: h-guo18 <67671475+h-guo18@users.noreply.github.com>
This was referenced Oct 24, 2026
This branch has not been deployed
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
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.
What does this PR do?
Type of change: Refactor, no behavior change
Rules the DFlash family wrote out more than once now live in one place each:
_block_targets,_block_tokensand_block_position_weights, driven by one class constant,_label_shift: DFlash, DFlash2 and LiLiCorr train positionkon the token atanchor + k, and DSpark and Domino onanchor + k + 1. DSpark and Domino still always use decay weighting, as today: they never readdflash_loss_objective. A follow-up makes them honor it.isit ati - 1. DFlash KD, DSpark and LiLiCorr each applied that shift before calling_teacher_logits. It now takes label positions and shifts them itself, and every caller passes the positions it supervises. A wrong shift trains silently, just worse._apply_base_final_norm(hidden, prenorm, norm)takes the flag as a bool._maybe_apply_base_final_normkeeps reading it from the producer's dict for EAGLE, and delegates._draft_zero_loss, in the cheap form (one element per parameter, not a fullsum()of every weight). DFlash's and Domino's zero branches inside the loss keep routing through the draft's output (flat_logits.sum() * 0.0): withDFlashModuleas its own FSDP2 unit (the MiniMax-M2.7 launcher layout), reading the resharded parameters directly there skips that unit's backward collectives on the rank and aborts the job. DSpark's degenerate-loss branch uses_draft_zero_loss, because its confidence head is not behind the output._masked_top1.hidden_states[lid + 1]for each target layer, in three forwards and both generation paths, becomes_target_hidden.pseudo_speculative_generateshared their first half; it is now_generation_base+_draft_block. The block's position ids come from the training_build_position_idsatanchor = ctx_len._build_generate_swa_maskhand-wrote the training visibility rule a second time. It now evaluates_draft_mask_modfor its one block, anchored atctx_len, through the dense helper the training mask also uses.getattr(config, name, <default>)reads, covering fifteen variant fields, duplicatedDFlashConfig's own defaults. One pair had drifted:dflash_fp32_master_weightsdefaults toTruein the config andFalsein the fallback. They are now read directly.projector_typedispatch goes through a table instead of an if/elif chain.Net change to source against main: +357 / −426 lines.
Usage
No API change.
Testing
tests/unit/torch/speculative(+tests/unit/torch/export/test_hf_spec_rope_export.py): 505 passed on transformers 5.5.4, 5.12.1 and 5.18.0 (5.5.4 and 5.18.0 fall within CI'stf_minandtf_latestpins).tests/gpu/torch/speculative/plugins: 17 passed, 1 skipped (the sharded-generation test needs 2 GPUs).fully_shardunit, one rank on the zero branch): both ranks finish, as before.Before your PR is "Ready for review"
CONTRIBUTING.md: N/AAdditional Information
Follow-up to the design review of #2279.
🤖 Generated with Claude Code