From f309b2787c502f42bdd94fc066cdaf1897e1fac1 Mon Sep 17 00:00:00 2001
From: derek <derek.yates@live.com>
Date: Wed, 22 Jul 2026 21:02:10 -0400
Subject: [PATCH] fix(moe): cooperatively launch fused W4A16 grids

---
 .../moe/_shared/kernels/w4a16/kernel.py       |  9 +++
 tests/moe/test_fused_moe.py                   | 66 +++++++++++++++++++
 2 files changed, 75 insertions(+)

diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
index 4b151d6a..a2bfc11e 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
@@ -4460,6 +4460,11 @@ def __call__(
             grid=grid,
             block=[self.cta_threads, 1, 1],
             min_blocks_per_mp=self.blocks_per_sm,
+            # The fused body crosses software all-CTA barriers between FC1,
+            # activation, and FC2. Require whole-grid admission so unrelated
+            # work cannot occupy an SM while resident CTAs wait for peers that
+            # have not been scheduled yet.
+            cooperative=True,
             stream=stream,
         )
 
@@ -5211,6 +5216,10 @@ def __call__(
             grid=(grid_x, 1, 1),
             block=[self.cta_threads, 1, 1],
             min_blocks_per_mp=self.blocks_per_sm,
+            # This fused two-tier body uses the same software all-CTA barriers
+            # as W4A16FusedMoeKernel and therefore has the same whole-grid
+            # residency requirement.
+            cooperative=True,
             stream=stream,
         )
 
diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py
index 6cd97924..1ebec2bc 100644
--- a/tests/moe/test_fused_moe.py
+++ b/tests/moe/test_fused_moe.py
@@ -130,3 +130,69 @@ def test_run_w4a16_replays_under_cuda_graph_with_frozen_resolution() -> None:
         "no kernel may compile during or after warm capture"
     )
     torch.testing.assert_close(captured, eager, rtol=0, atol=0)
+
+
+def test_run_w4a16_m9_graph_replay_with_prequeued_aux_work() -> None:
+    """The route-packed W4A16 grid must remain valid beside aux work.
+
+    M=9 is the first size above the fused-micro decode range. It enters the
+    route-packed W4A16 body, whose FC1/activation/FC2 phases synchronize every
+    CTA through software grid barriers. Prequeued shared-expert work must not
+    prevent whole-grid admission when the serving graph replays.
+    """
+    require_sparkinfer()
+    torch.manual_seed(20260723)
+
+    from sparkinfer.moe import fused_moe
+
+    device = torch.device("cuda")
+    global_e, hidden_size, intermediate_size = 16, 6144, 512
+    m, topk = 9, 8
+    a = (torch.randn(m, hidden_size, device=device) * 0.25).to(torch.bfloat16)
+    weights = make_modelopt_weights(
+        experts=global_e,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+    )
+    experts = prepare_experts(a, weights, torch.arange(global_e))
+    topk_ids = torch.randint(
+        0, global_e, (m, topk), dtype=torch.int32, device=device
+    )
+    topk_weights = torch.softmax(torch.randn(m, topk, device=device), dim=-1)
+    output = torch.empty_like(a)
+    binding = make_tp_moe_fp4_binding(
+        a=a,
+        experts=experts,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        output=output,
+        quant_mode="w4a16",
+    )
+
+    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)
+    torch.cuda.current_stream().wait_stream(capture_stream)
+    torch.cuda.synchronize()
+
+    aux_stream = torch.cuda.Stream()
+    aux_a = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device)
+    aux_b = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device)
+    aux_out = torch.empty_like(aux_a)
+    output.zero_()
+    with torch.cuda.stream(aux_stream):
+        for _ in range(16):
+            torch.mm(aux_a, aux_b, out=aux_out)
+    graph.replay()
+    torch.cuda.current_stream().wait_stream(aux_stream)
+    torch.cuda.synchronize()
+
+    assert output.isfinite().all()
+    assert output.abs().sum().item() > 0
+    torch.testing.assert_close(output, expected, atol=2e-3, rtol=0.0)