feat(grouped_gemm): add CK work-stealing variant with schedule API by wenchenvincent · Pull Request #348 · AMD-AGI/Primus-Turbo

wenchenvincent · 2026-05-27T18:16:11Z

Summary

Adds a work-stealing (WS) variant of the CK grouped-GEMM kernel (forward +
variable-K backward) for AMD MI355X / MI350. Under all-gather × grouped-GEMM
overlap the static-stride persistent kernel suffers a ~1.7–2.0× slowdown
when RCCL channels contend for the same CUs; with WS enabled, fast CUs
absorb work that the RCCL-throttled CUs would have done, recovering most of
the slowdown at ~2-3% isolated overhead on the 288-shape MoE sweep.

Rebased on top of current main (which now includes the Triton WS PR
#353, so the public API is schedule: str = "static" | "work_steal").
This PR wires CK into that same API, so users get WS on either backend by
setting schedule="work_steal" on primus_turbo.ops.grouped_gemm(...).

Fixes # (issue)

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Kernel (csrc/kernels/grouped_gemm/ck/grouped_gemm_kernel_ws.hpp):

Vendored work-stealing variant of ck_tile::GroupedGemmKernel. Static-
stride persistent loop replaced by an atomicAdd-based tile claim.
Two-phase scheme: phase 1 each CTA drains its XCD's local counter
for up to local_per_xcd claims (preserves L2 locality across XCDs);
phase 2 the remainder is claimed via a single global counter
(cross-XCD work-stealing).
Counter layout [xcd0..xcd7, global, done] — 10 int32 slots.
Last-out CTA self-resets the counter so the next launch sees a clean
slate; no host-side counter.zero_() needed.
FindGroupId binary search reads precomputed block_start/block_end
populated by the existing compute_grouped_gemm_args kernel.
Defensive bound — active_xcds = min(gridDim.x, NUM_XCDS_WS) used
for both xcd_id and phase1_total. Without this cap, launches where
gridDim.x < NUM_XCDS_WS (=8) would leave the tiles in the gap between
phase 1's actual end and phase 2's start unclaimed → silent wrong
output. Matches the equivalent fix on the Triton side (PR fix(grouped_gemm): guard work_steal + num_cu; harden WS per-XCD layout #402).

Binding + runner (csrc/pytorch/grouped_gemm/ck_grouped_gemm.cpp,
csrc/kernels/grouped_gemm/ck/ck_grouped_gemm_kernel_template.h):

ck_grouped_gemm and ck_grouped_gemm_variable_k extended with
(work_steal, ws_counter, ws_local_per_xcd) args.
CKGroupedGemmRunner exposes run_ws() alongside run().
PRIMUS_TURBO_CHECK on ws_local_per_xcd bounds [0, INT32_MAX] to
reject values that would silently truncate on the int64 → int32 cast
passed to the kernel.

Backend + heuristic
(primus_turbo/pytorch/kernels/grouped_gemm/grouped_gemm_impl.py,
primus_turbo/pytorch/kernels/grouped_gemm/ws_ck_heuristic.py):

GroupedGEMMCKBackend.can_handle now advertises
schedule in ("static", "work_steal"); execute translates
schedule="work_steal" to the internal work_steal=True and picks a
ws_local_per_xcd via the sync-free heuristic.
Heuristic (ws_ck_heuristic.py) resolves the ws_mode string to the
kernel's integer ws_local_per_xcd, using tpc = total_tiles / num_cus
as the density signal. Uses tensor metadata only (a.size(0),
group_lens.numel()); no GPU sync.
total_m computed as a.size(1) if trans_a else a.size(0) for shape
correctness under transposed inputs (can_handle already rejects
trans_a=True on this backend, but the shape logic is now explicit).
Wrapper-managed per-device counter cache (_get_ck_ws_counter); the
kernel self-resets so it's a one-shot torch.zeros per device.

Public API (primus_turbo/pytorch/ops/grouped_gemm.py):

No new kwargs. schedule="work_steal" on the existing grouped_gemm(...)
op now dispatches to the CK WS path (in addition to the Triton path
from PR feat: triton_grouped_gemm: add work-stealing variant with ws_mode API #353).
Guard: schedule="work_steal" combined with an explicit num_cu != None
raises ValueError. WS was designed and characterized for full-device
launches; partial-grid launches would silently under-utilize the WS
layout even with the ACTIVE_XCDS fix.

Tests:

tests/pytorch/ops/test_grouped_gemm.py:
- test_grouped_gemm_schedule_work_steal now parametrized over
  backend ∈ {TRITON, CK} — 64 cases per backend asserting forward +
  backward bit-equal vs the static reference.
- test_grouped_gemm_schedule_work_steal_single_group — G=1 smoke on
  both backends (dispatcher special-cases to gemm_impl).
- test_grouped_gemm_schedule_work_steal_rejects_hipblaslt — explicit
  HIPBLASLT + schedule="work_steal" fails at dispatch (only Triton /
  CK advertise WS via can_handle).
- test_ck_grouped_gemm_op_ws_num_cu_below_num_xcds — kernel-level
  regression for the ACTIVE_XCDS fix (num_cu ∈ {4, 6, 7, 8, 16}).
tests/pytorch/ops/test_ws_ck_heuristic.py (NEW): 15 pure-Python unit
tests pinning the CK heuristic's mode-resolution behavior (explicit
modes, auto thresholds for both kernel kinds, degenerate inputs,
ValueError path, total_tiles helpers, and the tile / chiplet
constants).

Perf (MI355X, 288-shape MoE sweep, isolated kernel, clean methodology)

CK WS isolated overhead vs the static-stride reference (schedule="static"):

Metric	Value
Median	+2.7%
p10	+0.0%
p90	+6.3%
Max	+14.5%

Under overlap on the canonical bench shape (G=32, M=267424, K=1280, N=2560,
NCCL_MAX_NCHANNELS=default), CK WS recovers >30% of the contended-regime
slowdown vs static.

Checklist

The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

Copilot

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds an opt-in CK work-stealing path for grouped GEMM (forward + backward) with a heuristic-driven mode selector exposed via a new work_steal/ws_mode API.

Changes:

Introduces a vendored GroupedGemmKernelWS CK kernel with per-XCD/global/hierarchical atomic tile claiming and self-resetting counter.
Extends C++ bindings, params, and Python impl/dispatcher to thread work_steal, ws_counter, and ws_local_per_xcd through, plus a per-device cached counter and a ws_ck_heuristic module for mode resolution.
Adds backward propagation of WS settings in the autograd Function and parametrized tests for the four ws_mode values.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 9 comments.

Show a summary per file

File	Description
tests/pytorch/ops/test_grouped_gemm.py	New parametrized work-stealing CK tests (multi-group + single-group).
primus_turbo/pytorch/ops/grouped_gemm.py	Plumbs `work_steal`/`ws_mode` through `GroupedGemmFunc` and public API.
primus_turbo/pytorch/kernels/grouped_gemm/ws_ck_heuristic.py	New module mapping `ws_mode` string to `ws_local_per_xcd` int.
primus_turbo/pytorch/kernels/grouped_gemm/grouped_gemm_impl.py	Per-device WS counter cache; dispatches WS args to CK; ignores in other backends.
csrc/pytorch/grouped_gemm/grouped_gemm_meta.cpp	Meta-function signatures updated to accept WS args.
csrc/pytorch/grouped_gemm/ck_grouped_gemm.cpp	Validates WS counter and forwards params to the kernel.
csrc/pytorch/extensions.h	Header signatures updated.
csrc/pytorch/bindings_pytorch.cpp	TORCH_LIBRARY schemas extended with WS args.
csrc/kernels/grouped_gemm/ck_grouped_gemm.cu	Computes block_start/end during args setup; dispatches `run_ws`.
csrc/kernels/grouped_gemm/ck/grouped_gemm_kernel_ws.hpp	New WS kernel with phase-1/phase-2 atomic claims and last-CTA self-reset.
csrc/kernels/grouped_gemm/ck/ck_grouped_gemm_kernel_template.h	Adds `run_ws` and `m_tile/n_tile` to the runner interface.
csrc/include/primus_turbo/grouped_gemm.h	Adds WS fields to `GroupedGemmParams`.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+    device = "cuda"
+    N, K = N_K
+    group_lens = generate_grouped_gemm_group_lens(B, M, balance=balance).to(device)
+    print(B, M, N, K, dtype, balance, trans_b, ws_mode)


+# Per-device cache for the CK work-stealing counter buffer.
+#
+# Layout: [xcd0..xcd7, global, done] -- 10 int32 slots. The kernel self-resets
+# every slot to 0 before exit (last-out CTA pattern, see
+# grouped_gemm_kernel_ws.hpp), so the only zero we ever need is the one
+# `torch.zeros` does at first allocation. Caching one buffer per device makes
+# that allocation a one-shot cost per process.
+_CK_WS_COUNTER_NUM_SLOTS = 10
+_ck_ws_counters: dict[torch.device, torch.Tensor] = {}
+
+
+def _get_ck_ws_counter(device: torch.device) -> torch.Tensor:
+    buf = _ck_ws_counters.get(device)
+    if buf is None:
+        buf = torch.zeros(_CK_WS_COUNTER_NUM_SLOTS, dtype=torch.int32, device=device)
+        _ck_ws_counters[device] = buf


+    for (int64_t g = 0; g < group_id; ++g) {
+        const ck_tile::index_t M_g       = static_cast<ck_tile::index_t>(group_lens_ptr[g]);
+        const ck_tile::index_t m_tiles_g = (M_g + m_tile - 1) / m_tile;
+        my_block_start += m_tiles_g * num_n_tiles * k_batch;
+    }


+        if(threadIdx.x == 0)
+        {
+            int32_t* const done_counter = tile_counter_ptr + NUM_XCDS_WS + 1;
+            const int32_t prev = atomicAdd(done_counter, 1);
+            if(prev == static_cast<int32_t>(gridDim.x - 1))
+            {
+                // I am the last CTA out. Zero every slot (work + done) so the
+                // next launch starts clean.
+                for(index_t i = 0; i < NUM_XCDS_WS + 2; ++i)
+                {
+                    tile_counter_ptr[i] = 0;
+                }
+            }
+        }


+// Counter layout (caller-allocated int32 buffer of size NUM_XCDS_WS + 2):
+//   [0..NUM_XCDS_WS-1] : per-XCD slots
+//   [NUM_XCDS_WS]      : global slot
+//   [NUM_XCDS_WS + 1]  : done slot (last-out CTA detection for self-reset)


wenchenvincent · 2026-06-02T04:16:38Z

        if len(group_lens) == 1:
            assert b.size(0) == 1, f"Expected first dimension to be 1, got {b.size(0)}"


The backward path also has the same if len(group_lens) == 1 short-circuit (lines 69-89 of primus_turbo/pytorch/ops/grouped_gemm.py) that dispatches to gemm_impl for both grad_a and grad_b -- work_steal and ws_mode are ignored in the backward when G=1, identical to the forward. So there's no asymmetry: both forward and backward bypass the WS path for the single-group case. The single-group pytest case (test_grouped_gemm_work_steal_ck_single_group) only validates the forward output as you noted, but the underlying dispatch logic is symmetric on both ends.

+                        const ck_tile::index_t /*local_per_xcd*/) {
+        // Should be overridden by WS runners. If we land here, the dispatch
+        // table picked the wrong runner.
+        std::abort();


+        # ~50% per-xcd, ~50% global tail.
+        return max(1, (total_tiles // 2) // CK_NUM_XCDS)


+            const index_t local_idx = s_block_id;
+            if(local_idx >= local_per_xcd)
+            {
+                break;
+            }


RuibinCheung · 2026-05-28T06:50:18Z

Hi @wenchenvincent , Thanks for your contribution. I noticed that code-lint CI failed. You can follow these steps to reformat your code to pass the code-lint CI.

# make sure you are in the Primus-Turbo root directory
pip install pre-commit
pre-commit install
git commit -m "message"

Addressing Copilot review comments on PR AMD-AGI#348: - tests/pytorch/ops/test_grouped_gemm.py: remove leftover debug ``print`` in the parametrized WS test (would spam stdout across 128 combinations). - ck_grouped_gemm_kernel_template.h: replace ``std::abort()`` in the default ``run_ws`` impl with PRIMUS_TURBO_CHECK(false, ...) so a dispatch mismatch produces an actionable error message instead of a silent SIGABRT. - grouped_gemm.h / ck_grouped_gemm_kernel_template.h: update stale counter- layout comments. Old comments said "NUM_XCDS_WS + 1 (9 ints)" and "caller zeroes the buffer before each launch"; the actual layout is NUM_XCDS_WS + 2 (10 ints on MI355X with the done slot), and the kernel self-resets so the binding no longer zeroes the buffer per launch. - primus_turbo/pytorch/ops/grouped_gemm.py: docstring now explicitly calls out the single-stream WS constraint -- the per-device singleton counter is corrupted under concurrent multi-stream WS launches on the same device. Safe for the typical single-stream autograd graph; multi-stream callers need to pass an explicit ``ws_counter`` via the lower-level impl API. - ws_ck_heuristic.py: document the edge case where total_tiles < NUM_XCDS*2 causes hierarchical to degenerate to per-xcd. Bounds-checked in the kernel. - grouped_gemm_kernel_ws.hpp: expand the self-reset block comment to explain why no __threadfence() is needed -- HIP atomic ordering provides the intra-launch barrier; stream ordering provides the kernel-to-kernel visibility guarantee for the next launch.

Copilot

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 6 comments.

+            n = b.size(1) if trans_b else b.size(2)
+            total_tiles = approximate_ck_standard_total_tiles(
+                a.size(0),
+                group_lens.numel(),
+                n,
+            )


+        params.work_steal       = work_steal;
+        params.ws_counter_ptr   = ws_counter_ptr;
+        params.ws_local_per_xcd = static_cast<int32_t>(ws_local_per_xcd);


+    const ck_tile::index_t num_n_tiles = (n + n_tile - 1) / n_tile;
+    ck_tile::index_t       my_block_start = 0;
+    for (int64_t g = 0; g < group_id; ++g) {
+        const ck_tile::index_t M_g       = static_cast<ck_tile::index_t>(group_lens_ptr[g]);
+        const ck_tile::index_t m_tiles_g = (M_g + m_tile - 1) / m_tile;
+        my_block_start += m_tiles_g * num_n_tiles * k_batch;
+    }
+    const ck_tile::index_t my_M       = static_cast<ck_tile::index_t>(group_lens_ptr[group_id]);
+    const ck_tile::index_t my_m_tiles = (my_M + m_tile - 1) / m_tile;
+    args_ptr[group_id].block_start    = my_block_start;
+    args_ptr[group_id].block_end      = my_block_start + my_m_tiles * num_n_tiles * k_batch;


+@pytest.mark.parametrize("B", [4, 8])
+@pytest.mark.parametrize("M", [1024, 4096])
+@pytest.mark.parametrize("N_K", [(2048, 1536), (4096, 7168)])
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
+@pytest.mark.parametrize("balance", [True, False])
+@pytest.mark.parametrize("trans_b", [True, False])
+@pytest.mark.parametrize("ws_mode", ["auto", "global", "per-xcd", "hierarchical"])


+
+from __future__ import annotations
+
+CK_NUM_XCDS = 8


+def resolve_ck_ws_local_per_xcd(
+    ws_mode: str,
+    total_tiles: int,
+    num_cus: int,
+    *,
+    kernel_kind: str = _KERNEL_STANDARD,
+) -> int:


…nt32 bounds, heuristic tests) Addressing Copilot review comments on PR AMD-AGI#348: - ``GroupedGEMMCKBackend.execute``: use ``a.size(1) if trans_a else a.size(0)`` when computing ``total_m`` for the heuristic. ``can_handle`` already rejects ``trans_a=True`` on this backend, so this is unreachable in practice, but keeps the shape logic honest if the constraint is ever relaxed. (Comment AMD-AGI#11.) - ``csrc/pytorch/grouped_gemm/ck_grouped_gemm.cpp``: reject ``ws_local_per_xcd`` values that would silently truncate on the ``int64_t -> int32_t`` cast passed to the CK kernel (``PRIMUS_TURBO_CHECK`` on ``[0, INT32_MAX]``). Applied to both the standard and variable-K bindings. (Comment AMD-AGI#12.) - ``tests/pytorch/ops/test_ws_ck_heuristic.py`` (NEW): pure-Python unit tests for the CK WS mode-resolution heuristic. Mirrors the Triton heuristic suite -- covers explicit modes, ``auto`` behavior for both standard and variable-K kernels around the thresholds, degenerate inputs (``total_tiles <= 0``, ``num_cus == 0``), the ``ValueError`` path for unknown modes, the ``total_tiles`` helpers, and pins the tile-shape / chiplet-count constants that must stay in sync with the C++ side. 15 tests, GPU-free. (Comments AMD-AGI#15, AMD-AGI#16.) Deferred (documented / noted, not fixed in this commit): - AMD-AGI#2 / stream-safety of the per-device singleton counter: doc'd on the cache with a clear "not stream-safe" caveat; the public API also rejects ``schedule="work_steal"`` + non-``None`` ``num_cu``. - AMD-AGI#3 / AMD-AGI#13: O(G) prefix-scan in the args-setup kernel is future work (irrelevant for current MoE shape ranges). - AMD-AGI#14: test parametrization size is workable now that the ``ws_mode`` parametrization is folded away by the ``schedule`` API rename. - AMD-AGI#15: CK_NUM_XCDS duplication with C++ NUM_XCDS_WS: single-source-of-truth would need a Python binding for a fixed hardware constant; kept parallel with a docstring note pointing at the C++ side.

Copilot

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 5 comments.

+_CK_WS_COUNTER_NUM_SLOTS = 10
+_ck_ws_counters: dict[torch.device, torch.Tensor] = {}
+
+
+def _get_ck_ws_counter(device: torch.device) -> torch.Tensor:
+    buf = _ck_ws_counters.get(device)
+    if buf is None:
+        buf = torch.zeros(_CK_WS_COUNTER_NUM_SLOTS, dtype=torch.int32, device=device)
+        _ck_ws_counters[device] = buf
+    return buf


            * ``"work_steal"``: work-stealing persistent kernel with a kernel-
              aware heuristic that picks per-XCD / global / hierarchical tile
-              claims based on tensor metadata. Triton backend only; explicit
-              selection of a non-Triton backend together with ``"work_steal"``
-              is rejected at dispatch. Internal WS sub-modes are not exposed
-              at this layer; tune via ``grouped_gemm_triton_kernel`` directly
-              when needed.
+              claims based on tensor metadata. Supported on the Triton and CK
+              backends; HIPBLASLT advertises only ``"static"``. Internal WS
+              sub-modes are not exposed at this layer; tune via the kernel-


+    // Prefix sum over prior groups' tile counts -> block_start. Required by the
+    // WS kernel's FindGroupId binary search (and benign for the static path,
+    // which doesn't read these fields). m_tile/n_tile are the runner's tile
+    // shape, supplied by the dispatch site. Each thread does an O(group_id)
+    // scan over `group_lens_ptr` (already populated on device); for G <= ~256


+        at::Tensor &counter = ws_counter.value();
+        PRIMUS_TURBO_CHECK(counter.scalar_type() == at::kInt,
+                           "ws_counter must be int32");
+        PRIMUS_TURBO_CHECK(counter.is_cuda(), "ws_counter must be on CUDA");
+        constexpr int64_t kCounterSlots = 8 + 2; // NUM_XCDS_WS + global + done
+        PRIMUS_TURBO_CHECK(counter.numel() >= kCounterSlots,
+                           "ws_counter must have at least ", kCounterSlots, " int32 slots");
+        ws_counter_ptr = static_cast<int32_t *>(counter.data_ptr());


+        at::Tensor &counter = ws_counter.value();
+        PRIMUS_TURBO_CHECK(counter.scalar_type() == at::kInt, "ws_counter must be int32");
+        PRIMUS_TURBO_CHECK(counter.is_cuda(), "ws_counter must be on CUDA");
+        constexpr int64_t kCounterSlots = 8 + 2;
+        PRIMUS_TURBO_CHECK(counter.numel() >= kCounterSlots,
+                           "ws_counter must have at least ", kCounterSlots, " int32 slots");
+        ws_counter_ptr = static_cast<int32_t *>(counter.data_ptr());


When the persistent grid is capped to fewer CUs than there are XCDs (gridDim.x < NUM_XCDS_WS=8), the WS kernel silently dropped tiles: ``xcd_id = blockIdx.x % NUM_XCDS_WS`` only populated XCD ids [0, gridDim.x), so atomic claims hit only those per-XCD slots, but ``phase1_total = local_per_xcd * NUM_XCDS_WS`` reserved the full ``[0, 8 * local_per_xcd)`` range for phase 1. Phase 2 then started at ``phase1_total`` -- past where phase 1 actually ended -- leaving the tiles in ``[gridDim.x * local_per_xcd, NUM_XCDS_WS * local_per_xcd)`` unclaimed. Same bug as PR AMD-AGI#402 fixed on the Triton WS kernel. The bug stayed hidden in practice because the ``auto`` heuristic returns ``local_per_xcd=0`` (global-only) whenever the shape is dense enough, and ``num_cu < 8`` always pushes the shape into a high tiles-per-CU regime. Explicit ``ws_mode="per-xcd"`` or ``"hierarchical"`` with small ``num_cu`` did trip it. Fix: introduce ``active_xcds = min(gridDim.x, NUM_XCDS_WS)`` and use it for both ``xcd_id`` and ``phase1_total``. When ``gridDim.x >= NUM_XCDS_WS`` (the common case), this evaluates to ``NUM_XCDS_WS`` and the generated code is unchanged. The public ``grouped_gemm()`` op already rejects ``schedule="work_steal"`` combined with non-``None`` ``num_cu``, so high-level callers are unaffected. This commit hardens the cpp op ``ck_grouped_gemm``, which still exposes ``num_cu`` -- belt-and-braces. Regression: ``test_ck_grouped_gemm_op_ws_num_cu_below_num_xcds`` exercises ``num_cu`` in {4, 6, 7, 8, 16} on the raw op (forcing the per-xcd heuristic via ``resolve_ck_ws_local_per_xcd("per-xcd", ...)``) and asserts bit-equal output vs. the static-stride CK kernel.

…nt32 bounds, heuristic tests) Addressing Copilot review comments on PR AMD-AGI#348: - ``GroupedGEMMCKBackend.execute``: use ``a.size(1) if trans_a else a.size(0)`` when computing ``total_m`` for the heuristic. ``can_handle`` already rejects ``trans_a=True`` on this backend, so this is unreachable in practice, but keeps the shape logic honest if the constraint is ever relaxed. (Comment AMD-AGI#11.) - ``csrc/pytorch/grouped_gemm/ck_grouped_gemm.cpp``: reject ``ws_local_per_xcd`` values that would silently truncate on the ``int64_t -> int32_t`` cast passed to the CK kernel (``PRIMUS_TURBO_CHECK`` on ``[0, INT32_MAX]``). Applied to both the standard and variable-K bindings. (Comment AMD-AGI#12.) - ``tests/pytorch/ops/test_ws_ck_heuristic.py`` (NEW): pure-Python unit tests for the CK WS mode-resolution heuristic. Mirrors the Triton heuristic suite -- covers explicit modes, ``auto`` behavior for both standard and variable-K kernels around the thresholds, degenerate inputs (``total_tiles <= 0``, ``num_cus == 0``), the ``ValueError`` path for unknown modes, the ``total_tiles`` helpers, and pins the tile-shape / chiplet-count constants that must stay in sync with the C++ side. 15 tests, GPU-free. (Comments AMD-AGI#15, AMD-AGI#16.) Deferred (documented / noted, not fixed in this commit): - AMD-AGI#2 / stream-safety of the per-device singleton counter: doc'd on the cache with a clear "not stream-safe" caveat; the public API also rejects ``schedule="work_steal"`` + non-``None`` ``num_cu``. - AMD-AGI#3 / AMD-AGI#13: O(G) prefix-scan in the args-setup kernel is future work (irrelevant for current MoE shape ranges). - AMD-AGI#14: test parametrization size is workable now that the ``ws_mode`` parametrization is folded away by the ``schedule`` API rename. - AMD-AGI#15: CK_NUM_XCDS duplication with C++ NUM_XCDS_WS: single-source-of-truth would need a Python binding for a fixed hardware constant; kept parallel with a docstring note pointing at the C++ side.

The CK static grouped-GEMM kernel already uses the full 64 KB LDS budget on both gfx942 and gfx950. Our WS kernel adds a 4-byte ``__shared__`` broadcast slot for the atomicAdd result (thread-0 atomicAdd -> all threads see it), which pushes total LDS to 65540 bytes. The lld diagnostic on gfx942 is: lld: error: local memory (65540) exceeds limit (65536) in function ck_tile::kentry<1, GroupedGemmKernelWS<...>> (In AMD/OpenCL terms "local memory" = LDS, not the CUDA meaning of per-thread scratch. Address space 3.) WS was designed and characterized for MI355X (gfx950) exclusively; gfx942 is only in our build matrix for CI. Rather than reduce the CK tile config on gfx942 (which would cost a real perf hit on that arch), stub out the WS kernel body on gfx942: - ``grouped_gemm_kernel_ws.hpp``: guard the ``operator()`` body with ``#if !defined(__gfx942__)``. Host compile and gfx950 device compile see the full body; gfx942 device compile sees an empty stub (uses 0 LDS) so the .o links successfully. - ``grouped_gemm_impl.py``: ``GroupedGEMMCKBackend.can_handle`` and ``GroupedGEMMVariableKCKBackend.can_handle`` now check ``is_gfx950()`` and refuse to dispatch ``schedule="work_steal"`` to CK on any other arch. Since the gfx942 device stub is a no-op that would produce wrong output if launched, this must be gated at the dispatch layer. Verified: local build with ``GPU_ARCHS="gfx942;gfx950"`` succeeds; the lld overflow no longer occurs. All 76 WS-related pytests (Triton + CK) still pass on gfx950.

Copilot

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 6 comments.

+        if schedule == "work_steal" and not is_gfx950():
+            supported = False


+        # See GroupedGEMMCKBackend.can_handle: the CK WS kernel is
+        # gfx950-only due to LDS budget.
+        if schedule == "work_steal" and not is_gfx950():
+            supported = False


+        PRIMUS_TURBO_CHECK(counter.scalar_type() == at::kInt,
+                           "ws_counter must be int32");
+        PRIMUS_TURBO_CHECK(counter.is_cuda(), "ws_counter must be on CUDA");
+        constexpr int64_t kCounterSlots = 8 + 2; // NUM_XCDS_WS + global + done
+        PRIMUS_TURBO_CHECK(counter.numel() >= kCounterSlots,
+                           "ws_counter must have at least ", kCounterSlots, " int32 slots");


+        PRIMUS_TURBO_CHECK(counter.scalar_type() == at::kInt, "ws_counter must be int32");
+        PRIMUS_TURBO_CHECK(counter.is_cuda(), "ws_counter must be on CUDA");
+        constexpr int64_t kCounterSlots = 8 + 2;
+        PRIMUS_TURBO_CHECK(counter.numel() >= kCounterSlots,
+                           "ws_counter must have at least ", kCounterSlots, " int32 slots");


+    // WS kernel's FindGroupId binary search (and benign for the static path,
+    // which doesn't read these fields). m_tile/n_tile are the runner's tile
+    // shape, supplied by the dispatch site. Each thread does an O(group_id)
+    // scan over `group_lens_ptr` (already populated on device); for G <= ~256
+    // total work is O(G^2) <= ~65k loads -- negligible vs the GEMM.


+    // Prefix sum for block_start/block_end (variable-K: every group has the
+    // same M and N, only K varies, so the per-group tile count is uniform --
+    // unless `effective_m` is 0 for empty groups, which contributes 0 tiles).
+    // Same pattern as the forward kernel; the WS kernel's FindGroupId reads
+    // these fields.


…fx950 Comments AMD-AGI#1/AMD-AGI#2 (multi-GPU safety of the arch check): - ``is_gfx950()`` inspects ``torch.cuda.current_device()``, not the tensor's device -- wrong when the tensor lives on a different device than the ambient one. Introduce ``_is_gfx950_device(device)`` that reads compute capability of the *tensor*'s device, and call it with ``a.device`` in both CK ``can_handle`` methods. Comments AMD-AGI#3/AMD-AGI#4 (ws_counter device-index validation): - ``ck_grouped_gemm`` / ``ck_grouped_gemm_variable_k``: add ``PRIMUS_TURBO_CHECK(counter.device() == a.device(), ...)`` and ``counter.is_contiguous()`` so passing a counter allocated on a different CUDA device (or a non-contiguous view) fails at the binding with a clear message instead of a kernel-side undefined pointer crash. Report device indices via ``.index()`` since ``c10::Device`` isn't ``to_string_like``-convertible for the check macro. Comments AMD-AGI#5/AMD-AGI#6 (O(G^2) args-setup runs even on the static path): - ``compute_grouped_gemm_args`` and ``compute_grouped_gemm_variable_k_args`` gain a ``work_steal`` bool param and only compute the ``block_start`` / ``block_end`` prefix-sum when it's true. The static path never reads those fields, so this makes args-setup O(G) instead of O(G^2) for the (common) static case. WS unchanged. CI fixes (gfx942 unittest-pytorch): - ``test_grouped_gemm_schedule_work_steal[BackendType.CK-...]``: my ``can_handle`` rejects the combination on non-gfx950 (WS kernel body is stubbed on gfx942 due to the LDS budget), so the dispatcher raises ``ValueError: cannot handle`` and the test recorded 32 failures. Skip the CK branch of that parametrization on non-gfx950 hardware. - ``test_ck_grouped_gemm_op_ws_num_cu_below_num_xcds``: hits the cpp op directly, bypassing ``can_handle``. On gfx942 the WS kernel body is a no-op stub, so the cpp op silently returned zero-ish output and the bit-exact check failed. Add ``@pytest.mark.skipif`` gated on compute capability != (9, 5). Verified locally: 64/64 ``test_grouped_gemm_schedule_work_steal`` CK cases pass on gfx950 (this host), and 5/5 ``test_ck_grouped_gemm_op_ws_num_cu_below_num_xcds`` cases pass; the same tests were the 37 failing on the gfx942 CI job.

s_block_id is a __shared__ broadcast slot written by thread 0 and read by all threads. The trailing block_sync_lds() after Run() serves as the "before" barrier for the next iteration's write on the normal path, but two exit paths bypass it: the ``continue`` on ``block_id >= total_tiles`` and the Phase-1 ``break`` that falls straight into Phase 2's atomicAdd. On either path a straggler warp still reading iter N's s_block_id could be clobbered by iter N+1's write. Add __syncthreads() at the top of each loop iteration to plug both. The break/continue conditions are derived from s_block_id alone, so they're block-uniform and the added barrier cannot diverge. Perf on gfx950 (G=32 M=267424 K=1280 N=2560, 200 iters): fwd static→WS: 2.258→2.324 ms (unchanged vs pre-fix 2.250→2.328) dgrad : 2.658→2.692 ms (unchanged vs 2.662→2.695) wgrad : 2.319→2.278 ms (unchanged vs 2.315→2.271) Barrier cost is well below MFMA time; no measurable regression. All 94 WS pytest cases pass. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VXwBs62MZQfBqhLxyJ9rnV

Copilot

Pull request overview

Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.

+        index_t total_tiles = 0;
+        for(index_t g = 0; g < group_count; ++g)
+        {
+            const auto& kargs_g = gemm_desc_ptr[g].group_karg;
+            total_tiles += TilePartitioner::GridSize(kargs_g.M, kargs_g.N) * kargs_g.k_batch;
+        }


+        for (int64_t g = 0; g < group_id; ++g) {
+            const ck_tile::index_t M_g       = static_cast<ck_tile::index_t>(group_lens_ptr[g]);
+            const ck_tile::index_t m_tiles_g = (M_g + m_tile - 1) / m_tile;
+            my_block_start += m_tiles_g * num_n_tiles * k_batch;
+        }


+        for (int64_t g = 0; g < group_id; ++g) {
+            const ck_tile::index_t M_g = (group_lens_ptr[g] > 0) ? m : 0;
+            const ck_tile::index_t m_tiles_g = (M_g + m_tile - 1) / m_tile;
+            my_block_start += m_tiles_g * num_n_tiles * k_batch;
+        }


Replace the per-CTA O(G) recompute of total_tiles with a single load of gemm_desc_ptr[group_count - 1].block_end, which is already populated as a prefix sum by compute_grouped_gemm_args in ck_grouped_gemm.cu. At the target G=32 workload the group_descs array is fully L2-resident, so the perf delta is not measurable (WS medians unchanged: fwd 2.328, dgrad 2.692, wgrad 2.266 ms). The change is worthwhile as a code- clarity + O(G) → O(1) reduction that scales for larger G (e.g. MoE configs with G in the hundreds). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VXwBs62MZQfBqhLxyJ9rnV

Copilot AI review requested due to automatic review settings May 27, 2026 18:16

wenchenvincent requested review from RuibinCheung, wenxie-amd and xiaobochen-amd as code owners May 27, 2026 18:16

Copilot AI reviewed May 27, 2026

View reviewed changes

wenchenvincent changed the title ~~ck_grouped_gemm: add work-stealing variant with ws_mode API~~ feat: ck_grouped_gemm: add work-stealing variant with ws_mode API May 27, 2026

wenchenvincent force-pushed the ws-ck-grouped-gemm branch from d27e5a8 to 6acafdc Compare June 2, 2026 04:05

Copilot AI review requested due to automatic review settings June 2, 2026 04:16

Copilot AI reviewed Jun 2, 2026

View reviewed changes

wenchenvincent force-pushed the ws-ck-grouped-gemm branch from d52f4a8 to 1be272e Compare June 30, 2026 19:23

Copilot AI review requested due to automatic review settings July 1, 2026 03:18

Copilot started reviewing on behalf of wenchenvincent July 1, 2026 03:19 View session

Copilot AI reviewed Jul 1, 2026

View reviewed changes

wenchenvincent changed the title ~~feat: ck_grouped_gemm: add work-stealing variant with ws_mode API~~ feat(grouped_gemm): add CK work-stealing variant with schedule API Jul 2, 2026

wenchenvincent added 3 commits July 2, 2026 04:13

ck_grouped_gemm: add work-stealing variant

33095f6

wenchenvincent force-pushed the ws-ck-grouped-gemm branch from e352e10 to 9387449 Compare July 2, 2026 04:20

xiaobochen-amd added the ci:gpu label Jul 2, 2026

Copilot AI review requested due to automatic review settings July 3, 2026 03:16

Copilot started reviewing on behalf of wenchenvincent July 3, 2026 03:17 View session

Copilot AI reviewed Jul 3, 2026

View reviewed changes

wenchenvincent and others added 2 commits July 3, 2026 15:50

Copilot AI review requested due to automatic review settings July 6, 2026 23:12

Copilot started reviewing on behalf of wenchenvincent July 6, 2026 23:13 View session

Copilot AI reviewed Jul 6, 2026

View reviewed changes

xiaobochen-amd approved these changes Jul 8, 2026

View reviewed changes

xiaobochen-amd merged commit 59a9840 into AMD-AGI:main Jul 8, 2026
7 of 8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

feat(grouped_gemm): add CK work-stealing variant with schedule API#348

feat(grouped_gemm): add CK work-stealing variant with schedule API#348
xiaobochen-amd merged 7 commits into
AMD-AGI:mainfrom
wenchenvincent:ws-ck-grouped-gemm

wenchenvincent commented May 27, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

wenchenvincent Jun 2, 2026

Uh oh!

RuibinCheung commented May 28, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

		if len(group_lens) == 1:
		assert b.size(0) == 1, f"Expected first dimension to be 1, got {b.size(0)}"

		# ~50% per-xcd, ~50% global tail.
		return max(1, (total_tiles // 2) // CK_NUM_XCDS)

		if schedule == "work_steal" and not is_gfx950():
		supported = False

Uh oh!

Conversation

wenchenvincent commented May 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Type of change

Changes

Perf (MI355X, 288-shape MoE sweep, isolated kernel, clean methodology)

Checklist

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

wenchenvincent Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

RuibinCheung commented May 28, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

wenchenvincent commented May 27, 2026 •

edited

Loading