[faresobeid]

Uses DeepSelect for the DSA indexer top-k in fp8_indexer when it is installed and built for the GPU (SM90, SM100, SM103). Everywhere else it falls back to torch.topk.

• select_topk passes the per-row causal end ke as DeepSelect's end, so a row is read only up to its end. The output is unsorted, int32 indices.
• Rows shorter than top-k are padded with the existing sentinel S_k. abort_when_nan_found=False makes a NaN row write an out-of-range index, which the existing range mask turns into the sentinel.
• KEY_ALIGNMENT goes from 16 to 256 keys, because DeepSelect needs a 1 KiB row stride. Padded keys are already masked to -inf.
• deep_select is not a dependency yet. The wheel (with sm_90a; upstream builds only sm_100a/sm_103a) comes from feat: build a DeepSelect wheel (sm_90a + sm_100a + sm_103a) prime-kernels#17, which is not released. Until then the import is optional and this path is inactive unless deep_select is importable. Once fix ac ckpt #17 is released, this PR pins deep-select, replaces the try/except ImportError with a plain import, and keeps only the GPU-capability check.

Exactness. DeepSelect picks exactly the same logit values as torch.topk in every row. The index sets differ only when several logits tie at the k-th place. Both choices are valid, and torch.topk does not define which one it picks either. Even with sorted=True the order differs from torch in 4-23% of rows because of ties inside the top-k, so bitwise-identical attention outputs are not reachable with any setting. That is why the output is unsorted. The new test checks equal value multisets in all rows and equal index sets in rows without a boundary tie, on integer logits where ties are common.

Numbers (H200).

Microbenchmark on fp32 logits from our Triton indexer kernel, top-2048 for GLM-5.3 and top-512 for DS-V4:

shape	torch.topk	DeepSelect	DeepSelect + end
GLM-5.3 cp4 last rank [32k x 131k]	54.4 ms	12.0 ms	10.9 ms
GLM-5.3 cp4 first rank	51.5 ms	11.3 ms	3.6 ms
16k sequence	4.2 ms	1.8 ms	1.1 ms
DS-V4 [64k x 16k], top-512	13.8 ms	4.1 ms	3.0 ms

Trainer setup: GLM-5.3 shape with random init, seq 131072, cp4, ep8, fp8 + deepgemm_fp8, AC + offload, compile, fake packed data, 1 node.

• The indexer runs once per full-indexer layer per micro-batch. AC recompute does not re-run it (checked in a trace).
• 8 layers (4 full-indexer): step 15.4 s -> 15.0 s (-2.6%).
• Extrapolated to 78 layers (21 full-indexer): top-k is ~2.2 s of a ~160 s step (~1.4%). This PR saves ~1.1% of the step. The share should be larger on Blackwell, where the GEMMs get faster.
• Entropy and mismatch KL are identical to 4 decimals between the A/B runs.

Ran: tests/unit/train/models/test_fp8_indexer.py on H200, with DeepSelect (4 passed) and without it (2 passed, 2 skipped). Also ruff check / ruff format --check, and pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py (301 passed).

Not validated: SM100/SM103 (no Blackwell hardware here). Upstream DeepSelect is developed on SM100.

---

Note

Medium Risk
Changes sparse token selection for DSA training on supported GPUs (alignment and tie-breaking vs torch.topk); fallback path preserves prior top-k behavior when DeepSelect is absent.

Overview
Adds a faster top-k path for the FP8 DSA indexer on Hopper/Blackwell when the optional deep_select package is importable, with unchanged behavior elsewhere.

New select_topk routes indexer logits through DeepSelect (per-row causal ke, int32 indices, sentinel padding) when the GPU is SM90/SM100/SM103 and topk ≤ 4096; otherwise it keeps torch.topk. fp8_indexer delegates to this helper instead of inlining top-k. KEY_ALIGNMENT increases from 16 to 256 so logit rows meet DeepSelect’s 1 KiB stride; padded keys stay masked as before.

A GPU test asserts DeepSelect and torch.topk return the same multiset of logit values (index sets match when there is no tie at the k-th boundary).

^{Reviewed by Cursor Bugbot for commit 7c42198. Bugbot is set up for automated code reviews on this repo. Configure here.}

[garrett361]

@faresobeid did you happen to test if deep select has better determinism for top_k? Know that's sometimes an issue, though maybe resolved in recent torch. Haven't checked

pytorch/pytorch#27542

[faresobeid]

@faresobeid did you happen to test if deep select has better determinism for top_k? Know that's sometimes an issue, though maybe resolved in recent torch. Haven't checked

pytorch/pytorch#27542

ya didnt see any noticeable differences in determinism