Repository navigation
GLM CUDA: a long prompt's DSA indexer is O(position) per chunk; CUDA pool scores and a one-pass select (byte-equal) cut a 1M prompt by ~20% on the Zig port #1430
Description
Version: native main (the measurement is from 67cc8fbce plus #1322, which has since landed; the indexer code is unchanged on 8c36359d5). Model: GLM-5.3-Flash EXL3 4 bpw (078455ff). GPU/OS: two NVIDIA DGX Spark (GB10), TP=2, --context 1048576 --parallel 8, TF_KV_PAGES=on, Ubuntu 24.04.
What happens
A 1,020,560-token prompt takes 1,277 s to its first token on the pair. Most of the excess over a linear prefill comes from the DSA indexer of each 2,048-row prompt chunk. Every row scores and selects over all the pools before it: tf_glm_scores (families/glm/cuda_sparse.zig:57-97) feeds tf_glm_select_rows (glm_glue_index.cuh:92+). So a chunk's cost grows with its position, and the total grows as N².
A measured reference
On the same pair and checkpoint, the GLM Zig engine in Mia's recipe (TensorFold v1.0.7 plus that recipe's patches) had the same O(position) indexer through Triton _scores and _select_rows. Two CUDA kernels from BadAd84's Python patches replaced those two steps there:
- 0105: pool scores in CUDA;
- 0101: a one-pass, strided-sample floor, then a compaction and a radix select over the survivors.
Results:
- Kernel check against the Triton kernels: byte-equal scores and pools at 1k-262k pools.
- At 147,456 pools, one 512-row block: scores 9.59 → 6.61 ms, selection 6.50 → 1.36 ms.
- Served: a 1,039,999-token prompt 894 → 712 s. Replies identical at 33k, 131k, 262k, 524k and 1M tokens (text and token digest).
- The Python engine with the same kernels: 671 s.
Proposal
Port the two kernels to the native indexer:
- byte-equal with
tf_glm_scoresandtf_glm_select_rows(checked on the device across pool counts, random and tied scores); - behind a switch, with the current kernels as the fallback;
- a served receipt at 131k / 524k / 1M (first token and identical replies).
This makes each row's work cheaper, which is separate from #1404, which divides a chunk's rows across the ranks, and from #1424, which widens decode's select blocks. So they should stack. I'll coordinate with @Bizuayeu on whichever lands first.
Activity
wojo commentedon Oct 27, 2026
/atc claim
wojo commentedon Oct 27, 2026
/atc presence 1f9f77d3fcd5131423b214f32e46a6fb7c0f512ea04b983720d5e1b8a36adf75
Bizuayeu commentedon Oct 27, 2026
#1430
Agreed, they should stack. Where the three meet:
selectTokensProfileinfamilies/glm/cuda_sparse.zig, the block loop you cite (57-97 on main): glm cuda: a prompt chunk's indexer rows split across the ranks, their picks all-gathered #1404 runs it over this rank's block of rows only and all-gatherspoolsafter it; glm cuda: a decode step's select gives each row a 1,024-thread block #1424 launchestf_glm_select_rows_wide(1,024 threads a row) instead oftf_glm_select_rowsfor at most 48 rows over 4,096 pools or more (wideSelect). Your switch goes into the same lines, so whichever PR goes second rebases there.kernels/cuda/glm_glue_index.cuh: glm cuda: a decode step's select gives each row a 1,024-thread block #1424 turnstf_glm_select_rowsintoglm_index::select_row<T>with two entry points, 128 and 1,024 threads a row (T only sets how many pools a pass reads).tf_glm_select_rows(glm_cuda_records.zig,contract_glm_select.zig); yours would sit beside it.All of them give
tf_glm_select_rows's bytes (ties to the lower pool, picks in pool order), so they stay interchangeable: under #1404 each rank runs your kernels on its own block, and for decode the launcher can keep whichever select is faster at a few rows over many pools. If yours is faster there too, the wide entry can go.The bytes stack; the speeds need measuring together. Both thresholds were tuned against the current kernels: #1404 splits from position 2,048 by default, which pays only while its all-gather of the picks costs less than the rows it saves (a thin margin at 32k), and #1424's wide rule picks by rows and pools against the 128-thread select. With cheaper rows, whichever lands second should re-measure them on top of the first.
If yours lands first, I'll rebase #1404 and #1424 onto it; if mine land first, I'm happy to help with yours. Once your branch is up, I can measure the stack on XSFP4 at TP=3 (three GB10) and TP=2: the 128k and 516k-token prompts and decode after 136k, against #1404 and #1424 alone. For scale on this checkpoint at TP=3: the indexer was 7.4% of a 128k prompt (scores 4.7%, select 2.7%), and the select was 4.1% of a decode step after a 136k prompt (before #1424).
wojo commentedon Oct 27, 2026
@Bizuayeu the branch is up as #1449 (on main b30d5b9). Thanks for laying out where the three meet; that matches what I see in
selectTokensProfileand the records. A stacked measurement at TP=3 and TP=2 would be very welcome. On my pair at TP=2, a 1M-token prompt goes 1,223 → 1,033 s alone, with the same tokens at 131k/524k/1M. I'll rebase onto whichever of #1404/#1424 lands first.By Robert Wojciechowski (wojo on GitHub, @rwojo on X)
glm cuda: a prompt chunk's DSA indexer in glm_index.cu, byte-equal to…
glm cuda: glm_index under the kernel contract; radix no-hit case as t…
glm cuda: register the glm_index kernels from the heap so contract_te…
glm cuda: glm-index-check prints one line per case and step (#1430)
glm cuda: glm-index-check names the path count, the pool compare and …
glm cuda: glm-index-check sizes in usize, orders its memsets before t…
glm cuda: the glm_index records carry nvcc 13.0's measured registers …
glm cuda: glm_index's comments say what the kernels compute, one line…
ashhart commentedon Oct 27, 2026
Thank you, @wojo. The two indexer kernels are on main (520662a onwards, changelog in 8a789e3): byte-equal to tf_glm_scores and tf_glm_select_rows, opt-in with TF_GLM_INDEX_CUDA=1, and the changelog carries your 131k, 524k and 1M first-token numbers. That is everything this issue proposed, so we're closing it as done; stacking with #1404 and #1424 can carry on in those PRs.