Repository navigation
lfm2: the DSpark confidence head schedules the verify (--draft-verify, --verify-costs) - #1392
Draft
BlivionIaG wants to merge 27 commits into
Draft
BlivionIaG wants to merge 27 commits into
BlivionIaG wants to merge 27 commits into
Conversation
…s.zig cli_capabilities.zig writes the capabilities document, but the Engines struct it reads still sat in cli.zig, which is 598 lines. The struct moves next to its writer and cli.zig re-exports it, so callers do not change. cli.zig drops to 589 lines. capabilities --json prints the same bytes as on main.
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.
… holds follow the grants A drafter that gives each held chain draft's chance (Drafter.VTable.chances) may set Facts.chances. Drafted.facts then sets Model.draft_probabilities and Model.granted_depth, so the window allocator prices each stream's verify rows by those chances, as it does a head's node chances. With Facts.confidence_floor it instead verifies the drafts at least that likely (a fixed window when 0). With granted_depth a stream's next hold is as deep as its last grant plus a slack: one more draft when the grant took every held draft, none when it stopped short of them (Stream.grant_short), and after a grant of none a probe of draft_slack drafts at the depth rule's probe interval. Facts.slack sets one slack for all three. A shared round's granted depths are cut so the drafting streams' blocks fit one hold pass of batch_rows. When the shared-round guard starts a drafted stretch, its first round grants every stream one draft, so the stretch measures drafted rounds rather than plain ones left by grants of none. All of it is behind Facts.chances and Model.granted_depth, off by default; windows.allocate records grant_short only for granted depths. Tests: chances price the chain's rows and drafted replies equal one-token decode alone and at 1, 4 and 8 shared streams, with the scheduled verify ragged across streams and verifying fewer rows than a fixed window; the facts reach the lane config; granted depths, their slacks and the one-pass cut; holds as wide as the grants.
…py, and holds kept only while they pay The depth rule asks one head draft for a stream with a copy ahead, and a shared round's granted holds were priced only by their verify rows. A drafter whose hold runs as a pass of its own (Facts.copy_first) pays that pass for the one draft, and the copy round then drops it; Model.copy_first now asks it for none. A drafter may also give its hold pass's timed cost by blocks and depth (Facts.hold_costs, priced by hold_price.zig). With granted depths a shared round then keeps granted holds, most expected drafts a row first, only while each raises the next round's expected tokens a ms over its verify rows and the pass; a probe after a grant of none keeps its hold. A stream's last grant records the drafts it expected to keep (Stream.grant_gain). Both are off by default (copy_first false, hold_costs empty). Tests: with copy_first no hold is asked for a copy round and replies equal one-token decode at 1 and 3 streams; a hold's cost between and past its timed passes; granted holds kept while they pay, none for a dear pass, and a probe's hold kept.
…re faster With granted depths a lone stream's verify rows are priced over the shared costs and the one-stream overhead, which is measured on drafted rounds and so holds the drafter's hold time; nothing compared a lone drafted round's tokens a ms with a plain round's. lone_guard.zig measures each lone stream's rounds of held drafts in stretches of 16 against plain rounds' tokens a ms, kept from stream to stream by sampling mode and context band and taken once in a plain stretch when unknown; copy rounds are left out of both. Drafting goes on while it beats plain by 5%; else the stream runs plain rounds and measures drafting again after 16 plain rounds, doubling to 512 while plain keeps winning. A drafted stretch's first hold drafts one and its slack, and a mode's first round is not measured. It runs only with granted depths and the shared-round guard on a timed engine (Facts.chances and Facts.shared_guard); shared rounds keep their own guard. Off by default. Tests: the guard's verdicts, rarer probes and kept plain rate; through the round loop a lone stream whose drafted rounds lose runs mostly plain and one whose drafted rounds win keeps drafting, replies equal to one-token decode.
…own_copy_rates) A shared round prices every copy window's rows at cfg.copy_rate, a fixed prior. A stream whose copies keep breaking (a different prompt on each stream, say) keeps paying verify rows for them. With Model.own_copy_rates a copy window is priced from the stream's own decayed copy history, by source (context or Recall), starting from four rows of copy_rate's prior; each verified copy window teaches its source's history. Off by default: without the flag allocation prices copies at copy_rate as before and the history is not kept. Test (copied_rate_test.zig): four streams' seven-row copy windows in one shared round, through windows.allocate. With the flag, after its context copies land one row of seven, a stream gets at most one copy row where the prior gave seven, the other streams keep theirs and its Recall copies keep copy_rate; without it every copy keeps seven rows and no history is kept.
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.
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.
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 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.
…bilities split under the LFM2 drafter
…ntract LiquidAI's DSpark drafter carries a confidence head: one linear over a draft's final block row (the row the target's head reads) and the W1 row of the token before it, 2,048 + 256 inputs, to one logit. lfm2_dspark_confidence (lfm2_dspark.metal) computes that logit for every row of a hold pass: 256 threads a row, fp32 products and sums, the eight simdgroups' partials summed in threadgroup memory. Nothing calls it yet. It registers with the kernel contract (metal_lfm2_confidence.zig): the served pass shapes against a float64 reference, the accumulation gate against reduced-precision sums, and gate 5 against every Apple chip we serve and this GPU.
…, --verify-costs) With the scheduled verify a hold also runs lfm2_dspark_confidence, and each lane keeps sigmoid(logit + bias) for its drafts. chances() gives their running product, the chance that a draft and every draft before it land, as Facts.chances: the lane core's window allocator then picks each stream's verify rows by those chances and the verify costs, and the next hold drafts as deep as the stream's grant. No second scheduler. While a hold still runs, its rows are priced by the lane's last settled confidences (draft j by the last hold's draft j, the deepest one repeated past it), so a verify does not wait on the GPU for them. A lane with none settled yet waits. Two serve flags on macOS, listed by capabilities --json and refused unless an lfm2_moe checkpoint is served with --drafter: - --draft-verify depth|fixed|scheduled: the core's depth rule, the whole block every round (the published fixed window), or the confidence schedule. scheduled is the default. - --verify-costs FILE: a table of verify forward ms by streams and rows a stream, written by tf-lfm2-verify-costs on the serving Mac, in place of the costs timed at load. Shared rounds keep the shared-round guard in every mode but fixed. Tests: test-lfm2-dspark reads a cost table, refuses malformed ones and checks the default mode; the server's parser tests take both flags and refuse an unknown mode. tf-lfm2-dspark checks the head against its fp32 reference.
…fied rows Each prompt chunk's rows were absorbed at once in a command buffer of their own, waited on before the first token, and the prompt's last row in a second one at the first draft request. The prefill pass is waited on before its rows are handed over, so a chunk that fits the lane's stage (128 rows) can be copied there on the host like a round's verified rows. The next hold, or a kept prompt state, absorbs the stage first, as before. Longer chunks are absorbed at once.
A DSpark hold is a pass of its own: the drafter runs its five layers and the target's head for every drafting stream before the verify. The lane core can now weigh that pass (Facts.copy_first, Facts.hold_costs, Model.own_copy_rates); the drafter gives it what it needs. - copy_first: a stream with a copy ahead verifies the copy, and the drafter is asked no hold for it. - hold_costs: after its warm-up the drafter times holds of 1, 2, 4, 8 and 16 blocks at one draft and at four (two at 16 blocks, one 32-row pass) on empty lanes, best of three. Each settled one-pass hold then moves the timed passes at its block counts a tenth of the way toward its own GPU ms, so the price follows the served contexts. - own_copy_rates: with a drafter, a stream's copy windows are priced by its own copies. Test: a 4-block hold at twice its price raises the 4-block passes by a tenth and leaves the others; one between 4 and 8 blocks moves both; a hold without GPU time teaches nothing.
…eal rows tf-lfm2-dspark times a shared round's hold over 2 to 16 streams, by the wall clock and on the GPU, and absorbs of 128 rows and of a whole absorb pass, so a hold's price can be checked against a verify's. tf-lfm2-verify-costs reads token id files after HISTORY: each of the 16 streams takes its history and then its verify rows from its own sixteenth of the files' ids (the history shrinks to fit), so a shared round's rows route the experts their own text routes. Without id files the rows are synthetic as before; the table's header names their kind.
The recipe names the confidence head, the three --draft-verify modes (scheduled by default), how the lane core prices a hold pass and a copy, the lone-stream guard, and --verify-costs with tf-lfm2-verify-costs.
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. |
Contributor
Author
Ah yes my bad the agent added pr mentions in the overlap resume check. Will remove them asap. |
This was referenced Oct 26, 2026
ashhart
requested changes
Oct 26, 2026
ashhart
left a comment
ashhart
left a comment
Owner
There was a problem hiding this comment.
Thanks @BlivionIaG. This builds on #1391, so it inherits that review's blockers (runtime lossy BF16 conversion, the 4-bit buffer mismatch, dropped async failure flags, the contract gaps). Here also: an explicit scheduled verify with no confidence head silently falls back to depth pricing; please refuse it by name. (Your cost-table non-finite admission fix is good.)
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-scheduled, head7dd1a2ef8. Sits on the drafter PR (#1391) and on #1382 (series/lanes-lone-guard,b66a7629e, which holds #1381 lanes-chances), #1380 (series/lanes-copy-rates,1e718458f) and #1375 (series/capabilities-split,06a5236ce), joined by the mergedc866a541(no code of its own). Part of #1374. Draft.What this changes
lfm2_dspark_confidence: the drafter's confidence head (one logit from a draft's final row andW1[previous token], 2,304 inputs, fp32 sums), registered with the kernel contract, gate 5 over every Apple chip we serve.Facts.chances); the lane core's window allocator pickseach stream's verify rows from those chances and the verify costs, and the next hold is as deep as the grant. No
second scheduler. A running hold is priced by the lane's last settled confidences, so a verify does not wait for it.
--draft-verify depth|fixed|scheduled(defaultscheduled) and--verify-costs FILE(a table from the newtf-lfm2-verify-costs), macOS serve flags listed bycapabilities --json, refused unless anlfm2_moecheckpoint isserved with
--drafter(the refusal sits inmetal.zig'srefusal(), so the file stays at 598 lines).Facts.hold_costs, timed at load, corrected toward eachsettled hold's GPU time), asks no hold for a stream with a copy ahead (
Facts.copy_first) and prices a draftingstream's copies by its own (
Model.own_copy_rates). Withscheduledthe lone-stream guard ofseries/lanes-lone-guardruns.
tf-lfm2-dsparktimes shared holds of 2 to 16 streams;tf-lfm2-verify-costsmay read real token rows.Receipt on this head
Decode, M5 Ultra, aggregate tok/s, min-max of the rounds' medians (median):
9df2cd98257608634b7dd1a2ef8exact.py8 of 8;--alone --serial1/4/8/16 streams, same and distinct prompts:alone 261/261, serial 37/37 over 12 cells, round 1; one token-hash list per cell over all builds and rounds: yes; cancellation
PASS; resumed turn B equal to fresh: yes (turn B cached 3859 resumed, 0 fresh).
processes during cells none.
Gates and checks
Gates on the new head
7dd1a2ef8(M4 Pro): lean 433, 433 on main40c6406b3, 0 new; fmt (57 files) and diff --check clean;test-lfm2+test-lfm2-dspark99/99; lanes 276/276; lane API 297/297; lane lone 19/19;test-server-cpu787 passed, 3 skipped;test-goldenok.nativeok; contract--only lfm2SUMMARY 8 kernels: 8 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems;--only tile_qmmSUMMARY 6 kernels: 6 pass gates 1-4, 0 fail, ... 0 limit problems.Exactness with
--drafter(defaultscheduled), M5 Ultra and M4 Pro:exact.py8 of 8;--alone --serialat 1, 4,8, 16 streams, same and distinct prompts, all equal (distinct 16 streams 48/48 and 16/16); one token-hash list per
cell over every build and round; cancellation PASS; resumed turn B equal to the fresh full-prompt run.
Confidence head against its fp32 reference (
tf-lfm2-dspark, 698-token context, an earlier build of this code): worstdifference 0.00123 (M5 Ultra), 0.00215 (M4 Pro).
Kernel contract, M4 Pro at this head:
SUMMARY 8 kernels: 8 pass gates 1-4, 0 fail, 0 known findings; 0 controls: 0 red as planted, 0 missed; 0 limit problems;--only tile_qmm6 pass, 0 limit problems. M5 Ultra (stack head):SUMMARY 8 kernels: 8 pass gates 1-4, 0 fail, 0 known findings; ... 0 limit problems.Failure the contract catches: with the simdgroup sum taken in half precision
(
d = float(simd_sum(half(d)))),FAIL shapes metal lfm2 dspark confidence: at 1x1x2304 output 0 is -3.258789e1 where float64 gives -3.259301e1: 5.59 of its bound,SUMMARY 1 kernels: 0 pass gates 1-4, 1 fail.Host tests (M4 Pro, first full run at an earlier build of this code with the refusal four lines longer in
metal.zig; thefinal run reused those results):
test-lfm2andtest-lfm2-dspark99/99;test-shared-lane-regressions-host266/266;
test-lane-api-host297/297;test-lane-lone-host19/19;test-lane-halt-host18/18;test-lane-interrupted-host18/18;test-root-copy-host31/31;test-server-cpu787 passed, 3 skipped;test-goldenok;test-cuda-host1044 passed, 9 skipped; HIP and CUDA host checks,native,dev-toolsok.Not measured
--draft-verify fixedanddepthcellsat this head;
--verify-costswith a measured table at this head; sampled replies on this head (measured at the stack head in lfm2: a cheaper DSpark hold (one concurrent encoder) #1394); 128k and the native window on this head (measured at the stack head in lfm2: a cheaper DSpark hold (one concurrent encoder) #1394); M1 to M3; Linux CI; a plainzig build test.Related PRs
Sits on #1391, #1382, #1380 and #1375. The capabilities split (#1375) keeps
cli.zigunder 600 lines with the two new flags (595). #1393 sits on this PR.