Repository navigation
glm cuda: a decode step's select gives each row a 1,024-thread block - #1424
Conversation
|
Updated the description: it now opens with the measured effect on my hardware, prompt and decode before and after against main, taken from the receipt below. No code change. |
|
Thanks, Bizuayeu. The argument holds: with digit counts from per-warp histograms and the picks leaving through an On main One change, please. We'll then run Heads-up: #1449 touches the same contract files, so whichever of the two lands second will need a small rebase. |
Added in 351ede1, in
Static shared bytes are 1,072 and the bound 128 on every SM under both, so 13.3's set needs no entry for it. With the record in, both runs read it as current on all five SMs, and each run's problem count drops by exactly those five lines (13.0: 87 to 82, 13.3: 564 to 559); the rest are other launches' records. On 351ede1: If #1449 lands first, I'll rebase this one onto it. |
|
Thanks: the 128-thread The remaining gate is ours: GLM-5.3-Flash NVFP4 at TP=2 on our GB10 pair, main against this branch, with token hashes, concurrency and the prompt's final short chunk (the wide rule reaches it too). It's queued behind a long correctness run on the pair, so expect it later this week. |
… tf_glm_scores and tf_glm_select_rows (TF_GLM_INDEX_CUDA) The kernels are ported from BadAd84's Python patches 0101 (one-pass top-K pools) and 0105 (pool scores), for #1430. tf_glm_index_scores keeps tf_glm_scores' arithmetic exactly (the mma chain over the 8 k16 steps from zero, the relu, the head sum and its shuffle order) and changes only the data movement: a CTA of 8 rows shares each staged 64-pool key tile instead of every row staging its own. tf_glm_index_pools replaces tf_glm_select_rows for a prompt block: a strided-sample floor, one order-preserving compaction into shared memory, then the same radix select and ascending write over the candidates, or over the whole row when the candidates miss [512, cap]; the candidate set holds the 512 best and every tie at the 512th, so the pools are the select's, ties to the lower pool. TF_GLM_INDEX_CUDA=1 takes both for chunks of 64 or more rows, =scores only the scores; unset, nothing changes. A build without the glm_index fatbin logs and keeps the existing launches. Decode and #1404/#1424's areas are untouched; the hook is two branches in selectTokensProfile and an optional field of glue.Set. Host: tools/glm_index_emulate runs both kernels as host C++ (real barriers, warp exchanges, ldmatrix and mma models) against plain models, in test-glm-index-host and test-cuda-host. GPU: tf-cuda-test glm-index-check compares selectTokens with and without the kernels byte for byte at 1k-262k pools, plus tied scores.
|
#1449 (the GLM prompt indexer) landed on main, and this branch now conflicts with it in The engine change itself merges cleanly, so our GB10 pair run (main against this branch, as above) is queued on a merge that keeps main's contract files. The contract run will follow on your rebased head. |
351ede1 to
557221b
Compare
Rebased onto main (a1e6d34): the same three commits, now 8108f9f, 40aed7e and 557221b. The engine change didn't conflict: with On the rebased head, |
|
Our GB10 pair run is done: GLM-5.3-Flash NVFP4 at TP=2 on two GB10s, main against your select change merged onto main (its kernel and engine files are byte-identical to your rebased head), alternating main, this branch, this branch, main, with a 65,536-token window and two streams.
The gain you measured after a 136k-token prompt is beyond this run's window, so it rests on your receipt, which fits the launch timings. What's left on our side is the kernel contract on a GB10 for your rebased head, once the pair finishes its current run; then we land it. Thank you, @Bizuayeu. |
tf_glm_select_rows runs a 128-thread block a row. A decode step has a few rows, so over a long context's bucket four warps a row read every pool while most SMs idle. tf_glm_select_rows_wide runs the same radix select with 1,024 threads: T only sets how many pools a pass reads (8 a thread), the digits are counts and the picks leave through an index-order scan, so every T gives the same pools in the same order. cuda_sparse launches it for at most 48 rows (one wave on GB10's 48 SMs) over buckets of 4,096 pools or more. glm-ref checks both kernels' bytes and the host's rule from the 1,024- to the 131,072-pool bucket at 1 to 128 rows, and prints both launches' times.
…r byte The wide select registers as a replacement of the 128-thread kernel (its bytes on the same scores), with a stated NaN rule checked against the host's picks at gate 1: order keys descending, -0 as +0 and any NaN above +inf, as the GPU's s + 0 gives one positive NaN. The index scores and the select register through one owned group, so the runner stays within the source-size limit, as the Gemma tower's do. Gate 5's records come from kernel-contract-limits over sm_80 to sm_121 fatbins of nvcc 13.0 and 13.3; 13.3 differs on sm_120 and sm_121.
…ged kernel It is now select_row<128>, and it still serves every prompt chunk and every decode step outside the wide rule, so gate 5 checks it as it checks the wide form. The numbers come from kernel-contract-limits over sm_80 to sm_121 fatbins of this tree: nvcc 13.0 and 13.3 give the same registers and static shared bytes on all five SMs, so 13.3's set needs no entry for it.
557221b to
b8f755d
Compare
|
Thanks @Bizuayeu, this is merged. The kernel contract on a GB10 (sm_121) passes for your rebased head: the wide select is registered as a byte-for-byte replacement of |
|
Thanks @Bizuayeu: the 1,024-thread DSA select shipped in TensorFold 1.1.9. |
What this changes
A decode step's DSA select at long positions now gives each row a 1,024-thread block, so decode over a long context
gets faster; short prompts, the prompt and the select's bytes do not change. On my hardware (three GB10 at TP=3,
GLM-5.3-Flash XSFP4,
--parallel 4):1b4f2f979bench_cache.py --tokens 256, mean of six)fibonacci-rawttft_s(mean of six)With drafts off, a build that switches between the two kernels gave +2.4% after the same prompt (28.06 → 28.74 tok/s,
every wide run faster). Details in the receipt.
GLM-5.3-Flash's DSA select (
tf_glm_select_rows) gives each row a 128-thread block, so a decode step's few rows overa long context read every pool on a handful of SMs.
tf_glm_select_rows_wideruns the same radix select with 1,024threads a row and gives the same bytes;
cuda_sparselaunches it for at most 48 rows over buckets of 4,096 pools ormore. It is registered with the kernel contract as a replacement of
tf_glm_select_rows.are counts, and the picks leave through an index-order block scan, so every block size gives the same pools in
the same order. The 128-thread entry point is the same template at 128.
128-thread kernel is as fast or faster; past 48 rows the wide form needs a second wave and the gap narrows, so the
rule keeps the 128-thread kernel there (it stays served, and as
select_row<128>it has its own launch record).Times of one launch on a GB10 (
tf-cuda-test glm-ref, which also checks both kernels' bytes and the host's rule atevery row):
In a profiled decode step at 136k tokens the 128-thread select took 1.44 ms of the 35.3 ms step (about 11 launches),
so the wide form should save about 1 ms a step there.
Receipt
91bd6bef5against main1b4f2f979(head351ede169adds only the 128-thread select'slaunch record); three GB10 (DGX Spark class, sm_121, 128 GB unified
memory) at TP=3 over a switchless RoCE ring; CUDA driver 580.178.04, nvcc 13.3.1 for the served build and nvcc
13.0 and 13.3 for gate 5; GLM-5.3-Flash XSFP4 (
Bizuayeu/GLM-5.3-Flash-XSFP4at88736f9). Each build servedwith
--parallel 4 --prompt-cache-gib 0, launched in the order main, this head, another branch twice, this head,main.
bench_concurrent.py --alone --serialat levels 1 and 4 (one rep, not the example's three, to fit thewindow): every cell equal alone and serial on both builds.
token_sha256is the same in every launch for bothdecode prompts and all three phases of the 136k prompt. Before the PR, on a measuring build with the same kernel
behind a switch, a 136,182-token prompt's 600-token reply (drafts off) had the same
token_sha256in all eightruns, wide and 128-thread.
bench_openai.py --tokens 256 --reps 3 --temperatures 0, medians): main 57.58 and57.55 tok/s, this head 57.58 and 57.50 on
fibonacci-raw; main 44.83 and 44.77, this head 44.83 and 44.75 ongpu-chat-no-think. After the 136k prompt (bench_cache.py --tokens 256, its three phases): main 57.75 tok/smean (57.58 to 57.95, n 6), this head 58.80 (57.80 to 59.21, n 6), +1.8%. The measuring build above (drafts off)
gave 28.73, 28.71, 28.77 and 28.76 tok/s wide against 28.07, 28.06, 28.05 and 28.06 with the 128-thread kernel
(+2.4%, alternated, every wide run faster).
bench_cache.py --require-cold --streams 1on a 136k prompt; all three phases arecold at
--prompt-cache-gib 0): main 87.91 s mean (87.25 to 88.62, n 6), this head 87.69 s (87.34 to 88.71, n 6).A prompt's last chunk can have 48 rows or fewer, so its select can take the wide form.
kernel-contracton a GB10 (nvcc 13.3, sm_121, the runner at the build's default):SUMMARY 341 kernels: 341 pass gates 1-4, 0 fail, 5 known findings; 94 controls: 94 red as planted, 0 missed; 231 limit problems. The wideselect passes nonfinite, repeat and shapes ("8 shapes bit-equal to the kernel it replaces") and declares no sums;
its live limits match the record ("1024 threads launched, the driver allows 1024; 62 registers, 1296 + 0 shared
bytes (recorded 62, 1296)"). None of the limit problems names it: they are other kernels' live registers against
their records, and the records against sm_80, 86, 89 and 120.
kernel-contract-limitsover fatbins for sm_80, sm_86, sm_89, sm_120 and sm_121): the wide select fitson all five with nvcc 13.0 (32, 61, 61, 54 and 61 registers, 1,296 static shared bytes, no local memory; sm_80
keeps an 8-byte stack frame at 32 registers) and with nvcc 13.3 (62 registers on sm_120 and sm_121). The
128-thread select fits on all five with nvcc 13.0 and 13.3 alike (45, 54, 54, 48 and 48 registers, 1,072 static
shared bytes), and with its record both read it as current.
tf-cuda-test glm-ref: pass (the 12 select cases above, byte-equal and the host's rule).zig build test-kernel-contract111/111,zig build compile-cuda-tests10/10,zig fmt --check, the leancheck (no new problems against main),
zig build test4,249 of 4,266 passed, 17 skipped, none failed.351ede169:zig build test-kernel-contract111/111,test-glm-contract-host7/7,compile-cuda-tests10/10,
zig build test4,248 of 4,266 passed, 17 skipped, 1 failed (cuda_cost_identity, which fails on mainthe same way on that host's coarse file times),
zig fmt --checkand the lean check (no new problems againstmain).
sm_120 have gate 5's compile records only).
zig build test-gradient-serving, which already fails on main at itsfirst stale pin (
zig/kernels/cuda/contract_controls.cu); its pin forglm_glue_index.cuhwas already stale too.Checklist
select's launch changes.)
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.(Concurrent equals solo in the receipt, and the select's bytes are the 128-thread kernel's.)
select's bytes are the 128-thread kernel's; the contract checks them.)
SUMMARYline is in the receipt. (CUDA only.)entry and the glm-ref cases need the wide kernel; without a GPU the runner is not run.)