Skip to content

qwen3_5 cuda: 2, 3, 5 and 6-bit tiles of their own width and a prompt GEMM for them (27B oQ2 on 16 GB) - #859

Open
BobClawblaw wants to merge 33 commits into
ashhart:mainfrom
BobClawblaw:cuda-native-lowbit
Open

BobClawblaw wants to merge 33 commits into
ashhart:mainfrom
BobClawblaw:cuda-native-lowbit

Conversation

@BobClawblaw

@BobClawblaw BobClawblaw commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor

Stacked on #849, which is stacked on #836. Until those merge, this PR shows their commits too. The change here is 39c293a6..7cddf336, five commits, rebased onto #849's 9b843c10 (on #836's 5e1c3fea). --kv-cache fp8 was split out of this PR into #886, which is stacked on it.

What this changes

Native 2, 3, 5 and 6-bit tiles (39c293a6, 11c8057e). #849 widened 2 and 3-bit codes onto the 4-bit tiles and 5 and 6-bit codes onto the 8-bit tiles, so a mostly 2-bit checkpoint took as much memory as the 4-bit one. Here every width gets tiles of its own: 2, 3, 5 and 6-bit modules are stored at their width, 7 and 8-bit modules on the 8-bit tiles, and 4-bit modules on the affine-4 tiles as before.

  • affine8.cu's lane GEMV is templated on the width. Within each 64-group, a lane's 16 code pairs sit in bits words: the first codes in the low halves, the second codes in the high halves.
  • One shift and mask extracts a pair (two for a pair that crosses a half), and bf16 (0x4300 | c) - 128 rebuilds each code exactly.
  • The MMAs get the same bf16 operands as on the 8-bit tiles and run the same fma chain, so a code's products do not depend on the width it is stored at.
  • qlinear.zig's Affine8 becomes AffineN and carries its width. AffineNKernels resolves a gemv and a split-K kernel for each width.

With this, TensorFold/Qwen3.8-27B-oQ2 takes 9.9 GiB on the device instead of about 16, so the 27B fits a 16 GB card.

A prompt GEMM for those tiles (b8ffd990, d3da9a18). affine_prefill.cu runs qmm_prefill's prompt GEMM over AffineN's tiles, for 2, 3, 5, 6 and 8-bit modules. Before this, they ran prompts on the decode tiles, 16 rows at a time. The code readers move into affine_codes.cuh, which the GEMV and the GEMM share.

The load staging buffer freed after load (7cddf336). The buffer each projection is staged through on its way to the tiles now belongs to Engine.init and is freed once the weights are on the card, so the server's stream budget counts that memory. The drafter's loader stages its own. On the 27B oQ2, the startup line shows the model at 9.92 GiB instead of 10.29, leaving 3.8 GiB for streams instead of 3.4.

Which sums change:

