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.
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 assistant's hold on the device, as CUDA graphs) ends at 585a04a, and is covered by the MTP cells above, where drafted output equals plain in every cell. - 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
tf_round_pick_ids(tree_round.cu). It turns a round's picks into the next forward's int32 token ids: a pick's token, or a fallback id for a pick with none (or past the vocabulary). The pick's flag stays for the caller to read.TENSORFOLD_GEMMA4_GRAPHS=0keeps them eager.Stacking. Based on #1184 (
g4-dflash-31b,ee72eb35). This PR adds two commits on top of it:732ff66b(the kernel and its contract) and3e146b2f(the assistant). Branchg4-assistant-graphs.Receipt
google/gemma-4-12B-it-assistant@46d4c6f1.ee72eb35against this PR, four arms with a server restart each, the assistant on, fixed prompts:Code is +1.7 % in both rounds; everything else is within the runs' spread. I predicted +0-3 %, so this is the low end. The case for the PR is mostly structural: no host syncs in a hold.
tf_round_pick_idsis checked exactly at the assistant's hold sizes, with a control without the fallback that goes red. SUMMARY 111 kernels: 111 pass; 40 controls: 40 red; limit problems 1533, unchanged. tree_round.cu's other kernels have the same SASS on sm_80, 86, 89, 120 and 121.zig build test -Dnvcc=…: 229/229 steps, 1438/1445 passed, 7 skipped, exit 0. lean_check: 0 new.Not run, and why
Gotchas
DraftNonfinite, as before, but at its end instead of at that step.