Repository navigation
Conversation
0ec193f to
21ae60f
Compare
|
🤖 MLXServe-PRBot — code review
The Generated by Claude Code |
|
🤖 MLXServe-PRBot — code review Delta since None of the earlier review is addressed:
No conflict with main. Generated by Claude Code |
…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.
bb87e0a to
22f3e8e
Compare
|
🤖 MLXServe-PRBot approved — reviewed at head All four earlier findings are fixed:
CI on Generated by Claude Code |
|
Addressed in the force-push (
Laptop: |
… leaf deferred; a characterization test pins the chain verify `Generator.verifyForwardDeferred` is the one place that sets `pipeline_build` and `ple_defer`, runs the forward (plain or capture-all), clears the flags, and fills the deferred n-gram leaf before anything downstream or any `rollbackSsmFromCapture` can see it. The MTP verify and the tree round built this sequence by hand (two copies of one ordering contract); the chain verify built eagerly and stopped at the layer-1 host gather. All three call the helper now. `MLX_SERVE_VERIFY_PIPELINE=0` (and `verify_pipeline_override`) builds eagerly; the MTP batched arm keeps its eager build through the same switch. This is a refactor, not a speedup: exact-acceptance A B B A x2 on the fixed-8 lane measured +1.7% greedy and -5.4% on sampled rows against the eager chain verify, so no claim is made for the pipelined chain verify. The ops are the same either way; only their submission differs. Real-pack test (QWEN4_TEST_MODEL + QWEN4_DFLASH, GPU): greedy DFlash decode at fixed block 4 on the repeating prompt, pipelined against eager, emits byte-identical ids with the same round and accepted counts (96 ids, 8 rounds, 15 accepted on the M5 Max). The qwen4 test rig loads through the server's loader so `MLX_SERVE_NOCACHE_LOAD` applies.
|
🤖 MLXServe-PRBot — code review Delta since
The characterization test ( Generated by Claude Code |
Why
incoai published a DFlash2 drafter for Qwen3.8-Flash-Next (
incoai/Qwen3.8-Flash-Next-DFlash2: 5 layers, block 8, selector rank 256 / top-16, two-tap grouped dynamic convolutions, mask token 248077, target layers 5/15/25/35/45). The host already parses, loads and runs DFlash2 drafters (GLM-5.3-Flash-DFlash2, Qwen3.8-27B-DFlash2), and Flash-Next passes every bind check (hidden 2560, vocab 248320, 48 layers, capture-capable). Three things stood in the way:forwardQwen4Withhanded a drafter the raw[B, L, hc*hidden]hyper-connection stream, four times the drafter's encoder width. DFlash2 is trained on the per-token mean of the four streams (its SGLang capture patch:unflatten(-1, (hc, hidden)).mean(-2)), the same tapforwardGlm5Withalready gives its drafter.dflash_policy.zig) was gated to GLM-5.3. Flash-Next is the same kind of sparse target: every extra verify row routes to more distinct experts.specInitWiring) forced DFlash off for every module-class arch, and Flash-Next is one (its n-gram table is module-shared state). With--drafterthe sidecar loaded, every request loggeddrafter=enabled, and the slot decoded plain at serial speed (65-69 tok/s on the 1K and 16K cells) 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), andDflashModel.bindrefuses the archs with no capture seam. A bound sidecar now outranks the head there too, as it does on a shell-owned arch; the gemma drafter and PLD stay off. Four wiring tests added.What
CaptureLayers.hc_mean: a DFlash consumer asks for the hyper-connection mean ([B, L, hidden], viaTransformer.hcStreamMean, unit-tested on a CPU stream); the qwen4 parity oracles keep the raw[B, L, hc*hidden]stream in the slot. Every DFlash capture site in generate.zig sets the flag. A new real-pack guard (QWEN4_TEST_MODEL) checks both shapes from one layer, prefill width and the one-token path, and that the mean slot equals the mean over the raw slot's streams.dflashPolicyFordrives Flash-Next too (dflashPolicyWanted, unit-tested);MLX_SERVE_DFLASH_POLICY=0keeps the fixed block.specInitWiring: a module-class arch with state rollback runs a bound DFlash sidecar (tests cover both sidecar states, the opt-out, and the no-rollback case).Measurements
128 GB M5 Max, the Q4 mixed-bit Flash-Next pack (
Youssofal/Qwen3.8-Flash-Next-MTPLX-Optimized-Speed, converted to the mlx-serve layout), production contract:--ctx-size 262144 --max-concurrent 1 --temp 1.0 --top-p 0.95 --top-k 20, thinking on,MLX_SERVE_NOCACHE_LOAD=1 MLX_SERVE_NGRAM_WARM=0. Prompts are the standard code-ctx seeds (1K = 1,052 tokens, 16K = 16,296),max_tokens 512, one fresh server per configuration, rows in the order 1K card, 1K greedy, 16K card, 16K greedy (so the greedy rows prefill from the card row's prefix cache; decode is unaffected). Card = T 1.0 / top-p 0.95 / top-k 20 / seed 20260829; greedy = T 0. Rates and acceptance are the server's own<- N+512 tokensand[spec-stats]lines. MTP control: the pack's own head,--mtp-max-depth 3 --mtp-typical 0.2. DFlash rows:--drafter <DFlash2> --no-mtp --draft-block-size N, fixed =MLX_SERVE_DFLASH_POLICY=0, adaptive = the round policy (default on Flash-Next with this PR).1. Block-size sweep (one pass)
Reading it:
dflashGateMinimum: the sparse-target floor is 1.80 accepted/round, and two drafts/round accept 1.5-1.67). That is the gate working as designed, not a bug; block 3 falls back to plain decode at 58-70 tok/s.2. Verification on the 1K / 16K track (two passes each, same window)
Pass 2 of every configuration ran slower than pass 1 (the box warms over a 6-minute window), so compare within a pass. The adaptive policy wins 7 of 8 cells against MTP and all 8 against fixed block 8; its mean over both passes is 109 tok/s against 97 for MTP and 97 for fixed block 4. Pick: DFlash2 adaptive, block 8 (the drafter's trained block). Fixed block 4 ties MTP; fixed block 8 loses to it.
Greedy outputs are not byte-identical across configurations: on the 1K cell MTP, fixed 4 and fixed 5 produce one text (a thinking preamble), fixed 6/7 diverge at character 820, and fixed 8 / adaptive open straight into code. The verify width changes the bf16 reduction order, and the opening token of this prompt is a near-tie; the 16K greedy rows agree through the first 86-341 characters the same way. Same class of drift as the GLM-5.3 DFlash rows.
3. Context ladder, 1K..128K, three tracks
One window, one fresh server per track (ctx 262,144 for every rung), the standard code-ctx prompts 1K..128K (the model's window is 262,144 tokens, so 256K and up do not fit with headroom), card then greedy at
max_tokens 512per rung. MTP = the pack's head at typical 0.2; fixed = DFlash2 block 4 withMLX_SERVE_DFLASH_POLICY=0; adaptive = DFlash2 block 8 under the round policy. Logs:reports/qwen-ladder-20261021-1154/<track>/server.log.Decode tok/s, card / greedy, with accepted drafts per round in brackets:
Card row (cold prefix): prefill tok/s and wall seconds for prompt + 512 tokens:
Over the 16 rows per track:
Reading it:
--no-mtpthe prompt pass does not also run the MTP head. At 128K that is 34.1 s against 39.8 s wall for prompt + 512 tokens.Reading: on Flash-Next the DFlash2 drafter under the round policy (
--drafter <DFlash2> --no-mtp --draft-block-size 8, policy on by default with this PR) is a second decode lane beside the MTP head, not a replacement: +4% mean decode over the ladder, +17% at 128K, +20-35% prefill, no gate trips, but the sampled rows at 2K and 32K and the 256K card row still go to MTP (see section 5). The MTP lane stays the default;--drafteropts a server in. The fixed block is not worth shipping on this target without lowering the MoE floor, which this PR does not touch.4. DFlash2 with typical acceptance at 0.2 (
--dflash-typical 0.2, side branchqwen4-dflash2-typical, NOT in this PR)The DFlash verify had no typical mode: on this row-exact trunk a sampled request drew every verify row with the keyed sampler and kept a draft only while it matched, so sections 1-3 compare MTP under
--mtp-typical 0.2against DFlash2 under an exact rule.--dflash-typical <d>gives the DFlash verify the MTP lane's floor (keep a draft while its filtered target probability clearsmin(1, d * exp(-H)), a miss resamples the target row). Same window shape as section 2, two passes, binary with this commit. Greedy rows are unaffected by construction and reproduce section 2's counts exactly.Reading it:
Takeaway: on Flash-Next the exact rule costs DFlash2 nothing measurable, so the recommendation stays the adaptive policy under exact acceptance.
--dflash-typicalis there for a matched comparison against--mtp-typicaland for targets where the verify is cheap enough for acceptance to be the bound.5. Standard coding cell, the YaRN ladder to 256K, and where the round's time goes
Standard cell (the GLM program's coding workload: a 1,024-token task on a CPython module, 16,387 rendered tokens, max_tokens 1024, T 1.0 / top-p 0.95, seeds 1234 and 4321, one fresh server per run, lane order balanced across the seeds; rates from the server's own line, peak RSS sampled at 2 s):
On this cell DFlash2 adaptive ties the MTP head on decode (100.8 vs 102.9) and wins time-to-first-token (6.5 vs 8.1 s: the prompt pass skips the head). Acceptance is lower than on the code-ctx cells because the output is a long reasoning preamble (2.8-3.5 accepted/round under the policy, 40-50% per draft).
Ladder to 256K. The 256K prompt is 262,144 tokens, the whole shipped window, so this ladder runs under Qwen's extension recipe (
--config-overridesYaRN factor 4, 262,144 -> 1,048,576; the server advertises 270,336 as its session window) with the hot prefix cache capped at 1 GB so the 256K rung fits beside the 83 GB model under a 112 GiB guard. Decode tok/s, card / greedy, accepted drafts per round in brackets:Reading it:
--dflash-typical 0.2shows up in tok/s here where it did not in section 4: the card rows accept 4.1-5.0 per round and run 5-15% faster than exact adaptive at 16K-128K.Where a round's time goes (M5 Max, 16K KV,
MLX_SERVE_DECODE_FWD_UBENCH: the trunk's pure forward, no host work):The pack is 77 GB of experts (3.07 MB per expert at 4-bit with group-32 scales and biases) and about 4.8 GB of dense weights per forward. One row routes to 10 experts per layer; 8 rows route to about 74 distinct experts per layer, so the verify reads 2.5x the bytes of a serial step for 8x the positions, at roughly 68% of a 546 GB/s bandwidth floor. A traced DFlash round on the coding cell (
MLX_SERVE_DFLASH_TRACE=1, which adds its own evals) is 42-52 ms of verify plus about 6 ms of drafter forward, draft head, acceptance and context append, for 3.2 to 5.6 accepted tokens; the adaptive policy's confident stretches reach ~145 tok/s and its narrow rounds pull the mean to 95-117. The draft-tree verify never engages on this arch (the GDN layers hold no per-branch state), so tokens per round are bounded by the single 8-token path. The n-gram table page-warm (29.8 GB) does not fit beside the model under the guard on a 128 GB box.Existing switches on the coding cell (one fresh server per run; seed 1234, and 4321 where two rows show). None moves the ceiling: a deeper MTP round accepts more but costs more (round 30 -> 43 ms at depth 4-5), pinning the policy at 8 rows loses a quarter, and the NAX verify lane's width threshold is already right.
Reaching 140 tok/s with exact acceptance therefore needs a cheaper 8-row verify (fewer expert bytes per round: a row-aware MoE verify that computes only the rows routed to each expert, and the group-32 scale/bias overhead) plus overlapping the drafter forward with the verify. That is follow-up kernel work, outside this PR.