Repository navigation
Conversation
…d MLX's rounded rows addRmsTail extends addRmsH for Gemma 4's sandwich tails: an optional RMS norm on the branch before the residual add, an optional bf16 scalar on the sum (rounded first), and a rounding flag that gives each norm MLX's bf16(w * bf16(x * inv)). With no option set it is addRmsH itself, and the existing kernels' code is unchanged. The kernels are registered with the kernel contract: their own check compares every output with float64's value at the rounding points (a range over the fp32 inverse's spread), the 64-value sums bit for bit with the fp32 tree, and four controls each drop one rounding point.
tf_geglu_rows computes mlx_lm's gelu_approx(gate) * up as MLX's bf16 ops (Gemma 4's MLP), and tf_swiglu_rows_round SwiGLU in tf_draft_swiglu's bf16 steps (Gemma 4's DFlash drafter), both with the 64-value group sums of the stored outputs. swiglu keeps its kernel and launch. Both are registered with the kernel contract: every output against float64's value at the rounding points (a range over tanhf's or expf's ulps), the sums bit for bit with the fp32 tree, and twelve controls each drop one rounding point.
tf_qk_rows_once carries tf_qk_rows' arithmetic to head dim 512 and can norm v by a unit RMS norm; tf_qk_rows_round adds MLX's rounding points, q and k = bf16(w * bf16(x * inv)) before RoPE, as Metal's gemma_norm_rope computes them. Shape.round and Shape.v_norm select them, as does a head dim past 256; every other call keeps its kernel and its launch, and tf_qk_rows and tf_qk_rows_w keep their code. Both are registered with the kernel contract: their own check compares every output with float64's value at the rounding points (a range over the fp32 inverse's spread and cosf's ulps), and three controls each drop one option.
bf16(-2^-135) gave +0 where bf16 rounding gives -0 (0x8000). The tests now compare u16 bits: signed zeros from underflow, subnormal ties, the normal boundary, overflow to both infinities, and the random fp32 sweep against the fp32 path's bits.
EmbedRows.affine reads MLX's affine tables at 2, 3, 4, 5, 6 or 8 bits (groups of 32, 64 or 128) and, given a bf16 scale, multiplies each dequantized bf16 value by it and rounds again: bf16(bf16(q * s + b) * scale), Gemma 4's embedding times bf16(sqrt(hidden)). With 4 bits, no scale and groups of 32 or 64 it is affine4Groups itself, and the existing kernels' code is unchanged. The 4-bit rows and the new kernel at each width are registered with the kernel contract: every output bit-equal to float64's value at the rounding points, and two controls, one without the bf16 rounding before the scale and one folding the scale into the group's scale and bias.
…or one-value ranges A range with an infinite end fell back to the float64 point's class: it accepted a zero below a finite lower end and refused the infinite end itself. Every range is now checked by its ordered ends (-0 below +0), a one-value range by its bits, NaN only for NaN. The control is the 64-value input-norm row whose first output's range is [largest bf16, +Inf].
bf16(-2^-135) gave +0 where bf16 rounding gives -0 (0x8000). The tests now compare u16 bits: signed zeros from underflow, subnormal ties, the normal boundary, overflow to both infinities, and the random fp32 sweep against the fp32 path's bits.
bf16(-2^-135) gave +0 where bf16 rounding gives -0 (0x8000). The tests now compare u16 bits: signed zeros from underflow, subnormal ties, the normal boundary, overflow to both infinities, and the random fp32 sweep against the fp32 path's bits.
bf16(-2^-135) gave +0 where bf16 rounding gives -0 (0x8000). The tests now compare u16 bits: signed zeros from underflow, subnormal ties, the normal boundary, overflow to both infinities, and the random fp32 sweep against the fp32 path's bits.
…32 probabilities window_rows.cu: decode by 256-key chunks and a merge in chunk order, prompt tiles over 128-dim slices; Window.p32 keeps P to 2^-16 as two bf16 halves (Metal's rounding points), promptRows takes each row's last key. window_attention.cu and its callers are unchanged. core/attention_range.zig: float64 at the rounding points and the range fp32 allows.
…ttention window_rows.cu at head dim 512: decode partials with K and V tiles in shared memory, a block for 4 query heads and a warp for each key run, merged in chunk order as at 256; prompt tiles as at 256. No window over a linear cache to 262144 keys, GQA 8 and 16.
# Conflicts: # docs/kernel-contract.md # zig/build/cuda.zig # zig/src/core/contract/registry.zig
# Conflicts: # docs/kernel-contract.md # zig/build/cuda.zig # zig/kernels/cuda/contract_controls.cu # zig/src/core/root.zig # zig/tests/cuda/contract_tests.zig
# Conflicts: # docs/kernel-contract.md # zig/build/cuda.zig # zig/kernels/cuda/contract_controls.cu # zig/src/core/contract/registry.zig # zig/src/core/root.zig # zig/tests/cuda/contract_tests.zig
# Conflicts: # zig/kernels/cuda/contract_controls.cu # zig/src/core/root.zig # zig/tests/cuda/main.zig
…line limit The merges put build/cuda.zig at 614 lines and the shared CUDA registry at 603. The five Gemma operator contract tests (window rows, MLX activations, embed rows, qk rows) are one loop, and embed_rows' and qk_rows' records move to gemma_cuda_records.zig beside act_rows'. Nothing a test runs changes.
Step 7 of ashhart#1048: a soft-capped head's logits, cap * tanh(x / cap), as MLX computes them in bf16 ops (Metal gemma4/glue.metal's gemma_softcap): bf16(bf16(tanh(bf16(x / cap))) * cap), the division and the product correctly rounded in fp32 and rounded to bf16 at each point, tanhf for tanh. Nothing in it is Gemma's: any final_logit_softcapping head can call ActRows.softcap(s, x, out, rows, n, cap), in place or not. It sits in act_rows.cu beside the MLX activations; the module's 30 other functions keep their SASS on sm_80, 86, 89, 120 and 121. core/act_ref.softcap is the float64 reference at those points (fp32 then bf16, as the kernel rounds) and softcapRange its ends with tanhf within 2^-21 either way. Contract: tf_softcap_rows at Gemma 4's head (1 and 16 rows of 262,144 logits, 3 of 4,096), inputs to +-80 around the cap; two controls, one dropping x / cap's rounding, one tanh's (tf_ctl_softcap), both red at the shapes gate. Records for 256 and 128 threads from nvcc 13.0 (24 registers, 25 on sm_120 and 121). GB10: SUMMARY 108 kernels: 108 pass gates 1-4, 0 fail; 37 controls: 37 red as planted, 0 missed; 1533 limit problems (the base: 107, 35, 1533); both launches pass limits live.
…ne 4-bit) Gemma 4 (gemma4_unified 12B, gemma4 31B) on the native CUDA engine, every launch a shared operator from ashhart#1048 steps 1-7: qlinear's affine-4 tiles, embed_rows, the MLX-rounded norm tails, qk_rows, window_attention (sliding ring 1536, global linear), act_rows' GeGLU and soft cap, tree_round's bf16 argmax and the keyed draws. No new kernels. Config and widths are the Metal family's files, copied because a file belongs to one module; a host test holds the copies equal but for their first lines. Prompts run in chunks of up to 512 rows cut where Metal cuts, decode windows up to 16 rows a stream and 128 a forward; prompt-cache states keep global layers' positions and each sliding layer's last 1023. Other widths, images and drafters (steps 9-12) are refused by name.
A decode forward's launches depend on its row count alone: positions and slots are read on the device. Each row count (1 to 128) is captured at its first use and replayed for every slot and position after; TENSORFOLD_GEMMA4_GRAPHS=0 keeps every forward eager. Replies are bit-identical to the eager forward (36 of 36 token hashes, solo, four streams, copy drafts, resumed prompts, greedy and seeded sampling). On a GB10, alternating arms: 1 stream +1.3 %, 4 streams +1.6 %.
…p 10) --drafter loads google/gemma-4-12B-it-assistant over the 12B on CUDA. Its admission is Metal's, refusal for refusal, and its bf16 linears and tied head are made MLX affine 4-bit on the host as Metal makes them. A hold runs every stream's chain of up to six drafts through shared operators: the target's embedding beside the absorbed row, the pre projection, four layers whose queries read the target's last sliding and last global caches at the kept row's position, the head and the bf16 argmax, then the post projection into the next step. The backend serves the drafter's tap, the last layer's rows before the final norm, from the forward that made them, and streams each prompt chunk's rows to it. The target's windows of 1 to 16 rows are timed at load so the depth rule prices drafts against them. Drafted replies equal the plain ones: 36 of 36 token hashes at one and four streams, resumed prompts, greedy and seeded sampling. On a GB10, alternating arms: one stream 47.8-52.9 against 30.6 tokens a second, four streams 170-176 against 113.5.
…tep 9) Each projection packs at its own width from the config's table: 4-bit into qlinear's affine-4 tiles as before, 2, 3, 5, 6 and 8 bits into AffineN's tiles (decode in 16-row tiles, the prompt GEMM over the same tiles). The embedding is read at its width and the tied head packed at the same one. The native registry lists mlx-q2g64 to mlx-q8g64. On unigilby/gemma-4-12B-it-oQ8e (uniform 8-bit): concurrent, resumed and assistant-drafted replies equal the plain ones (24 of 24 token hashes, greedy and seeded), and the 4-bit 12B's are unchanged (12 of 12). Teacher-forced against HF transformers in fp32 on the dequantized weights: 99.35 % top-1 over 2,292 rows. GB10: 17.2 tokens a second for one stream, 65.7 for four, 36.6 with the assistant; an 8,144-token prompt in 3.91 s.
tf_vit_layer_norm holds 8 values a thread, rows to 2048. Its body moves into vit_norm.cuh as a template over the values a thread; tf_vit_layer_norm keeps 8 (the same SASS on sm_80, 86, 89, 120 and 121), and tf_vit_layer_norm_wide takes 32, rows to 8192, for Gemma 4 12B's image embedder (6912 and 3840 wide). Both kernels are now under the contract (contract_layer_norm.zig): a float64 reference with the mean's and the squares' errors carried to each output, stress rows of zero-sum integers for gate 4, and two controls on the wide rows, one taking the variance in one pass and one rounding the normalized product before the bias. SUMMARY 110 kernels, all pass; 39 controls, all red; limit problems unchanged.
--vision on a 12B checkpoint loads its image embedder: the merged 48-px patches through LayerNorm, patch_dense (cuBLASLt, its bias rounded once as addmm's), LayerNorm, the two position rows, pos_norm, the weightless RMSNorm and embed_vision's projection at its own width, each step as torch rounds it. The processor is the Metal family's image.zig (copied as cuda_image.zig, kept equal by the host test). A quantized patch_dense and the 31B's tower are refused by name. An image's rows replace its placeholder rows unscaled, and on sliding layers each sees all of them (window_attention's last-key rows); global layers stay causal. A prompt chunk never cuts through an image, and a kept state never resumes inside one. On the oQ8e 12B: the embedder's rows are 0.342 % from torch's fp32 (torch bf16: 0.342 %); tools/zig/vision_served.py passes every check this box can run (WebP and JPEG decoders are not installed); drafted, concurrent and resumed image conversations equal fresh plain ones, greedy and seeded, with copy drafts and with the assistant. Teacher-forced against HF transformers on the same images: 96.3 % top-1 in fp32, 96.5 % in bf16.
--drafter with a DFlashDraftModel directory loads z-lab/gemma-4-31B-it-DFlash over the 31B. Its admission is Metal's, and its bf16 linears are made MLX affine 4-bit on the host as Metal makes them. The target keeps the drafter's tapped layers (1, 12, 23, 35, 46, 57) every forward: any set of layers now goes to a tap buffer, and decode graphs are captured again after a set changes. An absorb projects those rows through fc and hidden_norm into each layer's keys and values, normed and roped, at their positions. A hold runs up to eight 16-row blocks through the shared block_attention, one call a key head: causal with a 2047-key window on the sliding layers, the whole block on the full one, as Metal masks them. Then the target's tied head and the bf16 argmax. A kept prompt state carries the drafter's context, and a state without it is not resumed while the drafter is attached. States are marked after the chunk's absorb. Drafted replies equal the plain ones: 36 of 36 token hashes at one and four streams, resumed prompts (5,080 tokens cached with the drafter's context), greedy and seeded sampling. On a GB10, alternating arms over fixed prompts at one stream: 75.8, 23.7, 34.4 and 31.5 tokens a second (code, essay, explanation, history) against 12.5 undrafted; four streams 75.2 against 46.8.
Owner
|
Thank you @jschmied. The seven Gemma 4 CUDA stages landed together on main (4167b87..850c692), with your records kept as authored and maintainer children for main's contract registrations and line limits. Every prefix passes its host gates. This stage (z-lab's DFlash drafter on the 31B) ends at 6f1caa0. On a GB10: greedy chat 12.6 to 29.9 tok/s on one stream (2.4x), 47 to 91 at 4 and 86 to 108 at 8, with drafted output equal to plain everywhere; code at temperature 1.0 gains 7-9% at 1-4 streams and loses 6% at 8. - Synthema, Ash's AI Assistant. |
Owner
|
This shipped in TensorFold 1.1.4: https://github.com/ashhart/TensorFold/releases/tag/v1.1.4 Thank you @jschmied for the contribution. - Synthema, Ash's AI Assistant. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this changes
--drafterwith aDFlashDraftModeldirectory loadsz-lab/gemma-4-31B-it-DFlashover the 31B (#1048 step 11).Admission. Metal's:
NotDFlashV1,UnsupportedDraftMask,DraftForAnotherTargetand the rest. Its bf16 linears are made MLX affine 4-bit on the host, as Metal does.Taps. The target now keeps any set of layers' outputs every forward, in a tap buffer: here layers 1, 12, 23, 35, 46 and 57. Decode graphs are captured again after a set changes.
Absorb. Those rows go through fc and hidden_norm into each layer's keys and values, normed and roped, at their positions.
Hold. Up to eight 16-row blocks run through the shared
block_attention, one call per key head, masked as Metal masks them:Then the target's tied head and the bf16 argmax.
Prompt cache. A kept state carries the drafter's context, and a state without it is not resumed while the drafter is attached. States are now marked after the chunk's absorb: the first run failed resumed prompts with
DraftContextGap, because a state was copied before the drafter held its rows.Stacking. Based on #1183 (
g4-images-12b,50834d91). This PR adds one commit on top of it:ee72eb35(branchg4-dflash-31b).Receipt
mlx-community/gemma-4-31b-it-4bit@696d436c,z-lab/gemma-4-31B-it-DFlash@eabd6483.I predicted 25-35 tok/s at one stream: code is far above (6.1x), and essay prose just below (1.9x).
block_attentionis shared and registered).zig build test -Dnvcc=…: 227/227 steps, 884/886 passed, 2 skipped, exit 0. lean_check: 0 new.Not run, and why
DraftForAnotherTargetrefuses it on the 12B (host test).Gotchas
stream_bytes.block_attentionreads a context's heads packed as [hk, len, 128], while ours are capacity apart, hence one call per key head.