Repository navigation
lfm2: draft with LiquidAI's DSpark drafter (--drafter) - #1391
Open
BlivionIaG wants to merge 39 commits into
Open
BlivionIaG wants to merge 39 commits into
BlivionIaG wants to merge 39 commits into
Conversation
With Drafted.held_on_target set, Drafted.verify gives the target its windows with Window.held as the round loop made them, so a target that reads held drafts from the drafter on the device skips the host readback. A window of a stream with fillers still gets host tokens. The flag is off by default; without it verify runs as before and the new check is not evaluated. drafted_held_test.zig checks the three cases: held windows pass through, a stream with fillers gets its pending token as host tokens, and with the flag off the drafter's drafts are read back as host tokens.
engine.zig is 592 lines and the shared-round guard needs room in it. pipelinedRound and landInflight move unchanged to pipelined.zig as round and land; Engine.familyRound and Result become public for them. No behaviour change.
…el.shared_guard) A drafter whose hold is a pass of its own costs a shared round more than the depth rule's head steps price: the rule treats the rest of a round's time as overhead that plain rounds pay too. Facts.shared_guard (and Model.shared_guard for a backend) lets the engine compare the two modes directly. shared_guard.zig keeps, for each stream count, committed tokens over the round's wall time in plain stretches (no budgets, so no holds) and in drafted stretches. It starts plain, measures plain rounds for 16 rounds and drafted ones for 4, runs the faster mode for 512 rounds and doubles that, up to 64 times, while the same mode wins again; then it measures the other mode again. Drafted rounds must beat plain ones by 5%. A mode's first round still carries the other mode's drafts or holds and is not measured. An entry is allocated once a stream count; later rounds take no allocation. A plain stretch changes only how many drafts a window carries, never what the target draws, so replies stay the same tokens. Without the flag the engine never consults the guard. Tests: the guard's choice and re-measures, ever rarer re-measures of a losing mode, its skipped first rounds, no allocation after a stream count's first round, and replies equal to one-token decode with the guard on and off.
… stream joining one holds no drafts With the shared-round guard on, a plain round still verified the lane core's context and Recall copies, which are drafts too, priced from a prior; and a stream joining other streams held its first drafts even while the guard ran plain rounds at that stream count, a hold pass for each join. The engine now marks a shared round plain when the guard runs that stream count plain, and copyProposal proposes nothing for it, so the guard's plain mode is the drafts-off round. Open.firstDepth gives no drafts to a stream that joins live streams while the guard runs their count plus one plain. Both are behind Model.shared_guard; without it plain_round stays false and firstDepth is unchanged. Tests: a guarded plain round proposes no copies; a stream joining plain-guarded rounds holds no first drafts and, without the flag, holds them as before.
The host drew every verify row by scanning its 128,000 bf16 logits, after
the forward and with the GPU idle. lfm2_argmax picks each head row's first
largest logit on the GPU and flags a row with a nonfinite logit. verifyFn
takes the pick for greedy rows without a mask; sampled, masked and flagged
rows draw on the host as before.
The kernel is registered with the contract ("metal lfm2 argmax": the host
draw's pick, ties to the lower id, nonfinite rows flagged), and
tf-lfm2-exact checks every window's picks against the host draw.
…ctions lfm2_conv, lfm2_qk, lfm2_attn_mma and lfm2_attn_merge each call an inline body (lfm2_conv_row, lfm2_qk_row, lfm2_attn_rows, lfm2_merge_row) with the same arguments, so another kernel can run a row's arithmetic unchanged. The kernels' bits do not change.
The family's tests (config, draws, state) ran only inside `zig build test`. test-lfm2 runs the engine's test binary filtered to lfm2_moe, so a change to the family can run them without the whole suite.
…ll its streams A recorded call of several streams launched each stream's own short conv (30 layers) and its qk, attention and merge (10 layers): at 16 one-row streams 960 launches a forward for little work each. lfm2_segments.metal runs each of them once a layer over every row, from a table of the streams' cache addresses, row offsets and key counts (segments.zig); each row's arithmetic is the per-stream kernel's, called through the same inline bodies. A stream of more than 8 rows folds a chunk's sub-spans in its 8-row tiles, as its own 16-row tiles do. Prompt calls and one-stream windows keep their launches. Every cache buffer the table reaches is declared to the encoder: a new stream's caches join the server's residency set only at its next flush, so a kernel that reaches them through the table alone may read them before they are resident. The scratch's attention partials now hold each stream's apart (partialCount: a 32-row scratch's grow fourfold; the served scratch of 2,048 rows already holds them). The contract registers "metal lfm2 stream attention" and "metal lfm2 stream conv" as bit-equal replacements of the per-stream kernels, with float64 attention and a host replay of the conv steps. test-lfm2 checks the table, the row map, the tiles and the declared buffers, and the partials' count at 1, 32 and 2,048 rows.
… more From 8 rows the dense twin took its MMA window kernel, a threadgroup for 16 outputs over every row. A 2,048-wide matrix then has 128 threadgroups and runs slower than the scalar twin's 4-row passes; shared rounds of 8 to 32 rows call such matrices at every layer. The 6,144-wide short conv input and the head keep the window. Both kernels give the tiles' bits (the load check runs each matrix's twin at every width); the contract's twin now also serves 16 rows of 6,144 from 2,048.
… receipt README pointed at no native recipe for LFM2.5-24B-A2B. The page gives the checkpoint, the serve command, the shared-round kernels and their tests, and the measured M5 Ultra decode rates before and after them.
LiquidAI/LFM2.5-24B-A2B-DSpark is LiquidAI's speculative-decoding drafter for LFM2.5-24B-A2B. `tensorfold serve` now takes it with --drafter DIR for lfm2_moe checkpoints. The lane core verifies every draft, so a drafted reply equals its "draft": false reply. The drafter reads the target's rows after layers 2, 10, 18, 26 and 34 (the taps). fc and hidden_norm turn each row into the context keys and values of five Qwen3 layers. A block of the pending token and mask tokens runs the layers with non-causal attention, then the target's head. The Markov head adds W2 W1[previous token] to each row's logits and takes the argmax, one row after the other. A block is as long as the deepest chain its pass holds, at most the nine it was trained on. The drafter runs on the family's own kernels and core/tile_qmm (on GPUs with tensor units, core/affine_mm for absorb passes past the twin's rows): its bf16 linears are quantized to 8-bit MLX affine codes at load, and the q and k rows are reordered so that the target's half-split RoPE kernel gives the drafter's interleaved RoPE. New kernels: - lfm2_markov_scan and lfm2_markov_pick (lfm2_dspark.metal): a Markov step's W1 gather, W2 dot over 8-bit codes, bias add and argmax, ties to the lowest token, a NaN score flagged; - lfm2_attn_block and lfm2_attn_block_merge: the target's attention and merge kernels with a FULL template flag, every block row seeing the keys [0, len + 1). The target's own instantiations keep their causal count. Both register with the kernel contract (metal_lfm2_dspark.zig, metal_lfm2_attn.zig), gate 5 against every Apple chip we serve and this GPU. Rows reach the drafter in passes: a prompt chunk in one pass of up to 2,048 rows; a round's verified rows are staged on the host (128 a lane) and absorbed by the next hold, so a round without holds does no drafter GPU work. A hold pass takes as many streams as their blocks fit in 32 rows. The drafter sets Facts.shared_guard, so shared rounds draft only while drafted rounds measure faster than plain ones at that stream count; a lone stream keeps the depth rule's own plain guard. The drafter keeps its context with each kept prompt state, so a drafting stream resumes a prompt as before. Its lane's bytes join the streams the load limit admits. --drafter-bits 0 is refused for this drafter, and a refused open frees the model it loaded. Tests: test-lfm2-dspark checks the config's admission and refusals, the quantizer, the RoPE reordering through lfm2_qk, the Markov step and the block attention on the GPU. tf-lfm2-dspark runs the real drafter against an fp32 host reference of the released checkpoint and times its passes.
… the GPU A hold that fits in one pass signals a GPU event and returns without a wait. The next verify waits on that event on the GPU, and a new kernel, lfm2_draft_ids, copies each window's drafts from the drafter's picks into the verify's ids. The host reads the picks only after the verify, or when a grammar, an absorb, a keep or a release needs them. Multi-pass holds stay synchronous. The LFM2 backend takes held windows (Window.held) from a HeldSource the drafter gives, and Drafted.held_on_target passes windows without fillers to it as they are. lfm2_draft_ids registers with the kernel contract (metal lfm2 draft ids): each stream's first pick, or 0 for a token past the vocabulary, copied exactly; gate 5 against every Apple chip we serve and this GPU.
The lane core can put an external drafter in front of a target. The drafter reads the target's hidden rows after some layers (its taps). LFM2-MoE had no taps, so it could draft only with the core's own heads. prepare_features names the layers. From then on, each forward copies each row's output of those layers into the scratch (lfm2_tap, a plain copy of bf16 values). features() gives a stream's rows of the latest verify or one-token round. prefill_features hands each prompt chunk's rows to the drafter before the chunk's end is kept for reuse. The prompt's last row is kept beside the stream's caches for its first draft. A drafting stream resumes only a kept state that also holds its drafter's context (Snap.draft). Without a drafter nothing changes: no taps are kept and every forward runs as before. lfm2_tap registers with the kernel contract (metal_lfm2_tap.zig): the third of five taps' columns at 1, 10, 32 and 2,048 rows must hold the layer's bits and the other columns what the call found. Gate 5 checks its launch against every Apple chip we serve and this GPU. test-lfm2-dspark checks the same copy through the family's own kernels.
The recipe says where the drafter comes from (LiquidAI/LFM2.5-24B-A2B-DSpark at its revision), that its weights are under the LFM Open License v1.0 and not shipped with TensorFold, how to serve with --drafter, what the drafter runs and which kernels and checks cover it. The README row and the recipe index name the drafts.
This was referenced Oct 26, 2026
Owner
|
A small request, @BlivionIaG: this description links more than 20 issues and PRs, which is over Agent Traffic Control's limit, so ATC drops the PR from its board and collision checks (and marks every check incomplete for everyone). Could you trim the related-work list to the issues this PR actually fixes or depends on (20 at most)? Thanks. |
ashhart
requested changes
Oct 26, 2026
ashhart
left a comment
ashhart
left a comment
Owner
There was a problem hiding this comment.
Thanks @BlivionIaG, a lot of careful work in this stack. We mapped all 31 commits to their PRs and reviewed each on current main. For the DSpark drafter (this PR and its descendants):
- Loading changes the released BF16 weights.
dspark_weights.zigquantizes each dense matrix to 4/8-bit codes at startup (8-bit by default), and #1393 still discards the original. Target verification protects final tokens, but the engine doesn't create lossy copies of a checkpoint at serving time. Please run the released BF16 matrices, or refuse the format. Also honour each advertised--drafter-bitsmode or refuse it by name (the native adapter accepts the setting and then uses 8 regardless). - 4-bit mode reads past its buffer. With 4 bits the loader packs
markov_w2at that width, butlfm2_markov_scan(lfm2_dspark.metal) reads a byte per coefficient and the launcher passes no bit width. At 128000x256 the code buffer is 16,384,000 bytes and the shader can address byte 32,767,999 (found statically, not run on a GPU). Please refuse that mode before load or implement it correctly. - The held (asynchronous) path drops draft failure flags: the synchronous path refuses a NaN pick or invalid ID (
DraftNonfinite), butreadPicksafter an unwaited hold ignoresPick.nonfiniteand maps an invalid ID to 0, as does the GPU-held copy. Native serving enables that path, so please carry the flag through the fence and add planted NaN/Inf/invalid-ID controls for both paths. - Contract proof: the new operators report 0 planted controls and no recorded compiler/launch limits (the limit declarations use current-device checks with
offline = false). Please add the gates, red controls, served shapes and records. - Config admission doesn't check
hidden_actbefore the fixed SiLU path, and an explicit scheduled verify without a confidence head silently becomes depth mode; both should refuse by name. - The FP32 reference hard-codes a 6-bit target embedding and its host
matmuljoins workers only after all spawns succeed. - The device receipts are mostly on older bases; we'll need alternating A/B on current main per landing head (the stack's own notes report a 1.8-2.7% four-stream drop on Qwen3.8-27B + DFlash2 from the shared-lane changes, which needs explaining, not dismissing).
This was referenced Oct 26, 2026
…refused, and hidden_act must be silu The unwaited hold read its picks with nonfinite ignored and a bad id as 0; the GPU copy into a verify window did the same. Both now refuse as the synchronous path does (DraftNonfinite): readPicks errors, and lfm2_draft_ids sets a flag the verify reads. Planted NaN, Inf and invalid-id tests cover the host and the kernel. The config refuses a hidden_act other than silu by name.
… the checkpoint and joins its workers on every path
… 0, the default), no quantizing at load The drafter's linears and Markov W2 stay bf16: new rows and tile kernels for the linears and a bf16 scan. An explicit --drafter-bits other than 0 is refused by name. 8-bit codes stay behind Options.bits for tools.
…mits and four planted controls The five launches are in lfm2_records.zig, with the limits that --offline compiled, not this GPU's. lfm2_attn_segs reserves its 128 threads in its source, as M1 and M2 cannot build it for another count. Four controls edit the source: a tie to the highest id, a nonfinite logit not flagged, and every stream reading or writing the first stream's conv state or kv cache. Each goes red at its gate.
…planted controls Merges the shared-round contract records. lfm2_tap's launch is in lfm2_records.zig with the limits --offline compiled, not this GPU's. Two controls edit the source: rows packed at the hidden size, and a column of another tap overwritten. Each goes red at the shapes gate.
…d six planted controls Merges the tap and shared-round contract records. The Markov scan and pick, the draft ids and the block attention are in lfm2_records.zig with the limits --offline compiled, not this GPU's. lfm2_attn_block is its own kernel now, with its 128 threads reserved in its source, as M1 and M2 cannot build it for another count. Controls edit the source: a NaN score not flagged, every stream reading the first stream's previous token, a pick past the vocabulary keeping its token, and a block row missing the last key. Each goes red at its gate.
…, the README, the recipe and the startup log --drafter-bits takes 0, 4 and 8. LFM2's drafter runs bf16 as released when it is not given or 0, 8-bit codes quantized at load on 8, and refuses 4. The other families refuse 8 by name.
…y gates 1-4, recorded M1/M2 limits and six planted controls The few-rows and 64x64 tile kernels and the scan over a bf16 W2 pass gates 1 to 4. Six controls edit the source: a lane skipping half of k, a group's last row, an untransposed W tile, a lost last column, a NaN score not flagged and every stream reading the first stream's previous token. Each goes red at its gate. The tile kernel compiles for 832 threads on M1 and M2.
M1 and M2 cap a pipeline's threads as its registers rise, and no run covers them. The server refuses them by name at open. An M3 is served, with one log line that it is not yet tested with shared rounds.
…ogged as untested with the drafter The check is its own, after the engine's: it uses the same generationServed helper, with a test beside it. The drafter kernels' records keep their M1 and M2 limits.
The shared-round kernels carry recorded M1 and M2 launch limits in the kernel contract, and no run covers those chips. The server no longer refuses them.
Contributor
Author
|
Thanks @ashhart. Pushed
|
BlivionIaG
marked this pull request as ready for review
October 26, 2026 21:37
6 of 9 tasks
9 tasks done
Owner
|
Thanks @BlivionIaG. A lot of this is fixed, and fixed well:
What still needs to change:
Once these are in and the stack is rebased on current main without #1379 and #1382, we will run drafted against plain, |
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.
Branch
series/dspark-drafter, head35e9255c0. Sits on the taps PR (#1390) and on #1378 and #1379. Part of #1374.What this changes
tensorfold servetakes--drafter DIRforlfm2_moecheckpoints with LiquidAI's DSpark drafter,LiquidAI/LFM2.5-24B-A2B-DSparkate052e5b1dd103c38cdd4f8ddc068d7dea4586f55. Its weights are under the LFM OpenLicense v1.0 and are not shipped; the recipe links the licence and says to read it before use. The lane core verifies
every draft, so a drafted reply equals its
"draft": falsereply.families/lfm2_moe/dspark*.zig). It runs the released bf16 matrices as they are, on bf16 kernels (lfm2_bf16_rows,lfm2_bf16_mma,lfm2_markov_scan_bf16).--drafter-bits 8is an explicit option that quantizes them to 8-bit codes at load and runs them on the family's kernels andcore/tile_qmm; it is named in--helpand the startup log.4is refused on LFM2 and8on other families.lfm2_markov_scanandlfm2_markov_pick(metal lfm2 markov step),lfm2_attn_blockandlfm2_attn_block_merge(the target's attentionkernels with a
FULLtemplate flag; the target's instantiations keep their causal count),lfm2_draft_ids.(
Drafted.held_on_target, fromseries/lanes-held). Shared rounds take the shared-round guard (series/lanes-guard).metal_admitplans with; a refused open frees what it loaded.DraftNonfinite, as the synchronous path does; planted tests for both paths fail without the fix. Ahidden_actother thansiluis refused by name. The fp32 reference reads the target embedding bits from the checkpoint and its hostmatmuljoins its workers on every path.LFM2_DRAFT_TRACEandLFM2_DSPARK_SYNCof the source history are gone.Receipts at this head
M5 Ultra (Mac17,15, macOS 27.0 build 26A428), stack head
2e5966f68, three alternated rounds, a fresh server per cell, no other model process,--parallel 16.bf16is the default,8-bitis--drafter-bits 8."draft": falseon one server--context 131072)exact.py8 of 8 equal in all three rounds for bf16 and for 8-bit; all 38 short-prompt replies token-equal drafted against"draft": false;bench_concurrent --alone --serialequal in every cell, 0 failed. Drafter weights 625 MiB (bf16) and 361 MiB (8-bit); accepted tokens per round 2.827 and 2.859.tf-lfm2-dspark, 2,048 and 16,384 context, M5): 9 of 9 equal in both modes; at 8,192, 32,768, 65,536 and 130,816 context (bf16): 9 of 9, first-row logit cosine 0.99998, 0.99997, 0.99992, 0.99995.tf-lfm2-exactat the DSpark stack head2e5966f68(it holds this PR's kernels), M5 Ultra (Mac17,15, macOS 27.0 build 26A428):--only lfm2SUMMARY 11 kernels: 11 pass gates 1-4, 0 fail, 0 known findings; 17 controls: 17 red as planted, 0 missed; 0 limit problems.--only tile_qmm6 pass, 0 fail.tf-lfm2-exact(real weights, LFM2.5-24B-A2B-MLX-4bit667408d4): 72 forward/cache checks passed, 0 failed, 0 skipped.zig build test test-lfm2 test-lfm2-dsparkat the stack head (M4 Pro, main8c36359d5merged in): 534 of 534 steps passed;zig fmt --checkclean;tools/zig/lean_check.py428 problems, the same as main.lfm2_dspark.metalline 35), the family testa Markov step picks the argmax of bf16(base + bf16(W2 . W1[prev]))fails and the contract saysFAIL shapes metal lfm2 markov step: at 1x128000x256: the pick's score is below the row's best by more than its bound."draft": falseper stream there); the one-stream rows compare on one server."draft": false, resume, cancel, stress; see lfm2: a cheaper DSpark hold (one concurrent encoder) #1394), and one indicative round on a busy machine gave one-stream code 143.9 tok/s with the drafter against 140.6 plain and chat 121.2 against 121.1; it is not a receipt.Not measured
Related PRs
Sits on #1390, #1378 and #1379. The scheduled PR (#1392) sits on this PR.