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.
…pick_ids)
A draft chain's next step embeds the tokens its last step picked. Those
picks are {token, flags} pairs, and a row holding NaN or an infinity picks no
token. tf_round_pick_ids writes each pick's token as an int32 id, or a
fallback id when the pick has none (or names a token past the vocabulary).
The flag stays in the picks for the caller to read, so the next step needs no
host round trip.
It is under the contract (contract_pick_ids.zig): an exact check at the
assistant's hold sizes, and a control without the fallback that goes red.
SUMMARY 111 kernels, all pass; 40 controls, all red; limit problems
unchanged. tree_round.cu's other kernels have the same SASS on sm_80, 86, 89,
120 and 121.
Each draft step read its picks back on the host to refuse a row without a token before embedding the next ids: six syncs per hold. The next ids now come from tf_round_pick_ids on the device, and the flags are read once with the held drafts at the hold's end. With no host work between steps, a hold's steps run as one CUDA graph per (rows, steps), captured at first use; TENSORFOLD_GEMMA4_GRAPHS=0 keeps them eager. A 6-step chain takes 9.73 ms (10.0-10.8 eager): the steps are bound by bandwidth (the 4-bit head alone is 128 MB), not by launches. On a GB10, in alternating runs at one stream, code goes 114.3-114.8 -> 116.2-117.2 tokens a second; essay, explanation and history prompts, and four streams, stay within the runs' spread. Drafted replies equal the plain ones: 36 of 36.
…eGLU, pooling) Three kernels for the 31B's SigLIP-style tower, at torch's per-op bf16 rounding as transformers' Gemma4Vision* modules run it (and as Metal's vision.metal does): - tf_vit_g4_qkv: one warp per row, head and part of q|k|v [P, 3 H D], in place. q and k are normed by their RMSNorm weights, then the 2D RoPE turns each half of D by one axis; v gets a weightless RMSNorm. D is 72. - tf_vit_g4_gelu_mul: torch's tanh GeGLU on bf16 rows. - tf_vit_g4_pool: 3 x 3 block means in float32, rounded, times sqrt(C), then standardized by the tower's bias and scale. All three are under the contract (contract_vit_gemma.zig), checked bit for bit against the host's float32 steps (the GeGLU within tanhf's ulps). Each rounding point has a control that drops it, and every control goes red. SUMMARY 114 kernels, all pass; 44 controls, all red; limit problems unchanged. vit.cu's other kernels have the same SASS on sm_80, 86, 89, 120 and 121.
--vision on the 31B loads its 27-layer image tower in bf16. Its steps: - The 16-px patches go through the input projection (cuBLASLt), and the two position rows from the [2, 10240, 1152] table are added. - Each layer runs RMSNorm, the q|k|v GEMM, tf_vit_g4_qkv, the shared vision attention at head dim 72, the o projection, then the post-attention and pre-feedforward norms fused as a norm tail. Then gate and up, torch's GeGLU, down, and the post-feedforward norm fused with the next layer's input norm. - 3 x 3 pooling with the standardization, a weightless RMSNorm, and embed_vision's 4-bit projection close it. The image's rows go into the prompt as the 12B's do. Gemma 4's templates take image parts of type "image", so the CUDA vision offer passes a family's part_type through. The 31B's template wrote no placeholder for image_url parts. On the 31B: tools/zig/vision_served.py passes every check this box can run (the WebP decoder is not installed). A conversation with an image, resumed from the prompt cache and drafted by DFlash, equals a fresh plain run, 12 of 12 (greedy and seeded). Against transformers in fp32, the tower's rows are 5.20 % off per row on average; transformers' own bf16 run is 5.00 % off. The 12B's images are unchanged: 6 of 6.
9 tasks done
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 (images on the 31B, the image tower) ends at 850c692. Its image cells on a GB10 are still to run; text cells on the final stack qualify. - 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
--visionon the 31B (gemma4) loads its 27-layer image tower in bf16.tf_vit_g4_qkv: the per-head norms and the 2D RoPE, in place;tf_vit_g4_gelu_mul(torch's GeGLU), down, and the post-feedforward norm fused with the next layer's input norm.tf_vit_g4_pool(3 x 3 means, then the tower's standardization), a weightless RMSNorm, andembed_vision's 4-bit projection close it.Three new kernels, all in vit.cu (
vit_gemma.cuh). They round as transformers' Gemma4Vision* modules do (as Metal'svision.metaldoes).One change outside the family:
native/cuda_vision_offer.zigpasses a family'spart_typethrough, and Gemma 4 offers"image". The 31B's chat template wrote no image placeholder forimage_urlparts.Stacking. Based on #1185 (
g4-assistant-graphs,3e146b2f). This PR adds two commits on top of it:3fb53ada(the kernels and their contract) and0201f25b(the tower and the offer). Branchg4-images-31b.Receipt
mlx-community/gemma-4-31b-it-4bit@696d436c(tower in bf16, projection 4-bit),z-lab/gemma-4-31B-it-DFlash@eabd6483.vision_served.pypasses every runnable check, and an image conversation equals a fresh run, 6 of 6.Gemma4VisionModel+Gemma4MultimodalEmbedderin fp32 (266 rows): ours 5.20 % mean and 56.7 % max per-row relative error; transformers' own bf16 run 5.00 % and 48.0 %. Twenty-seven bf16 layers drift that much for both.tools/zig/vision_served.pyon the 31B: every check but the WebP decoder (libwebp not installed).zig build test -Dnvcc=…: 231/231 steps, 1417/1424 passed, 7 skipped, exit 0.Not run, and why
Gotchas
Question for @ashhart
part_typeon the CUDA vision offer the way you want families to name their template's image part? Metal's Gemma host already sets"image".