From 03d9a23fe46ecd6980a466543399fd2085a37208 Mon Sep 17 00:00:00 2001 From: Rodolfo Hansen Date: Sun, 12 Jul 2026 13:44:34 +0200 Subject: [PATCH] perf(moe): overlap NVMe expert loads with matmul via opt-in PIPE I/O pool (default off, oracle-exact) (#79) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit For streamed (non-resident) experts the MoE forward is disk-bound: each 64-expert block first blocks on a parallel pread of every miss, then runs the matmuls serially. Those two phases don't overlap, so the compute cores idle while the block's weights are still coming off NVMe. PIPE=1 hands the misses' expert_load() to a small persistent pool of I/O worker pthreads and lets the main thread start matmul immediately. The main thread walks the block's experts in routing order and waits on a per-slot ready flag only for the expert it needs right now, so loads of later experts in the block hide behind matmul of earlier ones. All matmul_qt stays on the main thread (it parallelises internally via OpenMP and gates GPU dispatch on !omp_in_parallel()), so the I/O pool never competes with the compute team for the matmul itself. Cross-generation correctness: generation-tagged lock-free cursor -------------------------------------------------------------------- The batch state (njobs/eids[]/layer/ready[]) is reused in place across 64-blocks, so a straggler or late-woken worker could touch the NEXT batch's state. Two prior fixes (a drain barrier, then an _Atomic active counter) each still left a cross-generation window. Both are now removed and replaced with a single generation-tagged cursor: _Atomic uint64_t cur = (gen << 8) | index; // gen main-only, index 0..njobs - dispatch writes njobs/layer/eids[]/ready-reset with RELAXED stores, then RELEASE-stores cur (bumping gen); that release publishes all batch state to any worker whose ACQUIRE-load of cur sees the new gen. - a worker grabs a job by CAS-advancing the index; it reads eids[i]/ layer only AFTER the winning CAS. The CAS comparand carries the generation, so if a new batch was published the stale CAS fails and the worker re-reads — it can never grab a wrong-generation job or read torn state, no matter where it was preempted (wake gap, post-cursor, anywhere). - gen is bumped only by the main thread and is monotonic ⇒ no ABA. - the per-expert pipe_wait(ready[q]) in the matmul loop (kept) makes every grabbed job complete before the block ends, so no grab outlives its generation. That is what makes the old `active` counter and the end-of-block drain barrier unnecessary — both are removed. ready[] is reset before the publishing release, so no stale flag survives into the new generation. - the mutex/condvar now exist ONLY to park and wake idle workers, not for correctness. This only reorders I/O, never the computation: greedy decode output is byte-identical to the blocking path. Default OFF (PIPE=0) so upstream behaviour is unchanged; PIPE=1 opts in and PIPE_WORKERS=n sizes the pool (default 8). No Makefile change: pure C (stdatomic.h + sched.h), builds with the default toolchain. Claude-Session: https://claude.ai/code/session_01DS7oc65c5RdA9V99otRCwt Co-authored-by: Claude Opus 4.8 --- c/glm.c | 126 ++++++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 119 insertions(+), 7 deletions(-) diff --git a/c/glm.c b/c/glm.c index 87706f4..1d74f78 100644 --- a/c/glm.c +++ b/c/glm.c @@ -24,6 +24,8 @@ #include #include #include /* thread I/O del PILOTA */ +#include /* PIPE: ready-flags / job queue */ +#include /* PIPE: sched_yield nello spin */ #include #if defined(__APPLE__) || defined(__linux__) #include @@ -972,6 +974,97 @@ static void expert_load(Model *m, int layer, int eid, ESlot *s){ s->eid=eid; } +/* ============================ PIPE: load ‖ matmul ============================ + * Overlap NVMe expert-weight loads with expert matmul. A small persistent pool + * of I/O worker pthreads runs the misses' pread (expert_load) into distinct + * ws[] slabs and sets a per-slot `ready` flag; the MAIN thread walks the block's + * experts in order, waiting on ready[q] only for the expert it needs right now, + * and does all matmul_qt on itself (matmul_qt parallelises internally via OpenMP + * and checks !omp_in_parallel() for GPU dispatch — so it must stay off the omp + * team and off these I/O threads). + * + * Cross-generation safety is provided by a single generation-tagged, lock-free + * cursor `cur = (gen<<8) | index`. The main thread is the sole writer of `gen` + * (monotonic bump, so no ABA); workers grab jobs by CAS-advancing the low 8-bit + * index. THE INVARIANT: a worker reads eids[i]/layer only AFTER its winning CAS, + * and that CAS's comparand carries the generation — so if `cur`'s gen advanced + * (a new batch was published), the CAS fails and the worker re-reads, seeing the + * new generation. A straggler preempted anywhere (wake gap, post-cursor) can + * therefore NEVER grab a wrong-generation job or read torn batch state: its + * first act is a gen-checked CAS. dispatch publishes all batch state with + * relaxed stores and then RELEASE-stores `cur`; each worker ACQUIRE-loads `cur`, + * so the ready[] reset + eids[]/njobs/layer are visible before any worker acts. + * The per-expert pipe_wait(ready[q]) in the matmul loop makes every grabbed job + * complete before the block ends, so no grab outlives its generation — which is + * why the old `active` counter AND the end-of-block drain barrier are gone (both + * were redundant with those per-slot waits + the gen-tagged cursor). The mutex/ + * condvar exist ONLY to park/wake idle workers, never for correctness. Gated + * behind PIPE=1; OFF => the original blocking-load + serial-matmul path runs + * byte-identically. */ +static int g_pipe=0; /* PIPE=1: async expert-load pipeline (default OFF) */ +static int g_pipe_nw=8; /* PIPE_WORKERS=n: I/O worker threads (disk-parallel reads) */ +typedef struct { + _Atomic uint64_t cur; /* (gen<<8)|index; gen main-only, index 0..njobs (≤64) */ + _Atomic int njobs; /* current batch job count */ + _Atomic int eids[64]; /* current batch expert ids */ + _Atomic int layer; /* current batch layer */ + _Atomic int ready[64]; /* per-slot load-done flag */ + pthread_mutex_t mx; pthread_cond_t cv; /* ONLY for parking/waking idle workers */ + Model *m; + pthread_t th[16]; int nw; int started; +} PipePool; +static PipePool g_pp; + +static void *pipe_worker(void *arg){ + (void)arg; PipePool *p=&g_pp; uint64_t seen=0; + for(;;){ + pthread_mutex_lock(&p->mx); + while((atomic_load_explicit(&p->cur,memory_order_relaxed)>>8)==seen) + pthread_cond_wait(&p->cv,&p->mx); + pthread_mutex_unlock(&p->mx); + for(;;){ + uint64_t c=atomic_load_explicit(&p->cur,memory_order_acquire); + seen=c>>8; + uint32_t i=(uint32_t)(c & 0xFF); + if(i >= (uint32_t)atomic_load_explicit(&p->njobs,memory_order_relaxed)) + break; /* batch drained → re-park */ + if(atomic_compare_exchange_weak_explicit(&p->cur,&c,c+1, + memory_order_acq_rel,memory_order_relaxed)){ + int L =atomic_load_explicit(&p->layer,memory_order_relaxed); + int eid=atomic_load_explicit(&p->eids[i],memory_order_relaxed); /* AFTER winning CAS */ + expert_load(p->m,L,eid,&p->m->ws[i]); + atomic_store_explicit(&p->ready[i],1,memory_order_release); + } + /* CAS failed → another worker advanced index (or gen advanced): re-loop */ + } + } + return NULL; +} +static void pipe_init(Model *m){ + if(g_pp.started) return; + g_pp.m=m; g_pp.nw=g_pipe_nw; if(g_pp.nw>16) g_pp.nw=16; if(g_pp.nw<1) g_pp.nw=1; + atomic_store(&g_pp.cur,0); atomic_store(&g_pp.njobs,0); + pthread_mutex_init(&g_pp.mx,NULL); pthread_cond_init(&g_pp.cv,NULL); + for(int i=0;i>8)+1; + atomic_store_explicit(&g_pp.cur,(g<<8),memory_order_release); /* PUBLISH */ + pthread_mutex_lock(&g_pp.mx); pthread_cond_broadcast(&g_pp.cv); pthread_mutex_unlock(&g_pp.mx); +} +static inline void pipe_wait(int q){ + while(!atomic_load_explicit(&g_pp.ready[q],memory_order_acquire)) sched_yield(); +} + /* prefetch asincrono dei pesi di un expert (e delle sue scale .qs): avvia il readahead * cosi' le letture sincrone successive trovano la page-cache calda. */ static void expert_prefetch(Model *m, int layer, int eid){ @@ -1234,18 +1327,26 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out){ int *rows=malloc(S*sizeof(int)); float *rw=malloc(S*sizeof(float)); for(int base=0;basepin[layer]; for(int z=0;znpin[layer];z++) if(P[z].eid==eid){ m->hits++; use[j]=&P[z]; break; } if(!use[j]){ ESlot *Sl=m->ecache[layer]; int nn=m->ecn[layer]; for(int z=0;zhits++; Sl[z].used=++m->eclock; use[j]=&Sl[z]; break; } } - if(!use[j]){ use[j]=&m->ws[nmiss]; missk[nmiss++]=j; m->miss++; } + if(!use[j]){ qof[j]=nmiss; use[j]=&m->ws[nmiss]; missk[nmiss++]=j; m->miss++; } + } + if(nmiss){ + if(g_pipe){ /* PIPE: launch loads async, matmul overlaps them */ + if(!g_pp.started) pipe_init(m); + double t0=now_s(); + int eids[64]; for(int q=0;qt_edisk += now_s()-t0; /* dispatch only; real reads hide behind matmul */ + } else { double t0=now_s(); /* ORIGINALE: blocking parallel load */ + #pragma omp parallel for schedule(dynamic,1) + for(int q=0;qws[q]); + m->t_edisk += now_s()-t0; } } - if(nmiss){ double t0=now_s(); - #pragma omp parallel for schedule(dynamic,1) - for(int q=0;qws[q]); - m->t_edisk += now_s()-t0; } /* I/O ASINCRONO: readahead (WILLNEED) del blocco SUCCESSIVO mentre calcoliamo * questo — il kernel legge in background, le pread dopo trovano cache calda */ if(base+64=1 routing invariant. */ + if(g_pipe && qof[j]>=0){ double tw=now_s(); pipe_wait(qof[j]); m->t_edisk += now_s()-tw; } int nr=0; /* righe (posizioni) che usano questo expert */ for(int s=0;st_emm += now_s()-t0; } + /* No drain barrier: the per-expert pipe_wait(qof[j]) above (issued for every + * dispatched miss slot, before the nr==0 skip) already waited on all ws[] loads + * for this block, so they are complete before the LRU swap — and the gen-tagged + * cursor keeps any still-spinning worker off a wrong-generation slot. */ { ESlot *Sl=m->ecache[layer]; int *nn=&m->ecn[layer]; /* promozione LRU (swap buffer) */ int promo = nmissecap ? nmiss : m->ecap; for(int a=0;a