Skip to content

perf(qwen3_5 cuda): deliver the first token before drafting - #932

Open
jregan-beasley-smc wants to merge 28 commits into
ashhart:mainfrom
jregan-beasley-smc:jon/pr836-first-token-handoff
Open

jregan-beasley-smc wants to merge 28 commits into
ashhart:mainfrom
jregan-beasley-smc:jon/pr836-first-token-handoff

Conversation

@jregan-beasley-smc

@jregan-beasley-smc jregan-beasley-smc commented Oct 16, 2026 •

Copy link
Copy Markdown

What this changes

Qwen CUDA already has the first output token after prefill, but currently waits for the initial speculative draft before handing that token to the client. This change commits and delivers the token first, then prepares the draft. It improves first-token responsiveness without removing the draft work.

The handoff preserves token/logprob ownership, cancellation and terminal cleanup, and prevents requeueing after a token has been emitted. The Qwen CUDA backend enables the capability; focused host tests cover the callback and error paths.

Stacked on #836 at 5e1c3feaa79e40f7659537d96b939129abe39df9. That dependency needs to land first. Both #836 and this stack currently conflict with main; rebase/integration will follow resolution of the dependency. GitHub's full diff currently includes #836; review this contribution alone. The 13 changed source files exactly match the candidate used for the measurements below. Both #836 and this stacked branch currently conflict with main; the eventual rebase will need checks against the integrated base.

Receipt

One DGX Spark (GB10), native Zig/CUDA, Qwen3.8-27B MLX 4-bit weights with DFlash2, BF16 KV, 32,768-token context and 16 server slots. Greedy short requests generate 256 tokens. Both arms restore zero prompt-cache tokens. This measures the handoff on #836, rather than reusing the old #747 cache-hit results.

Host driver 580.178.04 (nvidia-smi CUDA line 13.0). The build image is nvcr.io/nvidia/pytorch:26.07-py3 (sha256:15b8de8baabe…), CUDA toolkit 13.3.1.008, nvcc 13.3 V13.3.73. Native server SHA256 is 5700939f… baseline and e42ab4fc… candidate. Full pins are in the receipt. This is delivery ordering, not faster prefill arithmetic and not a startup claim.

Loaded-model fresh requests, baseline/candidate/candidate/baseline order; values are geometric means of the two per-process medians:

Workload Streams TTFT, baseline → candidate (ms) Decode, baseline → candidate (tok/s per stream) Full-response change
Coding 1 131.9 → 119.1 53.94 → 53.96 −0.30%
Coding 4 326.2 → 316.4 25.08 → 25.77 −2.11%
Coding 16 1111.2 → 1102.6 6.10 → 6.40 −4.92%
Chat 1 144.5 → 136.5 33.70 → 32.92 +2.23%
Chat 4 338.9 → 324.0 18.89 → 19.24 −1.72%
Chat 16 1152.3 → 1138.6 5.45 → 5.41 +1.37%

Lower latency/completion is better; higher token rate is better. TTFT falls 0.77–9.77% across all 12 loaded short fresh/repeat cells. Beside that, the slower completions stay in the table: fresh chat at 1 stream is +2.23% and fresh chat at 16 streams is +1.37%. Loaded single-stream chat decode falls 2.32%. This is not a universal performance win or a statistical significance claim.

Separate restart tests show mixed launch-to-ready and launch-inclusive results; first-batch chat at 16 streams has 4.05% higher request TTFT. Within that startup protocol, warm single-stream coding decode falls 5.19% and completion takes 5.21% longer. These are process-cold tests with warm/uncontrolled filesystem cache; no startup speedup is established.

Earlier delivery also changes the measured decode interval, and speculative work counts vary. Actual completion regressions remain; the receipt does not attribute them to a particular kernel or scheduler. All 37 CUDA fatbins and 331 kernel/build inputs are unchanged.

  • 376 paired timing requests match request bodies, generated token counts and output hashes; six additional smoke pairs pass.
  • Sampled draft-on/off equality, baseline/candidate output equality at 1/4/16 streams, disconnect recovery and terminal/error-path fixtures pass. Sampled performance was not measured.
  • Qualified candidate: 483/483 host tests passed and native build succeeded. Publication formatting/diff checks pass; lean findings remain 527 in each arm, with zero new findings.

Full receipt, source/model pins, all startup/repeat results and machine-readable observations. Other devices/models, image inputs and resumed-cache paths are unqualified.

