Skip to content

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

Closed
wojo wants to merge 10 commits into
ashhart:mainfrom
wojo:glm-index-cuda

Conversation

@wojo

@wojo wojo commented Oct 27, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

Fixes #1430.

A prompt chunk's DSA indexer can now run on two new CUDA kernels in glm_index.cu, behind TF_GLM_INDEX_CUDA (off by default):

  • tf_glm_index_scores computes the pool scores;
  • tf_glm_index_pools does 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_scores and tf_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 from tf_glm_select_rows. Unset or 0: off. Any other value stops the model load (BadIndexSetting) rather than silently running without it.
  • The hook is in selectTokensProfile (families/glm/cuda_sparse.zig). It takes launches of 64 rows or more over at least 512 pools, so prompt chunks.
  • Both kernels are registered with the kernel contract (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-check runs 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' .cuh as host C++ against a small CUDA shim. It checks them against a plain top-K and a model of tf_glm_scores without a GPU (about 45 s, its own step). test-glm-index-host covers 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 in selectTokensProfile and the contract records, so whichever lands second rebases there and re-measures its thresholds on top.

Receipt

  • Environment: the receipts ran on main 8c36359 plus this branch at d7551bc32. The two later commits change only the contract records and comments. The branch was then rebased onto main b30d5b9, which changed one line: glm_index joins contract_limits.zig's image list next to main's new nvfp4_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 same token_sha at 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, prompts code and chat, 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):

    Prompt Temperature Before After
    fibonacci-raw 0 49.1 48.8
    gpu-chat-no-think 0 27.5 27.4
    fibonacci-raw 1 32.4 32.3
    gpu-chat-no-think 1 27.4 27.3

    bench_concurrent aggregate 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:

    Prompt tokens Before After Change
    131,065 111.3 s 108.8 s -2%
    524,231 519.4 s 471.7 s -9%
    1,039,995 1,223.4 s 1,032.7 s -16%

    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 as tf_glm_scores (relu is fmaxf(x, 0), so a NaN dot scores 0).
    • tf_glm_index_pools: nonfinite, repeat and shapes pass; sums N/A.
    • Live limits pass for both at their records (124 and 59 registers).
    • None of the 180 limit problems names glm_index: the 12 live ones are qwen35 launches recorded at 0 registers. Details in the comment below.
  • Tests (macOS, no GPU): test-kernel-contract + test-glm-contract-host 117/117 (on b30d5b9); test-glm-index-host 9/9; test-glm-index-emulate, test-glm-fp8-kv-host and test-glm-prompt-profile-host pass. On the current head (82f0a80), test-kernel-contract, test-glm-contract-host, test-glm-index-host 9/9 and test-glm-index-emulate pass, and lean_check.py reports 428 problems before and after (adds 0).

  • Not run, and why: I measured the records for sm_80/86/120/121 with kernel-contract-limits and for sm_89 with cuobjdump -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's kernel-contract-limits run 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

  • Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
  • Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
  • No precision traded for speed. The bits do not change: scores and pools are byte-equal to the current kernels.
  • A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
  • Prompt processing is no slower than the last release.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • New tests fail before the change, pass after it, and skip cleanly without their dependency.
  • Comments and docstrings are one line. No measurements or history in the source.
  • No personal data, machine names, internal hosts or local paths. No AI attribution lines.

Credit to BadAd84 for the Python kernels (patches 0101 and 0105) these port.

By Robert Wojciechowski (wojo on GitHub, @rwojo on X)

