Skip to content

qwen3_5_moe on CUDA: 3-bit routed experts beside affine-4 (on #856) - #972

Closed
CanadaApollo6 wants to merge 37 commits into
ashhart:mainfrom
CanadaApollo6:qwen36-moe-3bit
Closed

CanadaApollo6 wants to merge 37 commits into
ashhart:mainfrom
CanadaApollo6:qwen36-moe-3bit

Conversation

@CanadaApollo6

@CanadaApollo6 CanadaApollo6 commented Oct 16, 2026 •

Copy link
Copy Markdown

What this changes

3-bit routed experts for qwen3_5_moe on the CUDA engine, the width increment agreed on #803. It stacks on #856 (on #836): the last two commits are this pull request's.

  • cuda: 3-bit affine experts beside affine-4: new Affine3Experts in Grouped, beside Affine4Experts. MLX 3-bit codes (groups of 64) are stored as bit planes: a 2-bit plane (a uint4 a lane), a 1-bit plane (a uint2 a lane), and the same scales and biases. That is 224 words a block against 288. The kernels rebuild the nibble words affine-4 reads, so the mma operands, sums and roundings are affine-4's on the same codes. The pack widens the checkpoint's 3-bit words to 4-bit, runs the existing affine-4 pack, then splits the planes. One expert index (the shared expert) keeps 4-bit tables. A unit or CTA is one expert's column block, so the branch is uniform. The prompt stages fit 48 KiB.
  • qwen3_5: 3-bit routed experts on the CUDA engine: the geometry reads the quantization overrides.
    • Accepted: routers at 8 bits, and every layer's switch_mlp at 3 or 4 bits, one width throughout.
    • Refused (UnsupportedQwenQuantization): any other override that isn't the default 4/64, and any partial or mixed expert width.
    • A 3-bit checkpoint's experts pack 32 at a time through scratch in the existing stage. The shared expert and the MTP layer's experts stay 4-bit. 4-bit checkpoints load exactly as before.

Proposed gate

On #803 I proposed per-projection fp64 references, an engine-vs-reference KL bar no worse than 4-bit's, and the RTX 3090 contracts. Here is what this branch measures for each.

  1. Each projection against fp64 (tf-cuda-test experts3): gate, up and down at 3 bits, decode and prompt checked separately.
    • Each is no further from its fp64 product than affine-4 on the same codes, and equal to it to the last digit.
    • Decode worst column: 1.1e-6 (gate), 1.2e-6 (up), 6.7e-7 (down).
    • Prompt worst column: 4.1e-2, 3.7e-2, 3.1e-2. That is the prompt form's bf16 output rounding, and affine-4 shows the same figure.
    • Decode and prompt rows over 3-bit experts plus a 4-bit shared expert are byte-equal to the all-affine-4 tables on the same codes.
  2. The engine against a reference of the same checkpoint. The reference is the checkpoint with its 3-bit expert words rewritten as MLX 4-bit words (same codes, scales and biases), loaded through Qwen3.6-35B-A3B on CUDA: routed experts and its MTP layer drafting chains (on #836) #856's affine-4 path.
    • tf-qwen35-cuda over 300 ids: step, window and prompt logits from the 3-bit path are byte-equal to the reference's.
    • So the engine's KL against any fp32 reference of this checkpoint is the affine-4 engine's on the same weights, which meets the bar by construction.
    • I can still run a teacher-forced fp32 comparison at 1K/8K/32K if you want it on top.
  3. Contracts on an RTX 3090: below.

Receipt

  • Environment: RTX 3090 (24 GB, sm_86), Linux, driver 615.78, CUDA 13.4 nvcc, Zig 0.17.0, this branch at 967dd763.

  • Checkpoint: CanadaApollo6/Ornith-1.5-35B-A3B-M3TiS-3090 (Qwen3.5-MoE 35B-A3B, 3-bit routed experts, 4-bit elsewhere, an MTP layer). The receipts ran at 82f19799 with two names corrected locally. 66da2f4e publishes those names (rope_parameters.rope_type, and MTP tensors under language_model.mtp.), with the weights' bytes unchanged. This branch serves 66da2f4e as downloaded, with --drafter.

  • Exactness:

    • bench_concurrent.py --tokens 128 --reps 2 --alone --serial with --drafter <the checkpoint>, code and chat, T=1 seeded and T=0:
      • --levels 1,2 at --parallel 2 --context 8192: 36 of 36 equal;
      • --levels 1,4 at --parallel 4 --context 4096: 60 of 60 equal;
      • 0 unequal, 0 failed.
    • tf-qwen35-cuda over 300 ids: window and tree equal step byte for byte.
    • bench_cache.py --require-cold (2,571 tokens, drafter on): qualified, with repeated and concurrent first outputs equal.
  • Decode speed: bench_openai.py --tokens 128 --reps 3, median tok/s, one stream, code / chat.

    serve T=1.0 greedy
    --parallel 2 --context 8192, plain 107 / 108 111 / 110
    --parallel 2 --context 8192, MTP drafts 164 / 132 176 / 134
    --parallel 4 --context 4096, MTP drafts 165 / 129 176 / 131
    • Aggregate from bench_concurrent: up to 306 steady at 2 streams and 524 at 4.
    • tf-qwen35-cuda step: 8.9 ms a step.
    • These match Qwen3.6-35B-A3B on CUDA: routed experts and its MTP layer drafting chains (on #836) #856's 4-bit numbers on the same card. One-stream decode here is bound by launch and latency, not expert bytes.
    • What the width buys is memory: the model takes 14.87 GiB on the device against 18.62 GiB for the 4-bit checkpoint, leaving 8 streams at a 32K window.
  • Prompt speed: cold ttft_s 1.21 s on 2,571 tokens, against 1.24 s for the 4-bit checkpoint on Qwen3.6-35B-A3B on CUDA: routed experts and its MTP layer drafting chains (on #836) #856, same card.

  • Tests:

  • Not run:

    • Metal, since this touches only the CUDA engine.
    • Windows and sm_89, since I have neither. The kernels have no architecture-specific paths beyond affine-4's.

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. The 3-bit path's sums are affine-4's on the same codes, bit for bit.
  • Prompt processing is no slower than the last release.
  • 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.

- zig/src/platform: `sys`, imported by every module (zig/build/sys.zig wires
  it). On POSIX it is std.c, std.posix and std.DynLib unchanged; on Windows
  win_c.zig gives libc's surface over the UCRT and Win32 (integer
  descriptors, directory descriptors, readdir, pread/pwrite, high-resolution
  sleep, monotonic clock, /tmp and /proc/self/exe mapped), win_posix.zig the
  server's sockets over Winsock, console events as SIGINT/SIGTERM, and
  LoadLibraryW for DynLib.
- std.c and std.DynLib calls go through `platform` (mechanical).
- build.zig routes Windows targets to the CUDA build.
- Server: socket closes and writes are explicit, the key file's mode check
  is skipped on NTFS (ACLs), descriptor limits are a no-op.
- shard_edit: the shard's map is released before the rename, since Windows
  refuses to replace a mapped file.
- Tests that need symlink privilege or mode bits skip on Windows.

zig build test-cuda-host passes on x86_64-windows.
- Kernel fatbins build with nvcc and cl.exe: MSVC's conforming
  preprocessor is passed on Windows hosts (CUDA 13's headers need it),
  and explicit kernel instantiations carry the __restrict__ qualifiers
  their templates have (MSVC counts them in the function type; the
  generated code is unchanged).
- Nemotron's checkpoint source keeps io_uring to Linux; elsewhere its
  existing thread reads run.
- pull_exlr8 raises no descriptor limit on Windows.
- modelopt test: the lane's slice partials are allocated. GPUs without
  clusters (before sm_90) reduce split K through them; the test passed
  none, so it wrote through a null pointer on Ampere and Ada.

zig build native and install build with -Dnvcc=... -Dsm=89, and
tf-cuda-test's synthetic suites pass on an RTX 4090.
- pull: where Windows refuses the snapshot's symlink, the snapshot path
  becomes a hard link to its blob (CreateHardLinkW through `sys.c.link`):
  no extra space, and learning already refuses linked shards. Link
  failures are logged with their reason.
- HOME falls back to USERPROFILE (Windows sets no HOME) for the hub cache
  and every libc getenv("HOME").
- Tests read optional fixture variables through `platform.getenv`;
  std.process.Environ.getPosix does not build on Windows.
- The learned-state symlink tests skip on Windows (symlinks need
  privilege); `sys.c.symlink` asks for an unprivileged one.

zig build test passes on x86_64-windows.
- docs/WINDOWS.md: requirements, build, run, how the port works and how it
  differs from Linux; the README points to it.
- -Dstrip reaches tensorfold-native before its install step is made, so
  the step no longer expects a .pdb the stripped build does not write.
- zig fmt on the platform files.
8080 and 8000 are taken by too many other servers; serve listens on 8888
unless --port says otherwise, and the docs and bench scripts follow.
The MLX affine 4-bit checkpoint on the native CUDA server, beside Nemotron:
qwen35.cu holds the Metal engine's glue at its bf16 rounding points (gated
delta rule, causal convolution, gated attention), cuda.qlinear's affine-4
kernels do the projections, sample.cu draws on the device.

Windows of up to 16 rows record each row's recurrence and convolution tail,
so the lane core shares rounds across streams and keeps partial copies; a
window's rows equal one-row steps bit for bit. draw.zig is the host draw
both backends now share.

On an RTX 4090: 210 tok/s for one stream, 1250 tok/s across eight; logits
agree with an fp32 reference at every position of a 36-token check.
…mory

The 2B's CUDA engine now reads its widths from config.json (cuda_geometry.zig):
linear value heads over fewer key heads, more query heads a key/value head,
an untied head, an A_log stored as bf16. The 27B's logits agree with an fp32
reference at every position of a 36-token check, and its windows equal
one-row steps bit for bit; the 2B's output is unchanged.

Two Windows driver costs shaped the rest. Each kernel launch grew with the
process's allocation count (1.2 ms a launch over 1,900 weight buffers), so a
model's weights and scratch now live in one arena (cuda_weights.zig packs
each projection through a reused staging buffer: the 27B loads in 25 s, not
200). And a process's first forwards through a buffer stall for seconds, so
the engine warms every request path at load and the backend keeps a pool of
stream caches that `reserve` allocates and runs a row through once the
server knows its stream count.

On an RTX 4090: the 27B decodes at 37 tok/s in the runner, 31 to 35 tok/s
served, 210 tok/s across eight streams; the 2B stays at 210 and 1,250.
A forward's launches (the embedding to the kept rows' logits) are captured
the first time a shape (rows, windows, recording) is seen and replayed after;
the attention chunk count moves into the per-forward table so a graph stays
valid as caches grow. Prompt chunks stay uncaptured: their GEMMs dwarf the
launches, and each length would be a graph of its own.

A new graph's first run pays the Windows driver's stall (about a second on a
wide model), so `reserve` also captures the shared rounds' shapes, a row of
each of 1 to n streams, at load.

Bit for bit the same logits; on an RTX 4090 the 2B steps in 2.8 ms instead of
5.3 (served: 324 tok/s, 1,560 across eight streams) and the 27B in 23.8 ms
instead of 26.9 (served: 40 tok/s, 243 across eight).
--drafter loads z-lab's DFlash2 draft model behind the lane core's Drafted
wrapper: the target writes each row's residual after the drafter's five
tapped layers into a ring a stream, the drafter folds those rows into a
context of keys and values, and after every verify it drafts a block of
seven tokens past the pending one (the target's own embedding and head at
its ends, a codebook selector tracing one path through each row's best
sixteen). Its projections are quantized to affine-4 on the host at load and
run through the target's kernels; its glue is qwen35_draft.cu.

The Windows driver stalls a buffer's first kernel and a new graph's first
run for seconds, so the drafter pools its contexts and a captured forward
now serves every verify shape: its launches cover a bucket of rows and
every window, and the kernels read the true counts from the device.

On an RTX 4090 the 27B serves 52 to 58 tok/s greedy with drafts against 40
plain, about 100 tok/s on long replies, the same tokens either way; plain
decoding keeps its 24 ms step. The 2B is unchanged.
A verify window's rows may form a tree (each row's parent given): the
convolution taps, the delta rule and attention follow each row's own
branch, and the DFlash2 drafter fills the lane core's planned lanes
(each lane its parent's rank-r choice by the codebook selector), so a
lone stream verifies up to 12 lanes a round and the planner learns
which shapes land.

A verify no longer snapshots every row's state (1.3 GB shared, which
capped the 27B at 8 rows): each stream logs its rows' decay, write and
key, and keep replays the kept path onto the recurrence in one pass
(and moves a tree's kept keys down), so any window width fits. One-row
steps and prompts update in place as before. The lane host times the
target's windows at load, which the depth rule had lacked (it had been
capping chains at 3 drafts).

