fix(moe): cooperatively launch fused W4A16 grids#73
Conversation
📝 WalkthroughWalkthroughW4A16 fused and hybrid MoE kernel launches now request cooperative CUDA execution. A new test validates CUDA graph replay at ChangesW4A16 cooperative execution
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Test as W4A16 graph replay test
participant Graph as CUDA graph
participant MoE as fused_moe.run
participant Aux as auxiliary CUDA stream
Test->>MoE: Compute eager expected output
Test->>Graph: Capture fused_moe.run
Test->>Aux: Queue auxiliary GEMMs
Test->>Graph: Replay captured graph
Graph->>MoE: Execute cooperative W4A16 launch
Test->>Graph: Synchronize and compare output
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
Runtime qualification update: the exact cooperative-launch source was byte-verified in the final v20 candidate. With shared-expert and MTP activity enabled, that process completed profiling and production graph capture for sizes 16 through 1, passed 624/624 diagnostic boundaries, then completed cold prefill plus ten decode cells through concurrency 16 with zero errors and RestartCount 0. Decode stayed within 5 percent of the v19 production baseline at the measured comparable cells. This is combined-stack rather than isolated A/B evidence, but it exercises the production concurrent-stream setting and closes the pending throughput/stability gate. The patch remains narrowly scoped to supplying the cooperative-admission contract required by the existing whole-grid barriers. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/moe/test_fused_moe.py (1)
172-182: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueWarmup runs on the default stream instead of the capture stream.
PyTorch's documented CUDA-graph pattern requires warmup to run on the same side stream used for capture (to stabilize allocator/algorithm-selection state before recording), but the eager
fused_moe.run(binding=binding)at line 172 runs on the default stream, while capture happens on a separatecapture_stream. Since this test is inference-only (no autograd), this is unlikely to break capture, but it doesn't mirror the officially recommended pattern nor necessarily the production warmup path this test is meant to validate.As per coding guidelines, "Treat CUDA graph capture/replay, warmup behavior, stable allocation, and fixed or preplanned workspace capacity as serving requirements."
♻️ Align warmup with the capture stream
- fused_moe.run(binding=binding) - torch.cuda.synchronize() - expected = output.clone() - - graph = torch.cuda.CUDAGraph() - capture_stream = torch.cuda.Stream() - capture_stream.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(capture_stream), torch.cuda.graph(graph): - fused_moe.run(binding=binding) + graph = torch.cuda.CUDAGraph() + capture_stream = torch.cuda.Stream() + capture_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(capture_stream): + fused_moe.run(binding=binding) + torch.cuda.current_stream().wait_stream(capture_stream) + torch.cuda.synchronize() + expected = output.clone() + + with torch.cuda.stream(capture_stream), torch.cuda.graph(graph): + fused_moe.run(binding=binding)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/moe/test_fused_moe.py` around lines 172 - 182, Move the warmup fused_moe.run call in the test to the capture_stream, ensuring the stream is initialized and synchronized appropriately before capture. Keep expected output generation and the existing capture/replay synchronization behavior unchanged, while making warmup use the same side stream as torch.cuda.graph.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/moe/test_fused_moe.py`:
- Around line 172-182: Move the warmup fused_moe.run call in the test to the
capture_stream, ensuring the stream is initialized and synchronized
appropriately before capture. Keep expected output generation and the existing
capture/replay synchronization behavior unchanged, while making warmup use the
same side stream as torch.cuda.graph.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 5e19bc02-ab58-4ba4-8672-fc8781db5e98
📒 Files selected for processing (2)
sparkinfer/moe/_shared/kernels/w4a16/kernel.pytests/moe/test_fused_moe.py
2 participants
Summary
Launch the two barrier-bearing fused W4A16 MoE grids cooperatively.
Both
W4A16FusedMoeKernelandW4A16FusedMoeHybridKernelsynchronize everyCTA between FC1, activation, output initialization and FC2. Their planners
already cap the grid to
SM count * blocks_per_sm, but the launch itself didnot request whole-grid cooperative admission.
This patch adds
cooperative=Trueto those two launch sites. It does notchange the standalone W4A16 GEMM, launch geometry, tile selection,
quantization, routing or output math.
Why
A software whole-grid barrier requires every participating CTA to be
resident. A bounded grid alone does not guarantee that: unrelated work on
another stream can occupy an SM after only part of the fused grid has been
admitted, leaving resident CTAs waiting for peers that cannot be scheduled.
SparkInfer already uses cooperative admission for the fused-micro W4A16 decode
path. The route-packed fallback entered above that micro range and the
two-tier hybrid path use the same barrier pattern but were missing the launch
contract.
The production workload that motivated the audit runs shared-expert and other
auxiliary-stream work beside MTP decode graph capture. This PR is deliberately
not presented as the fix for that workload's eventual CUDA illegal-address
failure: launch-blocking subsequently localized that separate failure to an
MLA query-absorption BMM, which has its own vLLM fix. Cooperative admission is
an independent correctness requirement visible directly in these fused
kernels.
Validation
Completed:
master;python -m py_compilefor both changed files;git diff --check;that completed profiling and production decode capture at sizes 16 through
1, with 624/624 diagnostic boundaries passing and no CUDA/cuBLAS/OOM/Xid
failure.
The engine result is combined-stack evidence, not an isolated A/B proof of
this change.
This PR adds a focused SM120 GPU regression:
It captures the first route-packed size above the fused-micro range using the
GLM shard geometry, prequeues large BF16 GEMMs on an auxiliary stream, replays
the W4A16 graph concurrently, and checks finite, nonzero, stable output.
The test cannot run on the contributor's macOS host because CUTLASS DSL and
SM120 CUDA are unavailable. This PR remains draft until the focused GPU test
and the ongoing decode-throughput qualification are recorded.
Tradeoff
Cooperative admission can alter scheduling latency when unrelated stream work
is queued. It does not serialize the streams globally or disable
shared-expert overlap; it only requires the complete bounded W4A16 grid to be
admitted as a unit, which is the execution contract required by its software
barriers.
AI assistance
OpenAI Codex assisted with source tracing, forward-porting, regression-test
design and PR drafting. The submitted code and runtime evidence are being
reviewed by the human contributor and target-system operator before the draft
is marked ready.
Summary by CodeRabbit
Bug Fixes
Tests