wojo added 8 commits October 27, 2026 05:49
… 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).
…he launches, and prints the first differing words (ashhart#1430)
@Bizuayeu

Copy link
Copy Markdown
Contributor

A stacked measurement at TP=3 and TP=2 would be very welcome.

#1430 (comment)

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 TF_GLM_INDEX_CUDA unset, then =1 twice, then unset again (means over each arm's two launches and their phases):

cold first token, ttft_s #1404 + #1424 + #1449
TP=2, 516,096 tokens 423.02 398.10 +6.3%
TP=2, 131,072 tokens 97.30 96.06 +1.3%
TP=2, 136,000 tokens 101.23 99.98 +1.3%
TP=3, 131,072 tokens 80.25 79.19 +1.3%
TP=3, 136,000 tokens 83.54 82.53 +1.2%

At these lengths every =1 launch beat every unset one. At TP=3 the 8,192- and 32,768-token prompts moved +0.6% and +0.2%, within the runs' spread. Decode (bench_openai, and after the 136,000-token prompt) did not change (-0.4% to +0.1%). No launch waited at the heat guard.

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.

@wojo

wojo commented Oct 27, 2026

Copy link
Copy Markdown
Contributor Author

The full kernel contract, on an idle GB10 (sm_121, nvcc 13.0), with this branch's tree:

SUMMARY 416 kernels: 416 pass gates 1-4, 0 fail, 6 known findings; 311 controls: 311 red as planted, 0 missed; 180 limit problems; gates 2138199 ms, limits 1459728 ms
  • The two new launches:
    • tf_glm_index_scores: repeat, shapes and sums pass. nonfinite is the same KNOWN finding as tf_glm_scores (relu is fmaxf(x, 0), so a NaN dot scores 0).
    • tf_glm_index_pools: nonfinite, repeat and shapes pass; sums N/A.
    • Live limits pass for both: 124 and 59 registers, 32768 and 33864 shared bytes, equal to their records.
  • The 180 limit problems are not in these launches. None of the log's limit lines names glm_index.
    • The 12 live ones are all qwen35 / qwen35_kv8 launches whose records read 0 registers and 0 shared bytes (tf_q35_attn_partial*, tf_q35_rope*, tf_q35k_compact*).
    • The log gives only the count for the 168 static ones (1300 registered CUDA launches ... 168 problems), not which launches they are.
    • I haven't run main alone to compare, so I can't say whether main gives the same count.

@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 kernel-contract-limits run there would settle that if you have the time.

@ashhart

ashhart commented Oct 27, 2026

Copy link
Copy Markdown
Owner

Thanks @wojo, this is a clean port, and the receipts are thorough. We read both kernels against tf_glm_scores and
tf_glm_select_rows, and agree they're byte-equal by construction:

  • Scores: the same MMA fragments, k order, epilogue and reduction.
  • Pools: a floor that only chooses candidates, kept in pool order, with the whole-row fallback.

fmt, lean, zig build test, test-glm-index-host, test-glm-index-emulate, test-kernel-contract and the aarch64
CUDA cross-compile all pass here.

Before landing:

  1. The kernel-contract SUMMARY line. We'll run it on an idle GB10 with nvcc 13.0 and 13.3. We'll also run
    kernel-contract-limits for the 13.3 records, unless you'd rather do both.
  2. The conflict with glm cuda: a decode step's select gives each row a 1,024-thread block #1424 in contract_tests.zig. glm cuda: a decode step's select gives each row a 1,024-thread block #1424 replaces the glm_scores registration with
    glm_index_k.Index, so keep that and add only your two Fast lines. Keeping both sides registers tf_glm_scores
    twice. Resolved that way, main + glm cuda: a prompt chunk's indexer rows split across the ranks, their picks all-gathered #1404 + glm cuda: a decode step's select gives each row a 1,024-thread block #1424 + this branch builds the CUDA tests for aarch64 and passes the
    contract host tests, with contract_tests.zig at 597 lines.
  3. NVFP4. We'll add a short NVFP4 run at TP=2 with the switch unset and set, checking that token_sha256 matches.
    Your runs and Bizuayeu's cover EXL3 and XSFP4.
  4. A small one. glm-index-check asserts that the whole-row and candidate paths occur, but not the retry after the
    candidates miss (glm_index_check.zig:168). Since the description counts all three, could it assert
    paths[2] > 0 too?

Since it's byte-equal and a pure prompt speedup, we may switch it on by default in a follow-up.

@Bizuayeu

Copy link
Copy Markdown
Contributor

Your nvcc 13.3 build is where the 13.0-only records would show as stale, so a kernel-contract-limits run there would settle that if you have the time.

#1449 (comment)

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:

registers a thread sm_80 sm_86 sm_89 sm_120 sm_121
tf_glm_index_scores, 13.3 128 128 128 126 126
its record (13.0) 128 128 128 124 124
tf_glm_index_pools, 13.3 62 64 64 54 54
its record (13.0) 62 64 64 59 59

Shared bytes (32,768 and 33,864) and the bound (256) are the same as recorded on every SM. Entries in the cuda_133 set like these would make all ten lines current:

.{ .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 glm_index. These are the 13.3 records in point 1 of #1449 (comment), in case they save that run.

On point 2 there: my stacked tree resolved contract_tests.zig that way, keeping #1424's glm_index_k.Index and adding only your two Fast lines (597 lines, tf_glm_scores registered once). "Keeping both sides" in my earlier comment was loose for that file; it holds for docs/kernel-contract.md and glm_cuda_records.zig, where both sides only add lines.

wojo added 2 commits October 27, 2026 12:02
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.
@wojo

wojo commented Oct 27, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @ashhart, and thanks @Bizuayeu for the 13.3 run. I pushed 82f0a80, two commits on 36b9542.

  1. SUMMARY. The full nvcc 13.0 run is in my comment above: 416 kernels pass, 0 fail. None of its 180 limit problems names glm_index. The branch now carries @Bizuayeu's 13.3 records in cuda_133 (sm_120/121: scores 126, pools 54; the other three SMs equal 13.0). docs/kernel-contract.md's 13.3 sentence now counts the two launches, so the 13.3 kernel-contract-limits run should show both current. The 13.3 full contract run is yours, as you offered.
  2. glm cuda: a decode step's select gives each row a 1,024-thread block #1424. Agreed: keep its glm_index_k.Index and add only our two Fast lines, so tf_glm_scores is registered once. The branch still merges cleanly on main ee89695. If this PR lands second, I'll resolve it that way.
  3. NVFP4. Thank you.
  4. The retry path. glm-index-check now requires all three paths (1246901). On a GB10 with nvcc 13.0 it passes twice: whole row 2,048, candidates 6,726, retry 1,466.

Host tests on the new head: test-kernel-contract, test-glm-contract-host, test-glm-index-host (9/9) and test-glm-index-emulate pass. lean_check.py reports 428 problems before and after. I've updated the description to match.

If you'd like it on by default, I'm happy to send that follow-up once this lands.

@ashhart

ashhart commented Oct 27, 2026

Copy link
Copy Markdown
Owner

Our kernel contract run on a GB10 (nvcc 13.0, your head before the 13.3-records commit): SUMMARY 420 kernels: 420 pass gates 1-4, 0 fail, 6 known findings; 312 controls: 312 red as planted, 0 missed. tf_glm_index_scores and tf_glm_index_pools pass repeat (five runs byte-identical) and shapes (bit-equal to the kernels they replace), the sums gate passes, and the nonfinite finding is the same KNOWN one as tf_glm_scores (relu as fmaxf(x, 0)). The run exits 1 only on limit problems that main has too. With the 13.3 records in, this is everything we asked for, so it is in our next landing batch. Thank you!

@ashhart

ashhart commented Oct 27, 2026

Copy link
Copy Markdown
Owner

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.

@ashhart

ashhart commented Oct 28, 2026

Copy link
Copy Markdown
Owner

This is in TensorFold 1.1.8, released today. Thank you, @wojo!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

landed Landed on main by the maintainers; the commits keep the contributor's authorship

Projects

None yet

3 participants