cuda: fix pipe2 profile double-count + add COLI_CUDA_MTP opt-in (#292)

Two issues found by the diagnostic sweep in #292:

1. Profile double-count in pipe_layer_sparse: the outer t_emm span wrapped
   both the shared-expert GPU dispatch AND the moe() call, but moe()
   self-times its own t_emm internally (6 accumulation sites). Routed-expert
   matmul was counted twice -> accounted > elapsed -> 'other' went negative,
   and expert-matmul showed an inflated ~120s. Fix: split into two narrow
   spans around only the GPU work moe() does NOT cover (shared-expert
   dispatch, routed upload+add), letting moe() self-time as all other
   callers already do. Verified: 'other' now positive (14-17s),
   expert-matmul realistic (5-6s).

2. MTP blanket-disabled under CUDA (g_draft=0 when g_cuda_enabled). This
   was a conservative guard from #163 before the root cause (cold-expert
   fused-pair + IDOT kernel divergence) was fully diagnosed. GPU-resident
   experts have no divergence; the cold subset still does but achieves
   30-50% acceptance anyway (matching non-CUDA builds). Add COLI_CUDA_MTP=1
   opt-in so users can test speculation under CUDA. Default unchanged.
   Verified: COLI_CUDA_MTP=1 -> MTP ACTIVE (draft=3), 44% acceptance,
   3 forwards for 8 tokens (vs 7 without MTP).

Refs #292 #163
This commit is contained in:
woolcoxm
2026-07-15 18:33:25 -04:00
parent e56d483a15
commit 78c675eb8b
+21 -4
View File
@@ -3093,18 +3093,26 @@ static int pipe_layer_sparse(Model *m, Layer *l, int li, float *x_dev, int S, in
* No pipe_sync at the end: the next layer's pipe_download (sync cudaMemcpy)
* provides the implicit sync point. The fallback path (caller downloads x_dev)
* also uses pipe_download which syncs. This lets GPU work chain across layers
* without a per-layer stall. */
* without a per-layer stall.
*
* Profiling: moe() self-times its own t_emm (routed expert matmul). We time only
* the GPU work that moe() does NOT cover: the shared-expert dispatch and the
* routed-expert upload+add. Previously a single outer span wrapped everything
* including moe(), double-counting the routed-expert time and driving the
* profile's "other" bucket negative (#292). */
double te=now_s();
if(!coli_cuda_pipe_gemm(l->sh_gate.cuda,sg_d,nrm_d,S)) return 0;
if(!coli_cuda_pipe_gemm(l->sh_up.cuda,su_d,nrm_d,S)) return 0;
if(!coli_cuda_pipe_silu_mul(dev,sg_d,su_d,(size_t)S*sI)) return 0;
if(!coli_cuda_pipe_gemm(l->sh_down.cuda,y_d,sg_d,S)) return 0;
if(!coli_cuda_pipe_add(dev,x_dev,y_d,(size_t)S*D)) return 0; /* shared residual (async) */
m->t_emm += now_s()-te; /* shared-expert GPU dispatch only */
/* expert routed su CPU/gruppi GPU come oggi (shared saltata: la fa il device) */
moe(m,l,li,nrm_host,S,out_host,0);
moe(m,l,li,nrm_host,S,out_host,0); /* self-times its own t_emm */
te=now_s();
if(!coli_cuda_pipe_upload(dev,y_d,out_host,xb)) return 0; /* sync: waits for moe */
if(!coli_cuda_pipe_add(dev,x_dev,y_d,(size_t)S*D)) return 0; /* routed residual (async) */
m->t_emm+=now_s()-te;
m->t_emm += now_s()-te; /* routed-expert upload + add only */
return 1;
}
#endif
@@ -5213,7 +5221,16 @@ int main(int argc, char **argv){
Model m; double t0=now_s(); model_init(&m,snap,cap,ebits,dbits);
if(g_draft<0){
#ifdef COLI_CUDA
g_draft = (m.has_mtp&&!g_cuda_enabled) ? 3 : 0;
/* MTP is disabled under CUDA by default: cold (streaming) experts still
* run on the CPU, where the S==1 fused-pair kernel and the S>=2 IDOT
* kernel diverge in FP accumulation order, collapsing draft acceptance
* (#163). GPU-resident experts have no divergence, but the cold subset
* always exists on a single 16 GB card. COLI_CUDA_MTP=1 opts in for
* users who want to test speculation under CUDA the #163 thread shows
* acceptance can still reach 30-50% even with the cold-expert mismatch.
* See #292 for the diagnostic sweep that identified this. */
int cuda_mtp = getenv("COLI_CUDA_MTP") ? atoi(getenv("COLI_CUDA_MTP")) : 0;
g_draft = (m.has_mtp && (!g_cuda_enabled || cuda_mtp)) ? 3 : 0;
#else
g_draft = m.has_mtp ? 3 : 0;
#endif