Skip to content

fix(cuda): pin speculative expert kernel family #1234

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
JustVugg merged 2 commits into from
Aug 28, 2026
Merged

fix(cuda): pin speculative expert kernel family #1234

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6af05c0
fix(cuda): pin speculative expert kernel family
ZacharyZcR Aug 26, 2026
9f4a8d1
test(loader): account for pinned CUDA group API
ZacharyZcR Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 25 additions & 7 deletions c/backend_cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1836,11 +1836,12 @@ static void f8_group_launch(DeviceContext *ctx,GroupDesc *dev,int I,int D,
grouped_down_f8<<<og,256,0,ctx->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];
Expand Down Expand Up @@ -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<count&&tc;c++)tc=rows[c]>=tc_min;
if(all_e8){
Expand All @@ -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<<<total,256,0,ctx->stream>>>(ctx->qx,ctx->qscale,ctx->gate,total,I);
grouped_s4_wmma<<<dim3((unsigned)((D+63)/64),(unsigned)count),256,0,ctx->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<count;c++) if(rows[c]>=tc16_min) return 1;
Expand Down Expand Up @@ -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
Expand Down
8 changes: 8 additions & 0 deletions c/backend_cuda.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
21 changes: 21 additions & 0 deletions c/backend_loader.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down
7 changes: 4 additions & 3 deletions c/colibri.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<g_cuda_ndev;di++) if(dev_nc[di]&&dev_ok[di]==0){
double td=g_prof?now_s():0;
dev_ok[di]=coli_cuda_expert_group(dev_g[di],dev_u[di],dev_d[di],dev_rows[di],dev_nc[di],
group_y+(int64_t)dev_off[di]*D,group_x+(int64_t)dev_off[di]*D);
dev_ok[di]=coli_cuda_expert_group_pinned(dev_g[di],dev_u[di],dev_d[di],
dev_rows[di],dev_nc[di],group_y+(int64_t)dev_off[di]*D,
group_x+(int64_t)dev_off[di]*D,spec_pinned());
if(g_prof)dev_time[di]=now_s()-td;
}
for(int di=0;di<g_cuda_ndev;di++){
Expand Down Expand Up @@ -6176,7 +6177,7 @@ static void moe(Model *m, Layer *l, int layer, float *x, int S, float *out, int
int shared_min=getenv("COLI_CUDA_SHARED_W4A16_MIN_ROWS")?
atoi(getenv("COLI_CUDA_SHARED_W4A16_MIN_ROWS")):32;
if(shared_min<16)shared_min=16;
if(shared_cuda==0&&S>=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)){
Expand Down
18 changes: 16 additions & 2 deletions c/tests/test_backend_cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
Expand All @@ -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) {
Expand Down
8 changes: 5 additions & 3 deletions c/tests/test_backend_loader.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -885,15 +885,17 @@ 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)
# attention_project_ragged: paged ragged KV runtime (#795).
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."""
Expand Down
Loading