Module select 

GLM-5.3 DSA token selection — the production launcher for the indexer pipeline.

Scoped to LibertAIDAI/GLM-5.3-Flash-NVFP4@9e0d74e3.

This is the examples/dsa_indexer_microtest.rs GATE-4 pipeline lifted out of the example and given a launcher a real layer can call. The kernels, their argument order and their numerics are already proven against HF 5.16.1 on real weights; nothing here re-derives them. What this module owns is the part the microtest did by hand and a layer cannot: geometry, capacity and refusal.

k_normed, gate, valid, ape  -> dsa_kpool_compress   -> pool keys / indices / valid
q, weights, q_pos           -> dsa_index_scores     -> [Q, P] scores + candidacy
                            -> dsa_topk_pools       -> [Q, select_k] pool ids
                            -> dsa_expand_selection -> [Q, out_width] token ids

🟢 The context ceiling this module used to impose is GONE

dsa_topk_pools no longer sorts the whole pool axis in shared memory. It walks the pools in fixed TOPK_TILE-wide tiles, keeping a running best-TOPK_TILE list, so shared memory is a constant 16 × TOPK_TILE bytes whatever the context. The result is bit-identical to the old whole-axis sort — the comparator (score DESC, pool index ASC) is a total order over unique indices, so the top-select_k prefix is unique and merge-and-truncate cannot reach a different set or order.

What survives is one requirement, checked in DsaSelectGeometry::plan: select_k <= TOPK_TILE. At index_topk = 2048 and index_kpool = 4 that is 512 against 2,048. DSA context is now bounded by the indexer cache allocation (state::max_dsa_context), not by this kernel. ANOMALIES A62.

🪤 Compaction is the identity here, and that is a derived fact, not an assumption

crate::layers::glm5next_dsa_ref::kept_pools keeps pool p only when every one of its kpool slots is in range and valid, with pooling starting at the first valid token. Over a contiguous, unpadded cache — every decode step at batch 1 — that set is exactly the prefix 0 .. seq / kpool, so the compacted array is a prefix of the full one and dsa_compact_pools would copy a buffer onto itself. This launcher therefore uses the full arrays in place and takes the prefix. contiguous_pool_count is proven equal to the reference for every sequence length in tests. A left-padded batch breaks the prefix property and genuinely needs the compaction arm — not built, and DsaSelectGeometry::plan is documented as contiguous-only.

Structs

DsaSelectGeometry

Launch geometry for one selection pass — every count the four kernels need, and every capacity check, decided before a single pointer is touched.

DsaSelectInputs

Device-side inputs to a selection pass. Every one is owned by the caller; this module allocates nothing but its own scratch.

DsaSelectScratch

Scratch the pipeline writes through, allocated once and reused across steps.

Enums

DsaSelectLaunch

How a pass is launched: exactly, or at the context ceiling so one graph serves any length.

Constants

TOPK_SMEM_CEILING

Runtime shared-memory ceiling the top-k select is budgeted against, matching SMEM_CEILING in examples/dsa_indexer_microtest.rs.

Functions

contiguous_pool_count

Pools kept over a contiguous, unpadded cache of seq tokens.

select_tokens

Run the four selection kernels, leaving [q_rows, out_width] token ids in DsaSelectScratch::tokens.

topk_smem_for_tile

Shared memory one dsa_topk_pools block needs for a tile of t pools.

topk_tile

Tile width dsa_topk_pools walks the pool axis in.