Skip to content

cuda gemma4: the dense family on the shared operators (text, MLX affine 4-bit), decode under CUDA graphs - #1174

Closed
jschmied wants to merge 19 commits into
ashhart:mainfrom
jschmied:cuda-gemma4
Closed

cuda gemma4: the dense family on the shared operators (text, MLX affine 4-bit), decode under CUDA graphs#1174
jschmied wants to merge 19 commits into
ashhart:mainfrom
jschmied:cuda-gemma4

Conversation

@jschmied

@jschmied jschmied commented Oct 21, 2026

Copy link
Copy Markdown
Contributor

What this changes

Gemma 4 text on the native CUDA engine, #1048 step 8. This covers gemma4_unified (12B) and gemma4 (31B) at MLX affine 4-bit, group 64. Every launch goes through a shared operator from steps 1-7, and the family adds no kernel of its own. Step 7 (the soft cap, tf_softcap_rows) is its own commit here, d3bfe17. Decode forwards run as CUDA graphs, one per row count. Part of #1048.

Stacking. This branch is cuda-attn-global (#1124, which itself carries #1123) with the other open step PRs merged in at their current heads:

Step Branch PR Head merged Merge commit
1 norms cuda-norms-gemma #1056 2b6cab85 3aa5a07d
2 act_rows cuda-act-gelu #1068 6472617d 7065e6b5
3 embed_rows cuda-embed-gemma #1070 82f714a2 7cb51714
4 qk_rows cuda-qk-gemma #1071 2a86f4fa 794b6d23
5 window head 256 cuda-window-gemma #1123 a2bed92a (base, through #1124)
6 window head 512 cuda-attn-global #1124 74248000 (base)

Then this branch's own commits:

  • 7d0a463: the build and records files had gone past 600 lines after the merges. This makes room.
  • d3bfe17: step 7, the soft cap.
  • 18f34b2: the family.
  • 3686e1c: graphs.

How it merges.

Receipt

  • Environment:
    • TensorFold 3686e1cd (this branch).
    • NVIDIA GB10 (sm_121, memory shared with the host), CUDA 13.0 nvcc, Zig 0.17.0.
    • Checkpoint: the MLX 4-bit (group 64) conversion of google/gemma-4-12B-it.
    • The 31B was not run.
  • Exactness: token_sha256 over 6 prompts, each greedy and seeded-sampled (temperature 0.8, top_p 0.95, seed 1234).
    • Baseline: 1 stream, --no-drafts, --prompt-cache-gib 0, graphs off.
    • Compared runs:
      • graphs on, same settings;
      • 4 streams with drafts on and a 4 GiB prompt cache, run twice. The second run resumes from cached prompts (5,080 / 3,144 / 975 tokens).
    • Result: 36 of 36 equal to the baseline. The same 24 from the four-stream runs were also equal with graphs off.
    • So drafted == "draft": false, resumed == fresh, concurrent == solo, and graphs == eager.
    • The drafts are copies of the prompt's own text; there is no drafter yet (see Gotchas). One prompt accepted 367 of 367 drafts, another 62 of 125.
  • Fidelity: HF transformers 5.19 (Gemma4UnifiedForCausalLM) in fp32 on the checkpoint's dequantized weights, with the embedding scale bf16(sqrt(3840)) = 62.0 as the bf16 model runs it. It was teacher-forced over 8 greedy replies: 2,282 rows, prompts up to 2,463 tokens.
    • Top-1 agreement: 98.95%.
    • Every disagreement is the reference's second choice, at most 0.33 logits behind.
    • The same HF model in bf16 agrees with these rows 98.47%, with gaps up to 0.625.
  • Decode speed (12B, GB10): 29.7 tok/s on one stream, 107 tok/s total over four. These come from a client script with ignore_eos and 256 tokens per reply, not tools/bench_openai.py.
    • Graphs on vs off, four alternating server arms, three runs each after a warm-up:
      • one stream: 30.15/30.51 vs 29.94/29.79 tok/s (+1.3%);
      • four streams: 112.2/113.3 vs 111.5/110.8 tok/s (+1.6%).
    • Graphs were faster in both rounds.
  • Prompt speed: cold first request, 8,144 tokens TTFT 3.63 s, 16,310 tokens 8.01 s. This is a new family, so there is no "before".
  • Kernel contract (sm_121): SUMMARY 108 kernels: 108 pass gates 1-4, 0 fail; 37 controls: 37 red as planted, 0 missed. That is the step PRs' entries plus the soft cap and its two controls. Limit problems are unchanged from main.
    • Soft-cap records are for sm_80/86/89/120/121.
    • act_rows' existing functions have identical SASS on all five.
  • Tests: zig build test -Dnvcc=…: 225/225 steps, 1248/1251 tests passed, 3 skipped, 0 failed.
    • test-gemma4-cuda-host covers: the copies equal Metal's files, RoPE tables, the embedding scale, snapshot runs over the sliding ring, worded refusals, and the stream fit.
    • lean_check: 444 problems, the same as the base, 0 new. zig fmt --check is clean.
  • Not run, and why:
    • The 31B: no checkpoint here.
    • Metal: no Mac in this run; the Metal family is untouched.
    • Images: step 12.
    • A real drafter: steps 10/11.
    • tools/bench_openai.py and tools/bench_cache.py: the numbers above come from equivalent client scripts.

Gotchas

  • The 31B is registered but untested. The gemma4 model_type opens through the same path as the 12B (same layer kinds, affine-4/64). It has not been run, so the receipt covers the 12B only.
  • Not listed as supported. Proposal: Gemma 4 on CUDA, built on the shared CUDA operators (shared extensions first) #1048 makes image input part of the bar before CUDA lists Gemma 4. This PR is text only: --vision is refused by name, and no docs or support list change.
  • Prompt-copy drafts only. No drafter loads yet. --drafter is refused with Gemma's own message, and draft requests get NoDraftHead. The drafted == plain gate above is over prompt-copy drafts. A drafter comes with step 10 (12B MTP assistant) and step 11 (31B DFlash).
  • Graphs gain about 1.5%. Decode is limited by weight bandwidth, so launch overhead was mostly hidden already. They are on by default; TENSORFOLD_GEMMA4_GRAPHS=0 turns them off, and the bits are the same either way.
  • The config and widths files are copied. families/gemma4/cuda_config.zig and cuda_widths.zig copy Metal's config.zig and widths.zig, because a file belongs to one Zig module. Only their doc line and imports differ, and a host test fails if anything past line 4 drifts. Admission and refusals are therefore Metal's, error names included.
  • Other widths are refused. Anything other than affine 4-bit/64 is refused at load (GemmaCudaWidthNotServed). That is step 9.
  • Not bit-equal to Metal. Metal's fast exp/cos/sin/pow and reduction orders differ (Proposal: Gemma 4 on CUDA, built on the shared CUDA operators (shared extensions first) #1048 says so), so fidelity is measured against the fp32 reference above.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
  • No precision traded for speed. Graphs replay the eager launches; same bits.
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
  • Prompt processing is no slower than the last release. It is a new family; no other family's path changed.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

jschmied added 19 commits October 19, 2026 01:38
@jschmied
cuda norms: bf16 residual tails with a branch norm, a layer scalar an…
88fd80f 
…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.
@jschmied
cuda act_rows: GeGLU and SwiGLU in MLX's bf16 steps
9f98fd9 
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.
@jschmied
cuda qk_rows: head dims to 512, a unit v norm and MLX's rounded rows
1446be3 
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.
@jschmied
bf16_round: an underflow to zero keeps the sign, checked by raw bits
668ab15 
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.
@jschmied
cuda embed_rows: MLX affine rows at 2 to 6 and 8 bits with a bf16 scale
360df33 
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.
@jschmied
norm tail check: ordered ends for ranges with an infinite end, bits f…
2b6cab8 
…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].
@jschmied
bf16_round: an underflow to zero keeps the sign, checked by raw bits
6472617 
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.
@jschmied
bf16_round: an underflow to zero keeps the sign, checked by raw bits
82f714a 
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.
@jschmied
bf16_round: an underflow to zero keeps the sign, checked by raw bits
2a86f4f 
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.
@jschmied
cuda window_attention: head dim 256, any GQA, a row's last key and fp…
a2bed92 
…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.
@jschmied
cuda window_attention: head dim 512 over a linear cache, for global a…
7424800 
…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.
@jschmied
Merge branch 'cuda-act-gelu' into cuda-gemma4
7065e6b 
# Conflicts:
#	docs/kernel-contract.md
#	zig/build/cuda.zig
#	zig/src/core/contract/registry.zig
@jschmied
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
@jschmied
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
@jschmied
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
build, records: room for the merged Gemma operator PRs under the 600-…
7d0a463 
…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.
@jschmied
cuda act_rows: soft-capped logits in MLX's bf16 steps (tf_softcap_rows)
d3bfe17 
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.
@jschmied
cuda gemma4: the dense family on the shared operators (text, MLX affi…
18f34b2 
…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.
@jschmied
cuda gemma4: decode forwards as CUDA graphs, one a row count
3686e1c 
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 %.
This was referenced Oct 22, 2026
Closed
Closed
@ashhart

ashhart commented Oct 24, 2026

Copy link
Copy Markdown
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 (the dense family on the shared operators) ends at 4987903. On a GB10 with mlx-community/gemma-4-12B-it-4bit: prefill 2,560 tok/s at 2k; decode 37 tok/s on one stream and about 200 tok/s at 8; every stream equals its run alone and its serial run; chat resume equals fresh, thinking on and off. A 32k needle miss on this 12B reproduces with the same reply on main's M5 Metal path, so it is the model's.

- Synthema, Ash's AI Assistant.

@ashhart ashhart closed this Oct 24, 2026
@jschmied jschmied mentioned this pull request Oct 24, 2026
Closed
@ashhart

ashhart commented Oct 24, 2026

Copy link
Copy Markdown
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.

@ashhart ashhart mentioned this pull request Oct 24, 2026
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@jschmied @ashhart