cuda gemma4: images on the 12B, and LayerNorm rows to 8192 wide (#1048 step 12) by jschmied · Pull Request #1183 · ashhart/TensorFold

jschmied · 2026-10-22T05:47:29Z

What this changes

#1048 step 12 for the 12B (gemma4_unified).

--vision loads its image embedder. Each step is rounded as torch rounds it:

the merged 48-px patches through LayerNorm;
patch_dense through cuBLASLt, its bias rounded once as addmm's is;
LayerNorm, then the two position rows;
pos_norm, the weightless RMSNorm, and embed_vision's projection at its own width.

The processor is the Metal family's image.zig, copied as cuda_image.zig and kept equal by a host test.

An image's rows:

replace its placeholder rows unscaled;
see each other on sliding layers, through window_attention's last-key rows; global layers stay causal.

A prompt chunk never cuts through an image, and a kept prompt state never resumes inside one.

Refused by name:

a quantized patch_dense: the local MLX 4-bit 12B has one, and Proposal: Gemma 4 on CUDA, built on the shared CUDA operators (shared extensions first) #1048 refuses it too;
the 31B's tower, which comes in a later PR of this stack.

One new kernel: tf_vit_layer_norm_wide, rows up to 8192 (the embedder's 6912 and 3840). tf_vit_layer_norm's body moved into vit_norm.cuh as a template over the values a thread holds. The existing 2048-wide kernel has the same SASS on sm_80, 86, 89, 120 and 121.

Stacking. Based on #1182 (g4-widths, ce493c47). This PR adds two commits on top of it: 2190a6d3 (the wide LayerNorm and its contract) and 50834d91 (images). Branch g4-images-12b.

Receipt

Environment. One GB10 (sm_121), CUDA 13.0, Zig 0.17.0. unigilby/gemma-4-12B-it-oQ8e @ a47ebf1a, whose patch_dense is bf16. Test images: generated PNGs (three red circles on blue; a five-bar chart).
Exactness.
- Drafted equals "draft": false with an image.
- Concurrent image requests equal their runs alone, 4 of 4.
- A later turn resumes past its image (2,420 tokens cached). Another image under the same placeholders resumes only the text before it (1,579 tokens).
- An image conversation resumed from the cache and drafted (copy drafts, and the assistant) equals a fresh plain run, 6 of 6, greedy and seeded.
Fidelity.
- Embedder rows against torch fp32: 0.342 % mean per-row relative error. Torch bf16 is also 0.342 %, and 73 % of our values are bit-identical to torch bf16.
- Teacher-forced image replies (707 rows) against HF: 96.32 % top-1 in fp32 and 96.46 % in bf16. Text replies on the same checkpoint reach 99.35 %; image prompts are simply more bf16-sensitive, and HF's own bf16 run lands at the same level.
Speed. First token for one image (about 284 prompt tokens): 0.16-0.20 s warm, 0.45 s on the first request (cuBLASLt plans).
Contract. SUMMARY 110 kernels: 110 pass gates 1-4; 39 controls: 39 red as planted; limit problems 1533, unchanged. Both LayerNorm kernels are now under the contract with a float64 reference: stress rows of zero-sum integers for gate 4, and on the wide rows one control taking the variance in one pass and one rounding the normalized product before the bias.
Tests.
- tools/zig/vision_served.py: every check passes but one, the WebP decoder. libwebp is not installed on this box; JPEG cases are skipped for the same reason.
- Full zig build test -Dnvcc=…: 227/227 steps, 1414/1421 passed, 7 skipped (the missing libjpeg/libwebp), exit 0.
- lean_check: 0 new.

Not run, and why

JPEG and WebP inputs: libjpeg-turbo and libwebp are not installed on this box.
Metal: no Mac.

Gotchas

The local MLX 4-bit 12B cannot take images: its patch_dense is quantized, and it is refused by name.

…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.

…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.

ashhart · 2026-10-24T13:55:46Z

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 12B, LayerNorm rows to 8192) ends at 823c709. Its image cells on a GB10 are still to run; text cells on the final stack qualify.

- Synthema, Ash's AI Assistant.

ashhart · 2026-10-24T19:38:38Z

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.

jschmied added 23 commits October 19, 2026 01:38

Merge branch 'cuda-act-gelu' into cuda-gemma4

7065e6b

# Conflicts: # docs/kernel-contract.md # zig/build/cuda.zig # zig/src/core/contract/registry.zig

Merge branch 'cuda-embed-gemma' into cuda-gemma4

7cb5171

# 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

Merge branch 'cuda-qk-gemma' into cuda-gemma4

794b6d2

# 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

Merge branch 'cuda-norms-gemma' into cuda-gemma4

3aa5a07

# Conflicts: # zig/kernels/cuda/contract_controls.cu # zig/src/core/root.zig # zig/tests/cuda/main.zig

jschmied mentioned this pull request Oct 22, 2026

cuda gemma4: z-lab's DFlash drafter on the 31B (#1048 step 11) #1184

Closed

ashhart closed this Oct 24, 2026

Navigation Menu

Repository navigation

cuda gemma4: images on the 12B, and LayerNorm rows to 8192 wide (#1048 step 12) - #1183

cuda gemma4: images on the 12B, and LayerNorm rows to 8192 wide (#1048 step 12)#1183
jschmied wants to merge 23 commits into
ashhart:mainfrom
jschmied:g4-images-12b

jschmied commented Oct 22, 2026

Uh oh!

ashhart commented Oct 24, 2026

Uh oh!

ashhart commented Oct 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

jschmied commented Oct 22, 2026

What this changes

Receipt

Not run, and why

Gotchas

Uh oh!

ashhart commented Oct 24, 2026

Uh oh!

ashhart commented Oct 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants