From 6af05c018ac4ce2624b6add4917cd085384d1933 Mon Sep 17 00:00:00 2001 From: ZacharyZcR Date: Wed, 26 Aug 2026 10:42:54 +0800 Subject: [PATCH 1/2] fix(cuda): pin speculative expert kernel family --- c/backend_cuda.cu | 32 +++++++++++++++++++++++++------- c/backend_cuda.h | 8 ++++++++ c/backend_loader.c | 21 +++++++++++++++++++++ c/colibri.c | 7 ++++--- c/tests/test_backend_cuda.cu | 18 ++++++++++++++++-- 5 files changed, 74 insertions(+), 12 deletions(-) diff --git a/c/backend_cuda.cu b/c/backend_cuda.cu index dff3eff3b..e484cdc4b 100644 --- a/c/backend_cuda.cu +++ b/c/backend_cuda.cu @@ -1836,11 +1836,12 @@ static void f8_group_launch(DeviceContext *ctx,GroupDesc *dev,int I,int D, grouped_down_f8<<stream>>>(ctx->y,ctx->gate,dev,D,I); } -extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, - ColiCudaTensor *const *ups, - ColiCudaTensor *const *downs, - const int *rows, int count, - float *y, const float *x) { +static int expert_group_impl(ColiCudaTensor *const *gates, + ColiCudaTensor *const *ups, + ColiCudaTensor *const *downs, + const int *rows, int count, + float *y, const float *x, + int pin_small_batch) { if (fault_injected()) return 0; if (!gates || !ups || !downs || !rows || !x || !y || count < 1) return 0; ColiCudaTensor *first=gates[0]; @@ -1906,7 +1907,7 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, * WMMA kernels are compiled out (COLI_HIP_NO_WMMA) the launch would * succeed with an EMPTY kernel and the output buffer would silently keep * stale data. Gate the branch like TC_W4A16 below does. */ - tc=tc&&COLI_GPU_HAS_WMMA&&all_s4&&D%32==0&&I%32==0&&D%8==0&&I%8==0; + tc=tc&&!pin_small_batch&&COLI_GPU_HAS_WMMA&&all_s4&&D%32==0&&I%32==0&&D%8==0&&I%8==0; int tc_min=getenv("COLI_CUDA_TC_MIN_ROWS")?atoi(getenv("COLI_CUDA_TC_MIN_ROWS")):8; for(int c=0;c=tc_min; if(all_e8){ @@ -1928,7 +1929,7 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, silu_mul<<<(unsigned)(((size_t)total*I+255)/256),256,0,ctx->stream>>>(ctx->gate,ctx->up,(size_t)total*I); quantize_s4_rows<<stream>>>(ctx->qx,ctx->qscale,ctx->gate,total,I); grouped_s4_wmma<<stream>>>(ctx->y,ctx->qx,ctx->qscale,dev,I,D,2); - }else if(all_s4&&COLI_GPU_HAS_WMMA&&ctx->compute_major>=7&&getenv("COLI_CUDA_TC_W4A16")&& + }else if(!pin_small_batch&&all_s4&&COLI_GPU_HAS_WMMA&&ctx->compute_major>=7&&getenv("COLI_CUDA_TC_W4A16")&& atoi(getenv("COLI_CUDA_TC_W4A16"))&& [&]{ int tc16_min=getenv("COLI_CUDA_TC_W4A16_MIN")?atoi(getenv("COLI_CUDA_TC_W4A16_MIN")):16; for(int c=0;c=tc16_min) return 1; @@ -2030,6 +2031,23 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, return 1; } +extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates, + ColiCudaTensor *const *ups, + ColiCudaTensor *const *downs, + const int *rows, int count, + float *y, const float *x) { + return expert_group_impl(gates,ups,downs,rows,count,y,x,0); +} + +extern "C" int coli_cuda_expert_group_pinned(ColiCudaTensor *const *gates, + ColiCudaTensor *const *ups, + ColiCudaTensor *const *downs, + const int *rows, int count, + float *y, const float *x, + int pin_small_batch) { + return expert_group_impl(gates,ups,downs,rows,count,y,x,pin_small_batch); +} + /* ---- 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 diff --git a/c/backend_cuda.h b/c/backend_cuda.h index 0373563a1..597f4273f 100644 --- a/c/backend_cuda.h +++ b/c/backend_cuda.h @@ -151,6 +151,14 @@ COLI_CUDA_DLLEXPORT int coli_cuda_expert_group(ColiCudaTensor *const *gates, ColiCudaTensor *const *downs, const int *rows, int count, float *y, const float *x); +/* Same operation, but force the small-batch grouped kernel family when + * pin_small_batch is nonzero. Speculative verification uses this to keep + * CUDA on the same numeric family as S=1 regardless of accepted draft depth. */ +COLI_CUDA_DLLEXPORT int coli_cuda_expert_group_pinned(ColiCudaTensor *const *gates, + ColiCudaTensor *const *ups, + ColiCudaTensor *const *downs, + const int *rows, int count, + float *y, const float *x, int pin_small_batch); /* Decode-only MLA weight-absorption core for one token. kv_b is [H*(Q+V),K]. */ COLI_CUDA_DLLEXPORT int coli_cuda_attention_absorb(ColiCudaTensor *kv_b,float *ctx,const float *q, diff --git a/c/backend_loader.c b/c/backend_loader.c index 22d91fdab..833e8e40b 100644 --- a/c/backend_loader.c +++ b/c/backend_loader.c @@ -77,6 +77,12 @@ typedef int (*fn_expert_mlp)(ColiCudaTensor *gate, ColiCudaTensor *up typedef int (*fn_expert_group)(ColiCudaTensor *const *gates, ColiCudaTensor *const *ups, ColiCudaTensor *const *downs, const int *rows, int count, float *y, const float *x); +typedef int (*fn_expert_group_pinned)(ColiCudaTensor *const *gates, + ColiCudaTensor *const *ups, + ColiCudaTensor *const *downs, + const int *rows, int count, + float *y, const float *x, + int pin_small_batch); typedef int (*fn_expert_group_issue)(ColiCudaTensor *const *gates, ColiCudaTensor *const *ups, ColiCudaTensor *const *downs, @@ -157,6 +163,7 @@ static struct { fn_group_stats_device group_stats_device; fn_expert_mlp expert_mlp; fn_expert_group expert_group; + fn_expert_group_pinned expert_group_pinned; fn_expert_group_issue expert_group_issue; fn_expert_group_take expert_group_take; fn_attention_absorb attention_absorb; @@ -1400,6 +1407,7 @@ static int coli_cuda_load(void){ RESOLVE(group_stats_device, fn_group_stats_device) RESOLVE(expert_mlp, fn_expert_mlp) RESOLVE(expert_group, fn_expert_group) + RESOLVE_OPT(expert_group_pinned, fn_expert_group_pinned) RESOLVE(expert_group_issue, fn_expert_group_issue) RESOLVE(expert_group_take, fn_expert_group_take) RESOLVE(attention_absorb, fn_attention_absorb) @@ -1541,6 +1549,19 @@ int coli_cuda_expert_group(ColiCudaTensor *const *gates, ColiCudaTensor *const * return g_cuda.expert_group(gates, ups, downs, rows, count, y, x); } +int coli_cuda_expert_group_pinned(ColiCudaTensor *const *gates, + ColiCudaTensor *const *ups, + ColiCudaTensor *const *downs, + const int *rows, int count, + float *y, const float *x, + int pin_small_batch){ + if(!g_cuda.available) return 0; + if(g_cuda.expert_group_pinned) + return g_cuda.expert_group_pinned(gates,ups,downs,rows,count,y,x,pin_small_batch); + if(pin_small_batch) return 0; /* old DLL: preserve SPEC_PIN via the CPU fallback */ + return g_cuda.expert_group(gates,ups,downs,rows,count,y,x); +} + int coli_cuda_expert_group_issue(ColiCudaTensor *const *gates, ColiCudaTensor *const *ups, ColiCudaTensor *const *downs, diff --git a/c/colibri.c b/c/colibri.c index 4a1d8f903..9c2ea2a8c 100644 --- a/c/colibri.c +++ b/c/colibri.c @@ -6115,8 +6115,9 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int #pragma omp parallel for if(g_cuda_ndev>1) schedule(static) for(int di=0;di=shared_min&&!l->shared_w4a16_failed&&!omp_in_parallel()&&g_cuda_enabled&& + if(shared_cuda==0&&!spec_pinned()&&S>=shared_min&&!l->shared_w4a16_failed&&!omp_in_parallel()&&g_cuda_enabled&& l->sh_gate.fmt==2&&l->sh_up.fmt==2&&l->sh_down.fmt==2&& getenv("COLI_CUDA_SHARED_W4A16")&&atoi(getenv("COLI_CUDA_SHARED_W4A16"))&& qt_cuda_upload(&l->sh_gate)&&qt_cuda_upload(&l->sh_up)&&qt_cuda_upload(&l->sh_down)){ diff --git a/c/tests/test_backend_cuda.cu b/c/tests/test_backend_cuda.cu index 40a4249de..6a8079424 100644 --- a/c/tests/test_backend_cuda.cu +++ b/c/tests/test_backend_cuda.cu @@ -375,7 +375,7 @@ int main(int argc, char **argv) { /* Native s4 WMMA path: compare the quantized-activation result against the existing FP32-activation/s4-weight grouped implementation. */ - uint8_t w4[32*32/2]; float ws4[32], gx4[64], scalar4[64], async4[64], tensor4[64]; + uint8_t w4[32*32/2]; float ws4[32], gx4[64], scalar4[64], async4[64], tensor4[64], pinned4[64]; for(int i=0;i<(int)sizeof(w4);i++){ int lo=((i%15)-7)&15,hi=(((i*3)%15)-7)&15; w4[i]=(uint8_t)(lo|(hi<<4)); @@ -400,12 +400,26 @@ int main(int argc, char **argv) { setenv("COLI_CUDA_TC_MIN_ROWS","1",1); if(!coli_cuda_expert_group(gg4,ug4,dg4,group_rows,2,tensor4,gx4)|| !relative_rms(tensor4,scalar4,64,0.30f))return 1; + if(!coli_cuda_expert_group_pinned(gg4,ug4,dg4,group_rows,2,pinned4,gx4,1)|| + std::memcmp(pinned4,scalar4,sizeof(pinned4))){ + std::fprintf(stderr,"pinned CUDA group did not bypass W4A4 Tensor Cores\n"); + return 1; + } unsetenv("COLI_CUDA_TC_INT4"); unsetenv("COLI_CUDA_TC_MIN_ROWS"); + setenv("COLI_CUDA_TC_W4A16","1",1); + setenv("COLI_CUDA_TC_W4A16_MIN","1",1); + if(!coli_cuda_expert_group_pinned(gg4,ug4,dg4,group_rows,2,pinned4,gx4,1)|| + std::memcmp(pinned4,scalar4,sizeof(pinned4))){ + std::fprintf(stderr,"pinned CUDA group did not bypass W4A16 Tensor Cores\n"); + return 1; + } + unsetenv("COLI_CUDA_TC_W4A16"); + unsetenv("COLI_CUDA_TC_W4A16_MIN"); coli_cuda_tensor_free(g4);coli_cuda_tensor_free(u4);coli_cuda_tensor_free(d4); uint64_t group_calls=0,group_experts=0,group_total_rows=0; coli_cuda_group_stats(&group_calls,&group_experts,&group_total_rows,nullptr,nullptr,nullptr); - if(group_calls!=4||group_experts!=8||group_total_rows!=8) return 1; + if(group_calls!=6||group_experts!=12||group_total_rows!=12) return 1; coli_cuda_stats(-1, &count, &bytes); if (count != 7 || bytes != 166) { From 9f4a8d1726df9c528edc1e2714095b24c40c574c Mon Sep 17 00:00:00 2001 From: ZacharyZcR Date: Wed, 26 Aug 2026 10:51:57 +0800 Subject: [PATCH 2/2] test(loader): account for pinned CUDA group API --- c/tests/test_backend_loader.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/c/tests/test_backend_loader.py b/c/tests/test_backend_loader.py index 60ef934a7..452f343b8 100644 --- a/c/tests/test_backend_loader.py +++ b/c/tests/test_backend_loader.py @@ -874,7 +874,7 @@ def tearDownClass(cls): cls.fixture = None def test_abi_is_derived_from_the_loader_source(self): - """47 mandatory + 3 optional, parsed from backend_loader.c. + """47 mandatory + 4 optional, parsed from backend_loader.c. The counts are a deliberate tripwire: adding a RESOLVE to the loader widens the ABI every Windows DLL must satisfy, and that should be a @@ -885,8 +885,8 @@ def test_abi_is_derived_from_the_loader_source(self): """ f = self.fixture self.assertEqual(len(f.mandatory), 47) - self.assertEqual(len(f.optional), 3) - self.assertEqual(len(f.exports), 50) + self.assertEqual(len(f.optional), 4) + self.assertEqual(len(f.exports), 51) self.assertEqual(len(f.exports), len(f.mandatory) + len(f.optional)) self.assertIn("coli_cuda_init", f.mandatory) self.assertIn("coli_cuda_e8_set_grid", f.optional) @@ -894,6 +894,8 @@ def test_abi_is_derived_from_the_loader_source(self): self.assertIn("coli_cuda_attention_project_ragged", f.mandatory) # fp8_set_lut: fmt=8 e4m3 dense/expert kernels (#817). self.assertIn("coli_cuda_fp8_set_lut", f.optional) + # expert_group_pinned: old DLLs remain usable outside SPEC_PIN (#689). + self.assertIn("coli_cuda_expert_group_pinned", f.optional) def test_both_runtimes_exist_with_the_production_basename(self): """Same basename, different directories — the conflict precondition."""