Repository navigation
glm cuda: a prompt chunk's DSA indexer in two CUDA kernels, byte-equal to tf_glm_scores and tf_glm_select_rows (TF_GLM_INDEX_CUDA), 1M-token first token 1,223 -> 1,033 s - #1449
Conversation
… 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 ashhart#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 ashhart#1404/ashhart#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.
…f_glm_select_rows; a fuller device check (ashhart#1430) Contract: launch records for tf_glm_index_scores and tf_glm_index_pools (regs estimated until nvcc measures them) and specs that must be bit-equal to tf_glm_scores (on a shuffled page table) and tf_glm_select_rows. Kernel: when no radix bin reaches need, the digit is 0 and need drops by every pool, as in tf_glm_select_rows; CAP is clamped to the shared-memory size; takes() needs at least 512 pools. Cleanups: unused constants, topk and the score loop threshold from sparse, the pool block named, the host emulator in its own step (test-glm-index-emulate). Device check: a paged case, heads below 32, NaN and signed zeros, the scores-only mode, every score block of a multi-block chunk, and glm_index_pools' paths counted (whole row and candidates must both occur).
…sts.zig stays within 600 lines (ashhart#1430)
…each read, and refuses a path outside 0..2 (ashhart#1430)
…he launches, and prints the first differing words (ashhart#1430)
… each and within 120 columns (ashhart#1430)
Stacked on #1404 and #1424 (XSFP4), your kernels take a 516,096-token prompt's first token from 423.0 to 398.1 s at TP=2 (6.3% faster), and make 131,072- and 136,000-token prompts about 1.3% faster at TP=2 and TP=3. The tokens are the same in every launch. One binary, launched with
At these lengths every Past position 2,048, #1404 already leaves each rank half (TP=2) or a third (TP=3) of a chunk's indexer rows, so at 131,072 tokens the gain sits below your -2% alone; on this checkpoint at TP=3 the indexer was 7.4% of a 131,072-token prompt before #1404. At 516k the indexer's share is larger, as in your receipts.
I haven't re-measured #1404's split threshold (2,048) on top of your kernels; the 8,192- and 32,768-token prompts above, both split, show no loss. |
|
The full kernel contract, on an idle GB10 (sm_121, nvcc 13.0), with this branch's tree:
@Bizuayeu, thank you for the stacked run. It's good to see #1404's split and these kernels compose cleanly at both TPs, with the same tokens. Your nvcc 13.3 build is where the 13.0-only records would show as stale, so a |
|
Thanks @wojo, this is a clean port, and the receipts are thorough. We read both kernels against
fmt, lean, Before landing:
Since it's byte-equal and a pure prompt speedup, we may switch it on by default in a follow-up. |
Ran it on this branch's head (36b9542), with nvcc 13.3 fatbins for sm_80, sm_86, sm_89, sm_120 and sm_121. Both launches fit on all five SMs. On sm_80, sm_86 and sm_89 the numbers equal the records; on sm_120 and sm_121 the registers differ, so those four lines read stale:
Shared bytes (32,768 and 33,864) and the bound (256) are the same as recorded on every SM. Entries in the .{ .name = "glm_index tf_glm_index_scores", .usage = .{ .{ .regs = 128, .static = 32768, .bound = 256 }, .{ .regs = 128, .static = 32768, .bound = 256 }, .{ .regs = 128, .static = 32768, .bound = 256 }, .{ .regs = 126, .static = 32768, .bound = 256 }, .{ .regs = 126, .static = 32768, .bound = 256 } } },
.{ .name = "glm_index tf_glm_index_pools", .usage = .{ .{ .regs = 62, .static = 33864, .bound = 256 }, .{ .regs = 64, .static = 33864, .bound = 256 }, .{ .regs = 64, .static = 33864, .bound = 256 }, .{ .regs = 54, .static = 33864, .bound = 256 }, .{ .regs = 54, .static = 33864, .bound = 256 } } },The run reports 194 problems in all; these four stale lines are the only ones that name On point 2 there: my stacked tree resolved |
The check asserted the whole-row and candidate paths of tf_glm_index_pools but not the whole-row retry after the candidates miss; it now requires all three.
nvcc 13.3 compiles tf_glm_index_scores at 126 registers and tf_glm_index_pools at 54 on sm_120 and sm_121, where 13.0 has 124 and 59. The other three SMs match 13.0. Measured by Bizuayeu with nvcc 13.3 on sm_80, sm_86, sm_89, sm_120 and sm_121.
|
Thanks @ashhart, and thanks @Bizuayeu for the 13.3 run. I pushed 82f0a80, two commits on 36b9542.
Host tests on the new head: If you'd like it on by default, I'm happy to send that follow-up once this lands. |
|
Our kernel contract run on a GB10 (nvcc 13.0, your head before the 13.3-records commit): |
|
Landed on main in 8a789e3, with your ten commits kept as they were. Thank you! It's opt-in with TF_GLM_INDEX_CUDA=1, as in your description, and the changelog carries your 131k/524k/1M first-token numbers. |
|
This is in TensorFold 1.1.8, released today. Thank you, @wojo! |
What this changes
Fixes #1430.
A prompt chunk's DSA indexer can now run on two new CUDA kernels in
glm_index.cu, behindTF_GLM_INDEX_CUDA(off by default):tf_glm_index_scorescomputes the pool scores;tf_glm_index_poolsdoes a one-pass top-512 pool select: a strided-sample floor, a compaction of the survivors, then a radix select over them.Their output is byte-equal to
tf_glm_scoresandtf_glm_select_rows: the same scores, and the same pools, with ties going to the lower pool and picks in pool order. When the switch is unset, the kernels fail to load, or a launch is outside their shapes (Kernels.takes), the current kernels run as before.The kernels port BadAd84's Python patches 0105 (pool scores) and 0101 (the one-pass select) from Mia's GLM recipe.
TF_GLM_INDEX_CUDA=1: scores and pools.=scores: scores only, with the pools fromtf_glm_select_rows. Unset or0: off. Any other value stops the model load (BadIndexSetting) rather than silently running without it.selectTokensProfile(families/glm/cuda_sparse.zig). It takes launches of 64 rows or more over at least 512 pools, so prompt chunks.glm_cuda_records.zig,contract_glm_index.zig,contract_glm_pools.zig). Their records are nvcc 13.0's measured registers on sm_80/86/89/120/121, plus nvcc 13.3's on sm_120/121 where they differ (scores 126, pools 54), measured by @Bizuayeu.tf-cuda-test glm-index-checkruns both paths on the device and compares every byte: random and tied scores, pool counts from 1k to 262k, and each pools path (whole row, candidates, and the whole-row retry after the candidates missed), and requires all three paths to occur.zig build test-glm-index-emulate(tools/glm_index_emulate) builds the two kernels'.cuhas host C++ against a small CUDA shim. It checks them against a plain top-K and a model oftf_glm_scoreswithout a GPU (about 45 s, its own step).test-glm-index-hostcovers the launch rules.Stacking with #1404 and #1424. As @Bizuayeu laid out on #1430, these kernels make each row cheaper; #1404 splits a chunk's rows across ranks, and #1424 widens decode's select. All three give
tf_glm_select_rows's bytes. They meet inselectTokensProfileand the contract records, so whichever lands second rebases there and re-measures its thresholds on top.Receipt
Environment: the receipts ran on
main8c36359 plus this branch at d7551bc32. The two later commits change only the contract records and comments. The branch was then rebased ontomainb30d5b9, which changed one line:glm_indexjoinscontract_limits.zig's image list next to main's newnvfp4_experts_tile. Two NVIDIA DGX Spark (GB10, sm_121), TP=2, CUDA 13.0 / nvcc 13.0, Ubuntu 24.04. GLM-5.3-Flash EXL3 4 bpw (078455ff), MTP drafts,--context 1048576 --parallel 8, prompt cache default.Exactness, prompt: needle prompts recorded with the switch unset, then replayed with
TF_GLM_INDEX_CUDA=1. Same prompt tokens, same text and sametoken_shaat all three sizes; the needle was found in both arms.Exactness, decode:
tools/bench_concurrent.py --levels 1,4,8 --tokens 256 --alone --serial, temperatures 1 and 0, promptscodeandchat, at--context 131072 --parallel 8(same tree; decode does not depend on the window). With the switch set: 156/156 concurrent replies equal their alone replies, and 54/54 alone replies equal"draft": false. 0 unequal, 0 failed. The switch-unset arm gave the same counts.Decode speed before and after (
tools/bench_openai.py, 256 tokens, 3 reps, median tok/s), unchanged within noise (decode steps here stay under the 64 rows the new kernels take):bench_concurrentaggregate tok/s at 8 streams: code 165.8 → 166.1 and chat 99.8 → 99.9 (temperature 0); code 72.5 → 72.6 and chat 63.8 → 63.9 (temperature 1).Prompt speed before and after: first token on a cold prompt, same server, switch unset vs
=1:The gain grows with length because the indexer's share grows with position. Lowest MemAvailable during the runs: 10/12 GiB (before) and 11/12 GiB (after) on ranks 0/1. Neither rank logged an error or a fallback.
Kernel check on the device (
glm-index-check, GB10, Debug and ReleaseFast, two runs each):PASS glm-index-check: scores, pools and token lists byte-equal at 10 cases, tied scores through both selects. At 262,144 pools, one block: 30.8 → 14.4 ms. On the current head, which also requires the retry path, two more runs pass with the paths at whole row 2,048, candidates 6,726, retry 1,466.Kernel contract (
tf-cuda-test kernel-contract, idle GB10, nvcc 13.0):SUMMARY 416 kernels: 416 pass gates 1-4, 0 fail, 6 known findings; 311 controls: 311 red as planted, 0 missed; 180 limit problems.tf_glm_index_scores: repeat, shapes and sums pass; nonfinite is the same KNOWN finding astf_glm_scores(relu isfmaxf(x, 0), so a NaN dot scores 0).tf_glm_index_pools: nonfinite, repeat and shapes pass; sums N/A.glm_index: the 12 live ones areqwen35launches recorded at 0 registers. Details in the comment below.Tests (macOS, no GPU):
test-kernel-contract+test-glm-contract-host117/117 (on b30d5b9);test-glm-index-host9/9;test-glm-index-emulate,test-glm-fp8-kv-hostandtest-glm-prompt-profile-hostpass. On the current head (82f0a80),test-kernel-contract,test-glm-contract-host,test-glm-index-host9/9 andtest-glm-index-emulatepass, andlean_check.pyreports 428 problems before and after (adds 0).Not run, and why: I measured the records for sm_80/86/120/121 with
kernel-contract-limitsand for sm_89 withcuobjdump -res-usage, but I could only run the kernels on sm_121 (GB10). I built only with nvcc 13.0. The 13.3 records are @Bizuayeu'skernel-contract-limitsrun on all five SMs. Nobody has measured nvcc 13.2, so a 13.2 build checks these launches against their 13.0 records and may show them as stale. Metal is unaffected: the change is CUDA-only.Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.SUMMARYline is in the receipt.Credit to BadAd84 for the Python kernels (patches 0101 and 0105) these port.
By Robert Wojciechowski (wojo on GitHub, @rwojo on X)