sampling: partial top-keep select in attention_rows DSA — O(nk) quickselect, not O(nk log nk) qsort (#356)

[woolcoxm]

Fixes #356

The DSA lightning indexer in attention_rows() selects the top-index_topk (2048 for GLM-5.2) context keys to attend to. It previously did this by full-qsorting all nk attention scores (O(nk log nk)) every layer, every token — just to read a single pivot value, the keep-th largest score, used as the threshold. It then scanned the original (unsorted) array in position order to build the kept set.

This is the attention-side twin of #335/#354 (sampling dist_build full-vocab qsort → heap partial-select), but it scales with context length, so the cost grows over the lifetime of a conversation and bites precisely in the long-context regime where decode latency matters most.

The fix

Replace the qsort with partial_select_desc (Hoare quickselect, median-of-three, descending) — O(nk) average to partition the keep largest into tmp[0..keep), then thr = min of that block:

float *tmp=falloc(nk); memcpy(tmp,isc,nk*sizeof(float));
partial_select_desc(tmp,nk,keep);            /* O(nk) avg, not O(nk log nk) */
float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];

The two position-order scans (>thr then ==thr) are unchanged → the kept-position set is bit-identical.

Correctness contract — stronger than #335

The qsort is unstable and only the threshold tmp[keep-1] is read from the sorted array; the actual kept-position set is determined entirely by the position-order scans, which are byte-for-byte identical before and after this PR. The quickselect pivot is, by definition, the keep-th largest element, so the new threshold equals the old tmp[keep-1] exactly. Therefore the kept set is element-wise identical — not just multiset-equal (which is all #335's unstable sampling heap could promise, since it also changed accumulation order).

Measured — bench_dsa_select (median of 2000 reps, keep=2048)

shape	nk (ctx)	old qsort	new quickselect	speedup
realistic	2049	118 µs	6.4 µs	18×
realistic	8192	629 µs	43 µs	15×
realistic	32768	3.4 ms	0.29 ms	12×
realistic	65536	6.3 ms	0.56 ms	11×
uniform	32768	3.0 ms	0.32 ms	9×
plateau (ties)	32768	0.67 ms	29 µs	23×
plateau (ties)	65536	1.4 ms	58 µs	25×

At 32K context × 78 layers, the qsort path was costing ~2.8 ms × 78 ≈ 220 ms/token spent sorting an array whose only purpose is to yield one threshold. The gap widens with context (linear vs n-log-n). DSA only activates past index_topk=2048, so this is invisible on short prompts but dominates long-conversation decode.

Test — test_dsa_select (in TEST_BINS)

129 cases. Drives the real partial_select_desc (via the include-glm.c pattern) and compares the full kept-position set against an independent reference that re-implements the old algorithm (qsort + tmp[keep-1] threshold + position scans) on a private buffer:

• Element-wise identical dst[] (not multiset) across shapes: random, peaked, strictly-decreasing (quickselect worst case), strictly-increasing (other worst case), tie-plateau, all-equal.
• Sweeps nk ∈ {1,2,8,64,2049,4097,8193} × keep ∈ {1,8,256,1024,2048}.
• Direct partition-invariant check (max(tail) ≤ min(top)) — catches a subtly broken quickselect even if the threshold happened to come out right.
• Edges: keep==nk, keep==1 (argmax), all-equal-scores (degenerate threshold, every kept slot is a tie → positions 0..keep-1).

test_dsa_select: 129 cases run, 0 failure(s)
test_dsa_select: ok

What I checked and did NOT change

Recorded in #356 so it isn't re-investigated — three other suspected attention-side hotspots, all measured to be already optimal at GLM-5.2 scale (E=256, K=8, index_topk=2048):

• MoE top-K (glm.c:2913, O(K·E) repeated linear-max-with-skip): the branch-predictable scan measures ~3 µs vs ~10.5 µs for a full argsort. Changing it would be a regression.
• softmax 3-pass vs 2-pass: within 1% (compiler hoists the loop-invariant reciprocal).
• tok_encode: looks O(L²) but is actually O(L·nsp·s) — the inner loop breaks on first hit.

Verification

• make tests/test_dsa_select && ./tests/test_dsa_select → 129 cases, 0 failures
• make tests/bench_dsa_select && ./tests/bench_dsa_select → reproduces the table above (6–25× across shapes)
• make glm.exe → clean (only the two pre-existing unrelated warnings: psapi pragma, g_numa_nodes)
• Full C test suite: 12/12 build+pass. (test_stops fails to build on dev — pre-existing mkdtemp/fix(test): make test_stops build on Windows (mkdtemp compat shim) #352 breakage, confirmed identical on clean origin/dev; unrelated to this PR.)

Scope / base

• Branch fix/dsa-partial-select-356 based off dev (per repo convention — dev is 42 commits ahead of main).
• Four files: glm.c (helper + one call-site), Makefile (test in TEST_BINS + bench on-demand rule), test_dsa_select.c, bench_dsa_select.c. No model files, no binaries, no benchmark artifacts.
• Cross-referenced on sampling: dist_build full-qsorts all 151K vocab entries per sampled token — partial top-p select is ~7x cheaper (measured) #335, sampling: partial top-p select in dist_build — O(V) heapify + k pops, not O(V log V) qsort (#335) #354, Prefill is attention-bound (not disk-bound), and int4 IDOT costs ~9-18% perplexity on the attention projections — ARM/POWER enable it by default #153, Diagnostic: token-exactness sweep + performance breakdown (pipe2 S=1 resident vs CPU vs CUDA — 4 prompts x 5 configs) #292.

🤖 Generated with ZCode

---

Credit / co-author: the partial-select algorithm this PR extends to the attention side is @KingIcyCreamProjects's work from issue #335 (the original max-heap partial top-p select + the 9950X3D measurement). This PR is the attention-side twin of #354, which implements the same idea for sampling. Recorded in-thread per the note below.

If proper git-native Co-Authored-By: attribution is wanted for any of these, it needs a maintainer to file it as a tracked issue/decision on the repo — the commits here are already merged into origin/dev, and amending them for trailers would rewrite shared integration history (every contributor's dev diverges).

[woolcoxm]

@KingIcyCreamProjects — you were right on the 8192 spike, and I verified it against the code rather than just trusting the read. Two bugs compounded:

1. One frozen input per cell — exactly as you said. partial_select_desc's median-of-three pivot is deterministic (no RNG), and the bench froze the input then ran 2000 reps on that identical array, so every rep hit the same pivot sequence. One lucky input → one spiked row, stable across re-runs because it's deterministic, not because it's real.
2. brng was never reset between cells — one you didn't mention. It's a static global initialized once and advanced by every prior cell's draws, so each (shape, nk) input depended on the whole history. Reordering nks[] silently shifted every later cell.

Fixed in #378: reseed per (shape, nk, seed), report median of 11 seeds × 2000 reps. On my Meteor Lake box the 8192 row drops to ~10× (realistic) / ~9× (uniform); the whole table is now a clean 6–26× band with no spikes. Would value a 9950X3D re-run of that branch to confirm your 75× row lands in-band.

Your whip-up bench was honestly fine for 2 seconds of work — the only thing missing was the seed sweep, which is a non-obvious gotcha with deterministic pivots.

---

On attribution (your note from #354): done. Credited you by name on #354, #357, and the #335 thread — the partial-select algorithm and the original 9950X3D measurement are your work and the PRs now say so. I did not add a `Co-Authored-By:` commit trailer, and I want to be straight about why: those commits are already merged into `origin/dev`, and amending them for the trailer would rewrite shared integration history — every contributor's `dev` would diverge, and that creates a real mess in the repo. The proper git-native route needs a maintainer to make a call on it (file it as a tracked decision). In-thread credit was the honest alternative; if you want the trailer too, say so and I'll raise it with the maintainer rather than force-push quietly.

On the coverage nit (V=151936 in `test_topp`): I owe you an honest retraction here. I said it was a reference-precision artifact and spent a while rewriting the reference to match production — that was circular, and you'd have been right to call it out. The real situation: at float-accumulation boundaries (uniform@0.5, V=151936) the heap gives keep=75969 where exact arithmetic says 75968 — which is actually sitting in your #354 bench table. That's not a regression (any float impl does this), but it means the bit-equivalence claim can't be extended to full V without either a looser tolerance or an honest caveat. I reverted my test changes; the gated test stays at V≤1519. The invariant checks (id-indexed `g_pbuf`, tail zeroed, head sums to 1) would extend cleanly to full V if you want that — those are real and cheap.

Serve-decode tok/s is still the open item — happy to run it, or if you've got the 9950X3D warm, that box is the one that matters.