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:
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user