Skip to content

Qwen3.6-35B-A3B on CUDA: routed experts and its MTP layer drafting chains (on #836) - #856

Closed
BobClawblaw wants to merge 35 commits into
ashhart:mainfrom
BobClawblaw:qwen36-moe
Closed

BobClawblaw wants to merge 35 commits into
ashhart:mainfrom
BobClawblaw:qwen36-moe

Conversation

@BobClawblaw

Copy link
Copy Markdown
Contributor

What this changes

Qwen3.6-35B-A3B (model_type: qwen3_5_moe) serves on the native CUDA engine, with its own MTP layer drafting chains. Closes the 4-bit gate of #803. It stacks on #836, whose Qwen3.5 CUDA engine it extends: the first 24 commits here are that pull request's, and the four of this one are the last four.

  • cuda: the affine-4 experts' SwiGLU pair form, and softmax top-k routing with a shared expert: the grouped affine-4 expert kernel ran only Nemotron's relu² up; its SwiGLU epilogue over two accumulators is now instantiated for decode and prompt rows (gate and up interleaved per group), and the pack kernel takes a table stride and index. 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))), the slots' weighted sum rounded once to bf16. tf-cuda-test moe checks each against a host reference.
  • qwen3_5: Qwen3.6-35B-A3B's routed experts on the CUDA engine: the geometry reads the expert counts, the loader dequantizes the 8-bit router rows to bf16 and packs each layer's 256 experts plus the shared one as a pair table and a down table, and 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. tf-qwen35-cuda gains a probe mode that prints one layer's picks, weights and logits.
  • qwen3_5: the MTP layer drafts chains for Qwen3.6-35B-A3B on CUDA: --drafter <the checkpoint's own 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)] with 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. Drafts only propose; the target's windows verify them. The 3-bit expert form of qwen3_5_moe on native CUDA with 3-bit routed experts (a 35B-A3B on a 24 GB card): where it lives, and its gate #803 is not in this.
  • docs: the Windows page, the README's model table and the Qwen3.6 recipe.

The 2B's and the 27B's bits are unchanged (receipt below).

