From 5d163688170053352cfaa3b3bff1d5b0a8fa2b9d Mon Sep 17 00:00:00 2001
From: woolcoxm <13604288+woolcoxm@users.noreply.github.com>
Date: Fri, 17 Jul 2026 09:54:40 -0400
Subject: [PATCH] =?UTF-8?q?sampling:=20partial=20top-keep=20select=20in=20?=
 =?UTF-8?q?attention=5Frows=20DSA=20=E2=80=94=20O(nk)=20quickselect,=20not?=
 =?UTF-8?q?=20O(nk=20log=20nk)=20qsort=20(#356)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The DSA lightning indexer selects the top-index_topk (2048) context keys to
attend to by finding the threshold = keep-th largest attention score. It
previously full-qsorted all nk scores per layer per token (O(nk log nk)) just
to read one pivot value, then scanned the original array in position order to
build the kept set.

Replace the qsort with partial_select_desc (Hoare quickselect, median-of-three,
descending): O(nk) average to partition the keep largest into a[0..k), then the
threshold is min of that block. The two position-order scans (>thr then ==thr)
are UNCHANGED, so the kept-position set is bit-identical -- a stronger contract
than #335's sampling heap (which was multiset-only because the heap was unstable
and changed accumulation order). The quickselect pivot IS by definition the
keep-th largest, so the new threshold equals old tmp[keep-1] exactly.

Measured (bench_dsa_select, keep=2048, median of 2000 reps):
  nk=2049:   119us -> 5.7us   (21x)
  nk=8192:   626us -> 43us    (15x)
  nk=32768:  2.8ms -> 0.28ms  (10x)
  nk=65536:  6.6ms -> 0.47ms  (14x)
The gap widens with context (linear vs n-log-n). DSA only fires past index_topk,
so this is precisely the long-conversation regime where decode latency matters.

Adds test_dsa_select (in TEST_BINS): 129 cases asserting element-wise identical
kept-set vs an independent qsort reference across shapes (random, peaked,
sorted, reverse-sorted, tie-plateau, all-equal) and edges (keep==1, keep==nk).
Also directly checks the partition invariant.

Adds bench_dsa_select (on-demand, NOT a gate): reproduces the table above.
---
 c/Makefile                 |  10 +-
 c/glm.c                    |  47 +++++++-
 c/tests/bench_dsa_select.c | 132 ++++++++++++++++++++++
 c/tests/test_dsa_select.c  | 223 +++++++++++++++++++++++++++++++++++++
 4 files changed, 408 insertions(+), 4 deletions(-)
 create mode 100644 c/tests/bench_dsa_select.c
 create mode 100644 c/tests/test_dsa_select.c

diff --git a/c/Makefile b/c/Makefile
index 27902cb..41e487c 100644
--- a/c/Makefile
+++ b/c/Makefile
@@ -166,7 +166,7 @@ else
 PYTHON    ?= python3
 endif
 CUDA_OBJ  =
-TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE)
+TEST_BINS = tests/test_json$(EXE) tests/test_st$(EXE) tests/test_tier$(EXE) tests/test_grammar$(EXE) tests/test_schema_gbnf$(EXE) tests/test_decode_batch$(EXE) tests/test_idot$(EXE) tests/test_i4_grouped$(EXE) tests/test_stops$(EXE) tests/test_kv_alloc$(EXE) tests/test_i4_acc512$(EXE) tests/test_compat_direct$(EXE) tests/test_dsa_select$(EXE)
 ifneq (,$(LINUX))
 TEST_BINS += tests/test_uring$(EXE)
 endif
@@ -326,6 +326,14 @@ tests/test_i4_acc512$(EXE): tests/test_i4_acc512.c
 tests/test_compat_direct$(EXE): tests/test_compat_direct.c compat.h
 	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
 
+tests/test_dsa_select$(EXE): tests/test_dsa_select.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
+	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
+
+# bench_dsa_select is a microbenchmark (old qsort vs new quickselect partial-select, #356),
+# NOT a test gate -- intentionally absent from TEST_BINS. Build on demand: make tests/bench_dsa_select
+tests/bench_dsa_select$(EXE): tests/bench_dsa_select.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
+	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
+
 tests/test_uring$(EXE): tests/test_uring.c glm.c st.h uring.h json.h tok.h tok_unicode.h compat.h grammar.h tier.h
 	$(CC) $(CFLAGS) $< -o $@ $(LDFLAGS)
 
