url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cuda: COLI_GROUP_ASYNC=1 — async expert-group issue/take with CPU/GPU overlap at decode (opt-in, +6-8% measured)

Browse Source
This commit is contained in:
ZacharyZcR
2026-07-17 15:32:09 +08:00
parent 61004dcb84
commit ab55f4900c
3 changed files with 213 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/backend_cuda.cu
+69
View File
@@ -38,6 +38,7 @@ typedef struct {
cudaStream_t stream;
void *group_desc; size_t group_desc_cap;
size_t tensor_count, tensor_bytes;
int group_pending; size_t group_pending_bytes; /* async expert-group in flight (Inc.4) */
} DeviceContext;
typedef struct {
@@ -757,6 +758,74 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates,
return 1;
}
/* ---- Async expert group (Inc.4): issue/take split of coli_cuda_expert_group ----
* The measured cost of the sync call at decode is ~0.45 ms/call of HOST-side wait
* (stream sync + staging), vs ~0.18 ms of actual GPU work — 70% tax, paid ~5x per
* layer because a token's 8 experts scatter across devices. issue() stages and
* launches on the device stream and returns immediately; take() syncs and hands
* back the pinned result rows. One issue may be outstanding per device; moe()
* takes at each layer end, which also orders the next layer's reuse of the ctx
* scratch buffers. Small batches only (decode/spec): bigger totals keep the sync
* path with its TC variants. Numerics are the sync path's small-batch kernels,
* so greedy output is byte-identical by construction. */
extern "C" int coli_cuda_expert_group_issue(ColiCudaTensor *const *gates,
ColiCudaTensor *const *ups,
ColiCudaTensor *const *downs,
const int *rows, int count,
const float *x) {
if (!gates || !ups || !downs || !rows || !x || count < 1 || count > 64) return 0;
ColiCudaTensor *first=gates[0];
if (!first) return 0;
int device=first->device,D=first->I,I=first->O,total=0;
GroupDesc host[64];
for(int c=0;c<count;c++){
ColiCudaTensor *g=gates[c],*u=ups[c],*d=downs[c];
if(!g||!u||!d||rows[c]<1||g->device!=device||u->device!=device||d->device!=device||
g->I!=D||u->I!=D||g->O!=I||u->O!=I||d->I!=I||d->O!=D) return 0;
host[c]={g->weights,u->weights,d->weights,g->scales,u->scales,d->scales,
g->fmt,u->fmt,d->fmt,rows[c],total};
total+=rows[c];
}
if(total>8) return 0; /* decode-scale only */
DeviceContext *ctx=find_ctx(device); if(!ctx||ctx->group_pending||!select_ctx(ctx)) return 0;
size_t xb=(size_t)total*D*sizeof(float), ib=(size_t)total*I*sizeof(float);
if(!reserve(&ctx->x,&ctx->x_cap,xb)||!reserve(&ctx->y,&ctx->y_cap,xb)||
!reserve(&ctx->gate,&ctx->gate_cap,ib)||!reserve(&ctx->up,&ctx->up_cap,ib)||
!reserve_pinned(&ctx->host_x,&ctx->host_x_cap,xb)||
!reserve_pinned(&ctx->host_y,&ctx->host_y_cap,xb)) return 0;
std::memcpy(ctx->host_x,x,xb);
if(!cuda_ok(cudaMemcpyAsync(ctx->x,ctx->host_x,xb,cudaMemcpyHostToDevice,ctx->stream),
"expert group issue upload")) return 0;
for(int c=0;c<count;c++){
int r=rows[c];
float *g16=ctx->gate+(size_t)host[c].offset*I,*u16=ctx->up+(size_t)host[c].offset*I;
float *x16=ctx->x+(size_t)host[c].offset*D,*y16=ctx->y+(size_t)host[c].offset*D;
quant_matmul<<<dim3((unsigned)I,(unsigned)r),256,0,ctx->stream>>>(g16,x16,
host[c].g,host[c].gs,host[c].gf,r,D,I,row_bytes(host[c].gf,D));
quant_matmul<<<dim3((unsigned)I,(unsigned)r),256,0,ctx->stream>>>(u16,x16,
host[c].u,host[c].us,host[c].uf,r,D,I,row_bytes(host[c].uf,D));
silu_mul<<<(unsigned)(((size_t)r*I+255)/256),256,0,ctx->stream>>>(g16,u16,(size_t)r*I);
quant_matmul<<<dim3((unsigned)D,(unsigned)r),256,0,ctx->stream>>>(y16,g16,
host[c].d,host[c].ds,host[c].df,r,I,D,row_bytes(host[c].df,I));
}
if(!cuda_ok(cudaGetLastError(),"expert group issue launch")||
!cuda_ok(cudaMemcpyAsync(ctx->host_y,ctx->y,xb,cudaMemcpyDeviceToHost,ctx->stream),
"expert group issue download")) return 0;
ctx->group_pending=1; ctx->group_pending_bytes=xb;
{ std::lock_guard<std::mutex> lock(g_group_stats_mu);
g_group_calls++; g_group_experts+=(uint64_t)count; g_group_rows+=(uint64_t)total; }
return 1;
}
extern "C" const float *coli_cuda_expert_group_take(int device) {
DeviceContext *ctx=find_ctx(device);
if(!ctx||!ctx->group_pending) return nullptr;
ctx->group_pending=0;
if(!select_ctx(ctx)) return nullptr;
if(!cuda_ok(cudaStreamSynchronize(ctx->stream),"expert group take")) return nullptr;
return ctx->host_y;
}
extern "C" int coli_cuda_attention_absorb(ColiCudaTensor *w,float *ctx,const float *q,
const float *latent,const float *rope,int H,int Q,
c/backend_cuda.h
+10
View File
@@ -67,6 +67,16 @@ COLI_CUDA_DLLEXPORT int coli_cuda_shared_mlp_w4a16(ColiCudaTensor *gate, ColiCud
/* Packed group of same-shaped experts. Inputs and outputs contain sum(rows)
* consecutive [D] rows in call order. */
/* Async issue/take split of the group call below (Inc.4): issue launches on the
* device stream and returns; take syncs and returns the pinned result rows (valid
* until the next issue on that device). Small totals only (<=8 rows); one
* outstanding issue per device. */
COLI_CUDA_DLLEXPORT int coli_cuda_expert_group_issue(ColiCudaTensor *const *gates,
ColiCudaTensor *const *ups,
ColiCudaTensor *const *downs,
const int *rows, int count, const float *x);
COLI_CUDA_DLLEXPORT const float *coli_cuda_expert_group_take(int device);
COLI_CUDA_DLLEXPORT int coli_cuda_expert_group(ColiCudaTensor *const *gates,
ColiCudaTensor *const *ups,
ColiCudaTensor *const *downs,
c/colibri.c
+134
View File
@@ -2516,6 +2516,11 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int
}
#ifdef COLI_CUDA
ESlot *group_e[64]; int group_n[64]; int ngroup=0;
/* Inc.4 overlap stash: pass-1 packing kept for the take phase after the CPU loop */
ESlot *eg_e[64]; int eg_n[64], eg_row[64][4], eg_npg=0; float eg_w[64][4];
int dev_nc0[COLI_CUDA_MAX_DEVICES], dev_off0[COLI_CUDA_MAX_DEVICES],
dev_total0[COLI_CUDA_MAX_DEVICES], dev_which0[COLI_CUDA_MAX_DEVICES][64];
memset(dev_nc0,0,sizeof(dev_nc0)); (void)eg_npg; (void)dev_total0; (void)dev_off0;
#endif
#ifdef COLI_METAL
if(g_metal_enabled){
@@ -2544,9 +2549,82 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int
free(mxg); free(mrows); free(mrw);
}
#undef MB_BUILD
#endif
#ifdef COLI_CUDA
/* Inc.4 pass 1: collect the VRAM-resident experts' groups and ISSUE them async
* BEFORE the CPU loop below, so the GPU computes its share while the CPU works
* through the RAM-tier/miss rows — t_emm becomes max(cpu, gpu) instead of the
* sum. Only resident experts are collected (misses are never cuda_eligible), so
* no pipe_wait is needed here; the CPU loop keeps its own waits. Any issue
* failure drops the layer back to the collect-in-loop + sync-group path. */
int early_issued=0, done_j[64]={0};
{
static int g_group_async2=-1;
if(g_group_async2<0) g_group_async2=getenv("COLI_GROUP_ASYNC")?atoi(getenv("COLI_GROUP_ASYNC")):0;
if(!metal_done && g_group_async2 && group_enabled && S<=4 && g_cuda_enabled &&
g_cuda_ndev>0 && !omp_in_parallel()){
ESlot *pg_e[64]; int pg_n[64], pg_j[64], npg=0;
int prow[64][4]; float pw[64][4];
for(int j=0;j<nb;j++){ ESlot *e=use[j]; int eid=uniq[base+j];
if(!(e->g.cuda_eligible&&e->u.cuda_eligible&&e->d.cuda_eligible)) continue;
int nr=0;
for(int s=0;s<S && nr<4;s++) for(int kk=0;kk<keff[s];kk++)
if(idxs[(int64_t)s*K+kk]==eid){ prow[npg][nr]=s; pw[npg][nr]=ws[(int64_t)s*K+kk]; nr++; break; }
if(!nr) continue;
pg_e[npg]=e; pg_n[npg]=nr; pg_j[npg]=j; npg++;
}
if(npg){
/* pack per device exactly like the sync path below */
ColiCudaTensor *pd_g[COLI_CUDA_MAX_DEVICES][64],*pd_u[COLI_CUDA_MAX_DEVICES][64],*pd_d[COLI_CUDA_MAX_DEVICES][64];
int pd_rows[COLI_CUDA_MAX_DEVICES][64],pd_which[COLI_CUDA_MAX_DEVICES][64];
int pd_nc[COLI_CUDA_MAX_DEVICES]={0},pd_total[COLI_CUDA_MAX_DEVICES]={0},pd_off[COLI_CUDA_MAX_DEVICES]={0};
for(int di=0;di<g_cuda_ndev;di++) for(int q=0;q<npg;q++)
if(pg_e[q]->g.cuda_device==g_cuda_devices[di]) pd_total[di]+=pg_n[q];
for(int di=1;di<g_cuda_ndev;di++) pd_off[di]=pd_off[di-1]+pd_total[di-1];
for(int di=0;di<g_cuda_ndev;di++){
int cursor=0,device=g_cuda_devices[di];
for(int q=0;q<npg;q++) if(pg_e[q]->g.cuda_device==device){
int nc=pd_nc[di]++; ESlot *e=pg_e[q];
pd_g[di][nc]=e->g.cuda; pd_u[di][nc]=e->u.cuda; pd_d[di][nc]=e->d.cuda;
pd_rows[di][nc]=pg_n[q]; pd_which[di][nc]=q;
for(int r=0;r<pg_n[q];r++) memcpy(group_x+(int64_t)(pd_off[di]+cursor+r)*D,
x+(int64_t)prow[q][r]*D,D*sizeof(float));
cursor+=pg_n[q];
}
}
double tg0=now_s();
int all=1, issued[COLI_CUDA_MAX_DEVICES]={0};
for(int di=0;di<g_cuda_ndev && all;di++) if(pd_nc[di])
all=issued[di]=coli_cuda_expert_group_issue(pd_g[di],pd_u[di],pd_d[di],
pd_rows[di],pd_nc[di],group_x+(int64_t)pd_off[di]*D);
if(all){
static int announced2;
if(!announced2){ announced2=1; fprintf(stderr,"[CUDA] expert group overlap active\n"); }
early_issued=1;
for(int q=0;q<npg;q++) done_j[pg_j[q]]=1;
/* stash packing for the take phase */
for(int di=0;di<g_cuda_ndev;di++){ dev_nc0[di]=pd_nc[di]; dev_off0[di]=pd_off[di]; dev_total0[di]=pd_total[di];
for(int q=0;q<pd_nc[di];q++) dev_which0[di][q]=pd_which[di][q]; }
for(int q=0;q<npg;q++){ eg_e[q]=pg_e[q]; eg_n[q]=pg_n[q];
for(int r=0;r<pg_n[q];r++){ eg_row[q][r]=prow[q][r]; eg_w[q][r]=pw[q][r]; } }
eg_npg=npg;
m->t_emm+=now_s()-tg0;
for(int q=0;q<npg;q++){ /* bookkeeping normally done in the loop */
m->gpu_expert_calls++;
}
} else {
for(int di=0;di<g_cuda_ndev;di++)
if(issued[di]) coli_cuda_expert_group_take(g_cuda_devices[di]);
}
}
}
}
#endif
if(!metal_done)
for(int j=0;j<nb;j++){ int eid=uniq[base+j]; ESlot *e=use[j];
#ifdef COLI_CUDA
if(early_issued && done_j[j]) continue; /* computing on the GPU right now */
#endif
/* Drain this miss's async load BEFORE the nr==0 early-exit below: every
* dispatched slot must be waited before the end-of-block LRU swap can reuse
* its ws[] slab, so correctness does not depend on the nr>=1 routing invariant.
@@ -2593,6 +2671,35 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int
m->cpu_expert_rows+=(uint64_t)nr;}
}
#ifdef COLI_CUDA
/* Inc.4 take phase: the CPU loop above ran while the GPU computed the issued
* groups — collect them now. A failed device recomputes its experts on the CPU
* (expert_host_ensure reloads slabs released by CUDA_RELEASE_HOST). */
if(early_issued){
double tg1=now_s();
for(int di=0;di<g_cuda_ndev;di++) if(dev_nc0[di]){
const float *hy=coli_cuda_expert_group_take(g_cuda_devices[di]);
int cur=0;
for(int q=0;q<dev_nc0[di];q++){
int gi=dev_which0[di][q], nr=eg_n[gi];
if(hy){
for(int r=0;r<nr;r++){ float *os=out+(int64_t)eg_row[gi][r]*D; float wgt=eg_w[gi][r];
const float *hr=hy+(int64_t)(cur+r)*D;
for(int d=0;d<D;d++) os[d]+=wgt*hr[d]; }
} else {
ESlot *e=eg_e[gi];
for(int r=0;r<nr;r++) memcpy(xg+(int64_t)r*D,x+(int64_t)eg_row[gi][r]*D,D*sizeof(float));
expert_host_ensure(m,layer,e);
expert_gate_up(gg,uu,xg,&e->g,&e->u,nr);
for(int64_t z=0;z<(int64_t)nr*I;z++) gg[z]=siluf(gg[z])*uu[z];
matmul_qt(hh,gg,&e->d,nr);
for(int r=0;r<nr;r++){ float *os=out+(int64_t)eg_row[gi][r]*D; float wgt=eg_w[gi][r];
for(int d=0;d<D;d++) os[d]+=wgt*hh[(int64_t)r*D+d]; }
}
cur+=nr;
}
}
m->t_emm+=now_s()-tg1;
}
ColiCudaTensor *dev_g[COLI_CUDA_MAX_DEVICES][64],*dev_u[COLI_CUDA_MAX_DEVICES][64];
ColiCudaTensor *dev_d[COLI_CUDA_MAX_DEVICES][64];
int dev_rows[COLI_CUDA_MAX_DEVICES][64],dev_which[COLI_CUDA_MAX_DEVICES][64];
@@ -2614,6 +2721,33 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int
}
}
double tg=now_s();
/* Inc.4: at decode scale, issue every device's group WITHOUT syncing, then take
* them all — one stream sync per device per layer instead of a full staged
* round-trip per call (measured: ~70% of the sync call is host-side wait).
* Any issue failure drains what was issued and the whole layer falls back to
* the sync path below, which recomputes from group_x (idempotent). */
int async_done=0;
static int g_group_async=-1;
if(g_group_async<0) g_group_async=getenv("COLI_GROUP_ASYNC")?atoi(getenv("COLI_GROUP_ASYNC")):0;
if(g_group_async && S<=4 && g_cuda_ndev>0){
int issued[COLI_CUDA_MAX_DEVICES]={0}, all=1;
for(int di=0;di<g_cuda_ndev && all;di++) if(dev_nc[di])
all=issued[di]=coli_cuda_expert_group_issue(dev_g[di],dev_u[di],dev_d[di],
dev_rows[di],dev_nc[di],group_x+(int64_t)dev_off[di]*D);
if(all){
static int announced;
if(!announced){ announced=1; fprintf(stderr,"[CUDA] expert group async path active\n"); }
async_done=1;
for(int di=0;di<g_cuda_ndev;di++) if(dev_nc[di]){
const float *hy=coli_cuda_expert_group_take(g_cuda_devices[di]);
if(hy){ dev_ok[di]=1;
memcpy(group_y+(int64_t)dev_off[di]*D,hy,(size_t)dev_total[di]*D*sizeof(float)); }
else dev_ok[di]=0; /* per-device sync failure: CPU fallback below */
}
} else for(int di=0;di<g_cuda_ndev;di++)
if(issued[di]) coli_cuda_expert_group_take(g_cuda_devices[di]);
}
if(!async_done)
#pragma omp parallel for if(g_cuda_ndev>1) schedule(static)
for(int di=0;di<g_cuda_ndev;di++) if(dev_nc[di]){
double td=g_prof?now_s():0;