Repository navigation
Qwen3.8-Flash-Next faster kernels in Zig engine - #788
Merged
Merged
Conversation
7 tasks
oQ checkpoints write per-module widths as language_model.model.layers.N..., while loaders ask for model.layers.N... or the bare path, so an exact-key lookup missed every override and read the tensor at the global width. The table keeps each entry under the config's own key, as before, and also indexes it under its canonical module key (no .weight suffix, no known prefix, the rule the 27B and Flash Next families use); `width` takes the key as given, then its module. Two spellings of one module that declare different widths refuse the table as unsupported, the error `detect` already returns, so its callers' error sets are unchanged; the same width twice is fine. Test that fails before: the fixture reads a 5-bit down_proj declared under the oQ prefix by the loader spelling; on main it fails the packed shape check (read at the global width). Also: every prefix form and a trailing .weight name one module; aliases must agree; an override in any spelling answers every spelling and stays under its own key. Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 539 problems (main 539), no new findings.
… tables repeated for a smaller group A projection of mixed-width members (the Flash Next hyper-connection down and inject pair, where 0.6.6's QWeights.of widens) reads as one matrix at their common width, the most bits and the smallest group. Narrowing and a group that does not divide the source group are refused. Used only where 0.6.6 widens; no caller on this commit. Tests (widen.zig, new): every source width and every wider target dequantize to the same fp64 values, bit for bit; a stack of experts widens row for row; the same width comes back as it is; narrowing bits, growing groups and widths the kernels do not run are refused; the common width takes the most bits and the smallest group; fused members keep each one's values in order. Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 539 problems (main 539), no new findings.
…wider output Chips without tensor units now decode every MLX affine width (2, 3, 4, 5, 6, 8 bits) in the row kernel: 3- and 5-bit rows unpack eight codes from three or five bytes, low code first, as MLX packs them. The 2-, 4-, 6- and 8-bit paths are unchanged. A 3- or 5-bit row must end on a word, as the checkpoint stores it. callStrided launches the row kernel with an output row stride, so one member of a fused stack (one run of equal width) writes its columns straight into the stack output. The default call, its dims and its kernels are unchanged. tf-row-widths-check (zig/build/core_checks.zig) runs it on the GPU: each width at groups of 32, 64 and 128 against an fp64 product from the MLX reference dequantization, within the fp32-sum bound; one row alone against the same row in 2-, 3-, 8- and 16-row windows, relu2, indexed slots and a strided output, bit for bit; a fused stack of 48-column members at mixed widths (4-bit among them, group 64) written through the stride at windows up to 48 rows, each member as its own call; and, given the unchanged kernel source, the widths it already served byte for byte against it. Tests that fail before: "shapes the row kernels serve" now validates 3 and 5 bits (main refuses them), "3- and 5-bit rows take whole words of codes", "every width has its own pipeline slot" and "a strided call writes rows at least a matrix wide" (no callStrided on main). tf-row-widths-check on an M3 Ultra (no tensor units), given main's row_projection.metal: PASS 64512 values within the fp32-sum bound (worst err/bound 0.9400), 408192 bit-equal pairs, 301056 outputs byte-equal to main's kernels at the widths main serves (3 and 5 bits have no main kernel to compare). Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings.
…d 16-row members Format accepts 2-, 3- and 5-bit codes beside 4, 6 and 8, in groups of 32, 64 or 128. The non-4-bit path widens 2- and 3-bit codes to nibbles for the bf16 x uint4 op and 5-, 6- and 8-bit codes to bytes for the uint8 op, the staging the Python lane matmul uses, so each width keeps that arithmetic. 3- and 5-bit groups of 32 are an odd number of words a lane; they are fetched a word at a time instead of in pairs, which would drop the last word. Layout.y_stride (default N) sets the elements between output rows in every kernel (tf_lane, tf_lane_pair, tf_lane_reg), so a member of a fused stack writes its columns of the fused rows; the output bound checks follow it. The store of the simdgroup-matrix layout (Macs without tensor units, ashhart#679) takes the same stride and keeps its column guard. N may be a multiple of 16 on the plain kernel: the last 32-column tile is packed with zero codes past N, its metadata reads and stores are already guarded by N, and weightWords counts the padded tiles (sourceWords is MLX's W). Cooperative and reg layouts still need N % 32 == 0. encodeSums checks its sums buffer by the group count, as ashhart#698 sizes the dispatch. Without tensor units the admitted layout is unchanged: 6-bit group 32 on the plain kernel (ashhart#679); every other width still answers TensorUnitsRequired there, now checked at 2, 3 and 5 bits too. With the defaults the 4-, 6- and 8-bit kernels preprocess (metal -E) to the same text as main's for every group, schedule, output, prepared-sums/reg layout and both fragment layouts tried (148 variants), so the uniform paths keep their kernels. Host tests that fail before: 2, 3 and 5 bits validate; y_stride reaches the source and below N is refused; a 48-row member validates (cooperative and reg refuse it); every width x group x {64, 48} rows packs MLX codes bit for bit with zero padding. The device program covers every width, 48-row members and strided output on tensor units; without them (M3 Ultra): PASS 14 layouts, 1904 bit-equal row/window pairs, 208 windows equal to fn_lane's recorded bytes, 7424 CPU values (max roundoff-bound fraction 0.0144), 10 refusals, and a 48-row member and two strided layouts at 6/32 equal the unstrided rows with the gaps untouched. fn_lane's recorded windows are compared only on the layouts they were recorded at. Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings.
affine_mm took 4- and 8-bit weights only, so a 2-bit matrix (oQ2's routed experts) had no tiled prompt product. At 2 bits a thread loads its 32 codes of a step as 8 bytes and writes each code into the bf16 tile as itself, with no offset (as at 8 bits; 0..3 is exact in bf16), so the group sums are the 4- and 8-bit path's. A first try with the 4-bit "+128" offset stayed within the float64 bound, but a reference with scales 1% off no longer exceeded it at K = 4096, so the codes enter exactly. The 4- and 8-bit sources preprocess as before. Test that fails before: "dense and gathered, each format, against exact sums on the host" adds 2/64 and 2/32, dense and gathered; source() refused 2 bits. The 4- and 8-bit sources preprocess (metal -E) to main's text at both groups, both outputs and both fragment layouts (16 variants). tf-affine-ref (next commit) holds the 2-bit products to float64. Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings.
…y row's bits in every window The device side of ashhart#633's quant-ref for the shared prompt products, item 5 of ashhart#716. tf-affine-ref takes no arguments and needs no model: seeded synthetic codes, scales, biases and inputs, packed as MLX affine. The engine module now names core/affine_mm so a check program can reach it. - affine_mm at 2, 4 and 8 bits, groups 32 and 64, bf16 and fp32 out, dense and as a one-expert gather, at [192, 1024] and [256, 4096]. - simd_qmm (4-bit, group 64) as tiles and as the scalar twin at 1 and 2 rows, at [40, 512], [256, 4096] and [6400, 256] (K chunks of 32, 16 and 8). The reference dequantizes the same codes in float64 and multiplies in float64. Each product prints its worst column (max abs error, and that error over the column's largest |y|) and its largest error as a fraction of the bound: (group + K / group + 8) fp32 roundings of the terms' magnitudes as the kernels sum them (4-bit affine_mm codes plus 128), plus one bf16 ulp for bf16 outputs. A reference with its scales 1% off must exceed the bound. Then every row of a 129-row call must equal, byte for byte, the same row computed in windows of 1, 2, 3, 8, 16, 17, 33, 64 and 65 rows (every full window over the 129 rows, so each row is seen at every position), which crosses the 16-row fragments, the 32- and 64-row tiles and the gather's switch between them. The twin must equal the tiles wherever twinExact says it does, since only then do callers take it. Receipts, M3 Ultra (no tensor units: nax.h's simdgroup-matrix layout), on main 0efbd44: - every product within 0.4935 of its bound; every row's bits equal at every window, both kernels, every format; the twin equals the tiles at 1 and 2 rows on all three shapes, as twinExact says. - worst columns, bf16 out (dense and gather equal), at [192, 1024] and [256, 4096]: affine_mm q2 g32 2.64e-2 (rel 3.23e-3) and 6.22e-2 (3.78e-3), q2 g64 2.87e-2 (3.51e-3) and 3.12e-2 (2.47e-3), q4 g32 2.45e-2 (2.65e-3) and 3.12e-2 (3.01e-3), q4 g64 1.56e-2 (2.89e-3) and 3.12e-2 (2.53e-3), q8 g32 0.496 (3.41e-3) and 0.968 (3.63e-3), q8 g64 0.453 (2.87e-3) and 0.956 (3.38e-3); simd_qmm 1.19e-2 (2.57e-3), 3.13e-2 (2.87e-3) and 1.54e-2 (3.74e-3). All under 2^-8 of the column's largest |y|; 0.20 to 0.49 of the bound. - fp32 out: at most 0.0081 of the bound (q2 1.5e-6 to 7.7e-6, rel 1.9e-7 to 6.0e-7). - scales 1% off: 2.3 to 29 times the bound for bf16 outputs, 2.7 to 391 for fp32, 44 to 273 for simd_qmm. Not run: an M5, where affine_mm takes the tensor units; `zig build tf-affine-ref && ./zig-out/bin/tf-affine-ref`. The fp32 bound is the worst case for any summation order, so a changed order still passes; the measured errors sit two orders below it, and the bit checks are what pin the order. Gate on main 0efbd44 (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test: `get` logs an error on the DeferredTensor path the test expects). lean_check: 536 problems (main 539), no new findings.
…unds buffers to engine_rounds.zig Moved, to make room under replay.zig's and engine.zig's lengths for the widths that follow: DepthRule, its test and copyDrafts (re-exported from replay.zig, the test reached from replay's test block), and the body of Engine.rounds (Engine keeps a one-line rounds, so the load test that reads the source still finds `try e.rounds();`; RING is now pub for it). Each moved doc comment is one line now; no behaviour changes. replay.zig 3321 -> 3265 lines, engine.zig 899 -> 857. Gate on main 0efbd44 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects); native and tf-flashnext-run (with flashnext_bench), -reuse, -pack and -think build. lean_check: 531 problems (main 539); the only findings not on main are the two long files' shorter line counts.
…ine still runs 6-bit group 32 only
Admission took one global 6-bit/group-32 width, and the packs refused group 128 and fused members of mixed widths.
- Widths (widths.zig): each module's width through core/quant (module keys under any prefix; disagreeing spellings
and half-declared overrides refused; `false` and `{}` mark plain modules). Admission reads MLX affine 2, 3, 4, 5, 6
or 8 bits in groups of 32, 64 or 128 for every text and MTP matrix, refuses the first other one under its
checkpoint tensor name, and an n-gram table whose group does not divide its 160-wide rows.
- Lane packs follow 0.6.6's _stacked: a fused projection's members grouped by bits in first-seen order, adjacent or
not, each group one run at its smallest group, regrouped exactly (core/quant widen; group 128 read as 64), with
.runs/.members tables. A run whose columns are not a multiple of 32 keeps MLX's layout untiled (decode.py's nt 0;
oQ4e's layer-21 GDN input). A projection of one width keeps its single run.
- Hyper-connections: a down and inject pair of two widths widens exactly to the wider (QWeights.of); the up
projection packs at its stored width. Prompt (pack_mlx) projections of several widths keep each member's bytes in
order with a .parts table. Experts, embedding and n-gram tables, which the engine reads straight from the
checkpoint, are shape-checked at their declared widths when the pack is built.
- oMLX's layout: an MTP head kept at top-level mtp.* is read where the index keeps it; an n-gram table's weight_scale
(oQ conversions of the FP8 checkpoint keep its scale, as oQ4e does) must be one finite, nonzero value, else the pack
is refused by name; the next commit applies it.
- Pack format 3, whose identity also hashes config.json, so a config-only change rebuilds the pack; packs under the
format-less identity are rebuilt, never reused. Pack building frees every buffer when an allocation fails.
The uniform 6/32 packs keep every tensor byte for byte (a digest pinned before the change). The engine refuses any
other width at load until the next commit gives the device its kernels.
Tests that fail before: widths_test.zig; load_admission_test.zig's overrides at other widths, global groups 64/128,
refusal under the tensor's name, oMLX's mtp.*; pack_lane_test.zig (runs by bits, a pack per run, ragged runs
untiled at oQ4e's real rows 10240/6144/48/48, allocation failures free everything); pack_hc_test.zig (oQ4e's
4/64 down with a 5/64 inject as one 5/64 matrix; the up projection at its stored width; mixed prompt members);
pack_widen_test.zig; pack_test.zig's half-declared and contradicting raw tensors; pack_source_test.zig's config-only
change; pack_cache_test.zig's format-less identities.
Gate on main 0efbd44 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has
309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's
split-checkpoint test logs an error on the path it expects); native and tf-flashnext-run (with flashnext_bench),
-reuse, -pack and -think build. lean_check: 531 problems (main 539); the only findings not on main are the two
long files' shorter line counts.
…nits, in decode and in prompt - Lanes (decode, the MTP head, lane rounds; lane_w.zig): the recorded lane kernel's text compiled per (N, K, bits, group): codes staged as bytes, group sums over the group, K slices as lane_qmm.split_k; one launch a run, columns placed through the .members map (FZ_MAP in the kernel's store), untiled runs as MLX lays them out. A 6/32 lane keeps its recorded kernels and launches; a lane at other widths holds no 6/32 pair, so a 6-bit kernel can't read one of its runs (LaneNot6Bit32). - Hyper-connections, experts, embedding and n-gram tables (variants.zig): 0.6.6's any-width qa_ kernels, recorded at 6/32, launched at the tensor's BITS/GS (WB/WG, SWB/SWG for the experts). The 6-bit fast paths (fz_xgu/fz_xdown, fz_dense, fz_hc_up) run only on 6/32 tensors; an expert whose gate and up differ in width, and n-gram shards of different widths, are refused (one kernel reads both). - Prompt products (prompt_w.zig, prefill/qmmw_nax.metal): qmm6_nax's tiles and fp32 sums with codes read at WB bits in groups of WG: dense, split-K (the hyper-connection down), sorted-expert gathers and the fused gate/up; a fused projection of several widths runs one product a member into its columns. Every prompt weight is registered with its width at load, and a launch on an unregistered one fails the pass by name (PromptWeightWidthUnknown). - Fail closed: every width kernel compiles at load (engine_warm.zig), so a compile failure refuses the load, not a request. Width tables are read on the host, never through device buffers. Speed-up mode (two Macs) splits the 6/32 kernels only and refuses any other checkpoint by name (SpeedUpNeeds6Bit32). The replay harness refuses other widths (HarnessNeeds6Bit32). The config and index are read once at load. - The n-gram table's weight_scale (ple_scale.zig): oQ conversions of the FP8 checkpoint store the unscaled rows and oMLX's runtime multiplies each looked-up row by the scale (bf16 row times bf16 scale, one rounding). One in-place kernel after the lookup, in decode and in prompt, compiled at load only when the scale is not 1. The uniform 6/32 checkpoint keeps every recorded kernel and launch. Tests that fail before (none of these files exist on the previous commit): lane_w_test.zig (split_k gives every recorded kernel's SK; every recorded lane text takes any width, a row stride and a column map and keeps its arithmetic; runs land at their members' columns; width tables read on the host; a lane at other widths gives a 6-bit kernel nothing to read), engine_widths_test.zig (oQ2's experts, embedding and n-gram tables at their own widths; the uniform checkpoint keeps every recorded kernel; mixed gate/up and n-gram shards refused; the 6/32-only harness refuses others), prompt_w_test.zig, ple_scale.zig's scale against the bf16 product on the host. Gate on main 0efbd44 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has 309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's split-checkpoint test logs an error on the path it expects); native and tf-flashnext-run (with flashnext_bench), -reuse, -pack and -think build. lean_check: 531 problems (main 539); the only findings not on main are the two long files' shorter line counts.
…loat64, each row's bits in every window
No model files: seeded synthetic codes, scales, biases and inputs, packed as the engine packs them. Every projection
class at bits {2,3,4,5,6,8} x group {32,64,128}, decode and prompt separately, plus each public checkpoint's mixed
stacks and oQ4e's ragged one: lanes at K slices 1, 2, 4 and 8; hyper-connections (down partials and up); experts
(gate/up through silu, down); the embedding and the n-gram tables (8-bit codes over their full range); prompt dense,
split-K, expert gathers and the fused gate/up. Each check reports its worst column against a float64 bound, and a
reference with scales 1% off must exceed it (negative control). Every row's bits must be equal alone and in every
window (1-16 rows in decode, 1-65 in prompt), and the generic texts at 6/32 must give the recorded kernels' bits.
`tf-flashnext-widths lane|hc|moe|prompt` runs one class.
tf-flashnext-widths (M3 Ultra, no tensor units): PASS 2722 checks, worst column 0.9947 of its float64 bound (bf16
output rounding); all 262 negative controls exceed the bound, the least by 1.012x; every row's bits equal alone
and in every window. The core lane check without tensor units passes as on the core commits.
Qwen3.8-Flash-Next-MLX-oQ2-MTP @0a1bf6f on an M3 Ultra (its cached pack; loads in 21.1 s from a cold file cache,
9.0 s warm): drafted == `"draft": false` 4/4 (MTP acceptance 55-59%; plain 45.0-54.1 tok/s, drafted 84.3-91.2);
resumed == fresh (a 5.3k-token prompt at 0 and 5262 cached tokens; the second turn resumed at 5262 == a server with
`--prompt-cache-gib 0`); concurrent == solo, 2 streams; greedy text sane (RYB and RGB primaries, a doctested
function, "5:45 PM").
Gate on main 0efbd44 + the core commits (M3 Ultra, Zig 0.17.0): `zig build test test-golden -Dcpu=apple_m1` has
309/319 goldens equal and 0 unexpected; every step passes but one that fails on main too (core.checkpoint_metal's
split-checkpoint test logs an error on the path it expects); native and tf-flashnext-run (with flashnext_bench),
-reuse, -pack and -think build. lean_check: 531 problems (main 539); the only findings not on main are the two
long files' shorter line counts.
…-row kernels without tensor units Without tensor units Flash Next's decode took the lane kernels at every window width: their -lanes twin reads a 1-row window's weights at 240-640 GB/s on the M3 Ultra. As 0.6.6 does before M5 (rows.py), decode now runs each dense projection one input row a simdgroup (fn_rows.metal: q4_qmv_rows at 4/32, qa_qmv_rows at any MLX width, with MLX's qmv_fast inner loop and fp32 sums), and for other widths than 4/32 from 4 rows a matrix-unit twin that keeps those sums (qa_qmv_rows_mma over qa_row_block_sums). Hyper-connections take fn_rows_hc.metal the same way, each row its own threadgroups, a 4/32 pair's 8-row tiles from 8 rows. rows_load.zig loads every dense lane's members in MLX layout at their own widths beside the lanes (TP and the 6-bit fast paths keep the lanes), checks each twin and tile against the per-row loop on 8 and 12 rows at load and keeps the loop where one differs. Model.lane and the hyper-connection calls take the rows when they are loaded and there is no TP. With FZ_M3_LANES set, decode keeps the lanes. A row's bits depend on its own inputs alone at every window width, so drafted equals plain, resumed equals fresh and concurrent equals solo as before. They differ from the lane kernels' bits, on oQ checkpoints and the uniform 6/32: the fp32 sums run in another order, within the same float64 bound. Tests: tf-flashnext-widths rows (dense) and rows-hc: every width 2/3/4/5/6/8 in groups of 32/64/128 on four shapes and the decode shapes at the public checkpoints' widths, against float64 with a 1%-off-scales control, every row's bits equal in windows of 1, 2, 3, 4, 5, 8, 9, 16, 17 and 33 rows against a 64-row call; the checkpoints' mixed stacks as the engine attaches them, each member's columns equal its own product's in every window; hyper-connections at every oQ width, rows 1-16, against float64 and the recorded 6/32 kernels' bits on the same values. rows-time prints the lane and row kernels' GPU time at decode shapes, 1 to 16 rows.
…t products at 2-4 rows, the same bits At 2-4 rows each simdgroup of qa_qmv_rows cut every code from its row again for its own input row, and the MMA twin from 4 rows ran 8-row tiles, so a 6/32 product at 4 rows took 1.8x its 1-row time and an 8/64 one 1.3x. qa_qmv_rows4 runs four input rows a simdgroup: each code is cut once and each row keeps qa_qmv_rows' lanes, fma chain and fp32 sums in the same order, so every row's bits are qa_qmv_rows' bits. rows_w picks the kernel by width and row count (all three give a row the same bits): 6 and 8 bits take the four-row kernel at 2-4 rows and the twin past them; 4/32 keeps the loop; the other widths keep the loop to 3 rows and the twin from 4, where the four-row kernel measured slower. The load check runs the four-row kernel beside the twin; a member stops taking whichever differs from the loop (the twin gives way to the four-row kernel, that one to the loop). Tests: tf-flashnext-widths rows: every window now crosses the four-row kernel at 2-4 rows against the 64-row call's twin, every width and decode shape. rows-time prints the loop, four-row and twin times apart.
…r units On the row kernels a round's time grows with its window (a 16-row copied window costs about twice a 1-row one, where the lane kernels' cost is nearly flat), so the depth rule tuned for the lanes drafts too deep. As 0.6.5 does for one stream before M5 (DraftDepth), drafts_cost.CostDepth keeps each depth's landing rate and each window width's measured round time (the GPU gap between consecutive read rounds) and picks the depth with the most expected tokens a millisecond: head rounds up to 3 drafts, measuring each depth once first and probing one deeper or shallower every 8 rounds; copied rounds up to 15 drafts at their landing rate. TP and tensor-unit GPUs keep the existing rule. The depth changes which rows share a window, never a row's bits: drafted output still equals "draft": false. Tests: drafts_cost.zig's unit tests: unmeasured depths first, then the best rate of tokens a millisecond; landing rates move toward kept drafts; costs interpolate between measured widths and extend past the widest by rows.
grantoverton
force-pushed
the
oqxe/speed-flashnext
branch
from
October 13, 2026 20:21
2759e70 to
b1a51ee
Compare
ThinkOffApp
pushed a commit
to ThinkOffApp/TensorFold
that referenced
this pull request
Oct 14, 2026
ThinkOffApp
pushed a commit
to ThinkOffApp/TensorFold
that referenced
this pull request
Oct 14, 2026
The contributor's three commits landed above with the same code. Our commit after them wraps the lines past 120 columns. This keeps the PR's head in the history with no change to the tree.
Owner
|
Thank you, @grantoverton. On Macs without tensor units, Qwen3.8-Flash-Next's dense projections and hyper-connections run one input row a simdgroup, with a matrix-unit twin from 4 rows, and the draft depth is chosen by measured round cost. This changes the uniform 6-bit checkpoint's tokens on those Macs. Its dense products' fp32 sums run in another order, inside the same float64 bound, and Your commits are in 1.0.10 under your name ( On your M3 Ultra, at |
Open
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
2 participants
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.
Rebased onto main c36b536 (1.0.9) and #774's rebased tip on 2026-10-13: no conflicts in this change's code; the tip passes the gate again. The receipts below are the pre-rebase runs.
Part of #716; stacks on #774 (which stacks on #755). This branch carries their ten commits until they merge, so please review the last three commits here (from "Flash Next decode runs dense projections and hyper-connections on per-row kernels" on). Without tensor units, Flash Next's decode runs every dense projection and hyper-connection on the lane kernels, whose
-lanestwin reads a one-row window's weights at 240-640 GB/s on an M3 Ultra. As 0.6.6 does before M5, this runs them one input row a simdgroup (with a matrix-unit twin and, for 6- and 8-bit products at 2-4 rows, four rows a simdgroup, all three giving a row the same bits), and picks the draft depth by measured round cost. On an M3 Ultra the uniform 6/32 checkpoint decodes 1.06× (code) and 1.19× (chat) faster, oQ2 1.02× and 1.14×, oQ3 1.04× and 1.11×, oQ4e 1.06× and 1.21×; prefill is unchanged. GPUs with tensor units keep every kernel and launch they run today.This changes the uniform 6/32 checkpoint's tokens on GPUs without tensor units (M1 to M4): its dense products' fp32 sums run in another order, within the same float64 bound, and fidelity against oMLX is below. With
FZ_M3_LANESset, decode keeps the lane kernels and their tokens.What this changes
fn_rows.metal:q4_qmv_rowsat 4/32 andqa_qmv_rowsat any MLX width, one input row a simdgroup with MLX'sqmv_fastinner loop and fp32 sums, and for other widths than 4/32 a matrix-unit twin from 4 rows that keeps those sums (qa_qmv_rows_mmaoverqa_row_block_sums).fn_rows_hc.metal: hyper-connections each row in its own threadgroups, a 4/32 pair's 8-row tiles from 8 rows.rows_load.zigloads every dense lane's members in MLX layout at their own widths beside the lanes (TP and the 6-bit fast paths keep the lanes), checks each twin and tile against the per-row loop at load and keeps the loop where one differs.Model.laneand the hyper-connection calls take the rows when they are loaded and there is no TP.qa_qmv_rows4cuts each code once for four input rows, each row keepingqa_qmv_rows' lanes, fma chain and fp32 sums in order.rows_wpicks the kernel by width and row count fromrows-time: 6 and 8 bits take it at 2-4 rows and the twin past them; 4/32 keeps the loop; the other widths the loop to 3 rows and the twin from 4, where the four-row kernel measured slower.drafts_cost.zig, 0.6.5'sDraftDepthfor one stream). On the row kernels a round's time grows with its window, so the lane-tuned depth rule drafts too deep.CostDepthkeeps each depth's landing rate and each window width's measured round time and picks the depth with the most expected tokens a millisecond: head rounds up to 3 drafts, copied rounds up to 15. TP and tensor-unit GPUs keep the existing rule. The depth changes which rows share a window, never a row's bits.Exactness
TensorFold/Qwen3.8-Flash-Next-MLX-oQ2-MTP @0a1bf6f, TensorFold/Qwen3.8-Flash-Next-MLX-oQ3-MTP @993186e, Jundot/Qwen3.8-Flash-Next-oQ4e-mtp @2615fc0 and #774's uniform 6/32, greedy, MTP drafts on unless
"draft": false;serve <checkpoint> --name flashnext --parallel 2 --context 16384 --prompt-cache-gib 0. Before is #774's tip b29d032; after is a19b2cb1, this branch's tip 2759e70 but for one doc comment joined onto one line.bench_concurrent's alone and serial checks at 1 and 4 requests, both prompts: 34/34 on each, both runs); cancellation, reset and partial acceptance, 25 steps, pass (uniformacd6c2129c4f/d2ed982f7f4b, oQ20fb078e5ec87/e7bf591da8fa). oQ3 and oQ4e, after: concurrent equals solo (34/34 each) and the 25 cancellation steps pass (oQ3a5c773efe57b/fcd1fba4ccc0, oQ4edc4869db94b0/96c25ce8bedb).cabc8bb121c8→af87d66e6194, chat7c3f61c8db2c→b294a0e6e95e; oQ2556545c4fd22→e6305acf5ab2,64b3971d709b→2710d23bcad7; oQ3 and oQ4e keep their code tokens (02a5b69658f0,581df06bae38) and change chat (cdf91657573d→d59e34de846e,33924006fc9d→03acda6865af).tf-flashnext-widthsat the tip: PASS, 4,486 checks, worst 0.9947 of the float64 bound. Therowsclass (1,038 checks, worst 0.9844): every width 2/3/4/5/6/8 in groups of 32/64/128 on four shapes and the decode shapes at the public checkpoints' widths, against float64 with a 1%-off-scales control; every row's bits equal in windows of 1, 2, 3, 4, 5, 8, 9, 16, 17 and 33 rows against a 64-row call, so the loop, the four-row kernel and the twin are checked against one another at every width; the checkpoints' mixed stacks as the engine attaches them, each member's columns equal its own product's in every window.rows-hc: hyper-connections at every oQ width, rows 1-16, against float64 and the recorded 6/32 kernels' bits on the same values.Performance
Mac Studio, M3 Ultra (80-core GPU), 256 GB, macOS 26.4.1 (25E253), Zig 0.17.0. Fresh builds,
zig build native -Dcpu=apple_m1;tensorfold-nativesha25692dca1896a14…(before) and0c8a03935857…(after). Every run starts from a fresh server, with other checkpoints' file cache evicted first; before and after alternate (before, after, after, before).Decode and concurrency.
bench_openai.py … --tokens 256 --reps 3 --temperatures 0(medians;fibonacci-raw, code, andgpu-chat-no-think, chat) andbench_concurrent.py --levels 1,4 --alone --serial(aggregate tok/s; 4 requests share the 2 streams):bench_openaicodeAgainst before: uniform 6/32 code 1.06× and chat 1.19×; oQ2 1.02× and 1.14×; oQ3 1.04× and 1.11×; oQ4e 1.06× and 1.21×; the concurrent cells move the same way. The chat prompt gains most: its windows are short (head drafts, 1-4 rows), where the row kernels are fastest. The code prompt copies long runs, whose wide windows the lanes ran almost for free; with the depth picked by cost it still gains. oQ4e now decodes code at 1.02× and chat at 1.03× of the uniform checkpoint, oQ3 at 0.92× and 1.04×. The oQ3 before-run is #774's receipt on the branch before the n-gram scale, which oQ3 (no
weight_scale) never reaches; oQ4e's is #774's on its tip.Cold prefill (uniform 6/32; before and after). Exact token arrays: upstream's public prefill paragraph ("A lighthouse keeper writes in the log as the boats cross the bay.",
tools/zig/qwen27_gate.py) repeated and tokenized once without special tokens;/v1/completionswith the ids,max_tokens 1,--parallel 1 --context 34816 --prompt-cache-gib 0, every request cold (cached_tokens0), a 2k warmup, one rep a size. The prompt kernels are unchanged:Per launch (
tf-flashnext-widths rows-time, synthetic weights at decode shapes over at least 1 GB of copies, µs a launch, median of three command buffers; lane is #774's path, after the kernelrows_wpicks; the receipt's run also timed a 4/32 four-row kernel, dropped as slower than the loop):The lane kernels' time is nearly flat in the row count and the row kernels' is not, which is why commit 3 picks the depth by measured cost: one-row and short windows gain, and the widest copied windows are drafted less often.
Memory
The row kernels read each dense projection in MLX layout beside the lane pack: +0.8 GB after load on the uniform 6/32 (184.25 → 185.02 GB physical footprint) and +0.9 GB on oQ2 (96.31 → 97.25 GB), oQ3 (118.19 → 119.13 GB) and oQ4e (131.97 → 132.91 GB). With
FZ_M3_LANESset, none of it is loaded.Fidelity
Teacher-forced against oMLX 0.7.1.dev1 on the same checkpoints and the same fixed histories as #774 (oMLX's reference reran here and is byte-identical to #774's on all four): upstream's four public teacher prompts, 1,024 positions, top two at every position from
tf-flashnext-teacher(a test program, not in the PR), which prefills each prompt through the prompt chunks and feeds every history token in a one-row window, so each position runs on the row kernels here.Two checkpoints move closer to oMLX, oQ2 is level and oQ3 has five more differences: a different fp32 sum order moves which near ties fall either way, In the 61 differences this branch's pick is oMLX's second token 54 times, its third 6 times and its fourth once (oQ4e's code position 142, as in #774).
Every difference (oMLX's pick and margin; this branch's pick and margin; its rank in oMLX's top 20)
Host suite and lean check
zig build test test-golden -Dcpu=apple_m1at each of the three commits, from a clean detached worktree, plusnativeand thetf-flashnext-*tools: 158 of 160 steps; goldens 309/319 equal, the 10 declared differences, 0 unexpected. The one failing step is main's owncore.checkpoint_metalsplit-checkpoint test (it logs an error on the path it expects), as in Core support for oQ in Zig engine #755 and Qwen3.8-Flash-Next support for oQ in Zig engine #774.tf-flashnext-widthsat the tip: PASS, 4,486 checks, every product within 0.9947 of its float64 bound, every row's bits equal in every window (Qwen3.8-Flash-Next support for oQ in Zig engine #774's classes plusrowsandrows-hc).drafts_cost.zig's unit tests: unmeasured depths first, then the best rate of tokens a millisecond; landing rates move toward kept drafts; costs interpolate and extend past the widest by rows.python3 tools/zig/lean_check.py: 531 problems at every commit, Qwen3.8-Flash-Next support for oQ in Zig engine #774's count, and no finding Qwen3.8-Flash-Next support for oQ in Zig engine #774 does not have. New files are at or under 600 lines;engine.zigandreplay.zigwere over 600 before this branch (870 and 3303 lines now, from 862 and 3294 at Qwen3.8-Flash-Next support for oQ in Zig engine #774's tip);zig fmtclean.Platforms
rows_load.onis false there, so no row kernel is compiled or loaded and the depth rule is the existing one.then, greedy: a request and the same with
"draft": false(same tokens), and two requests at once against each alone.Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo (served on the tip, every checkpoint above).Receipts: https://gist.github.com/grantoverton/09035b07768c87e8dd6112e1452e19d8