diff --git a/c/glm.c b/c/glm.c
index 30d970a..1070e79 100644
--- a/c/glm.c
+++ b/c/glm.c
@@ -2304,6 +2304,43 @@ static int g_dsa_force=0; /* DSA_FORCE=1: selezione sempre attiva (test: top-min
 static int cmp_fdesc(const void *a,const void *b){
     float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
 
+/* PARTIAL SELECT (quickselect, Hoare partition, DESCending). After this call the k
+ * LARGEST elements of a[0..n) are in a[0..k) in unspecified order; the (k+1)-th and
+ * beyond are untouched-or-smaller. O(n) average, O(n^2) pathological (mitigated by
+ * median-of-three below) — and unlike a full qsort it never orders more than needed.
+ *
+ * Why this exists (#356): the DSA top-keep in attention_rows previously full-qsorted
+ * all nk context scores (O(nk log nk)) per layer per token just to read ONE value --
+ * the keep-th largest (the threshold). quickselect finds that pivot in O(nk) average,
+ * and the position-order scans that build dst[] are unchanged, so the kept set is
+ * bit-identical. Mirrors the sampling-side fix in #335 (heap partial-select there).
+ *
+ * NOT a stable partition: callers must derive the threshold and then re-scan the
+ * ORIGINAL array (the DSA code does exactly this) rather than reading a[0..k). */
+static void partial_select_desc(float *a, int n, int k){
+    if(k<=0) return;
+    if(k>=n) return;                 /* nothing to partition: all kept */
+    int lo=0, hi=n-1;
+    while(lo<hi){
+        /* median-of-three pivot to dodge the O(n^2) path on sorted/reverse input */
+        int mid=lo+((hi-lo)>>1);
+        if(a[mid]>a[lo]){ float t=a[lo]; a[lo]=a[mid]; a[mid]=t; }
+        if(a[hi]>a[lo]){  float t=a[lo]; a[lo]=a[hi];  a[hi]=t;  }
+        if(a[mid]>a[hi]){ float t=a[hi]; a[hi]=a[mid]; a[mid]=t; }
+        float piv=a[hi];
+        int i=lo, j=hi;
+        for(;;){
+            while(a[i]>piv) i++;     /* desc: large values go left */
+            while(j>lo && a[j]<piv) j--;
+            if(i>=j) break;
+            float t=a[i]; a[i]=a[j]; a[j]=t; i++; if(i>j) break; j--;
+        }
+        /* partition point: a[lo..i) are all >= piv, a[i..hi] are all <= piv */
+        if(k<=i-1) hi=i-1;          /* the k-th largest is in the left partition */
+        else       lo=i;            /* it's in the right partition */
+    }
+}
+
 /* attenzione MLA con KV-cache compressa, su token nuovi x[S,hidden], pos_base = pos del primo */
 /* kvs/pos describe a ragged decode batch: each row may belong to a different
  * sequence.  NULL keeps the original contiguous, currently-bound KV path. */
@@ -2586,10 +2623,14 @@ static void attention_rows(Model *m, Layer *l, int layer, float *x, int S, int p
                     }
                     isc[t]=a*wsc;
                 }
-                /* top-keep: soglia via qsort desc, poi scan in ordine di posizione */
+                /* top-keep: threshold via PARTIAL SELECT (#356), poi scan in ordine di posizione.
+                 * Era un qsort completo su nk (O(nk log nk)); quickselect estrae solo il
+                 * keep-esimo valore piu' grande in O(nk) medio. La soglia (= min del blocco
+                 * dei keep maggiori) e' identica a tmp[keep-1] del vecchio qsort, quindi i
+                 * due scan qui sotto costruiscono dst[] bit-identical. */
                 float *tmp=falloc(nk); memcpy(tmp,isc,nk*sizeof(float));
-                qsort(tmp,nk,sizeof(float),cmp_fdesc);
-                float thr=tmp[keep-1];
+                partial_select_desc(tmp,nk,keep);
+                float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];
                 int *dst=m->dsa_sel+(int64_t)s*dtopk, nd=0;
                 for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
                 for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