Related work and attribution

Built on @BobClawblaw's #836 engine, following Ash's invitation to port focused TTFT changes. #608 and #726 provided the earlier engine/integration work behind #747.

Overlap review found complementary admission/cache work in #853/#888, planner work in #861, and fixes in #928/#929. If those admission/planner changes land first, their paths need explicit integration of this callback.

Authored by Jonathan Regan-Beasley in a private capacity. This contribution does not represent the views or positions of any organisation.

BobClawblaw and others added 27 commits October 14, 2026 17:04
- 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.

@cwschroeder cwschroeder left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reviewed the 13-file source increment from 5e1c3fea to d653cd50. I found no new correctness blocker in that increment: commit precedes delivery, cancellation/error cleanup retains ownership, and an emitted token prevents LaneMemory requeue. This is a source review of the pinned stack; the current-main conflicts and missing executed PR checks still need resolution.

I also checked the published receipt independently: all 13 candidate file hashes and available baseline hashes match; all 90 CSV geometric summaries recompute; JSON and CSV observations agree. Deduplicating metric rows gives the stated 336 short request pairs, without counting launch-inclusive rows again. This verifies the published summary arithmetic, not the 376 raw request/output joins or device execution.

Two useful host-test additions for the integration:

  • A real host sink that calls Engine.cancel(id) on its first token, then asserts exactly one completion/release and no initial draft.
  • Early/legacy output parity after a complete prompt-cache restore in the existing drafted fixture.

For the post-rebase hardware receipt, please link the runnable harness and per-request clock/hash rows, and record native binary SHA256 plus driver/CUDA versions. Those will let us replay the comparison on another GB10 while preserving the disclosed completion regressions and timing boundaries.

@ashhart

ashhart commented Oct 17, 2026

Copy link
Copy Markdown
Owner

Thank you, @jregan-beasley-smc, the opt-in first-token capability keeps delivery before initial drafting inside the lane engine.

We checked d653cd5. open_first.zig commits the selected token/logprob, invokes delivery and checks cancellation before drafting; the host callback retains admission ownership, and post-commit memory failure is not requeued. This is ready after #836's integration with the current registry preserved.

We can add early-versus-legacy restore parity and an actual host sink cancelling on first delivery at landing. Please pin the missing driver/toolkit versions with the refreshed receipt and retain the slower-completion cells beside the loaded-TTFT gain. The change is delivery ordering, not faster prefill arithmetic or a startup claim.

Posted by Ash's AI Assistant Synthena

@ashhart ashhart mentioned this pull request Oct 17, 2026
@jregan-beasley-smc

Copy link
Copy Markdown
Author

Thanks Ash - will have a look and get the details to you as requested soon. Cheers - Jon

The timing numbers are unchanged. The receipt now names the driver and the toolkit in the build image, and keeps the slower completion cells next to the loaded TTFT gain.
@jregan-beasley-smc

Copy link
Copy Markdown
Author

Thanks Ash. The receipt on 1471afeb pins the missing versions. Timing numbers are unchanged.

  • Host driver 580.178.04 (nvidia-smi CUDA line 13.0, the driver ceiling). That driver string was recorded on this Spark on 16 October 2026, and /usr/bin/nvidia-smi is dated 10 August 2026, before the runs.
  • Kernel 7.0.0-1019-nvidia was read on this Spark on 17 October 2026. It was not in the original timing logs.
  • Builds used nvcr.io/nvidia/pytorch:26.07-py3 (sha256:15b8de8baabe…, created 2026-07-25). Toolkit in that image: CUDA 13.3.1.008, nvcc release 13.3, V13.3.73. The host nvcc 13.0 was not the compiler.
  • Native server SHA256: baseline 5700939fd17c0ad52a84c7d4bdbeeba7975a6c96952375579fda0a5d70c306c6, candidate e42ab4fc5079bfdeafe61234de4d044fec62e18877568e87fb19bb7e9523d581.

Loaded short TTFT is still 0.77–9.77% lower across all 12 fresh/repeat cells. Beside that, fresh chat completion is 2.23% longer at 1 stream and 1.37% longer at 16 streams. This is delivery ordering after prefill, not faster prefill arithmetic and not a startup claim.

Happy for the early-versus-legacy restore parity check and the host sink that cancels on first delivery to be added at landing, as you offered.

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