diff --git a/c/backend_cuda.cu b/c/backend_cuda.cu index 236992b..4bab7c0 100644 --- a/c/backend_cuda.cu +++ b/c/backend_cuda.cu @@ -34,7 +34,7 @@ typedef struct { size_t qx_cap, qscale_cap; float *host_x,*host_y,*host_kv; size_t host_x_cap,host_y_cap,host_kv_cap; float *aq,*al,*ar,*ac; size_t aq_cap,al_cap,ar_cap,ac_cap; - float *pipe_buf[24]; size_t pipe_cap[24]; /* scratch persistenti del resident pipeline */ + float *pipe_buf[27]; size_t pipe_cap[27]; /* scratch persistenti del resident pipeline */ cudaStream_t stream; void *group_desc; size_t group_desc_cap; size_t tensor_count, tensor_bytes; @@ -455,7 +455,7 @@ extern "C" void coli_cuda_shutdown(void) { if (ctx->qx) cudaFree(ctx->qx); if (ctx->qscale) cudaFree(ctx->qscale); if(ctx->aq)cudaFree(ctx->aq);if(ctx->al)cudaFree(ctx->al);if(ctx->ar)cudaFree(ctx->ar);if(ctx->ac)cudaFree(ctx->ac); - for(int b=0;b<24;b++) if(ctx->pipe_buf[b]) cudaFree(ctx->pipe_buf[b]); + for(int b=0;b<27;b++) if(ctx->pipe_buf[b]) cudaFree(ctx->pipe_buf[b]); if (ctx->host_x) cudaFreeHost(ctx->host_x); if (ctx->host_y) cudaFreeHost(ctx->host_y); if (ctx->host_kv) cudaFreeHost(ctx->host_kv); @@ -994,7 +994,7 @@ __global__ static void pipe_rows_add(float *x,const float *partial,const int *ro * per layer (78 x ~10 alloc/richiesta erano puro churn). */ extern "C" float *coli_cuda_pipe_scratch(int device,int slot,size_t bytes){ DeviceContext *ctx=find_ctx(device); - if(slot<0||slot>=24||!select_ctx(ctx)) return NULL; + if(slot<0||slot>=27||!select_ctx(ctx)) return NULL; if(!reserve(&ctx->pipe_buf[slot],&ctx->pipe_cap[slot],bytes)) return NULL; return ctx->pipe_buf[slot]; } @@ -1046,6 +1046,82 @@ extern "C" int coli_cuda_pipe_rope_base(int device,float *v_dev,int pos_base,int pipe_rope_rows<<>>(v_dev,NULL,pos_base,stride,offset,R,heads,theta); return cuda_ok(cudaGetLastError(),"pipe rope base"); } +/* ---- device router (#431 PR-A) ------------------------------------------- + * Router for one decode row, entirely on the layer's home device: logits GEMV + * (E x D, tiny) + sigmoid, bias-augmented top-K selection, route-level TOPP + * truncation, norm_topk and routed_scale — a float-faithful clone of moe()'s + * plain routing path (colibri.c FASE A). Selection runs single-thread so the + * argmax order, tie-breaking (strict >, lowest index wins) and weight math + * match the CPU reference exactly; only the dot/expf rounding can differ, + * which is the documented kernel-family divergence class (#100/#163). + * Results are packed [idx[K] | w[K] | keff] in one scratch buffer and read + * back with a single tiny D2H. */ +__global__ void pipe_router_logits(const float *__restrict__ x, + const float *__restrict__ W, + const float *__restrict__ bias, + int D, float *logit, float *choice){ + int e = blockIdx.x; + const float *w = W + (size_t)e*D; + float acc = 0.f; + for(int i=threadIdx.x; i>1; s>0; s>>=1){ + if(threadIdx.xbv){bv=choice[e];best=e;} } + idx[kk]=best; w[kk]=logit[best]; + } + int Ke=Ksel; + if(topp>0.f && topp<1.f){ + for(int a=1;a=0 && w[b]=topp*tot){ Ke=kk+1; break; } } + } + if(norm_topk){ float sm=0.f; for(int kk=0;kk4096||Ksel<1||Ksel>64||!select_ctx(ctx)) return 0; + size_t pack=(size_t)Ksel*(sizeof(int)+sizeof(float))+sizeof(int); + float *logit=coli_cuda_pipe_scratch(device,22,(size_t)E*sizeof(float)); + float *chc =coli_cuda_pipe_scratch(device,23,(size_t)E*sizeof(float)); + char *out =(char*)coli_cuda_pipe_scratch(device,24,pack); + if(!logit||!chc||!out) return 0; + pipe_router_logits<<>>(x_dev,(const float*)rw_dev,(const float*)rb_dev,D,logit,chc); + pipe_router_select<<<1,1>>>(logit,chc,E,Ksel,topp,norm_topk,routed_scale,out); + if(!cuda_ok(cudaGetLastError(),"pipe router launch")) return 0; + char buf[64*(sizeof(int)+sizeof(float))+sizeof(int)]; + if(!cuda_ok(cudaMemcpy(buf,out,pack,cudaMemcpyDeviceToHost),"pipe router readback")) return 0; + memcpy(idx_host,buf,(size_t)Ksel*sizeof(int)); + memcpy(w_host,buf+Ksel*sizeof(int),(size_t)Ksel*sizeof(float)); + memcpy(keff_host,buf+Ksel*(sizeof(int)+sizeof(float)),sizeof(int)); + return 1; +} extern "C" int coli_cuda_pipe_copy2d(int device,float *dst,int dpitch,const float *src, int spitch,int width,int height){ DeviceContext *ctx=find_ctx(device); if(!select_ctx(ctx)) return 0; diff --git a/c/backend_cuda.h b/c/backend_cuda.h index 63c1f11..eb959c2 100644 --- a/c/backend_cuda.h +++ b/c/backend_cuda.h @@ -126,6 +126,10 @@ COLI_CUDA_DLLEXPORT int coli_cuda_pipe_rmsnorm_s(int device,float *y_dev,const f int xstride,int ystride); COLI_CUDA_DLLEXPORT int coli_cuda_pipe_rope_base(int device,float *v_dev,int pos_base,int rows, int stride,int offset,int R,int heads,float theta); +COLI_CUDA_DLLEXPORT int coli_cuda_pipe_router(int device,const float *x_dev, + const void *rw_dev,const void *rb_dev,int D,int E,int Ksel, + float topp,int norm_topk,float routed_scale, + int *idx_host,float *w_host,int *keff_host); COLI_CUDA_DLLEXPORT int coli_cuda_pipe_copy2d(int device,float *dst,int dpitch,const float *src, int spitch,int width,int height); COLI_CUDA_DLLEXPORT int coli_cuda_attention_project_batch_dev(ColiCudaTensor *kv_b,ColiCudaTensor *o_proj, diff --git a/c/backend_loader.c b/c/backend_loader.c index eedbd50..c87fbcd 100644 --- a/c/backend_loader.c +++ b/c/backend_loader.c @@ -76,6 +76,7 @@ typedef int (*fn_pipe_gemm)(ColiCudaTensor *t,float *y_dev,const float *x_dev,in typedef int (*fn_pipe_peer_copy)(int dst_dev,float *dst,int src_dev, const float *src,size_t bytes); typedef int (*fn_pipe_rmsnorm)(int device,float *y_dev,const float *x_dev, const float *w_dev,int S,int D,float eps); typedef int (*fn_pipe_rmsnorm_s)(int device,float *y_dev,const float *x_dev, const float *w_dev,int S,int D,float eps, int xstride,int ystride); +typedef int (*fn_pipe_router)(int device,const float *x_dev,const void *rw_dev,const void *rb_dev,int D,int E,int Ksel,float topp,int norm_topk,float routed_scale,int *idx_host,float *w_host,int *keff_host); typedef int (*fn_pipe_rope)(int device,float *v_dev,const int *pos_dev,int rows, int stride,int offset,int R,int heads,float theta); typedef int (*fn_pipe_rope_base)(int device,float *v_dev,int pos_base,int rows, int stride,int offset,int R,int heads,float theta); typedef int (*fn_pipe_rows_add)(int device,float *x_dev,const float *partial_dev, const int *rows_dev,int nrows,int D); @@ -123,6 +124,7 @@ static struct { fn_pipe_peer_copy pipe_peer_copy; fn_pipe_rmsnorm pipe_rmsnorm; fn_pipe_rmsnorm_s pipe_rmsnorm_s; + fn_pipe_router pipe_router; fn_pipe_rope pipe_rope; fn_pipe_rope_base pipe_rope_base; fn_pipe_rows_add pipe_rows_add; @@ -217,6 +219,7 @@ static int coli_cuda_load(void){ RESOLVE(pipe_peer_copy, fn_pipe_peer_copy) RESOLVE(pipe_rmsnorm, fn_pipe_rmsnorm) RESOLVE(pipe_rmsnorm_s, fn_pipe_rmsnorm_s) + RESOLVE(pipe_router, fn_pipe_router) RESOLVE(pipe_rope, fn_pipe_rope) RESOLVE(pipe_rope_base, fn_pipe_rope_base) RESOLVE(pipe_rows_add, fn_pipe_rows_add) @@ -407,6 +410,11 @@ int coli_cuda_pipe_rmsnorm(int device,float *y_dev,const float *x_dev, const flo return g_cuda.pipe_rmsnorm(device, y_dev, x_dev, w_dev, S, D, eps); } +int coli_cuda_pipe_router(int device,const float *x_dev,const void *rw_dev,const void *rb_dev,int D,int E,int Ksel,float topp,int norm_topk,float routed_scale,int *idx_host,float *w_host,int *keff_host){ + if(!g_cuda.available || !g_cuda.pipe_router){ return 0; } + return g_cuda.pipe_router(device, x_dev, rw_dev, rb_dev, D, E, Ksel, topp, norm_topk, routed_scale, idx_host, w_host, keff_host); +} + int coli_cuda_pipe_rmsnorm_s(int device,float *y_dev,const float *x_dev, const float *w_dev,int S,int D,float eps, int xstride,int ystride){ if(!g_cuda.available){ return 0; } return g_cuda.pipe_rmsnorm_s(device, y_dev, x_dev, w_dev, S, D, eps, xstride, ystride); diff --git a/c/colibri.c b/c/colibri.c index 1f4aa5b..a5044a5 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -69,10 +69,12 @@ static inline int omp_get_thread_num(void){ return 0; } #include static int g_metal_enabled; static int g_metal_gemm_min=16; /* COLI_METAL_GEMM_MIN: min rows to send a matmul_qt GEMM to GPU */ -/* routing precalcolata dalla GPU (layer CB): moe() la usa e salta la FASE A */ -static const int *g_pre_idx; static const float *g_pre_w; static const int *g_pre_keff; -static const float *g_pre_sh; /* output dello shared expert gia' calcolato su GPU */ +/* output dello shared expert gia' calcolato su GPU (solo Metal layer-CB) */ +static const float *g_pre_sh; #endif +/* routing precalcolata dalla GPU (Metal layer CB o device router CUDA, #431): + * moe() la usa e salta la FASE A. NULL = router su CPU. */ +static const int *g_pre_idx; static const float *g_pre_w; static const int *g_pre_keff; #ifdef __APPLE__ #include /* host_statistics64: MemAvailable di macOS */ #endif @@ -125,6 +127,10 @@ typedef struct { QT gate_proj, up_proj, down_proj; /* moe (sparse==1) */ float *router, *router_bias; /* router f32 (sensibile) */ +#ifdef COLI_CUDA + void *router_cuda, *router_bias_cuda; /* device router (#431 PR-A), lazy-uploaded */ + int router_cuda_bad; /* upload failed once: stay on the CPU router */ +#endif QT sh_gate, sh_up, sh_down; /* shared expert */ } Layer; @@ -1795,6 +1801,7 @@ static void qt_matvec_rows(const QT *t, int r0, int n, const float *x, float *y) static int g_absorb=-1; #ifdef COLI_CUDA static int g_cuda_pipe=0; /* COLI_CUDA_PIPE=1: prefill attention chain resident on the layer home device */ +static int g_cuda_router=0; /* COLI_CUDA_ROUTER=1 (#431 PR-A): router on the layer home device at decode */ #endif /* ABSORB: -1 auto (decode S<=4), 0 mai, 1 sempre (test) */ static int g_dsa_force=0; /* DSA_FORCE=1: selezione sempre attiva (test: top-min(k,T)=denso) */ static int cmp_fdesc(const void *a,const void *b){ @@ -2366,7 +2373,8 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int /* router in UN matmul batch: stessa matematica, via le S chiamate S=1 */ float *logits_all=falloc((int64_t)S*E); int pre_routed=0; (void)pre_routed; -#ifdef COLI_METAL + /* pre-routed shortcut: Metal layer-CB o device router CUDA (#431) — stessa + * contabilita' (#417: recency clock incluso), la selezione arriva dalla GPU */ if(g_pre_idx){ /* routing gia' calcolata dal layer CB (GPU) */ memcpy(idxs,g_pre_idx,(size_t)S*K*sizeof(int)); memcpy(ws,g_pre_w,(size_t)S*K*sizeof(float)); @@ -2396,7 +2404,6 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int } pre_routed=1; } -#endif if(!pre_routed) matmul(logits_all, x, l->router, S, D, E); if(!pre_routed) for(int s=0;seps)) return 0; + /* device router (#431 PR-A): route THIS row on the home device while the + * stream is still hot, then hand the selection to moe() through the same + * pre-routed shortcut the Metal layer-CB uses. Any failure (upload, launch, + * feature gate) falls back to the CPU router inside moe() — byte-identical + * behaviour, just slower. Gated to the plain routing path: CACHE_ROUTE / + * ROUTE_P / ROUTE_TRACE keep the CPU ranking they need. */ + static int lr_idx[64]; static float lr_w[64]; static int lr_keff[1]; + int dev_routed=0; + if(g_cuda_router && S==1 && !g_cache_route && g_route_p<=0.f && !g_route_fp + && c->n_experts<=4096 && c->topk<=64 && !l->router_cuda_bad){ + int E=c->n_experts, K=c->topk; + int Ksel = g_topk>0 ? (g_topk0 && g_topp<1.f) ? g_topp : 0.f; + if(!l->router_cuda){ + void *rw=coli_cuda_pipe_alloc(dev,(size_t)E*D*4); + void *rb=coli_cuda_pipe_alloc(dev,(size_t)E*4); + if(rw&&rb&&coli_cuda_pipe_upload(dev,rw,l->router,(size_t)E*D*4) + &&coli_cuda_pipe_upload(dev,rb,l->router_bias,(size_t)E*4)){ + l->router_cuda=rw; l->router_bias_cuda=rb; + } else { + if(rw)coli_cuda_pipe_free(dev,rw); if(rb)coli_cuda_pipe_free(dev,rb); + l->router_cuda_bad=1; + } + } + if(l->router_cuda && + coli_cuda_pipe_router(dev,nrm_d,l->router_cuda,l->router_bias_cuda, + D,E,Ksel,tp,c->norm_topk,c->routed_scale, + lr_idx,lr_w,lr_keff)) + dev_routed=1; + } if(!coli_cuda_pipe_download(dev,nrm_d,nrm_host,xb)) return 0; m->t_attn+=now_s()-ta; /* OVERLAP: issue the shared expert on the GPU BEFORE moe() runs on the CPU. @@ -3335,7 +3372,9 @@ static int pipe_layer_sparse(Model *m, Layer *l, int li, float *x_dev, int S, in if(!coli_cuda_pipe_add(dev,x_dev,y_d,(size_t)S*D)) return 0; /* shared residual (async) */ m->t_emm += now_s()-te; /* shared-expert GPU dispatch only */ /* expert routed su CPU/gruppi GPU come oggi (shared saltata: la fa il device) */ + if(dev_routed){ g_pre_idx=lr_idx; g_pre_w=lr_w; g_pre_keff=lr_keff; } moe(m,l,li,nrm_host,S,out_host,0); /* self-times its own t_emm */ + if(dev_routed){ g_pre_idx=NULL; g_pre_w=NULL; g_pre_keff=NULL; } te=now_s(); if(!coli_cuda_pipe_upload(dev,y_d,out_host,xb)) return 0; /* sync: waits for moe */ if(!coli_cuda_pipe_add(dev,x_dev,y_d,(size_t)S*D)) return 0; /* routed residual (async) */ @@ -5522,6 +5561,7 @@ int main(int argc, char **argv){ } g_cuda_dense=getenv("CUDA_DENSE")?atoi(getenv("CUDA_DENSE")):0; g_cuda_pipe=getenv("COLI_CUDA_PIPE")?atoi(getenv("COLI_CUDA_PIPE")):0; + g_cuda_router=getenv("COLI_CUDA_ROUTER")?atoi(getenv("COLI_CUDA_ROUTER")):0; const char *cuda_expert=getenv("CUDA_EXPERT_GB"); g_cuda_expert_auto=cuda_expert&&!strcmp(cuda_expert,"auto"); g_cuda_expert_gb=cuda_expert&&!g_cuda_expert_auto?atof(cuda_expert):0; diff --git a/c/tests/test_router_cuda.cu b/c/tests/test_router_cuda.cu new file mode 100644 index 0000000..4e714a6 --- /dev/null +++ b/c/tests/test_router_cuda.cu @@ -0,0 +1,79 @@ +/* Device-router kernel oracle (#431 PR-A). + * + * Feeds random activations/router weights through pipe_router_logits + + * pipe_router_select and checks against a CPU reference that replicates + * moe()'s plain routing path verbatim (sigmoid -> bias-augmented top-K by + * `choice`, weights from raw `logit`, route-level TOPP truncation, norm_topk, + * routed_scale). The dot/expf rounding may differ from libm at ~1e-6 rel, so + * a handful of near-tie index flips across trials is tolerated; the weight + * math itself must agree to 1e-4 rel on matching selections. + * + * Build: nvcc -O2 -std=c++17 -arch=native tests/test_router_cuda.cu -o tests/test_router_cuda + */ +#include +#include +#include +#include +#include + +/* pull in the kernel definitions (same idiom as the CPU tests' #include "../colibri.c") */ +#include "../backend_cuda.cu" + +static void cpu_ref(const float *x,const float *W,const float *bias,int D,int E, + int Ksel,float topp,int norm_topk,float rscale, + int *idx,float *w,int *keff){ + float *logit=(float*)malloc(E*sizeof(float)),*choice=(float*)malloc(E*sizeof(float)); + for(int e=0;ebv){bv=choice[e];best=e;} } + idx[kk]=best; w[kk]=logit[best]; } + int Ke=Ksel; + if(topp>0.f && topp<1.f){ + for(int a=1;a=0 && w[b]=topp*tot){ Ke=kk+1; break; } } } + if(norm_topk){ float sm=0; for(int kk=0;kk>>(x,W,b,D,lg,ch); + pipe_router_select<<<1,1>>>(lg,ch,E,K,topp,nt,rs,out); + char pack[K*8+4]; + if(cudaMemcpy(pack,out,sizeof(pack),cudaMemcpyDeviceToHost)!=cudaSuccess){ + printf("FAIL cuda\n"); return 1; } + int gidx[K],gkeff; float gw[K]; + memcpy(gidx,pack,K*4); memcpy(gw,pack+K*4,K*4); memcpy(&gkeff,pack+K*8,4); + int ridx[K],rkeff; float rw[K]; + cpu_ref(x,W,b,D,E,K,topp,nt,rs,ridx,rw,&rkeff); + int mism=0; for(int k2=0;k21e-4f*(fabsf(ref)+1e-6f)+1e-6f){ bad++; break; } + } + } + printf("router oracle: %d trials, %d near-tie flips, %d weight mismatches\n",TRIALS,flips,bad); + if(flips>4||bad){ printf("FAIL\n"); return 1; } + printf("OK\n"); return 0; +}