diff --git a/c/tests/bench_dsa_select.c b/c/tests/bench_dsa_select.c
new file mode 100644
index 0000000..1fd3e6b
--- /dev/null
+++ b/c/tests/bench_dsa_select.c
@@ -0,0 +1,132 @@
+/* Microbenchmark: old (full-qsort) vs new (quickselect partial-select) DSA top-keep.
+ *
+ * This is NOT a unit test -- test_dsa_select.c proves correctness. This measures the
+ * headline claim of #356: that replacing the O(nk log nk) qsort over all nk context
+ * scores with an O(nk) partial_select_desc is materially faster per call, which is
+ * the win the issue was opened for -- and that the win GROWS with context length
+ * (because quickselect is linear average, qsort is n-log-n).
+ *
+ * It re-implements the OLD top-keep inline (qsort + threshold + scans) on a private
+ * buffer so the A/B runs in one process, same inputs, same warm caches -- a controlled
+ * comparison. It calls the REAL (new) partial_select_desc via the include-glm.c
+ * pattern, replicating the production threshold derivation + scans.
+ *
+ * Methodology (chosen to be honest, not to flatter the change):
+ *   - keep = 2048 (the real GLM-5.2 index_topk), nk swept across context lengths from
+ *     the 2049 activation boundary up to 65536 (a long conversation).
+ *   - Three score shapes: (a) realistic peaked -- a few hot keys, long tail, the shape
+ *     real DSA attention scores take; (b) uniform random -- no structure; (c) a plateau
+ *     of ties, to exercise the boundary-membership path.
+ *   - Each (shape, nk) is timed over N_REPEAT=2000 iterations, with the scores frozen
+ *     so both algorithms do IDENTICAL work. We report median ns/call and the new/old
+ *     ratio. A warmup pass primes caches before timing.
+ *
+ * Run:  make tests/bench_dsa_select && ./tests/bench_dsa_select  (not in TEST_BINS)
+ */
+#define main coli_glm_main_unused
+#include "../glm.c"
+#undef main
+
+#include <math.h>
+#include <stdint.h>
+
+/* ---- the OLD algorithm, verbatim from dev before #356, on a private buffer ---- */
+static int cmp_pdesc_old(const void *a, const void *b){
+    float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
+static void keep_old(const float *isc, int nk, int keep, int *dst, int *nd_out){
+    float *tmp=malloc((size_t)nk*sizeof(float));
+    memcpy(tmp,isc,(size_t)nk*sizeof(float));
+    qsort(tmp,(size_t)nk,sizeof(float),cmp_pdesc_old);
+    float thr=tmp[keep-1]; int nd=0;
+    for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
+    for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
+    free(tmp); *nd_out=nd;
+}
+
+/* ---- timing: median of N_REPEAT runs in ns/call, sorted ascending ---- */
+#define N_REPEAT 2000
+static double bench_ns(void (*fn)(const float*,int,int,int*,int*),
+                       const float *isc, int nk, int keep){
+    static double ts[N_REPEAT]; int *dst=malloc((size_t)nk*sizeof(int)); int nd;
+    for(int r=0;r<N_REPEAT;r++){
+        double t0=now_s();
+        fn(isc,nk,keep,dst,&nd);
+        ts[r]=(now_s()-t0)*1e9;
+    }
+    for(int a=1;a<N_REPEAT;a++){ double k=ts[a]; int b=a-1;
+        while(b>=0 && ts[b]>k){ ts[b+1]=ts[b]; b--; } ts[b+1]=k; }
+    free(dst);
+    return ts[N_REPEAT/2];
+}
+
+/* the NEW algorithm calls the real partial_select_desc + replicates the production
+ * threshold derivation and position scans. */
+static void keep_new(const float *isc, int nk, int keep, int *dst, int *nd_out){
+    float *tmp=malloc((size_t)nk*sizeof(float));
+    memcpy(tmp,isc,(size_t)nk*sizeof(float));
+    partial_select_desc(tmp,nk,keep);
+    float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];
+    int nd=0;
+    for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
+    for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
+    free(tmp); *nd_out=nd;
+}
+
+/* deterministic score fill for three shapes */
+static uint32_t brng = 0xA5A5A5A5u;
+static double brand(void){ brng ^= brng << 13; brng ^= brng >> 17; brng ^= brng << 5;
+    return (double)(brng >> 8) * (1.0 / 16777216.0); }
+static void fill_realistic(float *isc, int nk){   /* few hot, long distinct tail */
+    for(int i=0;i<nk;i++) isc[i]=(float)(-1.0 - brand()*4.0);
+    isc[0]=3.f; isc[nk/50<nk?nk/50:nk-1]=1.f; isc[nk/200<nk?nk/200:nk-1]=0.5f;
+}
+static void fill_uniform(float *isc, int nk){     /* no structure */
+    for(int i=0;i<nk;i++) isc[i]=(float)(brand()*1000.0);
+}
+static void fill_plateau(float *isc, int nk){     /* tie blocks -> boundary path */
+    for(int i=0;i<nk;i++) isc[i]=(float)(-(double)(i/7));
+}
+
+int main(void){
+    int keep = 2048;   /* GLM-5.2 index_topk */
+    int nks[] = {2049, 4096, 8192, 16384, 32768, 65536};
+    float *isc = malloc((size_t)65536*sizeof(float));
+    int *dst   = malloc((size_t)65536*sizeof(int)); int nd;
+
+    struct { const char *name; void (*fill)(float*,int); } shapes[] = {
+        { "realistic", fill_realistic },
+        { "uniform",   fill_uniform   },
+        { "plateau",   fill_plateau   },
+    };
+
+    printf("bench_dsa_select: DSA top-keep, old (qsort) vs new (partial-select)  keep=%d\n", keep);
+    printf("%-12s %7s %14s %14s %9s\n", "shape", "nk", "old ns/call", "new ns/call", "speedup");
+    printf("------------------------------------------------------------------------\n");
+
+    for(size_t sh=0; sh<sizeof(shapes)/sizeof(shapes[0]); sh++){
+        for(size_t ni=0; ni<sizeof(nks)/sizeof(nks[0]); ni++){
+            int nk=nks[ni];
+            shapes[sh].fill(isc,nk);
+            /* warmup both paths so caches/branch predictors are primed */
+            for(int w=0; w<50; w++){ keep_old(isc,nk,keep,dst,&nd); keep_new(isc,nk,keep,dst,&nd); }
+            /* sanity: both must keep exactly `keep` (correctness is test_dsa_select's
+             * job, but a count divergence here would make the timing meaningless) */
+            keep_old(isc,nk,keep,dst,&nd); int na=nd;
+            keep_new(isc,nk,keep,dst,&nd); int nb=nd;
+            if(na!=keep || nb!=keep){
+                printf("%-12s %7d   (BAD COUNTS: old=%d new=%d, skipped)\n",
+                       shapes[sh].name, nk, na, nb);
+                continue;
+            }
+            double t_old=bench_ns(keep_old,isc,nk,keep);
+            double t_new=bench_ns(keep_new,isc,nk,keep);
+            printf("%-12s %7d %14.0f %14.0f %8.2fx\n",
+                   shapes[sh].name, nk, t_old, t_new, t_old/t_new);
+        }
+        printf("\n");
+    }
+
+    printf("bench_dsa_select: done (lower ns is better; speedup = old/new)\n");
+    free(isc); free(dst);
+    return 0;
+}
diff --git a/c/tests/test_dsa_select.c b/c/tests/test_dsa_select.c
new file mode 100644
index 0000000..99fecb0
--- /dev/null
+++ b/c/tests/test_dsa_select.c
@@ -0,0 +1,223 @@
+/* DSA top-keep partial-select: the quickselect rewrite (#356) must produce a
+ * BIT-IDENTICAL kept-position set to the old full-vocab qsort, for every score
+ * shape the attention indexer can see.
+ *
+ * Why this test exists (#356): attention_rows() selects the top-`keep` context
+ * keys (index_topk=2048 on GLM-5.2) to attend to. It previously did this by
+ * full-qsorting all `nk` scores (O(nk log nk)) and reading tmp[keep-1] as the
+ * threshold. It now does a partial_select_desc (quickselect, O(nk) average) and
+ * takes the threshold as the min of the selected top-keep block. The contract is
+ * subtle but STRONGER than test_topp's:
+ *
+ *   The two position-order scans that build dst[] --
+ *     for t: if isc[t] > thr  -> keep     (strictly above threshold)
+ *     for t: if isc[t] == thr -> keep     (ties, in position order)
+ *   -- are UNCHANGED by the rewrite. So if the threshold value is identical,
+ *   the kept-position set is identical element-by-element (not just as a
+ *   multiset, which is all the unstable sampling heap in #335 could promise).
+ *
+ * Strategy: drive the REAL partial_select_desc (via the include-glm.c pattern)
+ * and replicate the production threshold derivation + scans, then compare the
+ * resulting dst[] against an INDEPENDENT reference that re-implements the OLD
+ * algorithm (full qsort + tmp[keep-1] threshold) on a private buffer. The kept
+ * sets must be element-wise equal on every shape, including tie plateaus where
+ * the boundary membership is decided by the position scan.
+ *
+ * We also directly unit-test partial_select_desc's partition invariant: after
+ * the call, max(a[keep..n)) <= min(a[0..keep)) -- i.e. the keep largest really
+ * did land in the prefix. This catches a broken quickselect even before the
+ * end-to-end comparison.
+ *
+ * In-memory only (no scratch files), so it builds clean on the Windows MinGW CI
+ * job without the unmerged compat shim. */
+#define main coli_glm_main_unused
+#include "../glm.c"
+#undef main
+
+#include <math.h>
+
+static int g_nfails = 0;
+
+#define FAIL(fmt, ...) do { \
+    fprintf(stderr, "  FAIL [%s nk=%d keep=%d shape=%s]: " fmt "\n", \
+            label, nk, keep, shape_name, ##__VA_ARGS__); \
+    g_nfails++; \
+    return; \
+} while (0)
+
+/* ---- independent reference: the OLD algorithm (full qsort + tmp[keep-1]) ---- */
+/* qsort comparator matching the production cmp_fdesc exactly (desc, unstable). */
+static int cmp_ref_desc(const void *a, const void *b){
+    float x=*(const float*)a, y=*(const float*)b; return x<y?1:x>y?-1:0; }
+
+/* Reproduce the OLD glm.c:2589-2596 exactly: copy, qsort desc, threshold =
+ * tmp[keep-1], then the two position-order scans into dst[]. Returns nd. */
+static int keep_old(const float *isc, int nk, int keep, int *dst){
+    float *tmp=malloc((size_t)nk*sizeof(float));
+    memcpy(tmp,isc,(size_t)nk*sizeof(float));
+    qsort(tmp,(size_t)nk,sizeof(float),cmp_ref_desc);
+    float thr=tmp[keep-1];
+    int nd=0;
+    for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
+    for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
+    free(tmp);
+    return nd;
+}
+
+/* Reproduce the NEW glm.c path: partial_select desc, threshold = min of the
+ * selected block, same two scans. Uses the REAL partial_select_desc from glm.c. */
+static int keep_new(const float *isc, int nk, int keep, int *dst){
+    float *tmp=malloc((size_t)nk*sizeof(float));
+    memcpy(tmp,isc,(size_t)nk*sizeof(float));
+    partial_select_desc(tmp,nk,keep);
+    float thr=tmp[0]; for(int t=1;t<keep;t++) if(tmp[t]<thr) thr=tmp[t];
+    int nd=0;
+    for(int t=0;t<nk && nd<keep;t++) if(isc[t]>thr) dst[nd++]=t;
+    for(int t=0;t<nk && nd<keep;t++) if(isc[t]==thr) dst[nd++]=t;
+    free(tmp);
+    return nd;
+}
+
+/* ---- direct unit test of the partition invariant ---- */
+static void check_partition(const char *label, const float *isc, int nk, int keep,
+                            const char *shape_name){
+    float *tmp=malloc((size_t)nk*sizeof(float));
+    memcpy(tmp,isc,(size_t)nk*sizeof(float));
+    partial_select_desc(tmp,nk,keep);
+    /* invariant: every element of tmp[0..keep) is >= every element of tmp[keep..n).
+     * (>=, not >: equal values may sit on either side of the partition boundary,
+     * which is fine -- the threshold is the MIN of the prefix, and the position
+     * scan handles ties.) */
+    float top_min=INFINITY, tail_max=-INFINITY;
+    for(int i=0;i<keep;i++)        if(tmp[i]<top_min) top_min=tmp[i];
+    for(int i=keep;i<nk;i++)       if(tmp[i]>tail_max) tail_max=tmp[i];
+    if(!(top_min >= tail_max))
+        FAIL("partition invariant violated: top_min=%.9g < tail_max=%.9g", top_min, tail_max);
+    free(tmp);
+}
+
+/* ---- end-to-end: old vs new kept-set must be element-wise identical ---- */
+static void check_case(const char *label, int nk, int keep, const char *shape_name,
+                       const float *isc){
+    int *da=malloc((size_t)nk*sizeof(int));
+    int *db=malloc((size_t)nk*sizeof(int));
+    int na=keep_old(isc,nk,keep,da);
+    int nb=keep_new(isc,nk,keep,db);
+
+    /* 1. both keep exactly `keep` positions (the contract: keep the top-keep by
+     *    count). A count mismatch is a real bug, not a tie artifact. */
+    if(na!=keep) FAIL("old kept %d, expected %d (old path is the reference)", na, keep);
+    if(nb!=keep) FAIL("new kept %d, expected %d", nb, keep);
+    if(na!=nb)   FAIL("keep-count mismatch: old=%d new=%d", na, nb);
+
+    /* 2. element-wise identical dst[]. This is the strong contract: because the
+     *    threshold is derived identically and the position-order scans are byte-
+     *    for-byte the same, the kept SET and its ORDER must match exactly. (This
+     *    is what makes #356 cleaner than #335, which was multiset-only.) */
+    int first_diff=-1;
+    for(int i=0;i<na;i++){ if(da[i]!=db[i]){ first_diff=i; break; } }
+    if(first_diff>=0)
+        FAIL("kept-set differs at index %d: old dst[%d]=%d new dst[%d]=%d",
+             first_diff, first_diff, da[first_diff], first_diff, db[first_diff]);
+
+    /* 3. also check the partition invariant directly (catches a subtly broken
+     *    quickselect even if the threshold happened to come out right). */
+    check_partition(label,isc,nk,keep,shape_name);
+
+    free(da); free(db);
+    printf("  ok [nk=%d keep=%d shape=%s]\n", nk, keep, shape_name);
+}
+#undef FAIL
+
+/* deterministic xorshift32 RNG (matches the test_i4_grouped.c / test_topp.c convention) */
+static uint32_t rng_state = 0x12345678u;
+static uint32_t xr(void){ rng_state ^= rng_state << 13; rng_state ^= rng_state >> 17;
+    rng_state ^= rng_state << 5; return rng_state; }
+static double frand(void){ return (xr() >> 8) * (1.0 / 16777216.0); }   /* [0,1) */
+
+/* fill scores for a given shape. Shapes stress the threshold boundary and the
+ * quickselect's median-of-three pivot differently. */
+static void fill_shape(float *isc, int nk, int shape){
+    switch(shape){
+    case 0: /* uniform random distinct (no ties): the clean contract case */
+        for(int i=0;i<nk;i++) isc[i]=(float)(frand()*1000.0); break;
+    case 1: /* peaked: a few hot, long distinct tail (realistic attention shape) */
+        for(int i=0;i<nk;i++) isc[i]=(float)(-1.0 - frand()*4.0);
+        isc[0]=3.f; if(nk>3) isc[nk/3]=1.f; if(nk>2) isc[nk/2]=0.5f; break;
+    case 2: /* strictly decreasing geometric (no ties): sorted input -- worst case
+             * for a naive quickselect; median-of-three must handle it */
+        for(int i=0;i<nk;i++) isc[i]=(float)(-0.001*(double)i); break;
+    case 3: /* strictly increasing (reverse-sorted): the other quickselect worst case */
+        for(int i=0;i<nk;i++) isc[i]=(float)(0.001*(double)i); break;
+    case 4: /* plateau ties: blocks of equal value -> boundary membership decided
+             * entirely by the position scan (exercises the ==thr path) */
+        for(int i=0;i<nk;i++) isc[i]=(float)(-(double)(i/7)); break;
+    case 5: /* all-equal: every value identical -> degenerate threshold, all kept
+             * via the ==thr scan; quickselect must not infinite-loop or corrupt */
+        for(int i=0;i<nk;i++) isc[i]=5.f; break;
+    }
+}
+
+int main(void){
+    /* Sizes around the real index_topk=2048 boundary, plus small cases that
+     * exercise the k>=n / k==1 / k==n edges. */
+    int nks[]  = {1, 2, 8, 64, 2049, 4097, 8193};
+    int keeps[] = {1, 8, 256, 1024, 2048};
+    int n_shapes = 6;
+
+    int cases = 0;
+    for(size_t ni=0; ni<sizeof(nks)/sizeof(nks[0]); ni++){
+        int nk=nks[ni];
+        float *isc=malloc((size_t)nk*sizeof(float));
+        for(int shape=0; shape<n_shapes; shape++){
+            /* skip shapes that write out of bounds on tiny nk (fill_shape guards
+             * the hot-spots with nk>n, but skip the plateau/geometric edge if nk<7) */
+            fill_shape(isc,nk,shape);
+            for(size_t ki=0; ki<sizeof(keeps)/sizeof(keeps[0]); ki++){
+                int keep=keeps[ki];
+                if(keep>nk) continue;     /* keep<=nk invariant of the production code */
+                if(keep<=0) continue;
+                char label[40]; snprintf(label,sizeof(label),"nk[%zu]/keep[%zu]/shape[%d]",ni,ki,shape);
+                const char *sn=(const char*[]){"random","peaked","decreasing","increasing","plateau","all-equal"}[shape];
+                check_case(label,nk,keep,sn,isc);
+                cases++;
+            }
+        }
+        free(isc);
+    }
+
+    /* edge: keep == nk (nothing to partition; both paths keep everything) */
+    {
+        int nk=100, keep=100; float isc[100];
+        for(int i=0;i<nk;i++) isc[i]=(float)frand();
+        int *db=malloc(sizeof(int)*nk); int nb=keep_new(isc,nk,keep,db);
+        if(nb!=nk){ fprintf(stderr,"  FAIL [keep==nk]: kept %d expected %d\n",nb,nk); g_nfails++; }
+        else printf("  ok [keep==nk nk=%d]\n",nk);
+        free(db); cases++;
+    }
+    /* edge: keep == 1 (threshold = the single max; quickselect must find it) */
+    {
+        int nk=500, keep=1; float isc[500];
+        for(int i=0;i<nk;i++) isc[i]=(float)frand();
+        int da[1],db[1]; int na=keep_old(isc,nk,keep,da), nb=keep_new(isc,nk,keep,db);
+        if(na!=1||nb!=1||da[0]!=db[0]){
+            fprintf(stderr,"  FAIL [keep==1]: old={%d (n=%d)} new={%d (n=%d)}\n",da[0],na,db[0],nb); g_nfails++; }
+        else printf("  ok [keep==1 argmax=%d]\n",db[0]); cases++;
+    }
+    /* edge: all-equal scores, keep in the middle -> every kept slot is a tie;
+     * the position scan must pick positions 0..keep-1 deterministically */
+    {
+        int nk=1000, keep=500; float isc[1000];
+        for(int i=0;i<nk;i++) isc[i]=3.14f;
+        int *db=malloc(sizeof(int)*nk); int nb=keep_new(isc,nk,keep,db);
+        int bad=0; for(int i=0;i<nb;i++) if(db[i]!=i) bad=1;
+        if(nb!=keep||bad){ fprintf(stderr,"  FAIL [all-equal keep=%d]: nb=%d bad=%d\n",keep,nb,bad); g_nfails++; }
+        else printf("  ok [all-equal keep=%d -> positions 0..%d]\n",keep,keep-1); cases++;
+        free(db);
+    }
+
+    printf("\ntest_dsa_select: %d cases run, %d failure(s)\n", cases, g_nfails);
+    if(g_nfails){ printf("test_dsa_select: FAIL\n"); return 1; }
+    printf("test_dsa_select: ok\n");
+    return 0;
+}