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.
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 (every MLX affine width, and the 31B) ends at c740c69. On a GB10 the 31B 4-bit qualifies in every cell: prefill 1,020 / 938 / 725 tok/s at 2k / 8k / 32k, the 32k needle found, decode 13.7 tok/s on one stream and 88 at 8, every stream equal to its run alone. - 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
#1048 step 9 on CUDA.
AffineN's tiles (decode in 16-row tiles, the prompt GEMM over the same tiles).embed_rowstakes 2-6 and 8), and the tied head is packed at the same width.mlx-q2g64tomlx-q8g64.gemma4, MLX 4-bit) needed no code. Its receipt is here because it is the other half of step 9.Stacking. Based on #1181 (
g4-assistant,0813e6ea), which is based on #1174. This PR adds one commit on top of it:ce493c47(branchg4-widths).Receipt
unigilby/gemma-4-12B-it-oQ8e@a47ebf1a(uniform 8-bit);mlx-community/gemma-4-31b-it-4bit@696d436c;All within the ranges predicted beforehand: oQ8e 14-19 tok/s and 3.5-6 s; 31B 9-13 tok/s, 3-3.8x and 8-20 s.
zig build test -Dnvcc=…: 225/225 steps, exit 0. lean_check: 0 new. fmt clean.Not run, and why
AffineN's tested paths with 8-bit.Gotchas
metaand thento_emptyallocates in the default dtype (fp32) first and took the box out of memory; set the default dtype to bf16 first.