feat(dflash): scheduler/server verify-step glue — raw-argmax bypass, T>0 stochastic accept, EAGLE ctx-fix default-on#297
Open
Sujimoshi wants to merge 13 commits into
Conversation
…4 lm_head Split out of Avarok-Cybersecurity#218 (DFlash speculative decoding for Qwen3.6-27B-NVFP4). Adds the model-level plumbing DFlash's draft/verify path needs: trait default no-op hooks (prefill_phase1_verify, prefill_gdn_wy16) for future verify-path PRs, a row-major K-row dflash_hidden_save capture buffer (dflash_hidden_save_rows, try_dflash_capture_draft/_all) so every accepted verify position gets its own target hidden instead of only rows 0/1, ctx_slot_positions on DflashProposerState to track true sequence positions across thinking-token gaps, an FP32 GDN output sibling buffer (gdn_buf_out_f32 / ATLAS_DFLASH_VERIFY_GDN_F32 prototype) and ssm_verify_h_tmp scratch, a dense_gemm_pipelined kernel for the draft head lm_head, and a RoPE fix for the drafter (config.json has rope_scaling=null; the prior YaRN-interpolated inv_freq corrupted positions and caused 0% accept rate — replaced with the standard 1/theta^(2j/dim) formula). Conflicts against main (which has evolved independently since Avarok-Cybersecurity#218 branched, most notably via Avarok-Cybersecurity#229's NVFP4 MMQ FFN + FlashInfer-GDN work) were resolved by keeping main's own additions (default_loops module, VARLEN-FLA batched methods, NVFP4 lm_head) intact and layering the DFlash-specific fields/methods alongside them, rather than reapplying the original patch wholesale. impl_a1.rs's misc-scratch block was extracted into impl_a1_init::build_misc_scratch (shared with the new gdn_buf_out_f32/ssm_verify_h_tmp allocations) to keep the constructor readable. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…ot drafter's Found while GPU-testing this PR end-to-end on GB10: the DFlash lm_head NVFP4 fast path (w4a16_gemm) used the drafter's own (uncapped) vocab_size as the GEMM's N, but lm_head_nvfp4 is quantized against the target model's tokenizer-capped vocab_size, which can be smaller. On a checkpoint where the tokenizer caps vocab from 248320 to 248077, this read ~243 rows past the end of the packed weight/scale buffers, producing CUDA_ERROR_ILLEGAL_ADDRESS on the first DFlash proposal and poisoning the CUDA context for the rest of the process. Threads the target's capped config.vocab_size through build_model -> BlockDiffusionDraftHead::from_weights as a new lm_head_nvfp4_vocab_size field, and fixes both the GEMM call and the subsequent argmax loop (which shares the same row stride and had the same latent bug, though it wouldn't itself crash — just read the wrong row data). Re-verified end-to-end on GB10 with the same drafter checkpoint and target checkpoint that reproduced the original crash: 3 consecutive requests including the one that previously triggered the CUDA-700 poisoning now all complete cleanly. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
- Soft-load dense_gemm_bf16_pipelined via try_kernel (unused in this PR, shouldn't hard-fail model load) - Make drafter RoPE config-driven (standard vs YaRN) based on config.json's rope_scaling field instead of hardcoding standard-only - Guard ctx_slot_positions push against max_ctx_len to stay in lockstep with ctx_len clamping; downgrade expected-case "ctx_len update skipped" log from warn to trace - Remove LM_HEAD_GEMM_LOGGED debug-session leftover (static + one-shot GEMM pointer dump) - Gate gdn_buf_out_f32 (FP32 GDN verify buffer) allocation behind ATLAS_DFLASH_VERIFY_GDN_F32 env var instead of always allocating Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
The two new DFlash verify-scaffolding hooks pushed transformer_layer.rs to 535 LoC, tripping the file-size-cap CI gate. Condensed verbose doc comments (no logic changes) back to 492 LoC. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…1 conv, opt-in F32 GDN output Extracted from Avarok-Cybersecurity#218 (DFlash speculative decoding), stacked on Avarok-Cybersecurity#258 (model glue). SSM/GDN layer-level pieces of the K=γ verify path: - gated_delta_rule_wy16.cu: fused K=16 WY-chunkwise GDN verify kernel (+ F32-output variant) writing final h_state and Hi_0..Hi_14 intermediates in one pass; wired via prefill_gdn_wy16. - causal_conv1d_update_prefill_inter: batched phase-1 conv that also writes per-token conv_state intermediates for partial-accept rollback; wired via prefill_phase1_verify. - ATLAS_DFLASH_VERIFY_QKVZ_FIX (default ON): route the M=1 verify QKVZ projection through the same w4a16_gemv kernel serial decode uses (~2.29% rel_diff otherwise). - ATLAS_DFLASH_VERIFY_OUTPROJ_FIX (default ON): per-row w4a16_gemv out_proj for verify-window calls (≤17 tokens, native-NVFP4 builds only) to keep batched verify numerically aligned with serial decode (~2.26% rel_diff otherwise). Normal prefill chunks keep the shared GEMM dispatch. - ATLAS_DFLASH_VERIFY_GDN_F32 (opt-in): F32 GDN output in the split4 and wy16 verify branches, removing a BF16-truncation source. - decode_batched sequential K∉{2,3,4,17} fallback extracted to conv_gdn_sequential.rs; prof!-instrumented decode GDN/norm launches. - mod.rs: TransformerLayer impl moved to trait_transformer_impl.rs (500-LoC cap). Deviations from Avarok-Cybersecurity#218 as-authored, due to main's independent evolution: qkvz/out_proj projection ladders were replaced on main by shared dispatch helpers (prefill_qkvz_proj / prefill_out_proj_dispatch) that now carry the CUTLASS-NVFP4 and FP8 branches, so the verify fixes are layered as guarded pre-branches instead of reapplying Avarok-Cybersecurity#218's inline ladders; the out_proj fix is additionally gated to ≤17-token calls because phase3 is now also the hot batched-prefill path. NVTX instrumentation is left to Avarok-Cybersecurity#255. 218's trait_prefill_gdn_batched.rs and qkvz_projection.rs extractions were dropped — main since made its own supersets of both. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…contiguity fix Extracted from Avarok-Cybersecurity#218 (DFlash speculative decoding), stacked on Avarok-Cybersecurity#290 (SSM-layer verify kernels). Model-level orchestration of the K=γ verify path plus a correctness fix for the drafter's context accumulator: - verify_d_kgamma_layers.rs / verify_d_kgamma_attn_meta.rs: per-layer K=γ forward loop and batched attention metadata, extracted from verify_d.rs to keep it under the 500-LoC cap. - dflash_append.rs: EAGLE ctx-slot append dispatchers (dflash_accept_append / dflash_eagle_accept_append / dflash_eagle_kgamma_append) plus dflash_decode_ctx_append_dispatch, called unconditionally at the end of every decode() step so the drafter's ctx accumulator stays contiguous even across stretches where DFlash itself is gated off (e.g. inside a <think> block). Wires the accompanying `skip_next_decode_append` one-shot flag on `DflashProposerState` (new field) so propose()'s own internal decode-append doesn't double-append a row already appended by one of the eagle dispatchers. - verify_a.rs / verify_b.rs / verify_d.rs: wiring for the above plus the K=γ dispatch path (decode_verify_graphed_kgamma_dispatch). Deviations from Avarok-Cybersecurity#218 as-authored, due to main's independent evolution: NVTX instrumentation (verify_step, per-layer/per-phase ranges, d2h_argmax_readback) stripped throughout — left to Avarok-Cybersecurity#255, not yet merged into this stack. prefill_b/batch_kernel.rs and prefill_c.rs had their own `output_f32`/`output_f32_written` GdnPrefillBuffers fields already (from Avarok-Cybersecurity#290's F32-GDN-output work); Avarok-Cybersecurity#218's patch tried to add the same fields again verbatim — kept the existing (real, in prefill_c.rs's case non-null) values and dropped the duplicates. Test plan: cargo check/fmt/clippy/LoC-cap all green on this branch tip. GPU validation on GB10 (Qwen3.6-27B-NVFP4 + DFlash draft, T=0) pending — will report accept-rate/divergence vs the Avarok-Cybersecurity#290 baseline (byte 634) and whether this fixes the previously-observed 0/16 accept collapse on the forced K=γ path (ATLAS_DFLASH_DRAFT_CAP override), since that collapse predates this ctx-accumulator fix. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…FP4-safe lm_head, γ-1 noise0-skip argmax Next cluster from the Avarok-Cybersecurity#218 split, stacked on Avarok-Cybersecurity#292 (verify orchestration), which stacks on Avarok-Cybersecurity#290 (SSM verify kernels), which stacks on Avarok-Cybersecurity#258 (model glue). Contents: - `dflash_head/forward_block_layer.rs` — Q/K/V/O/MLP projections switched to `dense_gemm_bf16_pipelined` (tensor-core pipelined GEMM, ~faster than the scalar dense_gemm for the drafter's per-step M). - `dflash_head/forward_block_step0.rs` (new) — fc-projection step extracted from `forward_block.rs` to stay under the 500-LoC cap; batches the GEMM over all `eff_ctx` ctx slots in one launch instead of a per-slot loop (avoids O(eff_ctx) weight re-reads of the 262 MB fc matrix). - `dflash_head/forward_block.rs` — final logits GEMM now correctly skips noise0 (the conditioning slot, untrained per the block-diffusion paper — its logits are garbage) when extracting argmax drafts: `γ-1` valid drafts from noise1..noise{γ-1}, matching SGLang's reference behavior. Preserved the NVFP4-aware lm_head GEMM dispatch (w4a16_gemm for NVFP4 lm_head targets e.g. Holo) that already existed on this stack — Avarok-Cybersecurity#218's own version of this file predates that support and would have regressed it. - `dflash_head/propose.rs` — timing instrumentation (`t_propose_start`, per-step forward_us/total_us debug log) and optional `ATLAS_DFLASH_T=2` iterative denoising (off by default; tested upstream as acceptance-neutral, not recommended). Deviations from Avarok-Cybersecurity#218 as-authored: - Left `dense_ffn.rs`'s diff out entirely — its `prof!`-instrumentation hunk is superseded by later perf work already on this stack (split_silu / single-warp GEMV dispatch), and its other two hunks are NVTX (out of scope, see Avarok-Cybersecurity#255). Same policy applied to `qwen3_attention/prefill/paged_oproj.rs` and `paged_qkv.rs`: their only changes were NVTX `NvtxRange` tags around GEMM dispatch that's already more advanced on this stack (cutlass_nvfp4 routing) than Avarok-Cybersecurity#218's version — kept the current dispatch unchanged. - `dflash_head/from_weights.rs`: kept this stack's existing YaRN-aware RoPE inv_freq table (handles both `rope_scaling: yarn` and standard/null per config.json) over Avarok-Cybersecurity#218's blanket revert-to-standard-RoPE fix — the YaRN-aware version already supersedes it. - Did NOT adopt Avarok-Cybersecurity#218's flip of `ATLAS_DFLASH_DRAFT_CAP`'s default from 1 to `gamma - 1` (i.e. defaulting to the full K=γ path). Avarok-Cybersecurity#218's commit message cites the BF16 embed-stride fix (2e3a26a, already present in this stack's `verify_d.rs` since Avarok-Cybersecurity#292) as resolving K=γ's garbage-output bug — but our own GPU validation on this exact stack (GB10, Qwen3.6-27B-NVFP4+DFlash, forced `ATLAS_DFLASH_DRAFT_CAP=16`) still measures 0/16 accepted. The stride fix was real but not sufficient; flipping the default would silently degrade every DFlash request to full-γ speculation with 0% acceptance. Kept the safe default (1, forces step_verify_k2) with an updated comment. Test plan: - [x] cargo check -p spark-model / cargo check --workspace - [x] cargo fmt --all -- --check - [x] cargo clippy --workspace --tests - [x] file-size-cap (0 violations) - [ ] GPU validation on GB10 (Qwen3.6-27B-NVFP4 + DFlash draft, T=0) — pending, will report divergence/accept-rate vs Avarok-Cybersecurity#292's baseline (byte 899) before merge.
The comment written before GPU testing claimed K=γ still measures 0/16 accepted (true for Avarok-Cybersecurity#292, before this cluster's noise0-skip argmax fix). Actual validation on this commit shows real, non-zero acceptance (0-93% per step) with no corruption — update the comment to match, and record why the safe default (1) is still kept: verify cost outweighs measured acceptance on this hardware/model, a throughput question, not correctness.
… fix) GPU validation this session: on the identical docstring-template prompt, enabling the flag dropped verify steps 183→46 and raised mean K=γ accept 14.5%→35.5% (max 73%→93%) for the same output length, with no output regression (finish_reason=length, coherent text in both runs). The flag fixes a real ctx-undercount (propose() was appending exactly 1 ctx slot per verify step regardless of num_accepted) rather than being a risky experimental toggle, so it should be on by default. ATLAS_DFLASH_EAGLE_FIX=0 still opts back into the legacy single row-0 capture for A/B testing. This also explains a tester report of frequently seeing >50% (and some >90%) accept at wide gamma: they were very likely running with this ctx-accurate path already active.
…T>0 stochastic accept Threads a `dflash_verify_raw_argmax` flag from CLI args through scheduler::run and mtp_step into all four verify dispatchers (K2/K3/K4/dflash): when DFlash is active, the drafter's own raw argmax output is compared directly against the verify result instead of going through the rep_pen/DRY sampling pipeline first — the pipeline-processed token can legitimately differ from what the drafter proposed even when the target model agrees with the draft, which was silently collapsing accept rate. verify_k2_step.rs additionally gains a T>0 stochastic-acceptance path (rejection sampling against the draft distribution) for DFlash, used when sampling temperature > 0. Dropped from the upstream diff (fork/optimizations @ a7100e0) as debug/perf scaffolding, not feature content: an env-gated per-position acceptance diagnostic table in verify_dflash_step.rs, and an unconditional first-20-steps timing log in verify_k4_step.rs. Stacked on Avarok-Cybersecurity#294 (dflash_head core) -> Avarok-Cybersecurity#292 -> Avarok-Cybersecurity#290 -> Avarok-Cybersecurity#258.
…ers) Matches the model-side default flip in feat/dflash-head-core (c1bb0ad): verify_dflash_step.rs's dflash_eagle_kgamma_append call and verify_k2_step.rs's dflash_eagle_accept_append call now run by default. ATLAS_DFLASH_EAGLE_FIX=0 restores the legacy per-step single row-0 ctx-append path for A/B comparison.
Sujimoshi
requested review from
AzeezIsh,
SeedSource,
rsafier and
tbraun96
as code owners
7 tasks
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
2 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.
Summary
Next cluster from the #218 split. Wires the scheduler/server side of DFlash verify: threads a
dflash_verify_raw_argmaxflag from CLI throughscheduler::run/mtp_stepinto all four verify dispatchers (K2/K3/K4/dflash), adds T>0 stochastic (rejection-sampling) acceptance for DFlash in the K=2 path, and flipsATLAS_DFLASH_EAGLE_FIXto default-on across both this branch's server-side consumers and #294's model-side capture logic.Stacked on #294 (dflash_head core) → #292 → #290 → #258.
Contents
scheduler/mod.rs,mtp_step.rs,main_modules/serve.rs— threaddflash_verify_raw_argmaxthrough to all verify dispatchers.verify_dflash_step.rs/verify_k2_step.rs/verify_k3_step.rs/verify_k4_step.rs— when DFlash is active, compare the drafter's raw argmax directly against the verify result instead of routing through the rep_pen/DRY sampling pipeline first (the pipeline-processed token can legitimately differ from the draft even when the target agrees, silently collapsing accept rate).verify_k2_step.rs— newdflash_stochastic_accept(rejection sampling against the draft distribution) for T>0 DFlash sampling.ATLAS_DFLASH_EAGLE_FIXnow defaults on (both here and via a paired commit on feat(dflash): drafter head core — forward_block split, NVFP4-safe lm_head, γ-1 noise0-skip fix #294's branch): fixes a real ctx-undercount wherepropose()was appending exactly 1 ctx slot per verify step regardless ofnum_accepted, instead of one slot per accepted position.ATLAS_DFLASH_EAGLE_FIX=0still opts back into the legacy single-row-0 capture for A/B comparison.Dropped from the upstream #218 diff as debug/perf scaffolding, not feature content: an env-gated per-position acceptance diagnostic table in
verify_dflash_step.rs, and an unconditional first-20-steps timing log inverify_k4_step.rs.Why the EAGLE-fix default flip
A tester flagged seeing K=γ accept rates commonly >50%, some >90%, and asked to hold off on #294 pending investigation. GPU validation on the identical prompt showed:
finish_reason=lengthfinish_reason=lengthThe tester's report reproduces closely once this flag is on — but it's not a regression, it's the fix working as intended: correct per-position ctx accrual gives the drafter an accurate context, so it proposes better multi-token drafts and needs fewer verify iterations for the same output length. No corruption, no repetition, no degenerate loops observed in any run.
Test plan
cargo check --workspacecargo fmt --all -- --checkcargo clippy --workspace --testsATLAS_DFLASH_DRAFT_CAP=15, no other env overrides — exercising the new defaults as-shipped): 71 K=γ verify steps, mean accept 27.8%, max 100%,finish_reason=length, output coherent (valid Python + docstrings), cut off cleanly at the token limit.🤖 Generated with Claude Code