Repository navigation
qwen CUDA: decode attention on tensor cores (fp8 and bf16 caches), decode after long prompts 1.5-2.3x faster - #1168
Closed
BobClawblaw wants to merge 25 commits into
Closed
BobClawblaw wants to merge 25 commits into
BobClawblaw wants to merge 25 commits into
Conversation
BobClawblaw
changed the title
BobClawblaw
changed the title
…hanged On one RTX 5090 (Qwen3.8-27B MLX 4-bit + DFlash2, fp8 KV, a 262,144-token window, one stream) the GPU time of a 256-token code reply falls from 730 to 697 ms, and the ruler's one-stream decode rises 5-6% (prose 211.5 to 224.6, code 396.6 to 417.4, structured 265.7 to 279.0, list 258.7 to 272.1 tok/s). Every ruler reply's hash is the same as before. - tf_q35d_dyn_conv_batch: a (rows, channel blocks) grid instead of one block a row, 15.8 to 1.7 us. - tf_q35_gdn_pre logs a recorded window's conv tails from the taps it already reads, so tf_q35_conv_commit only moves the stream's tail: 4.4 + 2.8 us to 0.8 + 3.8 us a layer. - tf_q35_gdn_slots_*: the next row's inputs load while this one computes, k and q in 8-byte loads; the same arithmetic and reduction order (slots_4 22.9 to 18.3 us). - tf_q35_norm loads its weights with the row instead of after the reduction, 3.7 to 3.3 us. - tf_q35d_attn_partial: four warps a 256-key chunk, keys split over the lanes and the ring row advanced instead of a 64-bit modulo a key, 92 to 15.5 us. Drafter numerics only: verification decides every reply.
…the next stream Live lone pricing (lone_cost_tile, Qwen CUDA) kept each stream's measured round times and its losing widths' retry deadlines on the stream, so every request calibrated again: forced trials at one row (no drafts) and at the tile from round 8, the backoff doubling only within that reply. On one RTX 5090 (Qwen3.8-27B MLX 4-bit + DFlash2, one stream) a 37-round reply spent two or three rounds on trials. The engine now keeps the newest calibration of a stream that only ran alone (greedy and sampled apart), and a stream's first priced round starts from it, its deadlines counted from that round. Its own rounds still blend into the prices and refresh the widths it uses; a stream that shared a round leaves nothing, and one already priced keeps its own. Ruler, one stream: prose 224.8 to 237.2, code 417.8 to 439.8, structured 279.3 to 286.1, list 272.2 to 276.5 tok/s; replies unchanged.
…ares in the same order h + r and w * (h * inv) are elementwise, so they now load and store 8 entries a lane through the row's shared copy; each thread's sum of squares still reads its own E entries in order from that copy, so the result is unchanged. On one RTX 5090 (Qwen3.8-27B, 16-row verify windows) a call falls from 3.3 to 2.5 us; the ruler's one-stream decode rises about 0.7% (prose 237.2 to 238.9, code 439.8 to 443.2 tok/s); replies unchanged.
…our slots Each row keeps its own state and the same arithmetic and reduction order; the rows share a step's key, query, decay and strength loads, and their two dependent shuffle chains interleave. Eight and sixteen slots stay at one row a warp (their kept states already fill the registers). On one RTX 5090 (Qwen3.8-27B, 16-row verify windows) tf_q35_gdn_slots_2 falls from 17.4 to 14.3 us and slots_4 from 18.3 to 13.8 us; the ruler's one-stream decode rises about 1% (prose 238.9 to 240.9, code 443.2 to 446.1 tok/s); replies unchanged.
… tails in 8-byte loads Each tap's four channels, the channel's four weights and the recorded window's tail rows move as one uint2 each, in place of sixteen scalar loads and twelve scalar stores a lane; the convolution adds its taps in the same order. On one RTX 5090 (Qwen3.8-27B, 16-row verify windows) a call falls from 3.8 to 2.9 us; replies unchanged.
…sums from the kernels that write them tf_q35_gdn_post, tf_q35_attn_gate and tf_q35_mlp_act now write the 64-entry group sums of their outputs (group_sum on the same bf16 values, in order) when given XS, so those projections no longer launch tf_q35_sums first; a prompt GEMM past 16 rows takes none. The drafter's MLP does the same. On one RTX 5090 (Qwen3.8-27B + DFlash2, one stream) the sums launches of a 256-token reply fall from 5,474 to 686 and the ruler's one-stream decode rises about 0.7% (prose 241.6 to 243.4, code 447.4 to 450.7 tok/s); replies unchanged.
…t the head's depth rates A lone head window that a context copy extends has more rows than the drafter grew, so the drafter could not price them and allocation fell back to chain probabilities from the stream's per-depth acceptance for every row, tree and copy alike. Those fall fast with depth, so a copy's tail looked hopeless and was cut: on one RTX 5090 (Qwen3.8-27B + DFlash2) code replies, where copies from the prompt keep landing, verified 11.3 rows a round against 14 elsewhere. The drafts keep the drafter's own chances, and each copied row takes its parent's chance times the stream's copy rate: the rows its copies landed over those rows plus the chains that broke before their end (decayed 0.9 a round, primed from cfg.copy_rate's 0.94 as sixteen rows' worth). A copy counts only when its anchor row landed. On 134 greedy prompts (GSM8K, IFEval, HumanEval, MMLU, MGSM and a mixed set; 256 tokens, one stream), decode tok/s by domain went code 587 to 737, all 336 to 341, the others unchanged (gsm8k 482, ifeval 204-205, mgsm 318, mix 315, mmlu 334); every reply is unchanged. A fixed rate cannot do both: 0.99 gave code 738 but gsm8k 460 and mgsm 303.
…does The RTX 5090 verified a window past 16 rows as two 16-row tiles, each reading every weight: 32 rows of the 27B took 22.3 ms against 13.4 ms for 16. The 32-row tile (two m16 accumulators on each loaded panel, the narrow path's group FMAs and K-slice order) takes them in 17.8 ms. The startup width sweep measures it, so the lane planner widens code replies' context copies: on 134 greedy prompts, code decodes at 798 tok/s instead of 737, other domains within 1%, every reply unchanged. On the 5090 the device kernel contract's verify32 one-row bits gate finds all 18 shapes bit-equal to the one-row split-K decode it serves beside.
A grown tree took the hold's depth as its node count. A lone stream asks one draft when a context copy is ahead, and its own budget near a reply's end, so those rounds verified a one-node tree beside the copy and the allocation could admit no more than was grown. The block's forward costs the same at any depth and the allocation prices every node, so the tree now grows its whole budget whenever any draft is asked; the engine still admits what it prices. The draft equality test's hold hook checks the budget instead of the asked depth; its equality checks are unchanged and pass. On 134 prompts (25 each of GSM8K, IFEval, HumanEval and MMLU, 10 MGSM, a 24-prompt mix; 256 tokens, one stream), decode tok/s, every reply unchanged: greedy 343 to 354 (gsm8k 486 to 521, mmlu 339 to 352, code 793 to 807), at temperature 0.7 with fixed seeds 302 to 308 (gsm8k 422 to 442, code 663 to 653). Round times move under 1%.
The rope kernels ran a block a row: 16 blocks on a verify window, each rotating its 24 query heads one after another and then storing 4 key heads, two block reductions apiece for the e4m3 scales. A block now takes one query or key head (grid y: query heads, then key heads), the same arithmetic for every value: 7.6 to about 2 us a layer, and its shared staging is one head instead of 32. On 134 prompts (256 tokens, one stream), every reply unchanged: decode rounds 16.51 to 16.42 ms, greedy 354 to 356 tok/s, at temperature 0.7 with fixed seeds 308 to 310.
…ch, not one block a row Each sampled row launched tf_draw on one 1024-thread block in turn, about 97 us apiece: 1.5 ms of a 16-row round on an RTX 5090, so sampled rounds cost 17.6 ms where greedy ones cost 16.2. tf_draw_rows takes a run of rows in one launch, row r under rule[r] at meta[r] + 1, the arithmetic a one-row tf_draw at rule + r and meta + r does, so every token is unchanged. A kernel set built before it keeps the per-row launches.
…, then takes their softmax steps in order A warp walked its chunk one key at a time: the key's load, its warp sum, then the value's load, each waiting on the last, so a 128-key chunk was bound by load latency. The keys' rows and values now load eight at a time and their dots reduce together; the running max, sum and output still take each key in turn with the same operations, so every value is unchanged. On an RTX 5090 a 16-row round drops from 16.24 to 15.85 ms with bf16 rows and from 16.49 to 16.03 ms with e4m3 rows; four keys at a time gave 15.95, sixteen 15.91.
…a long match or out of the prompt early Every appended copy's rows were priced from cfg.copy_rate's 0.94, a prior worth sixteen rows that a stream's one or two copies never moved, so a lone window kept growing past 16 rows for copies that did not land. Traced over the 134-prompt set on Qwen3.8-27B: a copy after a context match of 20 or more tokens landed 85% of its rows (code's), one copied out of the prompt in the reply's first 64 tokens 89% on code (restating the asked-for function), and the rest 0 to 13%; outside code, rows past the 16th landed 11 of 777 priced. Those two kinds keep the prior; the others start from 0.08 with a history of their own, which trusted copies teach too (a stream sees too few of the others to learn them alone). The proposer reports where its last match's continuation starts (Proposer.lastEnd). Decode tok/s, three runs each, every reply unchanged: code 810 / 819 / 813 before and 803 / 817 / 813 after, gsm8k 525 to 533-534, mgsm 331-332 to 335, mmlu 354 to 357, rounds 16.06 to 15.92-15.95 ms.
…m's tokens leave them The drafter's logits came from the whole 248,320-row head, 450 us a round, where DFlash2's vocabulary (the Metal drafter's and the qwen27 engine's) is the head's first 98,304 rows and its last 288: 182 us, the first span a view of the head's tiles and the tail projected from the tile at 248,000. Drafting only those spans costs multilingual replies, whose tokens sit between them (MGSM 332 to 230 tok/s), so a stream drafts over the spans until its prompt or reply holds a token outside them, then over the whole head. A lane hold now carries the stream's tokens (lanes.drafter.Hold.tokens), and the drafter reads each stream's new ones once. Verification decides every token, so replies are unchanged: on the 134-prompt set all 364 to 367 tok/s (code 813-817 to 831, mgsm 335 kept), a round 15.92 to 15.73 ms.
… row for the sort and the race tf_draw_rows took each sampled row on one 1024-thread block passing over the whole vocabulary several times: 195 us a round, 16 blocks on a 170-SM GPU. Both of gpu_sample's candidate paths (a window near the max, or the 1,024 largest) reach the same thing for a top_k up to 1,024: the row's top_k by value then id, which its tail sums and races. So for a top_k up to 32 each 8192-entry slice now takes its own top_k (tf_draw_slice_top), and a block a row sorts the slices' candidates and runs the same tail (tf_draw_rows_finish); other rules keep the one-block draw. Every pick is unchanged: the community-v1 fixtures' sampled replies and the 134-prompt set sampled at 0.7 are byte-identical, and a sampled round drops from 15.59 to 15.39 ms (code at temperature 1 410 to 415 tok/s).
…utputs in one butterfly, its slots in shared A tree window's delta-rule scan reduced each value row's k . state sum and then its output with their own warp_sum, four butterflies a step for two rows, and kept its slots in registers that every step copied through. As the qwen27 engine's tf_gdn::tree_kernel does, a step's sums and the previous step's outputs now reduce in one spread_sum (the same lane pairs as warp_sum, so every sum is unchanged), the output written a step later, and the slots are the warp's rows of shared memory. Every reply is identical; on the community-v1 fixtures a round drops from 15.43-15.49 to 15.29-15.35 ms, and on the 134 prompts all 367 to 369 tok/s. The contract's static shared sizes follow.
…in 16-byte loads tf_q35_tap took a row on one 256-thread block of scalar loads (16 blocks, 9.3 us a launch) and tf_q35_mlp_act one entry a thread (2.0 us a launch, 70 a round). The tap now runs (rows, D / 2048) blocks and the activation 2048 entries a block, both eight entries a thread through uint4 loads and stores, each entry's arithmetic as before; the activation's 32 group sums a block still run their 64 entries in order. Every reply is identical; a round on the community-v1 fixtures drops from 15.29-15.35 to 15.23-15.29 ms.
…get draw's own Gumbel noise A sampled round accepts a draft when the target's draw picks it, and that draw is a Gumbel race keyed by the stream's seed, the position and the token. The drafts ranked the drafter's candidates by its logits alone, so at a temperature the tree followed the drafter's mode while the draw followed its noise. As the qwen27 engine's draft_tree.best_first does, a sampled hold now scores each row's candidates at the stream's temperature plus 0.7 times that draw's noise (lanes.gpu_rule.uniform24 at the draft's position), both for chains and lanes (choice) and for grown trees; a lane hold carries the stream's sampling (lanes.drafter.Hold.sampling). Verification decides every token, so replies are unchanged. On the community-v1 fixtures, every reply identical: code at temperature 1 198 to 195 rounds and 420 to 430 tok/s, chat at temperature 1 346 to 322 rounds and 236 to 250 tok/s.
…e activation takes four entries a thread For a recorded window of two rows or more, gdn_pre logs every row's conv tail and tf_q35_conv_commit returns at once, yet each of the 48 linear layers launched it: one window's rounds (the lone stream's) no longer do. The MLP activation now takes four entries a thread, 1024 a block (16 group sums, each in order), twice the blocks on an RTX 5090's 170 SMs. Every reply is identical; on the community-v1 fixtures a round drops from 15.25-15.50 to 15.15-15.40 ms (code at temperature 0 362 to 364 tok/s).
…base tf_draw_rows moves into sample_rows.cuh, and graftProspect into windows.zig beside the candidates() it reads; both unchanged. The verify-rows comment takes one line.
BobClawblaw
force-pushed
the
q35-decode-attn-tc
branch
from
October 21, 2026 22:04
d6260fe to
9a77cc7
Compare
Contributor
Author
|
Rebased with #1079 and #1080 onto 1.1.2, with no conflicts. One change in the second commit: the bf16 kernels' launch with their 64 KiB of dynamic shared memory moves from Checked on one RTX 5090 (Qwen3.8-27B MLX 4-bit + DFlash2, fresh servers):
|
…lines capped at 256 rounds The kept lone calibration (2c9871f) started every lone stream from the newest finished one's round times and trial deadlines. Two things in it let one long prompt slow every short reply after it until a restart: - the times carry no context length: a 253,039-token stream left 13 to 16-row rounds at 30 ms, and the short streams after it, whose 16-row round takes 15 ms, ran 10-row windows on those prices; - a losing trial (the one-row trial nearly always loses) at least doubles the spacing, and carried from stream to stream it grew without bound: 4,096 rounds after two minutes of soak, more after a long run, so no later stream re-timed the wide widths. Calibrations are now kept a context band each (under 8,192 tokens, then one a doubling, the last from 131,072 on), greedy and sampled apart, so a stream starts only from one measured at its own length and a long prompt leaves a short stream's calibration in place. A kept deadline or spacing reaches at most 256 rounds into the next stream. On one RTX 5090 (Qwen3.8-27B MLX 4-bit + DFlash2, fp8 caches, this stack's top with --prompt-precision int8): the quality suite's soak, then one greedy reply on a 252,783-token prompt, then the decode ruler's prose and code prompts. Before: 648-653 drafts proposed over 72 rounds (238 tok/s) and 379 over 41 (415 tok/s). After: 990 over 66-67 (253-257 tok/s) and 510 over 34 (495-496 tok/s), as on a fresh server. Replies are byte-identical either way.
…ule chain, the 2x2-warp prompt tile - prefill_attention_rows.cuh takes e4m3 cache rows with an fp32 scale each (F8: staged a step ahead and widened to bf16 in the slot) and head-major caches (HM: each key head's CAP rows in a run). hpattn_kernel and hpattn8_kernel instantiate them at head dim 256. pattn_kernel, tree_kernel, full_kernel and vision_kernel compile to the same SASS as before on sm_120. - gdn_prefill.cu: v0.6.6's prompt chain for the gated delta rule (two threads a value row, 32 steps staged), its strength bf16 or fp32. - qmm_prefill.cu adds the prompt tile on 2x2 warps and two stages; Affine4Kernels.half_prompt picks it.
…the delta-rule chain A prompt pass ran in 128-row chunks: each projection's GEMM launched about 80 blocks on a 170-SM GPU, attention ran the decode kernel's one-warp-a-row partials, and the delta rule stepped through rows one warp a value row. On one RTX 5090 (Qwen3.8-27B MLX 4-bit + DFlash2, fp8 KV, a 262,144-token window) a 75,616-token prompt took 144.5 s to its first token, against 26.6 s in the earlier qwen27 engine. - A prompt chunk is up to 4096 rows (Round.max_rows, the tables' and the activations' rows; the snap grid follows). Verify windows stay within most_verify, and the decode partials are sized for those rows only. - A one-part prompt chunk's attention runs hpattn(8)_kernel over the stream's own cache, and its delta rule gdn_prefill's chain, from pointers the host reads off the staged segment (prompt passes are not captured). - Prompt GEMMs take the 2x2-warp tile except on a GB10. - The drafters absorb a chunk's taps in pieces of most_verify rows, so their scratch does not grow with the chunk. The same prompt now reaches its first token in 26.9 s (56,143 tokens: 18.6 s); a needle at 40% of 60,000 tokens is found; decode rounds cost the same (16.13 to 16.20 ms). Prompt arithmetic changes with the kernels, so replies differ from 1.1.1's where its rounding decides a token; cache-on and cache-off passes still cut at the same rows.
… reading each key tile once tf_q35_attn_partial8 walked its chunk one key at a time in each warp (a warp a query head and row): a warp sum and a softmax step a key, every row and head reading the chunk again. At 109k cached keys a verify round spent 1.33 ms a layer there with about five rows, so long-context decode ran at half the qwen27 engine's speed. - tf_q35_attn_partial8: a block a (row, chunk, key head), the key head's query heads the lines of one m16 tile. Each 32-key tile of the row's logical keys is widened from e4m3 to bf16 in shared memory (exact); scores are mma dots of the codes times the key's scale, the softmax takes one step a tile, and P times each value's scale goes to bf16 for the value mma. It now serves streams under 8,192 committed keys. - tf_q35_attn_rows8 (qwen35_kv8_rows.cuh): a block a group of 64 / group rows (four m16 tiles), for streams of 8,192 keys and more. A stream's rows of one chunk width share each tile of committed keys; a tile holding window keys is taken row by row, its values accumulating with the other rows' weights zero. - Both kernels give a line the same arithmetic, so a row's partials are bit-identical alone or in any window, group or kernel: drafted replies still equal undrafted ones. tf_q35_attn_merge is unchanged. Replies change: attention outputs differ from the per-key kernel's by about 0.16% RMS (the weights round to bf16). Kernels alone (RTX 5090, 24 query and 4 key heads, both new launches), us: 109,000 keys at 6 rows 1445 -> 305, 10 rows 2361 -> 313, 16 rows 3634 -> 463; 300 keys at 6 rows 29.8 -> 10.6. Served (Qwen3.8-27B MLX 4-bit + DFlash2, fp8 caches, 262,144-token window, one stream, fresh servers the same day), decode tok/s before -> after: - a fixed 64,694-token prompt, 256 greedy tokens: 101.9 -> 180.4 (rounds 86 -> 65; the planner verifies wider windows once they are cheap); - longbench (a fresh nonce a run, so its reply and speed vary from run to run): 33,758 tokens 136.9 -> 203.1, about 104,610 tokens 87.8 -> 178.0; - short prompts, greedy: code 343.1 -> 344.0, chat 298.6 -> 299.9 (the sampled cells move within their unseeded run-to-run spread). Prefill is unchanged. Greedy replies match with and without the drafter (short prompts and the 64,694-token prompt), and three concurrent 30k-token streams each equal their solo replies. The contract records for both kernels were measured with nvcc 13.3; no 13.2 compiler was available.
…les shared tf_q35_attn_partial (--kv-cache bf16) still walked its chunk one key at a time in each warp, as the fp8 kernel did before the previous commit, so a bf16 cache decoded after a long prompt at half the fp8 cache's speed. - The previous commit's two kernels move into qwen35_attn_tc.cuh as templates over a tile loader: attn_one (a short stream's rows, a block each) and attn_rows (a long stream's rows in groups). Fp8Tiles is the e4m3 loader they had, unchanged: codes and scales into registers a tile ahead, widened into shared memory. Bf16Tiles copies a tile's rows with cp.async into the free one of two buffers of dynamic shared memory (64 KiB), a tile ahead; its scales are 1. - tf_q35_attn_partial and the new tf_q35_attn_rows are attn_one and attn_rows over Bf16Tiles, launched like the fp8 pair with their 64 KiB of dynamic shared memory (cuda_forward.zig); tf_q35_attn_partial8 and tf_q35_attn_rows8 are the same templates over Fp8Tiles. qwen35_kv8_rows.cuh is gone. - A row's partials are bit-identical alone, in any window, group or kernel, as for fp8. bf16 replies change: the weights round to bf16 for the value mma (outputs 1e-5 to 1.4e-4 RMS from the per-key kernel's on random caches). - fp8 is unchanged: both fp8 kernels' partials equal the previous commit's bit for bit (300 to 109,000 keys, 6 to 16 rows), and the served fp8 engine's greedy reply to a 64,811-token prompt is byte-identical. Kernels alone (RTX 5090, 24 query and 4 key heads, bf16 caches), us a launch, per-key -> both new launches: 300 keys 6 rows 27.4 -> 10.4; 8,190 keys 12 rows 253.5 -> 92.6; 33,000 keys 16 rows 1394.9 -> 174.4; 109,000 keys 6 rows 1975.4 -> 332.6, 16 rows 4742.2 -> 439.6. Served (Qwen3.8-27B MLX 4-bit + DFlash2, --kv-cache bf16, a 131,072-token window, one stream, fresh servers the same day), decode tok/s before -> after: - a fixed 64,811-token prompt, 256 greedy tokens: 96.9 -> 221.1 (rounds 91 -> 53); - longbench (a fresh nonce a run, so its replies and speeds vary): 34,094 tokens 112.5 -> 213.1, 104,864 tokens 73.4 -> 156.5 (an earlier build of the same kernels: 230.2 and 191.9); - short prompts, greedy: within their repeat-to-repeat spread (code 319-343 before, 303-353 after; chat 280-299, 282-300). Greedy replies match with and without the drafter (short prompts and the 64,811-token prompt), and three concurrent 30k-token streams on a 65,536-token window each equal their solo replies. The contract records were measured with nvcc 13.3; tf_q35_attn_partial's 1,024-thread record, which nvcc 13.3's registers no longer fit, is now 128 threads.
BobClawblaw
force-pushed
the
q35-decode-attn-tc
branch
from
October 22, 2026 09:29
9a77cc7 to
aac43f6
Compare
Contributor
Author
ashhart
added a commit
that referenced
this pull request
Oct 24, 2026
The lead's M54 review: the cost-open host test fails from the #1168 boundary on, because the admission refusal reads a geometry the test's probe never set. The refusal is right; the probe was wrong. The warm path's probe stands in for the target engine, whose geometry the real engine holds from the checkpoint's config (cuda_geometry.zig:248); the probe built its engine with `var engine: E.Engine = undefined` and set only `g.vocab` (cuda_warm_test.zig:56), so the guard read undefined query and KV heads. Under #1168's group bound (cuda_geometry.zig:148) that group is neither even nor in 2..16, and the one-row reservation the warm path opens (cuda_reserve.zig:16) is refused with error.UnsupportedQwenGeometry. Every served shape is inside the domain the guard states: the 27B's 24 query heads over 4 KV heads (cuda_geometry.zig:373), the 35B-A3B's 16 over 2 (cuda_geometry.zig:329), and the rest the repo carries (8/2, 32/8, 64/8, 4/2) all give an even group in 2..16. The refusal is not too strict; the probe had no geometry at all. The probe now reads the served target's query and KV heads (the 27B's 24 over 4, an even group the prompt tiles take) at cuda_warm_test.zig:58, and one case asserts the refusal still stands: a window past the partials' rows, and a group the prompt tiles cannot take, are refused at admission (qwen35_cost_open_host.zig:107). Before this change the step failed 1/3 from the #1168 boundary on (the #1080 boundary passed 3/3: the group bound arrived with #1168); after it the step passes 4/4 at every boundary.
ashhart
added a commit
that referenced
this pull request
Oct 24, 2026
…at admission Carried onto main: #1168's two records were written against the reviewed tree and replayed one commit at a time onto 8e49cf8. Their attention hunks name the forward's late refusal of the partials' rows as context, and that refusal was already gone: the rows are refused at admission since the earlier fix-up. The two records' own changes are unchanged; their context line is. The tensor-core decode attention narrows the partials' groups to 2 to 16, and the forward refused a group outside it late, after the embedding, the projections and the rotary writes had been queued. The supported domain now says so: cuda_geometry.zig's supported refuses a group outside 2 to 16 on the partials' path, and the forward's late refusal is gone. The host test refuses a group of 1 and a group of 32 at admission and admits a group of 2. Before this change the groups were admitted: the test found no error.
ashhart
added a commit
that referenced
this pull request
Oct 24, 2026
…t buffer again The attention merge kernel takes nine parameters: the geometry, the three partial buffers (the running maximum, the running sum and the output), the segment table, the row segment and ancestor tables, the parameters and the output. The launch #1168 added passes eight: the partial output buffer is left out, so the segment table lands on the output buffer's parameter and the kernel's own output pointer is read from past the argument list. The launch is refused with CUDA_ERROR_INVALID_VALUE before the model serves a row. The launch passes the partial output buffer again, between the running sum and the segment table, as the kernel declares it. The attention launches now name their kernel's parameter list, and a host test reads the .cu sources and checks every attention launch's argument count and kinds against the kernel signatures. A launch that drops the output buffer fails that test with the kernel that takes nine arguments; before this change it passed.
ashhart
added a commit
that referenced
this pull request
Oct 24, 2026
The stack's new symbols had no per-symbol entry: the row-sampling kernels, the prompt attention and delta-rule kernels, and the tensor-core decode attention were registered as one group or not at all, so the gates had no per-symbol spec, bound or served shape to hold them to. Each symbol now has its own spec, with its own name, its own rounding chain, its own tiles and its own served shapes: the three sampling kernels over a verify window's rows, the prompt kernels over a chunk's rows, and the decode kernels over a window's rows. The host test holds every symbol to its own spec: the spec names that symbol, its served shapes are not empty, and no two symbols share an entry. Before this change the symbols shared one entry: the test found a spec that named none of them.
Owner
|
Landed on main with the Qwen CUDA stack (through 6244c62), including the attention-merge launch fix (its kernel's nine arguments) that the first GB10 run caught, thank you @BobClawblaw. Measured on a GB10 with Qwen3.8-27B MLX 4-bit against main: cold prefill 834 to 1,281 tok/s at 2k (+54%), 612 to 1,266 at 8k (+107%), 326 to 1,165 at 32k (+258%), and earlier 203 to 1,045 at 64k (+416%). The decode sweep is +1.8 to +2.2% at 1 to 8 streams. Quality equals main's: 27 of 40 exact cases on both builds. Under our new-recipe ruling, the stack's own exactness holds: every stream equals its run alone, drafted equals plain, and resume equals fresh. - Synthema, Ash's AI Assistant. |
Owner
|
This shipped in TensorFold 1.1.4: https://github.com/ashhart/TensorFold/releases/tag/v1.1.4 Thank you @BobClawblaw for the contribution. - Synthema, Ash's AI Assistant. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on #1080 (which is on #1079): only the top two commits are new here.
e05dcbebmoves fp8-cache decode attention to tensor cores, andaac43f67does the same for bf16 caches with the same kernels.Why
With fp8 caches, decode after a long prompt on the qwen3_5 CUDA engine ran at about half the speed of the qwen27 engine it replaces: on one RTX 5090 (Qwen3.8-27B MLX 4-bit + DFlash2, 262,144-token window), 135 against 203 tok/s at a 34k-token prompt and 79 against 154 at 105k. Short prompts and prefill were already as fast or faster.
An Nsight Systems profile of one 512-token reply after a 109,532-token prompt put the gap in one kernel.
tf_q35_attn_partial8took 45% of the engine's GPU time: 1.33 ms a layer per verify round checking about four drafts, against 0.53 ms for the qwen27 engine's attention checking fifteen. The drafter's attention was 1%.The kernel walked its chunk one key at a time in each warp (a warp per query head and row). Every key cost a warp-wide sum and a softmax step before the next, and every row and query head of a key head read the chunk again. It was bound by that per-key chain, not by memory. Because verify rounds were expensive at long context, the lane planner also cut its windows to about four drafts a round, against fifteen with short prompts.
What changes
tf_q35_attn_partial8: one block per (row, chunk, key head), the key head's query heads the lines of one m16 tile.mma.syncdots of the codes, times the key's scale.mma.tf_q35_attn_rows8(new,qwen35_kv8_rows.cuh): one block per group of64 / groupforward rows (four m16 tiles), for streams of 8,192 committed keys and more.tf_q35_attn_mergeand the partial buffers are unchanged.cuda_forward.ziglaunches both kernels for fp8 caches. Each returns at once for the other's rows. The bf16-cache path is unchanged.Replies change. Unlike #1079's kernel changes, this one does not keep the per-key arithmetic: rounding the weights to bf16 moves attention outputs by about 0.16% RMS from the old kernel's (measured on random caches of 300 to 109,000 keys).
Measurements
Same RTX 5090, same day, fresh servers. Qwen3.8-27B MLX 4-bit + DFlash2,
--kv-cache fp8, a 262,144-token window, one stream. Base is2bd805a4(#1080's head). Decode tok/s:longbench, 512 tokens)longbench, 512 tokens)bench_openai.py --tokens 64 --reps 5)longbenchputs a fresh nonce in each prompt, so its replies and speeds vary between runs. The fixed prompt is deterministic: three runs on these kernels gave 180-185 tok/s, against 97-102 before.The kernels alone, on random caches with 24 query and 4 key heads (µs per launch, both new launches together):
Checks
longbenchneedles are found at 34k and 105k.zig build test,test-kernel-contract,test-golden(26/26),fuzz-cortheon,zig fmt --checkand the lean check pass.kernel-contract-limitspasses both kernels on sm_80, 86, 89, 120 and 121.tf_q35_attn_partialat its recorded 1,024 threads (122 registers). The second commit replaces that kernel; it now launches 128 threads.bf16 caches (
d6260feb)tf_q35_attn_partial(--kv-cache bf16) had the same per-key walk, so a bf16 cache decoded after a long prompt at about half the fp8 cache's speed.qwen35_attn_tc.cuhas templates over a tile loader:attn_one(a short stream's rows) andattn_rows(a long stream's rows, grouped).Fp8Tilesis the first commit's e4m3 loader, unchanged.Bf16Tilescopies a tile's rows withcp.asyncinto the free one of two buffers of dynamic shared memory (64 KiB), a tile ahead; its scales are 1.tf_q35_attn_partialand the newtf_q35_attn_rowsare the templates overBf16Tiles, launched like the fp8 pair (Engine.runShared).tf_q35_attn_partial8andtf_q35_attn_rows8are the same templates overFp8Tiles.qwen35_kv8_rows.cuhis gone.Kernels alone (bf16 caches, 24 query and 4 key heads), µs per launch:
Served with
--kv-cache bf16, a 131,072-token window, one stream, same day, fresh servers. Decode tok/s:longbench)longbench)longbench's nonce makes its numbers vary; an earlier build of the same kernels measured 230.2 and 191.9.zig build test,test-kernel-contract,test-golden,fuzz-cortheon,zig fmt --checkand the lean check pass.kernel-contract-limitspasses all four attention kernels on sm_80, 86, 89, 120 and 121 (records from nvcc 13.3).