Module select 

2026-09-25: GLM-5.3 DSA token selection: launch geometry, scratch, and the launcher for the selection kernels.

Owner: model-arch (GLM-5.3 DSA). Invariants:

• DsaSelectGeometry::plan returns Ok only when select_k <= topk_np2.
• select_tokens checks DsaSelectScratch::fits before its first launch and fails when the pass needs more scratch than was allocated.

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

Shared memory does not grow with the context

dsa_topk_pools walks the pool axis in tiles of topk_tile pools and keeps a running best list of one tile, so its shared memory (topk_smem_for_tile) is fixed. The one limit left is select_k of at most one tile, checked in DsaSelectGeometry::plan: at index_topk 2048 and index_kpool 4, select_k is 512 against a 2,048-pool tile. The context is bounded by the indexer cache (state::max_dsa_context).

No compaction pass

crate::glm5next_dsa_ref::kept_pools keeps pool p only when every one of its slots is in range and valid, counting from the first valid token. Over a contiguous cache with no padding that set is the prefix 0 .. seq / kpool (contiguous_pool_count, checked against kept_pools in tests), so the launcher uses the full arrays in place and dsa_compact_pools is not launched. A left-padded batch would need the compaction; DsaSelectGeometry::plan handles contiguous caches only.

Structs

DsaSelectGeometry

2026-09-25: Launch geometry for one selection pass, computed by Self::plan before any launch.

DsaSelectInputs

2026-09-25: Device-side inputs to a selection pass, all owned by the caller.

DsaSelectScratch

2026-09-25: Scratch the selection kernels write through, allocated once and reused across steps.

Enums

DsaSelectLaunch

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

Constants

TOPK_SMEM_CEILING

2026-09-25: Shared-memory ceiling the top-k select is budgeted against; the same value as SMEM_CEILING in examples/dsa_indexer_microtest.rs.

Functions

contiguous_pool_count

2026-09-25: Pools kept over a contiguous, unpadded cache of seq tokens.

select_tokens

2026-09-25: Run the selection kernels, leaving [q_rows, out_width] token ids in DsaSelectScratch::tokens.

topk_smem_for_tile

2026-09-25: Shared memory one dsa_topk_pools block needs for a tile of t pools.

topk_tile

2026-09-25: Tile width dsa_topk_pools walks the pool axis in.