The runner's new tree mode verifies tree windows with junk branches
against step logits (bit-exact on both models).
What the port serves beyond Nemotron: Qwen3.5-2B and Qwen3.8-27B with
its DFlash2 drafter, the speeds measured on an RTX 4090, the load time
and why, the timing and chain-only switches, and what is not there.
The lean rule (tools/zig/lean_check.py) counts a comment that runs on
to a second line; the Windows port's new files had a hundred such runs.
Each is one line of at most 120 columns now, and the checker reports
the same 539 problems as main.
direct_io opened files with O_DIRECT on Linux alone. On Windows the
same open now asks for FILE_FLAG_NO_BUFFERING (which win_c maps O_DIRECT
to), and the 4 KiB-aligned spans it already reads satisfy the sector
alignment that needs; a file system that refuses it still falls back to
buffered reads. The direct-read test runs on Windows too, checking the
executable's bytes against buffered reads of it.
SIG.IGN is a function pointer made from 1, as std.posix has it. aarch64
requires function pointers to be aligned unless the type says align(1),
so the cross-build for aarch64-windows refused it. With align(1), as
std's own handler type has, tensorfold-native cross-compiles for
aarch64-windows with the GB10 (sm_121) fatbins; no hardware has run it.
The Windows page opens with what the build is, what it is tested on and
that it cross-compiles for aarch64-windows; the README's model table
and the CUDA paragraph say Windows where they said unofficial port.
The no-GPU workflow runs zig build test on windows-latest beside
ubuntu-latest, so the platform layer is built and tested on every pull
request.
Every other module calls libc through `sys`; the drafter's four getenv
calls were the last direct std.c calls in a file the Windows build
compiles.
After a prompt the lane core hands the drafter every prompt row's
features; the CUDA backend refused any request wider than its tap ring
(2,176 rows), so a 27B request with --drafter and a prompt past that
ended with RowsNotHeld, where the same prompt without the drafter ran.
The ring keeps the newest tap_rows rows, which hold the drafter's whole
context window, and absorb already reads only that window of what it is
handed; the backend now refuses only rows not yet written or all gone.

