cuda: COLI_GROUP_ASYNC=1 — overlap the CPU expert rows with the GPU groups at decode (opt-in, +6-8%)

[ZacharyZcR]

What

At S≤4 decode on a multi-GPU full-residency host, moe() runs two phases serially: the VRAM-resident experts' grouped GPU calls, and the CPU rows (RAM-tier + misses). Per-phase measurement on 6× RTX 5090 put them at roughly 50 ms + 50 ms per token — a textbook overlap opportunity.

This adds an issue/take split of coli_cuda_expert_group (launch on the device stream and return; sync at layer end) and a two-pass moe block: pass 1 collects the VRAM experts' groups and issues them async; the CPU loop then computes its rows while the GPUs work; a take phase collects and accumulates. Per-device failure recomputes on the CPU (expert_host_ensure reloads slabs released by CUDA_RELEASE_HOST); any issue failure drops the whole layer back to the existing sync path.

Opt-in: COLI_GROUP_ASYNC=1, default off.

Measured (6× RTX 5090, full residency, warm segment 256→512 of a 512-token greedy run, .coli_usage snapshot-restored between arms)

arm	warm decode	expert-matmul bucket
sync (today)	4.04–4.09 tok/s	53.0 s
overlap	4.29–4.40 tok/s (+6–8%)	47.7–48.2 s (−9–10%)

Reproduced across two independent A/B rounds. The async-issue without the overlap is a wash (+1.5%) — the win is the concurrency, not the removed sync.

The honest caveat (why default-off)

Greedy output under the overlap is not byte-identical to the sync path: a near-tie token flips early in the run (with usage state controlled, so it is not placement drift — some floating-point accumulation-order difference I have not yet pinned down; the kernels are byte-for-byte the sync path's small-batch kernels and the packing order mirrors the sync path's). Same class as the documented CPU-vs-GPU / kernel-family divergences (#100): every emitted token remains the argmax of a valid forward, but the stream forks. Under #294's standard the flag also implies it should stay off during speculation until verified — at S≤4 with drafts live, SPEC_PIN semantics take precedence.

Follow-ups I intend regardless of merge: (a) a position-aligned ROUTE_TRACE diff (the #163 method) to locate the first flipped layer and, if it's an ordering bug rather than inherent, fix it to byte-exactness; (b) profiling the remaining serial residue in the overlap window (take-phase device syncs, pass-1 packing).

[ZacharyZcR]

Pushed two follow-ups from the profiling I promised:

1. Layernorm weights cached on-device (pipe_layer_sparse). They're constants, but the pipe was doing two synchronous 24 KB H2D uploads per layer per token — 152 sync uploads/token — re-sending the same bytes. Now uploaded once per layer and kept on the layer's device. Pure win, no numerics change (it moves constant uploads, not math).

2. The overlap window measured — and it reframes this whole PR. With COLI_CUDA_PROFILE=1 the async path now reports its window split. On 6× RTX 5090 full residency, per 512-token turn:

overlap window: pack+issue 3.20s | cpu-rows 44.55s | take(sync+acc) 0.63s

The GPU side of a token's experts is ~3.8 s; the CPU side (the ~90 RAM-tier experts each token routes to) is 44.5 s. So the overlap is working perfectly — it hides essentially all of the GPU time — but the end-to-end win is only +6% because the GPU was never the bottleneck at S=1. Decode on this class of host is CPU-bound (independently: 1.1% SM utilization). Full writeup in Discussion #208.

So this PR's real value is diagnostic: it's the instrument that proved the RAM tail, not the GPU, is the lever. The next moves it points to (streaming-weight GPU path for RAM experts, or a smaller container so the whole set fits VRAM — the #81 lattice line) are separate, larger PRs. Happy to keep this opt-in and default-off; the layernorm cache is worth taking on its own regardless.

[JustVugg]

check cuda and i merge it!

[ZacharyZcR]

Rebased onto current dev and fixed the Windows CUDA DLL bridge for the async expert-group API (issue/take).

Validation:

• local MinGW host build: make -C c glm CUDA_DLL=1 CC=x86_64-w64-mingw32-gcc
• upstream CI: all 8 checks pass, including CUDA build (Windows, MSVC host) and Windows make check

[JustVugg]

@ZacharyZcR — one gate left before this merges, and it's the same one #365/#372 passed: token-identity COLI_GROUP_ASYNC=0 vs =1 on real silicon, post-rebase. The build + CI validation you posted covers compile/link on every platform, but the async overlap changes when the CPU rows and the GPU groups land in the output accumulator — i.e. FP accumulation order — and that's exactly the class that flips near-tied tokens. (We re-learned this the hard way this week: an unpinned kernel-family switch flipped a tie at temp-0 and cost a full debugging cycle to trace. Order-of-accumulation changes are guilty until proven token-identical.)

The check is one A/B on your box: same prompt, TEMP=0, fixed seed/usage snapshot, GROUP_ASYNC=0 vs 1, diff the token streams (and the +6-8% number re-confirmed while you're there). If they're identical, this merges same-day — the opt-in gating means it's zero-risk for everyone else, and the overlap idea itself is exactly the right shape.

[ZacharyZcR]

Token-identity A/B: COLI_GROUP_ASYNC=0 vs =1 — IDENTICAL

Machine: rs-yuesheng-gpu (2× Xeon Silver 4510, 6× RTX 5090, 251 GiB RAM)
Binary: feat/expert-group-overlap tip (5978d05), make glm CUDA=1 CUDA_ARCH=native
Model: GLM-5.2 int4 (int8 MTP heads)
Config: full residency (CUDA_EXPERT_GB=auto CUDA_DENSE=1 COLI_CUDA_ATTN=1 PIN_GB=all RAM_GB=auto COLI_NUMA=1), TEMP=0 DRAFT=0, .coli_usage snapshot-restored between every arm.

Token identity

config	round 1	round 2
no pipe	IDENTICAL	IDENTICAL
COLI_CUDA_PIPE=2	IDENTICAL	—

All 64 generated token IDs match byte-for-byte across sync/async arms in every run. The [TOKENS] stderr dump was compared directly; text output differs only in wall-clock profiling numbers.

Token stream (round 1, both arms):

3555 374 198 840 20726 279 1376 11793 1948 15931 32 64529 323 15931 32 19854 19504 4938 23691 304 279 2266 315 1550 57084 24162 13 2823 3151 911 279 5068 24086 13 3555 374 279 3476 315 279 10345 1849 304 18122 15931 32 76618 30 2585 1558 279 10080 5978 429 11359 1598 29607 389 7299 12 9552 15931 32 5942

Performance (COLI_CUDA_PIPE=2, 128 tokens, NGEN=128)

arm	decode	expert-matmul	expert-wait	hit rate
sync	1.82 tok/s (70.16s)	12.11s	49.13s	96.4%
async	1.82 tok/s (70.46s)	12.03s	49.56s	96.4%

At 96% hit rate (9280 VRAM + 1574 RAM experts) the CPU expert tail dominates, so the overlap window is noise-level — consistent with the PR's own analysis (GPU side ~3.8s vs CPU 44.5s). The +6-8% from the original measurement was at a warmer cache state where the GPU fraction was proportionally larger.

Verdict

The gate is met: token-exact under controlled .coli_usage + TEMP=0 DRAFT=0, reproduced across two independent rounds and two pipeline configurations. The layernorm cache commit (39eabeb) is a pure win regardless.

[JustVugg]

Heads-up: c/glm.c was just split into c/colibri.c + header modules (#391, now merged into dev). This PR touches the old glm.c, so it will need a rebase onto current dev with its changes moved into colibri.c (or the relevant extracted header: quant.h, sample.h, kv_persist.h, telemetry.h, grammar.h). The make glm target still works (alias of make colibri), so no build scripts break. Apologies for the churn — ping me if the rebase gets thorny and I can help place the hunks.

[ZacharyZcR]

Rebased onto current dev — the async issue/take and layernorm-cache hunks now live in c/colibri.c (the moe/decode paths stayed in the main engine file, nothing landed in the extracted headers). Clean CPU build, 77/77 tests pass locally; CUDA paths compile-checked by CI.