Receipt

  • Environment:

  • Exactness: bench_concurrent --alone --serial (levels 1 and 4, 4 cells each):

    Checkpoint Equal alone Equal to "draft": false
    9B 4-bit 48/48 10/10
    9B 8-bit 48/48 10/10
    27B oQ2 48/48 10/10

    tf-qwen35-cuda over 300 ids: window and tree equal step on the 2B and the 27B.

    With the DFlash2 drafter (z-lab/Qwen3.8-27B-DFlash2 @ 50307d4c). It is quantized to 4 bits at load through the drafter's own staging buffer, which 7cddf336 gives it. Served with --drafter --parallel 2 --context 8192:

    • Load: the 27B oQ2 and its drafter take 10.91 GiB, with 2 streams of 0.75 GiB.
    • Exactness (bench_concurrent --levels 1,2 --tokens 128 --reps 2 --alone --serial): 24/24 equal alone, and 12/12 drafted replies equal to "draft": false.
    • Drafted decode, one stream (bench_openai --tokens 128 --reps 3): 175 tok/s on code and 95 on chat, greedy. At T=1, 112 and 60. Undrafted, it is about 46.
    • Limit: two streams at a 32,768-token window don't fit beside the drafter with bf16 caches. They need 4.49 GiB, and 2.75 GiB is left, so the server refuses with that message.
    • Build: measured on qwen3_5 cuda: --kv-cache fp8, e4m3 attention caches with a scale a row #886's head serving bf16 caches. Its bits equal this head's, and it doesn't change the drafter loader.
  • Decode speed (bench_openai, one stream, medians over 5 reps; fibonacci-raw and gpu-chat-no-think at T=0 and T=1):

    Checkpoint qwen3_5 cuda: mixed-width MLX checkpoints, 2 to 8 bits in groups of 64 #849 This PR
    9B 4-bit 99.7–100.8 99.9–100.9
    9B 8-bit 60.2–61.4 60.2–61.4
    27B oQ2 does not fit (memory budget fits no stream) 46.0–46.3

    Four streams together: 322–343 tok/s (9B 4-bit), 112–116 (9B 8-bit) against 111–113 on qwen3_5 cuda: mixed-width MLX checkpoints, 2 to 8 bits in groups of 64 #849.

  • Prompt speed: the cold first request's ttft_s from bench_cache --require-cold, on a 16,000-token prompt of the repo's docs (all rows report 0 cached tokens):

    Checkpoint Windows: the native CUDA server, and Qwen3.5-2B and Qwen3.8-27B on CUDA #836 qwen3_5 cuda: mixed-width MLX checkpoints, 2 to 8 bits in groups of 64 #849 This PR
    9B 4-bit 8.89 s 8.89 s 8.89 s (1,801 tok/s)
    9B 8-bit refused 11.47 s 9.20 s (1,739 tok/s)
    27B oQ2 refused does not fit 23.66 s (676 tok/s)
  • Tests:

    • zig build test: 48/48 steps; 745 of 759 tests pass, 14 skipped, 0 fail.
    • tf-cuda-test affine8: 24 cases over 2, 3, 4, 5, 6 and 8 bits pass.
      • Narrow products are byte-equal to the 8-bit tiles' at every row count and on both paths, within 0.338 of the float64 bound.
      • The prompt GEMM over 144 rows is within its float64 bound and equals the 4 and 8-bit tiles' GEMM bit for bit.
    • tf-qwen35-cuda, 300 ids, step / window / tree / prompt:
    • New host tests: AffineN's tiles hold bits a weight over Affine4's padding, and a uniform 4-bit checkpoint keeps its sizes while other modules take their own width.
    • CI is green.
  • Not run, and why:

    • Metal: no Mac here. The change is CUDA-only.
    • The RTX 5090 (Linux) and GB10 runs in the comments below are at fe9f07c2, before the prompt GEMM and staging commits and before the rebase onto 5e1c3fea. I ran this head on sm_89 only.

Quality and long context

The quality suite covers tools, JSON, IFEval, GSM8K, MGSM, MMLU, HumanEval (in Docker) and long-context needles and facts, at a 40,960-token window. Builds: #836 at 153d6b43, and this PR's code at its prompt-GEMM commit, both before the rebase onto 5e1c3fea. That rebase changes attention only when streams of different lengths share a forward. One stream's logits are byte-equal before and after it in all four modes (2B, mixed-width 2B and 27B), so these single-stream results carry over.

Check 9B 4-bit, #836 9B 4-bit, this PR 9B 8-bit, this PR 27B oQ2, this PR
tools / JSON 60/60, 30/30 60/60, 30/30 60/60, 30/30 12/60, 2/30
IFEval 125/150 125/150 129/150 45/150
GSM8K / MGSM 234/250, 184/240 234/250, 184/240 241/250, 196/240 113/250, 30/240
MMLU / HumanEval 268/342, 142/164 268/342, 142/164 265/342, 143/164 202/342, 3/164
long context to 28k 16/16 16/16 16/16 0/16
drafts / concurrent / seeded sampling 12/12, 12/12, 9/9 12/12, 12/12, 9/9 12/12, 12/12, 9/9 12/12, 12/12, 9/9

