Repository navigation
Conversation
|
Thank you @cshintov! Routed experts on the 27B's engine, with shared lanes and the prompt cache kept, are a big step toward #931. Please keep this open while four things get done:
Posted by Ash's AI Assistant Synthena |
|
Thanks Ash. Keeping it open. Item 3 is in It ran on an M4 Pro only, so the M5 paths (the prompt kernel, the prefix kernels and the lost-edge control) are not run. Reading the source, I expect gate 1 on prompt rows to go red on an M5: the prompt kernel multiplies a masked future key's -Inf value by its probability of 0, which gives NaN where float64 gives +Inf. That is a prediction, not an observation, so there is no known-finding entry for it. Gate 1 also plants no NaN or Inf in the fused prefix. Could you run these on an M5? zig build tf-kernel-contract -Dcpu=apple_m1
zig-out/bin/tf-kernel-contract --only 'metal qwen27 attention 27B prompt rows'
zig-out/bin/tf-kernel-contract --only 'metal qwen27 attention 35B-A3B prompt rows'
zig-out/bin/tf-kernel-contract --only 'metal qwen27 attention 27B decode rows'
zig-out/bin/tf-kernel-contract --only 'metal qwen27 attention 35B-A3B decode rows'
zig build kernel-contract -Dcpu=apple_m1Next: 2 (move to |
|
Item 4, first part: the shared-checkpoint CUDA comparison, logits only. Setup: mlx-community/Qwen3.6-35B-A3B-4bit at 38740b84, the same four shards on both machines (sha256 equal). CUDA is
The long prompt's four mlx_lm mismatches (24, 39, 44, 54) are the four positions with the smallest Metal margin CUDA's step and window modes gave bit-identical rows on both prompts. Still open in item 4: the staged per-layer references. No engine dumps every layer today, so that needs new code on |
157055b to
69c412d
Compare
|
Item 2: the branch is rebased onto main (d1018b3), and the 35B's config, loader and adapter now sit in What moved
Checks
What the rebase changed
A question for you: prefix_5 and _6 run at 320 and 384 threads. They are tensor-op kernels with no Not run here (an M4 Pro, no tensor units): a 35B prompt pass past 128 rows, and prefix_5/_6 through gates 1-4. The Next: the rest of item 4 (per-layer stage dumps on Metal, CUDA and mlx_lm against a float64 reference), then item 1. |
|
On the
None of these files are touched by this PR ( #1193 carries the three fixes against Once #1193 lands, rebasing this branch onto the new main should turn |
|
Item 4: the staged references fit the near-tie explanation. Every top-1 difference from a float64 forward, by The tools (commits
One stage's error (long prompt, 378 prompt rows plus 72 teacher-forced rows; median relative error per stage)
Top-1 against float64 (72 teacher-forced positions)
The 5 positions with a float64 margin under 2 bf16 steps (a step is 0.125 at these logit sizes):
Your positions 24, 39, 44 and 54 are four of these five. The two mlx_lm runs also disagree with each other, at 24 Metal's pick at 24 moved since my 10-20 comment. On Checks
Not run yet: the CUDA stage dump. Its code compiles, but I have not had a GB10 window for it since writing it; Next: item 1, MTP drafts in the shared lanes, then the image tower. |
|
Item 4, the CUDA stage dump: CUDA's per-stage error is in line with Metal's, a little lower on average, and its one Run (GB10,
One stage's error (long prompt, 378 prompt rows plus 72 teacher-forced rows; median relative error per stage,
CUDA is lower than Metal at 59 of the 80 stages and lower than mlx_lm at all 80. CUDA's and Metal's logits stage Top-1 against float64 (72 positions): CUDA 71, Metal 69, mlx_lm stage-dump run 68, the reference tokens 68. |
9047aa7 to
b318def
Compare
…hs under 4,096 Each lane takes four 32-wide steps, which assumed GLM's 4,096. At 2,048 the last lanes read past the router's rows. At 4,096 the guard is never taken, so GLM's bits are unchanged.
…ts (ashhart#931, slice 1) One engine for both models: each layer's MLP is a tagged union, and forward.post switches on it; the dense arm is the 27B's code as it was. Routing is ashhart#856's rule: fp32 logits, top-k by logit with ties to the lower id, exp(l - l_top) weights rounded to bf16, the shared expert as a slot weighted bf16(sigmoid(bf16(gate))). The experts run on core/tile_qmm's tiles and their decode twin. The attention head counts, the K/V unpack width, the K/V cache size and the shape check read the config. The attention kernels take Q27_GROUP and Q27_KV_HEADS defines that default to the 27B's 6 and 4, so the 27B's geometry is unchanged. The 35B refuses --drafter (slice 3) and --vision (slice 2).
…rnel at every served shape Select, sort and combine at the 35B-A3B's shapes against CPU oracles, with the selection's stated NaN rule; select and combine declare their fp32 sums for gate 4. M1 and M2 limits from the offline compile. tf_route_logits at GLM's router and indexer gate (decode and prompt chunks to 16,384 rows) and the 35B's router (to 128), against float64, with gate-5 launch records. Its 256 threads are reserved in the source: the translator gives it no M1/M2 limit, and every caller launches 256.
Decode rows and prompt rows for the 27B (24 query heads over 4 KV heads) and the 35B-A3B (16 over 2), driven through the engine's own Attention.encode so the contract checks the launches the engine makes. Float64 gated attention is the reference; gate 4's inputs make fp16 and bf16 score sums miss. Pack, append and keep compare bit-exact, keep also a second time on the caches the first left. Controls: a row seeing one key too far, a merge dropping chunks, an append one row late, and on tensor-unit chips a lost tile edge. A plan check on every chip shows the decode shapes plan prefix_1 to _4 for chains and _2 to _4 for trees, and a tree over two tail chunks; it reads the engine's plans, not an M5's launches. Gate-5 records for every kernel at both shapes; M1 and M2 limits from the offline compile. Run on an M4 Pro only. Not run: the prompt kernel, the prefix path and the lost-edge control, which need an M5's tensor units. Predicted from the source, not observed: on an M5, gate 1 on prompt rows may give NaN where float64 gives +Inf, a masked key's -Inf value times probability 0. Gate 1 plants no nonfinite input in the fused prefix.
… 27B and the 35B-A3B qwen27 had its own parse of config.json beside families/qwen3_5/config.zig, the reader the CUDA and HIP engines use. It now maps the family's Spec instead, reached through core like moe_route. The Metal engine keeps its own refusals: a tied head, a MoE field on a dense model type, a zero width, and the *_text model types it accepts. The family check refused uneven head groups by dividing, so a config with zero key heads or KV heads crashed it; both are refused first now. The real 27B, 35B-A3B and 27B-MTP configs parse to the same Config before and after. On 1,580 single-field edits of them, 219 differ, all on malformed values: floats such as 1.5 are truncated as the Python reference does (every width is then pinned by supported()); null reads as absent; values beside the ones used are checked too; and end ids are the root's, the text tower's and generation_config's, as on CUDA.
The experts' tensor checks, router packing and tile views move from qwen27/moe_weights.zig to qwen3_5/metal_experts.zig. They take the expert sizes and widths as plain values instead of qwen27's Config. The 27B's engine still loads them in its one layer loop (gpu_weights.zig's Mlp union), reaching the file through core like the family config. The 35B-A3B's teacher-forced top-2 is byte-identical to the step before on both prompts (72 and 60 positions). The qwen27 contract gates pass, and test-qwen27 and test-native pass.
…ense 27B's engine families/qwen3_5/metal_native_moe.zig names the qwen3_5_moe model type on Metal: its display name, its qualified format, and the flags refused until later slices (the checkpoint's MTP drafts, images). It is the twin of cuda_native_moe.zig. native/metal.zig reads it instead of comparing model-type strings in six places, and still opens the model through the dense 27B's openQwen27 and lane host. Refusal and log words are unchanged; test-native's refusal tests pass as before.
The rebase onto main brought registry.zig to 628 lines with qwen27's 38 records, past the 600-line rule. They move to registry_qwen27.zig, joined to registry.metal the way registry_media.zig joins registry.cuda. Two new family files' headers become one line each.
Main sizes the frame to max(rows asked, the prompt pass's rows): 512 by default on tensor units, up to 2,048 with TF_PROMPT_ROWS. The expert stage kept the 128 rows the native server asks for, so on an M5 a 35B prompt chunk past 128 rows failed with BadMoeRows. The stage now takes the frame's capacity. The contract's expert kernels and the 35B router's logits are checked at 512 and 2,048 rows too.
Main launches q27_attn_prefix_5 and _6 when a cache has 64 chunks or more. The contract registered and planned prefix_1 to _4 only, and its reach check clamped every plan to 4 tiles. attention.prefixTiles is now the one rule both the launcher and the reach check read; the decode shapes add 5 and 6 tiles at 8,190 and 8,191 cached keys (trees and chains) for both head counts. With the old 4-tile rule put back, reach fails for both head counts. Their 320- and 384-thread records are pending (registry_qwen27.pending, in all_metal beside Flash Next's): they are tensor-op kernels past 256 threads with no reservation in the source, as main launches them, so the limits gate flags them on M1/M2 until the source reserves the threads or the rule exempts launches that only run on tensor units.
…tions instead of crashing geometry() divided the hidden width by num_attention_heads when head_dim is absent, null or 0, so a zero head count crashed. mrope() added the three sections and doubled the sum unchecked, so sections near the 64-bit limit overflowed. Both now return the reader's ordinary errors. The Metal engine reaches this reader since it reads config.json through the family; its own parser had refused both before any arithmetic. The tests enumerate every mix of head count, hidden width and head_dim form, and every mix of small, exact and 64-bit-limit sections; with either fix taken out, they crash.
…orward One stage vocabulary for the family (families/qwen3_5/stages.zig): the residual stream after the embedding, after each layer's mixer and MLP, the final norm and the logits, in one on-disk layout. Off unless asked for: tf-qwen27-run --stages DIR and tf-qwen35-cuda <model> stages <ids> <dir>. Both engines claim planes through the same Round; Metal copies on the GPU, CUDA downloads once its stream drains (and runs uncaptured, as a learner). tools/zig/qwen36_stage_ref.py measures each stage against float64 from the checkpoint's exact weights: fed the engine's own input (that stage's error) and chained from the ids (the accumulated error). qwen36_mlx_stages.py dumps mlx_lm's generate_step path in the same layout. Long prompt (378 + 72 teacher rows), median relative error of one stage: Metal 0.23%, mlx_lm 0.32% (mean over stages). Every top-1 difference from float64, by Metal, mlx_lm, CUDA and the mlx_lm reference tokens, falls on the 5 of 72 positions whose float64 margin is under 2 bf16 steps.
…round whole or not at all A dump labels rows by consecutive position, so both engines now ask one rule (stages.admit) before a round runs: one stream, a linear chain of rows, and the head. Metal used to fence only the window count, after the round's state was prepared; a one-window tree was labelled as a chain. CUDA refused a beside stream only mid-forward and skipped the hooks for a hidden read. The recorder reserves every plane and position list before it appends any, and a round's marks are spent when it ends, so a repeated end is refused. A recording failure after Metal's GPU work ends the runner like a GPU error does, instead of leaving a retry that would replay committed recurrence. The sequence test now runs every four of: begin, a full round with each head, an odd head, a partial round, end, and an allocation failure at three depths, against a model of the bytes and positions that land; it reaches landed rounds and refused ones (the old four-mark bound never landed one). The comparer takes every plane's shape from the model and the positions, not from the dump's own index.
Engine.init builds the engine from undefined and sets each field, so the new stages field kept garbage, and the warm-up forward inside init called through it: tf-qwen35-cuda died with a bus error at its first forward on a GB10. Found by the first CUDA stage run; the field now starts null like learning and mlp beside it.
b318def to
47a6424
Compare
|
#1125 (Ornith, Qwen3.5-MoE in the qwen27 Metal engine), which you're carrying onto main, overlaps 16 of this PR's files, including the expert kernel, |
|
Review of the current head rebased onto main e9b69ed (it carries, with conflicts in native Metal admission resolved and two attention constructor calls updated for main's API). Four things before device qualification:
Then: the corrected stage comparison on current heads with all teacher-forced top-two margins and worst-column FP64 error by projection class, and old-family A/B (27B DFlash/cache identity, CUDA/HIP config). Text-without-MTP as the first slice is fine. |
What this changes
qwen3_5_moecheckpoints (Qwen3.6-35B-A3B) serve text on the native Metal engine through the 27B's engine,whose layers they share: each layer's MLP is now a tagged union of the 27B's dense MLP or routed experts. Slice 1
of #931. Images (slice 2) and MTP drafts (slice 3) are refused until their slices land.
exp(l - l_top)over the k rounded to bf16, the shared expert as one more slot weightedbf16(sigmoid(bf16(gate))), the slots' weighted sum rounded once to bf16. The logits run oncore/moe_route's kernel; selection, the by-expert sort and the weighted sum are three new kernels inkernels/metal/qwen27/moe.metal; the experts run oncore/tile_qmm's prompt tiles and decode twin.Q27_GROUP/Q27_KV_HEADSdefines whose defaults are the 27B's 6 and 4, so the 27B compiles without them),the K/V unpack width, the cache size check and the shape check.
core/moe_route.metal: a logits lane stops at the row's end. Each lane took four 32-wide steps, which assumedGLM's 4,096; at 2,048 the last lanes read past the router's rows. At 4,096 the guard is never taken.
--parallellanes come with the 27B's host, so they are qualified here (their receiptsbelow), and Qwen3.6-35B-A3B (qwen3_5_moe) on the native Metal engine: proposal and gate #931's slice 3 folds into this one.
refused at load (8-bit experts are slice 4).
--drafter,--visionand image parts are refused.build are not reused after this change.
families/qwen3_5/config.zig,metal_experts.zig,metal_native_moe.zig): qwen27 reads config.json through the family reader for the dense27B and the 35B alike, and the experts load in the 27B's one layer loop. The reader no longer crashes on zero
heads or
mrope_sectionvalues near the 64-bit limit; it returns its ordinary errors.Design: one engine, and the alternatives
The 35B is the 27B's architecture (Gated DeltaNet layers, gated full attention every fourth layer) at smaller
sizes with routed experts in place of the dense MLP, so it runs the 27B's engine with a different MLP step. The
union states in each layer's weights which MLP it runs; the dense arm is the 27B's code as it was.
Rejected:
qwen3_5_moefamily likelfm2_moe: a second copy of the DeltaNet layers, attention, rounds andhost (about 3,000 lines), or a cross-family import of the 27B's loop.
so slices 2 and 3 would be rebuilt there. The config, expert loader and adapter still belong in that family.
Runner(comptime M)): one copy, but a type parameter through about3,000 lines for a difference in the MLP step and a few pinned numbers.
The family's slices
merge 2, no deepstack) apart from its last projection (2,048 wide against 5,120), on the 27B's image path.
Receipt
Environment: this branch on
main6ee828c, before the rebase onto d1018b3 (what the rebase changed is inthe item 2 comment: the long prompt's logits move by bf16 steps from main's 128-key decode chunks); Mac mini M4
Pro (20-core GPU), 64 GB; macOS 26.5; Zig 0.17.0;
mlx-community/Qwen3.6-35B-A3B-4bitat38740b84; references from mlx_lm 0.31.3 (MLX 0.32.2) on the same Mac.The 27B checks use
TensorFold/Qwen3.8-27B-MLX-4bitat22d8d538andz-lab/Qwen3.8-27B-DFlash2at50307d4c.Exactness, 35B:
<|im_end|>. Teacher-forced, nativetop-1 equals mlx_lm's token at all 60 positions.
positions 24, 39, 44 and 54. At 24 and 39 both engines tie exactly (24: 20.0 and 20.0 native, 20.125 and
20.125 mlx_lm; 39: 21.75 and 21.75 in both); at 44 and 54 the two engines' top-2 logits are each one bf16 step
(0.125) apart, in opposite order. Free-running greedy diverges at the tie at 24.
tools/bench_concurrent.py --alone --serial, levels 1 and 8, temperatures 0 and 0.7, code and chatprompts, 128 tokens: 72 of 72 replies equal their solo run, 36 of 36 equal serial (an earlier run at levels 1
and 4, greedy: 30 of 30 and 4 of 4).
resumes at 6,581 of 6,660 tokens, first token in 0.28 s against 15.9 s with
--prompt-cache-gib 0, and itstoken SHA equals the cache-off pass; a second turn 2 resumes the same and is equal too.
tools/bench_cache.py --require-reuseon an identical raw 8,192-token completion reports no reuse, on the 27Bfrom this branch as on the 35B: the host keeps that prompt's state at its end, where an identical prompt
cannot resume. The repeated and 4 concurrent replies equal the first.
--paralleldefault, prompt cache on), five short sequences against a--prompt-cache-gib 0server's replies: a cancel in prefill then another request then the cancelled one whole; two streams in shared
rounds, one cancelled in decode and its slot taken by a third; a cancel in prefill of a prompt resuming part of
a kept state; an image refusal, then a text reply cancelled in decode, then that reply whole; four streams with
two cancelled in decode, then those two again. All 13 completed replies (greedy and seeded sampling, some
resuming kept states) equal the cache-off server's. Not run: a failed round mid-command-buffer (no fault hook).
tf-qwen27-streamson the 35B: 8 streams equal in a one-slot and an 8-slot runner over 48 tokens; 20 sharedrounds, 376 logits rows equal the streams alone (windows of 1 to 5 rows).
moe_test.zig: every sequence of up to three stage calls (prompt tiles, the decode twin andthe by-expert window at 1 to 16 rows, 258 sequences) gives each row its one-row call's bytes. It fails with a
stale sort position read in decode.
--drafterand--visionrefused before load (unit-tested innative/metal.zig); served, a chatrequest with an image part gets HTTP 400 ("this server accepts text parts only; image, audio and video inputs
are unsupported") and a text request on the same server then replies.
tensorfold-native inforeads thecheckpoint as
qwen3_5_moe, 4-bit in groups of 64, 19.03 GiB of weights.The 27B's bits:
tf-qwen27-runfrommainand from this branch on the 27B, a 378-token prompt, 64 tokens: thesame tokens, the same state-and-logits fingerprint after the prompt, every 32-row seam byte-equal between chunks
of 128 and 16 on both builds, teacher-forced top-2 ids and bf16 logits equal at all 64 steps.
tf-qwen27-streams,tf-qwen27-dflash-run --compare-plain,tf-qwen27-lanes-runandtf-qwen27-reusepasson both builds.
GLM: the logits guard is not taken at 4,096 (lane 31's last step is 127 of 128).
tf_route_logitscompiled frommain's source and from this branch's, at GLM's router (288 outputs) and indexer gate (128, bf16), on the samerandom inputs at 1, 3, 16, 2,047 and 16,384 rows: byte-equal output (a planted change to the source differs at
every height). No GLM checkpoint fits this Mac, so GLM's model-level bits were not rerun.
FP64 references (
moe_test.zig, a 256-wide, 8-expert stage against the host's rules and float64): routerlogits within 9.9e-8 of 1+|f64| at worst; the stage's output worst column 1.49e-2 at scale 3.08 (16 rows), the
prompt tiles and the decode twin equal on the same rows.
Kernel contract,
zig build kernel-contract -Dcpu=apple_m1:SUMMARY 36 kernels: 36 pass gates 1-4, 0 fail, 2 known findings; 8 controls: 8 red as planted, 0 missed; 0 limit problems(the two known findings areaffine_mm's, onmain). Without the logits guard the 35B router fails gates 1, 2 and 4 and GLM's two shapespass; without select's NaN rule select fails gate 1. New registrations:
q27_moe_select,q27_moe_sort,q27_moe_combineat the 35B's shapes (rows 1, 16, 128), andtf_route_logitsat GLM's router (4096x288) and indexer gate (4096x128, bf16 out) at 1 to 16,384 rows and the35B's router (2048x260) at 1 to 128 rows, each with its gate-5 launch record. The logits kernel now reserves
its 256 threads in the source (
max_total_threads_per_threadgroup(256)), as the vision kernels do: Xcode'stranslator gave no M1/M2 limit for it either way, and every caller launches 256. Gate 4:
combine's stress inputs make fp16 and bf16 sums miss 256x. Select's gate 4 compares the GPU's weights with a
float64 bound, but its inputs cannot show the contract's 100x separation: fp16 sums would miss 0x and bf16
0.3x, because each of the eight exponentials in (0, 1] is rounded to a bf16 weight. Its operation allowance
(2^16 + 64) stands for that bf16 rounding of weights stored as fp32, not for a count of operations. A bf16
denominator can still move a weight by a step (an exponential of 1 beside seven of 0.001); fp32 is the rule.
Attention (
attention_mpp,attention_prompt,attention_io), registered at both head shapes (the 27B's24 query heads over 4, the 35B's 16 over 2). A round's rows go through the engine's own plan and
Attention.encodeagainst float64 gated attention; prompt rows through the prompt kernel on tensor units;packing, append and keep bit for bit (keep also a second time on the caches the first left). Gate 4's inputs
make fp16 score sums miss without limit and bf16 by 87,066x; they separate the score sums only. A plan check
shows the decode shapes plan every prefix launch (1 to 4 tiles for chains, 2 to 4 for trees) and a tree over
two tail chunks; it reads the engine's plans, not an M5's launches. Controls: a row seeing one key past its
own, a merge dropping later tail chunks, an append one row late (all red here), and the prompt kernel losing
its bounded last step (tensor units only). Gate 5 has 16 launch records per head shape; M1/M2 limits for the
non-tensor entries come from the offline compile.
Not run here: the attention kernels' M5 paths (the prompt kernel, the prefix kernels, the lost-edge control).
Predicted from the source, not observed: on an M5, gate 1 on prompt rows may give NaN where float64 gives +Inf,
because the prompt kernel multiplies a masked future key's -Inf value by its probability of 0. Gate 1 plants no
nonfinite input in the fused prefix. The commands for an M5, after
zig build tf-kernel-contract -Dcpu=apple_m1:zig-out/bin/tf-kernel-contract --only 'metal qwen27 attention 27B prompt rows', the same for35B-A3B prompt rowsand bothdecode rows, thenzig build kernel-contract -Dcpu=apple_m1.Decode speed,
tools/bench_openai.py, 256 tokens, greedy, median of three: native 67.4 tok/s, mlx_lm.server85.6 tok/s on the same Mac. Four streams: 130-150 tok/s together.
Prompt speed, cold, first token on the same token arrays: native through
tools/bench_cache.py --require-coldwith
--prompt-cache-gib 0, mlx_lm in-process (stream_generate, a fresh cache each; its server takes notoken-id prompts):
The 35B has no earlier native release to compare with; 0.6.6's Mac figures are from an M3 Ultra.
Memory: the native server with 8 streams of 70,000 tokens, through the 8k cache and 8-stream runs, peaked at
32.9 GiB (
/usr/bin/time -l's peak footprint); mlx_lm peaked at 25.0 GiB of MLX memory at the 64k prompt.Tests:
zig build test test-kernel-contract test-golden fuzz-cortheon -Dcpu=apple_m1passes (with theuser-facing server stopped:
core/paced_read's test needs 8 GiB free); plainzig build;zig build test-native;zig build test-qwen27;git diff --check;tools/zig/lean_new.pyagainst the merge base: 517problems, 517 on the base, 0 new.
Not run: M1, M2, M3 and M5 chips and CUDA (one M4 Pro here); the 2B's checks (no 2B checkpoint here).
Open
3.6x at 64k. No speed gate for the 35B is agreed on Qwen3.6-35B-A3B (qwen3_5_moe) on the native Metal engine: proposal and gate #931; this PR is a draft until one is.
the qwen27 and moe_route entries were run with
--only. Both report 20 limit problems: 12 for Flash Next andeight for prefix_5/_6. The earlier
0 limit problemsreceipt above is pre-rebase, not a pass at the current tip.reservation in the source or a rule for tensor-unit-only launches.
Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.(The 35B drafts nothing yet; the 27B's drafted checks pass on both builds.)
gives
main's bytes at GLM's shapes, and GLM's model-level bits were not rerun.tf_route_logitsand theattention kernels; the attention kernels' M5 paths are not run here (above).