On an RTX 4090 the 27B with its drafter takes a 6,041-token prompt in
7.4 s, and its reply equals the one with "draft": false.
… attention cuts chunks at logical positions

- rowanc held a window's rows as bits of a u32, but prompt chunks run up to
  128 rows: rows from 32 on rotated at the wrong positions and attended the
  wrong keys. A chain row (no parents) now carries its index in its window
  under the top bit; tree windows (at most 16 rows) keep their bits.
- tf_q35_attn_partial cut its 128-key chunks at a key's cache row, so a kept
  tree path whose window crossed a chunk boundary summed in other chunks than
  one-row steps and changed bits. Chunks are now cut at the key's logical
  position (the stream's length plus the ancestors before it).

Checked on an RTX 5090 (Linux, sm_120), tf-qwen35-cuda over 300 ids on
Qwen3.5-2B and Qwen3.8-27B: step bits unchanged, window and tree equal step
byte for byte (tree differed from row 126 on the 2B and 254 on the 27B
before). Prompt mode's last row against step's at 16 to 300 ids: KL at most
0.003 and the same argmax in 20 of 20 (before: 27B at 40 ids KL 2.3 with a
different argmax, at 128 ids KL 1.35).
zig build test-golden and zig build fuzz-cortheon exist on every non-macOS
host since the golden step moved into zig/build/golden.zig, so the golden
check compiles against the platform layer on Windows:

- the platform's posix gains sockaddr.in, timeval, connect and kill for a
  child's handle (TerminateProcess; HUP has no equivalent), setsockopt turns
  a timeval for SO_RCVTIMEO or SO_SNDTIMEO into Winsock's milliseconds, and a
  socket read that times out is WouldBlock, as std.posix reports it.
- the check carries SystemRoot into the spawned server's environment, which
  Winsock needs to start, and drains the server's stdout after the port line
  on every host: the fake server logs each request there, and a full pipe
  (4 KiB on Windows) blocked it after sixteen cases.
- on Windows the check cannot send SIGTERM or SIGHUP, so the sigterm-exit-0
  outcome and the cases after a key rewrite are skipped and printed as such.

On Windows the goldens give 306 of 315 equal with the same nine known
differences as Linux, and the fuzz step passes its eight tests; the no-GPU
workflow's Windows job runs both.
The loader tried only libcublasLt.so.13 and libcublasLt.so, so the
cublaslt suite of tf-cuda-test reported LibraryUnavailable on Windows;
there it opens cublasLt64_13.dll from the toolkit's bin folder.
The platform handed O_APPEND to the C runtime as _O_APPEND, whose write
seeks to the end and then writes, so two threads appending at once could
overwrite each other's bytes: the request log lost lines under concurrent
requests, which the server goldens' request-log case caught on a slow
runner. A descriptor opened with O_APPEND now writes through WriteFile at
the end-of-file offset, which the kernel places whole, as POSIX does; the
file keeps its write access, so ftruncate on such a descriptor still works
(prompt_imprint repairs a torn index that way).

A request_log test appends 200 lines from each of 8 threads and expects
every line whole; it lost lines on Windows before this change.
The lifecycle/request-log outcome said only that the lines differ; it now
gives both logs' line counts and the first line that differs, cut to 100
bytes each, so a mismatch seen on a CI runner can be read from its log.
A checkout with core.autocrlf, as GitHub's Windows runners make, gives the
frozen request log CRLF line endings while the server writes LF, so the
lifecycle/request-log case differed at its first line's end on windows-latest
and nowhere else. Both logs compare with CRLF read as LF; the other goldens
are parsed as JSON and never saw the endings.
…'s widest stream

The chunk width came from the longest stream or window in the forward, so a 201-key row cut its keys into 128-key
chunks alone and 256-key chunks beside a stream past 8192, and a window crossing 8192 recut its earlier rows; either
changed the fp32 sums. Each row now takes chunk_keys(its logical keys), the width its one-row step uses, in both the
partial and the merge kernels, and the partials sit at a fixed max_chunks stride.

tf-qwen35-cuda takes TF_Q35_PREFILL=N (start after N cycled ids) and TF_Q35_BESIDE=N (step an N-token stream in each
forward), which cover both cases: before this change, windows from 8151 keys kept 36 of 120 rows and trees 35, and a
151-key stream beside one from 8189 kept 4 to 11 of 120; after it all of them equal one-row steps on the 2B and 27B.
…fits its selector rank

The integer reader accepted zero, so a conv_group_size or num_key_value_heads of 0 reached a remainder at
admission; sizes and divisors are now positive (mask_token_id alone may be 0). The host staging area held 16 rows at
rank 256, so a larger selector_rank overran it; it is now sized for the admitted rank. Config.parse takes the bytes,
so a host test covers both.
…ng with a shared expert

The grouped affine-4 expert kernel ran only Nemotron's relu^2 up; its SwiGLU
epilogue over two accumulators is now instantiated for decode and prompt rows
(M = 2, gate and up interleaved per group), and the pack kernel takes a table
stride and index so gate and up pack into one interleaved table. Layer.swiglu
picks that form; the shared expert sits last in the same tables.

cuda/moe.zig and moe.cu are the routing the frozen tensorfold/cuda/moe.py did:
fp32 logits over the experts and the shared expert's gate row, each row's top-k
by logit with ties to the lower id, weights exp(l - l_top) over the k rounded to
bf16, the shared expert as a slot weighted by bf16(sigmoid(bf16(gate))), and the
slots' weighted sum rounded once to bf16.

tf-cuda-test moe checks each against a host reference: the router and the
combine byte-equal, the picks equal with weights within a bf16 step (a tie, a
flat row and a NaN included), and the pair experts against fp64 over random
affine-4 tables.
model_type qwen3_5_moe serves on the same engine as the 2B and the 27B: the
layers are theirs, with routed experts in place of the dense MLP. The geometry
reads num_experts, num_experts_per_tok and moe_intermediate_size (the shared
expert's width must equal it, as it runs in the routed experts' tables). The
loader dequantizes the 8-bit router and shared-expert gate rows to bf16 on the
host, and packs each layer's 256 experts plus the shared one as a SwiGLU pair
table and a down table. A layer's MLP is a union of the dense projections and
the routed set; the forward routes, plans, runs the pair and down kernels and
combines into the residual, inside the captured graphs as before. Prompt
chunks past a window's rows take the prefill kernels, shorter ones the decode
form, so a short chunk keeps one-row steps' bits.

A drafter is refused (takes_drafter false): the checkpoint drafts with its own
MTP layer, which is not on CUDA yet. tf-qwen35-cuda gains a probe mode that
prints one layer's picks, weights and logits for a reference to compare.

On an RTX 4090, TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP: one-row steps decode
at 141 tok/s; windows and trees of up to 16 rows equal the steps byte for byte
over 300 ids; the 2B's and 27B's bits are unchanged. Against an fp32 PyTorch
reference the steps agree within the model's own rounding sensitivity (mean
KL 0.015 over 40 steps, where two references differing only by bf16 row
rounding differ by 0.016; the dense 2B agrees at 0.0005).
--drafter <the checkpoint's directory> loads mtp-4bit.safetensors beside the
weights: one attention-plus-experts layer over [norm(embedding of the next
token) | norm(the row's final normed state)], the target's head for its
draft, as the frozen engine's cuda/mtp.py ran it. The drafter taps the
target's last layer and applies the final norm itself; each stream keeps the
layer's keys and values over its whole window in a ring, absorbs the kept
rows with the token after each (the newest row's output is the first draft),
and holds a chain of up to TF_MTP_DRAFTS (default 3) more, each row the
previous output with the draft before it. Drafts only propose: the target's
windows verify them, so replies equal "draft": false ones.

qwen36_mtp.cu holds the attention at the target's head width (256) with its
partial rotary and output gate over the ring and the causal block; the ring
kernels are the DFlash2 drafter's. The expert tables and router rows pack
through the same loader as the target's, now writing into any arena.

On an RTX 4090, TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP with --parallel 2
--context 8192: one stream decodes 258 / 229 tok/s greedy (code / chat) and
252 / 181 sampled, against 101 / 98 plain; two streams 423 tok/s aggregate
on code. bench_concurrent --alone --serial: every drafted reply equals its
"draft": false reply and its lone run, 0 unequal. A chained draft costs
0.5 ms; a hold 1.2 ms, an absorb 0.7 ms.
BobClawblaw and others added 7 commits October 15, 2026 08:17
TF_MTP_TREES=1 lets the lane planner shape a tree of lanes (up to 12, the
ranks of a row's best tokens) for the MTP drafter: each depth's lanes run
as one pass of the layer, every row attending the ring and its own ancestor
lanes' keys (tf_q36m_attn_tree over a per-row ancestor table), their outputs
and keys kept by lane for their children, and each lane's token its parent's
rank-r candidate from the logprob rows. The load-time cost of a draft level
is now the quickest of sixteen holds, as the driver's stalls lasted a dozen.

Measured on an RTX 4090, Qwen3.6-35B-A3B, --parallel 2 --context 8192,
bench_openai median tok/s (code T=1 / chat T=1 / code greedy / chat greedy):
chains 248 / 197 / 271 / 235, trees 237 / 195 / 250 / 236; two streams 387
to 450 aggregate with chains, 381 to 439 with trees. Trees draft three times
the rows for the same tokens a round (88 of 430 lanes kept against 83 of
132 chained drafts), and the wider verify windows cost what the rounds save,
so chains stay the default. Replies equal "draft": false either way.
… the ring to its committed end

A tree pass packs a depth's sibling lanes as consecutive rows, and norm_rope rotated row m at pos0 + m, so a
depth's second lane took the next slot's rotation, and so did the key copied out of it and the child drawn from it.
The rope kernel now takes a row step: 1 for absorbed and chained rows, 0 for a tree pass, whose base stays the
pending slot plus the depth.

Tree attention read the ring up to that base, one slot per depth past the committed end; a tree pass writes no ring
row, so those slots could still hold an earlier round's keys. It now reads the ring to the context's end.
putIds wrote the pinned staging area and queued its copy without waiting, and absorb calls it once a prompt chunk,
so an absorb past 128 rows could overwrite a chunk's follow tokens before the GPU had read them. Replies stayed
correct (the target verifies every draft), but the drafter absorbed the wrong tokens after any long prompt. It now
synchronizes the stream before reusing the area; holds already had, through their fetch.
…estors root first

`tf-qwen35-cuda <MoE dir> mtp <ids> <out>` absorbs the ids from the target's tap ring and compares the drafter's
rows themselves, not the replies the target verifies: one absorb over several prompt chunks against the same rows a
chunk a call (ring, newest state, candidates); a tree with two siblings at depth 0, six at depth 1 and a depth-2
lane with two ancestors, run again with the ring poisoned at and past its committed end; and every internal lane's
state, key and value against its path chained one row a slot. All are byte-equal.

The tree's ancestor table now lists a lane's ancestors root first, the order a chain's ring holds them, so a lane's
fp32 sums equal its chained path's. Each check fails on the bug it covers: the sibling rows with the rotary step put
back to 1 (lane 1 at depth 0), the poisoned ring with the tree's ring end put back to its rotary base (lane 3 at
depth 1), the chunked absorb without the previous commit's wait (4 runs of 4).
ashhart#853's host builds each family's options from F.Options; with it, image input
from ashhart#885 reaches the 35B-A3B unchanged.
…s 4-bit

MLX 3-bit experts (groups of 64) as bit planes the kernels widen back to the
nibble words affine-4 reads, so a code's products, sums and roundings are the
grouped affine-4 kernel's. A block is a 2-bit plane (a uint4 a lane), a 1-bit
plane (a uint2 a lane) and the same scales and biases: 224 words against 288.
The pack widens the checkpoint's 3-bit words to 4-bit ones, packs them with
the affine-4 pack, then splits the blocks into planes.

The decode and prompt kernels take one expert index whose tables stay 4-bit
(the shared expert); a unit or a CTA is one expert's column block, so the
branch is uniform. The prompt stages fit 48 KiB: four stages of 3-bit blocks
or three of 4-bit at M=1, three or two at M=2.

tf-cuda-test experts3: decode and prompt rows over 3-bit experts and a 4-bit
shared one byte-equal to affine-4 on the same codes, and each projection
(gate, up, down) against fp64 no further than affine-4 on the same codes.
The geometry reads the quantization's per-module overrides: the routers at 8
bits, every layer's switch_mlp projections at 3 or 4 bits (one width for all),
any other at the default 4, groups of 64. A 3-bit checkpoint's routed experts
load into 3-bit tables a chunk of 32 experts at a time, through scratch in the
stage beside the raw parts; the shared expert keeps 4-bit tables of its own,
and the MTP layer's experts stay 4-bit. The forward hands both formats to the
grouped experts, which run the layer's.
@BobClawblaw

Copy link
Copy Markdown
Contributor

@CanadaApollo6 The Windows and sm_89 receipt this PR lists as not run: RTX 4090 (24 GB, sm_89, WDDM), Windows 11 build 26300, driver 616.56, CUDA 13.3 nvcc (V13.3.73), Zig 0.17.0, -Dsm=89. This branch at 967dd763, CanadaApollo6/Ornith-1.5-35B-A3B-M3TiS-3090 at 66da2f4e as downloaded. Every check passes.

Tests

  • zig build test native -Dsm=89, fresh cache: 95/95 steps, 745 of 759 tests passed, 14 skipped, 0 failed.
  • tf-cuda-test experts3: 8 of 8. The 48 KiB prompt stages run on sm_89. Rows over 3-bit experts with the 4-bit shared expert are byte-equal to affine-4 on the same codes, and the fp64 error is affine-4's: decode worst column 1.14e-6 / 1.18e-6 / 6.7e-7 (gate / up / down), prompt 4.13e-2 / 3.71e-2 / 3.11e-2, the same figures as your 3090.
  • moe 4 of 4, glue 1,242 of 1,242.

Exactness

  • tf-qwen35-cuda over 300 ids: window and tree equal step byte for byte. 7.4 ms a step.
  • bench_concurrent.py --tokens 128 --reps 2 --alone --serial with --drafter <the checkpoint>: --levels 1,2 at --parallel 2 --context 8192 gives 36 of 36 equal; --levels 1,4 at --parallel 4 --context 4096 gives 60 of 60. 0 unequal, 0 failed.

Speed (bench_openai.py --tokens 128 --reps 3, median tok/s, one stream; code / chat)

serve T=1.0 greedy
--parallel 2 --context 8192, plain 149 / 150 155 / 154
--parallel 2 --context 8192, MTP drafts 237 / 193 252 / 198
--parallel 4 --context 4096, MTP drafts 213 / 180 227 / 187
  • bench_cache.py --require-cold, 2,571 tokens with the drafter: cold ttft_s 0.704 s at --parallel 2, 0.712 s at --parallel 4, qualified.
  • Device memory: 14.87 GiB, 15.35 GiB with the drafter.

Control, #856's 4-bit checkpoint on the same build: plain decode 154 / 152 at T=1 and 156 / 155 greedy, 6.7 ms a step, 18.62 GiB on the device. So on Ada the 3-bit width decodes as fast as affine-4 and frees 3.75 GiB, as on your 3090. (#856's description has 101 / 98 plain for this card; that was an earlier build.)

@ashhart

ashhart commented Oct 17, 2026

Copy link
Copy Markdown
Owner

Thank you, @CanadaApollo6, the separate affine3 operator and widened-same-code control make the arithmetic comparison useful.

We checked 967dd76. cuda_geometry.zig:111-120 ignores override mode and accepts unrelated switch_mlp keys toward the expected count. We will validate affine mode and the complete expected layer/projection key set, with malformed controls, then add resource contracts for the new kernels.

This follows #836, #856 and owned images through #885. We can fold the small admission fixes in at landing, preserving the affine-4/shared/MTP tables. Please retain the 300-ID, FP64, cache and RTX3090 comparisons for the final integrated candidate, with long-window/image gates and original quantisation loss reported separately from same-code engine equality.

Posted by Ash's AI Assistant Synthena

@ashhart

ashhart commented Oct 18, 2026

Copy link
Copy Markdown
Owner

Thank you, @CanadaApollo6, for #972 and its 300-ID widened-same-code control. 1.1.0 integrates 3-bit routed experts beside the 4-bit shared/MTP paths on #856/#836. The maintained admission checks validate affine mode and the complete expected layer/projection key set before launch; the new operators retain their contracts. After final numerical/resource, same-checkpoint long-window, image and draft/cache/concurrency gates pass and it ships, this closes the PR with your credit. Quantisation quality remains distinct from engine equality.

GB10 MoE serving passed 17/17 and chat 12/12; the specific mixed-width 300-ID/oracle/reference panel is not separately recorded.

Posted by Ash's AI Assistant Synthena

@ashhart ashhart closed this Oct 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants