From 6d0cc64cd833120ead5aef056b35a44cda303aa5 Mon Sep 17 00:00:00 2001 From: Philipp Guevorguian Date: Thu, 17 Sep 2026 14:31:51 -0700 Subject: [PATCH 1/3] Drop cooperative launch from combine prologue The combine prologue's CTAs are independent: each block owns the dedup groups `blockIdx.x + j * num_sms`, reads read-only duplicate rows, and writes only its own primary row. All synchronization is block-local (warp shuffles and the block-local PipelineTmaAsync mbarriers); the kernel never calls grid_sync or cross_rank_barrier. PDL does not add a whole-grid requirement either. griddepcontrol's launch_dependents lets dependents be scheduled once all other CTAs issue the same instruction *or have completed*, so producer CTAs may run in waves and exit. Cooperative launch only imposes an occupancy-bounded whole-grid residency constraint here, so drop it. CombineKernel keeps its cooperative launch: its cross_rank_barrier spins in grid_sync until every local CTA arrives. Not GPU-tested. --- moonep/combine_prologue.py | 1 - 1 file changed, 1 deletion(-) diff --git a/moonep/combine_prologue.py b/moonep/combine_prologue.py index dc25765..0decb55 100644 --- a/moonep/combine_prologue.py +++ b/moonep/combine_prologue.py @@ -189,7 +189,6 @@ def __call__( block=(self.num_threads, 1, 1), smem=smem_bytes, stream=stream, - cooperative=True, ) # ----------------------------------------------------------------- kernel From dd33a1c43180eeb74109e19f4adfe79aceb3eef1 Mon Sep 17 00:00:00 2001 From: phil Date: Fri, 18 Sep 2026 00:10:38 +0000 Subject: [PATCH 2/3] Drop cooperative launch from dispatch epilogue --- moonep/dispatch_epilogue.py | 1 - 1 file changed, 1 deletion(-) diff --git a/moonep/dispatch_epilogue.py b/moonep/dispatch_epilogue.py index 16a8495..8710cbe 100644 --- a/moonep/dispatch_epilogue.py +++ b/moonep/dispatch_epilogue.py @@ -159,7 +159,6 @@ def __call__( block=(self.num_threads, 1, 1), smem=smem_bytes, stream=stream, - cooperative=True, use_pdl=self.pdl_launch, ) From 587075790d1d7f38d51e1b9621b073249b465a2a Mon Sep 17 00:00:00 2001 From: phil Date: Fri, 18 Sep 2026 00:10:39 +0000 Subject: [PATCH 3/3] Add wave-progress regression for local kernels --- tests/test_local_kernel_waves.py | 166 +++++++++++++++++++++++++++++++ 1 file changed, 166 insertions(+) create mode 100644 tests/test_local_kernel_waves.py diff --git a/tests/test_local_kernel_waves.py b/tests/test_local_kernel_waves.py new file mode 100644 index 0000000..25316a7 --- /dev/null +++ b/tests/test_local_kernel_waves.py @@ -0,0 +1,166 @@ +"""Wave-progress regression for the CTA-independent local kernels. + +``dispatch_epilogue`` and ``combine_prologue`` own no grid-wide barrier: every +CTA walks its own round-robin groups and synchronizes only inside its block. +Neither kernel may therefore require whole-grid residency. This test pins that +property by launching both kernels with grids that exceed what the device can +hold at once, including a grid past the shared-memory residency bound, and +requiring exact results within a bounded deadline. + +A cooperative launch fails here: CUDA rejects a cooperative grid larger than the +occupancy bound, so the oversized cases error out instead of completing. + +Single GPU, no torchrun: + + pytest -s tests/test_local_kernel_waves.py + +Set ``MOONEP_WAVE_CASE`` to a printed case index to replay one case. The test +body runs in a spawned subprocess so a hung launch is killed by the deadline +rather than hanging the session. +""" + +from __future__ import annotations + +import os +import time +from dataclasses import dataclass +from typing import Literal + +import pytest +import torch +import torch.multiprocessing as mp + + +@dataclass(frozen=True) +class WaveCase: + kernel: Literal["dispatch_epilogue", "combine_prologue"] + H: int + blocks: int + pdl: bool + seed: int + groups: int + NvS: int + max_dups: int + + +def _make_plan(case: WaveCase): + """Disjoint duplicate groups plus a CPU oracle, with no undefined headers read.""" + from moonep.planning import MoonEPCommPlan + + rng = torch.Generator().manual_seed(case.seed) + sizes = torch.randint(1, case.max_dups + 1, (case.groups,), generator=rng) + n_dups = int(sizes.sum()) + slots = torch.randperm(case.NvS, generator=rng) + primary = slots[: case.groups] + dups = slots[case.groups : case.groups + n_dups] + owners = torch.repeat_interleave(torch.arange(case.groups), sizes) + + headers = torch.full((case.NvS, 3), -1, dtype=torch.int32) + headers[: case.groups, 0] = primary.to(torch.int32) + headers[: case.groups, 1] = (sizes.cumsum(0) - sizes).to(torch.int32) + headers[: case.groups, 2] = sizes.to(torch.int32) + loffs = torch.full((case.NvS,), -1, dtype=torch.int32) + loffs[:n_dups] = dups.to(torch.int32) + + plan = MoonEPCommPlan( + dst=torch.zeros(1, dtype=torch.int32, device="cuda"), + experts_to_copy=torch.zeros((1, 1), dtype=torch.int32, device="cuda"), + zero_fill_ranges=torch.zeros((2, 2), dtype=torch.int32, device="cuda"), + remote_stats=torch.zeros(2, dtype=torch.int32, device="cuda"), + N=1, R=1, E=1, B=1, NvS=case.NvS, K=case.max_dups + 1, + dup_groups=headers.cuda(), + dup_loffs=loffs.cuda(), + dup_counts=torch.tensor([case.groups, n_dups], dtype=torch.int32, device="cuda"), + ) + + # Small integers keep fp32 accumulation exact on both device and oracle. + rows = torch.arange(case.NvS, dtype=torch.int32)[:, None] + cols = torch.arange(case.H, dtype=torch.int32)[None, :] + source = ((rows * 7 + cols * 3 + case.seed) % 31 - 15).to(torch.bfloat16) + if case.kernel == "dispatch_epilogue": + source[dups] = -99 # A skipped copy cannot pass as success. + expected = source.clone() + if case.kernel == "dispatch_epilogue": + expected[dups] = source[primary[owners]] + else: + sums = source[primary].float() + sums.index_add_(0, owners, source[dups].float()) + expected[primary] = sums.to(torch.bfloat16) + return plan, source, expected + + +def _worker() -> None: + from moonep.combine_prologue import CombinePrologueKernel, launch_combine_prologue + from moonep.dispatch_epilogue import DispatchEpilogueKernel, launch_dispatch_epilogue + + torch.set_num_threads(1) + torch.cuda.set_device(0) + props = torch.cuda.get_device_properties(0) + sms = props.multi_processor_count + budget = props.shared_memory_per_block_optin - 1024 + + shapes = ((128, 2 * sms + 1), (1024, sms + 1), (2048, 4 * sms + 1), (4096, 2 * sms + 1)) + selected = os.environ.get("MOONEP_WAVE_CASE") + case_index = 0 + completed = 0 + + for kernel in ("dispatch_epilogue", "combine_prologue"): + cls = DispatchEpilogueKernel if kernel == "dispatch_epilogue" else CombinePrologueKernel + for H, blocks in shapes: + B, stages = cls._pick_geometry(H, budget) + smem = cls._smem_bytes(H, stages, B) + # Shared memory alone bounds co-residency, whatever the register usage. + residency_bound = sms * (props.shared_memory_per_multiprocessor // smem) + if blocks == 4 * sms + 1: + assert blocks > residency_bound, "wave case must exceed the residency bound" + max_groups = max(blocks * B + 3, 3 * stages * B + 1) + max_dups = 3 + NvS = max_groups * (max_dups + 1) + 17 + group_counts = [0, 1, B, max_groups // 2, max_groups] + for pdl in (False, True): + for trial, groups in enumerate(group_counts): + case = WaveCase(kernel, H, blocks, pdl, 34 + trial, groups, NvS, max_dups) + identity = str(case_index) + case_index += 1 + if selected is not None and identity != selected: + continue + print(f"WAVE_CASE {identity} residency_bound={residency_bound} {case}", flush=True) + plan, source, expected = _make_plan(case) + hidden = source.cuda() + ctx = {"H": H, "R": 1, "NvS": NvS, "num_sms_dedup": blocks, + "hidden_buf_local": hidden} + if kernel == "dispatch_epilogue": + launch_dispatch_epilogue(ctx, plan, pdl_launch=pdl) + else: + launch_combine_prologue(ctx, plan, pdl_trigger=pdl) + done = torch.cuda.current_stream().record_event() + deadline = time.monotonic() + 10 + while not done.query() and time.monotonic() < deadline: + time.sleep(0.01) + if not done.query(): + print(f"WAVE_STALL {identity} {case}", flush=True) + os._exit(1) + torch.testing.assert_close(hidden.cpu(), expected, rtol=0, atol=0) + completed += 1 + + assert completed > 0, "MOONEP_WAVE_CASE selected no case" + print(f"WAVE_PASS cases={completed}", flush=True) + + +def test_local_kernels_progress_in_waves() -> None: + """Both local kernels complete oversized grids and match the CPU oracle exactly.""" + if "RANK" in os.environ: + pytest.skip("this test owns its own subprocess") + if not torch.cuda.is_available(): + pytest.skip("requires a CUDA device") + pytest.importorskip("moonep") + proc = mp.get_context("spawn").Process(target=_worker) + proc.start() + try: + proc.join(timeout=600) + assert not proc.is_alive(), "wave worker exceeded 600 s" + assert proc.exitcode == 0, f"wave worker failed: exit={proc.exitcode}" + finally: + if proc.is_alive(): + proc.kill() + proc.join(timeout=5)