woolcoxm
5d16368817
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.
133 lines
6.1 KiB
C
133 lines
6.1 KiB
C
/* 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;
|
|
}
|