Skip to content

cuda qwen3_5: TF_QWEN_KV_PAGES_ALL puts the engine's own stream on pages too, --parallel lanes and no reserved window - #1750

Open
jschmied wants to merge 1 commit into
ashhart:mainfrom
jschmied:qwen-pages-engine
Open

jschmied wants to merge 1 commit into
ashhart:mainfrom
jschmied:qwen-pages-engine

Conversation

@jschmied

Copy link
Copy Markdown
Contributor

With TF_QWEN_KV_PAGES=n and TF_QWEN_KV_PAGES_ALL=1, main runs every stream on pool pages, but it still reserves the engine's own window, a full stream cache at --context, and leaves it idle. The drafters also keep one ring of their own. With this change the engine's own stream becomes a pool lane too (engine_paged), as on GLM, Gemma 4 (#1619) and Kolibri-1 (#1620). The plan sets one window's worth of pages aside for it, no window is reserved, and DFlash2 and MTP keep their rings in the lane runs. Without _ALL, nothing changes. This follows #1156 and #1357 (the lane count fix, d60cc24, is already on main).

  • attach makes streams + 1 lanes when the engine's stream is paged; Loaded.engine_paged tells the host's plan.
  • reserve makes no window then: shapes warm on the first pool lane.
  • DFlash2's window timing gets null and makes its own window_rows cache, as cost_adapter.resolve already does for MTP.
  • RingOf.all: neither drafter keeps a ring of its own.

Startup line at --parallel 8 --context 16384:

main _ALL this _ALL
Qwen3.8-27B all 8 on 959 pages (14.98 GiB), the engine's own window (1.90 GiB) idle the engine's own and 7 more on 1039 pages (16.23 GiB), no window
Qwen3.6-35B-A3B all 8 on 994 pages (4.85 GiB), the engine's own window (0.59 GiB) idle the engine's own and 7 more on 1074 pages (5.24 GiB), no window

No kernel changed. The 27B and 35B attend over pages through tf_q35_attn_partial_paged (qwen35_attn.cuh), not through window_rows, so the paged merge that slowed #1653 (tf_wrows_merge_256_pg, fixed there in 7c0a80b) is not on this path.

Gates (GB10, sm_121)

These are the #1357 gate suite: 37 greedy outputs per arm (solo, concurrent 1/2/4/8, resumes at 8192/8292/8448/4796), --parallel 8 --context 16384, a fresh server per arm. Models: Qwen3.8-27B MLX-4bit with DFlash2, and Qwen3.6-35B-A3B MLX-4bit with its MTP layer, each with bf16 and fp8 KV. Main is the merge base 5d19272.

gate 27B bf16 27B fp8 35B bf16 35B fp8
unset: this == main 37/37 37/37 37/37 37/37
_ALL: this == main _ALL (solo, resumes) 22/22 22/22 22/22 22/22
_ALL == contiguous (solo, resumes) 22/22 22/22 22/22 22/22
_ALL: concurrent == solo 15/15 15/15 15/15 15/15
_ALL: resumed == main _ALL 12/12 12/12 12/12 12/12
_ALL: drafted == plain 22/22 22/22 22/22 22/22
PAGES alone: concurrent == window or page solo 15/15 15/15 15/15 15/15

Speed against main (GB10, sm_121)

Setup: same models and drafters, bf16 KV, --parallel 8 --context 16384 --prompt-cache-gib 0, a fresh server per arm. The four arms (main, this, main _ALL, this _ALL) alternate in ABBA order: 3 rounds for the 27B, 4 for the 35B.

  • Prefill: token-id prompts of 2,048, 9,216 and 15,000 tokens, 3 each per server; a server's value is the median.
  • Decode: 8 fixed chat prompts × 256 tokens (ignore_eos), 1, 4 or 8 in flight; aggregate tok/s = completion tokens / wall.
  • Deltas: medians over rounds, with the per-round paired deltas in brackets.
  • Tokens: greedy token_sha equal across arms in every cell.

Default (no env):

cell 27B main 27B this delta (per round) 35B main 35B this delta (per round)
prefill 2,048 (tok/s) 1,407.9 1,414.1 +0.44 % (+0.5, +0.1, -0.3) 5,091.2 5,154.2 +1.24 % (-0.2, -0.1, +2.0, +3.2)
prefill 9,216 (tok/s) 1,365.3 1,372.1 +0.50 % (+0.7, +0.2, -0.7) 5,459.8 5,509.9 +0.92 % (+0.2, -0.2, +1.1, +3.4)
prefill 15,000 (tok/s) 1,348.0 1,342.8 -0.39 % (-0.4, -0.2, -0.6) 5,390.6 5,428.5 +0.70 % (+0.2, -0.0, +1.0, +1.9)
1 stream, code, T 1.0 66.3 66.9 +0.82 % (+3.4, +0.2, +0.7) 131.9 132.0 +0.06 % (+0.3, -0.3, -0.7, +1.3)
1 stream, code, greedy 73.3 73.8 +0.64 % (+3.5, +0.2, +0.6) 134.1 132.9 -0.91 % (-0.8, -1.7, -1.5, +0.4)
1 stream, chat, T 1.0 42.4 42.8 +0.74 % (+3.5, +0.3, +0.7) 113.7 112.4 -1.15 % (-2.6, -0.9, -2.2, +1.0)
1 stream, chat, greedy 45.8 46.9 +2.35 % (+3.5, +0.5, +2.3) 115.6 115.0 -0.56 % (-1.1, -0.2, -2.0, +0.0)
4 streams, code, T 1.0 177.7 180.2 +1.38 % (+1.8, -0.3, +1.4) 295.3 293.8 -0.52 % (-0.5, -0.3, -1.5, +0.4)
4 streams, code, greedy 181.9 185.6 +2.04 % (+3.0, +0.0, +1.3) 293.0 291.4 -0.56 % (-0.6, -0.3, -2.4, +0.3)
4 streams, chat, T 1.0 114.3 115.5 +1.11 % (+1.7, +0.7, +1.1) 256.4 255.0 -0.54 % (-0.9, -0.1, -1.1, +0.4)
4 streams, chat, greedy 130.6 129.4 -0.92 % (+1.9, -1.7, -1.0) 261.4 260.3 -0.42 % (-0.6, -0.3, -0.3, +0.4)
8 streams, code, T 1.0 242.4 242.9 +0.23 % (+1.4, +0.2, -0.8) 397.8 396.5 -0.32 % (-0.4, -0.2, -0.4, +0.2)
8 streams, code, greedy 255.0 256.3 +0.49 % (+1.5, +0.1, +0.5) 399.1 397.3 -0.46 % (-0.5, -0.5, -0.2, +0.2)
8 streams, chat, T 1.0 170.0 170.9 +0.53 % (+1.6, +0.0, +0.5) 347.2 345.9 -0.37 % (-0.2, -0.4, -0.7, +0.5)
8 streams, chat, greedy 181.6 182.5 +0.47 % (+1.8, +0.7, +0.4) 360.6 359.3 -0.36 % (-0.2, -0.7, -0.1, +0.2)

Without the env, the decode path is the same code (the new fields are read only at startup), and the 103 fatbins embedded in the two binaries are byte-identical (md5). Two cells still had one sign in every round, so I checked both:

  • 35B decode was below main in rounds 1-3. Round 4, run in a separate block with an A/A arm (main's binary on a second server), came out above main in every cell, so no cell keeps its sign. The A/A arm itself differed from main by -1.6 to +3.4 % per cell. A graph-level nsys trace (--cuda-graph-trace=graph, 1 and 8 streams of code greedy, main/this/this/main) shows the same GPU work per decode round: busy time per tf_q35_argmax launch was 4.83 / 4.78 / 4.80 / 4.83 ms. The rest is idle time between steps, and it varies with the server, not with the binary.
  • 27B prefill at 15,000 was below main in all three rounds (-0.4 %). A prefill-only rerun (3 more rounds, 5 prompts each, plus the A/A arm) gave -0.52 / -0.05 / -0.03 % against main, and the A/A arm itself gave -0.62 / +0.23 / +0.10 %. Under nsys, three 15,000-token prompts kept the GPU busy for 99.9 % of their span in every arm. The span was 37,321 and 36,635 ms on main against 36,842 and 36,807 ms here, with the same kernels.

TF_QWEN_KV_PAGES=256 TF_QWEN_KV_PAGES_ALL=1 (opt-in):

cell 27B main 27B this delta (per round) 35B main 35B this delta (per round)
prefill 2,048 (tok/s) 1,384.3 1,391.3 +0.51 % (+1.3, -2.3, +0.6) 5,105.3 5,081.0 -0.48 % (-0.1, -0.1, -1.1)
prefill 9,216 (tok/s) 1,335.8 1,339.6 +0.29 % (+1.1, -1.6, +0.3) 5,406.0 5,377.7 -0.52 % (-0.2, +0.1, -0.7)
prefill 15,000 (tok/s) 1,305.9 1,301.5 -0.34 % (+1.5, -1.8, -0.3) 5,275.8 5,240.2 -0.68 % (-0.3, -0.1, -0.8)
1 stream, code, T 1.0 67.0 66.4 -0.93 % (-1.2, -1.1, -0.1) 132.8 132.0 -0.60 % (-0.1, -0.5, -0.6)
1 stream, code, greedy 74.0 73.3 -0.98 % (-1.4, -0.5, -0.0) 133.8 132.2 -1.20 % (-1.7, -1.2, +0.5)
1 stream, chat, T 1.0 42.9 42.5 -0.91 % (-1.3, -0.4, +0.1) 113.9 113.5 -0.31 % (-0.9, -0.4, +2.0)
1 stream, chat, greedy 47.0 46.7 -0.80 % (-1.3, -0.5, +0.1) 114.3 114.4 +0.12 % (-0.3, -1.4, +0.6)
4 streams, code, T 1.0 178.9 176.6 -1.27 % (-1.4, +0.2, -0.5) 293.2 295.5 +0.76 % (-0.2, +0.2, +1.3)
4 streams, code, greedy 185.3 184.0 -0.65 % (-2.1, -0.2, +0.2) 290.7 293.1 +0.83 % (-0.1, +0.3, +1.3)
4 streams, chat, T 1.0 114.7 114.1 -0.48 % (-1.4, -0.5, -0.4) 254.7 257.0 +0.92 % (-0.0, +0.2, +1.4)
4 streams, chat, greedy 130.3 128.3 -1.57 % (-2.5, -0.1, -1.0) 259.3 261.3 +0.77 % (-0.1, +0.1, +1.3)
8 streams, code, T 1.0 241.6 242.2 +0.26 % (-0.0, +0.6, +0.7) 395.0 397.6 +0.66 % (-0.3, +0.2, +1.5)
8 streams, code, greedy 255.1 253.7 -0.57 % (-1.1, -0.5, +0.2) 394.6 398.2 +0.91 % (-0.4, +0.1, +1.3)
8 streams, chat, T 1.0 170.1 169.2 -0.51 % (-1.2, -0.3, +0.2) 345.7 346.7 +0.28 % (-0.6, +0.0, +1.1)
8 streams, chat, greedy 181.7 180.3 -0.74 % (-1.1, -0.1, -0.2) 358.1 360.1 +0.54 % (-0.2, +0.3, +1.0)

Under _ALL, a lone stream is now the engine's own stream on pages, where main ran it on the reserved window. The 27B's 1-stream cells came out 0.8-1.0 % lower in their medians, but no cell held its sign in all three rounds. I did not trace that case. The 35B's third _ALL main arm was rerun later, because the first run was lost when its loop was stopped.

Checks

  • zig build test test-kernel-contract test-golden test-server-security test-pull-security fuzz-cortheon and git diff --check: pass on this head. They also pass on this head merged into main 94b90c9. The branch merges cleanly there and into e5b50c8 too, so it was not rebased.
  • Lean: 0 new problems.
  • No kernel file changed.

Whether pages should be the default for qwen3_5, as #1619 and #1620 propose for their families, is left open here.

…ges too, --parallel lanes and no reserved window
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