Repository navigation
dflash: the chain verify builds pipelined and defers the qwen4 PLE gather, as the MTP verify does - #937
dflash: the chain verify builds pipelined and defers the qwen4 PLE gather, as the MTP verify does#937davidtai wants to merge 4 commits into
Conversation
…ection streams; the round policy drives Flash-Next incoai's Qwen3.8-Flash-Next-DFlash2 is trained on the per-token MEAN of the four hyper-connection streams at its target layers (its SGLang capture patch: unflatten(-1, (hc, hidden)).mean(-2)), the tap GLM-5.3's forward already gives its drafter. forwardQwen4With handed the raw [B, L, hc*hidden] stream, 4x the drafter's encoder width. `hcStreamMean` is the shared reshape + mean, unit-tested on a CPU stream. The DFlash round policy (dflash_policy.zig) now drives Flash-Next as it drives GLM-5.3: both are sparse targets where every extra verify row routes to more distinct experts. MLX_SERVE_DFLASH_POLICY=0 keeps the fixed block.
…decar (Flash-Next) `specInitWiring` forced DFlash off for every module-class arch. Flash-Next (qwen4_exp) is one (its n-gram table is module-shared state), so with --drafter the sidecar loaded, every request logged "drafter=enabled", and the slot decoded plain: block 3 on the 1K and 16K cells ran at 65-69 tok/s, serial speed, with no [spec-stats] line. The rule already arms the arch's own MTP head when it can rewind its state across a verify (`moduleStateSpecRollback`); the DFlash block verify rolls back through the same per-position capture (`rollbackSsmFromCapture`), and `DflashModel.bind` refuses the archs with no capture seam. So a bound sidecar now outranks the head there too, as on a shell-owned arch; the gemma drafter and PLD stay off. Tests cover the four cases.
… module-class arch with no head The upstream case "rollback capability alone never arms a head that is not loaded" asserted use_dflash false on a module-class arch with a bound sidecar and no MTP head, which was the old veto, not a property of the head. With a sidecar bound and the arch able to rewind, DFlash runs there whether or not a head is loaded (the sidecar outranks the head either way). A new case covers no head and no sidecar: nothing arms, and no intent bit.
…ther, as the MTP verify does `dflashVerify` built its `[1, 1+m]` verify forward whole and, on qwen4_exp, stopped at trunk layer 1 for the n-gram PLE host gather. `nextMtp` and the tree round already hand the graph to the GPU every VERIFY_PIPELINE_LAYERS layers (`pipeline_build`) and defer the PLE leaf (`ple_defer`) so the host builds while the GPU runs, filling the leaf once after the build and before anything downstream is evaluated (the tree round measured 52.8 -> 50.3 ms a round). The chain verify now does the same. The leaf's ids are host values on this path (the draft ids were read for the verify input), so the flush is a gather, not a sync. Outputs are unchanged by construction: the same ops, submitted earlier.
|
🤖 MLXServe-PRBot — code review
No conflict with main. Generated by Claude Code |
|
All four points taken. Re-measured on this branch's code, exact acceptance, DFlash2 fixed block 8, the standard coding cell (16K prefill, 1,024 tokens), one warm-up request per server, two interleaved A B B A quads (seeds 1234 / 4321), PR head (
The +4.6% in the description came from a binary carrying What is still worth keeping, as a refactor with no performance claim: |
|
Closed: the speed claim did not survive exact-acceptance interleaved reps (see above). The helper refactor and characterization test go to #935 as a no-claim refactor once the real-pack test passes. |
Stacked on #935 (this branch is #935's head plus one commit; the diff shows both until #935 merges).
Why
dflashVerifybuilt its[1, 1+m]verify forward whole and, on qwen4_exp, stopped at trunk layer 1 for the n-gram PLE host gather.nextMtpand the tree round already hand the graph to the GPU everyVERIFY_PIPELINE_LAYERSlayers (pipeline_build) and defer the PLE leaf (ple_defer), so the host builds while the GPU runs; the leaf is filled once after the build and before anything downstream is evaluated. The chain verify now does the same. Outputs are unchanged by construction: the same ops, submitted earlier.Measured
Flash-Next on an M5 Max (128 GB), DFlash2 adaptive block 8 with typical 0.2, card sampling (T 1.0, top-p 0.95, top-k 20, seed 20260829), the same seeded 1K-token task on every prefill rung, 512 output tokens, one warm-up request per server, one binary per set. Same draws on both sets (identical acceptance per rung), so the delta is round time alone:
An earlier A/B on the standard coding cell (16K prefill, 1,024 tokens, seeds 1234/4321, A B B A) read +4.8%. The MTP lane is untouched by this change (its verify already pipelines).
Tests
zig build test-build -Dtest-filter=dflash: 77 passed, 1 skipped.