Repository navigation
Gemma 4 31B served natively on Metal: the dense family, DFlash drafts, tensor-unit prompts past mlx-lm, and core affine_mm without its epilogue - #879
Open
BobClawblaw wants to merge 43 commits into
Conversation
…by's gemma4-dense Python engine config.zig pins the 31B 4-bit geometry and names each layout it refuses (MoE, the 12B unified, per-layer inputs, shared KV, keys apart from values, per-module widths); tools/zig/gemma4_gate.py records greedy ids and top-two margins from gilby's RowDecode (lane kernels) through either prompt path, and compares a native run with them.
…ane projection bit-equal on every linear tools/zig/gemma4_dump.py writes each decode kernel's inputs and outputs for a sliding and a global layer at 1, 2 and 16 rows; tf-gemma4-kernels runs core lane_projection with lane_qmm's K split on the 31B's stacked q/k/v, o, gate/up, down and tied head: every output's bits equal the Python engine's.
…de kernels and the MLX ops around them Ports gilby's qkv_prep, attend (partials and merge), attn_tail and dense_tail (ashhart#360) with their rounding points, plus MLX's looped RMS norm, the scaled 4-bit embedding, mlx_lm's geglu and logit soft-cap as bf16 ops, and the global layers' proportional RoPE table from MLX's precise pow. tf-gemma4-glue checks every output against the Python engine's at 1, 2 and 16 rows, sliding and global layers: all bits equal.
…it for logit with gilby's RowDecode weights.zig packs each layer's q/k/v, o, gate/up and down stacks and the tied head once for the core lane projection (seven shapes); state.zig keeps a ring per sliding layer and a growing cache per global one; forward.zig runs a window through the 60 layers in one command buffer. tf-gemma4-run on four chat prompts, 128 tokens each: every token and every top-two logit equals the Python engine's (gemma4_gate.py check).
…every width, and 8k/32k receipts tf-gemma4-run takes a prompt window width: windows of 16 give the one-row path's tokens and margins on every chat prompt. Against the oracle's mlx_lm prompt forward, the first token is equal at 8k and 32k and 49 tokens are equal at 32k; at 8k the two diverge at token 32 on a 1-2 bf16-step tie, where the Python engine's own decode-path prompt diverges from its mlx path too, and the native run equals that decode path's tokens and last-row logits bit for bit. gemma4_gate.py gains a prompts command (8k/32k chat prompts from the repo's text) and saves last-row logits.
…stream equal to its plain reply forward.run takes several streams' windows in one forward (q/k/v prep, attention and the cache write per stream at offsets, every other kernel over all rows); backend.zig puts it behind the lanes Backend, windows up to 16 rows, keep rolling a cache back by its length. tf-gemma4-lanes: plain and copy-drafted single streams and 4- and 8-stream shared rounds equal the oracle's 128 tokens; on copy-heavy prompts (11.1 and 3.3 tokens a round, drafts landing and rolled back) one stream and 8 mixed streams equal the engine's own plain replies.
…eedy replies equal to the Python reference native/gemma4.zig serves the dense 31B on the lane host (context 32,768 by default); metal.zig registers gemma4 in mlx-q4g64 and refuses --speed-up and --learn for it by name; the tests that took gemma4 as an unserved type use gemma4_unified, which still is. gemma4_gate.py served: on an M5 Ultra the four prompts' greedy token SHAs equal the oracle's alone, with drafts and four at once; thinking stays in reasoning_content and a tool call parses. docs/recipes/gemma-4-native.md is the recipe; docs/credits.md credits gilby's ashhart#360 and cshintov's #10.
…s at 821 GB/s The core lane projection's reg variant (4-bit groups of 64 in tensor-op registers, group sums precomputed) gives lane_qmm's bits on every linear and width (tf-gemma4-kernels --reg) and reads the 31B's 17.3 GB a step at 821 GB/s, against 662 for the plain lane kernel. One-row decode goes from 35 to 42 tokens a second (tf-gemma4-bench); the four prompts' tokens and top-two logits still equal the oracle's.
…ion, the same bits A window's keys and values go into the cache before attention, and attention runs over it in 16-row calls reading every key from the cache (the kernel takes its first new key from the meta), so a prompt window can take the reg projection's 128 rows in one weight read. Tokens, top-two logits and last-row logits equal the 16-row and one-row runs; 8k prompts go from 340 to 576 tokens a second and 32k from 237 to 330.
…fine_mm prompt GEMMs beside the lane kernel
…ned copies of the MLX tensors prompt.zig runs a chunk's projections on affine_mm (the decode glue for the rest, the head on its last row only); affine_mm reads words as u32, and checkpoint tensors can sit at any byte offset in a shard, so weights.zig copies each linear's tensors to aligned storage. The sliding ring holds the window plus a chunk. Against the oracle's mlx-lm prompt receipts: first tokens equal at 8k and 32k, 49 tokens at 32k, the short prompts as the Python engine's own decode path; prompt attention is still the per-row decode kernel (637 tokens a second at 8k).
…128-wide slices across simdgroups gemma_prompt_attention takes 16 query rows of one head a threadgroup: each simdgroup scores keys over its 128 columns of the head (two on sliding layers, four on global ones), the slices' partial scores add in slice order through threadgroup memory, every simdgroup runs the same online softmax under the causal and sliding-window masks, and each keeps its slice of the output. Prompts go from 637 to 1,184 tokens a second at 8k and from 351 to 1,008 at 32k (mlx-lm: 1,217 and 1,233); first tokens equal the oracle's mlx-lm receipts at 8k and 32k, 49 tokens at 32k.
…it prompt path backend.zig's prefill runs prompt.zig's chunks; the lanes check and the served gate now take the engine's own replies as the reference (its prompt path follows mlx-lm's arithmetic, not the oracle's decode path). Copy-drafted, 4- and 8-stream rounds equal the plain replies; served on an M5 Ultra, greedy SHAs equal tf-gemma4-run's alone, drafted and four at once, with thinking and tool calls parsed.
… reg kernel's own sums dispatch
… key head's query heads as one tile gemma_decode_attention packs a key head's eight query heads (two decode rows) into one 16-row tile, so each key and value is read once a token, splits the keys into 256-key chunks across threadgroups and leaves partials for gemma_attention_merge. Decode goes from 36 to 38 tokens a second at 8k and from 30 to 35 at 32k (mlx-lm: 42, 38); 16-row windows at 32k run 2.5 times faster. Every row keeps its bits at any window width, and drafted, 4- and 8-stream replies equal the plain ones; the oracle's decode now differs only at near-ties (one bf16 step).
…xt prompts up to 252k in the gate The global layers' decode chunks come from prepare (256 keys a chunk), so the partials' scratch is sized by them and by the window on sliding layers: about 1 GiB at the 262,144-token context, not 4. gemma4_gate.py prompts builds 8k, 32k, 64k, 128k and 252k chat prompts from the repo's docs and source, no passage repeated.
… cache, the same bits The norm, merge, attn_tail, geglu and dense_tail kernels write the next projection's group sums in tf_lane_sums' order, so the reg lane kernel runs on them without its own sums dispatch, and qkv_prep writes keys and values straight into the cache. GeGLU runs 64 threads a group (4.9 us a call at one row, against 23.9 for a thread a group). One-row decode goes from 42 to 43 tokens a second; tokens, margins and long-prompt logits are unchanged.
…s as it makes them Backend.prefill_features runs the prompt pass with a FeatureSink: the Drafted wrapper opens the drafter first and absorbs each chunk's rows but the prompt's last (that one goes with the first token, as before), so a long prompt's taps are never held whole. Targets without it keep the one absorb after the pass.
… one stream, the same tokens families/gemma4/dflash.zig runs z-lab/gemma-4-31B-it-DFlash (DFlash v1: five Qwen3-style layers on six tapped target layers, blocks of the pending token and up to 15 mask tokens, the target's embedding and tied head) behind the lane core's Drafted wrapper. dflash_weights.zig admits that one geometry and quantizes the bf16 linears at load to 4-bit groups of 64 by MLX's affine rule (a plain min/max rule lost half the acceptance on the taps' outlier channels); the linears run on the lane projection, the norms on the target's fused norm kernel with the projections' group sums, the SwiGLU on a sums-writing copy of core/draft_ops'. Each stream keeps context keys and values in a 2,047-slot ring for the sliding layers and a growing buffer for the full one; an absorb's command buffer is committed and settled by the hold that follows. The target writes its tapped layers' rows from dense_tail, hands prompt chunks' rows through prefill_features and verify rows through features, reports its window forward's cost by rows (timed at load) so the depth rule prices drafts, takes 128 rows of streams a forward, and draws greedy tokens with the core's GPU argmax. tensorfold-native serve takes --drafter for gemma4; tf-gemma4-lanes runs the drafted cases with GEMMA_DRAFTER (and GEMMA_DRAFT_TRACE prints drafts and timings). On the M5 Ultra one stream decodes 6.4, 6.4, 2.1 and 6.7 tokens a round on the four prompts (150-180 tok/s; plain 43) and 10.2 and 7.8 on copying replies (190 tok/s), equal to plain; eight streams 180-195 tok/s together. The served gate passes against the engine's receipts (the decode attention's 256-key order ties one near-tie the oracle splits). Credit to z-lab for the drafter and its MLX reference.
…b pass; wide forwards' costs timed at load holdFn concatenates the streams' blocks (up to 128 rows) and runs the drafter once, attention a block at its rows: eight streams' blocks take 9 ms instead of 27. absorbFn packs the streams' rows into passes the same way, fc a piece and the rest together. The drafter reports batched head costs. The backend times forwards at 32, 48, 64, 96 and 128 rows as shared_costs, since the rule's extension of the window table priced them at twice the truth; GEMMA_DRAFT_TRACE also prints each verify's rows and time and the test each case's engine steps. A two-blocks-a-threadgroup variant of the reg lane kernel was measured and dropped: wide forwards are bound by tensor-op compute (0.8 ms a row), not weight reads. Eight streams: 203-208 tok/s (from 180-195); drafted replies still equal plain, the served gate passes.
…in the verify loop, a debugging aid)
…ong-context table to 252k One forward of a 252k-token prompt through mlx_lm builds 252k x 252k masks and took the Mac down; the oracle's mlx prompt path now feeds 2048-token chunks as mlx_lm.generate does, and the 252k receipt peaks at a quarter of the RAM. tf-gemma4-run reads ids files up to 4 MiB (the 252k prompt is 1.7 MB of JSON). Native equals mlx-lm for 49 tokens at 32k, 64k, 128k and 252k; at 8k they part at token 25 on an exact tie in mlx-lm. The recipe tabulates prompt and decode speed for both at each size.
…ead; maxes and probabilities cross once a block The old kernel sliced the head across simdgroups and summed their partial scores through threadgroup memory every 32 keys, with two barriers each time: a third of the attention time at 32k (234 ms of a 512-row chunk's 620; the GEMMs are 376 and the glue 13). Now each simdgroup scores its own 64 keys of a block over the whole head, reading the tile's Q from L1, and the tile exchanges only its row maxes and fp32 probabilities once a block; each simdgroup then accumulates its own 128 output columns. The attention stage at 32k is 185-191 ms, what mlx-lm's fused attention costs at that shape. Fully unrolling the 32 fragment steps ran 5x slower (the code outgrew the instruction cache); bf16 probabilities were 3% faster but put an 8k prompt's logits 0.5 apart at an mlx-lm tie. Value reads stop at the tile's last written key: a value row never written multiplies a zero probability into NaN for every row of the tile, which the served gate hit once the block grew past 32 keys. tf-gemma4-bench GEMMA_STAGES[=depth] times a chunk's GEMMs, attention and glue apart (prompt.layer takes a Stages mask). The recipe's long-context table is remeasured: 8k, 32k and 64k equal mlx-lm for 49 tokens; 128k and 252k each part once where mlx-lm's top two logits tie.
…cale and bias fetched a step ahead The same multiplies in the same order, so every caller's bits stay: Gemma 4's last-row logits at 8k to 252k are identical before and after. A 512-row chunk's 60 Gemma 4 layers go from 382 to 344 ms (mlx-lm's quantized matmuls take 302): the tile's 16 x 16 fragments load as two 8-byte vectors a lane instead of eight scalars (6%), and each step's 16 bytes of codes with their scale and bias are fetched during the step before, so the barriers no longer wait on memory (4%). Both keep their scalar form under TF_SIMD_FRAGS, where a lane's columns are not contiguous. Measured and set aside: taller tiles of 128 and 256 rows (every simdgroup waits on the dequantizing ones: 414 and 582 ms), a 64 x 128 tile with 16 x 128 a simdgroup (registers spill: 752), eight simdgroups on a 64 x 128 tile (410), and a double-buffered tile (occupancy: 397). The exact-codes epilogue's scale and bias FMAs bound at about 75 ms.
…oring 32 keys and owning 64 output columns The kernel streams keys and values, so each load serving twice the rows is what moves it: a 512-row chunk's 60 layers at 32k go 194 -> 138 ms, at 128k 810 -> 527 (mlx-lm's fused attention about 161 at 32k). The probabilities cross threadgroup memory 128 keys at a time, fp32 as before. Vector fragment loads, fetching a step ahead, and sparing interior blocks the mask and the rescale all measured as noise on the 16-row kernel; a prefetched set of fragments cost occupancy and ran up to twice slower. Only the row sums regroup. Long prompts: 8k 1142 -> 1187 tok/s, 32k 996 -> 1054, 64k 844 -> 949, 128k 648 -> 806, 252k 390 -> 527 (mlx-lm 1305/1139/944/651/440). The two tie-breaks against mlx-lm's equal top logits move from 128k and 252k to 8k and 128k; decode receipts, lanes with and without the drafter, and the served gate are unchanged.
…nd the row sums go The tile held exact codes (4-bit ones plus 128) and every 64-deep group's partial product was scaled and bias-corrected in a scalar epilogue, the bias times fp32 row sums from a kernel of its own. Now each thread writes its 32 weights as bf16(fma(scale, code, bias)), one rounding, as mlx-lm's quantized matmul dequantizes, and the tensor unit accumulates over the whole depth with nothing after it: a 512-row chunk's 60 Gemma 4 layers take 284 ms on the GEMMs where they took 344 (mlx-lm's 302). tf_affine_row_sums, rowSums/rowSumsAt, the sums strides and every caller's sums buffer are removed (Gemma 4, GLM, the tests and mm_bench). The prompt path's weights now round as mlx-lm's do while the lane kernels keep the codes exact, so prompt receipts move: affine_ref checks affine_mm against the rounded weights (its own reference stays exact for simd_qmm) with the tensor unit's 16-deep chain as the bound, and glm_rows_classes' fp64 check takes the rounding. Gemma 4 long prompts: 8k 1362, 32k 1202, 64k 1047, 128k 892, 252k 564 tok/s (mlx-lm 1305/1139/944/651/440); equal at 32k and 64k, parting at 8k, 128k and 252k where mlx-lm's top two logits tie. The four decode receipts all equal the oracle now; lanes with and without the drafter and the served gate are unchanged.
…jections' prologues on one-row windows A plain token at short context goes from 23.7 to 20.3 ms on the M5 Ultra (42 to 49 tokens a second; mlx-lm 46), 45 at 8k and 40 at 32k. The lane projection's K split is now the one each shape streams fastest at on one row (o and down 4, the gate/up stack and the head 8), another fp32 summation order than the reference's: tf-gemma4-kernels keeps the reference split bit for bit and reports the decode split, where a handful of outputs a linear move by one bf16 step. The core reg lane kernel takes a family prologue and epilogue (Fused): Gemma's dense tail runs as the next layer's qkv prologue, the attention tail as the gate/up projection's with geglu its epilogue (the stack's rows interleaved in 16-row blocks), and the down projection sums its own input; seven dispatches a layer rather than ten, the same bits as apart, worth 0.2 ms. tf-gemma4-bench times the decode stages, each shape plain and fused, and the glue kernels alone. Receipts: decode prompts equal the oracle but the story at a tie it takes the other way; long prompts equal at 32k, 64k and 252k, apart at 8k and 128k on mlx-lm's own ties; lanes and DFlash streams equal; served gate passes with its thinking budget at 800 tokens (the shifted reasoning ran past 400).
…s and, with one key chunk, merges itself On one-row windows gemma_attention_partial takes gemma_qkv_prep's work (FUSE.x): its simdgroups norm and rope the threadgroup's query heads from the stacked row, and the last chunk's threadgroup writes the row's key and value to the cache; with one chunk (FUSE.y) it writes gemma_attention_merge's output and sums itself, exp(top - top) being 1. The prep's lane layout and reductions are the prep kernel's, so the bits are the same, and the dispatches a sliding layer go from seven to five or six: a plain token at short context is 19.7 ms (51 tokens a second; mlx-lm 46), 45 at 8k and 40 at 32k. Global layers keep their prep and merge. Receipts as before: decode, long, lanes, drafter, served.
…ing and first-fetch order The fused hook moved the first weight fetch ahead of the tile setup and indexed tiles by simdgroup for every reg pipeline. Unfused pipelines now compile main's code path again (one tile a threadgroup, the first group's codes, pairs and sums fetched together after the setup), so models without a prologue see no change from it; Gemma 4's fused one-row decode keeps its bits (tf-gemma4-bench) and its fixtures (tf-gemma4-kernels), and its rate is unchanged.
Owner
|
Thank you, @BobClawblaw, for bringing the 60-layer Gemma model onto the lane engine with draft support. This first pass found image input still missing: Please send a follow-up covering image input, those startup owners and the implementation's source history for the next review. Posted by Ash's AI Assistant Synthena |
…, releases what was made before it A drafter that failed to load returned null from open(), a normal return, so the errdefers for the model, the host and the backend never ran; and Drafted.init could fail before the drafter's own errdefer was installed. open() now calls start(), which returns error.Refused once `problem` says why, so every errdefer above it runs; the drafter's errdefer is installed as soon as it exists and the wrapper's after it. The decode and prompt scratches and the drafter's 24 buffers release the ones already made when a later one fails (each struct's release(made), which deinit also uses), and Metal.init frees the decode scratch when the prompt scratch fails. With a bad --drafter the server's checking allocator reported 1661 leaked allocations at exit; it now reports none, as the --context refusal already did.
…d an image may be pieces with tokens around them The server and engine_api half of step5-native's aa221b1, the same text, so the branches merge clean: Info.vision's part_type is the content part type the chat template reads ("image" for most HF templates, Gemma 4's among them), and a PreparedImage may carry pieces, each its tower input and the tokens its placeholder expands to around its rows. Gemma 4's template writes <|image|> alone, and its processor puts the image's rows between boi and eoi.
…e blocks on sliding layers, and prompt reuse serve --vision loads the checkpoint's image tower (vision_tower., embed_vision.: 1.1 GB, packed apart from the language model) and reads PNG and JPEG data URLs. families/gemma4/image.zig prepares an image as transformers' Gemma4ImageProcessorPil does: the largest aspect-preserving size, sides multiples of 48, whose 16-pixel patches fit 280 tokens, PIL's bicubic resize, then np_rescale's float64 1/255. families/gemma4/vision.zig runs Gemma4VisionModel on gemm_nax's tensor-unit GEMMs at torch's per-op bf16 rounding (kernels/metal/gemma4/vision.metal): the patch projection and the x and y position rows, fp32 RMS norms times their weights, per-head q, k and v norms, the axial 2D RoPE (theta 100, x on dims 0-35, y on 36-71), attention at scale 1 as batched GEMMs and a softmax, GELU-tanh gating, the 3 x 3 pooler and standardization, then the 4-bit projection on core affine_mm: about 90 ms an image. In the language model an image's rows replace its tokens' embeddings unscaled, and on the 50 sliding layers they see each other both ways (use_bidirectional_attention "vision"): gemma_prompt_attention takes each row's last key (HI, its own position for text, so text prompts keep their bits) and global layers stay causal. A prompt chunk never cuts an image: a cut that would moves before it, or past it when the chunk starts there. Prompts are reused between requests (snapshot.zig): a kept state is each sliding ring's last 1,024 keys and values, the global layers' so far, and the DFlash drafter's context for the stream (dflash.zig keepLane, resumed by its first absorb), kept and resumed only at the request's planned chunk starts. A drafting stream does not resume a state kept without the drafter's context, and none is resumed inside an image. The budget is cache_fit's beside every stream's caches, 16 GiB by default. Checked: tools/zig/gemma4_vision_fixtures.py (transformers 5.17 and torch) and tf-gemma4-vision on five drawn images, 20 x 15 to 1200 x 800 and 180 x 1400: patches equal the PIL processor's bit for bit; the tower's rows are as close to an fp32 run as torch's own bf16 run (7.6% relative error averaged over the images, torch's 7.8%). Served with the drafter on an M5 Ultra, an image described in 80 greedy tokens equals mlx-vlm 0.7.2's reply character for character from the same 287-token prompt, drafts on and off; tools/zig/vision_served.py passes its 18 checks on a fresh server (drafted equals plain, concurrent equals alone, a later turn resumes past its image, another image resumes only the text before it, the refusals); a 3,336-token text conversation resumed 3,329 tokens and one with an image 3,584 of 3,607, each answering as its first pass did. tf-gemma4-run's decode and 8k-252k receipts are byte-identical to 64f1895's. zig build test, fmt and the lean check pass.
…, MLX's and mlx-lm's ops, transformers' tower, z-lab's DFlash glue.metal ports gilby's gemma4-dense reference (a6ff795: RowDecode, _EMIT, qkv_prep, the tails and attention) and, where it calls MLX or mlx-lm, their kernels' arithmetic: gemma_rms_looped is MLX's rms_looped (rms_norm.metal, the kernel mx.fast.rms_norm runs past 4,096), gemma_embed MLX's dequantize, gemma_geglu and gemma_softcap mlx-lm's gemma4_text ops, gemma_rope_freqs its ProportionalRoPE frequencies (MIT, Apple). The image tower and image.zig follow transformers' modeling_gemma4 and Gemma4ImageProcessorPil (Apache 2.0, the HuggingFace Team); the drafter z-lab's MLX DFlash model (MIT, Z Lab).
…rafter's held drafts to the grammar Every response_format request on Gemma 4 with drafts on failed with HeldDraftsUnreadable: the lane core reads a grammar stream's held drafts through the backend's held hook, which the wrapper did not set. It now reads them from the drafter (fillers: the pending token, as verify writes them). A drafted_test case runs a grammar stream through the wrapper: drafted equals plain, every token inside the grammar; it fails without the hook.
…holds The checkpoint, the packed weights, both scratches, the drafter's weights and buffers and the image tower's join one set; a stream's caches come and go with it. Served on an M5 Ultra, a 1-token request after 2/6/30/120 s idle took 0.48/0.47/0.42/0.40 s before; with the set and the keepalive's kernel tick (ashhart#903, ashhart#909) 0.11/0.11/0.11/0.12 s, where the tick alone gives 0.35-0.61 s and the set alone 0.11-0.66 s.
…d, prompt reuse and image timings, the idle first token Structured output is enforced with drafts on (the harness's check: 18 of 18), so the API doc lists Gemma 4 among the native engines that enforce it; tool_choice required or a named function gets HTTP 400, as Gemma's template writes its own call markup. The recipe adds served first-token times: 7.7 s fresh and 0.22 s resumed for an 11,481-token conversation, 8.1 s and 0.12 s with an image, an image's 270 ms. Known issues: the first request after six or more idle seconds takes 0.4-0.7 s until the keepalive's tick does GPU work.
Contributor
Author
|
@ashhart thanks. The follow-up is pushed: Landing order. Image input needs #598's server path, so #879 merges Image input (
Structured output with drafts ( Idle first token ( Checks (M5 Ultra, on
Startup owners ( Sources (
Other Metal families. #879 changes core
|
Owner
|
Thank you, @BobClawblaw, for adding Gemma's own image tower and repairing the startup owners. At d48368c, the earlier missing-image and refused-drafter cleanup findings close in source. Two small cases remain: Please address those cases and clarify the unchanged RMS helper's implementation history against MLX's We'd run filler-to-grammar-to-verify and malformed-tensor tests, then independent PIL pixels/tower stages, image-boundary reuse, grammar with drafts, seeded sampling and concurrent/solo full-token comparisons. Posted by Ash's AI Assistant Synthena |
…see the pending token The drafts wrapper's held hook cleared a stream's filler count when a grammar read it, so a verify of the same held window would read the head's chain instead. The count now ends only at the verify or the next request. The grammar tests move to drafted_grammar_test.zig with a case that fails without the fix, and the two whole-restore cases share one fixture.
…ales and biases per group of 64 The tower took embed_vision.embedding_projection's scales and biases without checking their dtype, rank, shape or byte count; a malformed checkpoint is now refused by name before a kernel reads them. A unit test covers f32 scales, groups of 32, short bytes, a flat tensor and wrong words.
…m.metal, MIT) and how it is specialized
ashhart#598 has landed with main's own image path, so Gemma 4 takes it: the server's wrap writes boi and eoi around an image's rows (the pieces this branch added are gone), images decode through core's picture decoder, and the residency set is Resident.hold's, whose keepalive tick runs a copy on the GPU. Info.vision keeps part_type, now carried by the messages' Media, since Gemma 4's template writes <|image|> only for "image" parts. LFM2's prompt GEMMs take the core affine_mm without row sums, as every other caller does since its weights are dequantized to bf16 in the tile. Both sides load the tile's fragments in two 8-byte reads; this branch's am_tile_frag stays, since main's frag_get_tg form of the same loads made Gemma 4's prompt GEMMs 11% slower (315 against 283 ms for a 512-row chunk at depth 8k).
…utes: each weight bf16(fma(s, q, b)) in the tile The contract's affine family gains that form (MLX's quantized matmul loader), and affine_mm's specs take it without row sums: its launch records are the 9,216 threadgroup bytes it now uses, and the row-sums kernel's entry goes. Its known finding, NaN for an infinite input from the +128 code offset, passes now and is dropped. tf-tile-check holds the tensor rows to the float64 bound alone: they round each weight as MLX does, so they no longer beat the exact-code tiles (worst 2.46e-3 against 2.21e-3 of the terms' magnitude at 4 bits). A server test covers a template that reads "image" parts: normalized twice, the part keeps that type, and the server's wrap lands around the image's rows.
…set and main's keepalive, HTTPS image URLs, Pillow's pixels, timings measured again
Contributor
Author
|
@ashhart thanks. Both cases are fixed ( Fillers read by a grammar ( The image projection ( The RMS helper (
The kernel's comment now names the file, version and license, as The merge (
Three places in main read core
Main's 8-byte tile loads ( Checks (M5 Ultra, on
Other Metal families, main (46a08c1) against this branch, served:
|
Owner
|
Thank you, @BobClawblaw, the filler/grammar regression and malformed projection tests close both earlier code findings, and the MLX RMS attribution answers its history clearly. We checked 21e8040. The release integration also changes shared affine_mm from exact-code/fp32 scaling to rounded bf16 weight tiles, while e8b6054 removes the stronger tile comparison. LFM2's reported error and task-score changes make this broader than the Gemma adapter. Please isolate Gemma's required arithmetic behind an explicit shared variant, preserving released callers and their comparison, or leave that arithmetic change for Ash's decision. We can help extract the variant at landing. Your image, draft and other-family receipts remain useful, with current GLM images and the remaining shared-math checks kept in scope. Posted by Ash's AI Assistant Synthena |
This was referenced Oct 17, 2026
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.
Gemma 4 31B served natively on Metal
The dense 31B (
mlx-community/gemma-4-31B-it-4bit, groups of 64) as the native familyzig/src/families/gemma4, registered asgemma4, with the recipe indocs/recipes/gemma-4-native.md:zig/kernels/metal/gemma4/glue.metal, ported from gilby'sgemma4-densePython engine (feat(gemma4): the dense 31B layout and oQ per-layer mixed widths on the Mac decode kernels #360) with its rounding points. Each kernel equals the reference bit for bit at 1, 2 and 16 rows (tf-gemma4-kernels,tf-gemma4-glue), and greedy decode matches the oracle token for token and logit for logit. Global layers' decode attention runs a key head's query heads as one tensor-unit tile.--drafter, 150-190 tok/s for one stream, the same tokens); every stream equal to its plain reply (tf-gemma4-lanes). The target hands the drafter each prompt chunk's tapped rows as it makes them.core/affine_mmand a tensor-unit attention kernel (32 query rows a threadgroup, a simdgroup scoring 32 keys over the whole head and owning 64 output columns, fp32 probabilities crossing threadgroup memory 128 keys at a time).affine_mmnow dequantizes each step's weights to bf16 in its tile (scale times code plus bias in fp32, rounded once, as mlx-lm's quantized matmul does) and accumulates over the whole depth: the per-group epilogue, the row-sums kernel and its host plumbing go, across the GLM callers and tests too.tf-affine-refchecks against the rounded weights with the tensor unit's 16-deep chain as its bound and keeps the exact reference for the lane kernel;tf-glm-rows-checkpasses.tools/zig/gemma4_gate.py: the oracle's receipts, the token gate, the served gate, and the long-context prompts up to 252k;tf-gemma4-benchfor stage costs.Measured on an M5 Ultra (macOS 27), 49 decoded tokens after each prompt, against mlx-lm 0.32.0 prefilling in 2048-token chunks:
Decode receipts (arithmetic, code, copy) equal the oracle's tokens; the story prompt parts at token 56 on an exact native tie. The long prompts equal it at 32k, 64k and 252k and part at 8k (token 25) and 128k (30), each where mlx-lm's own top two logits are exactly tied. Served gate (thinking and a tool call) passes;
zig build test,zig fmt --checkandlean_check.pyclean.Plain one-row decode at short context is 19.7 ms a token, 51 tok/s against mlx-lm's 46 (45 at 8k against 39, 40 at 32k against 32): each shape's lane projection takes the K split it streams fastest at on one row (o and down 4, the gate/up stack and the head 8), another fp32 summation order than the reference's, so
tf-gemma4-kernelskeeps the reference split bit for bit and reports the decode split, where a handful of outputs a linear move by one bf16 step; and the core reg lane kernel takes a family prologue and epilogue (Fused), so a one-row layer's dense tail, attention tail and geglu run inside the next projections, and a sliding layer's attention kernel preps its own heads and merges itself at one key chunk (five to seven dispatches a layer, the same bits).