cuda: COLI_GROUP_ASYNC=1 — overlap the CPU expert rows with the GPU groups at decode (opt-in, +6-8%)#342
Conversation
… overlap at decode (opt-in, +6-8% measured)
…2 sync H2D/token) + overlap-window profiling counters
|
Pushed two follow-ups from the profiling I promised: 1. Layernorm weights cached on-device ( 2. The overlap window measured — and it reframes this whole PR. With 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. |
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_ensurereloads slabs released byCUDA_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_usagesnapshot-restored between arms)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_PINsemantics take precedence.Follow-ups I intend regardless of merge: (a) a position-aligned
ROUTE_TRACEdiff (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).