Repository navigation
GLM-5.3-Flash faster oQ kernels in Zig engine - #786
Open
grantoverton wants to merge 18 commits into
Open
grantoverton wants to merge 18 commits into
grantoverton wants to merge 18 commits into
Conversation
7 tasks
|
Posted by Claude (Anthropic's AI assistant) on behalf of @sethforprivacy, as part of our local-inference tuning work on a Mac Studio M5 Ultra. @grantoverton the synthetic M5 checks for this PR at The live |
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.
…t; a width-carrying matrix replaces Q4
Every GLM projection was typed 4-bit/group-64 (weights.zig's Q4), so an oQ checkpoint could not even be planned.
- Widths (quant.zig): each module's MLX affine width from config.json through core/quant (keys.canonical, then
GLM's own spellings: the MTP block is layer 45, forget_gate folds into the KDA gate). Aliases that disagree,
a missing table and other modes are refused; a tensor stored at another width than declared is refused by name,
never read at a guessed width.
- Affine{bits, group} in place of Q4 (affine.zig): stored shapes, byte spans and column cuts follow each tensor's
width; a fused matrix whose members differ in format plans one run a member at its output rows, as 0.6.6's QSplit
calls each part (glm5_next/linear.py); a one-format stack stays one matrix. The load plan's layer and expert bytes
are counted at those widths.
- mlx-lm's layout (names.zig), which the public oQ checkpoints use: the MTP block, hyper-connections, stacked
experts, the absorbed embed_q/unembed_out pair and a bf16 eh_proj; a fused conv1d.weight is cut into q/k/v; fp32
hyper-connection mixes and router narrow to bf16 only when exact (else refused naming fix, shard, offset, shape).
- The dry-run plan prints its matrices by width. A real load still refuses any matrix the generated 4/64 kernels
cannot read (the next commit gives the device every width); TP2 and expert parallel refuse any width but 4/64.
The uniform 4/64 checkpoint plans the same matrices, offsets and bytes as the 4-bit-only loader (a digest pinned
before the change).
Tests that fail before (no width reading, no mlx-lm names on main): weights_test.zig, oQ-shaped fake checkpoints in
both layouts, each tensor found at its own width; overrides under either prefix; a tensor stored at another width,
two spellings or the MTP block and layer 45 at two widths, an odd member and a routed gate/up pair at two widths
refused by name; oQ4's KDA input and MLA x_proj planned as one run a member; mlx-lm's conv1d and fp32 mixes planned
as the original layout's tensors. quant.zig, affine.zig, names.zig and weights_fix.zig have their own unit tests.
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). lean_check: 533 problems (main 539), no new findings.
…t stacks a product a member GLM has no tensor-unit-only kernel (its *_nax kernels take the simdgroup-matrix fragments without tensor units), so the pre-M5 path only lacked products at other widths. glm_rows.metal (wide.zig) runs them over core row_projection's fp32 group sums, 2/3/4/5/6/8 bits in groups of 32/64/128, every row its own sums (one row alone == the same row in any window): - dense with input and output row strides and per-head batching: each split-stack member at its own output rows (oQ4's KDA input q, v at 5 bits and k, f_a, g_a, b at 8 is six products), mlx-lm's embed_q and unembed_out a head; - by expert id (decode's routed experts, gate and up fused with SwiGLU, glm_glue's sigmoid_fast and bfloat overloads compiled in, not copied), and by sorted pairs (prompt's routed experts); then the existing combine; - bf16 (the MTP head's unquantized eh_proj); the embedding at 6 and 8 bits. KDA f_b and g_b at 4/64 keep the generated kernel's fused dots; at any other width they are row products first and one instance of the generated kda_rows reads them in place of its dots (0.6.6's unfused shape); the products borrow the dense MLP scratch, so state refuses a config whose KDA width outgrows it. A shared expert not at 4/64 runs its own products (no shared-in-routed launch); routed experts not at 4/64 run by id and by sorted pair. Prompt chunks: a one-width 8/64 matrix runs on core affine_mm, split stacks and other widths on the row kernels. The new pipelines load only when a checkpoint has a matrix the 4/64 kernels cannot read, and their source hash joins the kernels' identity; a failed compile releases every pipeline built before it. The uniform 4/64 checkpoint compiles and launches what it did. The previous commit's refusals stay: a member whose rows are not a multiple of 4 (the row kernels read four rows), routed gate and up at two widths, any width but 4/64 with TP2 or expert parallel. Tests that fail before (no wide.zig on the previous commit): 4/64 of one width alone takes the generated kernels; f_b and g_b at any other width (8/64 too) take the products made apart; a failed compile releases every pipeline built before it and a load counts each one (GPU); glm_rows's SwiGLU compiles glm_glue's helpers, no copy of its own. The previous commit's load-time refusal of every other width is gone. 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). lean_check: 533 problems (main 539), no new findings.
…float64, each row bits in every launch
No model files: seeded synthetic codes, scales, biases and inputs, packed as MLX affine.
- Every bits {2,3,4,5,6,8} x group {32,64,128} on three shapes, and every GLM projection class at the public oQ
checkpoints' widths at real shapes (KDA in, f_b/g_b, KDA o, MLA x/qr/o, embed_q/unembed_out, dense, shared,
routed, lm_head), decode rows and the prompt's product separately, against float64 on every column of two rows.
The bound is the fp32 sums' depth (k / 8 + 64 additions) plus a bf16 rounding; each class prints its worst column
(max |err| and its share of max |y|). A reference with scales 1% off must exceed the bound in every class and width.
- Bit-equal pairs: one row alone == the same row in 2-, 3- and 16-row windows; strided == contiguous; per head ==
each head dense; by id == dense; GLU by id == glm_swiglu(dense gate | up); sorted == by id (experts of 8, 5, 2 and 1
pairs); the row kernels == core row_projection; a split stack in one call == each member alone (oQ4's KDA input
5/8/5/8/8/8 and x_proj 6/6/5/5, at the decode pitch and the 64-padded prompt pitch); experts [2048, 4096] and
[4096, 2048] at 2 and 4 bits by pick and sorted == dense.
- KDA: the separate f_b/g_b products then kda_rows == the generated 4/64 kda_rows byte for byte (output, state, conv
window) on every 4-bit code, random per-group power-of-two scales and nonzero biases.
- The 6- and 8-bit embedding against float64 within two bf16 roundings; a neighbouring group's bias is caught.
The engine module exports core/quant (`tensorfold.quant`) for the check's MLX reader.
tf-glm-rows-check (M3 Ultra, no tensor units): PASS 1,328,072 values within the fp32-sum bound of float64 (worst
0.9839 of it, bf16 output rounding; worst-column |err| / max|y| 1.7e-3 to 3.6e-3 a class), 1,780,096 bit-equal
pairs; scales 1% off exceed the bound in every class and width.
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). lean_check: 533 problems (main 539), no new findings.
…e_mm oQ2's routed experts (2-bit, group 64) ran a prompt chunk on the sorted row kernels: 81% of a 2048-row chunk's GPU time. With core affine_mm reading 2-bit codes, wide.routed(gate, up, down) picks affine_mm's 4/64 pipelines, its 2/64 ones (wide.Pipes.mm2) or the sorted row kernels, and 2/64 experts take gatherGlu/gatherTo as 4/64 ones do; other widths stay on the row kernels. Decode is unchanged. Test that fails before: "prompt chunks gather routed experts on affine_mm at 4/64 and 2/64 only; other widths take the sorted row products" (no wide.routed). tf-glm-rows-check adds gatherCheck: 2/64 experts gathered on affine_mm == the dense product (GLU == glm_swiglu(dense gate | up), down == dense) at expert shapes, and the routed q2 prompt classes now run on affine_mm against float64. tf-glm-rows-check (M3 Ultra, no tensor units): PASS 1,328,072 values within the fp32-sum bound of float64 (worst 0.9839 of it, bf16 output rounding; worst-column |err| / max|y| 1.7e-3 to 3.6e-3 a class), 1,976,704 bit-equal pairs; scales 1% off exceed the bound in every class and width. TensorFold/GLM-5.3-Flash-MLX-oQ2-MTP @dc55627 on an M3 Ultra (plan 109.01 GB: 2-bit/64 x126, 4-bit/64 x3, 8-bit/64 x290; loads in 21-24 s): drafted == `"draft": false` 4/4 at 160 tokens (MTP acceptance 74-88%; plain 39.5-40.1 tok/s, drafted 50.9-54.4); resumed == fresh (a 2427-token prompt at 0 and 2425 cached tokens; the second turn resumed at 2425 == a server with `--prompt-cache-gib 0`); concurrent == solo, 2 streams. Cache off: prefill 2427 tokens 5.09-5.11 s (475-477 tok/s; 163 tok/s with the experts on the sorted row kernels), decode 256 tokens 40.8-41.0 tok/s. 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). lean_check: 533 problems (main 539), no new findings.
…shapes, GPU time and GB/s by class Times every decode product class on synthetic weights at the shapes and widths the uniform 4/64, oQ2 and oQ4 checkpoints read: the generated 4/64 kernels, glm_rows at 8, 6 and 5 bits, the KDA step (generated dots or .pre), the MoE launches (uniform's fused gather, oQ4's routed 4/64, oQ2's 2-bit picks, the 8-bit shared expert). Each class rotates over at least 768 MB of weights; a token's time is summed by class for each checkpoint. GLM_SPEED_CANDIDATES=1 also times glm_rows launch shapes that keep its sums (each checked bit-equal first); GLM_SPEED_PROMPT=1 a 2048-row chunk's per-head and 5/6-bit products.
…ffine_mm, every head in one batch A prompt chunk took the decode's per-head row launch over every row for the absorbed MLA pair (each row a weight read, 64 heads x 2048 rows). At a width affine_mm reads (8/64 in oQ4 and oQ2), wide.headsBatch now runs the heads as one affine_mm batch over the strided rows (x's slices a head apart, sums one row apart), as the 4-bit path does with kv_b's value half. Other widths keep the row launch. Uniform 4/64 checkpoints never reach it. Test: tf-glm-rows-check's headsCheck: embed_q [4 x 512, 256] and unembed_out [4 x 256, 512] at 8 bits, 33 rows with a strided x, the batch == each head's own affine_mm product bit for bit (classes hold that product against fp64).
A thread's half step (32 codes, 4 x bits bytes) loads as words and each code is cut from the little-endian bit stream and enters the tile as itself (offset 0, as at 2 and 8 bits; codes up to 63 are exact in bf16). source() admits every width 2-8 but 7. The 2-, 4- and 8-bit sources preprocess byte-identically to before (metal -E, both fragment layouts' defines, both output types), so their compiled kernels are unchanged. Tests: core/affine_mm.zig's dense-and-gathered test adds 3/64, 5/64, 6/64, 5/32 and 6/32 (codes read from the bit stream on the host; fails before: source() refused them); tf-affine-ref adds q3, q5, q6 at g32/g64 on [192, 1024] and [256, 4096] (own seeds, the earlier formats' receipts unchanged).
… members on affine_mm oQ4's 5- and 6-bit matrices (KDA in, f_b/g_b, o_proj, MLA, the MTP layer) and its split stacks (layer 6's KDA input, layer 7's and the MTP's x_proj and qr_proj) took the row kernels in prompt chunks, each row a weight read. wide.zig now compiles affine_mm at 5/64 and 6/64 (mm5, mm6) and wide.denseMm runs each member of a matrix on affine_mm at its width where its rows fill 64-row tiles (a single run: its padded rows, as at 8 bits), into its rows of the chunk's output, the x group sums made once; members of other widths or rows (the 32-row indexer weights_proj) stay on the row kernels. Uniform 4/64 matrices keep their launches. Tests: tf-glm-rows-check's classes now take prompt.zig's choice (wide.mm), so every 5- and 6-bit class's prompt product is held against fp64 on affine_mm; splitMmCheck runs layer 6's KDA input (5/8/5/8/8/8, every member on affine_mm) and layer 7's x_proj (6/6/5 on affine_mm, the 32-row 5-bit member on rows) through wide.denseMm at the 64-padded pitch, each member == its own product bit for bit.
…e width, the same bits as made apart f_b and g_b at a width other than 4/64 (8/64 in oQ2 and 26 of oQ4's KDA layers, 5/64 in 8) were two glm_rows launches a KDA layer, written to scratch and read back by kda_rows (`.pre`). When both share one width in groups of 64, a kda_rows instance now makes them in its own threadgroup: each simdgroup sums four of the head's f_b or g_b rows with core tf_rp_rows (glm_rows's sums, lanes and simd_sum) and rounds them to bf16 as glm_rows's output does, so f and g are .pre's values bit for bit and so is everything after. Two launches and a round trip a KDA layer fewer (68 a decode token). Other widths keep .pre; uniform 4/64 keeps the generated dots. Test: tf-glm-rows-check's fusedCheck: at 2, 3, 4, 5, 6 and 8 bits in groups of 64, every code, random scales, biases, inputs, state and conv window, 3 rows: outputs, state and conv window byte-equal to .pre's (fails before: no fused pipeline). wide.zig's unit test: one width g64 -> .fused; 4|8, 8|4 and g32 -> .pre.
…one launch at its width, the same bits The shared expert at a width other than 4/64 (8/64 in oQ4 and oQ2) took a glm_rows launch for [gate; up] and a glm_swiglu launch every MoE layer. Gate and up of one width are one run of the stack, so wide.glu now runs them through the routed experts' by-ids GLU kernel with every id 0 (expert 0 of one, up N rows on): its sums are glm_rows's and its SwiGLU glm_swiglu's, so acts are the same bits. One launch a MoE layer fewer (42 a decode token). A split [gate; up] keeps the two launches; uniform 4/64 keeps the shared slot of the routed gather. Test: tf-glm-rows-check's sharedGluCheck: [2 x 2048, 4096] at 8, 5, 4 and 2 bits, 16 rows: wide.glu == glm_swiglu(dense [gate; up]) bit for bit (fails before: no wide.glu).
… expert in two launches, the same bits oQ2's MoE layers (routed 2/64, shared 8/64) took four launches: the routed GLU and down by pick, the shared expert's GLU and dense down. glm_moe_rows<BITS, 8, GLU> now runs both in one grid (slots, then the rows through the shared expert), each output with the same tf_rp_rows sums and SwiGLU as the kernel it replaces, so act, acts, ye and ys are the same bits. Two launches a MoE layer fewer (84 a decode token) and the shared expert's small products share the routed grid. Taken when every expert matrix is one run in groups of 64, the routed of one width, the shared at 8 bits; uniform 4/64 keeps the generated gathers and oQ4 (routed 4/64) its generated routed kernels. Test: tf-glm-rows-check's glm_rows_moe.zig: at q2..q8 routed with an 8-bit shared expert, 2 rows of 8 picks over 4 experts (one id past them), wide.moe == byIds GLU and down, wide.glu and dense, bit for bit (fails before: no wide.moe). glm_rows_check's helpers are pub for it (the file stays under 600 lines).
…bbles: the same products tf_rp_qdot16 read each 2-bit code with a shift and a mask; it now masks the code in place (code * 4^i) and tf_rp_load16 scales the code's input by 4^-i, a power of two, so every product is the same float and every sum the same order. 4-bit codes already work this way. Every 2-bit row product takes it, GLM-5.3-Flash's routed experts by pick and in the fused MoE launch among them. tf-glm-rows-speed, oQ2's MoE layer on M3 Ultra (routed q2 + shared q8, both launches): 158.7 -> 132.4 us (1.20x), every output equal to the shift-and-mask kernel's.
grantoverton
force-pushed
the
oqxe/speed-glm
branch
from
October 13, 2026 20:21
347f88a to
67ab196
Compare
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 #773's rebased tip on 2026-10-13: no conflicts; the tip passes the gate again. The receipts below are the pre-rebase runs.
Part of #716; stacks on #773 (which stacks on #755). This branch carries their ten commits until they merge, so please review the last eight commits here (from "tf-glm-rows-speed" on). In #773 the oQ checkpoints' products at widths other than 4/64 run on the generic row kernels: prompt chunks run the absorbed MLA pair a head at a time and every 5- and 6-bit matrix row by row, decode spends extra launches on KDA's f_b/g_b, the shared expert and the MoE, and every 2-bit code is cut out with a shift and a mask. This moves those prompt products onto core
affine_mm, folds the decode launches together with the same sums, and has core row projections read 2-bit codes in place as they already read 4-bit nibbles. On an M3 Ultra, oQ2 decodes 1.22× (code) and 1.13× (chat) faster with the same decode bits, and up to 1.34× under concurrent load; oQ4 prefills 1.33-1.37× faster and oQ2 1.06×. The uniform 4/64 path keeps its kernels, its launches and its tokens.What this changes
tf-glm-rows-speed(a test program): GPU time and GB/s of every decode product class at the shapes and widths the uniform 4/64, oQ2 and oQ4 checkpoints read, each class over at least 768 MB of weights, summed into a token's products per checkpoint;GLM_SPEED_PROMPT=1adds a 2048-row prompt chunk's per-head and 5/6-bit products.embed_qandunembed_outonaffine_mm, every head in one batch. At a widthaffine_mmreads (8/64 in oQ4 and oQ2),wide.headsBatchruns the 64 heads as one batch over the strided rows, as the 4-bit path does withkv_b's value half. Other widths keep the row launch; uniform 4/64 never reaches it.affine_mmreads 3-, 5- and 6-bit MLX codes exactly. A thread's half step loads as words and each code enters the tile as itself (codes up to 63 are exact in bf16). The 2-, 4- and 8-bit sources preprocess byte-identically to before, so their compiled kernels are unchanged.affine_mm. oQ4's 5/64 and 6/64 matrices (KDA in, f_b/g_b, o_proj, MLA, the MTP layer) and each member of its split stacks (layer 6's KDA input, layer 7's and the MTP'sx_projandqr_proj) where its rows fill 64-row tiles; the 32-row indexerweights_projstays on the row kernels.kda_rowsat one width, the same bits as made apart. Each simdgroup sums four of the head's rows with coretf_rp_rowsand rounds them to bf16 asglm_rowsdoes. Two launches and a scratch round trip a KDA layer fewer (68 a token).glm_moe_rows, taken when every expert matrix is one run in groups of 64, the routed of one width, the shared at 8 bits). Two launches a MoE layer fewer (84 a token). Uniform 4/64 keeps the generated gathers and oQ4 (routed 4/64) its generated routed kernels.tf_rp_qdot16masks each 2-bit code where it sits (code × 4^i) andtf_rp_load16scales that code's input by 4^-i, a power of two, so every product is the same float and every sum keeps its order. Every 2-bit row product takes it, GLM-5.3-Flash's routed experts among them.M5 (tensor units): GLM has no tensor-unit-only kernel;
affine_mmcompiles its tensor-unit variant there and the new 3/5/6-bit sources take the same code path. Not run on M5 (no hardware).Exactness
TensorFold/GLM-5.3-Flash-MLX-oQ2-MTP @dc55627, TensorFold/GLM-5.3-Flash-MLX-oQ4-MTP @9872546 and the uniform 4/64 checkpoint #773 used, greedy, MTP drafts on unless
"draft": false;serve <checkpoint> --name glm --parallel 2 --context 16384 --prompt-cache-gib 0. Before is #773's tip 9777aeb; after is 7ea90233, this branch's tip 347f88a but for one test's expected count (commit 7's compile-failure test counts its new pipelines; no engine source differs).b7932757843f, oQ47cd4ab7084c3, uniform370c5ca7e2df).affine_mm's tiles in place of the row kernels), within the same float64 bound, so the chat prompt (32 tokens, a prompt chunk) continues differently: oQ28e524c1de404→19eac26b10c1, oQ4810fec766b0f→09e32b8f5a89. The uniform 4/64 chat hash is unchanged (e2cf3b4b33c0). Fidelity against oMLX is below.bench_concurrent's alone and serial checks at 1 and 4 requests, both prompts: 34/34 on each); cancellation, reset and partial acceptance, 25 steps, pass on each (a fresh"draft": falseprobe gives the reference hash and the drafted probe equals it, with partial acceptance seen; then after a stream dropped mid-decode, a long prompt dropped mid-prompt, a stream dropped beside a running probe and a request past the window, the drafted and plain probes equal the fresh hashes; oQ289faa950b278/ad8ca98b7043, oQ47152f0a8505e/c9b4aa96047f).tf-glm-rows-checkat the tip: PASS, worst 0.9839 of the float64 bound (the routed 2-bit decode classes' worst 0.8564, as before commit 8). New checks, each failing before its commit: the per-head batch equals each head's ownaffine_mmproduct; split stacks onaffine_mmequal each member alone at the 64-padded pitch; 5- and 6-bit prompt classes against float64 onaffine_mm; the fused f_b/g_b inkda_rowsequal.pre's outputs, state and conv window byte for byte at 2/3/4/5/6/8 bits;wide.gluequalsglm_swiglu(dense gate | up);wide.moeequals the by-id GLU and down pluswide.gluand dense.tf-affine-refadds q3, q5 and q6 at g32 and g64: PASS (bf16 out at most 0.4868 of the bound, every row's bits equal in windows of 1 to 65 rows), the earlier formats unchanged.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-nativesha256d508376cf1fd…(before) anddac50042463d…(after). Every run starts from a fresh server, with other checkpoints' file cache evicted first. Rows are in run order.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_openaicodeoQ2's code prompt gives the same tokens in the same 65 rounds before and after (chat: 113 rounds against 116, on its new tokens), so its gain is the rounds' time: 2-bit routed experts are oQ2's widest bytes in every verify window, and commit 8 reads them in place. oQ2 now decodes code at 1.20× and chat at 1.01× of the uniform checkpoint on the same build. oQ4's decode products are unchanged here (its routed experts are 4/64, on the generated kernels) and so is its code speed; its chat prompt now continues on other tokens (above) that take 144 rounds against 136, at the same 33.7-34.0 ms a round, hence 0.95× on chat. The uniform runs' kernels, launches and tokens are unchanged (
bench_openaidecode only).Cold prefill (oQ2 and oQ4; 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 uniform 4/64 prefills at 512.0 / 518.3 / 519.8 tok/s at these sizes in #773 (
--context 66560, three reps), and its prompt kernels are unchanged here: oQ2 and oQ4 now prefill at 0.96-0.97× of it.Per projection class (
tf-glm-rows-speed, synthetic weights at the checkpoints' shapes and widths, each class over at least 768 MB; before is commit 1 alone on #773's tip, after this branch; two alternated runs each, ranges):.pre)affine_mm)embed_q8/64, 64 heads (×11 a chunk)unembed_out8/64, 64 heads (×11)Byte roofline: a decode token reads 9.65 GB on the uniform checkpoint, 11.63 GB on oQ2 and 13.75 GB on oQ4 (oQ's non-expert matrices are mostly 8-bit), so at the ~705 GB/s these products now reach, oQ decode cannot pass the uniform checkpoint on one row; verify windows read each 2-bit expert they pick, which is where oQ2 gains.
Memory
Physical footprint (
darwin_rusage_info_v4, the server process,--parallel 2 --context 16384), the same before and after: oQ2 113.70 GB after load, peak 113.70-113.76 GB; oQ4 189.38 GB, peak 189.4 GB; the new pipelines add nothing measurable.Fidelity
Teacher-forced against oMLX 0.7.1.dev1 on the same checkpoints and the same fixed histories as #773 (oMLX's reference reran here and is byte-identical to #773's): upstream's four public teacher prompts, 1,024 positions, top two at every position from
tf-glm-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. The uniform 4/64 is unchanged (#773: 28 differences, 26 beyond the tie policy).The decode bits are #773's, so the change comes from the prompt chunks' new sum order (commits 2-4), which moves the prompt's state within the float64 bound and with it which near ties fall either way. Both checkpoints stay at or near the uniform checkpoint's 28, and in every difference this branch's pick is oMLX's second or third token. The widest is oQ4's copy position 181, where oMLX's margin is 2.625 and this branch picks oMLX's second token.
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 eight commits, from a clean detached worktree: 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 GLM-5.3-Flash support for oQ in Zig engine #773.tf-glm-rows-checkandtf-affine-refpass at the tip (above), with each commit's new check failing before it.python3 tools/zig/lean_check.py: 533 problems at every commit, GLM-5.3-Flash support for oQ in Zig engine #773's count (main 539), no new findings. Every touched file is at or under 600 lines;zig fmtclean.Platforms
tf-glm-rows-checkandtf-affine-refPASS at 347f88a.affine_mmcompiles its tensor-unit variant for the new widths. Commands:then, greedy: a request and the same with
"draft": false(same tokens), and two requests at once against each alone.Reproduction settings
As #773:
GLM_LOAD_LIMIT_GB=208(set as the machine's owner; oQ4's plan is over the default 70% of RAM),iogpu.wired_limit_mb240000, one model process at a time behind one GPU lock, other checkpoints' file cache evicted before each load.Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo (oQ2 and oQ4 served on the tip; uniform tokens equal GLM-5.3-Flash support for oQ in Zig engine #773's).affine_mm(above), within the same float64 bound.Receipts: https://gist.github.com/grantoverton/7f9b0a1b3ba7919e11ad8f49f0f0acf1