feat(speculative): add the Laguna DFlash drafter - #1771
Open
inureyes wants to merge 11 commits into
Open
Conversation
Poolside ships a DFlash speculator for every Laguna release, but the DFlash machinery in mlxcel-core was hard-wired to the Qwen 3.5 drafter shape (split q/k/v, no gate, dense context cache), so the Laguna checkpoint failed on missing keys and the Laguna target had no `SpeculativeTarget`. - `mlxcel_core::drafter::laguna_dflash`: config contract and validation, sanitizer (strip `model.`, fuse split q/k/v, exact key-set check), fused-QKV per-head-gated sliding-window attention over a temporal context window, `aux_hidden_norms` + `fc` + `hidden_norm` context path, and the `Drafter` adapter; `load_drafter` routes `model_type: laguna` to it and the DFlash probe recognizes `DFlashLagunaForCausalLM`. - Laguna target: `SpeculativeTarget` on `LagunaModel` / `LagunaWrapper` (verify forward with capture, rollback trimming dense and buffered rotating caches). - Offline `mlxcel generate` runs the DFlash round loop for Laguna; the server burst target trait is generalized and gains a Laguna arm that seeds the drafter with every prompt row. - Detection test for the drafter directory, docs row, `MLXCEL_PRINT_TOKEN_IDS` for id-level parity checks. The forward follows vLLM's `laguna_dflash.py` (context passes each layer's `input_layernorm` before its K/V projection; per-query sliding window), which differs from the issue text in those two places. Core unit tests and the binary type check pass with `--features cuda`; the greedy-invariant tests and the real-checkpoint runs follow in later commits. Refs #1351
The forced-rejection greedy-invariant test located the bonus token by value in the reference continuation, which is ambiguous on a 32-token vocabulary and made the accept-length assertion flake on a repeated token. The oracle now advances its reference index by the accept length it forces and asserts the bonus it receives sits there, which also turns the accept-length check into a per-round check of the rollback. Refs #1351
…ty test The real-checkpoint probe (`laguna_real_checkpoint_probe`, ignored) takes the classic decode's prompt and token ids and reports two things: (a) the block-versus-chain top-2 logits at every argmax disagreement between a `block_size`-wide verify block and single-token decode on the target alone, and (b) the real drafter's per-position accuracy along the reference path through a shadow drafter that records its proposals while an oracle keeps the round loop on that path. On Laguna XS 2.1 NVFP4 this is what attributed the one greedy-id divergence in 128 tokens to an exact bf16 tie (equal top-2 logits on the chain arm, one ulp apart on the block arm) rather than to the cache or mask logic. The oracle drafter now carries its target layer ids, `run_speculative_with` takes the block size, and `mlxcel generate` prints the prompt ids next to the generated ids under `MLXCEL_PRINT_TOKEN_IDS` so the probe can replay the exact sequence. The core test suite gains `draft_block_depends_on_rope`. Refs #1351
…eport The supported-models row now states what was measured on a GB10 with the NVFP4 target rather than an unqualified identity claim: greedy output equals classic decode except at bf16 logit ties, where the multi-row verify kernels and the single-token decode kernel can round the winner differently (one or two positions per 128 tokens, every one at a chain-arm margin of 0.0 or 0.125), acceptance is 2.6 to 3.9 per round on raw code and 1.1 to 1.8 inside the `<think>` channel, and throughput beats classic only at `--draft-block-size 8` on code (1.14x) while the default block of 16 is slower on this host. The bilingual technical report carries the per-prompt table, the block-size sweep, the drafter's per-position accuracy along the reference path, the server-path counters, the vLLM deviations from the issue text, and the follow-ups (a DFlash exactness gate, an Apple Silicon run of the `metal,accelerate` gate). Refs #1351
inureyes
added
status:review
7 tasks
Main's DFlash server path moved to the `DFlashTargetModel` trait in `src/server/batch/dflash_target.rs` (#1751, #1762) with a per-family exactness gate and drafter-family pairing check, and the drafter refusal message grew a family list (#1770). This merge drops the branch's own `DFlashBurstTarget` trait in favor of main's design and integrates Laguna into it: `LagunaWrapper` implements `DFlashTargetModel` (every prompt row as the first hidden, `block_size` rows of speculative slack on the rotating caches, no batched burst, a `Laguna` drafter family that the target requires and the Laguna drafter self-reports through `Drafter::is_laguna_dflash`), `LagunaVerifyOutput` implements `DFlashVerifyOutput`, and the burst's variant gate, `drive!` match and batched gate gain Laguna arms. `LagunaModel::dflash_exactness_allows` runs the same measured block-versus-chain probe LFM2 and Muse Glimmer run, memoized through `mtp_exactness_gate` with the `MLXCEL_MTP_ALLOW_INEXACT` override; both the server burst and the offline `mlxcel generate --draft-kind dflash` arm fail closed on it. The DFlash probe recognizes `DFlashLagunaForCausalLM` next to the DSpark marker, `load_drafter` checks the Muse assistant, then Laguna, then the default, and the standalone-model refusal names the Laguna drafter. Refs #1351
The DFlash row and the #1351 report now say that the block-versus-chain probe declines the GB10 with the NVFP4 target (107246 of 200704 logit bytes differ at the first verify position), that every measurement was taken with `MLXCEL_MTP_ALLOW_INEXACT=1`, and that the block-8 speed-up did not reproduce under concurrent load (24.54 against 30.00 tok/s), so the pairing is not a reliable speed-up on that host. Refs #1351
…ound copy The drafter forward reshapes on the config's head geometry, so a checkpoint whose `qkv_proj`, `o_proj`, `g_proj` or `fc` rows disagree with the config surfaced as an MLX reshape exception inside the first draft, which crosses the cxx bridge as a process abort. `from_weights` now checks those row counts against the config and returns a load error naming the tensor, with a test for the qkv and fc cases. `combine_hidden` copied the captured hidden input on every round before slicing it, which is a real MLX copy of `[1, T, 5 * hidden]`; it now slices the input directly and casts only when the dtype differs from the drafter's. Refs #1351
…the hot path Implementation review of PR #1771 found one HIGH and several MEDIUM/LOW gaps, all addressed here. - Offline `mlxcel generate --draft-kind dflash` now refuses a non-Laguna DFlash drafter by name before any forward (the server burst already did through `required_family_pairing_error`; a plain Qwen 3.5 DFlash drafter passes `validate_target_compat` and the mismatch was an MLX shape throw crossing the cxx bridge), and refuses a sampling request because the round loop is greedy-only. - `LagunaDFlashDrafter::prefer_requested_block_size` is true so a wider `--draft-block-size` does not engage the adaptive width controller at a width the exactness gate never probed. - The drafter forward drops context rows older than the widest window before `aux_hidden_norms` and `fc` run, advancing every layer's offset, instead of projecting the whole prompt and discarding most of it per layer. - `concat_hidden_for_drafter` uses `concatenate_many`; the rotating caches take the Muse Glimmer slack rule (`speculative_buffer_size`) instead of bare `block_size`, and a cache that cannot take a buffer warns per layer instead of panicking; a short rollback trim logs an error in release builds next to the debug assertion. - `load_drafter` routes on the full Laguna predicate (architecture marker or `model_type` plus `dflash_config`), matching what detection refuses as standalone; `model_variant_label` names Laguna; the config rejects `attention_bias: true` and checks `intermediate_size` against `mlp.gate_proj` at load; the unused `max_position_embeddings` is gone. - `MLXCEL_PRINT_TOKEN_IDS` keeps its one stderr contract, now with a `[prompt ids (N): ...]` line ahead of the token ids, and the duplicate stdout print is removed; `docs/environment-variables.md` updated. - Tests: incremental context appends draft exactly like one combined append (offset continuity), and permuting the fused q/k/v rows changes the logits. Refs #1351
…er, trim the hot path Security and performance review of PR #1771. - `LagunaDFlashDrafter::greedy_only` is true, so the server burst declines a sampling request to classic decode instead of returning the greedy continuation without notice (the round loop's verify is a raw per-position argmax). - `dflash_config.block_size` is bounded to `2..=512` and every sliding window to `2..=2^20` at config load (and the block bound again at draft time): `config.json` is untrusted, the block sizes the verify logits and the target caches' slack, the window enters `i32` mask arithmetic, and an MLX allocation or shape failure crosses the cxx bridge as a process abort. `attention_bias: true` is refused by name. - The sanitizer checks that split q/k/v tensors agree past axis 0 before fusing them and refuses a checkpoint that carries both a prefixed and an unprefixed copy of a tensor; the loader checks projection input widths and the per-head norm lengths on plain tensors, not only the row counts. - The context cache exposes its retained K/V instead of returning a copy each round; the attention no longer copies the context on the non-dropping branch; the context K/V comes from a row-sliced K/V-only projection on plain checkpoints instead of the full `qkv_proj` (the query rows were computed and discarded); `combine_hidden` builds one `concatenate_many` node. - The offline arm refuses a token bias it could honour only for the first token. Refs #1351
…agnostics A lock-held sweep on an idle GB10 (same binary, feature off versus on, block sizes 2 to 16, n=3 per width, n=5 off, 200-token raw code completions) shows no block size beats classic decode: best 0.97x at block 8 with its whole range inside the off arm's, 0.86x at the checkpoint's block 16, 0.51x at block 2. The earlier single-run 1.14x at block 8 did not reproduce and is withdrawn from the docs row and the report; the default block size stays at 16. The reason is now attributed rather than guessed: the offline diagnostic line splits the round loop's verify time into host-side graph construction and synchronized device work, which shows a verify block costing a fixed 77 ms plus 3.3 ms per row of device time against a 30.8 ms graph-replayed classic step, plus about 32 ms of drafter graph construction per round. The PR body, the docs row and the bilingual report carry the table, the spread, and the attribution, with graph capture of the multi-row forwards recorded as the runtime follow-up. Refs #1351
This was referenced Sep 11, 2026
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
1 participant
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
Adds the Poolside Laguna DFlash drafter (
poolside/Laguna-*-DFlash) asmlxcel_core::drafter::laguna_dflash, implementsSpeculativeTargeton the Laguna target, routesmodel_type: lagunadrafters throughload_drafter, integrates Laguna into the server'sDFlashTargetModeldesign (#1751, #1762) as its own drafter family, and wires the pairing into bothmlxcel-serverand offlinemlxcel generate --draft-kind dflash(the offline DFlash arm did not exist for any target before; it is Laguna-only here). Like the LFM2 and Muse Glimmer arms, the pairing runs behind the measured block-versus-chain exactness gate (LagunaModel::dflash_exactness_allows,MLXCEL_MTP_ALLOW_INEXACToverride) on both the server and the offline arm.Where the landed code and the issue text disagree
laguna_dflash.py(Add Laguna XS.2.1 DFlash drafter support vllm-project/vllm#46853), which is what the published checkpoint is served with: every drafter layer passes the projected context through its owninput_layernormbefore the K/V projection (the issue hasctx_qkv = qkv_proj(ctx_i)), and the sliding window is a per-query limit over[context | block]rather than a fixed 511-row context view. With the checkpoint's norm weights all exactly 1.0 the first difference is numerically a no-op today; the second only matters past 512 context tokens.FirstHiddenRows::EveryPromptRow), not the last row the Qwen 3.5 port uses.src/server/speculative_dispatch.rsneeded no Laguna entry: dispatch is by drafter kind; the target arm lives indflash_target.rsandspeculative_burst.rs.enable_speculative_buffer(block_size)so rollback is a tail trim; the drafter's context cache is a temporal buffer capped atwindow - 1, not a ring.Measurements (GB10 sm_121, NVFP4 target, greedy,
MLXCEL_MTP_ALLOW_INEXACT=1)The exactness probe declines this host (107246 of 200704 logit bytes differ at the first verify position), so every run below used the override; without it the offline arm exits with the verdict and the server burst declines to classic decode, and the feature ships inert here by default.
Throughput A/B, same binary, feature off (no
--draft-model) versus on per block size. Raw code prompts, no chat template, 200 tokens, GPU lock held for the whole sweep, host otherwise idle (load 0.10, no other model or cargo process, GPU at 0 percent at start). Decode tok/s as the CLI reports it; n runs, min to max.retry_with_backoff), offlru_get), offNo configuration is a net win on this host. The best width, block 8, is 0.97x on code 0 (its whole range, 31.10 to 31.98, sits inside the off arm's 31.24 to 33.58) and 0.93x on code 1; block 16, the checkpoint's own width, is 0.86x. The earlier single-run 1.14x at block 8 did not reproduce and is withdrawn. The default block size stays at the checkpoint's 16; nothing was changed to favor a width that does not win.
Why the ceiling is where it is. Per round at block 8: about 32 ms of host-side drafter graph construction, 3 ms of target graph construction, and 100 ms of synchronized device work for 4.33 emitted tokens, against 30.8 ms per classic token. The device cost of a verify block is a fixed 77 ms plus 3.3 ms per row (83 ms at 2 rows, 130 ms at 16), 2.7x a single-token step even at 2 rows: the classic step is graph-replayed while the multi-row verify runs eagerly and launch-bound, which is a runtime property of the CUDA backend on this host, not of the drafter. The drafter's own 32 ms per round is one classic token's worth. Even perfect acceptance at block 8 would reach about 17 ms per token here; at the measured 3.3 to 3.6 accepted, 0.86x to 0.97x. Thinking-channel prompts (1.1 to 1.8 accepted) are about 0.4x.
Every greedy-id divergence from classic decode was checked with the block-versus-chain probe in
laguna_real_checkpoint_probe: the chain arm's top-2 margin is 0.0 or 0.125 (one bf16 ulp) at each, so it is theM >= 2versusM = 1quantized-kernel property the exactness gate measures, not cache or mask logic; the tiny f32 model is token-exact with forced accept lengths 0, 1, 2 and full across the window wrap. The drafter's per-position accuracy along the reference path is 0.88, 1.00, 0.75, 0.62 for d_0 to d_3 on code and 0.88, 0.50, 0.25 inside<think>. Server path: same counters as offline (rounds=60 proposed=835 accepted=67on the issue's chat prompt).Acceptance criteria
load_drafterand is refused as a standalone model: verified (detection test plus real runs).docs/supported-models.mdand the report.greedy_invariant_with_forced_rejections(oracle drafter forcing accept lengths 0, 1, 2, full over a window of 6) andgreedy_invariant_with_real_drafter.<think>); 3.55 at block 16 and 3.33 at block 8 on the raw code body (n=3).docs/supported-models.mdand the detection test: done.cargo test --workspace --profile test-fast --features metal,accelerate: not runnable on this Linux/CUDA host; unrun. Ran instead on the merged tree: coredrafter::laguna_dflash,drafter::dflash,drafter::tests(123 passed); binarymodels::laguna_dflash_tests,models::detection_tests,models::laguna_tests,server::batch::dflash_target,server::speculative_dispatch_tests(27 passed, 1 ignored real-checkpoint probe).cargo clippy --profile test-fast --features cuda --lib --bins --tests -- -D warnings(both crates) andcargo fmt --all -- --check: clean.Changes during review
origin/main(nine PRs, including theDFlashTargetModelserver design from feat(speculative): LFM2 / LFM2.5 DSpark drafter on the DFlash loop #1751 and feat(speculative): Muse Glimmer DFlash assistant drafter on the round loop #1762 and the drafter-refusal wording from test: fix settings-schema roster and drafter-refusal drift #1770); the branch's own burst trait was dropped in favor of main's and Laguna became a fourth drafter family with the shared measured exactness gate.from_weightsvalidates theqkv_proj,o_proj,g_projandfcrow counts against the config so a mismatched checkpoint is a load error rather than an MLX reshape throw inside the first draft, and the per-roundcopyof the captured hidden input was removed.--draft-block-sizecannot engage the adaptive width controller at an unprobed width; context rows older than the window are dropped beforeaux_hidden_normsandfcinstead of after;concatenate_manyreplaces the pairwise fold; the rotating caches take the Muse Glimmer slack rule and warn instead of panicking; a short rollback trim logs an error in release builds;load_drafterroutes on the same Laguna predicate detection uses;attention_bias: trueis rejected andintermediate_sizeis checked at load;MLXCEL_PRINT_TOKEN_IDSkeeps its single stderr contract with a prompt-ids line added; two tests pin context-offset continuity and the fused q/k/v row order.greedy_onlyis true so the server burst declines sampling requests to classic decode instead of silently returning greedy output;dflash_config.block_size(2..=512) and the sliding windows (2..=2^20) are bounded at config load because the drafter'sconfig.jsonis untrusted and sizes allocations andi32mask arithmetic; the sanitizer checks split q/k/v shapes before fusing and refuses prefixed-plus-unprefixed duplicates; projection input widths and norm lengths are checked at load; the context K/V comes from a row-sliced K/V-only projection on plain checkpoints, the context cache exposes its buffer instead of returning a copy, andcombine_hiddenbuilds oneconcatenate_manynode; the offline arm refuses a token bias. The regression check after each fix commit reproduced the same 128 ids andaccepted=98on the code prompt.Closes #1351