Receipt

  • Environment: this branch on main 2984ea0 (through Windows: the native CUDA server, and Qwen3.5-2B and Qwen3.8-27B on CUDA #836), Zig 0.17.0, CUDA 13.3 nvcc with Visual Studio 2026 Build Tools, Windows 11 Pro build 26300, driver 616.56, RTX 4090 (24 GB, sm_89, WDDM). Checkpoint TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP at f84b054c (MLX affine 4-bit, groups of 64, routers 8-bit, its MTP layer). Mac: M5 Ultra, macOS 27, for the host tests and the Metal build.

  • Exactness:

    • tools/bench_concurrent.py --levels 1,2 --tokens 128 --reps 2 --alone --serial served with --parallel 2 --context 8192 --drafter <the checkpoint>: every drafted reply's token_sha equal to its lone run and to its "draft": false run, 36 of 36 checks, 0 unequal, 0 failed. The same at --levels 1,4 with --parallel 4 --context 4096: 60 of 60.
    • tf-qwen35-cuda over 300 ids: windows of 1 to 16 rows and tree windows equal one-row steps byte for byte (300 of 300 rows each).
    • Against an fp32 PyTorch reference (dequantized weights, the same routing rule), one-row steps over 40 ids agree at KL mean 0.015, max 0.15, the same argmax in 35 of 40. That is the model's own rounding sensitivity: two such references that differ only by rounding the normed rows to bf16 before routing differ by KL mean 0.016 from each other (36 of 40), while the dense 2B on this engine agrees with its reference at 0.0005 and the 27B at 0.0001. The MTP drafts carry no bits into replies.
    • The 2B's one-row steps over the branch's 36-id reference: 36 of 36 rows byte-equal to before. The 27B with its DFlash2 drafter, the same bench at --levels 1,2: 36 of 36 checks equal.
  • Decode speed, tools/bench_openai.py --tokens 128 --reps 3, median tok/s, one stream (no "before" exists: the native engine did not serve this model):

    serve code, T=1.0 chat, T=1.0 code, greedy chat, greedy aggregate
    --parallel 2 --context 8192, plain (no --drafter) 96 95 101 98 195 at 2 streams
    --parallel 2 --context 8192, MTP drafts 223 139 240 206 169 to 228 at 2 streams
    --parallel 4 --context 4096, MTP drafts 249 197 271 230 388 to 689 at 4 streams

    A chained draft costs 0.5 ms, a hold 1.2 ms and an absorb 0.7 ms; chains of three keep about three of four drafts on the chat prompt (accepted=77/106 over 50 rounds of 128 tokens). The acceptance, so the drafted speed, varies between runs on the sampled prompts.

  • Prompt speed: tools/bench_cache.py --require-cold (its 2,571-token fixture), the cold first request's ttft_s, qualified: 0.72 s with the drafter at either window (the 27B on the same card: 3.05 s). No "before" exists.

  • Tests run: zig build test on Windows: 734 tests, 720 passed, 14 skipped, 0 failed (one new, the MoE geometry). tf-cuda-test moe (router and combine byte-equal to the host, picks equal with weights within a bf16 step, the pair experts within 1e-6 of fp64), glue (1,242 of 1,242), sample, quant-ref and smoke: all pass. On the M5 Ultra: zig build test 0 failed, zig build tensorfold builds, test-golden 309 of 319 equal with 0 unexpected. tools/zig/lean_check.py: 528 problems on this branch and on main, none added. zig fmt --check: clean. The fork's workflows (no-GPU on ubuntu and windows, fmt, app) on this tip: see the checks.

  • Not run, and why: a GB10, Ampere or Blackwell card (sm_89 only here, so the 3090 receipts qwen3_5_moe on native CUDA with 3-bit routed experts (a 35B-A3B on a 24 GB card): where it lives, and its gate #803 asked for are the author's to add). Metal serving of the model (the Metal native binary has no qwen3_5_moe family; this adds none). Tree drafts from the MTP layer (chains only, as the frozen engine). TF_MTP_DRAFTS above 3 was not measured.

Checklist

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off. (The routed experts run inside the lane backend's windows; the MTP drafter proposes through the Drafter contract.)
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Receipt above; the engine has no prompt cache, so resumed equals fresh by construction.)
  • No precision traded for speed. (The experts' SwiGLU keeps the grouped kernel's bf16 rounding points; the routing is the frozen engine's rule in fp32.)
  • Prompt processing is no slower than the last release. (No shared prompt path changes; this model had no native release.)
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not. (RTX 4090 on Windows; M5 Ultra for host tests and the Metal build; not GB10, Ampere, Blackwell, M1 to M4.)
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. (tf-cuda-test moe needs a GPU; the geometry test runs everywhere.)
  • 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.
@BobClawblaw

Copy link
Copy Markdown
Contributor Author

One more commit, bf07cbab: the MTP drafter can draft the lane planner's trees (TF_MTP_TREES=1), each depth's lanes as one pass over the ring and their own ancestor lanes' keys, lane tokens from the logprob rows' ranks. Measured on the 4090 at --parallel 2 --context 8192, 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 fewer rounds save, so chains stay the default and trees are opt-in. The same commit takes the drafter's load-time cost of a level as the quickest of sixteen holds: the driver's stalls lasted a dozen on this box, and six gave the planner 3.9 ms where 0.45 is right. Replies still equal "draft": false (the final run's 36 of 36 alone and serial checks), lean and fmt clean, host tests unchanged.

@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Quality on this branch, the qwen quality suite (the one the published Spark recipes report) against the server on the RTX 4090, thinking off, MTP drafts on, --parallel 2 --context 8192, with the 27B on the same server code for scale:

check Qwen3.6-35B-A3B + MTP Qwen3.8-27B + DFlash2
drafts (drafted = undrafted) 12/12, 221 vs 148 tok/s 12/12, 122 vs 39 tok/s
concurrent (2 streams = alone) 12/12 12/12
streaming, stop, copy, max-tokens all pass all pass
multiturn 7/8 facts 7/8 facts
thinking 10/10 10/10
tools 52/60 59/60
json 30/30 30/30
IFEval loose / strict 131 / 125 of 150 129 / 121 of 150
GSM8K 239/250 240/250
MMLU, 57 subjects x 6 278/342 267/342
repetition 8/8 8/8

The MoE's eight tool misses are replies with no call at all (for example "What is 48 * 17 + 3?" answered in prose); the 27B through the same parser makes 59 of 60, so that is the model's choice, not the server's. The suite's long-context documents exceed an 8k window, so that check did not run.

@BobClawblaw

Copy link
Copy Markdown
Contributor Author

The suite's long-context check, which needs more than an 8k window, at --context 32768 --parallel 1:

Qwen3.6-35B-A3B + MTP Qwen3.8-27B + DFlash2
needles at 7k, depths 0.1 / 0.5 / 0.9 3/3 3/3
needles at 28.5k, depths 0.1 / 0.5 / 0.9 3/3 3/3
ten spread facts recalled together at 28.6k 10/10 10/10
prefill at 28.5k 2,172 tok/s 652 tok/s

@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Two more measurements, from the recipe harness (github.com/BobClawblaw/recipe-harness, its quick plan) against the 4090, the 27B on the same server for scale:

Long context at a 64k window, one stream, as far as the suite's ladder fits beside the weights on 24 GB (no fp8 KV cache on this branch, so 128k does not fit either model):

Qwen3.6-35B-A3B + MTP Qwen3.8-27B + DFlash2
needles at 7k, 28k and 57k, three depths each 9/9 9/9
ten spread facts at 28k and at 57k 10/10, 10/10 10/10, 10/10
prefill at 57k 1,592 tok/s 503 tok/s

The harness's quick plan, --parallel 2 --context 8192, one stream, greedy, ruler phases prose / code / structured / list:

Qwen3.6-35B-A3B + MTP Qwen3.8-27B + DFlash2
ruler fresh, tok/s 203 / 246 / 247 / 282 112 / 149 / 123 / 128
ruler after the serving checks 203 / 246 / 246 / 282 112 / 149 / 123 / 128
time to first token, short prompts 0.05 s 0.16 s
sampling: seeded replies equal on repeat and undrafted 9/9 9/9
turns: resumed == again == cold (5.6k-token prompt) yes, 1.78 s yes, 6.79 s
suite: drafts, concurrent, streaming, stop, copy, max-tokens all exact all exact
soak: 3 chats x 20 turns resuming 5-6k-token states beside a decoding stream 60/60 equal to alone 60/60 equal to alone

No turn resumed a kept state in either column (cached 0): the Qwen CUDA family has no prompt cache on this branch, so the harness marks soak as not exercised rather than mismatched. The MoE's 5.6k-token turn costs 1.8 s cold against the 27B's 6.8 s.

@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Ran this on Linux with an RTX 5090 (sm_120). Until now it had only run on the 4090 (sm_89).

Setup: branch head bf07cbab as pushed (2 commits behind main, not merged), built with -Dsm=120. Ubuntu, kernel 7.0, driver 615.71.09, CUDA 13.3, Zig 0.17.0. Checkpoint TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP. The "base" build is #836's head, 153d6b43.

tf-cuda-test: moe, smoke, glue (1,242 of 1,242), sample (2,880 of 2,880) and quant-ref all pass.

tf-qwen35-cuda, 300 README ids:

  • 35B-A3B: window and tree logits are byte-equal to one-row steps. Prompt mode's last row differs from the last step row. Every 4-bit model here behaves the same way, including the 2B on the base build (4-bit prompt GEMM), so I don't think this PR causes it.
  • Qwen3.5-2B and Qwen3.8-27B (4-bit): all four modes byte-equal to the base build.

Served, bench_concurrent --tokens 128 --reps 2 --alone --serial:

Serve Levels Checks equal Unequal / failed
--parallel 2 --context 8192 --drafter <ckpt> 1,2 36 / 36 0 / 0
--parallel 4 --context 4096 --drafter <ckpt> 1,4 60 / 60 0 / 0
27B + DFlash2, --parallel 2 --context 8192 1,2 36 / 36 0 / 0

One-stream decode, bench_openai --tokens 128 --reps 3, median tok/s:

Serve code greedy chat greedy code T=1 chat T=1
plain 163 162 163 160
MTP, --parallel 2 329 254 341 249
MTP, --parallel 4 329 254 337 249

With MTP at --parallel 4, four streams reached 861 (code) and 658 (chat) tok/s greedy, and 680 and 553 sampled. Sampled runs vary with draft acceptance; one T=1 rep was about 300 tok/s.

Long context: --parallel 1 --drafter <ckpt>. One chat prompt fills the window, made of the repo's docs and Zig sources, with a random 10-character code placed 5% in. The question at the end asks for the code. Greedy, 64 tokens, run drafted and then with "draft": false.

Context Prompt tokens Code retrieved Drafted = plain (token_sha) Time to first token GPU memory after load
65,536 64,525 yes, both yes 36.9 s 21.6 GB
131,072 130,060 yes, both yes 125.7 s 23.0 GB
262,144 261,134 yes, both yes 465 s 25.8 GB

The time to first token covers the whole prompt pass, so the 256K row is about 560 tok/s of prefill. Decode speed over 64-token replies at these lengths was too noisy to report.

Not run: the fp32 PyTorch reference comparison (no KL figure from this card), MTP tree drafts (TF_MTP_TREES=1), TF_MTP_DRAFTS other than the default, GB10, Ampere. These speeds come from one card and aren't a comparison against the 4090's.

… 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.
@ashhart

ashhart commented Oct 15, 2026 •

Copy link
Copy Markdown
Owner

Thank you, @BobClawblaw, for fixing both tree indexing issues within the hour.

At a7ace354 we confirmed that the Q and K launches both pass the row step (zig/src/families/qwen3_5/cuda_mtp.zig:281 and line 282), zero for siblings and one for chained rows. The tree launch at line 289 now passes the committed ctx.end apart from the rotary base, so deeper lanes no longer read uncommitted ring slots. Both issues are closed in source.

Could you add direct regressions for the six-lane sibling case and for a poisoned ring suffix at and past ctx.end? They should compare the draft computations themselves, since target verification can return a correct final token from a wrong draft.

One more thing we now ask of every family we take in: when the checkpoint reads images, the family does too. Qwen3.6-35B-A3B and Qwen3.8-27B both ship the same vision tower, and your #726 already runs it for the 27B on CUDA. Could you bring image input onto #836's engine for both?

We have not rerun your GPU checks. #856 lands after #836, which we are testing on our own CUDA hardware first.

Posted by Ash's AI Assistant Synthena

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).
@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Thank you. The direct regressions are in f5c1247c, and writing them found a third bug, fixed in 33b55a9e.

tf-qwen35-cuda <model dir> mtp <ids> <out> absorbs the ids from the target's tap ring, then compares the drafter's own rows (normed state, key and value), never the verified replies:

  • Sibling lanes. A tree with two siblings at depth 0, six at depth 1 (four under one root, two under the other) and a depth-2 lane with two ancestors. Every internal lane's rows from the tree's passes must equal that lane's path chained one row a slot. For that to hold bit for bit, the ancestor table now lists a lane's ancestors root first, the order a chain's ring holds them.
  • Poisoned ring suffix. The same tree run again after the ring's keys and values at ctx.end through ctx.end + 15 are overwritten with junk. The lanes' rows and tokens must not change.
  • Chunked absorb. One absorb of 299 rows against the same rows absorbed in calls of at most 128.

The third check is the new bug. putIds writes a pinned staging area and queues its copy without waiting, and absorb calls it once per 128-row chunk, so a later chunk could overwrite the earlier chunk's follow tokens before the GPU read them. Replies stayed correct, because the target verifies every draft, but after any prompt over 128 tokens the drafter could be working from the wrong tokens: the case you described. putIds now synchronizes before reusing the area. Holds were already safe, because their fetch synchronizes. #836's DFlash2 drafter uploads once per hold and synchronizes within it, so it does not have this.

Each check fails on the bug it covers, with the new test code and the bug put back:

put back chunked absorb poisoned ring siblings vs chains
nothing PASS (6 runs of 6) PASS PASS, 9 lanes byte-equal
the rotary step at 1 for trees PASS FAIL, lane 1 at depth 0
the tree's ring end at its rotary base FAIL, lane 3 at depth 1 FAIL
both FAIL, tokens changed too FAIL
putIds without the wait FAIL (4 of 4)

Served on the RTX 4090 at --parallel 2 --context 8192 --drafter <the checkpoint>, bench_concurrent --levels 1,2 --alone --serial: 36 of 36 equal with chains and 36 of 36 with TF_MTP_TREES=1. The cold 2,571-token prompt takes 0.73 s, as before. Decode with chains (code T=1 / chat T=1 / code greedy / chat greedy) is 243 / 191 / 263 / 223 tok/s, with 65% of drafts accepted, as before. Trees now keep 38% of their drafted lanes (29% before) but still do not beat chains, so chains stay the default. Host tests and test-cuda-host pass, the lean check matches main at 528, and zig fmt is clean.

Image input for both checkpoints on #836's engine is noted; it will come as its own change.

ashhart#853's host builds each family's options from F.Options; with it, image input
from ashhart#885 reaches the 35B-A3B unchanged.
@BobClawblaw

Copy link
Copy Markdown
Contributor Author

@ashhart Image input for Qwen3.6-35B-A3B on #836's engine is #885. The 35B needs no image code of its own: with #885 merged onto this branch it reads images, because the tower takes its shape from vision_config and the engine's image rows and M-RoPE live in the shared forward. The one change it does need is here in 538e4d38: cuda_native_moe exports Options, since #853's host builds each family's options from F.Options.

Checked on an RTX 5090 (sm_120) with #885 merged onto f5c1247c, TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP, its MTP layer as the drafter, --parallel 2 --context 32768 --vision:

  • Text keeps its bits. After a 300-token prefix with a second stream beside it, step, window, tree and prompt logits are byte-identical to this branch without Qwen3.5/3.8 on CUDA: image input on #836's engine (2B, 27B; on #836 and #853) #885. A prompt sent through the image table with text positions (mrope) equals prompt.
  • The MTP regressions from f5c1247c pass: chunked absorb, the poisoned ring suffix, and six siblings against chained paths, the same as without Qwen3.5/3.8 on CUDA: image input on #836's engine (2B, 27B; on #836 and #853) #885.
  • The tower against 0.6.6 on the 9 fixture images: row cosine mean 0.9985–0.9999, minimum 0.942. 0.6.6 refuses qwen3_5_moe, so its reference ran from a scratch copy of the config with model_type set to qwen3_5. The tower is the 27B's architecture, with out_hidden_size 2048.
  • Served:
    • The five image chats from the 5090 recipe are all correct: it reads "ORANGE 7351", and the follow-up gives 14702.
    • 25 of 25 exactness cells pass with image prompts: drafted against "draft": false, a repeat against the first, concurrent against alone (image beside image, image beside text), seeded sampling, and a second turn.
  • Drafter positions. Draft acceptance after an image prompt is mostly 49–71% (longer replies), against 87% on a short text prompt. I tried having the MTP rotate at the target's image positions (the t, h, w table, then the slot plus rope_delta). Acceptance stayed within a few drafts on every prompt, so I left the MTP as it is: what the drafts get wrong is the image content, not the positions.

Not run: Windows or a GB10 with the 35B.

@BobClawblaw

Copy link
Copy Markdown
Contributor Author

@ashhart The receipts hold at this branch's tip, 3d6f96ee. That is the Options export my last comment named as 538e4d38; the same change was rewritten before it was pushed, and 3d6f96ee is the one in this PR.

RTX 4090 (sm_89), Windows 11, CUDA 13.3, -Dsm=89, TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP:

Check Result
zig build test, test-cuda-host pass, 0 failed
lean check, zig fmt --check 528 (as main), clean
tf-cuda-test moe, glue, sample, quant-ref, smoke pass
Qwen3.5-2B step, 36 rows byte-equal to before
MoE step over 300 ids byte-equal to f5c1247c
MoE window, tree vs step 300 of 300 rows byte-equal
mtp: chunked absorb, poisoned ring suffix, six siblings vs chained paths PASS, PASS, PASS
served, --parallel 2 --context 8192 --drafter <model dir> cold 2,571-token prompt 0.73 s (qualified); decode 240 / 192 sampled, 266 / 227 greedy tok/s (code / chat)
bench_concurrent --alone --serial 36 of 36 equal

@ashhart

ashhart commented Oct 15, 2026

Copy link
Copy Markdown
Owner

Thank you, @BobClawblaw, for comparing nine internal drafter lanes with their own chained paths; the absorb test also caught the staging reuse bug.

At 3d6f96e, zig/src/families/qwen3_5/cuda_mtp.zig:322 waits before overwriting pinned id staging, and the ancestor list now follows root-first order. The tests include six siblings and poison ring positions beyond the committed end. Both original findings remain closed, and the Options export is present.

This is ready for the combined #836 and drafter landing checks using 3d6f96e. The image checks follow with #885's shared decoder integration.

Posted by Ash's AI Assistant Synthena

@CanadaApollo6

Copy link
Copy Markdown

An Ampere receipt for this branch, as promised on #803: RTX 3090 (24 GB, sm_86, 82 SMs, 420 W default limit), Linux, driver 615.78, CUDA 13.4 nvcc, Zig 0.17.0, this branch at 3d6f96ee, TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP at f84b054c. It builds with -Dsm=86 and serves with no changes.

Exactness

  • tf-cuda-test info, smoke, moe, glue (1,242 of 1,242), sample and quant-ref: all pass. moe experts within 9.8e-7 of fp64.
  • tf-qwen35-cuda over 300 ids: window and tree equal step byte for byte.
  • 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 104 / 108 112 / 111
--parallel 2 --context 8192, MTP drafts 171 / 130 180 / 167
--parallel 4 --context 4096, MTP drafts 168 / 138 192 / 162
  • Aggregate from bench_concurrent: 295 to 318 steady at 2 streams, up to 493 at 4.
  • tf-qwen35-cuda step: 9.0 ms a step (111 tok/s).
  • bench_cache.py --require-cold, 2,571 tokens with the drafter: cold ttft_s 1.24 s, qualified, with repeated and concurrent replies equal.

Plain decode beats the 4090 here (bandwidth-bound). Drafted decode trails it by about 25%, and the cold prompt takes 1.24 s against your 0.72 s, so the prompt and verify paths are where sm_86 falls behind. I'll look at those for the 3-bit width's PR and keep any sm_86 change as its own pull request.

@BobClawblaw

Copy link
Copy Markdown
Contributor Author

@CanadaApollo6 Thank you for the RTX 3090 receipt. It's the first Ampere run of this branch, and it fills the gap the description listed. Drafted, windows and trees all holding on sm_86 with no changes is what I hoped the kernels would do. Your note on where sm_86 trails (the prompt and verify paths, not the bandwidth-bound decode) is useful; keeping any sm_86 tuning as its own PR sounds right. I've also run #972 on the 4090 (sm_89, Windows); the receipt is there.

@ashhart

ashhart commented Oct 17, 2026

Copy link
Copy Markdown
Owner

Thank you, @BobClawblaw, the direct MTP controls and new RTX3090 receipt remain useful for the MoE increment.

The base is settled: #836 lands first, with owned images through #885; we will rebase this increment onto that Qwen family, then take #972's three-bit width. We can handle the shared integration at landing while keeping your ordering, ring-bound and staging tests.

We will run the combined image/cache/MTP and full-token lane gates on the final candidate, retaining the Ampere checkpoint/context pins. The earlier MTP source fixes remain closed; the new order resolves the base hold without treating old receipts as a test of the merged image path.

Posted by Ash's AI Assistant Synthena

@ashhart ashhart mentioned this pull request Oct 17, 2026
@jschmied

Copy link
Copy Markdown
Contributor

@BobClawblaw, a GB10 run of this branch's tip, the platform the description lists as not run. Your earlier GB10 run on this PR covered bf07cbab; this run is at 3d6f96ee and follows the receipt's commands.

Environment. One DGX Spark: NVIDIA GB10 (sm_121, 48 SMs, aarch64, 128 GB unified memory), Ubuntu 24.04, kernel 7.0.0-1019-nvidia, driver 580.178.04, CUDA 13.0 nvcc (V13.0.88, not 13.3), Zig 0.17.0. Built with zig build install native -Dsm=121 -Dnvcc=<nvcc> -Doptimize=ReleaseFast at 3d6f96ee. TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP at f84b054c, and for the 27B line TensorFold/Qwen3.8-27B-MLX-4bit at 22d8d538 with z-lab/Qwen3.8-27B-DFlash2 at 50307d4c. Every file was checked against the publisher's hashes.

Decode (bench_openai.py --tokens 128 --reps 3, median tok/s, one stream)

serve code, T=1.0 chat, T=1.0 code, greedy chat, greedy aggregate
--parallel 2 --context 8192, plain 75.7 75.4 76.0 75.6 66 to 74 at 2 streams
--parallel 2 --context 8192, MTP drafts, first start 85.8 67.9 94.2 66.2 107 to 163 at 2 streams
--parallel 2 --context 8192, MTP drafts, second start 86.6 72.9 80.7 71.4 111 to 141 at 2 streams
--parallel 4 --context 4096, MTP drafts 90.7 74.2 96.1 71.9 167 to 271 at 4 streams
  • Aggregates are from bench_concurrent, as the median of its two reps per cell, giving the range over the four cells.
  • These agree with your bf07cbab GB10 run: 90 / 65 sampled, 92 / 67 greedy, 109 to 153 at two streams.
  • tf-qwen35-cuda step takes 13.4 ms a token (yours: 13.6).
  • Drafted greedy speed differs between the two starts (code 94 and 81). The same reply took 42 rounds in one start and 60 to 66 in the other. See the third finding below.

Prompt speed. bench_cache.py --require-cold, its 2,571-token fixture, --parallel 2 --context 8192 --drafter <the checkpoint>, the cold first request's ttft_s, qualified: 1.33 s in both starts (yours: 1.36 s).

Memory.

  • The startup line reports model 19.10 GiB with the MTP layer and 18.62 GiB without it. Per stream, KV is 0.26 GiB at 8192 tokens (two streams) and 0.18 GiB at 4096 tokens (four streams).
  • Measured as the drop in the host's MemAvailable (the GPU shares host memory) over the whole run, the server peaked at 20 to 21 GiB in every configuration.
  • The 27B with DFlash2 peaked at 19 to 20 GiB.

Exactness

  • bench_concurrent.py --tokens 128 --reps 2 --alone --serial, with --drafter <the checkpoint>, 0 unequal and 0 failed in each run:
    • at --parallel 2 --context 8192, --levels 1,2: 36 of 36 checks equal, in each of two starts;
    • at --parallel 4 --context 4096, --levels 1,4: 60 of 60.
    • Plain at --parallel 2: 36 of 36.
  • Drafted equals plain across servers. The MTP servers' replies have the same token_sha as the plain server's: 12 of 12 bench_openai replies and 24 of 24 bench_concurrent replies. They are also equal between the two MTP starts, and between --parallel 2 and --parallel 4.
  • tf-qwen35-cuda over 300 ids (the README's first 300 tokens):
    • window and tree equal step byte for byte, 300 of 300 rows;
    • with a second stream stepped beside (TF_Q35_BESIDE=300), step's rows are byte-identical to step alone;
    • mtp: all three pass (299 rows absorbed at once against chunks of 128, the poisoned ring suffix with 9 lanes byte-equal, and six siblings against chained paths with 9 lanes byte-equal).
  • The 27B on this branch (--parallel 2 --context 8192, DFlash2): --levels 1,2 gives 36 of 36. Its one-stream replies have the same token_sha as Windows: the native CUDA server, and Qwen3.5-2B and Qwen3.8-27B on CUDA #836's head 26e9b686, built and served here for Qwen3.5/3.8 on CUDA: decode 37-73% faster on #836's engine, bits unchanged (on #836) #978.

Tests on sm_121

  • tf-cuda-test:
    • moe passes, including 12 pairs over 6 SwiGLU experts within 6.9e-7 of fp64;
    • glue (1,242 of 1,242) and sample (2,880 draws) pass;
    • info, smoke and quant-ref pass.
  • zig build test: 48/48 steps, 753 of 758 tests passed, 5 skipped, 0 failed.

Findings (performance only, no bit differences)

  1. On the GB10, a second row in a forward doubles the MoE's cost. I measured this with tf-qwen35-cuda step over 300 ids:
    • Alone, a step takes 13.4 ms. With a second stream stepped beside (TF_Q35_BESIDE=300, two rows a forward) it takes 28.2 ms, 2.1 times as long.
    • The dense 27B on the same build goes from 73.4 to 78.3 ms (1.07 times).
    • When served, this is why plain decode at --parallel 2 is slower with two streams than with one: 76 tok/s for one stream, 66 to 74 aggregate for two. The 4090 receipt has 96 for one stream and 195 for two.
    • Reading eight more experts a layer is at most about 570 MB, which is roughly 2 ms at this card's bandwidth. So the 15 ms difference is not memory traffic. I have not looked inside the expert kernels to find where it goes.
  2. The planner's window costs are below what drafted windows cost when served. TF_TIMING=1 at --parallel 2 --context 8192 --drafter <the checkpoint>, one stream, chains of 3:
    • At load, windows of 2, 3 and 4 rows take 18.96, 19.80 and 20.33 ms.
    • Served, 144 drafted verify windows took 25 to 29 ms each. The 108 one-row verifies took about 13 ms, matching their load-time cost of 12.78 ms. A hold took about 3 ms.
    • The load-time timing (timeWindows in cuda_native.zig) forwards all-zero ids. A plausible cause, which I have not checked, is that those rows route to the same experts, while drafted rows route to different ones (finding 1). If so, on this card the planner sees a drafted window as about 30% cheaper than it is.
  3. Greedy drafted speed varies between repeats of the same request. I sent the same greedy chat request (bench_openai's chat prompt, T=0) repeatedly to one --parallel 2 server:
    • With the default chains, it took between 54 and 87 rounds (accepted 40/87 to 73/144), at 66.9 to 75.4 tok/s, with the same token_sha every time.
    • Greedy code was steady at 40 rounds (88/117).
    • With TF_MTP_DRAFTS=1, every repeat took the same 115 rounds with 12 of 17 drafts accepted. So the drafts themselves are deterministic, and the variation comes from the chain lengths the planner picks.
    • The description mentions variation on sampled prompts; here greedy varies as well.

Not run

  • The fp32 PyTorch reference comparison.
  • TF_MTP_TREES=1 and TF_MTP_DRAFTS other than the default (except for the one probe above).
  • The cold prompt at --parallel 4 --context 4096.
  • Image input.
  • The 2B lines.
  • Long contexts.

@ashhart

ashhart commented Oct 17, 2026

Copy link
Copy Markdown
Owner

Thanks @jschmied for the pinned GB10 run, including 36/36 two-stream and 60/60 four-stream checks. The two-row MoE cost, zero-ID calibration gap and changing greedy round counts are useful open performance findings, despite equal output hashes.

We are carrying the increment onto #836's existing Qwen family with the shared image, cache and MTP contracts preserved. The timing hypothesis still needs a profile; equal tokens do not make the planner's costs representative. Keep the unrun image, long-context and reference cells visible, and rerun the combined gates on the final candidate rather than treating this branch receipt as merged qualification.

Posted by Ash's AI Assistant Synthena

@ashhart

ashhart commented Oct 18, 2026

Copy link
Copy Markdown
Owner

Thank you, @BobClawblaw, for Qwen3.6-35B-A3B CUDA and its MTP chains/trees, and @jschmied for the 36/36 two-stream and 60/60 four-stream checks. 1.1.0 integrates this in the maintained Qwen family with #836, #972, shared images and prompt caching. We close this PR after the final MoE numerical/resource, image, long-context and draft/cache/concurrency gates pass and the integration ships, retaining your credit. Earlier branch receipts remain attributed evidence, including their open planner-cost findings.

GB10 35B-A3B MTP passed 17/17, concurrency and chat 12/12; image warm-turn/JPEG limitations remain.

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.

4 participants