Skip to content

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

Open
ZacharyZcR wants to merge 2 commits into
JustVugg:devfrom
ZacharyZcR:feat/expert-group-overlap
Open

cuda: COLI_GROUP_ASYNC=1 — overlap the CPU expert rows with the GPU groups at decode (opt-in, +6-8%)#342
ZacharyZcR wants to merge 2 commits into
JustVugg:devfrom
ZacharyZcR:feat/expert-group-overlap

Conversation

@ZacharyZcR

Copy link
Copy Markdown
Contributor

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).

…2 sync H2D/token) + overlap-window profiling counters
@ZacharyZcR

Copy link
Copy Markdown
Contributor Author

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant