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
+4
View File
@@ -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,