Repository navigation
Conversation
…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.
|
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 |
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.
|
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 |
2 participants
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.
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.
qwen3_5family, where Qwen3.6-35B-A3B on CUDA: routed experts and its MTP layer drafting chains (on #836) #856 puts the experts on CUDA: on Metal it is the 2B's engine, with no promptcache, drafter or image path, so slices 2 and 3 would be rebuilt there.
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; 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.31.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.
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).