cuda: decode expert groups take the grouped kernels even under TC_W4A16 (#431)
The TC_W4A16 branch of coli_cuda_expert_group handled every expert in a per-expert loop: rows >= threshold got the Tensor Core path, everything below fell back to 4 naive launches per expert. At decode every expert has 1 row, so the whole group rode the fallback — ~981 quant_matmul micro-launches per token (#431's measured flood) — while the grouped 3-launch path (grouped_hidden_w4_dual + silu + grouped_down_w4) sat one else-if below, unreachable whenever the PREFILL tuning flag was set. Gate the branch on 'at least one expert reaches the TC row threshold': all-small groups (decode) now fall through to the grouped kernels. Measured on 6x RTX 5090 (full residency, 39 forwards under nsys): expert-side quant_matmul instances drop 981 -> 337 per forward, total launches ~1,490 -> ~850 per token (-43%). Wall-clock is parity at the A/B operating point — the win is structural (PR-C graph node count, launch-tax share at champion speed). Behavioural fix folded in: before this change, toggling TC_W4A16 — a prefill-only optimization — changed DECODE output text (kernel-family divergence, #100 class). After it, decode always uses the grouped family: TC_W4A16=1 and =0 now produce byte-identical decode text (verified, 96-token greedy A/B), and the flag affects only the prefill it was built for.
This commit is contained in:
+9
-1
@@ -689,7 +689,15 @@ extern "C" int coli_cuda_expert_group(ColiCudaTensor *const *gates,
|
||||
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&&ctx->compute_major>=7&&getenv("COLI_CUDA_TC_W4A16")&&
|
||||
atoi(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;
|
||||
return 0; }()){
|
||||
/* At least one expert has enough rows for a Tensor Core tile. Groups
|
||||
* where EVERY expert is below the threshold (decode: r=1) fall through
|
||||
* to the grouped-W4 path below — 3 launches for the whole group instead
|
||||
* of 4 per expert (#431: the launch flood measured at ~981 micro-kernels
|
||||
* per token came from decode riding this branch's per-expert fallback). */
|
||||
/* W4A16 Tensor Core per gruppo: attivazioni fp16 per tile (lossless al
|
||||
* contrario del path W4A4), un lancio per expert dentro lo stream —
|
||||
* l'overhead di lancio e' trascurabile rispetto ai GEMM. */
|
||||
|
||||
Reference in New Issue
Block a user