Drop cooperative launch from the CTA-independent local kernels

[philippguevorguian]

Drop cooperative launch from the CTA-independent local kernels

What changes

cooperative=True is removed from two launches, CombinePrologueKernel in moonep/combine_prologue.py and DispatchEpilogueKernel in moonep/dispatch_epilogue.py. A new single-GPU test, tests/test_local_kernel_waves.py, pins the property both kernels now rely on. Grid dimensions, shared-memory geometry, PDL behavior, and stream ordering are untouched. Transport and planning kernels that own real grid-wide barriers keep their cooperative launches.

Why it's worth doing

A cooperative launch tells CUDA the entire grid must be co-resident before any block starts. These two kernels have no use for that guarantee. Each CTA walks its own round-robin slice of dedup groups, writes only rows it owns, and synchronizes only inside its block through PipelineTmaAsync mbarriers and warp shuffles. Neither file contains grid_sync or cross_rank_barrier. Cross-rank publication is handled by CombineKernel on entry and by dispatch on exit. PDL imposes no whole-grid requirement either, because griddepcontrol.launch_dependents is issued per CTA and griddepcontrol.wait is satisfied by the predecessor grid rather than by peer CTAs.

The cost of the flag is real. Both kernels launch with grid = num_sms, which defaults to all 132 SMs on H100, and at H=7168 each CTA reserves about 224 KiB of shared memory, so only one block fits per SM. A cooperative grid then cannot begin until the GPU fully drains. That has two consequences: the kernel queues behind co-resident work instead of filling in around it, and any grid larger than the occupancy bound is rejected outright by the driver.

Measured impact (8xH100, S=4096, K=8, H=7168, 100 launches timed against a saturating GEMM stream)

kernel	master isolated	master contended	this PR isolated	this PR contended
combine_prologue	22 to 24 us	1443 to 1503 us	22 to 24 us	58 to 68 us
dispatch_epilogue	22 to 24 us	1397 to 1492 us	23 to 25 us	59 to 68 us

Standalone latency is unchanged. Under contention each kernel is roughly 24x faster, which is the regime a real training step runs in.

New test

tests/test_local_kernel_waves.py builds duplicate groups and a CPU oracle directly, so it needs one GPU and no torchrun. It sweeps both kernels across 4 hidden sizes, both PDL settings, and 5 group counts covering empty work, a partial batch, and pipeline reuse, for 80 cases. Grids are deliberately oversubscribed, and the largest case asserts that it exceeds the shared-memory residency bound before launching. Dispatch destinations are poisoned with a sentinel so a skipped copy cannot pass as success, and results are compared bit-exactly. Each launch carries a 10 second event deadline inside a spawned worker, so a hung launch is killed instead of hanging the session.

The test discriminates. Against the current master launches it fails at the first case with CUDA_ERROR_COOPERATIVE_LAUNCH_TOO_LARGE (error 720). On this branch all 80 cases pass in 36 s.

Tests and validations

• Environment: two 8xH100 80GB nodes, CUDA 12.9, torch 2.11+cu128, nvidia-cutlass-dsl==4.4.2, and cuda-python==12.9 (the default 13.x bindings fail against the 570.x driver with cudaErrorInsufficientDriver).
• A/B performed on one machine with one build, toggling only these launch flags: master @2bd860b against this branch.
• Full existing suite with both changes applied, torchrun --nproc_per_node=8 -m pytest per file: test_combine 14 passed, test_e2e 1 passed, test_dispatch 12 passed, test_planning 18 passed with 1 skipped, test_grad_reduce 12 passed, test_prefetch 14 passed. The pre-change baseline on the same machine gave identical counts for test_combine, test_e2e, and test_dispatch.
• New wave test: 80 cases passed on this branch, first case failed with error 720 on master.
• Independent reproduction on a second node with this branch: test_combine 14 passed, test_e2e 1 passed.
• Contention microbenchmark repeated across both nodes, including a run with the kernel order reversed. The first ordering under-reported the second kernel measured, because the GEMM stream had partly drained by then, so each kernel's headline number above comes from a run where it was measured first.
• Source audit backing the correctness argument: grep over both files finds no grid_sync, no cross_rank_barrier, and no cluster-wide synchronization. What remains is block-local mbarrier and shuffle traffic plus the per-CTA PDL calls.

Not covered: multi-node runs, and PDL chains where the predecessor grid has not completed.

[asp0ex]

this is solved in the latest release. see 33327eb