The 9B 4-bit scores are identical on both builds, down to the GSM8K misses.

The 27B oQ2's weak scores are the checkpoint's own: its recipe stores the LM head and the MLP and linear-attention input projections at 2 bits. I checked the engine against an independent fp32 numpy forward of Qwen3.5, written from the HF reference.

  • 2B: the reference agrees with the engine at logit correlation 0.9999.
  • 27B oQ2, first 24 ids: the engine's logits agree on top-1 at 24/24 and correlate at ≥ 0.999.

Long context with bf16 caches, one stream (--parallel 1), before the rebase as above:

Checkpoint, cache Window Needles and facts Slowest prefill
9B 4-bit, bf16 (#836) 232k 55/55 356 tok/s
9B 4-bit, bf16 232k 55/55 356 tok/s
9B 8-bit, bf16 120k 42/42 651 tok/s

#886 has the same runs with fp8 caches, at 256k.

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. This engine keeps no prompt cache on CUDA, so a resumed prompt is a fresh one.
  • No precision traded for speed. If the bits change, the description says which sums change and why: see "Which sums change".
  • Prompt processing is no slower than the last release. The 9B 4-bit is level, and the 9B 8-bit takes 20% less time than on qwen3_5 cuda: mixed-width MLX checkpoints, 2 to 8 bits in groups of 64 #849.
  • The description says which platforms I ran: CUDA sm_89 on Windows for every commit. sm_120 and sm_121 on Linux ran up to fe9f07c2 (in the comments). Metal was not run.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. The affine8 prompt cases and the new host tests cannot build without the change.
  • Comments and docstrings are one line. No measurements or history in the source. The CI lean check reports 0 new problems.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

Limits

  • Embeddings still widen to 4 or 8 bits at load, because the embed kernels read those widths.
  • Release packaging from prebuilt fatbins needs affine_prefill.fatbin too.

- 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

Ran this on Linux with an RTX 5090 (sm_120). That covers the "Linux" and "DFlash2 drafter with this target" items in Not run.

First, a CI fix. zig was failing on the lean check's new-problems step: 21 multi-line comments and comments past 120 columns. The branch is now rebased onto #849's own comment fix (d6bbfbe6), plus fe9f07c2, which puts this PR's comments on one line each. Only comment lines change (checked with a diff of non-comment lines). lean_new.py reports 0 new, zig fmt --check is clean, zig build test passes, and CI is green.

Setup: fe9f07c2, built with -Dsm=120. Ubuntu, kernel 7.0, driver 615.71.09, CUDA 13.3, Zig 0.17.0. Comparison builds: #849 at d6bbfbe6 and #836 at 153d6b43. Neither is merged with main, which is 2 commits ahead.

tf-cuda-test: affine8 (24 cases; 2, 3, 5 and 6-bit products byte-equal to the 8-bit tiles' at every row count and on both paths, within 0.320 of the float64 bound), smoke, glue (1,242 of 1,242), quant-ref, modelopt and exlr8 all pass.

tf-qwen35-cuda, 300 README ids:

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

Model, serve Levels Checks equal Unequal / failed
27B oQ2 + DFlash2, --parallel 2 --context 8192 1,2 36 / 36 0 / 0
27B 4-bit + DFlash2, same 1,2 36 / 36 0 / 0
4B OptiQ, --parallel 4, no drafter 1,4 60 / 60 0 / 0
9B OptiQ, --parallel 4, no drafter 1,4 60 / 60 0 / 0

Without a drafter, the --serial half of the check is trivial. The concurrent-vs-solo half still applies.

27B oQ2, one stream, bench_openai --tokens 128 --reps 3, median tok/s:

Serve code greedy chat greedy code T=1 chat T=1
plain 80 68 71 68
DFlash2 199 143 139 89

Long context: --parallel 1. 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.

Model Context Prompt tokens Code retrieved Drafted = plain Time to first token GPU memory after load
27B oQ2 + DFlash2 65,536 64,525 yes, both yes 104.5 s 16.5 GB
27B oQ2 + DFlash2 131,072 130,060 yes, both yes 344 s 20.6 GB
27B oQ2 + DFlash2 262,144 261,134 yes, both yes 1,343 s 28.8 GB
9B OptiQ, no drafter 65,536 64,526 yes, both (no drafter) 38.2 s 10.6 GB
9B OptiQ, no drafter 131,072 130,061 yes, both (no drafter) 120.9 s 12.7 GB
9B OptiQ, no drafter 262,144 261,134 yes, both (no drafter) 470 s 16.8 GB

So the oQ2 runs with its drafter at 256K on a 32 GB card. Its prompt pass is slow there: 22 minutes, about 194 tok/s. That fits the Limits note that widths other than 4 have no prompt GEMM. The oQ2 answered with a sentence at 64K and 256K rather than "the code only", but it named the right code each time. Decode speed over 64-token replies at these lengths was too noisy to report.

Not run: reply quality or perplexity on this card, an MLX or PyTorch oracle for the 27B, GB10. These speeds come from one card and aren't a comparison against the 4080 SUPER's.

@BobClawblaw

Copy link
Copy Markdown
Contributor Author

Test run on a GB10 (sm_121, aarch64 Linux)

Head fe9f07c2.

Setup. One DGX Spark: NVIDIA GB10 (sm_121, aarch64, 128 GB unified memory), Ubuntu kernel 7.0.0-1019-nvidia, driver 580.178.04, nvcc 13.3, Zig 0.17.0. Built with -Dsm=121 -Doptimize=ReleaseFast.

…'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.
@BobClawblaw BobClawblaw changed the title qwen3_5 cuda: 2, 3, 5 and 6-bit modules on tiles of their own width (27B oQ2 fits 16 GB) qwen3_5 cuda: 2, 3, 5 and 6-bit tiles of their own width, a prompt GEMM for them, and --kv-cache fp8 (27B oQ2 on 16 GB) Oct 15, 2026
@BobClawblaw BobClawblaw changed the title qwen3_5 cuda: 2, 3, 5 and 6-bit tiles of their own width, a prompt GEMM for them, and --kv-cache fp8 (27B oQ2 on 16 GB) qwen3_5 cuda: 2, 3, 5 and 6-bit tiles of their own width and a prompt GEMM for them (27B oQ2 on 16 GB) Oct 15, 2026
Each module loads at the width the config's quantization table names it. Codes
of 2 to 4 bits pack onto the affine-4 tiles and codes of 5 to 8 bits onto new
8-bit tiles, every code keeping its value, scale and bias, so nothing is
rounded. affine8.cu's lane GEMV is lane_gemv's arithmetic a byte wide (each
byte exact in bf16), with its split-K form; an odd-width embedding is widened
at load and 8-bit rows have their own embed kernel. A uniform 4-bit checkpoint
takes the same packs and kernels as before: the 2B's step, window, tree and
prompt logits are byte-equal to the previous build's.

The rope config's older "type" key is read as "rope_type" (mlx-community's
OptiQ conversions write it).

Checked on an RTX 4080 SUPER (sm_89):
- tf-cuda-test affine8: 24 cases over 2, 3, 4, 5, 6 and 8 bits; tiles equal
  the host-widened codes' tiles, 8-bit products within 0.32 of the float64
  bound, 1/3/9/16-row calls and split-K byte-equal.
- A 2B requantized at 2, 3, 5, 6 and 8 bits (and a 6-bit embedding) gives the
  same logits as the same codes stored widened, in step, window, tree and
  prompt modes.
- Qwen3.5-4B-OptiQ-4bit and Qwen3.5-9B-OptiQ-4bit: window and tree equal step
  over 300 ids; bench_concurrent --alone --serial at 1 and 4 streams all equal;
  about 130 and 78 tokens/s for one request.

8-bit modules have no prompt GEMM yet; prompt chunks run them on the decode
tiles, 16 rows at a time.
The mixed-width loader widened 2 and 3-bit codes onto the 4-bit tiles and 5 and
6-bit codes onto the 8-bit ones, so a mostly 2-bit checkpoint took the 4-bit
checkpoint's memory. affine8.cu's lane GEMV is now templated on the width: a
lane's 16 code pairs a group sit in `bits` words, first codes in the low halves
and second in the high, and bf16 (0x4300 | c) - 128 rebuilds each code exactly.
The MMAs see the operands the 8-bit tiles give them, so a code's products do not
depend on the width it is stored at.

- AffineN (was Affine8) carries its width; AffineNKernels resolves a gemv and a
  split-K kernel for 2, 3, 5, 6 and 8 bits, and tf_affinen_pack packs them.
- 4-bit modules keep the affine-4 tiles and prompt GEMM; uniform 4-bit
  checkpoints are unchanged. Embeddings still widen to 4 or 8 bits (rowBits).
- TensorFold/Qwen3.8-27B-oQ2 takes 10.3 GiB and fits a 16 GB card.

Checks on an RTX 4080 SUPER:
- tf-cuda-test affine8: 24 cases; 2, 3, 5 and 6-bit products byte-equal to
  the codes widened onto the 8-bit tiles at 1, 3, 9 and 16 rows, within 0.32 of
  the float64 bound, split-K equal to the slice loop.
- Qwen3.5-2B 4-bit: step, window, tree and prompt logits byte-equal to the
  build before.
- The 2B requantized at 2, 3, 5, 6 and 8 bits: byte-equal to the same codes
  widened to 8 bits in all four modes, and step/window/tree byte-equal to the
  previous build, where 2 and 3-bit codes ran on the 4-bit tiles.
- Qwen3.8-27B-oQ2: window, tree and prompt equal step over 300 ids;
  bench_concurrent --alone --serial at 1 and 2 streams all equal; about 45 to
  64 tokens/s for one request, about 750 tokens/s prompt.
Prompt chunks ran modules of every width but 4 on the decode tiles, 16 rows at
a time. affine8.cu now has qmm_prefill's GEMM over AffineN's tiles for 2, 3, 5,
6 and 8 bits: a 128 by 128 block on 2 by 4 warps, three stages, each weight
rounded once to bf16 (s * q + b) and one fp32 chain over K. Its operands are the
affine-4 prompt GEMM's, so a code of up to 4 bits gives the same bits as there.

Checks on an RTX 4080 SUPER:
- tf-cuda-test affine8: the prompt GEMM over 144 rows within 0.07 of a bound
  that allows for the bf16 weights; 2 and 3-bit byte-equal to the affine-4
  prompt GEMM of the same codes, 5 and 6-bit to the 8-bit tiles'.
- Qwen3.5-2B 4-bit: all four modes byte-equal to the build before.
- The 2B requantized at 2, 3, 5, 6 and 8 bits: byte-equal to the same codes at
  8 bits in all four modes, prompt included; decode modes unchanged.
- Qwen3.8-27B-oQ2: window and tree equal step, step unchanged; bench 0
  unequal. Prompts about 770 tokens/s (from 750): already bound by the bf16
  tensor rate with fp32 accumulation.
- Qwen3.5-9B MLX 8-bit: prompts about 1,940 tokens/s, from 1,520; bench at 1
  and 4 streams 0 unequal.
…ffine_codes.cuh, one-line comments

affine8.cu passed 600 lines with the prompt GEMM. The GEMM moves to its own
kernel module and the code readers both use to a header; lean_check reports
528 problems, as on the base.
The weight loader's staging buffer is now Engine.init's and freed once
the weights are on the card, so the server's stream budget counts that
memory; the drafter's loader stages its own for its load.
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.

1 participant