cuda: COLI_CUDA_ROUTER=1 — route the decode row on the layer's home device (#431 PR-A)

First increment of the #431 plan (device router -> indirect kernels ->
one-graph decode). At decode (S=1) on the pipe2 path, the router runs on
the layer's home device: a tiny E x D logits GEMV + sigmoid, then a
single-thread selection kernel that clones moe()'s plain routing path
verbatim — bias-augmented top-K by choice with strict-> tie-breaking,
weights from the raw logit, route-level TOPP truncation, norm_topk,
routed_scale. Results pack into one scratch buffer and come back in a
single ~68-byte D2H; moe() consumes them through the same pre-routed
shortcut the Metal layer-CB uses (g_pre_idx, #417 bookkeeping included),
so usage/heat/recency accounting is identical to the CPU router.

Structural value: routing becomes available ON the device timeline,
which is what PR-B (indirect expert kernels, static topology) and PR-C
(whole-decode CUDA Graph) build on.

Opt-in, default off. Gated to the plain routing path — CACHE_ROUTE,
ROUTE_P and ROUTE_TRACE keep the CPU ranking they need; any upload or
launch failure falls back to the CPU router silently. Router weights
(E x D f32, ~6.3 MB/layer) upload lazily to the layer's home device.

tests/test_router_cuda.cu: kernel-vs-CPU-reference oracle over 200
random trials (mixed TOPP/norm_topk/scale): 200/200 exact selections,
zero near-tie flips, zero weight mismatches on a 5090.
This commit is contained in:
ZacharyZcR
2026-07-20 02:53:43 +08:00
parent 61004dcb84
commit 113ece3bc7
5 changed files with 215 additions and 8 deletions
+8
View File
@@ -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);