From df40cf97d4e806476e74ad72491845ca892dba38 Mon Sep 17 00:00:00 2001 From: Jari Kolehmainen Date: Thu, 27 Aug 2026 17:44:06 +0000 Subject: [PATCH] Isolate SM90 and SM100 MoE kernel paths (#1412) ## Summary - duplicate the complete MoE GEMM dependency graph for SM90 and SM100 behind compile-time entrypoint selectors - isolate architecture-specific math, TMEM, barriers, dispatch menus, and tuning-table wrappers so neither build depends on the opposite architecture - preserve mainline M128/NC4 communication geometry and add tests enforcing the boundary ## Validation - `make checkstyle` - architecture boundary tests: 3 passed - B200 distributed MoE suite: 8 passed - H100 distributed MoE suite: 8 passed - SM100: full SASS and resource report identical to mainline - SM90: normalized SASS and resource report identical to mainline; only compiler-generated internal NVSHMEM symbol hashes differ - paired performance (difference vs mainline): - SM100 Qwen3-30B-A3B: +0.11% total - SM100 Mixtral-8x7B: -0.16% total - SM90 Qwen3-30B-A3B: -0.20% total - SM90 Mixtral-8x7B: -0.22% total Co-authored-by: jkolehma Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../csrc/core/src/moe/cta_barrier.cuh | 140 +- .../csrc/core/src/moe/cta_barrier_sm100.cuh | 135 ++ .../csrc/core/src/moe/cta_barrier_sm90.cuh | 135 ++ liger_cute_kernels/csrc/core/src/moe/math.cuh | 160 +- .../csrc/core/src/moe/math_sm100.cuh | 155 ++ .../csrc/core/src/moe/math_sm90.cuh | 155 ++ liger_cute_kernels/csrc/core/src/moe/mlp.cuh | 422 +--- .../csrc/core/src/moe/mlp1_fused.cuh | 937 +------- .../csrc/core/src/moe/mlp1_fused_act.cuh | 706 +----- .../core/src/moe/mlp1_fused_act_sm100.cuh | 703 ++++++ .../csrc/core/src/moe/mlp1_fused_act_sm90.cuh | 703 ++++++ .../csrc/core/src/moe/mlp1_fused_sm100.cuh | 934 ++++++++ .../csrc/core/src/moe/mlp1_fused_sm90.cuh | 934 ++++++++ .../csrc/core/src/moe/mlp2_fused.cuh | 833 +------ .../csrc/core/src/moe/mlp2_fused_sm100.cuh | 830 +++++++ .../csrc/core/src/moe/mlp2_fused_sm90.cuh | 830 +++++++ .../csrc/core/src/moe/mlp2_t.cuh | 547 +---- .../csrc/core/src/moe/mlp2_t_fused.cuh | 682 +----- .../csrc/core/src/moe/mlp2_t_fused_sm100.cuh | 679 ++++++ .../csrc/core/src/moe/mlp2_t_fused_sm90.cuh | 679 ++++++ .../csrc/core/src/moe/mlp2_t_sm100.cuh | 540 +++++ .../csrc/core/src/moe/mlp2_t_sm90.cuh | 540 +++++ liger_cute_kernels/csrc/core/src/moe/mlp3.cuh | 1785 +------------- .../csrc/core/src/moe/mlp3_sm100.cuh | 1784 ++++++++++++++ .../csrc/core/src/moe/mlp3_sm90.cuh | 1784 ++++++++++++++ liger_cute_kernels/csrc/core/src/moe/mlp4.cuh | 1840 +-------------- .../csrc/core/src/moe/mlp4_sm100.cuh | 1839 +++++++++++++++ .../csrc/core/src/moe/mlp4_sm90.cuh | 1839 +++++++++++++++ liger_cute_kernels/csrc/core/src/moe/mlp5.cuh | 606 +---- .../csrc/core/src/moe/mlp5_fused.cuh | 632 +---- .../csrc/core/src/moe/mlp5_fused_sm100.cuh | 629 +++++ .../csrc/core/src/moe/mlp5_fused_sm90.cuh | 629 +++++ .../csrc/core/src/moe/mlp5_sm100.cuh | 599 +++++ .../csrc/core/src/moe/mlp5_sm90.cuh | 599 +++++ .../csrc/core/src/moe/mlp_bwd.cuh | 1892 +-------------- .../csrc/core/src/moe/mlp_bwd_sm100.cuh | 1885 +++++++++++++++ .../csrc/core/src/moe/mlp_bwd_sm90.cuh | 1885 +++++++++++++++ .../csrc/core/src/moe/mlp_comms.cuh | 1350 +---------- .../csrc/core/src/moe/mlp_comms_bwd.cuh | 779 +----- .../csrc/core/src/moe/mlp_comms_bwd_sm100.cuh | 772 ++++++ .../csrc/core/src/moe/mlp_comms_bwd_sm90.cuh | 772 ++++++ .../csrc/core/src/moe/mlp_comms_sm100.cuh | 1343 +++++++++++ .../csrc/core/src/moe/mlp_comms_sm90.cuh | 1343 +++++++++++ .../csrc/core/src/moe/mlp_sm100.cuh | 417 ++++ .../csrc/core/src/moe/mlp_sm90.cuh | 417 ++++ .../csrc/core/src/moe/models.cuh | 194 +- .../csrc/core/src/moe/models_sm100.cuh | 189 ++ .../csrc/core/src/moe/models_sm90.cuh | 189 ++ liger_cute_kernels/csrc/core/src/moe/moe.cu | 1774 +------------- liger_cute_kernels/csrc/core/src/moe/moe.cuh | 205 +- .../csrc/core/src/moe/moe_bwd.cu | 2091 +---------------- .../csrc/core/src/moe/moe_bwd.cuh | 224 +- .../csrc/core/src/moe/moe_bwd_sm100.cuh | 219 ++ .../csrc/core/src/moe/moe_bwd_sm100.inc | 2091 +++++++++++++++++ .../csrc/core/src/moe/moe_bwd_sm90.cuh | 219 ++ .../csrc/core/src/moe/moe_bwd_sm90.inc | 2091 +++++++++++++++++ .../csrc/core/src/moe/moe_comm_config.cuh | 50 +- .../core/src/moe/moe_comm_config_sm100.cuh | 45 + .../core/src/moe/moe_comm_config_sm90.cuh | 45 + .../src/moe/moe_dispatch_configs_sm100.cuh | 2 +- .../src/moe/moe_dispatch_configs_sm90.cuh | 2 +- .../moe/moe_fwd_bwd_tune_configs_sm100.hpp | 182 ++ .../src/moe/moe_fwd_bwd_tune_configs_sm90.hpp | 182 ++ .../moe/moe_fwd_bwd_tuning_configs_sm100.cuh | 29 + .../moe/moe_fwd_bwd_tuning_configs_sm90.cuh | 29 + .../csrc/core/src/moe/moe_sm100.cuh | 200 ++ .../csrc/core/src/moe/moe_sm100.inc | 1774 ++++++++++++++ .../csrc/core/src/moe/moe_sm90.cuh | 200 ++ .../csrc/core/src/moe/moe_sm90.inc | 1774 ++++++++++++++ .../csrc/core/src/moe/moe_symm_config.cuh | 57 +- .../core/src/moe/moe_symm_config_sm100.cuh | 52 + .../core/src/moe/moe_symm_config_sm90.cuh | 52 + .../csrc/core/src/moe/silu_bwd_fused.cuh | 249 +- .../core/src/moe/silu_bwd_fused_sm100.cuh | 244 ++ .../csrc/core/src/moe/silu_bwd_fused_sm90.cuh | 244 ++ .../csrc/core/src/moe/tile_iterator.cuh | 532 +---- .../csrc/core/src/moe/tile_iterator_bwd.cuh | 364 +-- .../core/src/moe/tile_iterator_bwd_sm100.cuh | 359 +++ .../core/src/moe/tile_iterator_bwd_sm90.cuh | 359 +++ .../csrc/core/src/moe/tile_iterator_sm100.cuh | 527 +++++ .../csrc/core/src/moe/tile_iterator_sm90.cuh | 527 +++++ .../csrc/core/src/moe/tma_copy_atom.cuh | 40 +- .../csrc/core/src/moe/tma_copy_atom_sm100.cuh | 35 + .../csrc/core/src/moe/tma_copy_atom_sm90.cuh | 35 + .../csrc/core/src/moe/tmem_load_op.cuh | 25 +- .../csrc/core/src/moe/tmem_load_op_sm100.cuh | 20 + .../csrc/core/src/moe/tmem_load_op_sm90.cuh | 20 + .../test/test_moe_architecture_split.py | 86 + 88 files changed, 38712 insertions(+), 18914 deletions(-) create mode 100644 liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/math_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/math_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp3_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp3_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp4_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp4_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp5_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp5_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/mlp_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/models_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/models_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.inc create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.inc create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm100.hpp create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm90.hpp create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_sm100.inc create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_sm90.inc create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm90.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm100.cuh create mode 100644 liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm90.cuh create mode 100644 liger_cute_kernels/test/test_moe_architecture_split.py diff --git a/liger_cute_kernels/csrc/core/src/moe/cta_barrier.cuh b/liger_cute_kernels/csrc/core/src/moe/cta_barrier.cuh index 836192bf..11d9b6a5 100644 --- a/liger_cute_kernels/csrc/core/src/moe/cta_barrier.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/cta_barrier.cuh @@ -1,135 +1,11 @@ #pragma once -#include -#include -#include +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -namespace liger { - -// ═══════════════════════════════════════════════════════════════════ -// CtaCounterBarrier — target-based cross-CTA monotonic barrier -// ═══════════════════════════════════════════════════════════════════ -// -// Synchronizes a group of CTAs (or column of CTAs) by incrementing a -// shared device-memory counter and waiting until it has accumulated -// one increment from every participant of the current round. -// -// ── Why not modulo? ──────────────────────────────────────────────── -// The naive `while (counter % stride != 0)` variant deadlocks under -// CTA-execution skew: -// 1. counter advances 0 → stride as all CTAs add 1 each. -// 2. A fast CTA polls at counter==stride, sees %==0, exits. -// 3. It runs post-barrier work AND returns for round 2's add → -// counter==stride+1. -// 4. A slow CTA that hadn't yet polled at counter==stride now reads -// counter==stride+1. `1 % stride != 0`, keeps spinning. -// 5. As other fast CTAs each add for round 2 the counter walks up -// to 2*stride − 1 (everyone except the slow one). `(2*stride−1) -// % stride == stride−1`, never 0. Deadlock. -// -// ── Why not derive the target from atomicAdd's return? ───────────── -// Snapshotting `prev = atomicAdd(counter, 1)` and computing -// `target = ((prev/stride) + 1)*stride` looks similar but is racy: -// two CTAs entering the SAME conceptual round can land on different -// sides of a stride boundary (one gets prev=stride−1, the other -// prev=stride), and they then wait for DIFFERENT targets — one for -// round K, the other for round K+1. The "round" assignment depends -// on the order in which CTAs win the atomicAdd, which is exactly -// the kind of skew the barrier is supposed to mask. -// -// ── The fix ──────────────────────────────────────────────────────── -// Each barrier instance owns a per-CTA `target` in register memory. -// The constructor initializes target=0; every wait() bumps it by -// `stride` BEFORE the atomicAdd. Each CTA always knows exactly which -// round it is in (its own call count × stride), independent of the -// race between CTAs. -// -// ── Usage ────────────────────────────────────────────────────────── -// liger::CtaCounterBarrier barrier( -// counter_ptr, stride); -// barrier.wait(); // round 1 -// ... work ... -// barrier.wait(); // round 2 -// -// Each CTA must construct the barrier ONCE and reuse it for the -// matching sequence of waits — constructing a fresh barrier per -// wait() resets target=0 and is equivalent to the broken behavior -// above. -// -// `kNumThreads` and `kBarrierId` are template parameters for the -// intra-CTA NamedBarrier (so MLP-only barriers can exclude comm warps -// via a smaller NumThreads). The counter holds device-memory data -// and must be zeroed at kernel-launch time. -// -template -struct CtaCounterBarrier { - int* counter; // device memory - int stride; // number of CTAs that contribute per wait() - int target; // per-CTA monotonic target — register memory - - // ── Pre-condition: the host MUST cudaMemsetAsync `counter` to 0 - // before every kernel launch. The constructor sets target=0 and - // each wait() then advances target by `stride` while atomically - // incrementing counter. If the caller skips the zero-out and - // counter still holds a residual value from the previous launch, - // the very first wait() will see counter >> target and exit - // without actually waiting for anyone — silent corruption. - // - // DO NOT read counter here to "derive" the starting target. That - // read races with any thread/CTA that has begun wait() — there - // is no implicit sync between CTAs at kernel prologue. The only - // race-free contract is: host zeros the counter, kernel - // constructs barrier with target=0. Keep it that way. - __device__ __forceinline__ CtaCounterBarrier(int* counter_, int stride_) - : counter(counter_), stride(stride_), target(0) {} - - __device__ __forceinline__ void wait() { - // Bump the per-CTA target by `stride` first, so we know - // which round this wait() corresponds to independent of - // the race between CTAs at the atomicAdd below. - target += stride; - __threadfence(); - cutlass::arch::NamedBarrier::sync(NumThreads, BarrierId); - if (threadIdx.x == 0) { - cuda::atomic_ref r(*counter); - r.fetch_add(1, cuda::memory_order_release); - while (r.load(cuda::memory_order_acquire) < target) {} - } - cutlass::arch::NamedBarrier::sync(NumThreads, BarrierId); - } -}; - -// ═══════════════════════════════════════════════════════════════════ -// SyncThreadsCtaCounterBarrier — same idea, intra-CTA via __syncthreads -// ═══════════════════════════════════════════════════════════════════ -// -// Variant for callers that need to sync ALL threads in the CTA (not -// just an MLP-warp subset). Used by PeSync::barrier where every thread -// — including comm warps — must be at the barrier before the atomic -// increments and the cross-PE signal_op fire. -struct SyncThreadsCtaCounterBarrier { - int* counter; - int stride; - int target; - - // Same pre-condition as CtaCounterBarrier: the host MUST zero - // `counter` before every kernel launch. Constructor sets - // target=0 — DO NOT read counter here to derive the starting - // value, that races with any thread that already entered wait(). - __device__ __forceinline__ SyncThreadsCtaCounterBarrier(int* counter_, int stride_) - : counter(counter_), stride(stride_), target(0) {} - - __device__ __forceinline__ void wait() { - target += stride; - __threadfence(); - __syncthreads(); - if (threadIdx.x == 0) { - cuda::atomic_ref r(*counter); - r.fetch_add(1, cuda::memory_order_release); - while (r.load(cuda::memory_order_acquire) < target) {} - } - __syncthreads(); - } -}; - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "cta_barrier_sm90.cuh" +#else +#include "cta_barrier_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm100.cuh new file mode 100644 index 00000000..836192bf --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm100.cuh @@ -0,0 +1,135 @@ +#pragma once + +#include +#include +#include + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// CtaCounterBarrier — target-based cross-CTA monotonic barrier +// ═══════════════════════════════════════════════════════════════════ +// +// Synchronizes a group of CTAs (or column of CTAs) by incrementing a +// shared device-memory counter and waiting until it has accumulated +// one increment from every participant of the current round. +// +// ── Why not modulo? ──────────────────────────────────────────────── +// The naive `while (counter % stride != 0)` variant deadlocks under +// CTA-execution skew: +// 1. counter advances 0 → stride as all CTAs add 1 each. +// 2. A fast CTA polls at counter==stride, sees %==0, exits. +// 3. It runs post-barrier work AND returns for round 2's add → +// counter==stride+1. +// 4. A slow CTA that hadn't yet polled at counter==stride now reads +// counter==stride+1. `1 % stride != 0`, keeps spinning. +// 5. As other fast CTAs each add for round 2 the counter walks up +// to 2*stride − 1 (everyone except the slow one). `(2*stride−1) +// % stride == stride−1`, never 0. Deadlock. +// +// ── Why not derive the target from atomicAdd's return? ───────────── +// Snapshotting `prev = atomicAdd(counter, 1)` and computing +// `target = ((prev/stride) + 1)*stride` looks similar but is racy: +// two CTAs entering the SAME conceptual round can land on different +// sides of a stride boundary (one gets prev=stride−1, the other +// prev=stride), and they then wait for DIFFERENT targets — one for +// round K, the other for round K+1. The "round" assignment depends +// on the order in which CTAs win the atomicAdd, which is exactly +// the kind of skew the barrier is supposed to mask. +// +// ── The fix ──────────────────────────────────────────────────────── +// Each barrier instance owns a per-CTA `target` in register memory. +// The constructor initializes target=0; every wait() bumps it by +// `stride` BEFORE the atomicAdd. Each CTA always knows exactly which +// round it is in (its own call count × stride), independent of the +// race between CTAs. +// +// ── Usage ────────────────────────────────────────────────────────── +// liger::CtaCounterBarrier barrier( +// counter_ptr, stride); +// barrier.wait(); // round 1 +// ... work ... +// barrier.wait(); // round 2 +// +// Each CTA must construct the barrier ONCE and reuse it for the +// matching sequence of waits — constructing a fresh barrier per +// wait() resets target=0 and is equivalent to the broken behavior +// above. +// +// `kNumThreads` and `kBarrierId` are template parameters for the +// intra-CTA NamedBarrier (so MLP-only barriers can exclude comm warps +// via a smaller NumThreads). The counter holds device-memory data +// and must be zeroed at kernel-launch time. +// +template +struct CtaCounterBarrier { + int* counter; // device memory + int stride; // number of CTAs that contribute per wait() + int target; // per-CTA monotonic target — register memory + + // ── Pre-condition: the host MUST cudaMemsetAsync `counter` to 0 + // before every kernel launch. The constructor sets target=0 and + // each wait() then advances target by `stride` while atomically + // incrementing counter. If the caller skips the zero-out and + // counter still holds a residual value from the previous launch, + // the very first wait() will see counter >> target and exit + // without actually waiting for anyone — silent corruption. + // + // DO NOT read counter here to "derive" the starting target. That + // read races with any thread/CTA that has begun wait() — there + // is no implicit sync between CTAs at kernel prologue. The only + // race-free contract is: host zeros the counter, kernel + // constructs barrier with target=0. Keep it that way. + __device__ __forceinline__ CtaCounterBarrier(int* counter_, int stride_) + : counter(counter_), stride(stride_), target(0) {} + + __device__ __forceinline__ void wait() { + // Bump the per-CTA target by `stride` first, so we know + // which round this wait() corresponds to independent of + // the race between CTAs at the atomicAdd below. + target += stride; + __threadfence(); + cutlass::arch::NamedBarrier::sync(NumThreads, BarrierId); + if (threadIdx.x == 0) { + cuda::atomic_ref r(*counter); + r.fetch_add(1, cuda::memory_order_release); + while (r.load(cuda::memory_order_acquire) < target) {} + } + cutlass::arch::NamedBarrier::sync(NumThreads, BarrierId); + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// SyncThreadsCtaCounterBarrier — same idea, intra-CTA via __syncthreads +// ═══════════════════════════════════════════════════════════════════ +// +// Variant for callers that need to sync ALL threads in the CTA (not +// just an MLP-warp subset). Used by PeSync::barrier where every thread +// — including comm warps — must be at the barrier before the atomic +// increments and the cross-PE signal_op fire. +struct SyncThreadsCtaCounterBarrier { + int* counter; + int stride; + int target; + + // Same pre-condition as CtaCounterBarrier: the host MUST zero + // `counter` before every kernel launch. Constructor sets + // target=0 — DO NOT read counter here to derive the starting + // value, that races with any thread that already entered wait(). + __device__ __forceinline__ SyncThreadsCtaCounterBarrier(int* counter_, int stride_) + : counter(counter_), stride(stride_), target(0) {} + + __device__ __forceinline__ void wait() { + target += stride; + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) { + cuda::atomic_ref r(*counter); + r.fetch_add(1, cuda::memory_order_release); + while (r.load(cuda::memory_order_acquire) < target) {} + } + __syncthreads(); + } +}; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm90.cuh new file mode 100644 index 00000000..836192bf --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm90.cuh @@ -0,0 +1,135 @@ +#pragma once + +#include +#include +#include + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// CtaCounterBarrier — target-based cross-CTA monotonic barrier +// ═══════════════════════════════════════════════════════════════════ +// +// Synchronizes a group of CTAs (or column of CTAs) by incrementing a +// shared device-memory counter and waiting until it has accumulated +// one increment from every participant of the current round. +// +// ── Why not modulo? ──────────────────────────────────────────────── +// The naive `while (counter % stride != 0)` variant deadlocks under +// CTA-execution skew: +// 1. counter advances 0 → stride as all CTAs add 1 each. +// 2. A fast CTA polls at counter==stride, sees %==0, exits. +// 3. It runs post-barrier work AND returns for round 2's add → +// counter==stride+1. +// 4. A slow CTA that hadn't yet polled at counter==stride now reads +// counter==stride+1. `1 % stride != 0`, keeps spinning. +// 5. As other fast CTAs each add for round 2 the counter walks up +// to 2*stride − 1 (everyone except the slow one). `(2*stride−1) +// % stride == stride−1`, never 0. Deadlock. +// +// ── Why not derive the target from atomicAdd's return? ───────────── +// Snapshotting `prev = atomicAdd(counter, 1)` and computing +// `target = ((prev/stride) + 1)*stride` looks similar but is racy: +// two CTAs entering the SAME conceptual round can land on different +// sides of a stride boundary (one gets prev=stride−1, the other +// prev=stride), and they then wait for DIFFERENT targets — one for +// round K, the other for round K+1. The "round" assignment depends +// on the order in which CTAs win the atomicAdd, which is exactly +// the kind of skew the barrier is supposed to mask. +// +// ── The fix ──────────────────────────────────────────────────────── +// Each barrier instance owns a per-CTA `target` in register memory. +// The constructor initializes target=0; every wait() bumps it by +// `stride` BEFORE the atomicAdd. Each CTA always knows exactly which +// round it is in (its own call count × stride), independent of the +// race between CTAs. +// +// ── Usage ────────────────────────────────────────────────────────── +// liger::CtaCounterBarrier barrier( +// counter_ptr, stride); +// barrier.wait(); // round 1 +// ... work ... +// barrier.wait(); // round 2 +// +// Each CTA must construct the barrier ONCE and reuse it for the +// matching sequence of waits — constructing a fresh barrier per +// wait() resets target=0 and is equivalent to the broken behavior +// above. +// +// `kNumThreads` and `kBarrierId` are template parameters for the +// intra-CTA NamedBarrier (so MLP-only barriers can exclude comm warps +// via a smaller NumThreads). The counter holds device-memory data +// and must be zeroed at kernel-launch time. +// +template +struct CtaCounterBarrier { + int* counter; // device memory + int stride; // number of CTAs that contribute per wait() + int target; // per-CTA monotonic target — register memory + + // ── Pre-condition: the host MUST cudaMemsetAsync `counter` to 0 + // before every kernel launch. The constructor sets target=0 and + // each wait() then advances target by `stride` while atomically + // incrementing counter. If the caller skips the zero-out and + // counter still holds a residual value from the previous launch, + // the very first wait() will see counter >> target and exit + // without actually waiting for anyone — silent corruption. + // + // DO NOT read counter here to "derive" the starting target. That + // read races with any thread/CTA that has begun wait() — there + // is no implicit sync between CTAs at kernel prologue. The only + // race-free contract is: host zeros the counter, kernel + // constructs barrier with target=0. Keep it that way. + __device__ __forceinline__ CtaCounterBarrier(int* counter_, int stride_) + : counter(counter_), stride(stride_), target(0) {} + + __device__ __forceinline__ void wait() { + // Bump the per-CTA target by `stride` first, so we know + // which round this wait() corresponds to independent of + // the race between CTAs at the atomicAdd below. + target += stride; + __threadfence(); + cutlass::arch::NamedBarrier::sync(NumThreads, BarrierId); + if (threadIdx.x == 0) { + cuda::atomic_ref r(*counter); + r.fetch_add(1, cuda::memory_order_release); + while (r.load(cuda::memory_order_acquire) < target) {} + } + cutlass::arch::NamedBarrier::sync(NumThreads, BarrierId); + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// SyncThreadsCtaCounterBarrier — same idea, intra-CTA via __syncthreads +// ═══════════════════════════════════════════════════════════════════ +// +// Variant for callers that need to sync ALL threads in the CTA (not +// just an MLP-warp subset). Used by PeSync::barrier where every thread +// — including comm warps — must be at the barrier before the atomic +// increments and the cross-PE signal_op fire. +struct SyncThreadsCtaCounterBarrier { + int* counter; + int stride; + int target; + + // Same pre-condition as CtaCounterBarrier: the host MUST zero + // `counter` before every kernel launch. Constructor sets + // target=0 — DO NOT read counter here to derive the starting + // value, that races with any thread that already entered wait(). + __device__ __forceinline__ SyncThreadsCtaCounterBarrier(int* counter_, int stride_) + : counter(counter_), stride(stride_), target(0) {} + + __device__ __forceinline__ void wait() { + target += stride; + __threadfence(); + __syncthreads(); + if (threadIdx.x == 0) { + cuda::atomic_ref r(*counter); + r.fetch_add(1, cuda::memory_order_release); + while (r.load(cuda::memory_order_acquire) < target) {} + } + __syncthreads(); + } +}; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/math.cuh b/liger_cute_kernels/csrc/core/src/moe/math.cuh index 38021a91..05f2f38c 100644 --- a/liger_cute_kernels/csrc/core/src/moe/math.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/math.cuh @@ -1,155 +1,11 @@ #pragma once -#include -#include -#include -#include +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -namespace liger { - -// ═══════════════════════════════════════════════════════════════════ -// Sigmoid and SiLU (Swish) — exact, matches torch.sigmoid / F.silu -// ═══════════════════════════════════════════════════════════════════ - -__device__ __forceinline__ float fast_sigmoid(float x) { - // if (x >= 4.0f) return 1.0f; - // if (x <= -4.0f) return 0.0f; - // const float x2 = x * x; - // return 0.5f + x * (0.23897898f + x2 * (-0.01301264f + x2 * 0.00035572f)); - return 1.0f / (1.0f + expf(-x)); -} - -__device__ __forceinline__ float fast_silu(float x) { - return x * fast_sigmoid(x); -} - -__device__ __forceinline__ float grad_fast_silu(float x) { - const float f = fast_sigmoid(x); - return f * (1.0f + x * (1.0f - f)); -} - -// ═══════════════════════════════════════════════════════════════════ -// In-place smem tile transpose — single warp group (128 threads) -// ═══════════════════════════════════════════════════════════════════ -// -// Transposes a square [N, N] smem tile in place by swapping upper- -// triangle elements: for each (r, c) with r < c, swap tensor(r, c) -// with tensor(c, r). Diagonal elements are untouched. -// -// Race-free: each (r, c) pair is assigned to exactly one thread. -// No pair can collide — if thread A swaps (r, c) ↔ (c, r) with -// r < c, no other thread touches either element because that would -// require c < r (contradiction). -// -// CuTe's operator() handles swizzled addressing automatically. -// -// Usage: -// auto sW = make_tensor(make_smem_ptr(ptr), SmemLayoutW_stage{}); -// smem_transpose_inplace_wg<128>(sW); -// -// Requires: square tile (dim 0 == dim 1), called by one warp group. - -template -__device__ __forceinline__ void smem_transpose_inplace_wg(Tensor& tensor) { - constexpr int kWarpGroupSize = 128; - int tid = threadIdx.x % kWarpGroupSize; - - for (int r = 0; r < N - 1; ++r) { - for (int c = r + 1 + tid; c < N; c += kWarpGroupSize) { - auto tmp = tensor(r, c); - tensor(r, c) = tensor(c, r); - tensor(c, r) = tmp; - } - } -} - -// ═══════════════════════════════════════════════════════════════════ -// float → Element conversion (template specialization per type) -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ Element from_float(float x); - -template <> -__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float x) { - return __float2bfloat16(x); -} - -template <> -__device__ __forceinline__ __half from_float<__half>(float x) { - return __float2half(x); -} - -template <> -__device__ __forceinline__ cutlass::bfloat16_t from_float(float x) { - __nv_bfloat16 tmp = __float2bfloat16(x); - return reinterpret_cast(tmp); -} - -template <> -__device__ __forceinline__ cutlass::half_t from_float(float x) { - __half tmp = __float2half(x); - return reinterpret_cast(tmp); -} - -// ═══════════════════════════════════════════════════════════════════ -// Element → float conversion (template specialization per type) -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ float to_float(Element v); - -template <> -__device__ __forceinline__ float to_float<__nv_bfloat16>(__nv_bfloat16 v) { - return __bfloat162float(v); -} - -template <> -__device__ __forceinline__ float to_float<__half>(__half v) { - return __half2float(v); -} - -template <> -__device__ __forceinline__ float to_float(cutlass::bfloat16_t v) { - return __bfloat162float(reinterpret_cast(v)); -} - -template <> -__device__ __forceinline__ float to_float(cutlass::half_t v) { - return __half2float(reinterpret_cast(v)); -} - -// ═══════════════════════════════════════════════════════════════════ -// Cached global load (__ldg) for 2-byte element types -// ═══════════════════════════════════════════════════════════════════ -// -// CUDA provides __ldg for __nv_bfloat16 / __half directly; the cutlass -// wrappers (cutlass::bfloat16_t / cutlass::half_t) require a reinterpret -// since the byte layout is identical. - -template -__device__ __forceinline__ T ldg_elem(const T* p); - -template <> -__device__ __forceinline__ __nv_bfloat16 ldg_elem<__nv_bfloat16>(const __nv_bfloat16* p) { - return __ldg(p); -} - -template <> -__device__ __forceinline__ __half ldg_elem<__half>(const __half* p) { - return __ldg(p); -} - -template <> -__device__ __forceinline__ cutlass::bfloat16_t ldg_elem(const cutlass::bfloat16_t* p) { - __nv_bfloat16 raw = __ldg(reinterpret_cast(p)); - return reinterpret_cast(raw); -} - -template <> -__device__ __forceinline__ cutlass::half_t ldg_elem(const cutlass::half_t* p) { - __half raw = __ldg(reinterpret_cast(p)); - return reinterpret_cast(raw); -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "math_sm90.cuh" +#else +#include "math_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/math_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/math_sm100.cuh new file mode 100644 index 00000000..38021a91 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/math_sm100.cuh @@ -0,0 +1,155 @@ +#pragma once + +#include +#include +#include +#include + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// Sigmoid and SiLU (Swish) — exact, matches torch.sigmoid / F.silu +// ═══════════════════════════════════════════════════════════════════ + +__device__ __forceinline__ float fast_sigmoid(float x) { + // if (x >= 4.0f) return 1.0f; + // if (x <= -4.0f) return 0.0f; + // const float x2 = x * x; + // return 0.5f + x * (0.23897898f + x2 * (-0.01301264f + x2 * 0.00035572f)); + return 1.0f / (1.0f + expf(-x)); +} + +__device__ __forceinline__ float fast_silu(float x) { + return x * fast_sigmoid(x); +} + +__device__ __forceinline__ float grad_fast_silu(float x) { + const float f = fast_sigmoid(x); + return f * (1.0f + x * (1.0f - f)); +} + +// ═══════════════════════════════════════════════════════════════════ +// In-place smem tile transpose — single warp group (128 threads) +// ═══════════════════════════════════════════════════════════════════ +// +// Transposes a square [N, N] smem tile in place by swapping upper- +// triangle elements: for each (r, c) with r < c, swap tensor(r, c) +// with tensor(c, r). Diagonal elements are untouched. +// +// Race-free: each (r, c) pair is assigned to exactly one thread. +// No pair can collide — if thread A swaps (r, c) ↔ (c, r) with +// r < c, no other thread touches either element because that would +// require c < r (contradiction). +// +// CuTe's operator() handles swizzled addressing automatically. +// +// Usage: +// auto sW = make_tensor(make_smem_ptr(ptr), SmemLayoutW_stage{}); +// smem_transpose_inplace_wg<128>(sW); +// +// Requires: square tile (dim 0 == dim 1), called by one warp group. + +template +__device__ __forceinline__ void smem_transpose_inplace_wg(Tensor& tensor) { + constexpr int kWarpGroupSize = 128; + int tid = threadIdx.x % kWarpGroupSize; + + for (int r = 0; r < N - 1; ++r) { + for (int c = r + 1 + tid; c < N; c += kWarpGroupSize) { + auto tmp = tensor(r, c); + tensor(r, c) = tensor(c, r); + tensor(c, r) = tmp; + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// float → Element conversion (template specialization per type) +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ Element from_float(float x); + +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float x) { + return __float2bfloat16(x); +} + +template <> +__device__ __forceinline__ __half from_float<__half>(float x) { + return __float2half(x); +} + +template <> +__device__ __forceinline__ cutlass::bfloat16_t from_float(float x) { + __nv_bfloat16 tmp = __float2bfloat16(x); + return reinterpret_cast(tmp); +} + +template <> +__device__ __forceinline__ cutlass::half_t from_float(float x) { + __half tmp = __float2half(x); + return reinterpret_cast(tmp); +} + +// ═══════════════════════════════════════════════════════════════════ +// Element → float conversion (template specialization per type) +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ float to_float(Element v); + +template <> +__device__ __forceinline__ float to_float<__nv_bfloat16>(__nv_bfloat16 v) { + return __bfloat162float(v); +} + +template <> +__device__ __forceinline__ float to_float<__half>(__half v) { + return __half2float(v); +} + +template <> +__device__ __forceinline__ float to_float(cutlass::bfloat16_t v) { + return __bfloat162float(reinterpret_cast(v)); +} + +template <> +__device__ __forceinline__ float to_float(cutlass::half_t v) { + return __half2float(reinterpret_cast(v)); +} + +// ═══════════════════════════════════════════════════════════════════ +// Cached global load (__ldg) for 2-byte element types +// ═══════════════════════════════════════════════════════════════════ +// +// CUDA provides __ldg for __nv_bfloat16 / __half directly; the cutlass +// wrappers (cutlass::bfloat16_t / cutlass::half_t) require a reinterpret +// since the byte layout is identical. + +template +__device__ __forceinline__ T ldg_elem(const T* p); + +template <> +__device__ __forceinline__ __nv_bfloat16 ldg_elem<__nv_bfloat16>(const __nv_bfloat16* p) { + return __ldg(p); +} + +template <> +__device__ __forceinline__ __half ldg_elem<__half>(const __half* p) { + return __ldg(p); +} + +template <> +__device__ __forceinline__ cutlass::bfloat16_t ldg_elem(const cutlass::bfloat16_t* p) { + __nv_bfloat16 raw = __ldg(reinterpret_cast(p)); + return reinterpret_cast(raw); +} + +template <> +__device__ __forceinline__ cutlass::half_t ldg_elem(const cutlass::half_t* p) { + __half raw = __ldg(reinterpret_cast(p)); + return reinterpret_cast(raw); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/math_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/math_sm90.cuh new file mode 100644 index 00000000..38021a91 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/math_sm90.cuh @@ -0,0 +1,155 @@ +#pragma once + +#include +#include +#include +#include + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// Sigmoid and SiLU (Swish) — exact, matches torch.sigmoid / F.silu +// ═══════════════════════════════════════════════════════════════════ + +__device__ __forceinline__ float fast_sigmoid(float x) { + // if (x >= 4.0f) return 1.0f; + // if (x <= -4.0f) return 0.0f; + // const float x2 = x * x; + // return 0.5f + x * (0.23897898f + x2 * (-0.01301264f + x2 * 0.00035572f)); + return 1.0f / (1.0f + expf(-x)); +} + +__device__ __forceinline__ float fast_silu(float x) { + return x * fast_sigmoid(x); +} + +__device__ __forceinline__ float grad_fast_silu(float x) { + const float f = fast_sigmoid(x); + return f * (1.0f + x * (1.0f - f)); +} + +// ═══════════════════════════════════════════════════════════════════ +// In-place smem tile transpose — single warp group (128 threads) +// ═══════════════════════════════════════════════════════════════════ +// +// Transposes a square [N, N] smem tile in place by swapping upper- +// triangle elements: for each (r, c) with r < c, swap tensor(r, c) +// with tensor(c, r). Diagonal elements are untouched. +// +// Race-free: each (r, c) pair is assigned to exactly one thread. +// No pair can collide — if thread A swaps (r, c) ↔ (c, r) with +// r < c, no other thread touches either element because that would +// require c < r (contradiction). +// +// CuTe's operator() handles swizzled addressing automatically. +// +// Usage: +// auto sW = make_tensor(make_smem_ptr(ptr), SmemLayoutW_stage{}); +// smem_transpose_inplace_wg<128>(sW); +// +// Requires: square tile (dim 0 == dim 1), called by one warp group. + +template +__device__ __forceinline__ void smem_transpose_inplace_wg(Tensor& tensor) { + constexpr int kWarpGroupSize = 128; + int tid = threadIdx.x % kWarpGroupSize; + + for (int r = 0; r < N - 1; ++r) { + for (int c = r + 1 + tid; c < N; c += kWarpGroupSize) { + auto tmp = tensor(r, c); + tensor(r, c) = tensor(c, r); + tensor(c, r) = tmp; + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// float → Element conversion (template specialization per type) +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ Element from_float(float x); + +template <> +__device__ __forceinline__ __nv_bfloat16 from_float<__nv_bfloat16>(float x) { + return __float2bfloat16(x); +} + +template <> +__device__ __forceinline__ __half from_float<__half>(float x) { + return __float2half(x); +} + +template <> +__device__ __forceinline__ cutlass::bfloat16_t from_float(float x) { + __nv_bfloat16 tmp = __float2bfloat16(x); + return reinterpret_cast(tmp); +} + +template <> +__device__ __forceinline__ cutlass::half_t from_float(float x) { + __half tmp = __float2half(x); + return reinterpret_cast(tmp); +} + +// ═══════════════════════════════════════════════════════════════════ +// Element → float conversion (template specialization per type) +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ float to_float(Element v); + +template <> +__device__ __forceinline__ float to_float<__nv_bfloat16>(__nv_bfloat16 v) { + return __bfloat162float(v); +} + +template <> +__device__ __forceinline__ float to_float<__half>(__half v) { + return __half2float(v); +} + +template <> +__device__ __forceinline__ float to_float(cutlass::bfloat16_t v) { + return __bfloat162float(reinterpret_cast(v)); +} + +template <> +__device__ __forceinline__ float to_float(cutlass::half_t v) { + return __half2float(reinterpret_cast(v)); +} + +// ═══════════════════════════════════════════════════════════════════ +// Cached global load (__ldg) for 2-byte element types +// ═══════════════════════════════════════════════════════════════════ +// +// CUDA provides __ldg for __nv_bfloat16 / __half directly; the cutlass +// wrappers (cutlass::bfloat16_t / cutlass::half_t) require a reinterpret +// since the byte layout is identical. + +template +__device__ __forceinline__ T ldg_elem(const T* p); + +template <> +__device__ __forceinline__ __nv_bfloat16 ldg_elem<__nv_bfloat16>(const __nv_bfloat16* p) { + return __ldg(p); +} + +template <> +__device__ __forceinline__ __half ldg_elem<__half>(const __half* p) { + return __ldg(p); +} + +template <> +__device__ __forceinline__ cutlass::bfloat16_t ldg_elem(const cutlass::bfloat16_t* p) { + __nv_bfloat16 raw = __ldg(reinterpret_cast(p)); + return reinterpret_cast(raw); +} + +template <> +__device__ __forceinline__ cutlass::half_t ldg_elem(const cutlass::half_t* p) { + __half raw = __ldg(reinterpret_cast(p)); + return reinterpret_cast(raw); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp.cuh index fcfbabc3..5beb806e 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp.cuh @@ -1,417 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// Fused MLP: Y = (Z @ A^T) where Z = SiLU(B @ X) * (C @ X) -// ═══════════════════════════════════════════════════════════════════ -// -// Pipelines are constructed once in the prologue (requires all threads). -// Pipeline states are persistent across tiles — cooperative 2-WG + -// fused pipes reduce the state register footprint to two states per -// phase (producer + consumer), so hoisting them out of the tile loop -// still fits comfortably in registers. -// -// The tile loop is driven by a TileIterator (local or remote). -// The intermediate Z buffer is double-buffered: z_m alternates -// between blockIdx.x*ZBufferSlots + (iter % ZBufferSlots). -// -// WG0 warps 1-3 do NOT participate — reserved for NVSHMEM comm. -// -// ═══════════════════════════════════════════════════════════════════ +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -#include "mlp1_fused.cuh" -#include "mlp2_fused.cuh" -#include "tile_iterator.cuh" -#include "cta_barrier.cuh" // target-based CtaCounterBarrier - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// Fused shared memory -// ═══════════════════════════════════════════════════════════════════ - -template -struct MlpFusedSmem { - union { - Mlp1FusedSmem mlp1; - Mlp2FusedSmem mlp2; - }; - - // Pipeline barrier storage — persistent across phases. - // Both phases use a single fused pipe (one mbarrier per stage). - typename Mlp1MainloopPipelineFor::SharedStorage p1_pipe; - typename Mlp2MainloopPipelineFor::SharedStorage p2_pipe; - alignas(16) uint32_t tmem_base; -}; - -static constexpr int kMlpBarrierId = 14; - -// Target-based monotonic barrier for the FWD MLP warps. Count is Compute-aware: -// Hopper 288 (w0 + w4-11); Blackwell 320 (w0 + w3-11, warp 3 = MLP1 UMMA). Must -// match the in-fn kMlpBarrierThreads — both fire kMlpBarrierId. See cta_barrier.cuh -// for why the modulo variant deadlocks under CTA skew. -template -using MlpFwdCtaBarrierT = CtaCounterBarrier<(Compute == 100 ? 10 : 9) * 32, kMlpBarrierId>; - -// ═══════════════════════════════════════════════════════════════════ -// MlpDims — invariant MLP dimensions, passed as __grid_constant__ -// ═══════════════════════════════════════════════════════════════════ - -struct MlpDims { - int hidden_dim; - int intermediate_dim; - int num_experts; // local experts represented by the weight tensors - int total_n_rows_1; // physical rows spanned by the strided B/C views - int total_n_rows_2; // experts_per_pe * hidden_dim - int expert_n_stride_1; // B/C expert stride in TileN1 units - int num_n_tiles_1; // intermediate_dim / TileN1 - int num_n_tiles_2; // ceildiv(hidden_dim/TileN2, NSub) — per-WG outer steps - int num_k_tiles_1; // hidden_dim / TileK1 - int num_k_tiles_2; // intermediate_dim / TileK2 - int y_buf_m_tiles; // max_total_slots / TileM — row-tile extent of the - // symmetric local_output (y_buf) descriptor, used as - // num_y_m_tiles for the direct-to-peer Y store. - int* phase_counter; -}; - -// ═══════════════════════════════════════════════════════════════════ -// mlp_fused_fwd — generic over TileIterator -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void mlp_fused_fwd( - MlpFusedSmem& smem, - TileIter& iter, - // MLP1 TMA descriptors - TmaLoadX const& tma_load_x, - TmaLoadW1 const& tma_load_b, - TmaLoadW1 const& tma_load_c, - TmaStoreZ const& tma_store_z, - // MLP2 TMA descriptors - TmaLoadZ const& tma_load_z, - TmaLoadW2 const& tma_load_a, - TmaStoreY const& tma_store_y, - // Invariant MLP dims — const ref into __grid_constant__ .param space - const MlpDims& dims, - // Per-phase params (change between local and remote) - int num_tokens, - int num_m_tiles, - // Persistent per-CTA x_barrier — constructed once in moe_fused_fwd - // and threaded through here so target accumulates across all 3 - // mlp_fused_fwd calls (A1 / B / A2). Standalone callers (no - // remote phase) still pass their own barrier instance. - MlpFwdCtaBarrierT & x_barrier, - int local_expert_start = 0, - // Logical column / grid_x for the flat-grid fused launch. Defaults - // (-1) fall back to blockIdx.x / gridDim.x for standalone 2-D-grid - // callers (mlp.cu bench kernel). - int col = -1, - int grid_x = -1, - // N-split identity for the GEMM producers/consumers. Defaults (-1) - // fall back to blockIdx.y / gridDim.y (standalone 2-D-grid mlp.cu). - int split_idx = -1, - int num_splits = -1, - // Direct-to-peer Y store: array of per-peer TMA descriptors over - // nvshmem_ptr(local_output, peer) (indexed by MlpTileInfo::peer_rank). - // For is_local tiles MLP2 stores Y straight into the destination's - // symmetric local_output instead of dst_staging (tma_store_y). nullptr → - // staging path only (standalone mlp.cu callers have no peer descriptors). - const TmaStoreY* y_peer_desc = nullptr) { - - using Pipeline1 = Mlp1MainloopPipelineFor; - using PipeState1 = typename Traits1::PipelineState; - using Pipeline2 = Mlp2MainloopPipelineFor; - using PipeState2 = typename Traits2::PipelineState; - - int warp_id = threadIdx.x / Traits1::WarpSize; - - // MLP-warp NamedBarrier arrival count — must equal MlpFwdCtaBarrierT's - // template count (both fire kMlpBarrierId). Hopper: w0+w4-11 (288). Blackwell: - // w0+w3-11 (320, warp 3 = MLP1 UMMA). - constexpr int kMlpBarrierThreads = (Compute == 100 ? 10 : 9) * 32; - - cute::prefetch_tma_descriptor(tma_load_x.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_b.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_c.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_z.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_z.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_a.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_y.get_tma_descriptor()); - - // ── Prologue: construct all pipelines once ────────── - // Pipeline constructors set per-thread role (Producer / Consumer / - // NonParticipant) so warps 1-3 can safely skip the MLP body below. - auto p1_pipe = [&]() { - if constexpr (Compute == 100) - return mlp1_make_pipe_umma(smem.p1_pipe); - else - return mlp1_make_pipe(smem.p1_pipe); - }(); - auto p2_pipe = [&]() { - if constexpr (Compute == 100) - return mlp2_make_pipe_umma(smem.p2_pipe); - else - return mlp2_make_pipe(smem.p2_pipe); - }(); - - if constexpr (Compute == 90) { - if (warp_id == 3) { - return; - } - } - - // Sync MLP warps only (excludes comm warps 1-3 which may be in - // nvshmem_comm_main when called from moe_fused_fwd). - cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); - - cute::TMEM::Allocator1Sm tmem_alloc{}; - if constexpr (Compute == 100) { - constexpr int kMlp1Cols = Traits1::AccStages * (2 * Traits1::TileN); - constexpr int kMlp2Cols = Traits2::AccStages * Traits2::TileN; - constexpr int kTmemColumns = (kMlp1Cols > kMlp2Cols) - ? kMlp1Cols : kMlp2Cols; - if (warp_id == 3) { - tmem_alloc.allocate(kTmemColumns, &smem.tmem_base); - __syncwarp(); - } - cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); - } - - // Z buffer: [grid_x, ZBufferSlots, SubTiles, GemmTileM, intermediate_dim]. - // Indexed in GemmTileM-row tiles: SubTiles Z regions per comm-tile slot so - // mlp1 can stage all sub-tiles before mlp2 drains them. z_rows is unchanged - // (CommTileM == SubTiles·GemmTileM); only the indexing splits. - int z_col = (col >= 0) ? col : (int)blockIdx.x; - int z_grid_x = (grid_x >= 0) ? grid_x : (int)gridDim.x; - int num_z_m_tiles = z_grid_x * ZBufferSlots * SubTiles; - int z_base = z_col * ZBufferSlots * SubTiles; - // N-split identity for the GEMM producers/consumers (flat-grid: NOT - // blockIdx.y / gridDim.y, which are 0 / 1). Each CTA computes only its - // stripe of N-tiles (n = n_split, n_split+n_count, ...). - int n_split = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_count = (num_splits >= 0) ? num_splits : (int)gridDim.y; - int tile_iter = 0; - - // ── Persistent pipeline states (hoisted out of the tile loop) ── - // With single fused pipes + cooperative consumer, we have only - // 2 states per phase. They persist across tiles since k-tile - // counts are fixed — the barrier phase bit cycles back to 0 - // after each tile's k-loop completes. - PipeState1 p1_state = (warp_id == 0) ? - cutlass::make_producer_start_state() : PipeState1{}; - PipeState1 p1_cons_state; - PipeState2 p2_state = (warp_id == 0) ? - cutlass::make_producer_start_state() : PipeState2{}; - PipeState2 p2_cons_state; - - // x_barrier is passed in by reference (constructed once per CTA - // in moe_fused_fwd) so target advances monotonically across - // the A1 / B / A2 phases that all share dims.phase_counter. - - // ── Tile loop ─────────────────────────────────────── - while (iter.has_next()) { - // Barrier: ensure previous phase 2 is done before reusing smem union. - cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); - iter.acquire_src(); - auto tile = iter.next(); - - int expert = tile.expert - local_expert_start; - int first_sub = 0; - int last_sub = SubTiles; - if constexpr (SubTiles > 1) { - first_sub = tile.m_subtile; - last_sub = first_sub + 1; - } - // Base Z region for this comm tile. SubTiles GemmTileM-row regions per - // slot, indexed z_m_base + sub. tile.y_m is the comm slot index in - // CommTileM(=GemmTileM·SubTiles) units; the GemmTileM-unit row coordinate - // of sub-tile `sub` in that slot is tile.y_m·SubTiles + sub. - int z_m_base = z_base + (tile_iter % ZBufferSlots) * SubTiles; - - // Phase 1: for each GemmTileM sub-tile, Z[z_m_base+sub] = SiLU(B@X)·(C@X). - // One comm semaphore (acquire_src above) covers all SubTiles; the GEMM - // just steps finer. SubTiles==1 collapses to the original single tile. - for (int sub = first_sub; sub < last_sub; ++sub) { - int x_mt = tile.y_m * SubTiles + sub; - int z_m = z_m_base + sub; - if constexpr (Compute == 100) - if (warp_id == 3 && threadIdx.x % Traits1::WarpSize == 0) - smem.mlp1.tmem_base = smem.tmem_base; - if (warp_id == 0) { - if constexpr (Compute == 100) { - mlp1_fused_producer( - p1_pipe, p1_state, - smem.mlp1, tma_load_x, tma_load_b, tma_load_c, - x_mt, expert * dims.expert_n_stride_1, - num_tokens, dims.hidden_dim, dims.intermediate_dim, - dims.num_experts, dims.total_n_rows_1, - dims.num_n_tiles_1, dims.num_k_tiles_1, n_split, n_count); - } else { - mlp1_fused_producer( - p1_pipe, p1_state, - smem.mlp1, tma_load_x, tma_load_b, tma_load_c, - x_mt, expert, - num_tokens, dims.hidden_dim, dims.intermediate_dim, - dims.num_experts, dims.total_n_rows_1, - dims.num_n_tiles_1, dims.num_k_tiles_1, n_split, n_count); - } - } - - // Fused MoE CTA warp map: warp 0 = TMA, warps 1..3 = NVSHMEM comm - // (warp 3 = put), warps 4..11 = MLP GEMM consumers. Warp 3 is a comm - // warp here, so the fused MLP1 keeps the ORIGINAL warp-4 UMMA design - // (single-warp UMMA + warps 4..11 epilogue). The warp-3 dual-WG - // rewrite is used ONLY by the isolated MLP1 kernel (no comm warps), - // so it neither frees warp 3 from comm nor changes the shared MLP - // barrier count — MLP2 and the rest of the fused pipeline are - // untouched. Compile-time flag (no register cost). Enabling warp-3 - // here would require the comm-warp + barrier migration (pipeline.md - // §5) — intentionally out of scope. - - if (warp_id >= 3 && warp_id <= 11) { - if constexpr (Compute == 100) { - mlp1_fused_consumer( - p1_pipe, p1_cons_state, - smem.mlp1, tma_store_z, z_m, dims.intermediate_dim, - num_z_m_tiles, dims.num_n_tiles_1, dims.num_k_tiles_1, - n_split, n_count); - } else { - mlp1_fused_consumer( - p1_pipe, p1_cons_state, - smem.mlp1, tma_store_z, z_m, dims.intermediate_dim, - num_z_m_tiles, dims.num_n_tiles_1, dims.num_k_tiles_1, - n_split, n_count); - } - } - - if constexpr (Compute == 90) { - if (sub + 1 < SubTiles) - cutlass::arch::NamedBarrier::sync( - kMlpBarrierThreads, kMlpBarrierId); - } - } - - // One phase barrier covers the full communication tile. The persistent - // TMA/UMMA pipeline states sequence its GEMM sub-tiles internally. - cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); - - // Release src staging (remote: free slot; local: no-op). - // Pass threadIdx.x directly so only thread 0 (TMA producer warp - // leader) has lane==0 and fires the atomicAdd exactly once per - // CTA — StagePipe expects NumConsumers = NC releases per slot. - iter.release_src(threadIdx.x); - - // Cross-CTA sync: wait for ALL y-CTAs sharing this blockIdx.x - // to finish phase 1 for this M-tile (all sub-tiles' Z regions). - x_barrier.wait(); - - // Acquire dst staging (remote: wait for put slot; local: no-op). - iter.acquire_dst(); - - // Y-store target select. Same-host (is_local) tiles store Y DIRECTLY - // into the destination peer's symmetric local_output via its per-peer - // descriptor at the token-tile row; all other tiles store to the - // staging ring (tma_store_y) at the slot row. mlp1's X-load above still - // uses tile.y_m (the staging slot) regardless — input path unchanged. - // Tile counts are in CommTileM units → scale by SubTiles for GemmTileM boxes. - bool y_direct = (y_peer_desc != nullptr && tile.is_local); - const TmaStoreY& y_desc = y_direct ? y_peer_desc[tile.peer_rank] : tma_store_y; - int y_base = y_direct ? tile.y_store_m : tile.y_m; - int y_ntiles = (y_direct ? dims.y_buf_m_tiles : num_m_tiles) * SubTiles; - // Phase 2: for each sub-tile, Y[y_base·SubTiles+sub] = Z[z_m_base+sub] @ A^T. - // Cooperative-M-split: both consumer WGs cooperate on the same (m,n) tile. - for (int sub = first_sub; sub < last_sub; ++sub) { - int z_m = z_m_base + sub; - int y_coord = y_base * SubTiles + sub; - constexpr int kM64PerGemmTile = Traits2::TileM / 64; - int valid_m64_subtiles = max( - 0, min(kM64PerGemmTile, - tile.valid_m64_subtiles - sub * kM64PerGemmTile)); - if constexpr (Compute == 100) - if (warp_id == 3 && threadIdx.x % Traits2::WarpSize == 0) - smem.mlp2.tmem_base = smem.tmem_base; - if (warp_id == 0) { - if constexpr (Compute == 100) { - mlp2_fused_producer( - p2_pipe, p2_state, - smem.mlp2, tma_load_z, tma_load_a, - z_m, expert * dims.num_n_tiles_2, - num_z_m_tiles * Traits1::TileM, dims.hidden_dim, - dims.intermediate_dim, dims.num_experts, dims.total_n_rows_2, - dims.num_n_tiles_2, dims.num_k_tiles_2, n_split, n_count); - } else { - mlp2_fused_producer( - p2_pipe, p2_state, - smem.mlp2, tma_load_z, tma_load_a, - z_m, expert, - num_z_m_tiles * Traits1::TileM, dims.hidden_dim, - dims.intermediate_dim, dims.num_experts, dims.total_n_rows_2, - dims.num_n_tiles_2, dims.num_k_tiles_2, n_split, n_count); - } - } - - if (warp_id >= 3 && warp_id <= 11) { - if constexpr (Compute == 100) { - mlp2_fused_consumer( - p2_pipe, p2_cons_state, - smem.mlp2, y_desc, - y_coord, dims.hidden_dim, - y_ntiles, dims.num_n_tiles_2, dims.num_k_tiles_2, - n_split, n_count, valid_m64_subtiles); - } else { - mlp2_fused_consumer( - p2_pipe, p2_cons_state, - smem.mlp2, y_desc, - y_coord, dims.hidden_dim, - y_ntiles, dims.num_n_tiles_2, dims.num_k_tiles_2, - n_split, n_count); - } - } - - if constexpr (Compute == 90) { - if (sub + 1 < SubTiles) - cutlass::arch::NamedBarrier::sync( - kMlpBarrierThreads, kMlpBarrierId); - } - } - - // Release dst staging (remote: cross-CTA barrier + signal; local: no-op). - // Pass threadIdx.x directly — exactly one atomicAdd per CTA - // (threadIdx.x == 0), matching NumProducers = NC on dst_pipe. - // Device-scope: the system-scope flush of Y to HBM for the NIC put is - // done by RemoteIter::release_dst below (all threads). This is just the - // intra-CTA ordering before the MLP-warp barrier. - __threadfence(); - // As in phase 1, synchronize once after all GEMM sub-tiles. - cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); - iter.release_dst(threadIdx.x); - - tile_iter++; - } - - if constexpr (Compute == 100) { - constexpr int kMlp1Cols = Traits1::AccStages * (2 * Traits1::TileN); - constexpr int kMlp2Cols = Traits2::AccStages * Traits2::TileN; - constexpr int kTmemColumns = (kMlp1Cols > kMlp2Cols) - ? kMlp1Cols : kMlp2Cols; - cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); - if (warp_id == 3) { - tmem_alloc.release_allocation_lock(); - tmem_alloc.free(smem.tmem_base, kTmemColumns); - } - } -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp_sm90.cuh" +#else +#include "mlp_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp1_fused.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused.cuh index 0580b62b..098a3ab9 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp1_fused.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused.cuh @@ -1,934 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Phase 1 (single-tile, fused): Z = SiLU(B @ X) * (C @ X) -// ═══════════════════════════════════════════════════════════════════ -// -// Single (TileM, TileN) producer/consumer building blocks for callers -// that drive the outer M-tile loop themselves: the standalone MLP1 -// benchmark (mlp1.cu), the fused MLP kernel (mlp.cuh), the activation -// variants (mlp1_fused_act.cuh), and the backward path (mlp_bwd.cuh). -// -// Design: -// - Cooperative 2-WG consumer; split axis selected by TileM: -// TileM=128 → Layout<_2,_1,_1>, split along M (each WG = 64×TileN) -// TileM=64 → Layout<_1,_2,_1>, split along N (each WG = 64×TileN/2) -// - Two accumulators per WG (acc_B = U = B@X, acc_C = V = C@X) -// - Per-thread SiLU(U)*V fusion (no cross-WG sync) -// - Per-WG store_buf with EpiChunkN-round epilogue -// - K_PIPE_MMAS = 1 (CUTLASS-canonical, see -// sm90_mma_tma_gmma_ss_warpspecialized.hpp:264) -// - Single fused X+W1+W2 TMA pipe (1 mbarrier per stage, 3 TMA copies) -// -// CTA layout (384 threads = 12 warps): -// WG0 warp 0 : Producer (TMA loads for X / B / C) -// WG0 warps 1-2 : Idle / reserved for fused-kernel communication -// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 -// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 -// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 -// -// Per `expert_ids`: one int per TileM-row m-block. -// -// Mlp1FusedSmem omits the pipeline storage so callers can fuse it with -// other phases' smem (see mlp.cuh). The benchmark wrapper (mlp1.cu) is -// responsible for attaching MainloopPipeline::SharedStorage and running -// the outer for(m) loop over m-tiles. -// -// ═══════════════════════════════════════════════════════════════════ - -#include "models.cuh" -#include "math.cuh" -#include "tmem_load_op.cuh" - -#include -#include -#include -#include -#include - -// SM100 (Blackwell / UMMA) support for the Compute=100 consumer -// specialization. These headers are header-only and safe to include under -// an sm_90a build — the tcgen05 PTX they emit is guarded by -// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on -// __CUDA_ARCH__ >= 1000) ever instantiates it. -#include -#include - -#include -#include -#include // PipelineTmaUmmaAsync, PipelineUmmaAsync - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// Traits -// ═══════════════════════════════════════════════════════════════════ - -template < - typename Element_, - int TileM_ = 128, - int TileN_ = 128, - int TileK_ = 64, - int Stages_ = 4, - int EpiChunkN_ = 32, - int ActStoreTileM_ = 64 -> -struct Mlp1Traits { - using Element = Element_; - using ElementAccum = float; - - static constexpr int TileM = TileM_; - static constexpr int TileN = TileN_; - static constexpr int TileK = TileK_; - static constexpr int Stages = Stages_; - static constexpr int EpiChunkN = EpiChunkN_; - static constexpr int ActStoreTileM = ActStoreTileM_; - - static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); - - // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate - // along M (each owns 64 rows × TileN). With TileM=64 the WGs - // cooperate along N (each owns 64 rows × TileN/2). AtomTileM is the - // per-WG M extent — always 64. - static constexpr bool kMSplit = (TileM == 128); - static constexpr int AtomTileM = 64; - static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); - - static_assert(TileN <= 256, "TileN must fit TMA box limit"); - static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); - static_assert(ActStoreTileM == 32 || ActStoreTileM == 64, - "activation TMA stores support M32 or M64"); - static_assert(TileM % ActStoreTileM == 0, - "activation store TileM must divide the GEMM TileM"); - static_assert(kMSplit || (TileN % 2 == 0), - "TileN must be even when TileM=64 (cooperative N-split halves)"); - static constexpr int NumEpiRounds = WgTileN / EpiChunkN; - - // Cooperative 2-WG layout: M-split for TileM=128, N-split for TileM=64. - static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); - using GmmaAtom = typename GmmaSelector::Atom; - using WgLayout = cute::conditional_t< - kMSplit, - Layout>, - Layout>>; - using TiledMma = TiledMMA, WgLayout>; - - using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom; - - using SmemLayoutX_1 = decltype(tile_to_shape(SmemLayoutAtom{}, - Shape, Int>{})); - using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{}, - Shape, Int>{})); - - using SmemLayoutX = decltype(tile_to_shape(SmemLayoutAtom{}, - Shape, Int, Int>{})); - using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{}, - Shape, Int, Int>{})); - - // Per-WG store buffer slot: (AtomTileM, EpiChunkN). One slot per WG. - // Epi writes the WG's m-strip in EpiChunkN-wide rounds. - using SmemLayoutStoreSlot = Layout< - Shape, Int>, - Stride, _1>>; - using SmemLayoutStore_1 = SmemLayoutStoreSlot; - using SmemLayoutActStoreSlot = Layout< - Shape, Int>, - Stride, _1>>; - - using MainloopPipeline = cutlass::PipelineTmaAsync; - using PipelineState = cutlass::PipelineState; - - // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop - // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the - // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the TMEM - // accumulator from the MMA warp to the epilogue warps. See - // Mlp1FusedConsumerImpl<100>. The Blackwell launcher must build the - // mainloop pipe with num_consumers=1 (one umma_arrive per stage). - using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< - Stages, cute::Shape, - cute::Shape>; - // Double-buffered SM100 TMEM accumulators: each stage owns U and V, each - // TileN columns wide. This lets MMA(n+1) use the alternate stage while the - // epilogue drains MMA(n)'s TMEM stage. The launcher allocates - // AccStages·2·TileN columns; the dedicated epilogue warpgroup (warps 8..11) - // drains the previous stage while the epilogue-free MMA warp runs ahead. - static constexpr int AccStages = 2; - static_assert(AccStages * (2 * TileN) <= 512, - "SM100 MLP1 needs 2 accumulator stages x (U,V) to fit in 512 TMEM columns"); - static_assert(TileN <= 128, - "SM100 double-buffered MLP1 accumulator requires TileN <= 128"); - using AccumulatorPipeline = cutlass::PipelineUmmaAsync; - - static constexpr int TmaTransBytesX = - static_cast(size(SmemLayoutX_1{}) * sizeof(Element)); - static constexpr int TmaTransBytesW = - static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); - // Fused pipeline: X + W1 + W2 share one mbarrier per stage. - static constexpr int TmaTransBytes = TmaTransBytesX + 2 * TmaTransBytesW; - - static constexpr int WarpSize = 32; - static constexpr int WarpGroupSize = 128; - static constexpr int NumConsumers = 2; - static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; - static constexpr int NumThreads = 384; -}; - -template -using Mlp1MainloopPipelineFor = cute::conditional_t< - Compute == 100, - typename Traits::MainloopPipelineUmma, - typename Traits::MainloopPipeline>; - -// ═══════════════════════════════════════════════════════════════════ -// Expert weight pointers (device-side) -// ═══════════════════════════════════════════════════════════════════ - -template -struct ExpertWeights1 { - const Element* B; // [intermediate_dim, hidden_dim] - const Element* C; // [intermediate_dim, hidden_dim] -}; - -// ═══════════════════════════════════════════════════════════════════ -// Fused MLP1 shared memory — no pipeline storage (managed externally) -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp1FusedSmem { - using Element = typename Traits::Element; - - static constexpr int smem_X_size = cosize_v; - static constexpr int smem_W_size = cosize_v; - static constexpr int smem_store_size = cosize_v; - - alignas(128) Element smem_X[smem_X_size]; - alignas(128) Element smem_W1[smem_W_size]; - alignas(128) Element smem_W2[smem_W_size]; - // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. - alignas(128) Element store_buf[2 * smem_store_size]; - - // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM - // base address, plus the accumulator pipeline's barrier storage (UMMA→ - // epilogue handoff). The Hopper (Compute=90) path never touches these; - // SM100 launchers allocate/free TMEM outside the consumer and publish the - // base here. - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; - - CUTE_DEVICE Element* X_data() { return &smem_X[0]; } - CUTE_DEVICE Element* W1_data() { return &smem_W1[0]; } - CUTE_DEVICE Element* W2_data() { return &smem_W2[0]; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// Pipeline construction helper — single fused X+W1+W2 pipeline -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ auto mlp1_make_pipe( - typename Traits::MainloopPipeline::SharedStorage& storage) { - - using Pipeline = typename Traits::MainloopPipeline; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 4 && warp_id <= 11); - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = Traits::ConsumerThreads; - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// SM100 (Compute=100) mainloop pipe. Mirrors mlp1_make_pipe but builds the -// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues -// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer -// is released by a single umma_arrive per stage (the MMA warp), not per -// consumer thread. Params/PipelineState/SharedStorage alias the Hopper -// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer -// (mlp1_fused_producer) and the Compute=100 consumer drive it unchanged. -template -__device__ __forceinline__ auto mlp1_make_pipe_umma( - typename Traits::MainloopPipelineUmma::SharedStorage& storage) { - - using Pipeline = typename Traits::MainloopPipelineUmma; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA consumer - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// ═══════════════════════════════════════════════════════════════════ -// Fused producer — single (m), all (n) under blockIdx.y stride. -// Single fused X + W1 + W2 TMA pipe per k-step. -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void mlp1_fused_producer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp1FusedSmem& smem, - TmaLoadX const& tma_load_x, - TmaLoadW const& tma_load_b, - TmaLoadW const& tma_load_c, - int m, - int expert_or_n_offset, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int num_experts, - int total_n_rows, - int num_n_tiles, - int num_k_tiles, - // N-split identity (flat-grid launch). Defaults (-1) fall back to - // blockIdx.y / gridDim.y for standalone 2-D-grid callers (mlp1.cu). - int split_idx = -1, - int num_splits = -1) { - - auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); - auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); - - // int64_t cast: total_n_rows × hidden_dim can exceed INT_MAX at - // production shapes (epp · I · D > 2^31), and CUTE's layout math - // otherwise wraps to a negative byte offset → kernel IMA. - auto mX = tma_load_x.get_tma_tensor(make_shape( - static_cast(num_tokens), - static_cast(hidden_dim))); - auto mB = [&]() { - if constexpr (Expert3D) { - return tma_load_b.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(num_experts))); - } else { - return tma_load_b.get_tma_tensor(make_shape( - static_cast(total_n_rows), - static_cast(hidden_dim))); - } - }(); - auto mC = [&]() { - if constexpr (Expert3D) { - return tma_load_c.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(num_experts))); - } else { - return tma_load_c.get_tma_tensor(make_shape( - static_cast(total_n_rows), - static_cast(hidden_dim))); - } - }(); - - auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); - auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); - auto cta_tma_c = tma_load_c.get_slice(Int<0>{}); - - auto tXsX = cta_tma_x.partition_D(sX); - auto tW1sW1 = cta_tma_b.partition_D(sW1); - auto tW2sW2 = cta_tma_c.partition_D(sW2); - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - auto gB = [&]() { - if constexpr (Expert3D) { - return local_tile(mB, - make_tile(Int{}, Int{}), - make_coord(n, _, expert_or_n_offset)); - } else { - return local_tile(mB, - make_tile(Int{}, Int{}), - make_coord(expert_or_n_offset + n, _)); - } - }(); - auto gC = [&]() { - if constexpr (Expert3D) { - return local_tile(mC, - make_tile(Int{}, Int{}), - make_coord(n, _, expert_or_n_offset)); - } else { - return local_tile(mC, - make_tile(Int{}, Int{}), - make_coord(expert_or_n_offset + n, _)); - } - }(); - auto tBgB = cta_tma_b.partition_S(gB); - auto tCgC = cta_tma_c.partition_S(gC); - - CUTE_UNROLL - for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { - auto gX = local_tile(mX, - make_tile(Int{}, Int{}), - make_coord(m + m_sub, _)); - auto tXgX = cta_tma_x.partition_S(gX); - - for (int k = 0; k < num_k_tiles; ++k) { - pipe.producer_acquire(state); - if (threadIdx.x == 0) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_x.with(*bar, 0), - tXgX(_, _, _, k), tXsX(_, _, _, state.index())); - copy(tma_load_b.with(*bar, 0), - tBgB(_, _, _, k), tW1sW1(_, _, _, state.index())); - copy(tma_load_c.with(*bar, 0), - tCgC(_, _, _, k), tW2sW2(_, _, _, state.index())); - } - ++state; - } - } - } -} - -// ═══════════════════════════════════════════════════════════════════ -// Fused consumer — architecture-specialized on `int Compute` -// ═══════════════════════════════════════════════════════════════════ -// -// Primary template is undefined; one full specialization per supported -// compute capability provides a `run(...)` member (a function-template on -// Traits/Pipeline/TmaStoreZ, since function templates can't be partially -// specialized). The free function `mlp1_fused_consumer` below forwards to -// the right specialization — call sites stay in the existing style. -// -// Compute=90 → Hopper / WGMMA (cooperative 2-WG, dual register acc) -// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulators) -// -// z_m is the M-tile index into the output Z buffer. - -template -struct Mlp1FusedConsumerImpl; - -// ─────────────────────────────────────────────────────────────────── -// Compute=90 — Hopper. The two WGs split the (TileM, TileN) output along -// M: WG1 writes rows [0, AtomTileM) of tile z_m, WG2 writes [AtomTileM, -// TileM). (Verbatim from the original mlp1_fused_consumer.) -// ─────────────────────────────────────────────────────────────────── - -template <> -struct Mlp1FusedConsumerImpl<90> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp1FusedSmem& smem, - TmaStoreZ const& tma_store_z, - int z_m, // M-tile index into Z buffer - int intermediate_dim, - int num_z_m_tiles, // number of M-tile slots in Z buffer - int num_n_tiles, - int num_k_tiles, - // N-split identity (flat-grid launch); -1 → blockIdx.y / gridDim.y. - int split_idx, - int num_splits) { - - static_assert(MSubTiles == 1, - "Hopper MLP1 consumer processes one GEMM M-tile per invocation"); - using Element = typename Traits::Element; - typename Traits::TiledMma tiled_mma; - int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 - int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 - auto thr_mma = tiled_mma.get_slice(tid_in_mma); - - auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); - auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); - - auto tCsX = thr_mma.partition_A(sX); - auto tCsW1 = thr_mma.partition_B(sW1); - auto tCsW2 = thr_mma.partition_B(sW2); - - // Two accumulators per WG, each covering the WG's (AtomTileM, TileN) m-strip. - auto acc_B = partition_fragment_C(tiled_mma, - Shape, Int>{}); - auto acc_C = partition_fragment_C(tiled_mma, - Shape, Int>{}); - - // Identity coords for per-thread acc → store_buf mapping. - auto cC = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCcC = thr_mma.partition_C(cC); - - const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 - const int my_barrier_id = 1 + my_wg; // 1 or 2 - const bool is_my_wg_leader = (tid_in_wg == 0); - - constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; - Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - // int64_t cast: num_z_m_tiles · TileM · intermediate_dim exceeds - // INT_MAX once grid_x · ZBuf · TileM · I > 2^31 (e.g. ZBuf=16 - // at I=16384 wraps and IMAs). - auto mZ = tma_store_z.get_tma_tensor( - make_shape(static_cast(num_z_m_tiles) * Traits::TileM, - static_cast(intermediate_dim))); - auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - bool store_in_flight = false; - constexpr int K_PIPE_MMAS = 1; // CUTLASS-matching - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - - clear(acc_B); - clear(acc_C); - - auto state_release = state; - int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; - - // Prologue - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc_B); - warpgroup_fence_operand(acc_C); - warpgroup_arrive(); - gemm(tiled_mma, tCsX(_, _, _, state.index()), - tCsW1(_, _, _, state.index()), acc_B); - gemm(tiled_mma, tCsX(_, _, _, state.index()), - tCsW2(_, _, _, state.index()), acc_C); - warpgroup_commit_batch(); - ++state; - } - // Steady state - for (int k = prologue_count; k < num_k_tiles; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc_B); - warpgroup_fence_operand(acc_C); - warpgroup_arrive(); - gemm(tiled_mma, tCsX(_, _, _, state.index()), - tCsW1(_, _, _, state.index()), acc_B); - gemm(tiled_mma, tCsX(_, _, _, state.index()), - tCsW2(_, _, _, state.index()), acc_C); - warpgroup_commit_batch(); - - warpgroup_wait(); - warpgroup_fence_operand(acc_B); - warpgroup_fence_operand(acc_C); - pipe.consumer_release(state_release); - ++state; - ++state_release; - } - // Drain - warpgroup_wait<0>(); - warpgroup_fence_operand(acc_B); - warpgroup_fence_operand(acc_C); - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_release(state_release); - ++state_release; - } - - // SiLU fusion: acc_B = SiLU(U) * V (per-thread, no sync). - CUTE_UNROLL - for (int i = 0; i < size(acc_B); ++i) { - float u = acc_B(i); - float v = acc_C(i); - acc_B(i) = fast_silu(u) * v; - } - - // ── Epilogue ────────────────────────────────────── - // Each WG writes its (AtomTileM=64) × EpiChunkN slot to store_buf - // and issues TMA stores. NumEpiRounds rounds per WG. The per-thread - // acc element → (m_local, n_local) mapping depends on the split: - // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N - // TileM=64 (N-split): WG_w owns full M, cols [w*WgTileN, ...) - CUTE_UNROLL - for (int r = 0; r < Traits::NumEpiRounds; ++r) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - int chunk_start = r * Traits::EpiChunkN; - CUTE_UNROLL - for (int i = 0; i < size(acc_B); ++i) { - auto coord = tCcC(i); - int m_loc = get<0>(coord); - int n_loc = get<1>(coord); - int m_local, n_local; - if constexpr (Traits::kMSplit) { - m_local = m_loc - my_wg * Traits::AtomTileM; - n_local = n_loc; - } else { - m_local = m_loc; - n_local = n_loc - my_wg * Traits::WgTileN; - } - if (n_local >= chunk_start && - n_local < chunk_start + Traits::EpiChunkN) { - int chunk_n = n_local - chunk_start; - sStore(m_local, chunk_n) = static_cast(acc_B(i)); - } - } - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - if (is_my_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx, n_tile_idx; - if constexpr (Traits::kMSplit) { - // TileM=128: row tile space = 2 tiles per logical M - // (AtomTileM=64). WG_w writes (2*z_m + w, n*Rounds + r). - m_tile_idx = 2 * z_m + my_wg; - n_tile_idx = n * Traits::NumEpiRounds + r; - } else { - // TileM=64: 1 row tile per logical M. WG_w handles its - // N-half: chunk index = n*(TileN/EpiChunkN) - // + w*NumEpiRounds + r. - m_tile_idx = z_m; - n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) - + my_wg * Traits::NumEpiRounds + r; - } - auto gZ = local_tile(mZ, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_z, cta_tma_z.partition_S(sStore), - cta_tma_z.partition_D(gZ)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - if (store_in_flight) - cute::tma_store_wait<0>(); -} -}; // Mlp1FusedConsumerImpl<90> - -// ─────────────────────────────────────────────────────────────────── -// Compute=100 — Blackwell / UMMA. -// -// No cooperative-warpgroup machinery (UMMA has no warpgroup concept): -// - One 1SM UMMA atom with M=TileM covers the whole tile; accumulators -// U = X·B and V = X·C live in TMEM (no register/M-split). -// - Warp 3 issues both UMMAs and stays epilogue-free so it can run ahead to -// the alternate TMEM accumulator stage (AccStages=2). -// - Epilogue: aligned WGs in warps 4..11 split TileN, read TMEM→regs, -// compute SiLU while the previous TMA store drains, and store 64-row -// (AtomTileM) tiles via the existing reg→SMEM→TMA path. -// -// TileM is the template parameter as on Hopper (64 or 128 — both native -// 1SM UMMA M sizes). For TileM=128 the epilogue stores MSub=2 row-tiles. -// ─────────────────────────────────────────────────────────────────── - -template <> -struct Mlp1FusedConsumerImpl<100> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp1FusedSmem& smem, - TmaStoreZ const& tma_store_z, - int z_m, - int intermediate_dim, - int num_z_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using Element = typename Traits::Element; - constexpr int TileM = Traits::TileM; - constexpr int TileN = Traits::TileN; - constexpr int EpiChunkN = Traits::EpiChunkN; - constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile - static_assert(TileN % 2 == 0, - "Blackwell consumer splits TileN across the two epilogue warpgroups"); - constexpr int WgN = TileN / 2; // per-WG N width - static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); - constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per epilogue WG - constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) - - // ── Thread identity (Blackwell warp-specialised layout) ────────────── - // Warp 3 is the dedicated, epilogue-free UMMA producer (moved out of the - // consumer warpgroups so it can run AccStages tiles ahead). Warps 4..11 are - // the two aligned epilogue warpgroups (WG0 = warps 4..7, WG1 = warps 8..11); - // each drains one half of TileN from the current accumulator stage while the - // MMA warp fills the alternate TMEM stage. In the fused MoE CTA warp 0 = TMA - // and warps 1..2 = NVSHMEM comm; warp 3 is the freed comm warp repurposed to - // issue UMMA (see PR #19 + mlp.cuh). Total: 12 warps / 3 hardware warpgroups. - const int warp_id = threadIdx.x / Traits::WarpSize; - const bool is_mma_warp = (warp_id == 3); // dedicated UMMA warp - const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); // two aligned WGs - // Epilogue-local identity — valid only for warps 4..11. Warp 3 never enters - // the epilogue block, so its (garbage) tid_in_epi/wg/tid_wg are unused. - const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 - const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; // 0|1 - const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; // 0..127 - const int tmem_copy_tid = tid_wg; - const bool is_wg_leader = is_epilogue && (tid_wg == 0); - const int wg_barrier_id = 1 + wg; // 1 (WG0) or 2 (WG1) - // Warps 3..11 (MMA + both epilogue WGs) rendezvous for the acc-pipe init and - // the launcher-safe TMEM free. ConsumerThreads (256) covers only warps 4..11. - constexpr int kEpilogueThreads = Traits::ConsumerThreads; - constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; - static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); - static_assert(kEpilogueThreads == 8 * Traits::WarpSize); - static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); - - // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, K-major SS ── - auto tiled_mma = make_tiled_mma( - SM100_MMA_F16BF16_SS{}); - auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 - - auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); - auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); - - // ── Two TMEM accumulators: U = X·B, V = X·C ── - auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCgC = cta_mma.partition_C(cAccFull); - auto tCtAccU = cta_mma.make_fragment_C(tCgC); - auto tCtAccV = cta_mma.make_fragment_C(tCgC); - - // TMEM is allocated by the outer fused/standalone launcher once per CTA. - - // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers - // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the - // alternate stage while both epilogue WGs drain MMA(n). producer_commit - // issues the UMMA-gated "accumulator ready" arrival; one elected epilogue - // thread releases the stage (after a 256-thread rendezvous) so the MMA - // warp may reuse it AccStages tiles later. ── - using AccPipe = typename Traits::AccumulatorPipeline; - typename AccPipe::Params acc_params; - acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer - : AccPipe::ThreadCategory::Consumer; - acc_params.producer_arv_count = 1; // one umma_arrive per commit (warp 3) - acc_params.consumer_arv_count = 1; // one elected releaser per stage - acc_params.initializing_warp = 4; // warp 4 inits the acc barriers - AccPipe acc_pipe(smem.acc_pipe, acc_params, - cute::Shape{}); - auto acc_prod_state = cutlass::make_producer_start_state(); - typename AccPipe::PipelineState acc_cons_state; - - // Warps 3..11 rendezvous so the acc-pipe barriers (init'd by warp 4) are - // visible before the MMA warp (3) and both epilogue WGs (4..11) use them. - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); - const uint32_t tmem_base = smem.tmem_base; - tCtAccU.data() = tmem_base; - tCtAccV.data() = tmem_base + uint32_t(TileN); - - // ── Per-WG store slot (AtomTileM × EpiChunkN) — WG0 uses store_buf[0..S), - // WG1 uses [S..2S). Same smem layout as Hopper. ── - constexpr int store_slot_elems = AtomTileM * EpiChunkN; - Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - auto mZ = tma_store_z.get_tma_tensor( - make_shape(static_cast(num_z_m_tiles) * TileM, - static_cast(intermediate_dim))); - auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); - - // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local - // tid). Each epi tile is the full TileM rows × EpiChunkN cols; the - // 128-thread tcgen05.ld maps one thread per TMEM datapath row. - // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) - // modes, which is what make_tmem_copy's cotiled builder requires. ── - auto epi_tile = make_tile(Int{}, Int{}); - // Extract the flat (M,N) view from the UMMA C-fragment (MMA,MMA_M,MMA_N) - // before tiling — matches CUTLASS's accumulators(make_coord(_,_),_0,_0). - auto accU_mn = tCtAccU(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) - auto tAccU_epi = flat_divide(accU_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) - auto t2r = make_tmem_copy(TmemLoadOp{}, tAccU_epi(_, _, _0{}, _0{})); - auto thr_t2r = t2r.get_slice(tmem_copy_tid); - // Within-chunk (m,n) coords for the reg → smem-slot scatter. - auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); - auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) - // Register fragments sized from the DEST (partition_D) shape — per-thread - // (T2R,T2R_M,T2R_N). Sizing from partition_S(tmem) would wrongly include the - // warp-collective datapath-lane dim (the TMEM_LOAD atom distributes those - // lanes into per-thread registers internally). Matches CUTLASS's - // make_tensor(shape(tTR_sD)). - auto tTR_rU = make_tensor(shape(tTR_cChunk)); // f32 regs (U, reused) - auto tTR_rV = make_tensor(shape(tTR_cChunk)); // f32 regs (V) - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - bool store_in_flight = false; - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - CUTE_UNROLL - for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { - - // ── Mainloop (MMA warp = warp 3 only): bracket the k-loop with the - // accumulator pipeline. producer_acquire selects the free TMEM stage - // (AccStages=2 → MMA(n+1) fills the alternate stage while both epilogue - // WGs drain MMA(n)). consumer_release issues the UMMA-gated smem-buffer - // arrival on the mainloop pipe (no hand-rolled mbarrier). ── - if (is_mma_warp) { - acc_pipe.producer_acquire(acc_prod_state); // selected TMEM stage is free - int acc_stage = acc_prod_state.index(); - uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); - tCtAccU.data() = stage_base; - tCtAccV.data() = stage_base + uint32_t(TileN); - for (int k = 0; k < num_k_tiles; ++k) { - pipe.consumer_wait(state); - auto tCsX = cta_mma.partition_A(sX (_, _, state.index())); - auto tCsW1 = cta_mma.partition_B(sW1(_, _, state.index())); - auto tCsW2 = cta_mma.partition_B(sW2(_, _, state.index())); - auto tCrX = cta_mma.make_fragment_A(tCsX); - auto tCrW1 = cta_mma.make_fragment_B(tCsW1); - auto tCrW2 = cta_mma.make_fragment_B(tCsW2); - CUTE_UNROLL - for (int kb = 0; kb < size<2>(tCrX); ++kb) { - // First k-block of this n-tile clears TMEM; the rest accumulate. - tiled_mma.accumulate_ = (k == 0 && kb == 0) - ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; - gemm(tiled_mma, tCrX(_, _, kb), tCrW1(_, _, kb), tCtAccU); - gemm(tiled_mma, tCrX(_, _, kb), tCrW2(_, _, kb), tCtAccV); - } - pipe.consumer_release(state); // UMMA-gated smem-buffer release - ++state; - } - acc_pipe.producer_commit(acc_prod_state); // signal U/V ready (umma_arrive) - ++acc_prod_state; - } - - // ── Epilogue (two aligned WGs, warps 4..11): wait for the accumulator - // stage, then THIS WG drains its half of TileN (chunks - // [wg·NChunksHalf, +NChunksHalf)) while the MMA warp fills the - // alternate stage. ── - if (is_epilogue) { - acc_pipe.consumer_wait(acc_cons_state); - int acc_stage = acc_cons_state.index(); - uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); - tCtAccU.data() = stage_base; - tCtAccV.data() = stage_base + uint32_t(TileN); - auto accU_mn_stage = tCtAccU(make_coord(_, _), _0{}, _0{}); - auto accV_mn_stage = tCtAccV(make_coord(_, _), _0{}, _0{}); - auto tAccU_epi_stage = flat_divide(accU_mn_stage, epi_tile); - auto tAccV_epi_stage = flat_divide(accV_mn_stage, epi_tile); - auto tTR_tAccU = thr_t2r.partition_S(tAccU_epi_stage); - auto tTR_tAccV = thr_t2r.partition_S(tAccV_epi_stage); - - CUTE_UNROLL - for (int r = 0; r < NChunksHalf; ++r) { - int chunk = wg * NChunksHalf + r; // absolute n-chunk index - - // TMEM → registers (this chunk, full TileM rows). - copy(t2r, tTR_tAccU(_, _, _, _0{}, chunk), tTR_rU); - copy(t2r, tTR_tAccV(_, _, _, _0{}, chunk), tTR_rV); - cutlass::arch::fence_view_async_tmem_load(); - - // SiLU(U)·V fused in place (overwrites the U regs) so it overlaps - // the prior in-flight store. The cast to Element happens at the - // smem write below — no extra register set. - CUTE_UNROLL - for (int i = 0; i < size(tTR_rU); ++i) - tTR_rU(i) = fast_silu(tTR_rU(i)) * tTR_rV(i); - - // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). - CUTE_UNROLL - for (int ms = 0; ms < MSub; ++ms) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - CUTE_UNROLL - for (int i = 0; i < size(tTR_rU); ++i) { - int m_row = get<0>(tTR_cChunk(i)); // 0..TileM - int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN - if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) - sStore(m_row - ms * AtomTileM, n_col) = - static_cast(tTR_rU(i)); - } - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - - if (is_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx = MSub * (z_m + m_sub) + ms; - int n_tile_idx = n * (TileN / EpiChunkN) + chunk; - auto gZ = local_tile(mZ, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_z, cta_tma_z.partition_S(sStore), - cta_tma_z.partition_D(gZ)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - - // Both epilogue WGs done reading this TMEM stage → rendezvous, then - // one elected thread releases the stage so the MMA warp may reuse it - // when the AccStages accumulator pipeline wraps. - cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); - if (tid_in_epi == 0) - acc_pipe.consumer_release(acc_cons_state); - ++acc_cons_state; - } - } - } - if (is_epilogue && store_in_flight) - cute::tma_store_wait<0>(); - - // Final rendezvous of the MMA warp (3) + both epilogue WGs (4..11) so the - // outer launcher can safely free TMEM after this consumer returns. - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); -#else - // The Compute=100 specialization is never dispatched on non-SM100 targets - // (host picks Compute=90 there). Trap if it is ever reached. - __trap(); +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 #endif -} -}; // Mlp1FusedConsumerImpl<100> -// ─────────────────────────────────────────────────────────────────── -// Forwarder — keeps the existing call style. `Compute` defaults to 90 so -// current call sites (mlp1_fused_consumer(...)) are unchanged; -// pass mlp1_fused_consumer(...) to select the Blackwell path. -// ─────────────────────────────────────────────────────────────────── - -template -__device__ __forceinline__ void mlp1_fused_consumer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp1FusedSmem& smem, - TmaStoreZ const& tma_store_z, - int z_m, - int intermediate_dim, - int num_z_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx = -1, - int num_splits = -1) { - Mlp1FusedConsumerImpl::template run< - Traits, MSubTiles, Pipeline, TmaStoreZ>( - pipe, state, smem, tma_store_z, z_m, intermediate_dim, - num_z_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits); -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp1_fused_sm90.cuh" +#else +#include "mlp1_fused_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act.cuh index 8cbcec81..b826bb5f 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act.cuh @@ -1,703 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Phase 1 (fused tile + silu-backward coefficients): -// U = B @ X -// V = C @ X -// V' = silu(U) (coefficient for dV = dZ * V' later) -// U' = V * silu'(U) (coefficient for dU = dZ * U' later) -// Z = V' * V = silu(U) * V (forward output, consumed by mlp3) -// ═══════════════════════════════════════════════════════════════════ -// -// Cooperative 2-WG consumer producing three outputs (U', V', Z) per -// (TileM, TileN) tile. Split axis selected by TileM (kMSplit pattern, -// same as mlp1_fused.cuh / mlp2_t_fused.cuh / mlp5_fused.cuh): -// TileM=128 (M-split, Layout<_2,_1,_1>): each WG owns AtomTileM=64 -// rows × full TileN. WG_w writes rows [w*64, (w+1)*64). -// TileM=64 (N-split, Layout<_1,_2,_1>): each WG owns full M (=64) -// × WgTileN = TileN/2. WG_w writes cols [w*WgTileN, (w+1)*WgTileN). -// Dual acc_B/acc_C from a single fused X+W1+W2 TMA pipeline, -// EpiChunkN-wide per-WG epilogue rounds — each round writes three -// smem buffers and issues three TMA stores. -// -// Smem layout (per Mlp1Traits): three per-WG store slots. WG_w uses -// indices [w·S, (w+1)·S) within each buf_{u,v,z}, where S = -// AtomTileM·EpiChunkN. Stages=3 default (vs Stages=4 in fwd mlp1) to -// fit the extra two output buffers within the 228 KiB SM smem cap. -// -// Single-tile API: caller passes (m, expert_n_offset) and a single -// fused X+B+C pipeline. Used by the standalone mlp1_act launcher. - -#include "mlp1_fused.cuh" - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// Shared memory — three per-WG store buffers (U', V', Z) -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp1FusedActSmem { - using Element = typename Traits::Element; - - static constexpr int smem_X_size = cosize_v; - static constexpr int smem_W_size = cosize_v; - static constexpr int smem_store_size = - cosize_v; - - alignas(128) Element smem_X[smem_X_size]; - alignas(128) Element smem_W1[smem_W_size]; - alignas(128) Element smem_W2[smem_W_size]; - // 3 outputs × 2 WGs of per-WG (AtomTileM, EpiChunkN) slots. - alignas(128) Element store_buf_u[2 * smem_store_size]; - alignas(128) Element store_buf_v[2 * smem_store_size]; - alignas(128) Element store_buf_z[2 * smem_store_size]; - - // SM100 (Compute=100) only: tcgen05.alloc landing slot + accumulator - // pipeline storage (see Mlp1FusedSmem). Untouched by the Hopper path. - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; - - CUTE_DEVICE Element* X_data() { return &smem_X[0]; } - CUTE_DEVICE Element* W1_data() { return &smem_W1[0]; } - CUTE_DEVICE Element* W2_data() { return &smem_W2[0]; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// Producer — single (m), all (n) under split_idx/num_splits stride. -// Single fused X + W1 + W2 TMA pipe per k-step. -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void mlp1_fused_act_producer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp1FusedActSmem& smem, - TmaLoadX const& tma_load_x, - TmaLoadW const& tma_load_b, - TmaLoadW const& tma_load_c, - int m, - int expert_or_n_offset, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int num_experts, - int total_n_rows, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { - - auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); - auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); - - // int64_t cast for production shapes where E·I·D > 2^31. - auto mX = tma_load_x.get_tma_tensor(make_shape( - static_cast(num_tokens), - static_cast(hidden_dim))); - auto mB = [&]() { - if constexpr (Expert3D) { - return tma_load_b.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(num_experts))); - } else { - return tma_load_b.get_tma_tensor(make_shape( - static_cast(total_n_rows), - static_cast(hidden_dim))); - } - }(); - auto mC = [&]() { - if constexpr (Expert3D) { - return tma_load_c.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(num_experts))); - } else { - return tma_load_c.get_tma_tensor(make_shape( - static_cast(total_n_rows), - static_cast(hidden_dim))); - } - }(); - - auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); - auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); - auto cta_tma_c = tma_load_c.get_slice(Int<0>{}); - - auto tXsX = cta_tma_x.partition_D(sX); - auto tW1sW1 = cta_tma_b.partition_D(sW1); - auto tW2sW2 = cta_tma_c.partition_D(sW2); - - auto gX = local_tile(mX, - make_tile(Int{}, Int{}), - make_coord(m, _)); - auto tXgX = cta_tma_x.partition_S(gX); - - for (int n = split_idx; n < num_n_tiles; n += num_splits) { - auto gB = [&]() { - if constexpr (Expert3D) { - return local_tile(mB, - make_tile(Int{}, Int{}), - make_coord(n, _, expert_or_n_offset)); - } else { - return local_tile(mB, - make_tile(Int{}, Int{}), - make_coord(expert_or_n_offset + n, _)); - } - }(); - auto gC = [&]() { - if constexpr (Expert3D) { - return local_tile(mC, - make_tile(Int{}, Int{}), - make_coord(n, _, expert_or_n_offset)); - } else { - return local_tile(mC, - make_tile(Int{}, Int{}), - make_coord(expert_or_n_offset + n, _)); - } - }(); - auto tBgB = cta_tma_b.partition_S(gB); - auto tCgC = cta_tma_c.partition_S(gC); - - for (int k = 0; k < num_k_tiles; ++k) { - pipe.producer_acquire(state); - if (threadIdx.x == 0) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_x.with(*bar, 0), - tXgX(_, _, _, k), tXsX(_, _, _, state.index())); - copy(tma_load_b.with(*bar, 0), - tBgB(_, _, _, k), tW1sW1(_, _, _, state.index())); - copy(tma_load_c.with(*bar, 0), - tCgC(_, _, _, k), tW2sW2(_, _, _, state.index())); - } - ++state; - } - } -} - -// ═══════════════════════════════════════════════════════════════════ -// Consumer — architecture-specialized on `int Compute` (see -// Mlp1FusedConsumerImpl in mlp1_fused.cuh for the dispatch rationale). -// Compute=90 → Hopper / WGMMA (cooperative 2-WG) -// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulators) -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp1FusedActConsumerImpl; - -// ─────────────────────────────────────────────────────────────────── -// Compute=90 — Hopper. z_m is the M-tile index into the output buffers. -// Split axis selected by TileM (kMSplit pattern): -// TileM=128 (M-split): WG_w writes rows [w*AtomTileM, (w+1)*AtomTileM) -// TileM=64 (N-split): WG_w writes cols [w*WgTileN, (w+1)*WgTileN) -// Each WG issues NumEpiRounds × 3 TMA stores per N-tile (U', V', Z, -// each at (AtomTileM, EpiChunkN)). (Verbatim from the original consumer.) -// ─────────────────────────────────────────────────────────────────── - -template <> -struct Mlp1FusedActConsumerImpl<90> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp1FusedActSmem& smem, - TmaStore const& tma_store_du_coef, - TmaStore const& tma_store_dv_coef, - TmaStore const& tma_store_z, - int z_m, - int intermediate_dim, - int num_z_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { - - using Element = typename Traits::Element; - typename Traits::TiledMma tiled_mma; - int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 - int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 - auto thr_mma = tiled_mma.get_slice(tid_in_mma); - - auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); - auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); - - auto tCsX = thr_mma.partition_A(sX); - auto tCsW1 = thr_mma.partition_B(sW1); - auto tCsW2 = thr_mma.partition_B(sW2); - - auto acc_B = partition_fragment_C(tiled_mma, - Shape, Int>{}); - auto acc_C = partition_fragment_C(tiled_mma, - Shape, Int>{}); - - // Identity coords for per-thread acc → store_buf mapping. - auto cC = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCcC = thr_mma.partition_C(cC); - - const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 - const int my_barrier_id = 1 + my_wg; // 1 or 2 - const bool is_my_wg_leader = (tid_in_wg == 0); - - static_assert(Traits::AtomTileM % Traits::ActStoreTileM == 0, - "activation store TileM must divide the Hopper warpgroup strip"); - constexpr int StoreSlicesPerWg = - Traits::AtomTileM / Traits::ActStoreTileM; - constexpr int StoreTilesPerM = - Traits::TileM / Traits::ActStoreTileM; - constexpr int store_slot_elems = - Traits::ActStoreTileM * Traits::EpiChunkN; - Element* my_store_u = smem.store_buf_u + my_wg * store_slot_elems; - Element* my_store_v = smem.store_buf_v + my_wg * store_slot_elems; - Element* my_store_z = smem.store_buf_z + my_wg * store_slot_elems; - auto sStoreU = make_tensor(make_smem_ptr(my_store_u), - typename Traits::SmemLayoutActStoreSlot{}); - auto sStoreV = make_tensor(make_smem_ptr(my_store_v), - typename Traits::SmemLayoutActStoreSlot{}); - auto sStoreZ = make_tensor(make_smem_ptr(my_store_z), - typename Traits::SmemLayoutActStoreSlot{}); - - // int64_t cast: num_z_m_tiles · TileM · intermediate_dim can exceed INT_MAX. - auto out_shape = make_shape( - static_cast(num_z_m_tiles) * Traits::TileM, - static_cast(intermediate_dim)); - auto mU = tma_store_du_coef.get_tma_tensor(out_shape); - auto mV = tma_store_dv_coef.get_tma_tensor(out_shape); - auto mZ = tma_store_z.get_tma_tensor(out_shape); - auto cta_tma_u = tma_store_du_coef.get_slice(Int<0>{}); - auto cta_tma_v = tma_store_dv_coef.get_slice(Int<0>{}); - auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); - - bool store_in_flight = false; - constexpr int K_PIPE_MMAS = 1; // CUTLASS-matching - - for (int n = split_idx; n < num_n_tiles; n += num_splits) { - - clear(acc_B); - clear(acc_C); - - auto state_release = state; - int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; - - // Prologue - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc_B); - warpgroup_fence_operand(acc_C); - warpgroup_arrive(); - gemm(tiled_mma, tCsX(_, _, _, state.index()), - tCsW1(_, _, _, state.index()), acc_B); - gemm(tiled_mma, tCsX(_, _, _, state.index()), - tCsW2(_, _, _, state.index()), acc_C); - warpgroup_commit_batch(); - ++state; - } - // Steady state - for (int k = prologue_count; k < num_k_tiles; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc_B); - warpgroup_fence_operand(acc_C); - warpgroup_arrive(); - gemm(tiled_mma, tCsX(_, _, _, state.index()), - tCsW1(_, _, _, state.index()), acc_B); - gemm(tiled_mma, tCsX(_, _, _, state.index()), - tCsW2(_, _, _, state.index()), acc_C); - warpgroup_commit_batch(); - - warpgroup_wait(); - warpgroup_fence_operand(acc_B); - warpgroup_fence_operand(acc_C); - pipe.consumer_release(state_release); - ++state; - ++state_release; - } - // Drain - warpgroup_wait<0>(); - warpgroup_fence_operand(acc_B); - warpgroup_fence_operand(acc_C); - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_release(state_release); - ++state_release; - } - - // ── Epilogue: 3-store per-WG rounds ───────────── - // Per round each WG: write U', V', Z into 3 separate per-WG smem - // slots, then issue 3 TMA stores. No cross-WG sync — each WG - // owns its strip independently. - // TileM=128 (M-split): WG_w owns rows [w*AtomTileM, (w+1)*AtomTileM), - // full N. Store at row_tile = 2*z_m + my_wg, col_tile = n*NER + r. - // TileM=64 (N-split): WG_w owns full M (=64), cols - // [w*WgTileN, (w+1)*WgTileN). Store at row_tile = z_m, - // col_tile = n * (TileN/EpiChunkN) + my_wg*NER + r. - CUTE_UNROLL - for (int r = 0; r < Traits::NumEpiRounds; ++r) { - CUTE_UNROLL - for (int ms = 0; ms < StoreSlicesPerWg; ++ms) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync( - Traits::WarpGroupSize, my_barrier_id); - - int chunk_start = r * Traits::EpiChunkN; - CUTE_UNROLL - for (int i = 0; i < size(acc_B); ++i) { - auto coord = tCcC(i); - int m_loc = get<0>(coord); - int n_loc = get<1>(coord); - int m_local, n_local; - if constexpr (Traits::kMSplit) { - m_local = m_loc - my_wg * Traits::AtomTileM; - n_local = n_loc; - } else { - m_local = m_loc; - n_local = n_loc - my_wg * Traits::WgTileN; - } - if (m_local >= ms * Traits::ActStoreTileM - && m_local < (ms + 1) * Traits::ActStoreTileM - && n_local >= chunk_start - && n_local < chunk_start + Traits::EpiChunkN) { - int store_m = m_local - ms * Traits::ActStoreTileM; - int chunk_n = n_local - chunk_start; - float u = acc_B(i); - float v = acc_C(i); - float sig = fast_sigmoid(u); - float vprime = u * sig; - float sil_d = sig + vprime * (1.0f - sig); - sStoreU(store_m, chunk_n) = - static_cast(v * sil_d); - sStoreV(store_m, chunk_n) = - static_cast(vprime); - sStoreZ(store_m, chunk_n) = - static_cast(vprime * v); - } - } - - cutlass::arch::NamedBarrier::sync( - Traits::WarpGroupSize, my_barrier_id); - - if (is_my_wg_leader) { - cute::tma_store_fence(); - int row_tile, col_tile; - if constexpr (Traits::kMSplit) { - row_tile = StoreTilesPerM * z_m - + my_wg * StoreSlicesPerWg + ms; - col_tile = n * Traits::NumEpiRounds + r; - } else { - row_tile = StoreTilesPerM * z_m + ms; - col_tile = n * (Traits::TileN / Traits::EpiChunkN) - + my_wg * Traits::NumEpiRounds + r; - } - auto tile = make_tile( - Int{}, - Int{}); - auto gU = local_tile( - mU, tile, make_coord(row_tile, col_tile)); - auto gV = local_tile( - mV, tile, make_coord(row_tile, col_tile)); - auto gZ = local_tile( - mZ, tile, make_coord(row_tile, col_tile)); - copy(tma_store_du_coef, cta_tma_u.partition_S(sStoreU), - cta_tma_u.partition_D(gU)); - cute::tma_store_arrive(); - copy(tma_store_dv_coef, cta_tma_v.partition_S(sStoreV), - cta_tma_v.partition_D(gV)); - cute::tma_store_arrive(); - copy(tma_store_z, cta_tma_z.partition_S(sStoreZ), - cta_tma_z.partition_D(gZ)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - } - if (store_in_flight) - cute::tma_store_wait<0>(); -} -}; // Mlp1FusedActConsumerImpl<90> - -// ─────────────────────────────────────────────────────────────────── -// Compute=100 — Blackwell / UMMA. Same structure as the fused -// Mlp1FusedConsumerImpl<100>, but the epilogue derives the three act -// outputs from (U, V) and issues three TMA stores per 64-row tile: -// V' = silu(U), U' = V·silu'(U), Z = V'·V. -// ─────────────────────────────────────────────────────────────────── - -template <> -struct Mlp1FusedActConsumerImpl<100> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp1FusedActSmem& smem, - TmaStore const& tma_store_du_coef, - TmaStore const& tma_store_dv_coef, - TmaStore const& tma_store_z, - int z_m, - int intermediate_dim, - int num_z_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using Element = typename Traits::Element; - constexpr int TileM = Traits::TileM; - constexpr int TileN = Traits::TileN; - constexpr int EpiChunkN = Traits::EpiChunkN; - constexpr int StoreTileM = Traits::ActStoreTileM; - static_assert(TileN % 2 == 0, - "Blackwell consumer splits TileN across the two consumer warpgroups"); - constexpr int WgN = TileN / 2; - static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); - constexpr int NChunksHalf = WgN / EpiChunkN; - constexpr int MSub = TileM / StoreTileM; - - const int warp_id = threadIdx.x / Traits::WarpSize; - const bool is_mma_warp = (warp_id == 3); - const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); - const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 - const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; - const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; - const int tmem_copy_tid = tid_wg; - const bool is_wg_leader = is_epilogue && (tid_wg == 0); - const int wg_barrier_id = 1 + wg; - constexpr int kEpilogueThreads = Traits::ConsumerThreads; - constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; - static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); - static_assert(kEpilogueThreads == 8 * Traits::WarpSize); - static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); - - auto tiled_mma = make_tiled_mma( - SM100_MMA_F16BF16_SS{}); - auto cta_mma = tiled_mma.get_slice(0); - - auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); - auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); - - auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCgC = cta_mma.partition_C(cAccFull); - auto tCtAccU = cta_mma.make_fragment_C(tCgC); // U = X·B - auto tCtAccV = cta_mma.make_fragment_C(tCgC); // V = X·C - - // TMEM is allocated by the outer fused/standalone launcher once per CTA. - - // Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers - // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the - // alternate stage while both epilogue WGs drain MMA(n). - using AccPipe = typename Traits::AccumulatorPipeline; - typename AccPipe::Params acc_params; - acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer - : AccPipe::ThreadCategory::Consumer; - acc_params.producer_arv_count = 1; - acc_params.consumer_arv_count = 1; - acc_params.initializing_warp = 4; - AccPipe acc_pipe(smem.acc_pipe, acc_params, - cute::Shape{}); - auto acc_prod_state = cutlass::make_producer_start_state(); - typename AccPipe::PipelineState acc_cons_state; - - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); - const uint32_t tmem_base = smem.tmem_base; - tCtAccU.data() = tmem_base; - tCtAccV.data() = tmem_base + uint32_t(TileN); - - // Three per-WG store slots (U', V', Z) — same smem as Hopper. - constexpr int store_slot_elems = StoreTileM * EpiChunkN; - auto sStoreU = make_tensor(make_smem_ptr(smem.store_buf_u + wg * store_slot_elems), - typename Traits::SmemLayoutActStoreSlot{}); - auto sStoreV = make_tensor(make_smem_ptr(smem.store_buf_v + wg * store_slot_elems), - typename Traits::SmemLayoutActStoreSlot{}); - auto sStoreZ = make_tensor(make_smem_ptr(smem.store_buf_z + wg * store_slot_elems), - typename Traits::SmemLayoutActStoreSlot{}); - - auto out_shape = make_shape( - static_cast(num_z_m_tiles) * TileM, - static_cast(intermediate_dim)); - auto mU = tma_store_du_coef.get_tma_tensor(out_shape); - auto mV = tma_store_dv_coef.get_tma_tensor(out_shape); - auto mZ = tma_store_z.get_tma_tensor(out_shape); - auto cta_tma_u = tma_store_du_coef.get_slice(Int<0>{}); - auto cta_tma_v = tma_store_dv_coef.get_slice(Int<0>{}); - auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); - - auto epi_tile = make_tile(Int{}, Int{}); - // Flat (M,N) view of the UMMA C-fragment before tiling (see mlp1_fused.cuh). - auto accU_mn = tCtAccU(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) - auto accV_mn = tCtAccV(make_coord(_, _), _0{}, _0{}); - auto tAccU_epi = flat_divide(accU_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) - auto tAccV_epi = flat_divide(accV_mn, epi_tile); - auto t2r = make_tmem_copy(TmemLoadOp{}, tAccU_epi(_, _, _0{}, _0{})); - auto thr_t2r = t2r.get_slice(tmem_copy_tid); - auto tTR_tAccU = thr_t2r.partition_S(tAccU_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) - auto tTR_tAccV = thr_t2r.partition_S(tAccV_epi); - auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); - auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) - // Register fragments sized from the DEST (partition_D) per-thread shape, not - // partition_S(tmem) (warp-collective). See mlp1_fused.cuh for the rationale. - auto tTR_rU = make_tensor(shape(tTR_cChunk)); // f32 regs (U) - auto tTR_rV = make_tensor(shape(tTR_cChunk)); // f32 regs (V) - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - bool store_in_flight = false; - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - - // Mainloop (warp 3 only), bracketed by the accumulator pipeline. - if (is_mma_warp) { - acc_pipe.producer_acquire(acc_prod_state); - int acc_stage = acc_prod_state.index(); - uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); - tCtAccU.data() = stage_base; - tCtAccV.data() = stage_base + uint32_t(TileN); - for (int k = 0; k < num_k_tiles; ++k) { - pipe.consumer_wait(state); - auto tCsX = cta_mma.partition_A(sX (_, _, state.index())); - auto tCsW1 = cta_mma.partition_B(sW1(_, _, state.index())); - auto tCsW2 = cta_mma.partition_B(sW2(_, _, state.index())); - auto tCrX = cta_mma.make_fragment_A(tCsX); - auto tCrW1 = cta_mma.make_fragment_B(tCsW1); - auto tCrW2 = cta_mma.make_fragment_B(tCsW2); - CUTE_UNROLL - for (int kb = 0; kb < size<2>(tCrX); ++kb) { - tiled_mma.accumulate_ = (k == 0 && kb == 0) - ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; - gemm(tiled_mma, tCrX(_, _, kb), tCrW1(_, _, kb), tCtAccU); - gemm(tiled_mma, tCrX(_, _, kb), tCrW2(_, _, kb), tCtAccV); - } - pipe.consumer_release(state); - ++state; - } - acc_pipe.producer_commit(acc_prod_state); - ++acc_prod_state; - } - - if (is_epilogue) { - acc_pipe.consumer_wait(acc_cons_state); - int acc_stage = acc_cons_state.index(); - uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); - tCtAccU.data() = stage_base; - tCtAccV.data() = stage_base + uint32_t(TileN); - auto accU_mn_stage = tCtAccU(make_coord(_, _), _0{}, _0{}); - auto accV_mn_stage = tCtAccV(make_coord(_, _), _0{}, _0{}); - auto tAccU_epi_stage = flat_divide(accU_mn_stage, epi_tile); - auto tAccV_epi_stage = flat_divide(accV_mn_stage, epi_tile); - auto tTR_tAccU_stage = thr_t2r.partition_S(tAccU_epi_stage); - auto tTR_tAccV_stage = thr_t2r.partition_S(tAccV_epi_stage); - - CUTE_UNROLL - for (int r = 0; r < NChunksHalf; ++r) { - int chunk = wg * NChunksHalf + r; - - // U and V stay in registers (tTR_rU, tTR_rV); the three act - // outputs are derived and cast to Element inline at the smem write - // below — no extra register sets. Each reg's row falls in exactly - // one MSub slice, so this evaluates each element once. - copy(t2r, tTR_tAccU_stage(_, _, _, _0{}, chunk), tTR_rU); - copy(t2r, tTR_tAccV_stage(_, _, _, _0{}, chunk), tTR_rV); - cutlass::arch::fence_view_async_tmem_load(); - - CUTE_UNROLL - for (int ms = 0; ms < MSub; ++ms) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - CUTE_UNROLL - for (int i = 0; i < size(tTR_rU); ++i) { - int m_row = get<0>(tTR_cChunk(i)); - int n_col = get<1>(tTR_cChunk(i)); - if (m_row >= ms * StoreTileM && m_row < (ms + 1) * StoreTileM) { - int m_loc = m_row - ms * StoreTileM; - float u = tTR_rU(i); - float v = tTR_rV(i); - float sig = fast_sigmoid(u); - float vprime = u * sig; // silu(U) - float sil_d = sig + vprime * (1.0f - sig); // silu'(U) - float uprime = v * sil_d; - Element uprime_e = static_cast(uprime); - Element vprime_e = static_cast(vprime); - Element z_e = static_cast(vprime * v); - sStoreU(m_loc, n_col) = uprime_e; // V·silu'(U) - sStoreV(m_loc, n_col) = vprime_e; // silu(U) - sStoreZ(m_loc, n_col) = z_e; // V'·V - } - } - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - - if (is_wg_leader) { - cute::tma_store_fence(); - int row_tile = MSub * z_m + ms; - int col_tile = n * (TileN / EpiChunkN) + chunk; - auto gU = local_tile(mU, - make_tile(Int{}, Int{}), - make_coord(row_tile, col_tile)); - auto gV = local_tile(mV, - make_tile(Int{}, Int{}), - make_coord(row_tile, col_tile)); - auto gZ = local_tile(mZ, - make_tile(Int{}, Int{}), - make_coord(row_tile, col_tile)); - copy(tma_store_du_coef, cta_tma_u.partition_S(sStoreU), - cta_tma_u.partition_D(gU)); - cute::tma_store_arrive(); - copy(tma_store_dv_coef, cta_tma_v.partition_S(sStoreV), - cta_tma_v.partition_D(gV)); - cute::tma_store_arrive(); - copy(tma_store_z, cta_tma_z.partition_S(sStoreZ), - cta_tma_z.partition_D(gZ)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - - // Epilogue done reading TMEM; release the accumulator for the next - // n-tile's MMA (one elected consumer thread arrives). - cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); - if (tid_in_epi == 0) - acc_pipe.consumer_release(acc_cons_state); - ++acc_cons_state; - } - } - if (is_epilogue && store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); -#else - __trap(); +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 #endif -} -}; // Mlp1FusedActConsumerImpl<100> -// Forwarder — `Compute` defaults to 90 (existing call sites unchanged); -// pass for the Blackwell path. -template -__device__ __forceinline__ void mlp1_fused_act_consumer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp1FusedActSmem& smem, - TmaStore const& tma_store_du_coef, - TmaStore const& tma_store_dv_coef, - TmaStore const& tma_store_z, - int z_m, - int intermediate_dim, - int num_z_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { - Mlp1FusedActConsumerImpl::template run( - pipe, state, smem, tma_store_du_coef, tma_store_dv_coef, tma_store_z, - z_m, intermediate_dim, num_z_m_tiles, num_n_tiles, num_k_tiles, - split_idx, num_splits); -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp1_fused_act_sm90.cuh" +#else +#include "mlp1_fused_act_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm100.cuh new file mode 100644 index 00000000..bb61a215 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm100.cuh @@ -0,0 +1,703 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 1 (fused tile + silu-backward coefficients): +// U = B @ X +// V = C @ X +// V' = silu(U) (coefficient for dV = dZ * V' later) +// U' = V * silu'(U) (coefficient for dU = dZ * U' later) +// Z = V' * V = silu(U) * V (forward output, consumed by mlp3) +// ═══════════════════════════════════════════════════════════════════ +// +// Cooperative 2-WG consumer producing three outputs (U', V', Z) per +// (TileM, TileN) tile. Split axis selected by TileM (kMSplit pattern, +// same as mlp1_fused.cuh / mlp2_t_fused.cuh / mlp5_fused.cuh): +// TileM=128 (M-split, Layout<_2,_1,_1>): each WG owns AtomTileM=64 +// rows × full TileN. WG_w writes rows [w*64, (w+1)*64). +// TileM=64 (N-split, Layout<_1,_2,_1>): each WG owns full M (=64) +// × WgTileN = TileN/2. WG_w writes cols [w*WgTileN, (w+1)*WgTileN). +// Dual acc_B/acc_C from a single fused X+W1+W2 TMA pipeline, +// EpiChunkN-wide per-WG epilogue rounds — each round writes three +// smem buffers and issues three TMA stores. +// +// Smem layout (per Mlp1Traits): three per-WG store slots. WG_w uses +// indices [w·S, (w+1)·S) within each buf_{u,v,z}, where S = +// AtomTileM·EpiChunkN. Stages=3 default (vs Stages=4 in fwd mlp1) to +// fit the extra two output buffers within the 228 KiB SM smem cap. +// +// Single-tile API: caller passes (m, expert_n_offset) and a single +// fused X+B+C pipeline. Used by the standalone mlp1_act launcher. + +#include "mlp1_fused_sm100.cuh" + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Shared memory — three per-WG store buffers (U', V', Z) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp1FusedActSmem { + using Element = typename Traits::Element; + + static constexpr int smem_X_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = + cosize_v; + + alignas(128) Element smem_X[smem_X_size]; + alignas(128) Element smem_W1[smem_W_size]; + alignas(128) Element smem_W2[smem_W_size]; + // 3 outputs × 2 WGs of per-WG (AtomTileM, EpiChunkN) slots. + alignas(128) Element store_buf_u[2 * smem_store_size]; + alignas(128) Element store_buf_v[2 * smem_store_size]; + alignas(128) Element store_buf_z[2 * smem_store_size]; + + // SM100 (Compute=100) only: tcgen05.alloc landing slot + accumulator + // pipeline storage (see Mlp1FusedSmem). Untouched by the Hopper path. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* X_data() { return &smem_X[0]; } + CUTE_DEVICE Element* W1_data() { return &smem_W1[0]; } + CUTE_DEVICE Element* W2_data() { return &smem_W2[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Producer — single (m), all (n) under split_idx/num_splits stride. +// Single fused X + W1 + W2 TMA pipe per k-step. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp1_fused_act_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedActSmem& smem, + TmaLoadX const& tma_load_x, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + int m, + int expert_or_n_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_n_rows, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + // int64_t cast for production shapes where E·I·D > 2^31. + auto mX = tma_load_x.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(hidden_dim))); + auto mB = [&]() { + if constexpr (Expert3D) { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(hidden_dim))); + } + }(); + auto mC = [&]() { + if constexpr (Expert3D) { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(hidden_dim))); + } + }(); + + auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); + auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); + auto cta_tma_c = tma_load_c.get_slice(Int<0>{}); + + auto tXsX = cta_tma_x.partition_D(sX); + auto tW1sW1 = cta_tma_b.partition_D(sW1); + auto tW2sW2 = cta_tma_c.partition_D(sW2); + + auto gX = local_tile(mX, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tXgX = cta_tma_x.partition_S(gX); + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + auto gB = [&]() { + if constexpr (Expert3D) { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto gC = [&]() { + if constexpr (Expert3D) { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto tBgB = cta_tma_b.partition_S(gB); + auto tCgC = cta_tma_c.partition_S(gC); + + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_x.with(*bar, 0), + tXgX(_, _, _, k), tXsX(_, _, _, state.index())); + copy(tma_load_b.with(*bar, 0), + tBgB(_, _, _, k), tW1sW1(_, _, _, state.index())); + copy(tma_load_c.with(*bar, 0), + tCgC(_, _, _, k), tW2sW2(_, _, _, state.index())); + } + ++state; + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Consumer — architecture-specialized on `int Compute` (see +// Mlp1FusedConsumerImpl in mlp1_fused.cuh for the dispatch rationale). +// Compute=90 → Hopper / WGMMA (cooperative 2-WG) +// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulators) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp1FusedActConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. z_m is the M-tile index into the output buffers. +// Split axis selected by TileM (kMSplit pattern): +// TileM=128 (M-split): WG_w writes rows [w*AtomTileM, (w+1)*AtomTileM) +// TileM=64 (N-split): WG_w writes cols [w*WgTileN, (w+1)*WgTileN) +// Each WG issues NumEpiRounds × 3 TMA stores per N-tile (U', V', Z, +// each at (AtomTileM, EpiChunkN)). (Verbatim from the original consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp1FusedActConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedActSmem& smem, + TmaStore const& tma_store_du_coef, + TmaStore const& tma_store_dv_coef, + TmaStore const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + auto tCsX = thr_mma.partition_A(sX); + auto tCsW1 = thr_mma.partition_B(sW1); + auto tCsW2 = thr_mma.partition_B(sW2); + + auto acc_B = partition_fragment_C(tiled_mma, + Shape, Int>{}); + auto acc_C = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + // Identity coords for per-thread acc → store_buf mapping. + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const bool is_my_wg_leader = (tid_in_wg == 0); + + static_assert(Traits::AtomTileM % Traits::ActStoreTileM == 0, + "activation store TileM must divide the Hopper warpgroup strip"); + constexpr int StoreSlicesPerWg = + Traits::AtomTileM / Traits::ActStoreTileM; + constexpr int StoreTilesPerM = + Traits::TileM / Traits::ActStoreTileM; + constexpr int store_slot_elems = + Traits::ActStoreTileM * Traits::EpiChunkN; + Element* my_store_u = smem.store_buf_u + my_wg * store_slot_elems; + Element* my_store_v = smem.store_buf_v + my_wg * store_slot_elems; + Element* my_store_z = smem.store_buf_z + my_wg * store_slot_elems; + auto sStoreU = make_tensor(make_smem_ptr(my_store_u), + typename Traits::SmemLayoutActStoreSlot{}); + auto sStoreV = make_tensor(make_smem_ptr(my_store_v), + typename Traits::SmemLayoutActStoreSlot{}); + auto sStoreZ = make_tensor(make_smem_ptr(my_store_z), + typename Traits::SmemLayoutActStoreSlot{}); + + // int64_t cast: num_z_m_tiles · TileM · intermediate_dim can exceed INT_MAX. + auto out_shape = make_shape( + static_cast(num_z_m_tiles) * Traits::TileM, + static_cast(intermediate_dim)); + auto mU = tma_store_du_coef.get_tma_tensor(out_shape); + auto mV = tma_store_dv_coef.get_tma_tensor(out_shape); + auto mZ = tma_store_z.get_tma_tensor(out_shape); + auto cta_tma_u = tma_store_du_coef.get_slice(Int<0>{}); + auto cta_tma_v = tma_store_dv_coef.get_slice(Int<0>{}); + auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); + + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; // CUTLASS-matching + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + + clear(acc_B); + clear(acc_C); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), acc_B); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW2(_, _, _, state.index()), acc_C); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), acc_B); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW2(_, _, _, state.index()), acc_C); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue: 3-store per-WG rounds ───────────── + // Per round each WG: write U', V', Z into 3 separate per-WG smem + // slots, then issue 3 TMA stores. No cross-WG sync — each WG + // owns its strip independently. + // TileM=128 (M-split): WG_w owns rows [w*AtomTileM, (w+1)*AtomTileM), + // full N. Store at row_tile = 2*z_m + my_wg, col_tile = n*NER + r. + // TileM=64 (N-split): WG_w owns full M (=64), cols + // [w*WgTileN, (w+1)*WgTileN). Store at row_tile = z_m, + // col_tile = n * (TileN/EpiChunkN) + my_wg*NER + r. + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + CUTE_UNROLL + for (int ms = 0; ms < StoreSlicesPerWg; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc_B); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (m_local >= ms * Traits::ActStoreTileM + && m_local < (ms + 1) * Traits::ActStoreTileM + && n_local >= chunk_start + && n_local < chunk_start + Traits::EpiChunkN) { + int store_m = m_local - ms * Traits::ActStoreTileM; + int chunk_n = n_local - chunk_start; + float u = acc_B(i); + float v = acc_C(i); + float sig = fast_sigmoid(u); + float vprime = u * sig; + float sil_d = sig + vprime * (1.0f - sig); + sStoreU(store_m, chunk_n) = + static_cast(v * sil_d); + sStoreV(store_m, chunk_n) = + static_cast(vprime); + sStoreZ(store_m, chunk_n) = + static_cast(vprime * v); + } + } + + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int row_tile, col_tile; + if constexpr (Traits::kMSplit) { + row_tile = StoreTilesPerM * z_m + + my_wg * StoreSlicesPerWg + ms; + col_tile = n * Traits::NumEpiRounds + r; + } else { + row_tile = StoreTilesPerM * z_m + ms; + col_tile = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto tile = make_tile( + Int{}, + Int{}); + auto gU = local_tile( + mU, tile, make_coord(row_tile, col_tile)); + auto gV = local_tile( + mV, tile, make_coord(row_tile, col_tile)); + auto gZ = local_tile( + mZ, tile, make_coord(row_tile, col_tile)); + copy(tma_store_du_coef, cta_tma_u.partition_S(sStoreU), + cta_tma_u.partition_D(gU)); + cute::tma_store_arrive(); + copy(tma_store_dv_coef, cta_tma_v.partition_S(sStoreV), + cta_tma_v.partition_D(gV)); + cute::tma_store_arrive(); + copy(tma_store_z, cta_tma_z.partition_S(sStoreZ), + cta_tma_z.partition_D(gZ)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp1FusedActConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. Same structure as the fused +// Mlp1FusedConsumerImpl<100>, but the epilogue derives the three act +// outputs from (U, V) and issues three TMA stores per 64-row tile: +// V' = silu(U), U' = V·silu'(U), Z = V'·V. +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp1FusedActConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedActSmem& smem, + TmaStore const& tma_store_du_coef, + TmaStore const& tma_store_dv_coef, + TmaStore const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int StoreTileM = Traits::ActStoreTileM; + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two consumer warpgroups"); + constexpr int WgN = TileN / 2; + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; + constexpr int MSub = TileM / StoreTileM; + + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + const int tmem_copy_tid = tid_wg; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + auto tiled_mma = make_tiled_mma( + SM100_MMA_F16BF16_SS{}); + auto cta_mma = tiled_mma.get_slice(0); + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAccU = cta_mma.make_fragment_C(tCgC); // U = X·B + auto tCtAccV = cta_mma.make_fragment_C(tCgC); // V = X·C + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the + // alternate stage while both epilogue WGs drain MMA(n). + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; + acc_params.consumer_arv_count = 1; + acc_params.initializing_warp = 4; + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAccU.data() = tmem_base; + tCtAccV.data() = tmem_base + uint32_t(TileN); + + // Three per-WG store slots (U', V', Z) — same smem as Hopper. + constexpr int store_slot_elems = StoreTileM * EpiChunkN; + auto sStoreU = make_tensor(make_smem_ptr(smem.store_buf_u + wg * store_slot_elems), + typename Traits::SmemLayoutActStoreSlot{}); + auto sStoreV = make_tensor(make_smem_ptr(smem.store_buf_v + wg * store_slot_elems), + typename Traits::SmemLayoutActStoreSlot{}); + auto sStoreZ = make_tensor(make_smem_ptr(smem.store_buf_z + wg * store_slot_elems), + typename Traits::SmemLayoutActStoreSlot{}); + + auto out_shape = make_shape( + static_cast(num_z_m_tiles) * TileM, + static_cast(intermediate_dim)); + auto mU = tma_store_du_coef.get_tma_tensor(out_shape); + auto mV = tma_store_dv_coef.get_tma_tensor(out_shape); + auto mZ = tma_store_z.get_tma_tensor(out_shape); + auto cta_tma_u = tma_store_du_coef.get_slice(Int<0>{}); + auto cta_tma_v = tma_store_dv_coef.get_slice(Int<0>{}); + auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); + + auto epi_tile = make_tile(Int{}, Int{}); + // Flat (M,N) view of the UMMA C-fragment before tiling (see mlp1_fused.cuh). + auto accU_mn = tCtAccU(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto accV_mn = tCtAccV(make_coord(_, _), _0{}, _0{}); + auto tAccU_epi = flat_divide(accU_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto tAccV_epi = flat_divide(accV_mn, epi_tile); + auto t2r = make_tmem_copy(TmemLoadOp{}, tAccU_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tmem_copy_tid); + auto tTR_tAccU = thr_t2r.partition_S(tAccU_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) + auto tTR_tAccV = thr_t2r.partition_S(tAccV_epi); + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + // Register fragments sized from the DEST (partition_D) per-thread shape, not + // partition_S(tmem) (warp-collective). See mlp1_fused.cuh for the rationale. + auto tTR_rU = make_tensor(shape(tTR_cChunk)); // f32 regs (U) + auto tTR_rV = make_tensor(shape(tTR_cChunk)); // f32 regs (V) + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + // Mainloop (warp 3 only), bracketed by the accumulator pipeline. + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); + int acc_stage = acc_prod_state.index(); + uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); + tCtAccU.data() = stage_base; + tCtAccV.data() = stage_base + uint32_t(TileN); + for (int k = 0; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + auto tCsX = cta_mma.partition_A(sX (_, _, state.index())); + auto tCsW1 = cta_mma.partition_B(sW1(_, _, state.index())); + auto tCsW2 = cta_mma.partition_B(sW2(_, _, state.index())); + auto tCrX = cta_mma.make_fragment_A(tCsX); + auto tCrW1 = cta_mma.make_fragment_B(tCsW1); + auto tCrW2 = cta_mma.make_fragment_B(tCsW2); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrX); ++kb) { + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrX(_, _, kb), tCrW1(_, _, kb), tCtAccU); + gemm(tiled_mma, tCrX(_, _, kb), tCrW2(_, _, kb), tCtAccV); + } + pipe.consumer_release(state); + ++state; + } + acc_pipe.producer_commit(acc_prod_state); + ++acc_prod_state; + } + + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); + tCtAccU.data() = stage_base; + tCtAccV.data() = stage_base + uint32_t(TileN); + auto accU_mn_stage = tCtAccU(make_coord(_, _), _0{}, _0{}); + auto accV_mn_stage = tCtAccV(make_coord(_, _), _0{}, _0{}); + auto tAccU_epi_stage = flat_divide(accU_mn_stage, epi_tile); + auto tAccV_epi_stage = flat_divide(accV_mn_stage, epi_tile); + auto tTR_tAccU_stage = thr_t2r.partition_S(tAccU_epi_stage); + auto tTR_tAccV_stage = thr_t2r.partition_S(tAccV_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; + + // U and V stay in registers (tTR_rU, tTR_rV); the three act + // outputs are derived and cast to Element inline at the smem write + // below — no extra register sets. Each reg's row falls in exactly + // one MSub slice, so this evaluates each element once. + copy(t2r, tTR_tAccU_stage(_, _, _, _0{}, chunk), tTR_rU); + copy(t2r, tTR_tAccV_stage(_, _, _, _0{}, chunk), tTR_rV); + cutlass::arch::fence_view_async_tmem_load(); + + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rU); ++i) { + int m_row = get<0>(tTR_cChunk(i)); + int n_col = get<1>(tTR_cChunk(i)); + if (m_row >= ms * StoreTileM && m_row < (ms + 1) * StoreTileM) { + int m_loc = m_row - ms * StoreTileM; + float u = tTR_rU(i); + float v = tTR_rV(i); + float sig = fast_sigmoid(u); + float vprime = u * sig; // silu(U) + float sil_d = sig + vprime * (1.0f - sig); // silu'(U) + float uprime = v * sil_d; + Element uprime_e = static_cast(uprime); + Element vprime_e = static_cast(vprime); + Element z_e = static_cast(vprime * v); + sStoreU(m_loc, n_col) = uprime_e; // V·silu'(U) + sStoreV(m_loc, n_col) = vprime_e; // silu(U) + sStoreZ(m_loc, n_col) = z_e; // V'·V + } + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int row_tile = MSub * z_m + ms; + int col_tile = n * (TileN / EpiChunkN) + chunk; + auto gU = local_tile(mU, + make_tile(Int{}, Int{}), + make_coord(row_tile, col_tile)); + auto gV = local_tile(mV, + make_tile(Int{}, Int{}), + make_coord(row_tile, col_tile)); + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(row_tile, col_tile)); + copy(tma_store_du_coef, cta_tma_u.partition_S(sStoreU), + cta_tma_u.partition_D(gU)); + cute::tma_store_arrive(); + copy(tma_store_dv_coef, cta_tma_v.partition_S(sStoreV), + cta_tma_v.partition_D(gV)); + cute::tma_store_arrive(); + copy(tma_store_z, cta_tma_z.partition_S(sStoreZ), + cta_tma_z.partition_D(gZ)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // Epilogue done reading TMEM; release the accumulator for the next + // n-tile's MMA (one elected consumer thread arrives). + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + __trap(); +#endif +} +}; // Mlp1FusedActConsumerImpl<100> + +// Forwarder — `Compute` defaults to 90 (existing call sites unchanged); +// pass for the Blackwell path. +template +__device__ __forceinline__ void mlp1_fused_act_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedActSmem& smem, + TmaStore const& tma_store_du_coef, + TmaStore const& tma_store_dv_coef, + TmaStore const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + Mlp1FusedActConsumerImpl::template run( + pipe, state, smem, tma_store_du_coef, tma_store_dv_coef, tma_store_z, + z_m, intermediate_dim, num_z_m_tiles, num_n_tiles, num_k_tiles, + split_idx, num_splits); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm90.cuh new file mode 100644 index 00000000..46a6b1c6 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm90.cuh @@ -0,0 +1,703 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 1 (fused tile + silu-backward coefficients): +// U = B @ X +// V = C @ X +// V' = silu(U) (coefficient for dV = dZ * V' later) +// U' = V * silu'(U) (coefficient for dU = dZ * U' later) +// Z = V' * V = silu(U) * V (forward output, consumed by mlp3) +// ═══════════════════════════════════════════════════════════════════ +// +// Cooperative 2-WG consumer producing three outputs (U', V', Z) per +// (TileM, TileN) tile. Split axis selected by TileM (kMSplit pattern, +// same as mlp1_fused.cuh / mlp2_t_fused.cuh / mlp5_fused.cuh): +// TileM=128 (M-split, Layout<_2,_1,_1>): each WG owns AtomTileM=64 +// rows × full TileN. WG_w writes rows [w*64, (w+1)*64). +// TileM=64 (N-split, Layout<_1,_2,_1>): each WG owns full M (=64) +// × WgTileN = TileN/2. WG_w writes cols [w*WgTileN, (w+1)*WgTileN). +// Dual acc_B/acc_C from a single fused X+W1+W2 TMA pipeline, +// EpiChunkN-wide per-WG epilogue rounds — each round writes three +// smem buffers and issues three TMA stores. +// +// Smem layout (per Mlp1Traits): three per-WG store slots. WG_w uses +// indices [w·S, (w+1)·S) within each buf_{u,v,z}, where S = +// AtomTileM·EpiChunkN. Stages=3 default (vs Stages=4 in fwd mlp1) to +// fit the extra two output buffers within the 228 KiB SM smem cap. +// +// Single-tile API: caller passes (m, expert_n_offset) and a single +// fused X+B+C pipeline. Used by the standalone mlp1_act launcher. + +#include "mlp1_fused_sm90.cuh" + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Shared memory — three per-WG store buffers (U', V', Z) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp1FusedActSmem { + using Element = typename Traits::Element; + + static constexpr int smem_X_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = + cosize_v; + + alignas(128) Element smem_X[smem_X_size]; + alignas(128) Element smem_W1[smem_W_size]; + alignas(128) Element smem_W2[smem_W_size]; + // 3 outputs × 2 WGs of per-WG (AtomTileM, EpiChunkN) slots. + alignas(128) Element store_buf_u[2 * smem_store_size]; + alignas(128) Element store_buf_v[2 * smem_store_size]; + alignas(128) Element store_buf_z[2 * smem_store_size]; + + // SM100 (Compute=100) only: tcgen05.alloc landing slot + accumulator + // pipeline storage (see Mlp1FusedSmem). Untouched by the Hopper path. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* X_data() { return &smem_X[0]; } + CUTE_DEVICE Element* W1_data() { return &smem_W1[0]; } + CUTE_DEVICE Element* W2_data() { return &smem_W2[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Producer — single (m), all (n) under split_idx/num_splits stride. +// Single fused X + W1 + W2 TMA pipe per k-step. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp1_fused_act_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedActSmem& smem, + TmaLoadX const& tma_load_x, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + int m, + int expert_or_n_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_n_rows, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + // int64_t cast for production shapes where E·I·D > 2^31. + auto mX = tma_load_x.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(hidden_dim))); + auto mB = [&]() { + if constexpr (Expert3D) { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(hidden_dim))); + } + }(); + auto mC = [&]() { + if constexpr (Expert3D) { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(hidden_dim))); + } + }(); + + auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); + auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); + auto cta_tma_c = tma_load_c.get_slice(Int<0>{}); + + auto tXsX = cta_tma_x.partition_D(sX); + auto tW1sW1 = cta_tma_b.partition_D(sW1); + auto tW2sW2 = cta_tma_c.partition_D(sW2); + + auto gX = local_tile(mX, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tXgX = cta_tma_x.partition_S(gX); + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + auto gB = [&]() { + if constexpr (Expert3D) { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto gC = [&]() { + if constexpr (Expert3D) { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto tBgB = cta_tma_b.partition_S(gB); + auto tCgC = cta_tma_c.partition_S(gC); + + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_x.with(*bar, 0), + tXgX(_, _, _, k), tXsX(_, _, _, state.index())); + copy(tma_load_b.with(*bar, 0), + tBgB(_, _, _, k), tW1sW1(_, _, _, state.index())); + copy(tma_load_c.with(*bar, 0), + tCgC(_, _, _, k), tW2sW2(_, _, _, state.index())); + } + ++state; + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Consumer — architecture-specialized on `int Compute` (see +// Mlp1FusedConsumerImpl in mlp1_fused.cuh for the dispatch rationale). +// Compute=90 → Hopper / WGMMA (cooperative 2-WG) +// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulators) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp1FusedActConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. z_m is the M-tile index into the output buffers. +// Split axis selected by TileM (kMSplit pattern): +// TileM=128 (M-split): WG_w writes rows [w*AtomTileM, (w+1)*AtomTileM) +// TileM=64 (N-split): WG_w writes cols [w*WgTileN, (w+1)*WgTileN) +// Each WG issues NumEpiRounds × 3 TMA stores per N-tile (U', V', Z, +// each at (AtomTileM, EpiChunkN)). (Verbatim from the original consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp1FusedActConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedActSmem& smem, + TmaStore const& tma_store_du_coef, + TmaStore const& tma_store_dv_coef, + TmaStore const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + auto tCsX = thr_mma.partition_A(sX); + auto tCsW1 = thr_mma.partition_B(sW1); + auto tCsW2 = thr_mma.partition_B(sW2); + + auto acc_B = partition_fragment_C(tiled_mma, + Shape, Int>{}); + auto acc_C = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + // Identity coords for per-thread acc → store_buf mapping. + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const bool is_my_wg_leader = (tid_in_wg == 0); + + static_assert(Traits::AtomTileM % Traits::ActStoreTileM == 0, + "activation store TileM must divide the Hopper warpgroup strip"); + constexpr int StoreSlicesPerWg = + Traits::AtomTileM / Traits::ActStoreTileM; + constexpr int StoreTilesPerM = + Traits::TileM / Traits::ActStoreTileM; + constexpr int store_slot_elems = + Traits::ActStoreTileM * Traits::EpiChunkN; + Element* my_store_u = smem.store_buf_u + my_wg * store_slot_elems; + Element* my_store_v = smem.store_buf_v + my_wg * store_slot_elems; + Element* my_store_z = smem.store_buf_z + my_wg * store_slot_elems; + auto sStoreU = make_tensor(make_smem_ptr(my_store_u), + typename Traits::SmemLayoutActStoreSlot{}); + auto sStoreV = make_tensor(make_smem_ptr(my_store_v), + typename Traits::SmemLayoutActStoreSlot{}); + auto sStoreZ = make_tensor(make_smem_ptr(my_store_z), + typename Traits::SmemLayoutActStoreSlot{}); + + // int64_t cast: num_z_m_tiles · TileM · intermediate_dim can exceed INT_MAX. + auto out_shape = make_shape( + static_cast(num_z_m_tiles) * Traits::TileM, + static_cast(intermediate_dim)); + auto mU = tma_store_du_coef.get_tma_tensor(out_shape); + auto mV = tma_store_dv_coef.get_tma_tensor(out_shape); + auto mZ = tma_store_z.get_tma_tensor(out_shape); + auto cta_tma_u = tma_store_du_coef.get_slice(Int<0>{}); + auto cta_tma_v = tma_store_dv_coef.get_slice(Int<0>{}); + auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); + + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; // CUTLASS-matching + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + + clear(acc_B); + clear(acc_C); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), acc_B); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW2(_, _, _, state.index()), acc_C); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), acc_B); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW2(_, _, _, state.index()), acc_C); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue: 3-store per-WG rounds ───────────── + // Per round each WG: write U', V', Z into 3 separate per-WG smem + // slots, then issue 3 TMA stores. No cross-WG sync — each WG + // owns its strip independently. + // TileM=128 (M-split): WG_w owns rows [w*AtomTileM, (w+1)*AtomTileM), + // full N. Store at row_tile = 2*z_m + my_wg, col_tile = n*NER + r. + // TileM=64 (N-split): WG_w owns full M (=64), cols + // [w*WgTileN, (w+1)*WgTileN). Store at row_tile = z_m, + // col_tile = n * (TileN/EpiChunkN) + my_wg*NER + r. + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + CUTE_UNROLL + for (int ms = 0; ms < StoreSlicesPerWg; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc_B); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (m_local >= ms * Traits::ActStoreTileM + && m_local < (ms + 1) * Traits::ActStoreTileM + && n_local >= chunk_start + && n_local < chunk_start + Traits::EpiChunkN) { + int store_m = m_local - ms * Traits::ActStoreTileM; + int chunk_n = n_local - chunk_start; + float u = acc_B(i); + float v = acc_C(i); + float sig = fast_sigmoid(u); + float vprime = u * sig; + float sil_d = sig + vprime * (1.0f - sig); + sStoreU(store_m, chunk_n) = + static_cast(v * sil_d); + sStoreV(store_m, chunk_n) = + static_cast(vprime); + sStoreZ(store_m, chunk_n) = + static_cast(vprime * v); + } + } + + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int row_tile, col_tile; + if constexpr (Traits::kMSplit) { + row_tile = StoreTilesPerM * z_m + + my_wg * StoreSlicesPerWg + ms; + col_tile = n * Traits::NumEpiRounds + r; + } else { + row_tile = StoreTilesPerM * z_m + ms; + col_tile = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto tile = make_tile( + Int{}, + Int{}); + auto gU = local_tile( + mU, tile, make_coord(row_tile, col_tile)); + auto gV = local_tile( + mV, tile, make_coord(row_tile, col_tile)); + auto gZ = local_tile( + mZ, tile, make_coord(row_tile, col_tile)); + copy(tma_store_du_coef, cta_tma_u.partition_S(sStoreU), + cta_tma_u.partition_D(gU)); + cute::tma_store_arrive(); + copy(tma_store_dv_coef, cta_tma_v.partition_S(sStoreV), + cta_tma_v.partition_D(gV)); + cute::tma_store_arrive(); + copy(tma_store_z, cta_tma_z.partition_S(sStoreZ), + cta_tma_z.partition_D(gZ)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp1FusedActConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. Same structure as the fused +// Mlp1FusedConsumerImpl<100>, but the epilogue derives the three act +// outputs from (U, V) and issues three TMA stores per 64-row tile: +// V' = silu(U), U' = V·silu'(U), Z = V'·V. +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp1FusedActConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedActSmem& smem, + TmaStore const& tma_store_du_coef, + TmaStore const& tma_store_dv_coef, + TmaStore const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int StoreTileM = Traits::ActStoreTileM; + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two consumer warpgroups"); + constexpr int WgN = TileN / 2; + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; + constexpr int MSub = TileM / StoreTileM; + + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + const int tmem_copy_tid = tid_wg; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + auto tiled_mma = make_tiled_mma( + SM100_MMA_F16BF16_SS{}); + auto cta_mma = tiled_mma.get_slice(0); + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAccU = cta_mma.make_fragment_C(tCgC); // U = X·B + auto tCtAccV = cta_mma.make_fragment_C(tCgC); // V = X·C + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the + // alternate stage while both epilogue WGs drain MMA(n). + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; + acc_params.consumer_arv_count = 1; + acc_params.initializing_warp = 4; + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAccU.data() = tmem_base; + tCtAccV.data() = tmem_base + uint32_t(TileN); + + // Three per-WG store slots (U', V', Z) — same smem as Hopper. + constexpr int store_slot_elems = StoreTileM * EpiChunkN; + auto sStoreU = make_tensor(make_smem_ptr(smem.store_buf_u + wg * store_slot_elems), + typename Traits::SmemLayoutActStoreSlot{}); + auto sStoreV = make_tensor(make_smem_ptr(smem.store_buf_v + wg * store_slot_elems), + typename Traits::SmemLayoutActStoreSlot{}); + auto sStoreZ = make_tensor(make_smem_ptr(smem.store_buf_z + wg * store_slot_elems), + typename Traits::SmemLayoutActStoreSlot{}); + + auto out_shape = make_shape( + static_cast(num_z_m_tiles) * TileM, + static_cast(intermediate_dim)); + auto mU = tma_store_du_coef.get_tma_tensor(out_shape); + auto mV = tma_store_dv_coef.get_tma_tensor(out_shape); + auto mZ = tma_store_z.get_tma_tensor(out_shape); + auto cta_tma_u = tma_store_du_coef.get_slice(Int<0>{}); + auto cta_tma_v = tma_store_dv_coef.get_slice(Int<0>{}); + auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); + + auto epi_tile = make_tile(Int{}, Int{}); + // Flat (M,N) view of the UMMA C-fragment before tiling (see mlp1_fused.cuh). + auto accU_mn = tCtAccU(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto accV_mn = tCtAccV(make_coord(_, _), _0{}, _0{}); + auto tAccU_epi = flat_divide(accU_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto tAccV_epi = flat_divide(accV_mn, epi_tile); + auto t2r = make_tmem_copy(TmemLoadOp{}, tAccU_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tmem_copy_tid); + auto tTR_tAccU = thr_t2r.partition_S(tAccU_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) + auto tTR_tAccV = thr_t2r.partition_S(tAccV_epi); + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + // Register fragments sized from the DEST (partition_D) per-thread shape, not + // partition_S(tmem) (warp-collective). See mlp1_fused.cuh for the rationale. + auto tTR_rU = make_tensor(shape(tTR_cChunk)); // f32 regs (U) + auto tTR_rV = make_tensor(shape(tTR_cChunk)); // f32 regs (V) + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + // Mainloop (warp 3 only), bracketed by the accumulator pipeline. + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); + int acc_stage = acc_prod_state.index(); + uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); + tCtAccU.data() = stage_base; + tCtAccV.data() = stage_base + uint32_t(TileN); + for (int k = 0; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + auto tCsX = cta_mma.partition_A(sX (_, _, state.index())); + auto tCsW1 = cta_mma.partition_B(sW1(_, _, state.index())); + auto tCsW2 = cta_mma.partition_B(sW2(_, _, state.index())); + auto tCrX = cta_mma.make_fragment_A(tCsX); + auto tCrW1 = cta_mma.make_fragment_B(tCsW1); + auto tCrW2 = cta_mma.make_fragment_B(tCsW2); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrX); ++kb) { + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrX(_, _, kb), tCrW1(_, _, kb), tCtAccU); + gemm(tiled_mma, tCrX(_, _, kb), tCrW2(_, _, kb), tCtAccV); + } + pipe.consumer_release(state); + ++state; + } + acc_pipe.producer_commit(acc_prod_state); + ++acc_prod_state; + } + + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); + tCtAccU.data() = stage_base; + tCtAccV.data() = stage_base + uint32_t(TileN); + auto accU_mn_stage = tCtAccU(make_coord(_, _), _0{}, _0{}); + auto accV_mn_stage = tCtAccV(make_coord(_, _), _0{}, _0{}); + auto tAccU_epi_stage = flat_divide(accU_mn_stage, epi_tile); + auto tAccV_epi_stage = flat_divide(accV_mn_stage, epi_tile); + auto tTR_tAccU_stage = thr_t2r.partition_S(tAccU_epi_stage); + auto tTR_tAccV_stage = thr_t2r.partition_S(tAccV_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; + + // U and V stay in registers (tTR_rU, tTR_rV); the three act + // outputs are derived and cast to Element inline at the smem write + // below — no extra register sets. Each reg's row falls in exactly + // one MSub slice, so this evaluates each element once. + copy(t2r, tTR_tAccU_stage(_, _, _, _0{}, chunk), tTR_rU); + copy(t2r, tTR_tAccV_stage(_, _, _, _0{}, chunk), tTR_rV); + cutlass::arch::fence_view_async_tmem_load(); + + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rU); ++i) { + int m_row = get<0>(tTR_cChunk(i)); + int n_col = get<1>(tTR_cChunk(i)); + if (m_row >= ms * StoreTileM && m_row < (ms + 1) * StoreTileM) { + int m_loc = m_row - ms * StoreTileM; + float u = tTR_rU(i); + float v = tTR_rV(i); + float sig = fast_sigmoid(u); + float vprime = u * sig; // silu(U) + float sil_d = sig + vprime * (1.0f - sig); // silu'(U) + float uprime = v * sil_d; + Element uprime_e = static_cast(uprime); + Element vprime_e = static_cast(vprime); + Element z_e = static_cast(vprime * v); + sStoreU(m_loc, n_col) = uprime_e; // V·silu'(U) + sStoreV(m_loc, n_col) = vprime_e; // silu(U) + sStoreZ(m_loc, n_col) = z_e; // V'·V + } + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int row_tile = MSub * z_m + ms; + int col_tile = n * (TileN / EpiChunkN) + chunk; + auto gU = local_tile(mU, + make_tile(Int{}, Int{}), + make_coord(row_tile, col_tile)); + auto gV = local_tile(mV, + make_tile(Int{}, Int{}), + make_coord(row_tile, col_tile)); + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(row_tile, col_tile)); + copy(tma_store_du_coef, cta_tma_u.partition_S(sStoreU), + cta_tma_u.partition_D(gU)); + cute::tma_store_arrive(); + copy(tma_store_dv_coef, cta_tma_v.partition_S(sStoreV), + cta_tma_v.partition_D(gV)); + cute::tma_store_arrive(); + copy(tma_store_z, cta_tma_z.partition_S(sStoreZ), + cta_tma_z.partition_D(gZ)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // Epilogue done reading TMEM; release the accumulator for the next + // n-tile's MMA (one elected consumer thread arrives). + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + __trap(); +#endif +} +}; // Mlp1FusedActConsumerImpl<100> + +// Forwarder — `Compute` defaults to 90 (existing call sites unchanged); +// pass for the Blackwell path. +template +__device__ __forceinline__ void mlp1_fused_act_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedActSmem& smem, + TmaStore const& tma_store_du_coef, + TmaStore const& tma_store_dv_coef, + TmaStore const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + Mlp1FusedActConsumerImpl::template run( + pipe, state, smem, tma_store_du_coef, tma_store_dv_coef, tma_store_z, + z_m, intermediate_dim, num_z_m_tiles, num_n_tiles, num_k_tiles, + split_idx, num_splits); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm100.cuh new file mode 100644 index 00000000..57b905ce --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm100.cuh @@ -0,0 +1,934 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 1 (single-tile, fused): Z = SiLU(B @ X) * (C @ X) +// ═══════════════════════════════════════════════════════════════════ +// +// Single (TileM, TileN) producer/consumer building blocks for callers +// that drive the outer M-tile loop themselves: the standalone MLP1 +// benchmark (mlp1.cu), the fused MLP kernel (mlp.cuh), the activation +// variants (mlp1_fused_act.cuh), and the backward path (mlp_bwd.cuh). +// +// Design: +// - Cooperative 2-WG consumer; split axis selected by TileM: +// TileM=128 → Layout<_2,_1,_1>, split along M (each WG = 64×TileN) +// TileM=64 → Layout<_1,_2,_1>, split along N (each WG = 64×TileN/2) +// - Two accumulators per WG (acc_B = U = B@X, acc_C = V = C@X) +// - Per-thread SiLU(U)*V fusion (no cross-WG sync) +// - Per-WG store_buf with EpiChunkN-round epilogue +// - K_PIPE_MMAS = 1 (CUTLASS-canonical, see +// sm90_mma_tma_gmma_ss_warpspecialized.hpp:264) +// - Single fused X+W1+W2 TMA pipe (1 mbarrier per stage, 3 TMA copies) +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for X / B / C) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Per `expert_ids`: one int per TileM-row m-block. +// +// Mlp1FusedSmem omits the pipeline storage so callers can fuse it with +// other phases' smem (see mlp.cuh). The benchmark wrapper (mlp1.cu) is +// responsible for attaching MainloopPipeline::SharedStorage and running +// the outer for(m) loop over m-tiles. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm100.cuh" +#include "math_sm100.cuh" +#include "tmem_load_op_sm100.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer +// specialization. These headers are header-only and safe to include under +// an sm_90a build — the tcgen05 PTX they emit is guarded by +// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on +// __CUDA_ARCH__ >= 1000) ever instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 128, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 32, + int ActStoreTileM_ = 64 +> +struct Mlp1Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + static constexpr int ActStoreTileM = ActStoreTileM_; + + static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); + + // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate + // along M (each owns 64 rows × TileN). With TileM=64 the WGs + // cooperate along N (each owns 64 rows × TileN/2). AtomTileM is the + // per-WG M extent — always 64. + static constexpr bool kMSplit = (TileM == 128); + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); + + static_assert(TileN <= 256, "TileN must fit TMA box limit"); + static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); + static_assert(ActStoreTileM == 32 || ActStoreTileM == 64, + "activation TMA stores support M32 or M64"); + static_assert(TileM % ActStoreTileM == 0, + "activation store TileM must divide the GEMM TileM"); + static_assert(kMSplit || (TileN % 2 == 0), + "TileN must be even when TileM=64 (cooperative N-split halves)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + // Cooperative 2-WG layout: M-split for TileM=128, N-split for TileM=64. + static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); + using GmmaAtom = typename GmmaSelector::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom; + + using SmemLayoutX_1 = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int>{})); + using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int>{})); + + using SmemLayoutX = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int, Int>{})); + using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int, Int>{})); + + // Per-WG store buffer slot: (AtomTileM, EpiChunkN). One slot per WG. + // Epi writes the WG's m-strip in EpiChunkN-wide rounds. + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore_1 = SmemLayoutStoreSlot; + using SmemLayoutActStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop + // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the + // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the TMEM + // accumulator from the MMA warp to the epilogue warps. See + // Mlp1FusedConsumerImpl<100>. The Blackwell launcher must build the + // mainloop pipe with num_consumers=1 (one umma_arrive per stage). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Double-buffered SM100 TMEM accumulators: each stage owns U and V, each + // TileN columns wide. This lets MMA(n+1) use the alternate stage while the + // epilogue drains MMA(n)'s TMEM stage. The launcher allocates + // AccStages·2·TileN columns; the dedicated epilogue warpgroup (warps 8..11) + // drains the previous stage while the epilogue-free MMA warp runs ahead. + static constexpr int AccStages = 2; + static_assert(AccStages * (2 * TileN) <= 512, + "SM100 MLP1 needs 2 accumulator stages x (U,V) to fit in 512 TMEM columns"); + static_assert(TileN <= 128, + "SM100 double-buffered MLP1 accumulator requires TileN <= 128"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + static constexpr int TmaTransBytesX = + static_cast(size(SmemLayoutX_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesW = + static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); + // Fused pipeline: X + W1 + W2 share one mbarrier per stage. + static constexpr int TmaTransBytes = TmaTransBytesX + 2 * TmaTransBytesW; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; +}; + +template +using Mlp1MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Expert weight pointers (device-side) +// ═══════════════════════════════════════════════════════════════════ + +template +struct ExpertWeights1 { + const Element* B; // [intermediate_dim, hidden_dim] + const Element* C; // [intermediate_dim, hidden_dim] +}; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP1 shared memory — no pipeline storage (managed externally) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp1FusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_X_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_X[smem_X_size]; + alignas(128) Element smem_W1[smem_W_size]; + alignas(128) Element smem_W2[smem_W_size]; + // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these; + // SM100 launchers allocate/free TMEM outside the consumer and publish the + // base here. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* X_data() { return &smem_X[0]; } + CUTE_DEVICE Element* W1_data() { return &smem_W1[0]; } + CUTE_DEVICE Element* W2_data() { return &smem_W2[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Pipeline construction helper — single fused X+W1+W2 pipeline +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ auto mlp1_make_pipe( + typename Traits::MainloopPipeline::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp1_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer +// (mlp1_fused_producer) and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ auto mlp1_make_pipe_umma( + typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA consumer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused producer — single (m), all (n) under blockIdx.y stride. +// Single fused X + W1 + W2 TMA pipe per k-step. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp1_fused_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedSmem& smem, + TmaLoadX const& tma_load_x, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + int m, + int expert_or_n_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_n_rows, + int num_n_tiles, + int num_k_tiles, + // N-split identity (flat-grid launch). Defaults (-1) fall back to + // blockIdx.y / gridDim.y for standalone 2-D-grid callers (mlp1.cu). + int split_idx = -1, + int num_splits = -1) { + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + // int64_t cast: total_n_rows × hidden_dim can exceed INT_MAX at + // production shapes (epp · I · D > 2^31), and CUTE's layout math + // otherwise wraps to a negative byte offset → kernel IMA. + auto mX = tma_load_x.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(hidden_dim))); + auto mB = [&]() { + if constexpr (Expert3D) { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(hidden_dim))); + } + }(); + auto mC = [&]() { + if constexpr (Expert3D) { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(hidden_dim))); + } + }(); + + auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); + auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); + auto cta_tma_c = tma_load_c.get_slice(Int<0>{}); + + auto tXsX = cta_tma_x.partition_D(sX); + auto tW1sW1 = cta_tma_b.partition_D(sW1); + auto tW2sW2 = cta_tma_c.partition_D(sW2); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + auto gB = [&]() { + if constexpr (Expert3D) { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto gC = [&]() { + if constexpr (Expert3D) { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto tBgB = cta_tma_b.partition_S(gB); + auto tCgC = cta_tma_c.partition_S(gC); + + CUTE_UNROLL + for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { + auto gX = local_tile(mX, + make_tile(Int{}, Int{}), + make_coord(m + m_sub, _)); + auto tXgX = cta_tma_x.partition_S(gX); + + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_x.with(*bar, 0), + tXgX(_, _, _, k), tXsX(_, _, _, state.index())); + copy(tma_load_b.with(*bar, 0), + tBgB(_, _, _, k), tW1sW1(_, _, _, state.index())); + copy(tma_load_c.with(*bar, 0), + tCgC(_, _, _, k), tW2sW2(_, _, _, state.index())); + } + ++state; + } + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused consumer — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported +// compute capability provides a `run(...)` member (a function-template on +// Traits/Pipeline/TmaStoreZ, since function templates can't be partially +// specialized). The free function `mlp1_fused_consumer` below forwards to +// the right specialization — call sites stay in the existing style. +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, dual register acc) +// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulators) +// +// z_m is the M-tile index into the output Z buffer. + +template +struct Mlp1FusedConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. The two WGs split the (TileM, TileN) output along +// M: WG1 writes rows [0, AtomTileM) of tile z_m, WG2 writes [AtomTileM, +// TileM). (Verbatim from the original mlp1_fused_consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp1FusedConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedSmem& smem, + TmaStoreZ const& tma_store_z, + int z_m, // M-tile index into Z buffer + int intermediate_dim, + int num_z_m_tiles, // number of M-tile slots in Z buffer + int num_n_tiles, + int num_k_tiles, + // N-split identity (flat-grid launch); -1 → blockIdx.y / gridDim.y. + int split_idx, + int num_splits) { + + static_assert(MSubTiles == 1, + "Hopper MLP1 consumer processes one GEMM M-tile per invocation"); + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + auto tCsX = thr_mma.partition_A(sX); + auto tCsW1 = thr_mma.partition_B(sW1); + auto tCsW2 = thr_mma.partition_B(sW2); + + // Two accumulators per WG, each covering the WG's (AtomTileM, TileN) m-strip. + auto acc_B = partition_fragment_C(tiled_mma, + Shape, Int>{}); + auto acc_C = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + // Identity coords for per-thread acc → store_buf mapping. + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // int64_t cast: num_z_m_tiles · TileM · intermediate_dim exceeds + // INT_MAX once grid_x · ZBuf · TileM · I > 2^31 (e.g. ZBuf=16 + // at I=16384 wraps and IMAs). + auto mZ = tma_store_z.get_tma_tensor( + make_shape(static_cast(num_z_m_tiles) * Traits::TileM, + static_cast(intermediate_dim))); + auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; // CUTLASS-matching + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc_B); + clear(acc_C); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), acc_B); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW2(_, _, _, state.index()), acc_C); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), acc_B); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW2(_, _, _, state.index()), acc_C); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // SiLU fusion: acc_B = SiLU(U) * V (per-thread, no sync). + CUTE_UNROLL + for (int i = 0; i < size(acc_B); ++i) { + float u = acc_B(i); + float v = acc_C(i); + acc_B(i) = fast_silu(u) * v; + } + + // ── Epilogue ────────────────────────────────────── + // Each WG writes its (AtomTileM=64) × EpiChunkN slot to store_buf + // and issues TMA stores. NumEpiRounds rounds per WG. The per-thread + // acc element → (m_local, n_local) mapping depends on the split: + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N + // TileM=64 (N-split): WG_w owns full M, cols [w*WgTileN, ...) + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc_B); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc_B(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: row tile space = 2 tiles per logical M + // (AtomTileM=64). WG_w writes (2*z_m + w, n*Rounds + r). + m_tile_idx = 2 * z_m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M. WG_w handles its + // N-half: chunk index = n*(TileN/EpiChunkN) + // + w*NumEpiRounds + r. + m_tile_idx = z_m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_z, cta_tma_z.partition_S(sStore), + cta_tma_z.partition_D(gZ)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp1FusedConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. +// +// No cooperative-warpgroup machinery (UMMA has no warpgroup concept): +// - One 1SM UMMA atom with M=TileM covers the whole tile; accumulators +// U = X·B and V = X·C live in TMEM (no register/M-split). +// - Warp 3 issues both UMMAs and stays epilogue-free so it can run ahead to +// the alternate TMEM accumulator stage (AccStages=2). +// - Epilogue: aligned WGs in warps 4..11 split TileN, read TMEM→regs, +// compute SiLU while the previous TMA store drains, and store 64-row +// (AtomTileM) tiles via the existing reg→SMEM→TMA path. +// +// TileM is the template parameter as on Hopper (64 or 128 — both native +// 1SM UMMA M sizes). For TileM=128 the epilogue stores MSub=2 row-tiles. +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp1FusedConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedSmem& smem, + TmaStoreZ const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two epilogue warpgroups"); + constexpr int WgN = TileN / 2; // per-WG N width + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per epilogue WG + constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) + + // ── Thread identity (Blackwell warp-specialised layout) ────────────── + // Warp 3 is the dedicated, epilogue-free UMMA producer (moved out of the + // consumer warpgroups so it can run AccStages tiles ahead). Warps 4..11 are + // the two aligned epilogue warpgroups (WG0 = warps 4..7, WG1 = warps 8..11); + // each drains one half of TileN from the current accumulator stage while the + // MMA warp fills the alternate TMEM stage. In the fused MoE CTA warp 0 = TMA + // and warps 1..2 = NVSHMEM comm; warp 3 is the freed comm warp repurposed to + // issue UMMA (see PR #19 + mlp.cuh). Total: 12 warps / 3 hardware warpgroups. + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); // dedicated UMMA warp + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); // two aligned WGs + // Epilogue-local identity — valid only for warps 4..11. Warp 3 never enters + // the epilogue block, so its (garbage) tid_in_epi/wg/tid_wg are unused. + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; // 0|1 + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; // 0..127 + const int tmem_copy_tid = tid_wg; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; // 1 (WG0) or 2 (WG1) + // Warps 3..11 (MMA + both epilogue WGs) rendezvous for the acc-pipe init and + // the launcher-safe TMEM free. ConsumerThreads (256) covers only warps 4..11. + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, K-major SS ── + auto tiled_mma = make_tiled_mma( + SM100_MMA_F16BF16_SS{}); + auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + // ── Two TMEM accumulators: U = X·B, V = X·C ── + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAccU = cta_mma.make_fragment_C(tCgC); + auto tCtAccV = cta_mma.make_fragment_C(tCgC); + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the + // alternate stage while both epilogue WGs drain MMA(n). producer_commit + // issues the UMMA-gated "accumulator ready" arrival; one elected epilogue + // thread releases the stage (after a 256-thread rendezvous) so the MMA + // warp may reuse it AccStages tiles later. ── + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; // one umma_arrive per commit (warp 3) + acc_params.consumer_arv_count = 1; // one elected releaser per stage + acc_params.initializing_warp = 4; // warp 4 inits the acc barriers + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + // Warps 3..11 rendezvous so the acc-pipe barriers (init'd by warp 4) are + // visible before the MMA warp (3) and both epilogue WGs (4..11) use them. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAccU.data() = tmem_base; + tCtAccV.data() = tmem_base + uint32_t(TileN); + + // ── Per-WG store slot (AtomTileM × EpiChunkN) — WG0 uses store_buf[0..S), + // WG1 uses [S..2S). Same smem layout as Hopper. ── + constexpr int store_slot_elems = AtomTileM * EpiChunkN; + Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + auto mZ = tma_store_z.get_tma_tensor( + make_shape(static_cast(num_z_m_tiles) * TileM, + static_cast(intermediate_dim))); + auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); + + // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local + // tid). Each epi tile is the full TileM rows × EpiChunkN cols; the + // 128-thread tcgen05.ld maps one thread per TMEM datapath row. + // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) + // modes, which is what make_tmem_copy's cotiled builder requires. ── + auto epi_tile = make_tile(Int{}, Int{}); + // Extract the flat (M,N) view from the UMMA C-fragment (MMA,MMA_M,MMA_N) + // before tiling — matches CUTLASS's accumulators(make_coord(_,_),_0,_0). + auto accU_mn = tCtAccU(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto tAccU_epi = flat_divide(accU_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto t2r = make_tmem_copy(TmemLoadOp{}, tAccU_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tmem_copy_tid); + // Within-chunk (m,n) coords for the reg → smem-slot scatter. + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + // Register fragments sized from the DEST (partition_D) shape — per-thread + // (T2R,T2R_M,T2R_N). Sizing from partition_S(tmem) would wrongly include the + // warp-collective datapath-lane dim (the TMEM_LOAD atom distributes those + // lanes into per-thread registers internally). Matches CUTLASS's + // make_tensor(shape(tTR_sD)). + auto tTR_rU = make_tensor(shape(tTR_cChunk)); // f32 regs (U, reused) + auto tTR_rV = make_tensor(shape(tTR_cChunk)); // f32 regs (V) + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + CUTE_UNROLL + for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { + + // ── Mainloop (MMA warp = warp 3 only): bracket the k-loop with the + // accumulator pipeline. producer_acquire selects the free TMEM stage + // (AccStages=2 → MMA(n+1) fills the alternate stage while both epilogue + // WGs drain MMA(n)). consumer_release issues the UMMA-gated smem-buffer + // arrival on the mainloop pipe (no hand-rolled mbarrier). ── + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); // selected TMEM stage is free + int acc_stage = acc_prod_state.index(); + uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); + tCtAccU.data() = stage_base; + tCtAccV.data() = stage_base + uint32_t(TileN); + for (int k = 0; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + auto tCsX = cta_mma.partition_A(sX (_, _, state.index())); + auto tCsW1 = cta_mma.partition_B(sW1(_, _, state.index())); + auto tCsW2 = cta_mma.partition_B(sW2(_, _, state.index())); + auto tCrX = cta_mma.make_fragment_A(tCsX); + auto tCrW1 = cta_mma.make_fragment_B(tCsW1); + auto tCrW2 = cta_mma.make_fragment_B(tCsW2); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrX); ++kb) { + // First k-block of this n-tile clears TMEM; the rest accumulate. + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrX(_, _, kb), tCrW1(_, _, kb), tCtAccU); + gemm(tiled_mma, tCrX(_, _, kb), tCrW2(_, _, kb), tCtAccV); + } + pipe.consumer_release(state); // UMMA-gated smem-buffer release + ++state; + } + acc_pipe.producer_commit(acc_prod_state); // signal U/V ready (umma_arrive) + ++acc_prod_state; + } + + // ── Epilogue (two aligned WGs, warps 4..11): wait for the accumulator + // stage, then THIS WG drains its half of TileN (chunks + // [wg·NChunksHalf, +NChunksHalf)) while the MMA warp fills the + // alternate stage. ── + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); + tCtAccU.data() = stage_base; + tCtAccV.data() = stage_base + uint32_t(TileN); + auto accU_mn_stage = tCtAccU(make_coord(_, _), _0{}, _0{}); + auto accV_mn_stage = tCtAccV(make_coord(_, _), _0{}, _0{}); + auto tAccU_epi_stage = flat_divide(accU_mn_stage, epi_tile); + auto tAccV_epi_stage = flat_divide(accV_mn_stage, epi_tile); + auto tTR_tAccU = thr_t2r.partition_S(tAccU_epi_stage); + auto tTR_tAccV = thr_t2r.partition_S(tAccV_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; // absolute n-chunk index + + // TMEM → registers (this chunk, full TileM rows). + copy(t2r, tTR_tAccU(_, _, _, _0{}, chunk), tTR_rU); + copy(t2r, tTR_tAccV(_, _, _, _0{}, chunk), tTR_rV); + cutlass::arch::fence_view_async_tmem_load(); + + // SiLU(U)·V fused in place (overwrites the U regs) so it overlaps + // the prior in-flight store. The cast to Element happens at the + // smem write below — no extra register set. + CUTE_UNROLL + for (int i = 0; i < size(tTR_rU); ++i) + tTR_rU(i) = fast_silu(tTR_rU(i)) * tTR_rV(i); + + // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rU); ++i) { + int m_row = get<0>(tTR_cChunk(i)); // 0..TileM + int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN + if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) + sStore(m_row - ms * AtomTileM, n_col) = + static_cast(tTR_rU(i)); + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx = MSub * (z_m + m_sub) + ms; + int n_tile_idx = n * (TileN / EpiChunkN) + chunk; + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_z, cta_tma_z.partition_S(sStore), + cta_tma_z.partition_D(gZ)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // Both epilogue WGs done reading this TMEM stage → rendezvous, then + // one elected thread releases the stage so the MMA warp may reuse it + // when the AccStages accumulator pipeline wraps. + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + // Final rendezvous of the MMA warp (3) + both epilogue WGs (4..11) so the + // outer launcher can safely free TMEM after this consumer returns. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + // The Compute=100 specialization is never dispatched on non-SM100 targets + // (host picks Compute=90 there). Trap if it is ever reached. + __trap(); +#endif +} +}; // Mlp1FusedConsumerImpl<100> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — keeps the existing call style. `Compute` defaults to 90 so +// current call sites (mlp1_fused_consumer(...)) are unchanged; +// pass mlp1_fused_consumer(...) to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp1_fused_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedSmem& smem, + TmaStoreZ const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx = -1, + int num_splits = -1) { + Mlp1FusedConsumerImpl::template run< + Traits, MSubTiles, Pipeline, TmaStoreZ>( + pipe, state, smem, tma_store_z, z_m, intermediate_dim, + num_z_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm90.cuh new file mode 100644 index 00000000..ce4c9de2 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm90.cuh @@ -0,0 +1,934 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 1 (single-tile, fused): Z = SiLU(B @ X) * (C @ X) +// ═══════════════════════════════════════════════════════════════════ +// +// Single (TileM, TileN) producer/consumer building blocks for callers +// that drive the outer M-tile loop themselves: the standalone MLP1 +// benchmark (mlp1.cu), the fused MLP kernel (mlp.cuh), the activation +// variants (mlp1_fused_act.cuh), and the backward path (mlp_bwd.cuh). +// +// Design: +// - Cooperative 2-WG consumer; split axis selected by TileM: +// TileM=128 → Layout<_2,_1,_1>, split along M (each WG = 64×TileN) +// TileM=64 → Layout<_1,_2,_1>, split along N (each WG = 64×TileN/2) +// - Two accumulators per WG (acc_B = U = B@X, acc_C = V = C@X) +// - Per-thread SiLU(U)*V fusion (no cross-WG sync) +// - Per-WG store_buf with EpiChunkN-round epilogue +// - K_PIPE_MMAS = 1 (CUTLASS-canonical, see +// sm90_mma_tma_gmma_ss_warpspecialized.hpp:264) +// - Single fused X+W1+W2 TMA pipe (1 mbarrier per stage, 3 TMA copies) +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for X / B / C) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Per `expert_ids`: one int per TileM-row m-block. +// +// Mlp1FusedSmem omits the pipeline storage so callers can fuse it with +// other phases' smem (see mlp.cuh). The benchmark wrapper (mlp1.cu) is +// responsible for attaching MainloopPipeline::SharedStorage and running +// the outer for(m) loop over m-tiles. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm90.cuh" +#include "math_sm90.cuh" +#include "tmem_load_op_sm90.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer +// specialization. These headers are header-only and safe to include under +// an sm_90a build — the tcgen05 PTX they emit is guarded by +// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on +// __CUDA_ARCH__ >= 1000) ever instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 128, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 32, + int ActStoreTileM_ = 64 +> +struct Mlp1Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + static constexpr int ActStoreTileM = ActStoreTileM_; + + static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); + + // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate + // along M (each owns 64 rows × TileN). With TileM=64 the WGs + // cooperate along N (each owns 64 rows × TileN/2). AtomTileM is the + // per-WG M extent — always 64. + static constexpr bool kMSplit = (TileM == 128); + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); + + static_assert(TileN <= 256, "TileN must fit TMA box limit"); + static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); + static_assert(ActStoreTileM == 32 || ActStoreTileM == 64, + "activation TMA stores support M32 or M64"); + static_assert(TileM % ActStoreTileM == 0, + "activation store TileM must divide the GEMM TileM"); + static_assert(kMSplit || (TileN % 2 == 0), + "TileN must be even when TileM=64 (cooperative N-split halves)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + // Cooperative 2-WG layout: M-split for TileM=128, N-split for TileM=64. + static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); + using GmmaAtom = typename GmmaSelector::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom; + + using SmemLayoutX_1 = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int>{})); + using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int>{})); + + using SmemLayoutX = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int, Int>{})); + using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int, Int>{})); + + // Per-WG store buffer slot: (AtomTileM, EpiChunkN). One slot per WG. + // Epi writes the WG's m-strip in EpiChunkN-wide rounds. + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore_1 = SmemLayoutStoreSlot; + using SmemLayoutActStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop + // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the + // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the TMEM + // accumulator from the MMA warp to the epilogue warps. See + // Mlp1FusedConsumerImpl<100>. The Blackwell launcher must build the + // mainloop pipe with num_consumers=1 (one umma_arrive per stage). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Double-buffered SM100 TMEM accumulators: each stage owns U and V, each + // TileN columns wide. This lets MMA(n+1) use the alternate stage while the + // epilogue drains MMA(n)'s TMEM stage. The launcher allocates + // AccStages·2·TileN columns; the dedicated epilogue warpgroup (warps 8..11) + // drains the previous stage while the epilogue-free MMA warp runs ahead. + static constexpr int AccStages = 2; + static_assert(AccStages * (2 * TileN) <= 512, + "SM100 MLP1 needs 2 accumulator stages x (U,V) to fit in 512 TMEM columns"); + static_assert(TileN <= 128, + "SM100 double-buffered MLP1 accumulator requires TileN <= 128"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + static constexpr int TmaTransBytesX = + static_cast(size(SmemLayoutX_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesW = + static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); + // Fused pipeline: X + W1 + W2 share one mbarrier per stage. + static constexpr int TmaTransBytes = TmaTransBytesX + 2 * TmaTransBytesW; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; +}; + +template +using Mlp1MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Expert weight pointers (device-side) +// ═══════════════════════════════════════════════════════════════════ + +template +struct ExpertWeights1 { + const Element* B; // [intermediate_dim, hidden_dim] + const Element* C; // [intermediate_dim, hidden_dim] +}; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP1 shared memory — no pipeline storage (managed externally) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp1FusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_X_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_X[smem_X_size]; + alignas(128) Element smem_W1[smem_W_size]; + alignas(128) Element smem_W2[smem_W_size]; + // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these; + // SM100 launchers allocate/free TMEM outside the consumer and publish the + // base here. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* X_data() { return &smem_X[0]; } + CUTE_DEVICE Element* W1_data() { return &smem_W1[0]; } + CUTE_DEVICE Element* W2_data() { return &smem_W2[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Pipeline construction helper — single fused X+W1+W2 pipeline +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ auto mlp1_make_pipe( + typename Traits::MainloopPipeline::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp1_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer +// (mlp1_fused_producer) and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ auto mlp1_make_pipe_umma( + typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA consumer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused producer — single (m), all (n) under blockIdx.y stride. +// Single fused X + W1 + W2 TMA pipe per k-step. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp1_fused_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedSmem& smem, + TmaLoadX const& tma_load_x, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + int m, + int expert_or_n_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_n_rows, + int num_n_tiles, + int num_k_tiles, + // N-split identity (flat-grid launch). Defaults (-1) fall back to + // blockIdx.y / gridDim.y for standalone 2-D-grid callers (mlp1.cu). + int split_idx = -1, + int num_splits = -1) { + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + // int64_t cast: total_n_rows × hidden_dim can exceed INT_MAX at + // production shapes (epp · I · D > 2^31), and CUTE's layout math + // otherwise wraps to a negative byte offset → kernel IMA. + auto mX = tma_load_x.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(hidden_dim))); + auto mB = [&]() { + if constexpr (Expert3D) { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(hidden_dim))); + } + }(); + auto mC = [&]() { + if constexpr (Expert3D) { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(hidden_dim))); + } + }(); + + auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); + auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); + auto cta_tma_c = tma_load_c.get_slice(Int<0>{}); + + auto tXsX = cta_tma_x.partition_D(sX); + auto tW1sW1 = cta_tma_b.partition_D(sW1); + auto tW2sW2 = cta_tma_c.partition_D(sW2); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + auto gB = [&]() { + if constexpr (Expert3D) { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto gC = [&]() { + if constexpr (Expert3D) { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto tBgB = cta_tma_b.partition_S(gB); + auto tCgC = cta_tma_c.partition_S(gC); + + CUTE_UNROLL + for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { + auto gX = local_tile(mX, + make_tile(Int{}, Int{}), + make_coord(m + m_sub, _)); + auto tXgX = cta_tma_x.partition_S(gX); + + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_x.with(*bar, 0), + tXgX(_, _, _, k), tXsX(_, _, _, state.index())); + copy(tma_load_b.with(*bar, 0), + tBgB(_, _, _, k), tW1sW1(_, _, _, state.index())); + copy(tma_load_c.with(*bar, 0), + tCgC(_, _, _, k), tW2sW2(_, _, _, state.index())); + } + ++state; + } + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused consumer — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported +// compute capability provides a `run(...)` member (a function-template on +// Traits/Pipeline/TmaStoreZ, since function templates can't be partially +// specialized). The free function `mlp1_fused_consumer` below forwards to +// the right specialization — call sites stay in the existing style. +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, dual register acc) +// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulators) +// +// z_m is the M-tile index into the output Z buffer. + +template +struct Mlp1FusedConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. The two WGs split the (TileM, TileN) output along +// M: WG1 writes rows [0, AtomTileM) of tile z_m, WG2 writes [AtomTileM, +// TileM). (Verbatim from the original mlp1_fused_consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp1FusedConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedSmem& smem, + TmaStoreZ const& tma_store_z, + int z_m, // M-tile index into Z buffer + int intermediate_dim, + int num_z_m_tiles, // number of M-tile slots in Z buffer + int num_n_tiles, + int num_k_tiles, + // N-split identity (flat-grid launch); -1 → blockIdx.y / gridDim.y. + int split_idx, + int num_splits) { + + static_assert(MSubTiles == 1, + "Hopper MLP1 consumer processes one GEMM M-tile per invocation"); + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + auto tCsX = thr_mma.partition_A(sX); + auto tCsW1 = thr_mma.partition_B(sW1); + auto tCsW2 = thr_mma.partition_B(sW2); + + // Two accumulators per WG, each covering the WG's (AtomTileM, TileN) m-strip. + auto acc_B = partition_fragment_C(tiled_mma, + Shape, Int>{}); + auto acc_C = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + // Identity coords for per-thread acc → store_buf mapping. + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // int64_t cast: num_z_m_tiles · TileM · intermediate_dim exceeds + // INT_MAX once grid_x · ZBuf · TileM · I > 2^31 (e.g. ZBuf=16 + // at I=16384 wraps and IMAs). + auto mZ = tma_store_z.get_tma_tensor( + make_shape(static_cast(num_z_m_tiles) * Traits::TileM, + static_cast(intermediate_dim))); + auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; // CUTLASS-matching + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc_B); + clear(acc_C); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), acc_B); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW2(_, _, _, state.index()), acc_C); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + warpgroup_arrive(); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW1(_, _, _, state.index()), acc_B); + gemm(tiled_mma, tCsX(_, _, _, state.index()), + tCsW2(_, _, _, state.index()), acc_C); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc_B); + warpgroup_fence_operand(acc_C); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // SiLU fusion: acc_B = SiLU(U) * V (per-thread, no sync). + CUTE_UNROLL + for (int i = 0; i < size(acc_B); ++i) { + float u = acc_B(i); + float v = acc_C(i); + acc_B(i) = fast_silu(u) * v; + } + + // ── Epilogue ────────────────────────────────────── + // Each WG writes its (AtomTileM=64) × EpiChunkN slot to store_buf + // and issues TMA stores. NumEpiRounds rounds per WG. The per-thread + // acc element → (m_local, n_local) mapping depends on the split: + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N + // TileM=64 (N-split): WG_w owns full M, cols [w*WgTileN, ...) + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc_B); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc_B(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: row tile space = 2 tiles per logical M + // (AtomTileM=64). WG_w writes (2*z_m + w, n*Rounds + r). + m_tile_idx = 2 * z_m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M. WG_w handles its + // N-half: chunk index = n*(TileN/EpiChunkN) + // + w*NumEpiRounds + r. + m_tile_idx = z_m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_z, cta_tma_z.partition_S(sStore), + cta_tma_z.partition_D(gZ)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp1FusedConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. +// +// No cooperative-warpgroup machinery (UMMA has no warpgroup concept): +// - One 1SM UMMA atom with M=TileM covers the whole tile; accumulators +// U = X·B and V = X·C live in TMEM (no register/M-split). +// - Warp 3 issues both UMMAs and stays epilogue-free so it can run ahead to +// the alternate TMEM accumulator stage (AccStages=2). +// - Epilogue: aligned WGs in warps 4..11 split TileN, read TMEM→regs, +// compute SiLU while the previous TMA store drains, and store 64-row +// (AtomTileM) tiles via the existing reg→SMEM→TMA path. +// +// TileM is the template parameter as on Hopper (64 or 128 — both native +// 1SM UMMA M sizes). For TileM=128 the epilogue stores MSub=2 row-tiles. +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp1FusedConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedSmem& smem, + TmaStoreZ const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two epilogue warpgroups"); + constexpr int WgN = TileN / 2; // per-WG N width + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per epilogue WG + constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) + + // ── Thread identity (Blackwell warp-specialised layout) ────────────── + // Warp 3 is the dedicated, epilogue-free UMMA producer (moved out of the + // consumer warpgroups so it can run AccStages tiles ahead). Warps 4..11 are + // the two aligned epilogue warpgroups (WG0 = warps 4..7, WG1 = warps 8..11); + // each drains one half of TileN from the current accumulator stage while the + // MMA warp fills the alternate TMEM stage. In the fused MoE CTA warp 0 = TMA + // and warps 1..2 = NVSHMEM comm; warp 3 is the freed comm warp repurposed to + // issue UMMA (see PR #19 + mlp.cuh). Total: 12 warps / 3 hardware warpgroups. + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); // dedicated UMMA warp + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); // two aligned WGs + // Epilogue-local identity — valid only for warps 4..11. Warp 3 never enters + // the epilogue block, so its (garbage) tid_in_epi/wg/tid_wg are unused. + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; // 0|1 + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; // 0..127 + const int tmem_copy_tid = tid_wg; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; // 1 (WG0) or 2 (WG1) + // Warps 3..11 (MMA + both epilogue WGs) rendezvous for the acc-pipe init and + // the launcher-safe TMEM free. ConsumerThreads (256) covers only warps 4..11. + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, K-major SS ── + auto tiled_mma = make_tiled_mma( + SM100_MMA_F16BF16_SS{}); + auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 + + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto sW1 = make_tensor(make_smem_ptr(smem.W1_data()), typename Traits::SmemLayoutW{}); + auto sW2 = make_tensor(make_smem_ptr(smem.W2_data()), typename Traits::SmemLayoutW{}); + + // ── Two TMEM accumulators: U = X·B, V = X·C ── + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAccU = cta_mma.make_fragment_C(tCgC); + auto tCtAccV = cta_mma.make_fragment_C(tCgC); + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the + // alternate stage while both epilogue WGs drain MMA(n). producer_commit + // issues the UMMA-gated "accumulator ready" arrival; one elected epilogue + // thread releases the stage (after a 256-thread rendezvous) so the MMA + // warp may reuse it AccStages tiles later. ── + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; // one umma_arrive per commit (warp 3) + acc_params.consumer_arv_count = 1; // one elected releaser per stage + acc_params.initializing_warp = 4; // warp 4 inits the acc barriers + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + // Warps 3..11 rendezvous so the acc-pipe barriers (init'd by warp 4) are + // visible before the MMA warp (3) and both epilogue WGs (4..11) use them. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAccU.data() = tmem_base; + tCtAccV.data() = tmem_base + uint32_t(TileN); + + // ── Per-WG store slot (AtomTileM × EpiChunkN) — WG0 uses store_buf[0..S), + // WG1 uses [S..2S). Same smem layout as Hopper. ── + constexpr int store_slot_elems = AtomTileM * EpiChunkN; + Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + auto mZ = tma_store_z.get_tma_tensor( + make_shape(static_cast(num_z_m_tiles) * TileM, + static_cast(intermediate_dim))); + auto cta_tma_z = tma_store_z.get_slice(Int<0>{}); + + // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local + // tid). Each epi tile is the full TileM rows × EpiChunkN cols; the + // 128-thread tcgen05.ld maps one thread per TMEM datapath row. + // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) + // modes, which is what make_tmem_copy's cotiled builder requires. ── + auto epi_tile = make_tile(Int{}, Int{}); + // Extract the flat (M,N) view from the UMMA C-fragment (MMA,MMA_M,MMA_N) + // before tiling — matches CUTLASS's accumulators(make_coord(_,_),_0,_0). + auto accU_mn = tCtAccU(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto tAccU_epi = flat_divide(accU_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto t2r = make_tmem_copy(TmemLoadOp{}, tAccU_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tmem_copy_tid); + // Within-chunk (m,n) coords for the reg → smem-slot scatter. + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + // Register fragments sized from the DEST (partition_D) shape — per-thread + // (T2R,T2R_M,T2R_N). Sizing from partition_S(tmem) would wrongly include the + // warp-collective datapath-lane dim (the TMEM_LOAD atom distributes those + // lanes into per-thread registers internally). Matches CUTLASS's + // make_tensor(shape(tTR_sD)). + auto tTR_rU = make_tensor(shape(tTR_cChunk)); // f32 regs (U, reused) + auto tTR_rV = make_tensor(shape(tTR_cChunk)); // f32 regs (V) + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + CUTE_UNROLL + for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { + + // ── Mainloop (MMA warp = warp 3 only): bracket the k-loop with the + // accumulator pipeline. producer_acquire selects the free TMEM stage + // (AccStages=2 → MMA(n+1) fills the alternate stage while both epilogue + // WGs drain MMA(n)). consumer_release issues the UMMA-gated smem-buffer + // arrival on the mainloop pipe (no hand-rolled mbarrier). ── + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); // selected TMEM stage is free + int acc_stage = acc_prod_state.index(); + uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); + tCtAccU.data() = stage_base; + tCtAccV.data() = stage_base + uint32_t(TileN); + for (int k = 0; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + auto tCsX = cta_mma.partition_A(sX (_, _, state.index())); + auto tCsW1 = cta_mma.partition_B(sW1(_, _, state.index())); + auto tCsW2 = cta_mma.partition_B(sW2(_, _, state.index())); + auto tCrX = cta_mma.make_fragment_A(tCsX); + auto tCrW1 = cta_mma.make_fragment_B(tCsW1); + auto tCrW2 = cta_mma.make_fragment_B(tCsW2); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrX); ++kb) { + // First k-block of this n-tile clears TMEM; the rest accumulate. + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrX(_, _, kb), tCrW1(_, _, kb), tCtAccU); + gemm(tiled_mma, tCrX(_, _, kb), tCrW2(_, _, kb), tCtAccV); + } + pipe.consumer_release(state); // UMMA-gated smem-buffer release + ++state; + } + acc_pipe.producer_commit(acc_prod_state); // signal U/V ready (umma_arrive) + ++acc_prod_state; + } + + // ── Epilogue (two aligned WGs, warps 4..11): wait for the accumulator + // stage, then THIS WG drains its half of TileN (chunks + // [wg·NChunksHalf, +NChunksHalf)) while the MMA warp fills the + // alternate stage. ── + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + uint32_t stage_base = tmem_base + uint32_t(acc_stage * (2 * TileN)); + tCtAccU.data() = stage_base; + tCtAccV.data() = stage_base + uint32_t(TileN); + auto accU_mn_stage = tCtAccU(make_coord(_, _), _0{}, _0{}); + auto accV_mn_stage = tCtAccV(make_coord(_, _), _0{}, _0{}); + auto tAccU_epi_stage = flat_divide(accU_mn_stage, epi_tile); + auto tAccV_epi_stage = flat_divide(accV_mn_stage, epi_tile); + auto tTR_tAccU = thr_t2r.partition_S(tAccU_epi_stage); + auto tTR_tAccV = thr_t2r.partition_S(tAccV_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; // absolute n-chunk index + + // TMEM → registers (this chunk, full TileM rows). + copy(t2r, tTR_tAccU(_, _, _, _0{}, chunk), tTR_rU); + copy(t2r, tTR_tAccV(_, _, _, _0{}, chunk), tTR_rV); + cutlass::arch::fence_view_async_tmem_load(); + + // SiLU(U)·V fused in place (overwrites the U regs) so it overlaps + // the prior in-flight store. The cast to Element happens at the + // smem write below — no extra register set. + CUTE_UNROLL + for (int i = 0; i < size(tTR_rU); ++i) + tTR_rU(i) = fast_silu(tTR_rU(i)) * tTR_rV(i); + + // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rU); ++i) { + int m_row = get<0>(tTR_cChunk(i)); // 0..TileM + int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN + if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) + sStore(m_row - ms * AtomTileM, n_col) = + static_cast(tTR_rU(i)); + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx = MSub * (z_m + m_sub) + ms; + int n_tile_idx = n * (TileN / EpiChunkN) + chunk; + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_z, cta_tma_z.partition_S(sStore), + cta_tma_z.partition_D(gZ)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // Both epilogue WGs done reading this TMEM stage → rendezvous, then + // one elected thread releases the stage so the MMA warp may reuse it + // when the AccStages accumulator pipeline wraps. + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + // Final rendezvous of the MMA warp (3) + both epilogue WGs (4..11) so the + // outer launcher can safely free TMEM after this consumer returns. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + // The Compute=100 specialization is never dispatched on non-SM100 targets + // (host picks Compute=90 there). Trap if it is ever reached. + __trap(); +#endif +} +}; // Mlp1FusedConsumerImpl<100> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — keeps the existing call style. `Compute` defaults to 90 so +// current call sites (mlp1_fused_consumer(...)) are unchanged; +// pass mlp1_fused_consumer(...) to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp1_fused_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp1FusedSmem& smem, + TmaStoreZ const& tma_store_z, + int z_m, + int intermediate_dim, + int num_z_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx = -1, + int num_splits = -1) { + Mlp1FusedConsumerImpl::template run< + Traits, MSubTiles, Pipeline, TmaStoreZ>( + pipe, state, smem, tma_store_z, z_m, intermediate_dim, + num_z_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp2_fused.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp2_fused.cuh index 3e8dd89d..770dc872 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp2_fused.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp2_fused.cuh @@ -1,830 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Phase 2 (single-tile, fused): Y = Z @ A^T (down-projection) -// ═══════════════════════════════════════════════════════════════════ -// -// Single (TileM, TileN) producer/consumer building blocks for callers -// that drive the outer M-tile loop themselves: the standalone MLP2 -// benchmark (mlp2.cu) and the fused MLP kernel (mlp.cuh). -// -// Design: -// - Cooperative 2-WG consumer; split axis selected by TileM: -// TileM=128 → Layout<_2,_1,_1>, split along M (each WG = 64×TileN) -// TileM=64 → Layout<_1,_2,_1>, split along N (each WG = 64×TileN/2) -// - One accumulator per WG -// - Per-WG store_buf with EpiChunkN-round epilogue -// - K_PIPE_MMAS = 1 (CUTLASS-canonical, see -// sm90_mma_tma_gmma_ss_warpspecialized.hpp:264) -// - Single fused Z+W TMA pipe (1 mbarrier per stage, 2 TMA copies) -// -// CTA layout (384 threads = 12 warps): -// WG0 warp 0 : Producer (TMA loads for Z / A) -// WG0 warps 1-2 : Idle / reserved for fused-kernel communication -// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 -// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 -// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 -// -// Per `expert_ids`: one int per TileM-row m-block. -// -// Mlp2FusedSmem omits the pipeline storage so callers can fuse it with -// other phases' smem (see mlp.cuh). The benchmark wrapper (mlp2.cu) is -// responsible for attaching MainloopPipeline::SharedStorage and running -// the outer for(m) loop over m-tiles. -// -// ═══════════════════════════════════════════════════════════════════ - -#include "models.cuh" -#include "math.cuh" -#include "tmem_load_op.cuh" - -#include -#include -#include -#include -#include - -// SM100 (Blackwell / UMMA) support for the Compute=100 consumer -// specialization. These headers are header-only and safe to include under -// an sm_90a build — the tcgen05 PTX they emit is guarded by -// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on -// __CUDA_ARCH__ >= 1000) ever instantiates it. -#include -#include - -#include -#include -#include // PipelineTmaUmmaAsync, PipelineUmmaAsync - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// Traits -// ═══════════════════════════════════════════════════════════════════ - -template < - typename Element_, - int TileM_ = 128, - int TileN_ = 128, - int TileK_ = 64, - int Stages_ = 4, - int EpiChunkN_ = 32 -> -struct Mlp2Traits { - using Element = Element_; - using ElementAccum = float; - - static constexpr int TileM = TileM_; - static constexpr int TileN = TileN_; - static constexpr int TileK = TileK_; - static constexpr int Stages = Stages_; - static constexpr int EpiChunkN = EpiChunkN_; - - static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); - - // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate - // along M (each owns 64 rows × TileN). With TileM=64 the WGs - // cooperate along N (each owns 64 rows × TileN/2). AtomTileM is the - // per-WG M extent — always 64. - static constexpr bool kMSplit = (TileM == 128); - static constexpr int AtomTileM = 64; - static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); - - static_assert(TileN <= 256, "TileN must fit TMA box limit"); - static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); - static_assert(kMSplit || (TileN % 2 == 0), - "TileN must be even when TileM=64 (cooperative N-split halves)"); - static constexpr int NumEpiRounds = WgTileN / EpiChunkN; - - // Cooperative 2-WG layout: M-split for TileM=128, N-split for TileM=64. - // AtomN is the per-WG WGMMA atom N — combined across both WGs it must - // cover exactly TileN cols (= sW's N extent). With WgLayout=<1,2,1> - // and AtomN=TileN, the TiledMMA would cover 2·TileN cols and WG1's - // partition_B(sW) slice would read past sW's TileN-wide buffer → IMA. - static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); - using GmmaAtom = typename GmmaSelector::Atom; - using WgLayout = cute::conditional_t< - kMSplit, - Layout>, - Layout>>; - using TiledMma = TiledMMA, WgLayout>; - - using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom; - - using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtom{}, - Shape, Int>{})); - using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{}, - Shape, Int>{})); - - using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtom{}, - Shape, Int, Int>{})); - using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{}, - Shape, Int, Int>{})); - - // Per-WG store buffer slot: (AtomTileM, EpiChunkN). - using SmemLayoutStoreSlot = Layout< - Shape, Int>, - Stride, _1>>; - using SmemLayoutStore_1 = SmemLayoutStoreSlot; - - using MainloopPipeline = cutlass::PipelineTmaAsync; - using PipelineState = cutlass::PipelineState; - - // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop - // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the - // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the single - // TMEM accumulator from the MMA warp to the epilogue warps. See - // Mlp2FusedConsumerImpl<100>. The Blackwell launcher must build the mainloop - // pipe with num_consumers=1 (one umma_arrive per stage). - using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< - Stages, cute::Shape, - cute::Shape>; - // Optional accumulator double-buffering: each stage owns one TileN-wide - // TMEM accumulator, allowing MMA(n+1) to fill one stage while epilogue - // drains/TMA-stores the previous stage. - static constexpr int AccStages = 2; - static_assert(AccStages * TileN <= 512, - "SM100 MLP2 accumulator stages must fit in 512 TMEM columns"); - using AccumulatorPipeline = cutlass::PipelineUmmaAsync; - - static constexpr int TmaTransBytesZ = - static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); - static constexpr int TmaTransBytesW = - static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); - // Fused pipeline: Z + W share one mbarrier per stage. - static constexpr int TmaTransBytes = TmaTransBytesZ + TmaTransBytesW; - - static constexpr int WarpSize = 32; - static constexpr int WarpGroupSize = 128; - static constexpr int NumConsumers = 2; - static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; - static constexpr int NumThreads = 384; -}; - -template -using Mlp2MainloopPipelineFor = cute::conditional_t< - Compute == 100, - typename Traits::MainloopPipelineUmma, - typename Traits::MainloopPipeline>; - -// ═══════════════════════════════════════════════════════════════════ -// Fused MLP2 shared memory — no pipeline storage (managed externally) -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp2FusedSmem { - using Element = typename Traits::Element; - - static constexpr int smem_Z_size = cosize_v; - static constexpr int smem_W_size = cosize_v; - static constexpr int smem_store_size = cosize_v; - - alignas(128) Element smem_Z[smem_Z_size]; - alignas(128) Element smem_W[smem_W_size]; - // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. - alignas(128) Element store_buf[2 * smem_store_size]; - - // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM - // base address, plus the accumulator pipeline's barrier storage (UMMA→ - // epilogue handoff). The Hopper (Compute=90) path never touches these; - // SM100 launchers allocate/free TMEM outside the consumer and publish the - // base here. - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; - - CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } - CUTE_DEVICE Element* W_data() { return &smem_W[0]; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// Pipeline construction helper — single fused Z+W pipeline -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ auto mlp2_make_pipe( - typename Traits::MainloopPipeline::SharedStorage& storage) { - - using Pipeline = typename Traits::MainloopPipeline; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 4 && warp_id <= 11); - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = Traits::ConsumerThreads; - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// SM100 (Compute=100) mainloop pipe. Mirrors mlp2_make_pipe but builds the -// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues -// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer -// is released by a single umma_arrive per stage (the MMA warp), not per -// consumer thread. Params/PipelineState/SharedStorage alias the Hopper -// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer -// (mlp2_fused_producer) and the Compute=100 consumer drive it unchanged. -template -__device__ __forceinline__ auto mlp2_make_pipe_umma( - typename Traits::MainloopPipelineUmma::SharedStorage& storage) { - - using Pipeline = typename Traits::MainloopPipelineUmma; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// ═══════════════════════════════════════════════════════════════════ -// Fused producer — single (m), all (n) under blockIdx.y stride. -// Single fused Z + W TMA pipe per k-step. -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void mlp2_fused_producer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2FusedSmem& smem, - TmaLoadZ const& tma_load_z, - TmaLoadW const& tma_load_a, - int m, - int expert_or_n_offset, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int num_experts, - int total_n_rows, - int num_n_tiles, - int num_k_tiles, - // N-split identity (flat-grid launch); -1 → blockIdx.y / gridDim.y. - int split_idx = -1, - int num_splits = -1) { - - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); - - // int64_t cast: total_n_rows × intermediate_dim and num_tokens × - // intermediate_dim both exceed INT_MAX at production shapes. CUTE's - // layout math otherwise wraps to a negative byte offset → IMA. - auto mZ = tma_load_z.get_tma_tensor(make_shape( - static_cast(num_tokens), - static_cast(intermediate_dim))); - auto mA = [&]() { - if constexpr (Expert3D) { - return tma_load_a.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(intermediate_dim), - static_cast(num_experts))); - } else { - return tma_load_a.get_tma_tensor(make_shape( - static_cast(total_n_rows), - static_cast(intermediate_dim))); - } - }(); - - auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); - auto cta_tma_a = tma_load_a.get_slice(Int<0>{}); - - auto tZsZ = cta_tma_z.partition_D(sZ); - auto tWsW = cta_tma_a.partition_D(sW); - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - auto gW = [&]() { - if constexpr (Expert3D) { - return local_tile(mA, - make_tile(Int{}, Int{}), - make_coord(n, _, expert_or_n_offset)); - } else { - return local_tile(mA, - make_tile(Int{}, Int{}), - make_coord(expert_or_n_offset + n, _)); - } - }(); - auto tWgW = cta_tma_a.partition_S(gW); - - CUTE_UNROLL - for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { - auto gZ = local_tile(mZ, - make_tile(Int{}, Int{}), - make_coord(m + m_sub, _)); - auto tZgZ = cta_tma_z.partition_S(gZ); - - for (int k = 0; k < num_k_tiles; ++k) { - pipe.producer_acquire(state); - if (threadIdx.x == 0) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_z.with(*bar, 0), - tZgZ(_, _, _, k), tZsZ(_, _, _, state.index())); - copy(tma_load_a.with(*bar, 0), - tWgW(_, _, _, k), tWsW(_, _, _, state.index())); - } - ++state; - } - } - } -} - -// ═══════════════════════════════════════════════════════════════════ -// Fused consumer — architecture-specialized on `int Compute` -// ═══════════════════════════════════════════════════════════════════ -// -// Primary template is undefined; one full specialization per supported -// compute capability provides a `run(...)` member (a function-template on -// Traits/Pipeline/TmaStoreY, since function templates can't be partially -// specialized). The free function `mlp2_fused_consumer` below forwards to -// the right specialization — call sites stay in the existing style. -// -// Compute=90 → Hopper / WGMMA (cooperative 2-WG, single register acc) -// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulator) -// -// y_m is the M-tile index into the output Y buffer. - -template -struct Mlp2FusedConsumerImpl; - -// ─────────────────────────────────────────────────────────────────── -// Compute=90 — Hopper. The two WGs split the (TileM, TileN) output along -// M: WG1 writes rows [0, AtomTileM), WG2 writes [AtomTileM, TileM). -// (Verbatim from the original mlp2_fused_consumer.) -// ─────────────────────────────────────────────────────────────────── - -template <> -struct Mlp2FusedConsumerImpl<90> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2FusedSmem& smem, - TmaStoreY const& tma_store_y, - int y_m, // M-tile index into Y buffer - int hidden_dim, - int num_y_m_tiles, // number of M-tile slots in Y buffer - int num_n_tiles, - int num_k_tiles, - // N-split identity (flat-grid launch); -1 → blockIdx.y / gridDim.y. - int split_idx, - int num_splits, - int valid_m64_subtiles) { - - static_assert(MSubTiles == 1, - "Hopper MLP2 consumer processes one GEMM M-tile per invocation"); - using Element = typename Traits::Element; - constexpr int kM64Subtiles = Traits::TileM / 64; - int valid_store_subtiles = valid_m64_subtiles < 0 - ? kM64Subtiles : min(kM64Subtiles, valid_m64_subtiles); - typename Traits::TiledMma tiled_mma; - int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; - int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; - auto thr_mma = tiled_mma.get_slice(tid_in_mma); - - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); - - auto tCsZ = thr_mma.partition_A(sZ); - auto tCsW = thr_mma.partition_B(sW); - - auto acc = partition_fragment_C(tiled_mma, - Shape, Int>{}); - - auto cC = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCcC = thr_mma.partition_C(cC); - - const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; - const int my_barrier_id = 1 + my_wg; - const bool is_my_wg_leader = (tid_in_wg == 0); - - constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; - Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - // int64_t cast: num_y_m_tiles · TileM · hidden_dim can exceed - // INT_MAX at large output buffers; cast keeps CUTE's layout math - // in 64-bit so byte offsets don't wrap. - auto mY = tma_store_y.get_tma_tensor( - make_shape(static_cast(num_y_m_tiles) * Traits::TileM, - static_cast(hidden_dim))); - auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - bool store_in_flight = false; - constexpr int K_PIPE_MMAS = 1; // CUTLASS-matching - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - - clear(acc); - - auto state_release = state; - int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; - - // Prologue - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - ++state; - } - // Steady state - for (int k = prologue_count; k < num_k_tiles; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - - warpgroup_wait(); - warpgroup_fence_operand(acc); - pipe.consumer_release(state_release); - ++state; - ++state_release; - } - // Drain - warpgroup_wait<0>(); - warpgroup_fence_operand(acc); - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_release(state_release); - ++state_release; - } - - // ── Epilogue ────────────────────────────────────── - // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. - // TileM=64 (N-split): WG_w owns full M, cols [w*WgTileN, ...). - bool store_output = Traits::kMSplit - ? (my_wg < valid_store_subtiles) - : (valid_store_subtiles > 0); - if (store_output) { - CUTE_UNROLL - for (int r = 0; r < Traits::NumEpiRounds; ++r) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - int chunk_start = r * Traits::EpiChunkN; - CUTE_UNROLL - for (int i = 0; i < size(acc); ++i) { - auto coord = tCcC(i); - int m_loc = get<0>(coord); - int n_loc = get<1>(coord); - int m_local, n_local; - if constexpr (Traits::kMSplit) { - m_local = m_loc - my_wg * Traits::AtomTileM; - n_local = n_loc; - } else { - m_local = m_loc; - n_local = n_loc - my_wg * Traits::WgTileN; - } - if (n_local >= chunk_start && - n_local < chunk_start + Traits::EpiChunkN) { - int chunk_n = n_local - chunk_start; - sStore(m_local, chunk_n) = static_cast(acc(i)); - } - } - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - if (is_my_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx, n_tile_idx; - if constexpr (Traits::kMSplit) { - // TileM=128: 2 row tiles per logical M tile. - m_tile_idx = 2 * y_m + my_wg; - n_tile_idx = n * Traits::NumEpiRounds + r; - } else { - // TileM=64: 1 row tile per logical M tile; WG_w handles - // its N-half. - m_tile_idx = y_m; - n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) - + my_wg * Traits::NumEpiRounds + r; - } - auto gY = local_tile(mY, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_y, cta_tma_y.partition_S(sStore), - cta_tma_y.partition_D(gY)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - } - if (store_in_flight) - cute::tma_store_wait<0>(); -} -}; // Mlp2FusedConsumerImpl<90> - -// ─────────────────────────────────────────────────────────────────── -// Compute=100 — Blackwell / UMMA. -// -// No cooperative-warpgroup machinery (UMMA has no warpgroup concept): -// - One 1SM UMMA atom with M=TileM covers the whole tile; the single -// accumulator Y = Z·Aᵀ lives in TMEM (no register/M-split). -// - Warp 3 issues UMMA and stays epilogue-free; warps 4-11 form the two -// aligned epilogue WGs. -// - Epilogue: the two consumer warpgroups split TileN; each reads its -// N-half from TMEM→regs in EpiChunkN chunks, casts to bf16, and -// stores 64-row (AtomTileM) tiles via the existing reg→SMEM→TMA path -// (so the host TMA atom is unchanged). No activation fusion. -// -// TileM is the template parameter as on Hopper (64 or 128 — both native -// 1SM UMMA M sizes). For TileM=128 the epilogue stores MSub=2 row-tiles. -// ─────────────────────────────────────────────────────────────────── - -template <> -struct Mlp2FusedConsumerImpl<100> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2FusedSmem& smem, - TmaStoreY const& tma_store_y, - int y_m, - int hidden_dim, - int num_y_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits, - int valid_m64_subtiles) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using Element = typename Traits::Element; - constexpr int TileM = Traits::TileM; - constexpr int TileN = Traits::TileN; - constexpr int EpiChunkN = Traits::EpiChunkN; - constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile - static_assert(TileN % 2 == 0, - "Blackwell consumer splits TileN across the two consumer warpgroups"); - constexpr int WgN = TileN / 2; // per-warpgroup N width - static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); - constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per warpgroup - constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) - int valid_store_subtiles = valid_m64_subtiles < 0 - ? MSub : min(MSub, valid_m64_subtiles); - - // ── Thread identity ───────────────────────────────────── - // Warp 3 is the dedicated, epilogue-free UMMA producer. Warps 4..11 are - // the two aligned epilogue WGs. This mirrors MLP1's Blackwell path. - const int warp_id = threadIdx.x / Traits::WarpSize; - const bool is_mma_warp = (warp_id == 3); - const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); - const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 - const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; - const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; - const int tmem_copy_tid = tid_wg; - const bool is_wg_leader = is_epilogue && (tid_wg == 0); - const int wg_barrier_id = 1 + wg; // 1 or 2 - constexpr int kEpilogueThreads = Traits::ConsumerThreads; - constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; - static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); - static_assert(kEpilogueThreads == 8 * Traits::WarpSize); - static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); - - // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, K-major SS. - // Operand A = Z, operand B = A (the down weight) — both K-major, the - // SM100-native orientation, so no transpose subtlety. ── - auto tiled_mma = make_tiled_mma( - SM100_MMA_F16BF16_SS{}); - auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 - - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); - - // ── One TMEM accumulator: Y = Z·Aᵀ ── - auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCgC = cta_mma.partition_C(cAccFull); - auto tCtAcc = cta_mma.make_fragment_C(tCgC); - - // TMEM is allocated by the outer fused/standalone launcher once per CTA. - - // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers - // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the - // alternate stage while epilogue drains/TMA-stores MMA(n). ── - using AccPipe = typename Traits::AccumulatorPipeline; - typename AccPipe::Params acc_params; - acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer - : AccPipe::ThreadCategory::Consumer; - acc_params.producer_arv_count = 1; // one umma_arrive per commit - acc_params.consumer_arv_count = 1; // one elected releaser - acc_params.initializing_warp = 4; // warp 4 inits the acc barriers - AccPipe acc_pipe(smem.acc_pipe, acc_params, - cute::Shape{}); - auto acc_prod_state = cutlass::make_producer_start_state(); - typename AccPipe::PipelineState acc_cons_state; - - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); - const uint32_t tmem_base = smem.tmem_base; - tCtAcc.data() = tmem_base; - - // ── Per-WG store slot (AtomTileM × EpiChunkN) — same smem as Hopper ── - constexpr int store_slot_elems = AtomTileM * EpiChunkN; - Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - // int64_t cast: num_y_m_tiles · TileM · hidden_dim can exceed INT_MAX at - // large output buffers; cast keeps CUTE's layout math in 64-bit. - auto mY = tma_store_y.get_tma_tensor( - make_shape(static_cast(num_y_m_tiles) * TileM, - static_cast(hidden_dim))); - auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); - - // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local - // tid). Each epi tile is the full TileM rows × EpiChunkN cols; the - // 128-thread tcgen05.ld maps one thread per TMEM datapath row. - // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) - // modes, which is what make_tmem_copy's cotiled builder requires. ── - auto epi_tile = make_tile(Int{}, Int{}); - // Extract the flat (M,N) view from the UMMA C-fragment (MMA,MMA_M,MMA_N) - // before tiling — matches CUTLASS's accumulators(make_coord(_,_),_0,_0). - auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) - auto tAcc_epi = flat_divide(acc_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) - auto t2r = make_tmem_copy(TmemLoadOp{}, tAcc_epi(_, _, _0{}, _0{})); - auto thr_t2r = t2r.get_slice(tmem_copy_tid); - auto tTR_tAcc = thr_t2r.partition_S(tAcc_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) - // Within-chunk (m,n) coords for the reg → smem-slot scatter. - auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); - auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) - // Register fragment sized from the DEST (partition_D) shape — per-thread - // (T2R,T2R_M,T2R_N). Sizing from partition_S(tmem) would wrongly include the - // warp-collective datapath-lane dim. Matches CUTLASS's - // make_tensor(shape(tTR_sD)). - auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); // f32 regs - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - bool store_in_flight = false; - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - CUTE_UNROLL - for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { - - // ── Mainloop (warp 3 only): bracket the k-loop with the accumulator - // pipeline. Each k-stage waits/releases the TMA→UMMA mainloop - // pipeline; consumer_release issues the UMMA-gated smem-buffer - // arrival (no hand-rolled mbarrier). (Conservative: no MMA/k overlap.) - if (is_mma_warp) { - acc_pipe.producer_acquire(acc_prod_state); // TMEM acc free (prev epilogue done) - int acc_stage = acc_prod_state.index(); - tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); - for (int k = 0; k < num_k_tiles; ++k) { - pipe.consumer_wait(state); - auto tCsZ = cta_mma.partition_A(sZ(_, _, state.index())); - auto tCsW = cta_mma.partition_B(sW(_, _, state.index())); - auto tCrZ = cta_mma.make_fragment_A(tCsZ); - auto tCrW = cta_mma.make_fragment_B(tCsW); - CUTE_UNROLL - for (int kb = 0; kb < size<2>(tCrZ); ++kb) { - // First k-block of this n-tile clears TMEM; the rest accumulate. - tiled_mma.accumulate_ = (k == 0 && kb == 0) - ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; - gemm(tiled_mma, tCrZ(_, _, kb), tCrW(_, _, kb), tCtAcc); - } - pipe.consumer_release(state); // UMMA-gated smem-buffer release - ++state; - } - acc_pipe.producer_commit(acc_prod_state); // signal acc ready (umma_arrive) - ++acc_prod_state; - } - - // ── Epilogue: wait for the accumulator, then this WG processes its - // n-chunks [wg·NChunksHalf, +NChunksHalf). ── - if (is_epilogue) { - acc_pipe.consumer_wait(acc_cons_state); - int acc_stage = acc_cons_state.index(); - tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); - auto acc_mn_stage = tCtAcc(make_coord(_, _), _0{}, _0{}); - auto tAcc_epi_stage = flat_divide(acc_mn_stage, epi_tile); - auto tTR_tAcc_stage = thr_t2r.partition_S(tAcc_epi_stage); - - CUTE_UNROLL - for (int r = 0; r < NChunksHalf; ++r) { - int chunk = wg * NChunksHalf + r; // absolute n-chunk index - - // TMEM → registers (this chunk, full TileM rows). - copy(t2r, tTR_tAcc_stage(_, _, _, _0{}, chunk), tTR_rAcc); - cutlass::arch::fence_view_async_tmem_load(); - - // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). - CUTE_UNROLL - for (int ms = 0; ms < MSub; ++ms) { - if (ms >= valid_store_subtiles) - continue; - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - CUTE_UNROLL - for (int i = 0; i < size(tTR_rAcc); ++i) { - int m_row = get<0>(tTR_cChunk(i)); // 0..TileM - int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN - if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) - sStore(m_row - ms * AtomTileM, n_col) = - static_cast(tTR_rAcc(i)); - } - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - - if (is_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx = MSub * (y_m + m_sub) + ms; - int n_tile_idx = n * (TileN / EpiChunkN) + chunk; - auto gY = local_tile(mY, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_y, cta_tma_y.partition_S(sStore), - cta_tma_y.partition_D(gY)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - - // All epilogue TMEM reads are done; release the accumulator so the next - // n-tile's MMA may reuse it (one elected consumer thread arrives). - cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); - if (tid_in_epi == 0) - acc_pipe.consumer_release(acc_cons_state); - ++acc_cons_state; - } - } - } - if (is_epilogue && store_in_flight) - cute::tma_store_wait<0>(); - - // Keep the final MMA+epilogue barrier so the outer launcher can safely free - // TMEM after this consumer returns. - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); -#else - // The Compute=100 specialization is never dispatched on non-SM100 targets - // (host picks Compute=90 there). Trap if it is ever reached. - __trap(); +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 #endif -} -}; // Mlp2FusedConsumerImpl<100> -// ─────────────────────────────────────────────────────────────────── -// Forwarder — keeps the existing call style. `Compute` defaults to 90 so -// current call sites (mlp2_fused_consumer(...)) are unchanged; -// pass mlp2_fused_consumer(...) to select the Blackwell path. -// ─────────────────────────────────────────────────────────────────── - -template -__device__ __forceinline__ void mlp2_fused_consumer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2FusedSmem& smem, - TmaStoreY const& tma_store_y, - int y_m, - int hidden_dim, - int num_y_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx = -1, - int num_splits = -1, - int valid_m64_subtiles = -1) { - Mlp2FusedConsumerImpl::template run( - pipe, state, smem, tma_store_y, y_m, hidden_dim, - num_y_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits, - valid_m64_subtiles); -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp2_fused_sm90.cuh" +#else +#include "mlp2_fused_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm100.cuh new file mode 100644 index 00000000..968019f7 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm100.cuh @@ -0,0 +1,830 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 2 (single-tile, fused): Y = Z @ A^T (down-projection) +// ═══════════════════════════════════════════════════════════════════ +// +// Single (TileM, TileN) producer/consumer building blocks for callers +// that drive the outer M-tile loop themselves: the standalone MLP2 +// benchmark (mlp2.cu) and the fused MLP kernel (mlp.cuh). +// +// Design: +// - Cooperative 2-WG consumer; split axis selected by TileM: +// TileM=128 → Layout<_2,_1,_1>, split along M (each WG = 64×TileN) +// TileM=64 → Layout<_1,_2,_1>, split along N (each WG = 64×TileN/2) +// - One accumulator per WG +// - Per-WG store_buf with EpiChunkN-round epilogue +// - K_PIPE_MMAS = 1 (CUTLASS-canonical, see +// sm90_mma_tma_gmma_ss_warpspecialized.hpp:264) +// - Single fused Z+W TMA pipe (1 mbarrier per stage, 2 TMA copies) +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for Z / A) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Per `expert_ids`: one int per TileM-row m-block. +// +// Mlp2FusedSmem omits the pipeline storage so callers can fuse it with +// other phases' smem (see mlp.cuh). The benchmark wrapper (mlp2.cu) is +// responsible for attaching MainloopPipeline::SharedStorage and running +// the outer for(m) loop over m-tiles. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm100.cuh" +#include "math_sm100.cuh" +#include "tmem_load_op_sm100.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer +// specialization. These headers are header-only and safe to include under +// an sm_90a build — the tcgen05 PTX they emit is guarded by +// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on +// __CUDA_ARCH__ >= 1000) ever instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 128, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 32 +> +struct Mlp2Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + + static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); + + // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate + // along M (each owns 64 rows × TileN). With TileM=64 the WGs + // cooperate along N (each owns 64 rows × TileN/2). AtomTileM is the + // per-WG M extent — always 64. + static constexpr bool kMSplit = (TileM == 128); + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); + + static_assert(TileN <= 256, "TileN must fit TMA box limit"); + static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); + static_assert(kMSplit || (TileN % 2 == 0), + "TileN must be even when TileM=64 (cooperative N-split halves)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + // Cooperative 2-WG layout: M-split for TileM=128, N-split for TileM=64. + // AtomN is the per-WG WGMMA atom N — combined across both WGs it must + // cover exactly TileN cols (= sW's N extent). With WgLayout=<1,2,1> + // and AtomN=TileN, the TiledMMA would cover 2·TileN cols and WG1's + // partition_B(sW) slice would read past sW's TileN-wide buffer → IMA. + static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); + using GmmaAtom = typename GmmaSelector::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom; + + using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int>{})); + using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int>{})); + + using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int, Int>{})); + using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int, Int>{})); + + // Per-WG store buffer slot: (AtomTileM, EpiChunkN). + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore_1 = SmemLayoutStoreSlot; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop + // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the + // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the single + // TMEM accumulator from the MMA warp to the epilogue warps. See + // Mlp2FusedConsumerImpl<100>. The Blackwell launcher must build the mainloop + // pipe with num_consumers=1 (one umma_arrive per stage). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Optional accumulator double-buffering: each stage owns one TileN-wide + // TMEM accumulator, allowing MMA(n+1) to fill one stage while epilogue + // drains/TMA-stores the previous stage. + static constexpr int AccStages = 2; + static_assert(AccStages * TileN <= 512, + "SM100 MLP2 accumulator stages must fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + static constexpr int TmaTransBytesZ = + static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesW = + static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); + // Fused pipeline: Z + W share one mbarrier per stage. + static constexpr int TmaTransBytes = TmaTransBytesZ + TmaTransBytesW; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; +}; + +template +using Mlp2MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP2 shared memory — no pipeline storage (managed externally) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp2FusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element smem_W[smem_W_size]; + // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these; + // SM100 launchers allocate/free TMEM outside the consumer and publish the + // base here. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } + CUTE_DEVICE Element* W_data() { return &smem_W[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Pipeline construction helper — single fused Z+W pipeline +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ auto mlp2_make_pipe( + typename Traits::MainloopPipeline::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp2_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer +// (mlp2_fused_producer) and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ auto mlp2_make_pipe_umma( + typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused producer — single (m), all (n) under blockIdx.y stride. +// Single fused Z + W TMA pipe per k-step. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_fused_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2FusedSmem& smem, + TmaLoadZ const& tma_load_z, + TmaLoadW const& tma_load_a, + int m, + int expert_or_n_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_n_rows, + int num_n_tiles, + int num_k_tiles, + // N-split identity (flat-grid launch); -1 → blockIdx.y / gridDim.y. + int split_idx = -1, + int num_splits = -1) { + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + // int64_t cast: total_n_rows × intermediate_dim and num_tokens × + // intermediate_dim both exceed INT_MAX at production shapes. CUTE's + // layout math otherwise wraps to a negative byte offset → IMA. + auto mZ = tma_load_z.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(intermediate_dim))); + auto mA = [&]() { + if constexpr (Expert3D) { + return tma_load_a.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(num_experts))); + } else { + return tma_load_a.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(intermediate_dim))); + } + }(); + + auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); + auto cta_tma_a = tma_load_a.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_z.partition_D(sZ); + auto tWsW = cta_tma_a.partition_D(sW); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + auto gW = [&]() { + if constexpr (Expert3D) { + return local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto tWgW = cta_tma_a.partition_S(gW); + + CUTE_UNROLL + for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m + m_sub, _)); + auto tZgZ = cta_tma_z.partition_S(gZ); + + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_z.with(*bar, 0), + tZgZ(_, _, _, k), tZsZ(_, _, _, state.index())); + copy(tma_load_a.with(*bar, 0), + tWgW(_, _, _, k), tWsW(_, _, _, state.index())); + } + ++state; + } + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused consumer — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported +// compute capability provides a `run(...)` member (a function-template on +// Traits/Pipeline/TmaStoreY, since function templates can't be partially +// specialized). The free function `mlp2_fused_consumer` below forwards to +// the right specialization — call sites stay in the existing style. +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, single register acc) +// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulator) +// +// y_m is the M-tile index into the output Y buffer. + +template +struct Mlp2FusedConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. The two WGs split the (TileM, TileN) output along +// M: WG1 writes rows [0, AtomTileM), WG2 writes [AtomTileM, TileM). +// (Verbatim from the original mlp2_fused_consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp2FusedConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2FusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, // M-tile index into Y buffer + int hidden_dim, + int num_y_m_tiles, // number of M-tile slots in Y buffer + int num_n_tiles, + int num_k_tiles, + // N-split identity (flat-grid launch); -1 → blockIdx.y / gridDim.y. + int split_idx, + int num_splits, + int valid_m64_subtiles) { + + static_assert(MSubTiles == 1, + "Hopper MLP2 consumer processes one GEMM M-tile per invocation"); + using Element = typename Traits::Element; + constexpr int kM64Subtiles = Traits::TileM / 64; + int valid_store_subtiles = valid_m64_subtiles < 0 + ? kM64Subtiles : min(kM64Subtiles, valid_m64_subtiles); + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; + const int my_barrier_id = 1 + my_wg; + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // int64_t cast: num_y_m_tiles · TileM · hidden_dim can exceed + // INT_MAX at large output buffers; cast keeps CUTE's layout math + // in 64-bit so byte offsets don't wrap. + auto mY = tma_store_y.get_tma_tensor( + make_shape(static_cast(num_y_m_tiles) * Traits::TileM, + static_cast(hidden_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; // CUTLASS-matching + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. + // TileM=64 (N-split): WG_w owns full M, cols [w*WgTileN, ...). + bool store_output = Traits::kMSplit + ? (my_wg < valid_store_subtiles) + : (valid_store_subtiles > 0); + if (store_output) { + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: 2 row tiles per logical M tile. + m_tile_idx = 2 * y_m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M tile; WG_w handles + // its N-half. + m_tile_idx = y_m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp2FusedConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. +// +// No cooperative-warpgroup machinery (UMMA has no warpgroup concept): +// - One 1SM UMMA atom with M=TileM covers the whole tile; the single +// accumulator Y = Z·Aᵀ lives in TMEM (no register/M-split). +// - Warp 3 issues UMMA and stays epilogue-free; warps 4-11 form the two +// aligned epilogue WGs. +// - Epilogue: the two consumer warpgroups split TileN; each reads its +// N-half from TMEM→regs in EpiChunkN chunks, casts to bf16, and +// stores 64-row (AtomTileM) tiles via the existing reg→SMEM→TMA path +// (so the host TMA atom is unchanged). No activation fusion. +// +// TileM is the template parameter as on Hopper (64 or 128 — both native +// 1SM UMMA M sizes). For TileM=128 the epilogue stores MSub=2 row-tiles. +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp2FusedConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2FusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, + int hidden_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits, + int valid_m64_subtiles) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two consumer warpgroups"); + constexpr int WgN = TileN / 2; // per-warpgroup N width + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per warpgroup + constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) + int valid_store_subtiles = valid_m64_subtiles < 0 + ? MSub : min(MSub, valid_m64_subtiles); + + // ── Thread identity ───────────────────────────────────── + // Warp 3 is the dedicated, epilogue-free UMMA producer. Warps 4..11 are + // the two aligned epilogue WGs. This mirrors MLP1's Blackwell path. + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + const int tmem_copy_tid = tid_wg; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; // 1 or 2 + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, K-major SS. + // Operand A = Z, operand B = A (the down weight) — both K-major, the + // SM100-native orientation, so no transpose subtlety. ── + auto tiled_mma = make_tiled_mma( + SM100_MMA_F16BF16_SS{}); + auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + // ── One TMEM accumulator: Y = Z·Aᵀ ── + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the + // alternate stage while epilogue drains/TMA-stores MMA(n). ── + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; // one umma_arrive per commit + acc_params.consumer_arv_count = 1; // one elected releaser + acc_params.initializing_warp = 4; // warp 4 inits the acc barriers + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + // ── Per-WG store slot (AtomTileM × EpiChunkN) — same smem as Hopper ── + constexpr int store_slot_elems = AtomTileM * EpiChunkN; + Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // int64_t cast: num_y_m_tiles · TileM · hidden_dim can exceed INT_MAX at + // large output buffers; cast keeps CUTE's layout math in 64-bit. + auto mY = tma_store_y.get_tma_tensor( + make_shape(static_cast(num_y_m_tiles) * TileM, + static_cast(hidden_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local + // tid). Each epi tile is the full TileM rows × EpiChunkN cols; the + // 128-thread tcgen05.ld maps one thread per TMEM datapath row. + // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) + // modes, which is what make_tmem_copy's cotiled builder requires. ── + auto epi_tile = make_tile(Int{}, Int{}); + // Extract the flat (M,N) view from the UMMA C-fragment (MMA,MMA_M,MMA_N) + // before tiling — matches CUTLASS's accumulators(make_coord(_,_),_0,_0). + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto tAcc_epi = flat_divide(acc_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto t2r = make_tmem_copy(TmemLoadOp{}, tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tmem_copy_tid); + auto tTR_tAcc = thr_t2r.partition_S(tAcc_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) + // Within-chunk (m,n) coords for the reg → smem-slot scatter. + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + // Register fragment sized from the DEST (partition_D) shape — per-thread + // (T2R,T2R_M,T2R_N). Sizing from partition_S(tmem) would wrongly include the + // warp-collective datapath-lane dim. Matches CUTLASS's + // make_tensor(shape(tTR_sD)). + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); // f32 regs + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + CUTE_UNROLL + for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { + + // ── Mainloop (warp 3 only): bracket the k-loop with the accumulator + // pipeline. Each k-stage waits/releases the TMA→UMMA mainloop + // pipeline; consumer_release issues the UMMA-gated smem-buffer + // arrival (no hand-rolled mbarrier). (Conservative: no MMA/k overlap.) + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); // TMEM acc free (prev epilogue done) + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + for (int k = 0; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + auto tCsZ = cta_mma.partition_A(sZ(_, _, state.index())); + auto tCsW = cta_mma.partition_B(sW(_, _, state.index())); + auto tCrZ = cta_mma.make_fragment_A(tCsZ); + auto tCrW = cta_mma.make_fragment_B(tCsW); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrZ); ++kb) { + // First k-block of this n-tile clears TMEM; the rest accumulate. + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrZ(_, _, kb), tCrW(_, _, kb), tCtAcc); + } + pipe.consumer_release(state); // UMMA-gated smem-buffer release + ++state; + } + acc_pipe.producer_commit(acc_prod_state); // signal acc ready (umma_arrive) + ++acc_prod_state; + } + + // ── Epilogue: wait for the accumulator, then this WG processes its + // n-chunks [wg·NChunksHalf, +NChunksHalf). ── + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc_stage = thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; // absolute n-chunk index + + // TMEM → registers (this chunk, full TileM rows). + copy(t2r, tTR_tAcc_stage(_, _, _, _0{}, chunk), tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + + // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (ms >= valid_store_subtiles) + continue; + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rAcc); ++i) { + int m_row = get<0>(tTR_cChunk(i)); // 0..TileM + int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN + if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) + sStore(m_row - ms * AtomTileM, n_col) = + static_cast(tTR_rAcc(i)); + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx = MSub * (y_m + m_sub) + ms; + int n_tile_idx = n * (TileN / EpiChunkN) + chunk; + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // All epilogue TMEM reads are done; release the accumulator so the next + // n-tile's MMA may reuse it (one elected consumer thread arrives). + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + // Keep the final MMA+epilogue barrier so the outer launcher can safely free + // TMEM after this consumer returns. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + // The Compute=100 specialization is never dispatched on non-SM100 targets + // (host picks Compute=90 there). Trap if it is ever reached. + __trap(); +#endif +} +}; // Mlp2FusedConsumerImpl<100> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — keeps the existing call style. `Compute` defaults to 90 so +// current call sites (mlp2_fused_consumer(...)) are unchanged; +// pass mlp2_fused_consumer(...) to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp2_fused_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2FusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, + int hidden_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx = -1, + int num_splits = -1, + int valid_m64_subtiles = -1) { + Mlp2FusedConsumerImpl::template run( + pipe, state, smem, tma_store_y, y_m, hidden_dim, + num_y_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits, + valid_m64_subtiles); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm90.cuh new file mode 100644 index 00000000..155e6ae3 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm90.cuh @@ -0,0 +1,830 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 2 (single-tile, fused): Y = Z @ A^T (down-projection) +// ═══════════════════════════════════════════════════════════════════ +// +// Single (TileM, TileN) producer/consumer building blocks for callers +// that drive the outer M-tile loop themselves: the standalone MLP2 +// benchmark (mlp2.cu) and the fused MLP kernel (mlp.cuh). +// +// Design: +// - Cooperative 2-WG consumer; split axis selected by TileM: +// TileM=128 → Layout<_2,_1,_1>, split along M (each WG = 64×TileN) +// TileM=64 → Layout<_1,_2,_1>, split along N (each WG = 64×TileN/2) +// - One accumulator per WG +// - Per-WG store_buf with EpiChunkN-round epilogue +// - K_PIPE_MMAS = 1 (CUTLASS-canonical, see +// sm90_mma_tma_gmma_ss_warpspecialized.hpp:264) +// - Single fused Z+W TMA pipe (1 mbarrier per stage, 2 TMA copies) +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for Z / A) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Per `expert_ids`: one int per TileM-row m-block. +// +// Mlp2FusedSmem omits the pipeline storage so callers can fuse it with +// other phases' smem (see mlp.cuh). The benchmark wrapper (mlp2.cu) is +// responsible for attaching MainloopPipeline::SharedStorage and running +// the outer for(m) loop over m-tiles. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm90.cuh" +#include "math_sm90.cuh" +#include "tmem_load_op_sm90.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer +// specialization. These headers are header-only and safe to include under +// an sm_90a build — the tcgen05 PTX they emit is guarded by +// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on +// __CUDA_ARCH__ >= 1000) ever instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 128, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 32 +> +struct Mlp2Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + + static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); + + // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate + // along M (each owns 64 rows × TileN). With TileM=64 the WGs + // cooperate along N (each owns 64 rows × TileN/2). AtomTileM is the + // per-WG M extent — always 64. + static constexpr bool kMSplit = (TileM == 128); + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); + + static_assert(TileN <= 256, "TileN must fit TMA box limit"); + static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); + static_assert(kMSplit || (TileN % 2 == 0), + "TileN must be even when TileM=64 (cooperative N-split halves)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + // Cooperative 2-WG layout: M-split for TileM=128, N-split for TileM=64. + // AtomN is the per-WG WGMMA atom N — combined across both WGs it must + // cover exactly TileN cols (= sW's N extent). With WgLayout=<1,2,1> + // and AtomN=TileN, the TiledMMA would cover 2·TileN cols and WG1's + // partition_B(sW) slice would read past sW's TileN-wide buffer → IMA. + static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); + using GmmaAtom = typename GmmaSelector::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom; + + using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int>{})); + using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int>{})); + + using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int, Int>{})); + using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{}, + Shape, Int, Int>{})); + + // Per-WG store buffer slot: (AtomTileM, EpiChunkN). + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore_1 = SmemLayoutStoreSlot; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop + // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the + // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the single + // TMEM accumulator from the MMA warp to the epilogue warps. See + // Mlp2FusedConsumerImpl<100>. The Blackwell launcher must build the mainloop + // pipe with num_consumers=1 (one umma_arrive per stage). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Optional accumulator double-buffering: each stage owns one TileN-wide + // TMEM accumulator, allowing MMA(n+1) to fill one stage while epilogue + // drains/TMA-stores the previous stage. + static constexpr int AccStages = 2; + static_assert(AccStages * TileN <= 512, + "SM100 MLP2 accumulator stages must fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + static constexpr int TmaTransBytesZ = + static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesW = + static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); + // Fused pipeline: Z + W share one mbarrier per stage. + static constexpr int TmaTransBytes = TmaTransBytesZ + TmaTransBytesW; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; +}; + +template +using Mlp2MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP2 shared memory — no pipeline storage (managed externally) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp2FusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element smem_W[smem_W_size]; + // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these; + // SM100 launchers allocate/free TMEM outside the consumer and publish the + // base here. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } + CUTE_DEVICE Element* W_data() { return &smem_W[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Pipeline construction helper — single fused Z+W pipeline +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ auto mlp2_make_pipe( + typename Traits::MainloopPipeline::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp2_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer +// (mlp2_fused_producer) and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ auto mlp2_make_pipe_umma( + typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused producer — single (m), all (n) under blockIdx.y stride. +// Single fused Z + W TMA pipe per k-step. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_fused_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2FusedSmem& smem, + TmaLoadZ const& tma_load_z, + TmaLoadW const& tma_load_a, + int m, + int expert_or_n_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_n_rows, + int num_n_tiles, + int num_k_tiles, + // N-split identity (flat-grid launch); -1 → blockIdx.y / gridDim.y. + int split_idx = -1, + int num_splits = -1) { + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + // int64_t cast: total_n_rows × intermediate_dim and num_tokens × + // intermediate_dim both exceed INT_MAX at production shapes. CUTE's + // layout math otherwise wraps to a negative byte offset → IMA. + auto mZ = tma_load_z.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(intermediate_dim))); + auto mA = [&]() { + if constexpr (Expert3D) { + return tma_load_a.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(num_experts))); + } else { + return tma_load_a.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(intermediate_dim))); + } + }(); + + auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); + auto cta_tma_a = tma_load_a.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_z.partition_D(sZ); + auto tWsW = cta_tma_a.partition_D(sW); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + auto gW = [&]() { + if constexpr (Expert3D) { + return local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_n_offset)); + } else { + return local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(expert_or_n_offset + n, _)); + } + }(); + auto tWgW = cta_tma_a.partition_S(gW); + + CUTE_UNROLL + for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m + m_sub, _)); + auto tZgZ = cta_tma_z.partition_S(gZ); + + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_z.with(*bar, 0), + tZgZ(_, _, _, k), tZsZ(_, _, _, state.index())); + copy(tma_load_a.with(*bar, 0), + tWgW(_, _, _, k), tWsW(_, _, _, state.index())); + } + ++state; + } + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused consumer — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported +// compute capability provides a `run(...)` member (a function-template on +// Traits/Pipeline/TmaStoreY, since function templates can't be partially +// specialized). The free function `mlp2_fused_consumer` below forwards to +// the right specialization — call sites stay in the existing style. +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, single register acc) +// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulator) +// +// y_m is the M-tile index into the output Y buffer. + +template +struct Mlp2FusedConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. The two WGs split the (TileM, TileN) output along +// M: WG1 writes rows [0, AtomTileM), WG2 writes [AtomTileM, TileM). +// (Verbatim from the original mlp2_fused_consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp2FusedConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2FusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, // M-tile index into Y buffer + int hidden_dim, + int num_y_m_tiles, // number of M-tile slots in Y buffer + int num_n_tiles, + int num_k_tiles, + // N-split identity (flat-grid launch); -1 → blockIdx.y / gridDim.y. + int split_idx, + int num_splits, + int valid_m64_subtiles) { + + static_assert(MSubTiles == 1, + "Hopper MLP2 consumer processes one GEMM M-tile per invocation"); + using Element = typename Traits::Element; + constexpr int kM64Subtiles = Traits::TileM / 64; + int valid_store_subtiles = valid_m64_subtiles < 0 + ? kM64Subtiles : min(kM64Subtiles, valid_m64_subtiles); + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; + const int my_barrier_id = 1 + my_wg; + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // int64_t cast: num_y_m_tiles · TileM · hidden_dim can exceed + // INT_MAX at large output buffers; cast keeps CUTE's layout math + // in 64-bit so byte offsets don't wrap. + auto mY = tma_store_y.get_tma_tensor( + make_shape(static_cast(num_y_m_tiles) * Traits::TileM, + static_cast(hidden_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; // CUTLASS-matching + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. + // TileM=64 (N-split): WG_w owns full M, cols [w*WgTileN, ...). + bool store_output = Traits::kMSplit + ? (my_wg < valid_store_subtiles) + : (valid_store_subtiles > 0); + if (store_output) { + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: 2 row tiles per logical M tile. + m_tile_idx = 2 * y_m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M tile; WG_w handles + // its N-half. + m_tile_idx = y_m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp2FusedConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. +// +// No cooperative-warpgroup machinery (UMMA has no warpgroup concept): +// - One 1SM UMMA atom with M=TileM covers the whole tile; the single +// accumulator Y = Z·Aᵀ lives in TMEM (no register/M-split). +// - Warp 3 issues UMMA and stays epilogue-free; warps 4-11 form the two +// aligned epilogue WGs. +// - Epilogue: the two consumer warpgroups split TileN; each reads its +// N-half from TMEM→regs in EpiChunkN chunks, casts to bf16, and +// stores 64-row (AtomTileM) tiles via the existing reg→SMEM→TMA path +// (so the host TMA atom is unchanged). No activation fusion. +// +// TileM is the template parameter as on Hopper (64 or 128 — both native +// 1SM UMMA M sizes). For TileM=128 the epilogue stores MSub=2 row-tiles. +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp2FusedConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2FusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, + int hidden_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits, + int valid_m64_subtiles) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two consumer warpgroups"); + constexpr int WgN = TileN / 2; // per-warpgroup N width + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per warpgroup + constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) + int valid_store_subtiles = valid_m64_subtiles < 0 + ? MSub : min(MSub, valid_m64_subtiles); + + // ── Thread identity ───────────────────────────────────── + // Warp 3 is the dedicated, epilogue-free UMMA producer. Warps 4..11 are + // the two aligned epilogue WGs. This mirrors MLP1's Blackwell path. + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + const int tmem_copy_tid = tid_wg; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; // 1 or 2 + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, K-major SS. + // Operand A = Z, operand B = A (the down weight) — both K-major, the + // SM100-native orientation, so no transpose subtlety. ── + auto tiled_mma = make_tiled_mma( + SM100_MMA_F16BF16_SS{}); + auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + // ── One TMEM accumulator: Y = Z·Aᵀ ── + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the + // alternate stage while epilogue drains/TMA-stores MMA(n). ── + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; // one umma_arrive per commit + acc_params.consumer_arv_count = 1; // one elected releaser + acc_params.initializing_warp = 4; // warp 4 inits the acc barriers + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + // ── Per-WG store slot (AtomTileM × EpiChunkN) — same smem as Hopper ── + constexpr int store_slot_elems = AtomTileM * EpiChunkN; + Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // int64_t cast: num_y_m_tiles · TileM · hidden_dim can exceed INT_MAX at + // large output buffers; cast keeps CUTE's layout math in 64-bit. + auto mY = tma_store_y.get_tma_tensor( + make_shape(static_cast(num_y_m_tiles) * TileM, + static_cast(hidden_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local + // tid). Each epi tile is the full TileM rows × EpiChunkN cols; the + // 128-thread tcgen05.ld maps one thread per TMEM datapath row. + // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) + // modes, which is what make_tmem_copy's cotiled builder requires. ── + auto epi_tile = make_tile(Int{}, Int{}); + // Extract the flat (M,N) view from the UMMA C-fragment (MMA,MMA_M,MMA_N) + // before tiling — matches CUTLASS's accumulators(make_coord(_,_),_0,_0). + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto tAcc_epi = flat_divide(acc_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto t2r = make_tmem_copy(TmemLoadOp{}, tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tmem_copy_tid); + auto tTR_tAcc = thr_t2r.partition_S(tAcc_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) + // Within-chunk (m,n) coords for the reg → smem-slot scatter. + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + // Register fragment sized from the DEST (partition_D) shape — per-thread + // (T2R,T2R_M,T2R_N). Sizing from partition_S(tmem) would wrongly include the + // warp-collective datapath-lane dim. Matches CUTLASS's + // make_tensor(shape(tTR_sD)). + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); // f32 regs + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + CUTE_UNROLL + for (int m_sub = 0; m_sub < MSubTiles; ++m_sub) { + + // ── Mainloop (warp 3 only): bracket the k-loop with the accumulator + // pipeline. Each k-stage waits/releases the TMA→UMMA mainloop + // pipeline; consumer_release issues the UMMA-gated smem-buffer + // arrival (no hand-rolled mbarrier). (Conservative: no MMA/k overlap.) + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); // TMEM acc free (prev epilogue done) + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + for (int k = 0; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + auto tCsZ = cta_mma.partition_A(sZ(_, _, state.index())); + auto tCsW = cta_mma.partition_B(sW(_, _, state.index())); + auto tCrZ = cta_mma.make_fragment_A(tCsZ); + auto tCrW = cta_mma.make_fragment_B(tCsW); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrZ); ++kb) { + // First k-block of this n-tile clears TMEM; the rest accumulate. + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrZ(_, _, kb), tCrW(_, _, kb), tCtAcc); + } + pipe.consumer_release(state); // UMMA-gated smem-buffer release + ++state; + } + acc_pipe.producer_commit(acc_prod_state); // signal acc ready (umma_arrive) + ++acc_prod_state; + } + + // ── Epilogue: wait for the accumulator, then this WG processes its + // n-chunks [wg·NChunksHalf, +NChunksHalf). ── + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc_stage = thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; // absolute n-chunk index + + // TMEM → registers (this chunk, full TileM rows). + copy(t2r, tTR_tAcc_stage(_, _, _, _0{}, chunk), tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + + // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (ms >= valid_store_subtiles) + continue; + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rAcc); ++i) { + int m_row = get<0>(tTR_cChunk(i)); // 0..TileM + int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN + if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) + sStore(m_row - ms * AtomTileM, n_col) = + static_cast(tTR_rAcc(i)); + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx = MSub * (y_m + m_sub) + ms; + int n_tile_idx = n * (TileN / EpiChunkN) + chunk; + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // All epilogue TMEM reads are done; release the accumulator so the next + // n-tile's MMA may reuse it (one elected consumer thread arrives). + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + // Keep the final MMA+epilogue barrier so the outer launcher can safely free + // TMEM after this consumer returns. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + // The Compute=100 specialization is never dispatched on non-SM100 targets + // (host picks Compute=90 there). Trap if it is ever reached. + __trap(); +#endif +} +}; // Mlp2FusedConsumerImpl<100> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — keeps the existing call style. `Compute` defaults to 90 so +// current call sites (mlp2_fused_consumer(...)) are unchanged; +// pass mlp2_fused_consumer(...) to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp2_fused_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2FusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, + int hidden_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx = -1, + int num_splits = -1, + int valid_m64_subtiles = -1) { + Mlp2FusedConsumerImpl::template run( + pipe, state, smem, tma_store_y, y_m, hidden_dim, + num_y_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits, + valid_m64_subtiles); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp2_t.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp2_t.cuh index ecf01a8a..be66ce7a 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp2_t.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp2_t.cuh @@ -1,540 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Phase 2 (transposed): Y = Z @ A -// ═══════════════════════════════════════════════════════════════════ -// -// Cooperative 2-WG consumer; split axis selected by TileM: -// TileM=128 → Layout<_2,_1,_1>, split along M (each WG owns 64×TileN) -// TileM=64 → Layout<_1,_2,_1>, split along N (each WG owns 64×TileN/2) -// One accumulator per WG covering (AtomTileM=64, AtomN). Each WG writes -// its strip independently — no cross-WG NamedBarrier sync in the epilogue. -// -// Difference from mlp2_fused.cuh: A is read through a COLUMN-MAJOR TMA view -// (the physical layout is the same per-expert [hidden_dim, -// intermediate_dim] row-major buffer, but the column-major view makes -// intermediate_dim the row axis so it lines up as the N axis of the -// GEMM). Output Y has shape [num_tokens, intermediate_dim]. -// - num_n_tiles = intermediate_dim / TileN (the N axis) -// - num_k_tiles = hidden_dim / TileK (the K axis) -// - expert offset is along K: expert · num_k_tiles -// -// A single fused TMA pipeline carries Z + W per k-step; transaction = -// bytes(Z) + bytes(W). Producer issues two TMA copies against the same -// mbarrier per acquire. -// -// CUTLASS-canonical patterns: -// - K_PIPE_MMAS = 1 -// - __launch_bounds__(NumThreads, 1) — applied by the caller kernel -// -// CTA layout (384 threads = 12 warps): -// WG0 warp 0 : Producer (TMA loads for Z / A) -// WG0 warps 1-3 : Idle on SM90; warp 3 issues UMMA on SM100 -// WG1 (warps 4-7) : Consumer atom 0 — owns m=[0, AtomTileM) -// WG2 (warps 8-11) : Consumer atom 1 — owns m=[AtomTileM, TileM) -// -// Per `expert_ids`: one int per TileM=128 m-block. -// -// ═══════════════════════════════════════════════════════════════════ - -#include "models.cuh" -#include "math.cuh" -#include "tmem_load_op.cuh" - -#include -#include -#include -#include -#include - -// SM100 (Blackwell / UMMA) support for the Compute=100 consumer -// specialization in mlp2_t_fused.cuh. These headers are header-only and safe -// to include under an sm_90a build — the tcgen05 PTX they emit is guarded by -// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on -// __CUDA_ARCH__ >= 1000) ever instantiates it. -#include -#include - -#include -#include -#include // PipelineTmaUmmaAsync, PipelineUmmaAsync - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// Traits -// ═══════════════════════════════════════════════════════════════════ - -template < - typename Element_, - int TileM_ = 128, - int TileN_ = 128, - int TileK_ = 64, - int Stages_ = 4, - int EpiChunkN_ = 32 -> -struct Mlp2TTraits { - using Element = Element_; - using ElementAccum = float; - - static constexpr int TileM = TileM_; - static constexpr int TileN = TileN_; - static constexpr int TileK = TileK_; - static constexpr int Stages = Stages_; - static constexpr int EpiChunkN = EpiChunkN_; - - static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); - - // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate - // along M (each owns 64 rows × TileN). With TileM=64 the WGs cooperate - // along N (each owns 64 rows × TileN/2). AtomTileM is the per-WG M - // extent — always 64. Mirrors mlp2_fused.cuh's kMSplit/N-split pattern. - static constexpr bool kMSplit = (TileM == 128); - static constexpr int AtomTileM = 64; - static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); - - static_assert(TileN <= 256, "TileN must fit TMA box limit"); - static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); - static_assert(kMSplit || (TileN % 2 == 0), - "TileN must be even when TileM=64 (cooperative N-split halves)"); - static constexpr int NumEpiRounds = WgTileN / EpiChunkN; - - // A is loaded through a column-major TMA view (stride=(1, I)), so the - // W operand in shared memory must be MN-major to match. The wgmma atom - // for operand B must accept an MN-major source (GmmaSelectorKMN — note - // the trailing N in "KMN" denotes MN-major B). Z stays K-major (Z is - // row-major in gmem with K=H contiguous-ish given the layout CUTE - // produces for shape (M, K), stride (K, 1)), so its atom is the regular - // K-major SW128 atom. - // - // AtomN is the per-WG WGMMA atom N — combined across both WGs it must - // cover exactly TileN cols (= sW's N extent). With WgLayout=<1,2,1> - // and AtomN=TileN, the TiledMMA would cover 2·TileN cols and WG1's - // partition_B(sW) slice would read past sW's TileN-wide buffer → IMA. - static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); - using GmmaAtom = typename GmmaSelectorKMN::Atom; - using WgLayout = cute::conditional_t< - kMSplit, - Layout>, - Layout>>; - using TiledMma = TiledMMA, WgLayout>; - - using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom; - using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom; - - using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{}, - Shape, Int>{})); - // Step<_2, _1>: tile_to_shape interprets the inner SmemLayoutAtomW as - // MN-major (M is contiguous in smem), so passing the tile shape with - // the swap step makes the resulting layout's M dim contiguous — which - // is what the column-major A gmem feed expects. - using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtomW{}, - Shape, Int>{}, - Step<_2, _1>{})); - - using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{}, - Shape, Int, Int>{})); - using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{}, - Shape, Int, Int>{}, - Step<_2, _1, _3>{})); - - // Per-WG store buffer slot: (AtomTileM, EpiChunkN). - using SmemLayoutStoreSlot = Layout< - Shape, Int>, - Stride, _1>>; - using SmemLayoutStore_1 = SmemLayoutStoreSlot; - - using MainloopPipeline = cutlass::PipelineTmaAsync; - using PipelineState = cutlass::PipelineState; - - // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop - // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the - // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the single - // TMEM accumulator from the MMA warp to the epilogue warps. See - // Mlp2TFusedConsumerImpl<100> in mlp2_t_fused.cuh. The Blackwell launcher - // must build the mainloop pipe with num_consumers=1 (one umma_arrive per - // stage). - using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< - Stages, cute::Shape, - cute::Shape>; - // Double-buffered SM100 TMEM accumulators: each stage owns one TileN-wide - // accumulator, allowing MMA(n+1) to fill the alternate stage while the - // epilogue drains/TMA-stores MMA(n). - static constexpr int AccStages = 2; - static_assert(AccStages * TileN <= 512, - "SM100 MLP2-T accumulator stages must fit in 512 TMEM columns"); - using AccumulatorPipeline = cutlass::PipelineUmmaAsync; - - // SM100 UMMA TiledMMA — one 1SM tcgen05 atom spanning the whole - // (TileM, TileN) tile with the accumulator in TMEM. THE mlp2_t-specific - // choice: operand A = Z is K-major (UMMA::Major::K), operand B = A (the - // weight) is MN-major (UMMA::Major::MN) — mlp2_fused uses . The B - // operand descriptor is produced by make_umma_desc reading - // SmemLayoutW's (MN-contiguous, SW128) stride; SmemLayoutW stays built from - // Layout_MN_SW128_Atom + Step<_2,_1>, which is exactly the canonical - // Major-MN layout make_umma_desc accepts (mma_traits_sm100.hpp). - using TiledMmaUmma = decltype(make_tiled_mma( - SM100_MMA_F16BF16_SS{})); - - static constexpr int TmaTransBytesZ = - static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); - static constexpr int TmaTransBytesW = - static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); - // Fused pipeline: Z + W share one mbarrier per stage. - static constexpr int TmaTransBytes = TmaTransBytesZ + TmaTransBytesW; - - static constexpr int WarpSize = 32; - static constexpr int WarpGroupSize = 128; - static constexpr int NumConsumers = 2; - static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; - static constexpr int NumThreads = 384; -}; - -template -using Mlp2TMainloopPipelineFor = cute::conditional_t< - Compute == 100, - typename Traits::MainloopPipelineUmma, - typename Traits::MainloopPipeline>; - -// ═══════════════════════════════════════════════════════════════════ -// Shared Memory -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp2TSmem { - using Element = typename Traits::Element; - - static constexpr int smem_Z_size = cosize_v; - static constexpr int smem_W_size = cosize_v; - static constexpr int smem_store_size = cosize_v; - - alignas(128) Element smem_Z[smem_Z_size]; - alignas(128) Element smem_W[smem_W_size]; - // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. - alignas(128) Element store_buf[2 * smem_store_size]; - - typename Traits::MainloopPipeline::SharedStorage pipe_storage; - - CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } - CUTE_DEVICE Element* W_data() { return &smem_W[0]; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// TMA Producer (Warp 0) — one acquire per k-step, issues Z + W -// against the same mbarrier. -// - Z view: [num_tokens, hidden_dim] row-major (Z's K axis = H). -// - A view: [intermediate_dim, total_k_cols] column-major, where -// total_k_cols = num_experts · hidden_dim. Per-expert tile at -// k = expert · num_k_tiles + k_local. -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void mlp2_t_producer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2TSmem& smem, - TmaLoadZ const& tma_load_z, - TmaLoadW const& tma_load_a, - const int* expert_ids, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int total_k_cols, // num_experts · hidden_dim - int num_m_tiles, - int num_n_tiles, // intermediate_dim / TileN - int num_k_tiles) { // hidden_dim / TileK - - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); - - auto mZ = tma_load_z.get_tma_tensor(make_shape( - static_cast(num_tokens), - static_cast(hidden_dim))); - // Column-major view of A: shape (intermediate_dim, total_k_cols) - // — same underlying buffer as the row-major [E·H, I] tensor in - // mlp2_fused.cuh, but the box geometry makes intermediate_dim the row - // axis so it lines up with the GEMM N dimension. - auto mA = tma_load_a.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(total_k_cols))); - - auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); - auto cta_tma_a = tma_load_a.get_slice(Int<0>{}); - - auto tZsZ = cta_tma_z.partition_D(sZ); - auto tWsW = cta_tma_a.partition_D(sW); - - int n_start = blockIdx.y; - int n_stride = gridDim.y; - - for (int m = blockIdx.x; m < num_m_tiles; m += gridDim.x) { - int expert = expert_ids[m]; - // In mlp2_t the expert is contiguous in the K direction - // (= hidden_dim), so the offset is along k_tile, not n_tile. - int expert_k_offset = expert * num_k_tiles; - - auto gZ = local_tile(mZ, - make_tile(Int{}, Int{}), - make_coord(m, _)); - auto tZgZ = cta_tma_z.partition_S(gZ); - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - // A tile: (TileN rows of I, TileK rows of K=H), positioned - // at (n along I, expert_k_offset+k along total_k_cols). - auto gW = local_tile(mA, - make_tile(Int{}, Int{}), - make_coord(n, _)); - auto tWgW = cta_tma_a.partition_S(gW); - - for (int k = 0; k < num_k_tiles; ++k) { - pipe.producer_acquire(state); - if (threadIdx.x == 0) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_z.with(*bar, 0), - tZgZ(_, _, _, k), tZsZ(_, _, _, state.index())); - copy(tma_load_a.with(*bar, 0), - tWgW(_, _, _, expert_k_offset + k), - tWsW(_, _, _, state.index())); - } - ++state; - } - } - } -} - -// ═══════════════════════════════════════════════════════════════════ -// Cooperative 2-WG Consumer — M-split, single acc per WG -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void mlp2_t_consumer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2TSmem& smem, - TmaStoreY const& tma_store_y, - int intermediate_dim, - int num_m_tiles, - int num_n_tiles, - int num_k_tiles) { - - (void)Compute; - using Element = typename Traits::Element; - typename Traits::TiledMma tiled_mma; - int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 - int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 - auto thr_mma = tiled_mma.get_slice(tid_in_mma); - - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); - - auto tCsZ = thr_mma.partition_A(sZ); - auto tCsW = thr_mma.partition_B(sW); - - auto acc = partition_fragment_C(tiled_mma, - Shape, Int>{}); - - auto cC = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCcC = thr_mma.partition_C(cC); - - const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 - const int my_barrier_id = 1 + my_wg; // 1 or 2 - const bool is_my_wg_leader = (tid_in_wg == 0); - - constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; - Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - // Y view: [num_tokens, intermediate_dim] row-major. - auto mY = tma_store_y.get_tma_tensor(make_shape( - static_cast(num_m_tiles) * Traits::TileM, - static_cast(intermediate_dim))); - auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); - - int n_start = blockIdx.y; - int n_stride = gridDim.y; - bool store_in_flight = false; - constexpr int K_PIPE_MMAS = 1; - - for (int m = blockIdx.x; m < num_m_tiles; m += gridDim.x) { - for (int n = n_start; n < num_n_tiles; n += n_stride) { - - clear(acc); - - auto state_release = state; - int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; - - // Prologue - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - ++state; - } - // Steady state - for (int k = prologue_count; k < num_k_tiles; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - - warpgroup_wait(); - warpgroup_fence_operand(acc); - pipe.consumer_release(state_release); - ++state; - ++state_release; - } - // Drain - warpgroup_wait<0>(); - warpgroup_fence_operand(acc); - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_release(state_release); - ++state_release; - } - - // ── Epilogue ────────────────────────────────────── - // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. - // TileM=64 (N-split): WG_w owns full M (=64), cols [w*WgTileN, ...). - CUTE_UNROLL - for (int r = 0; r < Traits::NumEpiRounds; ++r) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - int chunk_start = r * Traits::EpiChunkN; - CUTE_UNROLL - for (int i = 0; i < size(acc); ++i) { - auto coord = tCcC(i); - int m_loc = get<0>(coord); - int n_loc = get<1>(coord); - int m_local, n_local; - if constexpr (Traits::kMSplit) { - m_local = m_loc - my_wg * Traits::AtomTileM; - n_local = n_loc; - } else { - m_local = m_loc; - n_local = n_loc - my_wg * Traits::WgTileN; - } - if (n_local >= chunk_start && - n_local < chunk_start + Traits::EpiChunkN) { - int chunk_n = n_local - chunk_start; - sStore(m_local, chunk_n) = static_cast(acc(i)); - } - } - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - if (is_my_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx, n_tile_idx; - if constexpr (Traits::kMSplit) { - // TileM=128: 2 row tiles per logical M tile. - m_tile_idx = 2 * m + my_wg; - n_tile_idx = n * Traits::NumEpiRounds + r; - } else { - // TileM=64: 1 row tile per logical M tile; WG_w handles - // its N-half. - m_tile_idx = m; - n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) - + my_wg * Traits::NumEpiRounds + r; - } - auto gY = local_tile(mY, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_y, cta_tma_y.partition_S(sStore), - cta_tma_y.partition_D(gY)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - } - if (store_in_flight) - cute::tma_store_wait<0>(); -} - -// ═══════════════════════════════════════════════════════════════════ -// mlp2_t_fwd — device orchestrator -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void mlp2_t_fwd( - Mlp2TSmem& smem, - TmaLoadZ const& tma_load_z, - TmaLoadW const& tma_load_a, - TmaStoreY const& tma_store_y, - const int* expert_ids, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int total_k_cols, - int num_m_tiles) { - - using Pipeline = typename Traits::MainloopPipeline; - using PipeState = typename Traits::PipelineState; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - int num_n_tiles = intermediate_dim / Traits::TileN; - int num_k_tiles = hidden_dim / Traits::TileK; - - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 4 && warp_id <= 11); - - cute::prefetch_tma_descriptor(tma_load_z.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_a.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_y.get_tma_descriptor()); - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = Traits::ConsumerThreads; - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - Pipeline pipe(smem.pipe_storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); - - __syncthreads(); - - PipeState state = is_producer ? - cutlass::make_producer_start_state() : PipeState{}; - PipeState cons_state; - - if (is_producer) { - mlp2_t_producer( - pipe, state, - smem, tma_load_z, tma_load_a, - expert_ids, - num_tokens, hidden_dim, intermediate_dim, total_k_cols, - num_m_tiles, num_n_tiles, num_k_tiles); - } - - if (is_consumer) { - mlp2_t_consumer( - pipe, cons_state, - smem, tma_store_y, - intermediate_dim, - num_m_tiles, num_n_tiles, num_k_tiles); - } - - __syncthreads(); -} - -} // namespace liger +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif + +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp2_t_sm90.cuh" +#else +#include "mlp2_t_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused.cuh index c3a5f264..797c713e 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused.cuh @@ -1,679 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Phase 2 Transpose (fused tile): Y = Z @ A -// ═══════════════════════════════════════════════════════════════════ -// -// Single M-tile version for embedding in the fused MLP kernel -// (mlp_bwd.cuh Phase 1b'). Matches the production design in mlp2_t.cuh: -// - Cooperative 2-WG consumer; split axis selected by TileM (kMSplit -// in Mlp2TTraits): TileM=128 → M-split, TileM=64 → N-split. -// - One accumulator per WG covering (AtomTileM=64, AtomN) -// - Per-WG store_buf with EpiChunkN-round epilogue -// - K_PIPE_MMAS = 1 (CUTLASS-canonical) -// - Single fused Z+W TMA pipe (1 mbarrier per stage, 2 TMA copies) -// -// Differs from mlp2_t.cuh only in: -// - No outer for(m) loop — caller passes a single m (y_m for the -// output buffer index, expert_k_offset for the input weights) -// - Mlp2TFusedSmem omits pipeline storage (caller manages it) -// -// Mirrors mlp2_fused.cuh's single-tile style. - -#include "mlp2_t.cuh" - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// Fused MLP2-T shared memory — no pipeline storage (managed externally) -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp2TFusedSmem { - using Element = typename Traits::Element; - - static constexpr int smem_Z_size = cosize_v; - static constexpr int smem_W_size = cosize_v; - static constexpr int smem_store_size = cosize_v; - - alignas(128) Element smem_Z[smem_Z_size]; - alignas(128) Element smem_W[smem_W_size]; - // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. - alignas(128) Element store_buf[2 * smem_store_size]; - - // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM - // base address, plus the accumulator pipeline's barrier storage (UMMA→ - // epilogue handoff). The Hopper (Compute=90) path never touches these, and - // keeping them here leaves the consumer signature identical across Compute - // values (TMEM base is externally owned; the acc pipeline lives here). - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; - - CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } - CUTE_DEVICE Element* W_data() { return &smem_W[0]; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// Pipeline construction helper — single fused Z+W pipeline -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ auto mlp2_t_make_pipe( - typename Traits::MainloopPipeline::SharedStorage& storage) { - - using Pipeline = typename Traits::MainloopPipeline; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = Traits::ConsumerThreads; - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// SM100 (Compute=100) mainloop pipe. Mirrors mlp2_t_make_pipe but builds the -// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues -// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer -// is released by a single umma_arrive per stage (the MMA warp), not per -// consumer thread. Params/PipelineState/SharedStorage alias the Hopper -// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer -// (mlp2_t_fused_producer) and the Compute=100 consumer drive it unchanged. -template -__device__ __forceinline__ auto mlp2_t_make_pipe_umma( - typename Traits::MainloopPipelineUmma::SharedStorage& storage) { - - using Pipeline = typename Traits::MainloopPipelineUmma; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 4 && warp_id <= 11); - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// ═══════════════════════════════════════════════════════════════════ -// Fused producer — single (m), all (n) under blockIdx.y stride. -// Single fused Z + W TMA pipe per k-step. -// -// Per mlp2_t: the expert is contiguous along K (= hidden_dim), so the -// caller supplies expert_k_offset = expert · num_k_tiles. A is read -// through a column-major view (shape (intermediate_dim, total_k_cols)) -// with N along intermediate_dim and K along E·hidden_dim. -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void mlp2_t_fused_producer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2TFusedSmem& smem, - TmaLoadZ const& tma_load_z, - TmaLoadW const& tma_load_a, - int m, - int expert_or_k_offset, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int num_experts, - int total_k_cols, // num_experts · hidden_dim - int num_n_tiles, - int num_k_tiles, - // Flat-grid N-split: split_idx = flat_id % NSplit, num_splits = NSplit. - // Defaults (-1) fall back to blockIdx.y / gridDim.y for legacy callers. - int split_idx = -1, - int num_splits = -1) { - - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); - - // int64_t casts: num_tokens × hidden_dim and intermediate_dim × - // total_k_cols both exceed INT_MAX at production shapes. - auto mZ = tma_load_z.get_tma_tensor(make_shape( - static_cast(num_tokens), - static_cast(hidden_dim))); - auto mA = [&]() { - if constexpr (Expert3D) { - return tma_load_a.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(num_experts))); - } else { - return tma_load_a.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(total_k_cols))); - } - }(); - - auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); - auto cta_tma_a = tma_load_a.get_slice(Int<0>{}); - - auto tZsZ = cta_tma_z.partition_D(sZ); - auto tWsW = cta_tma_a.partition_D(sW); - - auto gZ = local_tile(mZ, - make_tile(Int{}, Int{}), - make_coord(m, _)); - auto tZgZ = cta_tma_z.partition_S(gZ); - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - auto gA = [&]() { - if constexpr (Expert3D) { - return local_tile(mA, - make_tile(Int{}, Int{}), - make_coord(n, _, expert_or_k_offset)); - } else { - return local_tile(mA, - make_tile(Int{}, Int{}), - make_coord(n, _)); - } - }(); - auto tAgA = cta_tma_a.partition_S(gA); - - for (int k = 0; k < num_k_tiles; ++k) { - if constexpr (ThrottleTma) { - if (state.count() >= 1) { - using FullBarrier = typename Pipeline::FullBarrier; - auto completed = - cutlass::make_producer_start_state(); - completed.advance(state.count() - 1); - FullBarrier::wait( - pipe.producer_get_barrier(completed), - completed.phase() ^ 1); - } - } - pipe.producer_acquire(state); - if (threadIdx.x == 0) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_z.with(*bar, 0), - tZgZ(_, _, _, k), tZsZ(_, _, _, state.index())); - if constexpr (Expert3D) { - copy(tma_load_a.with(*bar, 0), - tAgA(_, _, _, k), tWsW(_, _, _, state.index())); - } else { - copy(tma_load_a.with(*bar, 0), - tAgA(_, _, _, expert_or_k_offset + k), - tWsW(_, _, _, state.index())); - } - if constexpr (ThrottleTma) { - using FullBarrier = typename Pipeline::FullBarrier; - FullBarrier::wait(bar, state.phase() ^ 1); - } - } - ++state; - } - } -} - -// ═══════════════════════════════════════════════════════════════════ -// Fused consumer — architecture-specialized on `int Compute` -// ═══════════════════════════════════════════════════════════════════ -// -// Primary template is undefined; one full specialization per supported compute -// capability provides a `run(...)` member (a function-template on -// Traits/Pipeline/TmaStoreY, since function templates can't be partially -// specialized). The free function `mlp2_t_fused_consumer` below forwards to the -// right specialization — call sites stay in the existing style. -// -// Compute=90 → Hopper / WGMMA (cooperative 2-WG, single register acc) -// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulator) -// -// y_m is the M-tile index into the output Y buffer. - -template -struct Mlp2TFusedConsumerImpl; - -// ─────────────────────────────────────────────────────────────────── -// Compute=90 — Hopper. Split axis selected by TileM: -// TileM=128 (M-split): WG_w writes rows [w*AtomTileM, (w+1)*AtomTileM). -// TileM=64 (N-split): WG_w writes cols [w*WgTileN, (w+1)*WgTileN). -// (Verbatim from the original mlp2_t_fused_consumer.) -// ─────────────────────────────────────────────────────────────────── - -template <> -struct Mlp2TFusedConsumerImpl<90> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2TFusedSmem& smem, - TmaStoreY const& tma_store_y, - int y_m, // M-tile index into Y buffer - int intermediate_dim, - int num_y_m_tiles, - int num_n_tiles, - int num_k_tiles, - // Flat-grid N-split (see producer). Defaults (-1) → blockIdx.y/gridDim.y. - int split_idx, - int num_splits) { - - using Element = typename Traits::Element; - typename Traits::TiledMma tiled_mma; - int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; - int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; - auto thr_mma = tiled_mma.get_slice(tid_in_mma); - - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); - - auto tCsZ = thr_mma.partition_A(sZ); - auto tCsW = thr_mma.partition_B(sW); - - auto acc = partition_fragment_C(tiled_mma, - Shape, Int>{}); - - auto cC = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCcC = thr_mma.partition_C(cC); - - const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; - const int my_barrier_id = 1 + my_wg; - const bool is_my_wg_leader = (tid_in_wg == 0); - - constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; - Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - // Y view: [num_y_m_tiles · TileM, intermediate_dim] row-major. - auto mY = tma_store_y.get_tma_tensor(make_shape( - static_cast(num_y_m_tiles) * Traits::TileM, - static_cast(intermediate_dim))); - auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - bool store_in_flight = false; - constexpr int K_PIPE_MMAS = 1; - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - - clear(acc); - - auto state_release = state; - int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; - - // Prologue - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - ++state; - } - // Steady state - for (int k = prologue_count; k < num_k_tiles; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - - warpgroup_wait(); - warpgroup_fence_operand(acc); - pipe.consumer_release(state_release); - ++state; - ++state_release; - } - // Drain - warpgroup_wait<0>(); - warpgroup_fence_operand(acc); - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_release(state_release); - ++state_release; - } - - // ── Epilogue ────────────────────────────────────── - // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. - // TileM=64 (N-split): WG_w owns full M (=64), cols [w*WgTileN, ...). - CUTE_UNROLL - for (int r = 0; r < Traits::NumEpiRounds; ++r) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - int chunk_start = r * Traits::EpiChunkN; - CUTE_UNROLL - for (int i = 0; i < size(acc); ++i) { - auto coord = tCcC(i); - int m_loc = get<0>(coord); - int n_loc = get<1>(coord); - int m_local, n_local; - if constexpr (Traits::kMSplit) { - m_local = m_loc - my_wg * Traits::AtomTileM; - n_local = n_loc; - } else { - m_local = m_loc; - n_local = n_loc - my_wg * Traits::WgTileN; - } - if (n_local >= chunk_start && - n_local < chunk_start + Traits::EpiChunkN) { - int chunk_n = n_local - chunk_start; - sStore(m_local, chunk_n) = static_cast(acc(i)); - } - } - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - if (is_my_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx, n_tile_idx; - if constexpr (Traits::kMSplit) { - // TileM=128: 2 row tiles per logical M tile. - m_tile_idx = 2 * y_m + my_wg; - n_tile_idx = n * Traits::NumEpiRounds + r; - } else { - // TileM=64: 1 row tile per logical M tile; WG_w handles - // its N-half. - m_tile_idx = y_m; - n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) - + my_wg * Traits::NumEpiRounds + r; - } - auto gY = local_tile(mY, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_y, cta_tma_y.partition_S(sStore), - cta_tma_y.partition_D(gY)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - if (store_in_flight) - cute::tma_store_wait<0>(); -} -}; // Mlp2TFusedConsumerImpl<90> - -// ─────────────────────────────────────────────────────────────────── -// Compute=100 — Blackwell / UMMA. -// -// No cooperative-warpgroup machinery (UMMA has no warpgroup concept): -// - One 1SM UMMA atom with M=TileM covers the whole tile; the single -// accumulator Y = Z·A lives in TMEM (no register/M-split). -// - Warp 3 issues UMMA and remains epilogue-free; warps 4-11 run the -// dual-WG epilogue. -// - Epilogue: the two consumer warpgroups split TileN; each reads its -// N-half from TMEM→regs in EpiChunkN chunks, casts to bf16, and stores -// 64-row (AtomTileM) tiles via the existing reg→SMEM→TMA path (so the host -// TMA atom is unchanged). No activation fusion. -// -// mlp2_t-specific: operand B (the weight A) is consumed MN-major -// (Traits::TiledMmaUmma is SS) — the ONLY difference vs the mlp2_fused -// Compute=100 consumer (SS). The epilogue is byte-for-byte identical. -// ─────────────────────────────────────────────────────────────────── - -template <> -struct Mlp2TFusedConsumerImpl<100> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2TFusedSmem& smem, - TmaStoreY const& tma_store_y, - int y_m, - int intermediate_dim, - int num_y_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using Element = typename Traits::Element; - constexpr int TileM = Traits::TileM; - constexpr int TileN = Traits::TileN; - constexpr int EpiChunkN = Traits::EpiChunkN; - constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile - static_assert(TileN % 2 == 0, - "Blackwell consumer splits TileN across the two consumer warpgroups"); - constexpr int WgN = TileN / 2; // per-warpgroup N width - static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); - constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per warpgroup - constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) - - // ── Thread identity ───────────────────────────────────── - // Warp 3 is the dedicated, epilogue-free UMMA producer. Warps 4..11 are - // the two aligned epilogue WGs, mirroring the forward MLP2 SM100 path. - const int warp_id = threadIdx.x / Traits::WarpSize; - const bool is_mma_warp = (warp_id == 3); - const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); - const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 - const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; - const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; - const int tmem_copy_tid = tid_wg; - const bool is_wg_leader = is_epilogue && (tid_wg == 0); - const int wg_barrier_id = 1 + wg; // 1 or 2 - constexpr int kEpilogueThreads = Traits::ConsumerThreads; - constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; - static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); - static_assert(kEpilogueThreads == 8 * Traits::WarpSize); - static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); - - // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, K-major A / MN-major - // B. Operand A = Z (K-major), operand B = A (the weight, MN-major) — the - // mlp2_t transpose orientation. Traits::TiledMmaUmma is - // make_tiled_mma(SM100_MMA_F16BF16_SS<...,Major::K,Major::MN>{}). ── - typename Traits::TiledMmaUmma tiled_mma; - auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 - - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); - - // ── One TMEM accumulator: Y = Z·A ── - auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCgC = cta_mma.partition_C(cAccFull); - auto tCtAcc = cta_mma.make_fragment_C(tCgC); - - // TMEM is allocated by the outer fused/standalone launcher once per CTA. - - // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers - // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the - // alternate stage while epilogue drains/TMA-stores MMA(n). ── - using AccPipe = typename Traits::AccumulatorPipeline; - typename AccPipe::Params acc_params; - acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer - : AccPipe::ThreadCategory::Consumer; - acc_params.producer_arv_count = 1; // one umma_arrive per commit - acc_params.consumer_arv_count = 1; // one elected releaser - acc_params.initializing_warp = 4; // warp 4 inits the acc barriers - AccPipe acc_pipe(smem.acc_pipe, acc_params, - cute::Shape{}); - auto acc_prod_state = cutlass::make_producer_start_state(); - typename AccPipe::PipelineState acc_cons_state; - - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); - const uint32_t tmem_base = smem.tmem_base; - tCtAcc.data() = tmem_base; - - // ── Per-WG store slot (AtomTileM × EpiChunkN) — same smem as Hopper ── - constexpr int store_slot_elems = AtomTileM * EpiChunkN; - Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - // Y view: [num_y_m_tiles · TileM, intermediate_dim] row-major. int64_t cast - // keeps CUTE's layout math in 64-bit at production shapes. - auto mY = tma_store_y.get_tma_tensor( - make_shape(static_cast(num_y_m_tiles) * TileM, - static_cast(intermediate_dim))); - auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); - - // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local tid). - // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) modes, - // which is what make_tmem_copy's cotiled builder requires. ── - auto epi_tile = make_tile(Int{}, Int{}); - auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) - auto tAcc_epi = flat_divide(acc_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) - auto t2r = make_tmem_copy(TmemLoadOp{}, tAcc_epi(_, _, _0{}, _0{})); - auto thr_t2r = t2r.get_slice(tmem_copy_tid); - auto tTR_tAcc = thr_t2r.partition_S(tAcc_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) - auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); - auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) - // Register fragment sized from the DEST (partition_D) shape, not partition_S. - auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); // f32 regs - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - bool store_in_flight = false; - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - - // ── Mainloop (warp 3 only): bracket the k-loop with the accumulator - // pipeline. Each k-stage waits/releases the TMA→UMMA mainloop - // pipeline; consumer_release issues the UMMA-gated smem-buffer - // arrival. ── - if (is_mma_warp) { - acc_pipe.producer_acquire(acc_prod_state); // TMEM acc free (prev epilogue done) - int acc_stage = acc_prod_state.index(); - tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); - for (int k = 0; k < num_k_tiles; ++k) { - pipe.consumer_wait(state); - auto tCsZ = cta_mma.partition_A(sZ(_, _, state.index())); - auto tCsW = cta_mma.partition_B(sW(_, _, state.index())); - auto tCrZ = cta_mma.make_fragment_A(tCsZ); - auto tCrW = cta_mma.make_fragment_B(tCsW); - CUTE_UNROLL - for (int kb = 0; kb < size<2>(tCrZ); ++kb) { - // First k-block of this n-tile clears TMEM; the rest accumulate. - tiled_mma.accumulate_ = (k == 0 && kb == 0) - ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; - gemm(tiled_mma, tCrZ(_, _, kb), tCrW(_, _, kb), tCtAcc); - } - pipe.consumer_release(state); // UMMA-gated smem-buffer release - ++state; - } - acc_pipe.producer_commit(acc_prod_state); // signal acc ready (umma_arrive) - ++acc_prod_state; - } - - // ── Epilogue: wait for the accumulator, then this WG processes its - // n-chunks [wg·NChunksHalf, +NChunksHalf). ── - if (is_epilogue) { - acc_pipe.consumer_wait(acc_cons_state); - int acc_stage = acc_cons_state.index(); - tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); - auto acc_mn_stage = tCtAcc(make_coord(_, _), _0{}, _0{}); - auto tAcc_epi_stage = flat_divide(acc_mn_stage, epi_tile); - auto tTR_tAcc_stage = thr_t2r.partition_S(tAcc_epi_stage); - - CUTE_UNROLL - for (int r = 0; r < NChunksHalf; ++r) { - int chunk = wg * NChunksHalf + r; // absolute n-chunk index - - // TMEM → registers (this chunk, full TileM rows). - copy(t2r, tTR_tAcc_stage(_, _, _, _0{}, chunk), tTR_rAcc); - cutlass::arch::fence_view_async_tmem_load(); - - // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). - CUTE_UNROLL - for (int ms = 0; ms < MSub; ++ms) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - CUTE_UNROLL - for (int i = 0; i < size(tTR_rAcc); ++i) { - int m_row = get<0>(tTR_cChunk(i)); // 0..TileM - int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN - if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) - sStore(m_row - ms * AtomTileM, n_col) = - static_cast(tTR_rAcc(i)); - } - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - - if (is_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx = MSub * y_m + ms; - int n_tile_idx = n * (TileN / EpiChunkN) + chunk; - auto gY = local_tile(mY, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_y, cta_tma_y.partition_S(sStore), - cta_tma_y.partition_D(gY)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - - // All epilogue TMEM reads are done; release the accumulator so the next - // n-tile's MMA may reuse it (one elected consumer thread arrives). - cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); - if (tid_in_epi == 0) - acc_pipe.consumer_release(acc_cons_state); - ++acc_cons_state; - } - } - if (is_epilogue && store_in_flight) - cute::tma_store_wait<0>(); - - // Keep the final MMA+epilogue barrier so the outer launcher can safely free - // TMEM after this consumer returns. - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); -#else - // The Compute=100 specialization is never dispatched on non-SM100 targets - // (host picks Compute=90 there). Trap if it is ever reached. - __trap(); +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 #endif -} -}; // Mlp2TFusedConsumerImpl<100> -// ─────────────────────────────────────────────────────────────────── -// Forwarder — keeps the existing call style. `Compute` defaults to 90 so -// current call sites (mlp2_t_fused_consumer(...)) are unchanged; pass -// mlp2_t_fused_consumer(...) to select the Blackwell path. -// ─────────────────────────────────────────────────────────────────── - -template -__device__ __forceinline__ void mlp2_t_fused_consumer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp2TFusedSmem& smem, - TmaStoreY const& tma_store_y, - int y_m, - int intermediate_dim, - int num_y_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx = -1, - int num_splits = -1) { - Mlp2TFusedConsumerImpl::template run( - pipe, state, smem, tma_store_y, y_m, intermediate_dim, - num_y_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits); -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp2_t_fused_sm90.cuh" +#else +#include "mlp2_t_fused_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm100.cuh new file mode 100644 index 00000000..b7ee9fe3 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm100.cuh @@ -0,0 +1,679 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 2 Transpose (fused tile): Y = Z @ A +// ═══════════════════════════════════════════════════════════════════ +// +// Single M-tile version for embedding in the fused MLP kernel +// (mlp_bwd.cuh Phase 1b'). Matches the production design in mlp2_t.cuh: +// - Cooperative 2-WG consumer; split axis selected by TileM (kMSplit +// in Mlp2TTraits): TileM=128 → M-split, TileM=64 → N-split. +// - One accumulator per WG covering (AtomTileM=64, AtomN) +// - Per-WG store_buf with EpiChunkN-round epilogue +// - K_PIPE_MMAS = 1 (CUTLASS-canonical) +// - Single fused Z+W TMA pipe (1 mbarrier per stage, 2 TMA copies) +// +// Differs from mlp2_t.cuh only in: +// - No outer for(m) loop — caller passes a single m (y_m for the +// output buffer index, expert_k_offset for the input weights) +// - Mlp2TFusedSmem omits pipeline storage (caller manages it) +// +// Mirrors mlp2_fused.cuh's single-tile style. + +#include "mlp2_t_sm100.cuh" + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP2-T shared memory — no pipeline storage (managed externally) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp2TFusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element smem_W[smem_W_size]; + // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these, and + // keeping them here leaves the consumer signature identical across Compute + // values (TMEM base is externally owned; the acc pipeline lives here). + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } + CUTE_DEVICE Element* W_data() { return &smem_W[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Pipeline construction helper — single fused Z+W pipeline +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ auto mlp2_t_make_pipe( + typename Traits::MainloopPipeline::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp2_t_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer +// (mlp2_t_fused_producer) and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ auto mlp2_t_make_pipe_umma( + typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused producer — single (m), all (n) under blockIdx.y stride. +// Single fused Z + W TMA pipe per k-step. +// +// Per mlp2_t: the expert is contiguous along K (= hidden_dim), so the +// caller supplies expert_k_offset = expert · num_k_tiles. A is read +// through a column-major view (shape (intermediate_dim, total_k_cols)) +// with N along intermediate_dim and K along E·hidden_dim. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_t_fused_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TFusedSmem& smem, + TmaLoadZ const& tma_load_z, + TmaLoadW const& tma_load_a, + int m, + int expert_or_k_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_k_cols, // num_experts · hidden_dim + int num_n_tiles, + int num_k_tiles, + // Flat-grid N-split: split_idx = flat_id % NSplit, num_splits = NSplit. + // Defaults (-1) fall back to blockIdx.y / gridDim.y for legacy callers. + int split_idx = -1, + int num_splits = -1) { + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + // int64_t casts: num_tokens × hidden_dim and intermediate_dim × + // total_k_cols both exceed INT_MAX at production shapes. + auto mZ = tma_load_z.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(hidden_dim))); + auto mA = [&]() { + if constexpr (Expert3D) { + return tma_load_a.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_a.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(total_k_cols))); + } + }(); + + auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); + auto cta_tma_a = tma_load_a.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_z.partition_D(sZ); + auto tWsW = cta_tma_a.partition_D(sW); + + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgZ = cta_tma_z.partition_S(gZ); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + auto gA = [&]() { + if constexpr (Expert3D) { + return local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_k_offset)); + } else { + return local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(n, _)); + } + }(); + auto tAgA = cta_tma_a.partition_S(gA); + + for (int k = 0; k < num_k_tiles; ++k) { + if constexpr (ThrottleTma) { + if (state.count() >= 1) { + using FullBarrier = typename Pipeline::FullBarrier; + auto completed = + cutlass::make_producer_start_state(); + completed.advance(state.count() - 1); + FullBarrier::wait( + pipe.producer_get_barrier(completed), + completed.phase() ^ 1); + } + } + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_z.with(*bar, 0), + tZgZ(_, _, _, k), tZsZ(_, _, _, state.index())); + if constexpr (Expert3D) { + copy(tma_load_a.with(*bar, 0), + tAgA(_, _, _, k), tWsW(_, _, _, state.index())); + } else { + copy(tma_load_a.with(*bar, 0), + tAgA(_, _, _, expert_or_k_offset + k), + tWsW(_, _, _, state.index())); + } + if constexpr (ThrottleTma) { + using FullBarrier = typename Pipeline::FullBarrier; + FullBarrier::wait(bar, state.phase() ^ 1); + } + } + ++state; + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused consumer — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported compute +// capability provides a `run(...)` member (a function-template on +// Traits/Pipeline/TmaStoreY, since function templates can't be partially +// specialized). The free function `mlp2_t_fused_consumer` below forwards to the +// right specialization — call sites stay in the existing style. +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, single register acc) +// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulator) +// +// y_m is the M-tile index into the output Y buffer. + +template +struct Mlp2TFusedConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. Split axis selected by TileM: +// TileM=128 (M-split): WG_w writes rows [w*AtomTileM, (w+1)*AtomTileM). +// TileM=64 (N-split): WG_w writes cols [w*WgTileN, (w+1)*WgTileN). +// (Verbatim from the original mlp2_t_fused_consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp2TFusedConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TFusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, // M-tile index into Y buffer + int intermediate_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + // Flat-grid N-split (see producer). Defaults (-1) → blockIdx.y/gridDim.y. + int split_idx, + int num_splits) { + + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; + const int my_barrier_id = 1 + my_wg; + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // Y view: [num_y_m_tiles · TileM, intermediate_dim] row-major. + auto mY = tma_store_y.get_tma_tensor(make_shape( + static_cast(num_y_m_tiles) * Traits::TileM, + static_cast(intermediate_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. + // TileM=64 (N-split): WG_w owns full M (=64), cols [w*WgTileN, ...). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: 2 row tiles per logical M tile. + m_tile_idx = 2 * y_m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M tile; WG_w handles + // its N-half. + m_tile_idx = y_m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp2TFusedConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. +// +// No cooperative-warpgroup machinery (UMMA has no warpgroup concept): +// - One 1SM UMMA atom with M=TileM covers the whole tile; the single +// accumulator Y = Z·A lives in TMEM (no register/M-split). +// - Warp 3 issues UMMA and remains epilogue-free; warps 4-11 run the +// dual-WG epilogue. +// - Epilogue: the two consumer warpgroups split TileN; each reads its +// N-half from TMEM→regs in EpiChunkN chunks, casts to bf16, and stores +// 64-row (AtomTileM) tiles via the existing reg→SMEM→TMA path (so the host +// TMA atom is unchanged). No activation fusion. +// +// mlp2_t-specific: operand B (the weight A) is consumed MN-major +// (Traits::TiledMmaUmma is SS) — the ONLY difference vs the mlp2_fused +// Compute=100 consumer (SS). The epilogue is byte-for-byte identical. +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp2TFusedConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TFusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, + int intermediate_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two consumer warpgroups"); + constexpr int WgN = TileN / 2; // per-warpgroup N width + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per warpgroup + constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) + + // ── Thread identity ───────────────────────────────────── + // Warp 3 is the dedicated, epilogue-free UMMA producer. Warps 4..11 are + // the two aligned epilogue WGs, mirroring the forward MLP2 SM100 path. + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + const int tmem_copy_tid = tid_wg; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; // 1 or 2 + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, K-major A / MN-major + // B. Operand A = Z (K-major), operand B = A (the weight, MN-major) — the + // mlp2_t transpose orientation. Traits::TiledMmaUmma is + // make_tiled_mma(SM100_MMA_F16BF16_SS<...,Major::K,Major::MN>{}). ── + typename Traits::TiledMmaUmma tiled_mma; + auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + // ── One TMEM accumulator: Y = Z·A ── + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the + // alternate stage while epilogue drains/TMA-stores MMA(n). ── + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; // one umma_arrive per commit + acc_params.consumer_arv_count = 1; // one elected releaser + acc_params.initializing_warp = 4; // warp 4 inits the acc barriers + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + // ── Per-WG store slot (AtomTileM × EpiChunkN) — same smem as Hopper ── + constexpr int store_slot_elems = AtomTileM * EpiChunkN; + Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // Y view: [num_y_m_tiles · TileM, intermediate_dim] row-major. int64_t cast + // keeps CUTE's layout math in 64-bit at production shapes. + auto mY = tma_store_y.get_tma_tensor( + make_shape(static_cast(num_y_m_tiles) * TileM, + static_cast(intermediate_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local tid). + // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) modes, + // which is what make_tmem_copy's cotiled builder requires. ── + auto epi_tile = make_tile(Int{}, Int{}); + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto tAcc_epi = flat_divide(acc_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto t2r = make_tmem_copy(TmemLoadOp{}, tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tmem_copy_tid); + auto tTR_tAcc = thr_t2r.partition_S(tAcc_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + // Register fragment sized from the DEST (partition_D) shape, not partition_S. + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); // f32 regs + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + // ── Mainloop (warp 3 only): bracket the k-loop with the accumulator + // pipeline. Each k-stage waits/releases the TMA→UMMA mainloop + // pipeline; consumer_release issues the UMMA-gated smem-buffer + // arrival. ── + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); // TMEM acc free (prev epilogue done) + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + for (int k = 0; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + auto tCsZ = cta_mma.partition_A(sZ(_, _, state.index())); + auto tCsW = cta_mma.partition_B(sW(_, _, state.index())); + auto tCrZ = cta_mma.make_fragment_A(tCsZ); + auto tCrW = cta_mma.make_fragment_B(tCsW); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrZ); ++kb) { + // First k-block of this n-tile clears TMEM; the rest accumulate. + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrZ(_, _, kb), tCrW(_, _, kb), tCtAcc); + } + pipe.consumer_release(state); // UMMA-gated smem-buffer release + ++state; + } + acc_pipe.producer_commit(acc_prod_state); // signal acc ready (umma_arrive) + ++acc_prod_state; + } + + // ── Epilogue: wait for the accumulator, then this WG processes its + // n-chunks [wg·NChunksHalf, +NChunksHalf). ── + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc_stage = thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; // absolute n-chunk index + + // TMEM → registers (this chunk, full TileM rows). + copy(t2r, tTR_tAcc_stage(_, _, _, _0{}, chunk), tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + + // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rAcc); ++i) { + int m_row = get<0>(tTR_cChunk(i)); // 0..TileM + int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN + if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) + sStore(m_row - ms * AtomTileM, n_col) = + static_cast(tTR_rAcc(i)); + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx = MSub * y_m + ms; + int n_tile_idx = n * (TileN / EpiChunkN) + chunk; + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // All epilogue TMEM reads are done; release the accumulator so the next + // n-tile's MMA may reuse it (one elected consumer thread arrives). + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + // Keep the final MMA+epilogue barrier so the outer launcher can safely free + // TMEM after this consumer returns. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + // The Compute=100 specialization is never dispatched on non-SM100 targets + // (host picks Compute=90 there). Trap if it is ever reached. + __trap(); +#endif +} +}; // Mlp2TFusedConsumerImpl<100> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — keeps the existing call style. `Compute` defaults to 90 so +// current call sites (mlp2_t_fused_consumer(...)) are unchanged; pass +// mlp2_t_fused_consumer(...) to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp2_t_fused_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TFusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, + int intermediate_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx = -1, + int num_splits = -1) { + Mlp2TFusedConsumerImpl::template run( + pipe, state, smem, tma_store_y, y_m, intermediate_dim, + num_y_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm90.cuh new file mode 100644 index 00000000..cb6bb78c --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm90.cuh @@ -0,0 +1,679 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 2 Transpose (fused tile): Y = Z @ A +// ═══════════════════════════════════════════════════════════════════ +// +// Single M-tile version for embedding in the fused MLP kernel +// (mlp_bwd.cuh Phase 1b'). Matches the production design in mlp2_t.cuh: +// - Cooperative 2-WG consumer; split axis selected by TileM (kMSplit +// in Mlp2TTraits): TileM=128 → M-split, TileM=64 → N-split. +// - One accumulator per WG covering (AtomTileM=64, AtomN) +// - Per-WG store_buf with EpiChunkN-round epilogue +// - K_PIPE_MMAS = 1 (CUTLASS-canonical) +// - Single fused Z+W TMA pipe (1 mbarrier per stage, 2 TMA copies) +// +// Differs from mlp2_t.cuh only in: +// - No outer for(m) loop — caller passes a single m (y_m for the +// output buffer index, expert_k_offset for the input weights) +// - Mlp2TFusedSmem omits pipeline storage (caller manages it) +// +// Mirrors mlp2_fused.cuh's single-tile style. + +#include "mlp2_t_sm90.cuh" + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP2-T shared memory — no pipeline storage (managed externally) +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp2TFusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element smem_W[smem_W_size]; + // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these, and + // keeping them here leaves the consumer signature identical across Compute + // values (TMEM base is externally owned; the acc pipeline lives here). + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } + CUTE_DEVICE Element* W_data() { return &smem_W[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Pipeline construction helper — single fused Z+W pipeline +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ auto mlp2_t_make_pipe( + typename Traits::MainloopPipeline::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp2_t_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer +// (mlp2_t_fused_producer) and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ auto mlp2_t_make_pipe_umma( + typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused producer — single (m), all (n) under blockIdx.y stride. +// Single fused Z + W TMA pipe per k-step. +// +// Per mlp2_t: the expert is contiguous along K (= hidden_dim), so the +// caller supplies expert_k_offset = expert · num_k_tiles. A is read +// through a column-major view (shape (intermediate_dim, total_k_cols)) +// with N along intermediate_dim and K along E·hidden_dim. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_t_fused_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TFusedSmem& smem, + TmaLoadZ const& tma_load_z, + TmaLoadW const& tma_load_a, + int m, + int expert_or_k_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_k_cols, // num_experts · hidden_dim + int num_n_tiles, + int num_k_tiles, + // Flat-grid N-split: split_idx = flat_id % NSplit, num_splits = NSplit. + // Defaults (-1) fall back to blockIdx.y / gridDim.y for legacy callers. + int split_idx = -1, + int num_splits = -1) { + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + // int64_t casts: num_tokens × hidden_dim and intermediate_dim × + // total_k_cols both exceed INT_MAX at production shapes. + auto mZ = tma_load_z.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(hidden_dim))); + auto mA = [&]() { + if constexpr (Expert3D) { + return tma_load_a.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_load_a.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(total_k_cols))); + } + }(); + + auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); + auto cta_tma_a = tma_load_a.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_z.partition_D(sZ); + auto tWsW = cta_tma_a.partition_D(sW); + + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgZ = cta_tma_z.partition_S(gZ); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + auto gA = [&]() { + if constexpr (Expert3D) { + return local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(n, _, expert_or_k_offset)); + } else { + return local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(n, _)); + } + }(); + auto tAgA = cta_tma_a.partition_S(gA); + + for (int k = 0; k < num_k_tiles; ++k) { + if constexpr (ThrottleTma) { + if (state.count() >= 1) { + using FullBarrier = typename Pipeline::FullBarrier; + auto completed = + cutlass::make_producer_start_state(); + completed.advance(state.count() - 1); + FullBarrier::wait( + pipe.producer_get_barrier(completed), + completed.phase() ^ 1); + } + } + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_z.with(*bar, 0), + tZgZ(_, _, _, k), tZsZ(_, _, _, state.index())); + if constexpr (Expert3D) { + copy(tma_load_a.with(*bar, 0), + tAgA(_, _, _, k), tWsW(_, _, _, state.index())); + } else { + copy(tma_load_a.with(*bar, 0), + tAgA(_, _, _, expert_or_k_offset + k), + tWsW(_, _, _, state.index())); + } + if constexpr (ThrottleTma) { + using FullBarrier = typename Pipeline::FullBarrier; + FullBarrier::wait(bar, state.phase() ^ 1); + } + } + ++state; + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Fused consumer — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported compute +// capability provides a `run(...)` member (a function-template on +// Traits/Pipeline/TmaStoreY, since function templates can't be partially +// specialized). The free function `mlp2_t_fused_consumer` below forwards to the +// right specialization — call sites stay in the existing style. +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, single register acc) +// Compute=100 → Blackwell / UMMA (single-warp issue, TMEM accumulator) +// +// y_m is the M-tile index into the output Y buffer. + +template +struct Mlp2TFusedConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. Split axis selected by TileM: +// TileM=128 (M-split): WG_w writes rows [w*AtomTileM, (w+1)*AtomTileM). +// TileM=64 (N-split): WG_w writes cols [w*WgTileN, (w+1)*WgTileN). +// (Verbatim from the original mlp2_t_fused_consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp2TFusedConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TFusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, // M-tile index into Y buffer + int intermediate_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + // Flat-grid N-split (see producer). Defaults (-1) → blockIdx.y/gridDim.y. + int split_idx, + int num_splits) { + + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; + const int my_barrier_id = 1 + my_wg; + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // Y view: [num_y_m_tiles · TileM, intermediate_dim] row-major. + auto mY = tma_store_y.get_tma_tensor(make_shape( + static_cast(num_y_m_tiles) * Traits::TileM, + static_cast(intermediate_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. + // TileM=64 (N-split): WG_w owns full M (=64), cols [w*WgTileN, ...). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: 2 row tiles per logical M tile. + m_tile_idx = 2 * y_m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M tile; WG_w handles + // its N-half. + m_tile_idx = y_m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp2TFusedConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. +// +// No cooperative-warpgroup machinery (UMMA has no warpgroup concept): +// - One 1SM UMMA atom with M=TileM covers the whole tile; the single +// accumulator Y = Z·A lives in TMEM (no register/M-split). +// - Warp 3 issues UMMA and remains epilogue-free; warps 4-11 run the +// dual-WG epilogue. +// - Epilogue: the two consumer warpgroups split TileN; each reads its +// N-half from TMEM→regs in EpiChunkN chunks, casts to bf16, and stores +// 64-row (AtomTileM) tiles via the existing reg→SMEM→TMA path (so the host +// TMA atom is unchanged). No activation fusion. +// +// mlp2_t-specific: operand B (the weight A) is consumed MN-major +// (Traits::TiledMmaUmma is SS) — the ONLY difference vs the mlp2_fused +// Compute=100 consumer (SS). The epilogue is byte-for-byte identical. +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp2TFusedConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TFusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, + int intermediate_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two consumer warpgroups"); + constexpr int WgN = TileN / 2; // per-warpgroup N width + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per warpgroup + constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) + + // ── Thread identity ───────────────────────────────────── + // Warp 3 is the dedicated, epilogue-free UMMA producer. Warps 4..11 are + // the two aligned epilogue WGs, mirroring the forward MLP2 SM100 path. + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + const int tmem_copy_tid = tid_wg; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; // 1 or 2 + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, K-major A / MN-major + // B. Operand A = Z (K-major), operand B = A (the weight, MN-major) — the + // mlp2_t transpose orientation. Traits::TiledMmaUmma is + // make_tiled_mma(SM100_MMA_F16BF16_SS<...,Major::K,Major::MN>{}). ── + typename Traits::TiledMmaUmma tiled_mma; + auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + // ── One TMEM accumulator: Y = Z·A ── + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages TMEM stages let MMA(n+1) fill the + // alternate stage while epilogue drains/TMA-stores MMA(n). ── + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; // one umma_arrive per commit + acc_params.consumer_arv_count = 1; // one elected releaser + acc_params.initializing_warp = 4; // warp 4 inits the acc barriers + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + // ── Per-WG store slot (AtomTileM × EpiChunkN) — same smem as Hopper ── + constexpr int store_slot_elems = AtomTileM * EpiChunkN; + Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // Y view: [num_y_m_tiles · TileM, intermediate_dim] row-major. int64_t cast + // keeps CUTE's layout math in 64-bit at production shapes. + auto mY = tma_store_y.get_tma_tensor( + make_shape(static_cast(num_y_m_tiles) * TileM, + static_cast(intermediate_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local tid). + // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) modes, + // which is what make_tmem_copy's cotiled builder requires. ── + auto epi_tile = make_tile(Int{}, Int{}); + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto tAcc_epi = flat_divide(acc_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto t2r = make_tmem_copy(TmemLoadOp{}, tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tmem_copy_tid); + auto tTR_tAcc = thr_t2r.partition_S(tAcc_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + // Register fragment sized from the DEST (partition_D) shape, not partition_S. + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); // f32 regs + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + // ── Mainloop (warp 3 only): bracket the k-loop with the accumulator + // pipeline. Each k-stage waits/releases the TMA→UMMA mainloop + // pipeline; consumer_release issues the UMMA-gated smem-buffer + // arrival. ── + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); // TMEM acc free (prev epilogue done) + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + for (int k = 0; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + auto tCsZ = cta_mma.partition_A(sZ(_, _, state.index())); + auto tCsW = cta_mma.partition_B(sW(_, _, state.index())); + auto tCrZ = cta_mma.make_fragment_A(tCsZ); + auto tCrW = cta_mma.make_fragment_B(tCsW); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrZ); ++kb) { + // First k-block of this n-tile clears TMEM; the rest accumulate. + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrZ(_, _, kb), tCrW(_, _, kb), tCtAcc); + } + pipe.consumer_release(state); // UMMA-gated smem-buffer release + ++state; + } + acc_pipe.producer_commit(acc_prod_state); // signal acc ready (umma_arrive) + ++acc_prod_state; + } + + // ── Epilogue: wait for the accumulator, then this WG processes its + // n-chunks [wg·NChunksHalf, +NChunksHalf). ── + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc_stage = thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; // absolute n-chunk index + + // TMEM → registers (this chunk, full TileM rows). + copy(t2r, tTR_tAcc_stage(_, _, _, _0{}, chunk), tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + + // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rAcc); ++i) { + int m_row = get<0>(tTR_cChunk(i)); // 0..TileM + int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN + if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) + sStore(m_row - ms * AtomTileM, n_col) = + static_cast(tTR_rAcc(i)); + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx = MSub * y_m + ms; + int n_tile_idx = n * (TileN / EpiChunkN) + chunk; + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // All epilogue TMEM reads are done; release the accumulator so the next + // n-tile's MMA may reuse it (one elected consumer thread arrives). + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + // Keep the final MMA+epilogue barrier so the outer launcher can safely free + // TMEM after this consumer returns. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + // The Compute=100 specialization is never dispatched on non-SM100 targets + // (host picks Compute=90 there). Trap if it is ever reached. + __trap(); +#endif +} +}; // Mlp2TFusedConsumerImpl<100> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — keeps the existing call style. `Compute` defaults to 90 so +// current call sites (mlp2_t_fused_consumer(...)) are unchanged; pass +// mlp2_t_fused_consumer(...) to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp2_t_fused_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TFusedSmem& smem, + TmaStoreY const& tma_store_y, + int y_m, + int intermediate_dim, + int num_y_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx = -1, + int num_splits = -1) { + Mlp2TFusedConsumerImpl::template run( + pipe, state, smem, tma_store_y, y_m, intermediate_dim, + num_y_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm100.cuh new file mode 100644 index 00000000..10af3b99 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm100.cuh @@ -0,0 +1,540 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 2 (transposed): Y = Z @ A +// ═══════════════════════════════════════════════════════════════════ +// +// Cooperative 2-WG consumer; split axis selected by TileM: +// TileM=128 → Layout<_2,_1,_1>, split along M (each WG owns 64×TileN) +// TileM=64 → Layout<_1,_2,_1>, split along N (each WG owns 64×TileN/2) +// One accumulator per WG covering (AtomTileM=64, AtomN). Each WG writes +// its strip independently — no cross-WG NamedBarrier sync in the epilogue. +// +// Difference from mlp2_fused.cuh: A is read through a COLUMN-MAJOR TMA view +// (the physical layout is the same per-expert [hidden_dim, +// intermediate_dim] row-major buffer, but the column-major view makes +// intermediate_dim the row axis so it lines up as the N axis of the +// GEMM). Output Y has shape [num_tokens, intermediate_dim]. +// - num_n_tiles = intermediate_dim / TileN (the N axis) +// - num_k_tiles = hidden_dim / TileK (the K axis) +// - expert offset is along K: expert · num_k_tiles +// +// A single fused TMA pipeline carries Z + W per k-step; transaction = +// bytes(Z) + bytes(W). Producer issues two TMA copies against the same +// mbarrier per acquire. +// +// CUTLASS-canonical patterns: +// - K_PIPE_MMAS = 1 +// - __launch_bounds__(NumThreads, 1) — applied by the caller kernel +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for Z / A) +// WG0 warps 1-3 : Idle on SM90; warp 3 issues UMMA on SM100 +// WG1 (warps 4-7) : Consumer atom 0 — owns m=[0, AtomTileM) +// WG2 (warps 8-11) : Consumer atom 1 — owns m=[AtomTileM, TileM) +// +// Per `expert_ids`: one int per TileM=128 m-block. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm100.cuh" +#include "math_sm100.cuh" +#include "tmem_load_op_sm100.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer +// specialization in mlp2_t_fused.cuh. These headers are header-only and safe +// to include under an sm_90a build — the tcgen05 PTX they emit is guarded by +// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on +// __CUDA_ARCH__ >= 1000) ever instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 128, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 32 +> +struct Mlp2TTraits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + + static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); + + // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate + // along M (each owns 64 rows × TileN). With TileM=64 the WGs cooperate + // along N (each owns 64 rows × TileN/2). AtomTileM is the per-WG M + // extent — always 64. Mirrors mlp2_fused.cuh's kMSplit/N-split pattern. + static constexpr bool kMSplit = (TileM == 128); + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); + + static_assert(TileN <= 256, "TileN must fit TMA box limit"); + static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); + static_assert(kMSplit || (TileN % 2 == 0), + "TileN must be even when TileM=64 (cooperative N-split halves)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + // A is loaded through a column-major TMA view (stride=(1, I)), so the + // W operand in shared memory must be MN-major to match. The wgmma atom + // for operand B must accept an MN-major source (GmmaSelectorKMN — note + // the trailing N in "KMN" denotes MN-major B). Z stays K-major (Z is + // row-major in gmem with K=H contiguous-ish given the layout CUTE + // produces for shape (M, K), stride (K, 1)), so its atom is the regular + // K-major SW128 atom. + // + // AtomN is the per-WG WGMMA atom N — combined across both WGs it must + // cover exactly TileN cols (= sW's N extent). With WgLayout=<1,2,1> + // and AtomN=TileN, the TiledMMA would cover 2·TileN cols and WG1's + // partition_B(sW) slice would read past sW's TileN-wide buffer → IMA. + static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); + using GmmaAtom = typename GmmaSelectorKMN::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom; + using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom; + + using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{}, + Shape, Int>{})); + // Step<_2, _1>: tile_to_shape interprets the inner SmemLayoutAtomW as + // MN-major (M is contiguous in smem), so passing the tile shape with + // the swap step makes the resulting layout's M dim contiguous — which + // is what the column-major A gmem feed expects. + using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtomW{}, + Shape, Int>{}, + Step<_2, _1>{})); + + using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{}, + Shape, Int, Int>{})); + using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{}, + Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + + // Per-WG store buffer slot: (AtomTileM, EpiChunkN). + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore_1 = SmemLayoutStoreSlot; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop + // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the + // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the single + // TMEM accumulator from the MMA warp to the epilogue warps. See + // Mlp2TFusedConsumerImpl<100> in mlp2_t_fused.cuh. The Blackwell launcher + // must build the mainloop pipe with num_consumers=1 (one umma_arrive per + // stage). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Double-buffered SM100 TMEM accumulators: each stage owns one TileN-wide + // accumulator, allowing MMA(n+1) to fill the alternate stage while the + // epilogue drains/TMA-stores MMA(n). + static constexpr int AccStages = 2; + static_assert(AccStages * TileN <= 512, + "SM100 MLP2-T accumulator stages must fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + // SM100 UMMA TiledMMA — one 1SM tcgen05 atom spanning the whole + // (TileM, TileN) tile with the accumulator in TMEM. THE mlp2_t-specific + // choice: operand A = Z is K-major (UMMA::Major::K), operand B = A (the + // weight) is MN-major (UMMA::Major::MN) — mlp2_fused uses . The B + // operand descriptor is produced by make_umma_desc reading + // SmemLayoutW's (MN-contiguous, SW128) stride; SmemLayoutW stays built from + // Layout_MN_SW128_Atom + Step<_2,_1>, which is exactly the canonical + // Major-MN layout make_umma_desc accepts (mma_traits_sm100.hpp). + using TiledMmaUmma = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS{})); + + static constexpr int TmaTransBytesZ = + static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesW = + static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); + // Fused pipeline: Z + W share one mbarrier per stage. + static constexpr int TmaTransBytes = TmaTransBytesZ + TmaTransBytesW; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; +}; + +template +using Mlp2TMainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Shared Memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp2TSmem { + using Element = typename Traits::Element; + + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element smem_W[smem_W_size]; + // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + typename Traits::MainloopPipeline::SharedStorage pipe_storage; + + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } + CUTE_DEVICE Element* W_data() { return &smem_W[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// TMA Producer (Warp 0) — one acquire per k-step, issues Z + W +// against the same mbarrier. +// - Z view: [num_tokens, hidden_dim] row-major (Z's K axis = H). +// - A view: [intermediate_dim, total_k_cols] column-major, where +// total_k_cols = num_experts · hidden_dim. Per-expert tile at +// k = expert · num_k_tiles + k_local. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_t_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TSmem& smem, + TmaLoadZ const& tma_load_z, + TmaLoadW const& tma_load_a, + const int* expert_ids, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int total_k_cols, // num_experts · hidden_dim + int num_m_tiles, + int num_n_tiles, // intermediate_dim / TileN + int num_k_tiles) { // hidden_dim / TileK + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + auto mZ = tma_load_z.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(hidden_dim))); + // Column-major view of A: shape (intermediate_dim, total_k_cols) + // — same underlying buffer as the row-major [E·H, I] tensor in + // mlp2_fused.cuh, but the box geometry makes intermediate_dim the row + // axis so it lines up with the GEMM N dimension. + auto mA = tma_load_a.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(total_k_cols))); + + auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); + auto cta_tma_a = tma_load_a.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_z.partition_D(sZ); + auto tWsW = cta_tma_a.partition_D(sW); + + int n_start = blockIdx.y; + int n_stride = gridDim.y; + + for (int m = blockIdx.x; m < num_m_tiles; m += gridDim.x) { + int expert = expert_ids[m]; + // In mlp2_t the expert is contiguous in the K direction + // (= hidden_dim), so the offset is along k_tile, not n_tile. + int expert_k_offset = expert * num_k_tiles; + + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgZ = cta_tma_z.partition_S(gZ); + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + // A tile: (TileN rows of I, TileK rows of K=H), positioned + // at (n along I, expert_k_offset+k along total_k_cols). + auto gW = local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(n, _)); + auto tWgW = cta_tma_a.partition_S(gW); + + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_z.with(*bar, 0), + tZgZ(_, _, _, k), tZsZ(_, _, _, state.index())); + copy(tma_load_a.with(*bar, 0), + tWgW(_, _, _, expert_k_offset + k), + tWsW(_, _, _, state.index())); + } + ++state; + } + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Cooperative 2-WG Consumer — M-split, single acc per WG +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_t_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TSmem& smem, + TmaStoreY const& tma_store_y, + int intermediate_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles) { + + (void)Compute; + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // Y view: [num_tokens, intermediate_dim] row-major. + auto mY = tma_store_y.get_tma_tensor(make_shape( + static_cast(num_m_tiles) * Traits::TileM, + static_cast(intermediate_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + int n_start = blockIdx.y; + int n_stride = gridDim.y; + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; + + for (int m = blockIdx.x; m < num_m_tiles; m += gridDim.x) { + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. + // TileM=64 (N-split): WG_w owns full M (=64), cols [w*WgTileN, ...). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: 2 row tiles per logical M tile. + m_tile_idx = 2 * m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M tile; WG_w handles + // its N-half. + m_tile_idx = m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} + +// ═══════════════════════════════════════════════════════════════════ +// mlp2_t_fwd — device orchestrator +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_t_fwd( + Mlp2TSmem& smem, + TmaLoadZ const& tma_load_z, + TmaLoadW const& tma_load_a, + TmaStoreY const& tma_store_y, + const int* expert_ids, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int total_k_cols, + int num_m_tiles) { + + using Pipeline = typename Traits::MainloopPipeline; + using PipeState = typename Traits::PipelineState; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + int num_n_tiles = intermediate_dim / Traits::TileN; + int num_k_tiles = hidden_dim / Traits::TileK; + + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + cute::prefetch_tma_descriptor(tma_load_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_y.get_tma_descriptor()); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + Pipeline pipe(smem.pipe_storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); + + __syncthreads(); + + PipeState state = is_producer ? + cutlass::make_producer_start_state() : PipeState{}; + PipeState cons_state; + + if (is_producer) { + mlp2_t_producer( + pipe, state, + smem, tma_load_z, tma_load_a, + expert_ids, + num_tokens, hidden_dim, intermediate_dim, total_k_cols, + num_m_tiles, num_n_tiles, num_k_tiles); + } + + if (is_consumer) { + mlp2_t_consumer( + pipe, cons_state, + smem, tma_store_y, + intermediate_dim, + num_m_tiles, num_n_tiles, num_k_tiles); + } + + __syncthreads(); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm90.cuh new file mode 100644 index 00000000..5a5d82d2 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm90.cuh @@ -0,0 +1,540 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 2 (transposed): Y = Z @ A +// ═══════════════════════════════════════════════════════════════════ +// +// Cooperative 2-WG consumer; split axis selected by TileM: +// TileM=128 → Layout<_2,_1,_1>, split along M (each WG owns 64×TileN) +// TileM=64 → Layout<_1,_2,_1>, split along N (each WG owns 64×TileN/2) +// One accumulator per WG covering (AtomTileM=64, AtomN). Each WG writes +// its strip independently — no cross-WG NamedBarrier sync in the epilogue. +// +// Difference from mlp2_fused.cuh: A is read through a COLUMN-MAJOR TMA view +// (the physical layout is the same per-expert [hidden_dim, +// intermediate_dim] row-major buffer, but the column-major view makes +// intermediate_dim the row axis so it lines up as the N axis of the +// GEMM). Output Y has shape [num_tokens, intermediate_dim]. +// - num_n_tiles = intermediate_dim / TileN (the N axis) +// - num_k_tiles = hidden_dim / TileK (the K axis) +// - expert offset is along K: expert · num_k_tiles +// +// A single fused TMA pipeline carries Z + W per k-step; transaction = +// bytes(Z) + bytes(W). Producer issues two TMA copies against the same +// mbarrier per acquire. +// +// CUTLASS-canonical patterns: +// - K_PIPE_MMAS = 1 +// - __launch_bounds__(NumThreads, 1) — applied by the caller kernel +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for Z / A) +// WG0 warps 1-3 : Idle on SM90; warp 3 issues UMMA on SM100 +// WG1 (warps 4-7) : Consumer atom 0 — owns m=[0, AtomTileM) +// WG2 (warps 8-11) : Consumer atom 1 — owns m=[AtomTileM, TileM) +// +// Per `expert_ids`: one int per TileM=128 m-block. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm90.cuh" +#include "math_sm90.cuh" +#include "tmem_load_op_sm90.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer +// specialization in mlp2_t_fused.cuh. These headers are header-only and safe +// to include under an sm_90a build — the tcgen05 PTX they emit is guarded by +// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on +// __CUDA_ARCH__ >= 1000) ever instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 128, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 32 +> +struct Mlp2TTraits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + + static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); + + // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate + // along M (each owns 64 rows × TileN). With TileM=64 the WGs cooperate + // along N (each owns 64 rows × TileN/2). AtomTileM is the per-WG M + // extent — always 64. Mirrors mlp2_fused.cuh's kMSplit/N-split pattern. + static constexpr bool kMSplit = (TileM == 128); + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); + + static_assert(TileN <= 256, "TileN must fit TMA box limit"); + static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); + static_assert(kMSplit || (TileN % 2 == 0), + "TileN must be even when TileM=64 (cooperative N-split halves)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + // A is loaded through a column-major TMA view (stride=(1, I)), so the + // W operand in shared memory must be MN-major to match. The wgmma atom + // for operand B must accept an MN-major source (GmmaSelectorKMN — note + // the trailing N in "KMN" denotes MN-major B). Z stays K-major (Z is + // row-major in gmem with K=H contiguous-ish given the layout CUTE + // produces for shape (M, K), stride (K, 1)), so its atom is the regular + // K-major SW128 atom. + // + // AtomN is the per-WG WGMMA atom N — combined across both WGs it must + // cover exactly TileN cols (= sW's N extent). With WgLayout=<1,2,1> + // and AtomN=TileN, the TiledMMA would cover 2·TileN cols and WG1's + // partition_B(sW) slice would read past sW's TileN-wide buffer → IMA. + static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); + using GmmaAtom = typename GmmaSelectorKMN::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom; + using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom; + + using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{}, + Shape, Int>{})); + // Step<_2, _1>: tile_to_shape interprets the inner SmemLayoutAtomW as + // MN-major (M is contiguous in smem), so passing the tile shape with + // the swap step makes the resulting layout's M dim contiguous — which + // is what the column-major A gmem feed expects. + using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtomW{}, + Shape, Int>{}, + Step<_2, _1>{})); + + using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{}, + Shape, Int, Int>{})); + using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{}, + Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + + // Per-WG store buffer slot: (AtomTileM, EpiChunkN). + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore_1 = SmemLayoutStoreSlot; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop + // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the + // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the single + // TMEM accumulator from the MMA warp to the epilogue warps. See + // Mlp2TFusedConsumerImpl<100> in mlp2_t_fused.cuh. The Blackwell launcher + // must build the mainloop pipe with num_consumers=1 (one umma_arrive per + // stage). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Double-buffered SM100 TMEM accumulators: each stage owns one TileN-wide + // accumulator, allowing MMA(n+1) to fill the alternate stage while the + // epilogue drains/TMA-stores MMA(n). + static constexpr int AccStages = 2; + static_assert(AccStages * TileN <= 512, + "SM100 MLP2-T accumulator stages must fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + // SM100 UMMA TiledMMA — one 1SM tcgen05 atom spanning the whole + // (TileM, TileN) tile with the accumulator in TMEM. THE mlp2_t-specific + // choice: operand A = Z is K-major (UMMA::Major::K), operand B = A (the + // weight) is MN-major (UMMA::Major::MN) — mlp2_fused uses . The B + // operand descriptor is produced by make_umma_desc reading + // SmemLayoutW's (MN-contiguous, SW128) stride; SmemLayoutW stays built from + // Layout_MN_SW128_Atom + Step<_2,_1>, which is exactly the canonical + // Major-MN layout make_umma_desc accepts (mma_traits_sm100.hpp). + using TiledMmaUmma = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS{})); + + static constexpr int TmaTransBytesZ = + static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesW = + static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); + // Fused pipeline: Z + W share one mbarrier per stage. + static constexpr int TmaTransBytes = TmaTransBytesZ + TmaTransBytesW; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; +}; + +template +using Mlp2TMainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Shared Memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp2TSmem { + using Element = typename Traits::Element; + + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element smem_W[smem_W_size]; + // Per-WG store buffer: WG1 uses store_buf[0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + typename Traits::MainloopPipeline::SharedStorage pipe_storage; + + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } + CUTE_DEVICE Element* W_data() { return &smem_W[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// TMA Producer (Warp 0) — one acquire per k-step, issues Z + W +// against the same mbarrier. +// - Z view: [num_tokens, hidden_dim] row-major (Z's K axis = H). +// - A view: [intermediate_dim, total_k_cols] column-major, where +// total_k_cols = num_experts · hidden_dim. Per-expert tile at +// k = expert · num_k_tiles + k_local. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_t_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TSmem& smem, + TmaLoadZ const& tma_load_z, + TmaLoadW const& tma_load_a, + const int* expert_ids, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int total_k_cols, // num_experts · hidden_dim + int num_m_tiles, + int num_n_tiles, // intermediate_dim / TileN + int num_k_tiles) { // hidden_dim / TileK + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + auto mZ = tma_load_z.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(hidden_dim))); + // Column-major view of A: shape (intermediate_dim, total_k_cols) + // — same underlying buffer as the row-major [E·H, I] tensor in + // mlp2_fused.cuh, but the box geometry makes intermediate_dim the row + // axis so it lines up with the GEMM N dimension. + auto mA = tma_load_a.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(total_k_cols))); + + auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); + auto cta_tma_a = tma_load_a.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_z.partition_D(sZ); + auto tWsW = cta_tma_a.partition_D(sW); + + int n_start = blockIdx.y; + int n_stride = gridDim.y; + + for (int m = blockIdx.x; m < num_m_tiles; m += gridDim.x) { + int expert = expert_ids[m]; + // In mlp2_t the expert is contiguous in the K direction + // (= hidden_dim), so the offset is along k_tile, not n_tile. + int expert_k_offset = expert * num_k_tiles; + + auto gZ = local_tile(mZ, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgZ = cta_tma_z.partition_S(gZ); + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + // A tile: (TileN rows of I, TileK rows of K=H), positioned + // at (n along I, expert_k_offset+k along total_k_cols). + auto gW = local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(n, _)); + auto tWgW = cta_tma_a.partition_S(gW); + + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_z.with(*bar, 0), + tZgZ(_, _, _, k), tZsZ(_, _, _, state.index())); + copy(tma_load_a.with(*bar, 0), + tWgW(_, _, _, expert_k_offset + k), + tWsW(_, _, _, state.index())); + } + ++state; + } + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Cooperative 2-WG Consumer — M-split, single acc per WG +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_t_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp2TSmem& smem, + TmaStoreY const& tma_store_y, + int intermediate_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles) { + + (void)Compute; + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.W_data()), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + // Y view: [num_tokens, intermediate_dim] row-major. + auto mY = tma_store_y.get_tma_tensor(make_shape( + static_cast(num_m_tiles) * Traits::TileM, + static_cast(intermediate_dim))); + auto cta_tma_y = tma_store_y.get_slice(Int<0>{}); + + int n_start = blockIdx.y; + int n_stride = gridDim.y; + bool store_in_flight = false; + constexpr int K_PIPE_MMAS = 1; + + for (int m = blockIdx.x; m < num_m_tiles; m += gridDim.x) { + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc); + + auto state_release = state; + int prologue_count = (num_k_tiles < K_PIPE_MMAS) ? num_k_tiles : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < num_k_tiles; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. + // TileM=64 (N-split): WG_w owns full M (=64), cols [w*WgTileN, ...). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: 2 row tiles per logical M tile. + m_tile_idx = 2 * m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M tile; WG_w handles + // its N-half. + m_tile_idx = m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gY = local_tile(mY, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_y, cta_tma_y.partition_S(sStore), + cta_tma_y.partition_D(gY)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} + +// ═══════════════════════════════════════════════════════════════════ +// mlp2_t_fwd — device orchestrator +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp2_t_fwd( + Mlp2TSmem& smem, + TmaLoadZ const& tma_load_z, + TmaLoadW const& tma_load_a, + TmaStoreY const& tma_store_y, + const int* expert_ids, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int total_k_cols, + int num_m_tiles) { + + using Pipeline = typename Traits::MainloopPipeline; + using PipeState = typename Traits::PipelineState; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + int num_n_tiles = intermediate_dim / Traits::TileN; + int num_k_tiles = hidden_dim / Traits::TileK; + + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + cute::prefetch_tma_descriptor(tma_load_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_y.get_tma_descriptor()); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + Pipeline pipe(smem.pipe_storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); + + __syncthreads(); + + PipeState state = is_producer ? + cutlass::make_producer_start_state() : PipeState{}; + PipeState cons_state; + + if (is_producer) { + mlp2_t_producer( + pipe, state, + smem, tma_load_z, tma_load_a, + expert_ids, + num_tokens, hidden_dim, intermediate_dim, total_k_cols, + num_m_tiles, num_n_tiles, num_k_tiles); + } + + if (is_consumer) { + mlp2_t_consumer( + pipe, cons_state, + smem, tma_store_y, + intermediate_dim, + num_m_tiles, num_n_tiles, num_k_tiles); + } + + __syncthreads(); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp3.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp3.cuh index df72f67d..247ef16d 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp3.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp3.cuh @@ -1,1784 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Phase 3: dA = dY^T @ Z (weight gradient, backward) -// ═══════════════════════════════════════════════════════════════════ -// -// TMA_REDUCE_ADD variant. Cooperative 2-WG consumer; the LARGER of -// (TileM, TileN) is the cooperation axis. Only two configs supported: -// (TileM=256, TileN=128) → M-split (Layout<_2,_1,_1>) -// (TileM=128, TileN=256) → N-split (Layout<_1,_2,_1>) ← default -// -// Per-WG output footprint = (WgTileM, WgTileN) = (128, 128) in both -// cases — each WG covers 128 rows × 128 cols per K-step via 2 wgmma -// m=64 issues, matching today's non-coop dual-WG per-WG work. -// -// Single fused DYT + Z TMA pipe (one mbarrier per stage, 2 TMA copies). -// Consumer computes a per-WG partial dA[m,n] in registers, casts to -// smem store_buf, then issues SM90_TMA_REDUCE_ADD into gmem — hardware -// atomically adds the partial. dA starts zero (torch::zeros / bwd host). -// -// CTA layout (384 threads = 12 warps): -// WG0 warp 0 : Producer (TMA loads for DYT / Z) -// WG0 warps 1-2 : Idle / reserved for fused-kernel communication -// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 -// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 -// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 -// -// Chunk-fixed walk: each CTA owns ONE (chunk, lane) cell. We hold the -// SHARED (cooperative) operand chunk-constant and walk the cooperation -// (split) axis, maximizing L2 reuse of the operand both WGs consume: -// kMSplit=true (M-split, Z shared): chunk = (e, n_tile), walks m_tile; holds Z(n,*) -// kMSplit=false (N-split, dY^T shared): chunk = (e, m_tile), walks n_tile; holds dY^T(m,*) -// ═══════════════════════════════════════════════════════════════════ +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -#include "models.cuh" -#include "tmem_load_op.cuh" - -#include -#include -#include -#include -#include - -// SM100 (Blackwell / UMMA) support for the Compute=100 consumer specialization -// (Mlp3ConsumerImpl<100>). Header-only and safe to include under an sm_90a -// build — the tcgen05 PTX they emit is guarded by CUTE_ARCH_TCGEN05_* and only -// the Compute=100 body (itself gated on __CUDA_ARCH__ >= 1000) instantiates it. -#include -#include - -#include -#include -#include // PipelineTmaUmmaAsync, PipelineUmmaAsync - -// SM100 paired-CTA MLP3 support. Header-only, with the same guarding -// rationale as the other SM100 includes above. -#include -#include -#include - -namespace liger { - -using namespace cute; - -template -struct Mlp3UmmaTiledMmaSelector { - using type = void; -}; - -template -struct Mlp3UmmaTiledMmaSelector { - using type = decltype(make_tiled_mma( - SM100_MMA_F16BF16_SS{})); -}; - -// ═══════════════════════════════════════════════════════════════════ -// Traits -// ═══════════════════════════════════════════════════════════════════ - -template < - typename Element_, - int TileM_ = 128, - int TileN_ = 256, - int TileK_ = 64, - int Stages_ = 4, - int EpiChunkN_ = 64 -> -struct Mlp3Traits { - using Element = Element_; - using ElementAccum = float; - - static constexpr int TileM = TileM_; - static constexpr int TileN = TileN_; - static constexpr int TileK = TileK_; - static constexpr int Stages = Stages_; - static constexpr int EpiChunkN = EpiChunkN_; - - // Cooperative 2-WG layout. Only two configs supported: - // (TileM=256, TileN=128) → kMSplit=true, WgLayout<_2,_1,_1> - // (TileM=128, TileN=256) → kMSplit=false, WgLayout<_1,_2,_1> - // Per-WG output footprint = (WgTileM=128, WgTileN=128) in both. Each - // WG covers 2 wgmma m=64 atoms per K-step (atom_count_M = 2). - static_assert((TileM == 256 && TileN == 128) || - (TileM == 128 && TileN == 256), - "Mlp3Traits supports only (TileM, TileN) = (256, 128) M-split " - "or (128, 256) N-split"); - - static constexpr bool kMSplit = (TileM > TileN); - static constexpr int AtomTileM = 64; - static constexpr int WgTileM = 128; - static constexpr int WgTileN = 128; - - static_assert(WgTileN % EpiChunkN == 0, - "EpiChunkN must divide WgTileN (=128)"); - static constexpr int NumEpiRounds = WgTileN / EpiChunkN; - - static constexpr int AtomN = WgTileN; - using GmmaAtom = typename GmmaSelectorMNMN::Atom; - using WgLayout = cute::conditional_t< - kMSplit, - Layout>, - Layout>>; - using TiledMma = TiledMMA, WgLayout>; - - using SmemLayoutAtom = GMMA::Layout_MN_SW128_Atom; - - using SmemLayoutDYT_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, - Step<_2, _1>{})); - using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, - Step<_2, _1>{})); - using SmemLayoutDYT = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, - Step<_2, _1, _3>{})); - using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, - Step<_2, _1, _3>{})); - - using MainloopPipeline = cutlass::PipelineTmaAsync; - using PipelineState = cutlass::PipelineState; - - // ── SM100 (Compute=100) pipelines — single CTA, no cluster. ── - // MainloopPipelineUmma: UMMA-aware TMA pipeline whose consumer_release - // issues the UMMA-gated smem-buffer arrival. The Blackwell launcher must - // build it with num_consumers=1 (one umma_arrive per stage releases the - // buffer). Params/PipelineState/SharedStorage alias the Hopper - // PipelineTmaAsync's, so producer + Compute=100 consumer drive it unchanged. - using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< - Stages, cute::Shape, - cute::Shape>; - // Double-buffered TMEM accumulators: MMA for the next (cell, walk) can fill - // the alternate stage while warps 4..11 drain/reduce-add the previous stage. - static constexpr int AccStages = 2; - static_assert(AccStages * TileN <= 512, - "SM100 MLP3 accumulator stages must fit in 512 TMEM columns"); - using AccumulatorPipeline = cutlass::PipelineUmmaAsync; - - // SM100 UMMA TiledMMA — one 1SM tcgen05 atom spanning (TileM, TileN) with - // the accumulator in TMEM. mlp3 is the FIRST ported kernel with BOTH - // operands MN-major: A = dY^T (M=H contiguous) and B = Z (N=I contiguous), - // exactly matching SmemLayoutDYT / SmemLayoutZ (both Layout_MN_SW128_Atom + - // Step<_2,_1,_3>). make_umma_desc reads the MN-contiguous stride - // from each smem layout, so no transpose is needed (cf. mlp2_t's ). - using TiledMmaUmma = typename Mlp3UmmaTiledMmaSelector< - Element, ElementAccum, TileM, TileN>::type; - - static constexpr int TmaTransBytesDYT = - static_cast(size(SmemLayoutDYT_1{}) * sizeof(Element)); - static constexpr int TmaTransBytesZ = - static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); - // Fused pipeline: DYT + Z share one mbarrier per stage. Both WGs - // consume the same Z via cooperative TiledMma slicing. - static constexpr int TmaTransBytes = TmaTransBytesDYT + TmaTransBytesZ; - - static constexpr int WarpSize = 32; - static constexpr int WarpGroupSize = 128; - static constexpr int NumConsumers = 2; - static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; // 256 - static constexpr int NumThreads = 384; - - // Per-WG store buffer slot: (WgTileM, EpiChunkN). One slot per WG; it - // holds kAtomsPerWg = WgTileM/AtomTileM contiguous atom-rows after the - // acc→smem fold. - static constexpr int kAtomsPerWg = WgTileM / AtomTileM; // 2 - using SmemLayoutStoreSlot = Layout< - Shape, Int>, - Stride, _1>>; - // TMA-reduce descriptor box = ONE wgmma atom-row (AtomTileM × EpiChunkN). - // At TileM=256 each WG owns 2 *interleaved* atoms (rows {0-63,128-191} / - // {64-127,192-255}) that map to non-contiguous gmem rows, so a WG issues - // kAtomsPerWg separate TMA stores per epi round (mirrors mlp1_fused.cuh, - // which uses an AtomTileM box but only ever has 1 atom/WG). - using SmemLayoutStore = Layout< - Shape, Int>, - Stride, _1>>; -}; - -template -using Mlp3MainloopPipelineFor = cute::conditional_t< - Compute == 100, - typename Traits::MainloopPipelineUmma, - typename Traits::MainloopPipeline>; - -// ═══════════════════════════════════════════════════════════════════ -// Shared Memory -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp3Smem { - using Element = typename Traits::Element; - - static constexpr int smem_DYT_size = cosize_v; - static constexpr int smem_Z_size = cosize_v; - static constexpr int smem_store_size = cosize_v; - - alignas(128) Element smem_DYT[smem_DYT_size]; - alignas(128) Element smem_Z[smem_Z_size]; - // Per-WG store buffer: WG1 uses [0..S-1], WG2 uses [S..2S-1]. - alignas(128) Element store_buf[2 * smem_store_size]; - - typename Traits::MainloopPipeline::SharedStorage pipe_storage; - - // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM - // base address, plus the accumulator pipeline's barrier storage (UMMA→ - // epilogue handoff). The Hopper (Compute=90) path never touches these; they - // keep the consumer signature identical across Compute values. - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; - - CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } - CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// Fused MLP3 shared memory (data only — pipe storage lives in the -// outer MlpFusedBwdSmem union; otherwise identical to Mlp3Smem). -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp3FusedSmem { - using Element = typename Traits::Element; - - static constexpr int smem_DYT_size = cosize_v; - static constexpr int smem_Z_size = cosize_v; - static constexpr int smem_store_size = cosize_v; - - alignas(128) Element smem_DYT[smem_DYT_size]; - alignas(128) Element smem_Z[smem_Z_size]; - alignas(128) Element store_buf[2 * smem_store_size]; - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; - - CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } - CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } -}; - - -// ═══════════════════════════════════════════════════════════════════ -// Pipe construction helper -// ═══════════════════════════════════════════════════════════════════ -// -// Mirrors mlp2/mlp5's make_pipe pattern: hides the role/transaction- -// byte setup so callers (mlp3_fwd, or any fused parent kernel) get a -// consistent Pipeline. The plain SM90 helper uses warp 0 as TMA producer and -// warps 4-11 as WGMMA consumers; the SM100 helper additionally enrolls warp 3 -// as the UMMA consumer of the TMA pipeline. - -template -__device__ __forceinline__ typename Traits::MainloopPipeline -mlp3_make_pipe(typename Traits::MainloopPipeline::SharedStorage& storage) { - using Pipeline = typename Traits::MainloopPipeline; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 4 && warp_id <= 11); - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = Traits::ConsumerThreads; - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// SM100 (Compute=100) mainloop pipe. Mirrors mlp3_make_pipe but builds the -// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues -// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer -// is released by a single umma_arrive per stage (the MMA warp), not per -// consumer thread. The ctor arg pattern is identical to the Hopper pipe, so the -// arch-agnostic mlp3_producer drives it unchanged. -template -__device__ __forceinline__ typename Traits::MainloopPipelineUmma -mlp3_make_pipe_umma(typename Traits::MainloopPipelineUmma::SharedStorage& storage) { - using Pipeline = typename Traits::MainloopPipelineUmma; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// ═══════════════════════════════════════════════════════════════════ -// Chunk-fixed cooperative producer/consumer -// ═══════════════════════════════════════════════════════════════════ -// -// Used by the standalone mlp3 kernel (mlp3.cu) and mlp_bwd Phase 2. -// Each CTA owns ONE (chunk, lane) cell. We hold the SHARED operand and -// walk the cooperation (split) axis: -// kMSplit=true (M-split): chunk = (e, n_tile), lane partitions -// num_m_tiles, walks m_tile. Z(n_tile, *) — the shared -// N-side operand — is the L2-held chunk-constant. -// kMSplit=false (N-split): chunk = (e, m_tile), lane partitions -// num_n_tiles, walks n_tile. dY^T(m_tile, *) — the -// shared M-side operand — is L2-held. -// -// Cell decode (chunk-major: cell_idx = chunk_idx × outer_split + lane): -// total_cells = total_chunks × outer_split -// cell → chunk_idx = cell_idx / outer_split -// lane = cell_idx % outer_split - -template -struct Mlp3ProducerImpl; - -template <> -struct Mlp3ProducerImpl<90> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaLoadDYT const& tma_load_dyt, - TmaLoadZ const& tma_load_z, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int hidden_dim, - int intermediate_dim, - int num_tokens, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - // Staging-ring depth in K-blocks. > 0 means batch_kb_*/the per-expert - // k-ranges are UNWRAPPED and the gmem K-tile coordinate wraps by % ring_kb - // (grouped grouped Phase-2 may straddle the ring). 0 = no wrap. - int ring_kb = 0) { - - auto sDYT = make_tensor(make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - - auto mDYT = tma_load_dyt.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(num_tokens))); - auto mZT = tma_load_z.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(num_tokens))); - - auto cta_tma_dyt = tma_load_dyt.get_slice(Int<0>{}); - auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); - - auto tDYTsDYT = cta_tma_dyt.partition_D(sDYT); - auto tZsZ = cta_tma_z.partition_D(sZ); - - int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); - // K-split: each (chunk, walk-lane) is replicated k_split times; the - // k_slice handles a 1/k_split sub-range of the expert's K-loop. Partial - // reductions accumulate into dA via the TMA_REDUCE_ADD epilogue. - int total_cells = total_chunks * outer_split * k_split; - - for (int cell_idx = cell_start; - cell_idx < total_cells; - cell_idx += cell_stride) { - - int k_slice = cell_idx % k_split; - int cell_om = cell_idx / k_split; - int chunk_idx = cell_om / outer_split; - int lane = cell_om - chunk_idx * outer_split; - - int e, m_chunk, n_chunk, walk_begin, walk_end; - if constexpr (Traits::kMSplit) { - // M-split: hold Z(n) shared. chunk = (e, n_tile), walks m_tile. - e = chunk_idx / num_n_tiles; - n_chunk = chunk_idx - e * num_n_tiles; - m_chunk = 0; - int m_per = num_m_tiles / outer_split; - walk_begin = lane * m_per; - walk_end = walk_begin + m_per; - } else { - // N-split: hold dY^T(m) shared. chunk = (e, m_tile), walks n_tile. - e = chunk_idx / num_m_tiles; - m_chunk = chunk_idx - e * num_m_tiles; - n_chunk = 0; - int n_per = num_n_tiles / outer_split; - walk_begin = lane * n_per; - walk_end = walk_begin + n_per; - } - - // Intersect this expert's global K range with the current - // batch's K-window. Global k_starts/ends were filled once per - // pass; per-batch we just clip. - int kb_lo = max(expert_k_starts[e], batch_kb_start); - int kb_hi = min(expert_k_ends[e], batch_kb_end); - if (kb_hi <= kb_lo) continue; - - // K-split: clip to this cell's k_slice sub-range of [kb_lo, kb_hi). - { - int k_total = kb_hi - kb_lo; - int k_per = (k_total + k_split - 1) / k_split; - int s_lo = kb_lo + k_slice * k_per; - int s_hi = min(s_lo + k_per, kb_hi); - if (s_hi <= s_lo) continue; // empty slice (e.g. straggler) - kb_lo = s_lo; - kb_hi = s_hi; - } - - for (int w = walk_begin; w < walk_end; ++w) { - int m_tile = Traits::kMSplit ? w : m_chunk; - int n_tile = Traits::kMSplit ? n_chunk : w; - auto gDYT = local_tile(mDYT, - make_tile(Int{}, Int{}), - make_coord(m_tile, _)); - auto gZ = local_tile(mZT, - make_tile(Int{}, Int{}), - make_coord(n_tile, _)); - auto tDYTgDYT = cta_tma_dyt.partition_S(gDYT); - auto tZgZ = cta_tma_z.partition_S(gZ); - - for (int kb = kb_lo; kb < kb_hi; ++kb) { - int kb_g = (ring_kb > 0) ? (kb % ring_kb) : kb; // physical ring slot - pipe.producer_acquire(state); - if (threadIdx.x == 0) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_dyt.with(*bar, 0), - tDYTgDYT(_, _, _, kb_g), tDYTsDYT(_, _, _, state.index())); - copy(tma_load_z.with(*bar, 0), - tZgZ(_, _, _, kb_g), tZsZ(_, _, _, state.index())); - } - ++state; - } - } - } -} -}; // Mlp3ProducerImpl<90> - -template -__device__ __forceinline__ void mlp3_producer( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaLoadDYT const& tma_load_dyt, - TmaLoadZ const& tma_load_z, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int hidden_dim, - int intermediate_dim, - int num_tokens, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - int ring_kb = 0) { - Mlp3ProducerImpl::template run( - pipe, state, smem, tma_load_dyt, tma_load_z, - expert_k_starts, expert_k_ends, num_experts, - hidden_dim, intermediate_dim, num_tokens, - num_m_tiles, num_n_tiles, outer_split, - cell_start, cell_stride, - batch_kb_start, batch_kb_end, k_split, ring_kb); -} - -// ═══════════════════════════════════════════════════════════════════ -// Consumer — arch-split into Compute=90 (WGMMA) / Compute=100 (UMMA) -// ═══════════════════════════════════════════════════════════════════ -// Compute=90 → Hopper / WGMMA. Both WGs cooperate in ONE TiledMMA that -// partitions (TileM,TileN); each thread holds a register acc -// slice, scatters into store_buf, and the WG leader issues -// SM90_TMA_REDUCE_ADD. -// Compute=100 → Blackwell / UMMA. One 1SM tcgen05 atom (warp 3) fills ONE -// TMEM accumulator over the whole (TileM,TileN); both WGs pull -// their N-half from TMEM into the SAME store_buf slots, then the -// IDENTICAL REDUCE_ADD leader block stores them. - -template -struct Mlp3ConsumerImpl; - -// ─────────────────────────────────────────────────────────────────── -// Compute=90 — Hopper WGMMA. Verbatim from the original mlp3_consumer. -// ─────────────────────────────────────────────────────────────────── -template <> -struct Mlp3ConsumerImpl<90> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaReduceAddDA const& tma_reduce_da, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int hidden_dim, - int intermediate_dim, - int total_n_rows, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split) { - - (void)Compute; - using Element = typename Traits::Element; - typename Traits::TiledMma tiled_mma; - - const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 - const int my_barrier_id = 1 + my_wg; // 1 or 2 - const int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 - const int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 - auto thr_mma = tiled_mma.get_slice(tid_in_mma); - - auto sDYT = make_tensor(make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); - auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - - auto tCsDYT = thr_mma.partition_A(sDYT); - auto tCsZ = thr_mma.partition_B(sZ); - - auto acc = partition_fragment_C(tiled_mma, - Shape, Int>{}); - - // Identity coords for per-thread acc → store_buf mapping. - auto cC = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCcC = thr_mma.partition_C(cC); - - constexpr int store_slot_elems = Traits::WgTileM * Traits::EpiChunkN; - Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - const bool is_my_wg_leader = (tid_in_wg == 0); - - auto mdA = [&]() { - if constexpr (Expert3D) { - return tma_reduce_da.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(intermediate_dim), - static_cast(num_experts))); - } else { - return tma_reduce_da.get_tma_tensor(make_shape( - static_cast(total_n_rows), - static_cast(intermediate_dim))); - } - }(); - auto cta_tma_da = tma_reduce_da.get_slice(Int<0>{}); - - int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); - int total_cells = total_chunks * outer_split * k_split; - - constexpr int K_PIPE_MMAS = 1; - - // Carries across walk iters AND across cells for cross-tile K-loop / - // store overlap. Drained once at the end. - bool store_in_flight = false; - - for (int cell_idx = cell_start; - cell_idx < total_cells; - cell_idx += cell_stride) { - - int k_slice = cell_idx % k_split; - int cell_om = cell_idx / k_split; - int chunk_idx = cell_om / outer_split; - int lane = cell_om - chunk_idx * outer_split; - - int e, m_chunk, n_chunk, walk_begin, walk_end; - if constexpr (Traits::kMSplit) { - // M-split: hold Z(n) shared. chunk = (e, n_tile), walks m_tile. - e = chunk_idx / num_n_tiles; - n_chunk = chunk_idx - e * num_n_tiles; - m_chunk = 0; - int m_per = num_m_tiles / outer_split; - walk_begin = lane * m_per; - walk_end = walk_begin + m_per; - } else { - // N-split: hold dY^T(m) shared. chunk = (e, m_tile), walks n_tile. - e = chunk_idx / num_m_tiles; - m_chunk = chunk_idx - e * num_m_tiles; - n_chunk = 0; - int n_per = num_n_tiles / outer_split; - walk_begin = lane * n_per; - walk_end = walk_begin + n_per; - } - - // Intersect this expert's global K range with the current - // batch's K-window. - int kb_lo = max(expert_k_starts[e], batch_kb_start); - int kb_hi = min(expert_k_ends[e], batch_kb_end); - if (kb_hi <= kb_lo) continue; - - // K-split: clip to this cell's k_slice sub-range (must match producer). - { - int k_total = kb_hi - kb_lo; - int k_per = (k_total + k_split - 1) / k_split; - int s_lo = kb_lo + k_slice * k_per; - int s_hi = min(s_lo + k_per, kb_hi); - if (s_hi <= s_lo) continue; // empty slice - kb_lo = s_lo; - kb_hi = s_hi; - } - - int k_count = kb_hi - kb_lo; - - for (int w = walk_begin; w < walk_end; ++w) { - int m_tile = Traits::kMSplit ? w : m_chunk; - int n_tile = Traits::kMSplit ? n_chunk : w; - - clear(acc); - auto state_release = state; - int prologue_count = (k_count < K_PIPE_MMAS) ? k_count : K_PIPE_MMAS; - - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsDYT(_, _, _, state.index()), - tCsZ(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - ++state; - } - - for (int k = prologue_count; k < k_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsDYT(_, _, _, state.index()), - tCsZ(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - - warpgroup_wait(); - warpgroup_fence_operand(acc); - pipe.consumer_release(state_release); - - ++state; - ++state_release; - } - - warpgroup_wait<0>(); - warpgroup_fence_operand(acc); - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_release(state_release); - ++state_release; - } - - // ── Epilogue ────────────────────────────────────── - // Each WG writes its (WgTileM × EpiChunkN) slot to store_buf - // and issues TMA stores. NumEpiRounds rounds per WG. - // kMSplit=true : M-split. WG_w owns WgTileM contiguous - // rows of TileM. With WgLayout_M=2 and - // atom_count_M = WgTileM/AtomTileM, the - // natural CUTE m-atom positions interleave - // between WGs — formula below maps m_loc - // back to a contiguous local row. - // kMSplit=false : N-split. WG_w owns full TileM, cols - // [w*WgTileN, (w+1)*WgTileN). - CUTE_UNROLL - for (int r = 0; r < Traits::NumEpiRounds; ++r) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - int chunk_start = r * Traits::EpiChunkN; - CUTE_UNROLL - for (int i = 0; i < size(acc); ++i) { - auto coord = tCcC(i); - int m_loc = get<0>(coord); - int n_loc = get<1>(coord); - int m_local, n_local; - if constexpr (Traits::kMSplit) { - // CUTE atom-iter interleaves WGs in M direction - // when atom_count_M > 1 (TileM=256). Unified - // formula folds 64-row WG strips into a - // contiguous WgTileM-row local range: - // m_atom_global = m_loc / 64 (atom row idx) - // m_atom_in_wg = m_atom_global / 2 - // m_local = m_atom_in_wg*64 + (m_loc % 64) - // = (m_loc / 128)*64 + (m_loc % 64) - // For tile_count_M=1 (TileM=128) the / 128 term - // is 0 so m_local = m_loc % 64. - m_local = (m_loc / 128) * 64 + (m_loc % 64); - n_local = n_loc; - } else { - m_local = m_loc; - n_local = n_loc - my_wg * Traits::WgTileN; - } - if (n_local >= chunk_start && - n_local < chunk_start + Traits::EpiChunkN) { - int chunk_n = n_local - chunk_start; - sStore(m_local, chunk_n) = static_cast(acc(i)); - } - } - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - if (is_my_wg_leader) { - cute::tma_store_fence(); - int da_m = Expert3D ? m_tile : e * num_m_tiles + m_tile; - // Store box = ONE atom-row (AtomTileM). Each WG issues - // kAtomsPerWg stores, one per atom it owns. The store_buf - // holds the WG's atoms in contiguous local rows - // [a*AtomTileM, +AtomTileM); their gmem atom-rows differ: - // M-split: the 2 atoms interleave → 4*da_m + my_wg + 2a - // (gmem M-grid is TileM/AtomTileM rows per tile). - // N-split: WG owns the full TileM (2 contiguous atoms) and - // its own N-half → 2*da_m + a, n offset by my_wg. - int n_tile_idx; - if constexpr (Traits::kMSplit) { - n_tile_idx = n_tile * Traits::NumEpiRounds + r; - } else { - n_tile_idx = n_tile * (Traits::TileN / Traits::EpiChunkN) - + my_wg * Traits::NumEpiRounds + r; - } - constexpr int kRowAtoms = Traits::TileM / Traits::AtomTileM; // 4 (M-split) / 2 (N-split) - CUTE_UNROLL - for (int a = 0; a < Traits::kAtomsPerWg; ++a) { - int m_atom_row = Traits::kMSplit - ? (kRowAtoms * da_m + my_wg + Traits::kAtomsPerWg * a) - : (kRowAtoms * da_m + a); - auto sStore_a = local_tile(sStore, - make_tile(Int{}, Int{}), - make_coord(a, 0)); - auto gdA = [&]() { - if constexpr (Expert3D) { - return local_tile(mdA, - make_tile(Int{}, Int{}), - make_coord(m_atom_row, n_tile_idx, e)); - } else { - return local_tile(mdA, - make_tile(Int{}, Int{}), - make_coord(m_atom_row, n_tile_idx)); - } - }(); - copy(tma_reduce_da, cta_tma_da.partition_S(sStore_a), - cta_tma_da.partition_D(gdA)); - } - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - } - if (store_in_flight) - cute::tma_store_wait<0>(); -} -}; // Mlp3ConsumerImpl<90> - -// ─────────────────────────────────────────────────────────────────── -// Forwarder — Compute defaults to 90 so existing call sites -// (mlp3_consumer(...)) are unchanged; pass mlp3_consumer -// for the Blackwell path. -// ─────────────────────────────────────────────────────────────────── -template -__device__ __forceinline__ void mlp3_consumer( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaReduceAddDA const& tma_reduce_da, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int hidden_dim, - int intermediate_dim, - int total_n_rows, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, - uint32_t* pair_init_phase = nullptr) { - if constexpr (Compute == 100) { - Mlp3ConsumerImpl::template run( - pipe, state, smem, tma_reduce_da, - expert_k_starts, expert_k_ends, num_experts, - hidden_dim, intermediate_dim, total_n_rows, - num_m_tiles, num_n_tiles, outer_split, - cell_start, cell_stride, - batch_kb_start, batch_kb_end, k_split, - pair_init_barrier, pair_init_phase); - } else { - Mlp3ConsumerImpl::template run( - pipe, state, smem, tma_reduce_da, - expert_k_starts, expert_k_ends, num_experts, - hidden_dim, intermediate_dim, total_n_rows, - num_m_tiles, num_n_tiles, outer_split, - cell_start, cell_stride, - batch_kb_start, batch_kb_end, k_split); - } -} - -// ═══════════════════════════════════════════════════════════════════ -// Host launcher (raw-pointer; defined in mlp3.cu) -// ═══════════════════════════════════════════════════════════════════ - -// dA = dY^T @ Z, accumulated atomically into the [E, H, I] output via -// TMA_REDUCE_ADD. dA MUST be zero-initialized by the caller. -// -// dY : [num_tokens, hidden_dim] -// Z : [num_tokens, intermediate_dim] -// dA : [num_experts, hidden_dim, intermediate_dim] -// expert_for_k_block: [num_tokens / TileK] int32, TileK=64 -void mlp3_fwd_bf16_launch( - const cutlass::bfloat16_t* dY, - const cutlass::bfloat16_t* Z, - cutlass::bfloat16_t* dA, - const int* expert_for_k_block, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int num_experts, - int device, - cudaStream_t stream); - -// ═══════════════════════════════════════════════════════════════════ -// SM100 2-SM (paired-CTA) MLP3 — basic ClusterM=2 path -// ═══════════════════════════════════════════════════════════════════ -// -// Joins two adjacent CTAs into one cta_group::2 UMMA atom covering a -// 256x256x64 joined tile, doubling the M extent a single CTA-pair can -// produce per wave versus the 1SM path above. Callers must launch with -// cudaLaunchKernelEx, clusterDim=(2,1,1), an even gridDim.x, and -// pair-aware SM100 2SM TMA load descriptors for dY^T / Z. The dA -// descriptor remains an ordinary (non-paired) TMA_REDUCE_ADD. -// -// TileK, Stages, and EpiChunkN stay trait parameters — the validated -// production candidate is TileK=64/Stages=5/EpiChunkN=64 (see -// 2sm.md / 2sm_stage.md), but that tuple is a caller-side choice, not -// hardcoded here. ClusterM is fixed at 2 (one paired-CTA "pair"): the -// ClusterM=4 (two-pair, N-split-across-pairs) and shared-operand TMA -// multicast variants explored alongside this trait did not clear the -// required performance bar (see the mlp3_multicast bench results) and -// are intentionally not ported, nor are the CompactEpilogue and -// EarlyTmemRelease epilogue experiments (see mlp3_s6_bounce / -// mlp3_s5_early_release) — both were rejected and are omitted here. -// -// Host-launcher contract (no launcher lives in this header — these are -// hard requirements on any .cu translation unit that selects Compute=100, -// e.g. an mlp_bwd.cuh fused entry point): -// 1. Grid: gridDim.x must be even (mlp3_fwd __traps if it is not); -// launch must use cudaLaunchKernelEx with cudaLaunchConfig_t:: -// cudaLaunchAttributeClusterDimension = {ClusterM, 1, 1}. -// 2. The kernel function must first be registered via -// cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) -// before the cudaLaunchKernelEx call — the driver rejects a 2-CTA -// cluster launch of a kernel that exceeds the "portable" static -// shared-memory budget without this opt-in. -// 3. The dY^T / Z TMA operands passed in as TmaLoadDYT / TmaLoadZ MUST -// be constructed with cute::make_tma_copy_A_sm100 / -// cute::make_tma_copy_B_sm100 (NOT the ordinary 1SM make_tma_copy) -// — only the _sm100 constructors produce the cta_group::2, -// pair-aware copy atoms that create_tma_multicast_mask / -// tma_partition below assume. The dA reduce-add operand stays an -// ordinary (non-paired) TMA_REDUCE_ADD, per above. -// 4. Dynamic shared memory for Mlp3Smem2Sm / Mlp3FusedSmem2Sm -// must be sized and opted in at runtime by the launcher — -// query cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemory -// PerBlockOptin, device), then cudaFuncSetAttribute(kernel, -// cudaFuncAttributeMaxDynamicSharedMemorySize, requested) with -// requested <= optin. This header intentionally has no compile- -// time static_assert on a fixed smem byte ceiling (only the TMEM -// column assert on AccStages*TileN below) — the ceiling is a -// runtime, device-dependent quantity and must not be hardcoded. -template < - typename Element_, - int TileM_ = 256, - int TileN_ = 256, - int TileK_ = 64, - int Stages_ = 3, - int EpiChunkN_ = 64, - int AccStages_ = 2, - int ClusterM_ = 2> -struct Mlp3Traits2Sm { - using Element = Element_; - using ElementAccum = float; - - static constexpr int TileM = TileM_; - static constexpr int TileN = TileN_; - static constexpr int TileK = TileK_; - static constexpr int Stages = Stages_; - static constexpr int EpiChunkN = EpiChunkN_; - static constexpr int AccStages = AccStages_; - static constexpr int ClusterM = ClusterM_; - static constexpr int NumPairs = ClusterM / 2; - static constexpr int CtaTileM = TileM / 2; - static constexpr int AtomTileM = 64; - static constexpr int WgTileN = TileN / 2; - static constexpr int NumEpiRounds = WgTileN / EpiChunkN; - static constexpr int kAtomsPerCta = CtaTileM / AtomTileM; - - static_assert(TileM == 256, - "MLP3 2SM requires joined TileM=256; joined TileM=128 is invalid"); - static_assert(TileN == 128 || TileN == 256, - "MLP3 2SM supports TileN=128 or TileN=256"); - static_assert(TileK == 32 || TileK == 64, - "MLP3 2SM pilot supports TileK=32 or 64"); - static_assert(Stages >= 2 && Stages <= 12, - "MLP3 2SM supports two to twelve mainloop stages"); - static_assert(CtaTileM == 128, "each peer CTA must own 128 M rows"); - static_assert(WgTileN % EpiChunkN == 0, - "EpiChunkN must divide each warpgroup's N half"); - static_assert(kAtomsPerCta == 2, - "each CTA must reduce-add two 64-row atoms"); - static_assert(AccStages >= 2, "MLP3 2SM requires accumulator double buffering"); - static_assert(AccStages * TileN <= 512, - "MLP3 2SM accumulator stages must fit in 512 TMEM columns"); - static_assert(ClusterM == 2, - "MLP3 2SM only ports the validated single-pair path; ClusterM=4 " - "(two-pair N-split, evaluated together with shared-operand " - "multicast) was rejected -- see the mlp3_multicast bench results"); - - using TileShape = Shape, Int, Int>; - using ClusterShape = Shape, _1, _1>; - using AtomThrShape = Shape<_2, _1, _1>; - - using TiledMma2Sm = decltype(make_tiled_mma( - SM100_MMA_F16BF16_2x1SM_SS< - Element, Element, ElementAccum, TileM, TileN, - UMMA::Major::MN, UMMA::Major::MN>{})); - static_assert(size(typename TiledMma2Sm::AtomThrID{}) == 2, - "MLP3 2SM requires a two-CTA MMA atom"); - - using ClusterLayoutVMNK = decltype(tiled_divide( - make_layout(ClusterShape{}), - make_tile(typename TiledMma2Sm::AtomThrID{}))); - - using MmaShapeA_MK = decltype(partition_shape_A( - TiledMma2Sm{}, make_shape(Int{}, Int{}))); - using MmaShapeB_NK = decltype(partition_shape_B( - TiledMma2Sm{}, make_shape(Int{}, Int{}))); - - using SmemLayoutAtom = UMMA::Layout_MN_SW128_Atom; - using SmemLayoutDYT = decltype(UMMA::tile_to_mma_shape( - SmemLayoutAtom{}, append(MmaShapeA_MK{}, Int{}), - Step<_2, _1, _3>{})); - using SmemLayoutZ = decltype(UMMA::tile_to_mma_shape( - SmemLayoutAtom{}, append(MmaShapeB_NK{}, Int{}), - Step<_2, _1, _3>{})); - using SmemLayoutDYT_1 = - decltype(SmemLayoutDYT{}(_, _, _, Int<0>{})); - using SmemLayoutZ_1 = - decltype(SmemLayoutZ{}(_, _, _, Int<0>{})); - - static constexpr int TmaTransBytesDYT = - static_cast(cosize_v * sizeof(Element)); - static constexpr int TmaTransBytesZ = - static_cast(cosize_v * sizeof(Element)); - static constexpr int TmaTransBytes = - 2 * (TmaTransBytesDYT + TmaTransBytesZ); - - using MainloopPipelineUmma2Sm = cutlass::PipelineTmaUmmaAsync< - Stages, ClusterShape, AtomThrShape>; - using MainloopPipeline = MainloopPipelineUmma2Sm; - using MainloopPipelineUmma = MainloopPipelineUmma2Sm; - using PipelineState = typename MainloopPipelineUmma2Sm::PipelineState; - using AccumulatorPipeline2Sm = cutlass::PipelineUmmaAsync< - AccStages, AtomThrShape>; - - static constexpr int WarpSize = 32; - static constexpr int WarpGroupSize = 128; - static constexpr int NumConsumers = 2; - static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; - static constexpr int NumThreads = 384; - - // One (CtaTileM, EpiChunkN) store slot per epilogue warpgroup; WG0 - // uses [0..S-1], WG1 uses [S..2S-1] (S = smem_store_size, see - // Mlp3Smem2Sm below). Mirrors the 1SM Mlp3Traits::SmemLayoutStoreSlot, - // sized CtaTileM (=128) instead of WgTileM. - using SmemLayoutStoreSlot = Layout< - Shape, Int>, - Stride, _1>>; - using SmemLayoutStore = Layout< - Shape, Int>, - Stride, _1>>; -}; - -// ═══════════════════════════════════════════════════════════════════ -// SM100 2-SM Shared Memory -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp3Smem2Sm { - using Element = typename Traits::Element; - - static constexpr int smem_DYT_size = - cosize_v; - static constexpr int smem_Z_size = - cosize_v; - static constexpr int smem_store_size = - cosize_v; - - alignas(128) Element smem_DYT[smem_DYT_size]; - alignas(128) Element smem_Z[smem_Z_size]; - alignas(128) Element store_buf[2 * smem_store_size]; - alignas(16) typename Traits::MainloopPipelineUmma2Sm::SharedStorage pipe_storage; - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline2Sm::SharedStorage acc_pipe; - - CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } - CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } -}; - -// Fused MLP3 2SM shared memory (data only — the mainloop pipe's storage -// lives in the outer fused-kernel smem union; otherwise identical to -// Mlp3Smem2Sm). Mirrors the 1SM Mlp3FusedSmem / Mlp3Smem split. -template -struct Mlp3FusedSmem2Sm { - using Element = typename Traits::Element; - - static constexpr int smem_DYT_size = - cosize_v; - static constexpr int smem_Z_size = - cosize_v; - static constexpr int smem_store_size = - cosize_v; - - alignas(128) Element smem_DYT[smem_DYT_size]; - alignas(128) Element smem_Z[smem_Z_size]; - alignas(128) Element store_buf[2 * smem_store_size]; - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline2Sm::SharedStorage acc_pipe; - - CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } - CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// SM100 2-SM pipe construction helper -// ═══════════════════════════════════════════════════════════════════ -// -// Mirrors mlp3_make_pipe's role/transaction-byte setup, adjusted for the -// paired-CTA cluster: warp 0 is the TMA producer (leader CTA elects the -// arrival), warps 3-11 are UMMA/epilogue consumers. -template -__device__ __forceinline__ typename Traits::MainloopPipelineUmma2Sm -mlp3_make_pipe_umma_2sm( - typename Traits::MainloopPipelineUmma2Sm::SharedStorage& storage) { - using Pipeline = typename Traits::MainloopPipelineUmma2Sm; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = warp_id == 0; - bool is_consumer = warp_id >= 3 && warp_id <= 11; - - typename Pipeline::Params params; - params.transaction_bytes = Traits::TmaTransBytes; - params.num_producers = 1; - params.num_consumers = 1; - params.initializing_warp = 0; - if (is_producer) { - params.role = Category::Producer; - params.is_leader = - threadIdx.x == 0 && - cute::block_rank_in_cluster() % 2 == 0; - } else if (is_consumer) { - params.role = Category::Consumer; - } else { - params.role = Category::NonParticipant; - } - return Pipeline( - storage, params, typename Traits::ClusterShape{}, - cute::true_type{}, cute::true_type{}); -} - -// ═══════════════════════════════════════════════════════════════════ -// SM100 2-SM producer / consumer -// ═══════════════════════════════════════════════════════════════════ -// -// Chunk-fixed walk, same principle as mlp3_producer/mlp3_consumer above: -// each CTA-pair owns ONE (chunk, lane) cell and walks m_tile × kb -// internally. Within a pair, the leader CTA (pair_rank==0) drives the -// paired UMMA and owns the accumulator producer role; both CTAs run the -// epilogue over their own CtaTileM=128 half of the joined 256-row tile. -template <> -struct Mlp3ProducerImpl<100> { -template < - typename Traits, - typename Pipeline, - typename SmemType, - typename TmaLoadDYT, - typename TmaLoadZ> -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaLoadDYT const& tma_load_dyt, - TmaLoadZ const& tma_load_z, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int hidden_dim, - int intermediate_dim, - int num_tokens, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - int ring_kb = 0) { - // Compute!=100 must stay 1SM: this body issues SM100-only paired-CTA - // TMA ops (pair-aware tma_partition / create_tma_multicast_mask on a - // TiledMma2Sm/MainloopPipelineUmma2Sm operand) and must never be live - // code for an sm_90a compilation pass of this same translation unit. -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - auto sDYT = make_tensor( - make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); - auto sZ = make_tensor( - make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - - auto mDYT = tma_load_dyt.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(num_tokens))); - auto mZT = tma_load_z.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(num_tokens))); - - typename Traits::TiledMma2Sm tiled_mma; - int cta_rank = static_cast(cute::block_rank_in_cluster()); - int pair_rank = cta_rank % 2; - auto cta_mma = tiled_mma.get_slice(pair_rank); - auto pair_layout_vmnk = tiled_divide( - make_layout(typename Traits::ClusterShape{}), - make_tile(typename Traits::TiledMma2Sm::AtomThrID{})); - auto pair_coord_vmnk = - pair_layout_vmnk.get_flat_coord(pair_rank); - uint16_t mcast_mask_dyt = - create_tma_multicast_mask<2>(pair_layout_vmnk, pair_coord_vmnk); - uint16_t mcast_mask_z = - create_tma_multicast_mask<1>(pair_layout_vmnk, pair_coord_vmnk); - - int total_chunks = num_experts * num_n_tiles; - int total_cells = total_chunks * outer_split * k_split; - - for (int cell_idx = cell_start; - cell_idx < total_cells; - cell_idx += cell_stride) { - int k_slice = cell_idx % k_split; - int cell_om = cell_idx / k_split; - int chunk_idx = cell_om / outer_split; - int lane = cell_om - chunk_idx * outer_split; - - int e = chunk_idx / num_n_tiles; - int n_tile = chunk_idx - e * num_n_tiles; - int walk_begin = lane * num_m_tiles / outer_split; - int walk_end = (lane + 1) * num_m_tiles / outer_split; - - int kb_lo = max(expert_k_starts[e], batch_kb_start); - int kb_hi = min(expert_k_ends[e], batch_kb_end); - if (kb_hi <= kb_lo) - continue; - - int k_total = kb_hi - kb_lo; - int k_per = (k_total + k_split - 1) / k_split; - int slice_lo = kb_lo + k_slice * k_per; - int slice_hi = min(slice_lo + k_per, kb_hi); - if (slice_hi <= slice_lo) - continue; - kb_lo = slice_lo; - kb_hi = slice_hi; - - for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { - auto coord = make_coord(m_tile, n_tile, _); - auto gDYT = local_tile( - mDYT, typename Traits::TileShape{}, coord, - Step<_1, X, _1>{}); - auto gZ = local_tile( - mZT, typename Traits::TileShape{}, coord, - Step{}); - auto tCgDYT = cta_mma.partition_A(gDYT); - auto tCgZ = cta_mma.partition_B(gZ); - - auto [tDYTgDYT, tDYTsDYT] = tma_partition( - tma_load_dyt, - get<2>(pair_coord_vmnk), - make_layout(size<2>(pair_layout_vmnk)), - group_modes<0, 3>(sDYT), - group_modes<0, 3>(tCgDYT)); - auto [tZgZ, tZsZ] = tma_partition( - tma_load_z, - get<1>(pair_coord_vmnk), - make_layout(size<1>(pair_layout_vmnk)), - group_modes<0, 3>(sZ), - group_modes<0, 3>(tCgZ)); - - for (int kb = kb_lo; kb < kb_hi; ++kb) { - int kb_g = ring_kb > 0 ? kb % ring_kb : kb; - pipe.producer_acquire(state); - if (cute::elect_one_sync()) { - auto* barrier = pipe.producer_get_barrier(state); - copy( - tma_load_dyt.with( - *barrier, mcast_mask_dyt), - tDYTgDYT(_, kb_g), - tDYTsDYT(_, state.index())); - copy( - tma_load_z.with( - *barrier, mcast_mask_z), - tZgZ(_, kb_g), - tZsZ(_, state.index())); - } - ++state; - } - } - } - pipe.producer_tail(state); +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp3_sm90.cuh" #else - __trap(); +#include "mlp3_sm100.cuh" #endif -} -}; // Mlp3ProducerImpl<100> - -template <> -struct Mlp3ConsumerImpl<100> { -template < - typename Traits, - bool Expert3D = false, - typename Pipeline, - typename SmemType, - typename TmaReduceAddDA> -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaReduceAddDA const& tma_reduce_da, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int hidden_dim, - int intermediate_dim, - int total_n_rows, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - const cutlass::arch::ClusterBarrier* pair_init_barrier, - uint32_t* pair_init_phase) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - (void)hidden_dim; - using Element = typename Traits::Element; - constexpr int TileM = Traits::TileM; - constexpr int TileN = Traits::TileN; - constexpr int CtaTileM = Traits::CtaTileM; - constexpr int EpiChunkN = Traits::EpiChunkN; - constexpr int NumEpiRounds = Traits::NumEpiRounds; - constexpr int kMmaEpiThreads = - Traits::ConsumerThreads + Traits::WarpSize; - - int cta_rank = static_cast(cute::block_rank_in_cluster()); - int pair_rank = cta_rank % 2; - bool is_leader_cta = pair_rank == 0; - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_mma_warp = warp_id == 3; - bool is_epilogue = warp_id >= 4 && warp_id <= 11; - int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; - int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; - int tid_in_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; - int wg_barrier_id = 1 + wg; - bool is_wg_leader = is_epilogue && tid_in_wg == 0; - - typename Traits::TiledMma2Sm tiled_mma; - auto cta_mma = tiled_mma.get_slice(pair_rank); - auto sDYT = make_tensor( - make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); - auto sZ = make_tensor( - make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); - auto tCrDYT = cta_mma.make_fragment_A(sDYT); - auto tCrZ = cta_mma.make_fragment_B(sZ); - - auto cAccFull = make_identity_tensor( - make_shape(Int{}, Int{})); - auto tCgC = cta_mma.partition_C(cAccFull); - auto tCtAcc = cta_mma.make_fragment_C(tCgC); - - if (pair_init_barrier != nullptr) { - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); - if (warp_id == 3) { - if (cute::elect_one_sync()) { - uint32_t rank = cute::block_rank_in_cluster(); - uint32_t peer = rank ^ 1; - bool leader = (rank % 2) == 0; - pair_init_barrier->arrive(peer, !leader); - pair_init_barrier->wait(*pair_init_phase); - pair_init_barrier->arrive(peer, leader); - *pair_init_phase ^= 1; - } - __syncwarp(); - } - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); - } - - using AccPipe = typename Traits::AccumulatorPipeline2Sm; - typename AccPipe::Params acc_params; - acc_params.role = is_mma_warp && is_leader_cta - ? AccPipe::ThreadCategory::Producer - : AccPipe::ThreadCategory::Consumer; - acc_params.producer_arv_count = 1; - acc_params.consumer_arv_count = 2; - acc_params.initializing_warp = 4; - AccPipe acc_pipe( - smem.acc_pipe, acc_params, typename Traits::ClusterShape{}); - auto acc_prod_state = - cutlass::make_producer_start_state(); - typename AccPipe::PipelineState acc_cons_state; - - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); - if (pair_init_barrier != nullptr && warp_id == 3) { - if (cute::elect_one_sync()) { - uint32_t rank = cute::block_rank_in_cluster(); - uint32_t peer = rank ^ 1; - bool leader = (rank % 2) == 0; - pair_init_barrier->arrive(peer, !leader); - pair_init_barrier->wait(*pair_init_phase); - pair_init_barrier->arrive(peer, leader); - *pair_init_phase ^= 1; - } - __syncwarp(); - } - if (pair_init_barrier != nullptr) - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); - uint32_t tmem_base = smem.tmem_base; - tCtAcc.data() = tmem_base; - - constexpr int store_slot_elems = CtaTileM * EpiChunkN; - Element* store_ptr = smem.store_buf + wg * store_slot_elems; - auto sStore = make_tensor( - make_smem_ptr(store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - auto mdA = tma_reduce_da.get_tma_tensor(make_shape( - static_cast(total_n_rows), - static_cast(intermediate_dim))); - auto cta_tma_da = tma_reduce_da.get_slice(Int<0>{}); - - auto epi_tile = - make_tile(Int{}, Int{}); - auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); - auto tAcc_epi = flat_divide(acc_mn, epi_tile); - auto t2r = make_tmem_copy( - TmemLoadOp{}, - tAcc_epi(_, _, _0{}, _0{})); - auto thr_t2r = t2r.get_slice(tid_in_wg); - auto cChunk = make_identity_tensor( - make_shape(Int{}, Int{})); - auto tTR_cChunk = thr_t2r.partition_D(cChunk); - auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); - - Layout tmem_warp_layout = - typename decltype(make_tmem_warp_partitioner( - tAcc_epi(_, _, _0{}, _0{})))::TiledLayout_TV{}; - constexpr bool predicate_tmem_load = - size(tmem_warp_layout) != cosize(tmem_warp_layout); - - int total_chunks = num_experts * num_n_tiles; - int total_cells = total_chunks * outer_split * k_split; - bool store_in_flight = false; - - for (int cell_idx = cell_start; - cell_idx < total_cells; - cell_idx += cell_stride) { - int k_slice = cell_idx % k_split; - int cell_om = cell_idx / k_split; - int chunk_idx = cell_om / outer_split; - int lane = cell_om - chunk_idx * outer_split; - - int e = chunk_idx / num_n_tiles; - int n_tile = chunk_idx - e * num_n_tiles; - int walk_begin = lane * num_m_tiles / outer_split; - int walk_end = (lane + 1) * num_m_tiles / outer_split; - - int kb_lo = max(expert_k_starts[e], batch_kb_start); - int kb_hi = min(expert_k_ends[e], batch_kb_end); - if (kb_hi <= kb_lo) - continue; - - int k_total = kb_hi - kb_lo; - int k_per = (k_total + k_split - 1) / k_split; - int slice_lo = kb_lo + k_slice * k_per; - int slice_hi = min(slice_lo + k_per, kb_hi); - if (slice_hi <= slice_lo) - continue; - kb_lo = slice_lo; - kb_hi = slice_hi; - - for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { - if (is_mma_warp && is_leader_cta) { - acc_pipe.producer_acquire(acc_prod_state); - int acc_stage = acc_prod_state.index(); - tCtAcc.data() = - tmem_base + uint32_t(acc_stage * TileN); - - bool first = true; - for (int kb = kb_lo; kb < kb_hi; ++kb) { - pipe.consumer_wait(state); - CUTE_UNROLL - for (int ks = 0; ks < size<2>(tCrDYT); ++ks) { - tiled_mma.accumulate_ = first - ? UMMA::ScaleOut::Zero - : UMMA::ScaleOut::One; - first = false; - gemm( - tiled_mma, - tCrDYT(_, _, ks, state.index()), - tCrZ(_, _, ks, state.index()), - tCtAcc); - } - pipe.consumer_release(state); - ++state; - } - acc_pipe.producer_commit(acc_prod_state); - ++acc_prod_state; - } - - if (is_epilogue) { - acc_pipe.consumer_wait(acc_cons_state); - int acc_stage = acc_cons_state.index(); - tCtAcc.data() = - tmem_base + uint32_t(acc_stage * TileN); - auto acc_mn_stage = - tCtAcc(make_coord(_, _), _0{}, _0{}); - auto tAcc_epi_stage = - flat_divide(acc_mn_stage, epi_tile); - auto tTR_tAcc = - thr_t2r.partition_S(tAcc_epi_stage); - - CUTE_UNROLL - for (int round = 0; - round < NumEpiRounds; - ++round) { - int chunk = wg * NumEpiRounds + round; - auto tAccChunk = - tTR_tAcc(_, _, _, _0{}, chunk); - bool issue_tmem_load = true; - if constexpr (predicate_tmem_load) { - int subpart = - (tAccChunk.data().dp_ / 32) % 4; - issue_tmem_load = - tid_in_wg / Traits::WarpSize == subpart; - } - if (issue_tmem_load) { - copy(t2r, tAccChunk, tTR_rAcc); - cutlass::arch::fence_view_async_tmem_load(); - } - - int da_m = e * num_m_tiles + m_tile; - int n_chunk = - n_tile * - (Traits::TileN / Traits::EpiChunkN) + - wg * Traits::NumEpiRounds + round; - - if (store_in_flight) - cute::tma_store_wait<0>(); - cutlass::arch::NamedBarrier::sync( - Traits::WarpGroupSize, - wg_barrier_id); - - if (issue_tmem_load) { - CUTE_UNROLL - for (int i = 0; - i < size(tTR_rAcc); - ++i) { - int m_local = - get<0>(tTR_cChunk(i)); - int n_local = - get<1>(tTR_cChunk(i)); - sStore(m_local, n_local) = - static_cast( - tTR_rAcc(i)); - } - } - cutlass::arch::NamedBarrier::sync( - Traits::WarpGroupSize, - wg_barrier_id); - - if (is_wg_leader) { - cute::tma_store_fence(); - CUTE_UNROLL - for (int atom = 0; - atom < Traits::kAtomsPerCta; - ++atom) { - int m_atom = - (Traits::TileM / - Traits::AtomTileM) * - da_m + - Traits::kAtomsPerCta * - pair_rank + - atom; - auto sStoreAtom = - local_tile( - sStore, - make_tile( - Int{}, - Int{}), - make_coord(atom, 0)); - auto gdA = local_tile( - mdA, - make_tile( - Int{}, - Int{}), - make_coord(m_atom, n_chunk)); - copy( - tma_reduce_da, - cta_tma_da.partition_S( - sStoreAtom), - cta_tma_da.partition_D(gdA)); - } - cute::tma_store_arrive(); - } - store_in_flight = true; - } - - // TMEM reads done → release the acc so the next - // (cell, walk) MMA may reuse it (one elected consumer - // thread arrives). Always released after all epilogue - // rounds' register→SMEM fold + TMA_REDUCE_ADD store - // work has been issued for this m_tile — matches the - // 1SM mlp3_consumer's release placement. (The donor's - // EarlyTmemRelease experiment moved this earlier for - // one accumulator stage's worth of extra overlap; it - // was rejected — see mlp3_s5_early_release — so is not - // ported.) - cutlass::arch::NamedBarrier::sync( - Traits::ConsumerThreads, 0); - if (tid_in_epi == 0) - acc_pipe.consumer_release(acc_cons_state); - ++acc_cons_state; - } - } - } - - if (is_epilogue && store_in_flight) - cute::tma_store_wait<0>(); - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); -#else - __trap(); -#endif -} -}; // Mlp3ConsumerImpl<100> - -// ═══════════════════════════════════════════════════════════════════ -// Single-launch driver, chunk-fixed and architecture-dispatched -// ═══════════════════════════════════════════════════════════════════ -// -// Compute=90 uses the original per-CTA Hopper pipeline and cell mapping. -// Compute=100 uses a paired-CTA pipeline, pair-granular cell mapping, a -// shared Allocator2Sm allocation, and cluster synchronization. -template < - typename Traits, - int Compute = 90, - bool Expert3D = false, - typename TmaLoadDYT, - typename TmaLoadZ, - typename TmaReduceAddDA> -__device__ __forceinline__ void mlp3_fwd( - cute::conditional_t< - Compute == 100, - Mlp3Smem2Sm, - Mlp3Smem>& smem, - TmaLoadDYT const& tma_load_dyt, - TmaLoadZ const& tma_load_z, - TmaReduceAddDA const& tma_reduce_da, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int hidden_dim, - int intermediate_dim, - int num_tokens, - int total_n_rows, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int batch_kb_start = 0, - int batch_kb_end = -1, - int k_split = 1, - int ring_kb = 0) { - static_assert(Compute == 90 || Compute == 100, - "mlp3_fwd supports Compute=90 or Compute=100"); - using PipeState = typename Traits::PipelineState; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = warp_id == 0; - constexpr int kFirstConsumerWarp = Compute == 100 ? 3 : 4; - bool is_consumer = - warp_id >= kFirstConsumerWarp && warp_id <= 11; - - cute::prefetch_tma_descriptor( - tma_load_dyt.get_tma_descriptor()); - cute::prefetch_tma_descriptor( - tma_load_z.get_tma_descriptor()); - cute::prefetch_tma_descriptor( - tma_reduce_da.get_tma_descriptor()); - - auto pipe = [&]() { - if constexpr (Compute == 100) - return mlp3_make_pipe_umma_2sm( - smem.pipe_storage); - else - return mlp3_make_pipe(smem.pipe_storage); - }(); - using Pipeline = decltype(pipe); - - if constexpr (Compute == 100) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - cute::TMEM::Allocator2Sm tmem_allocator; - if (gridDim.x % Traits::ClusterM != 0) - __trap(); - cute::cluster_sync(); - constexpr int kTmemColumns = - Traits::AccStages * Traits::TileN; - if (warp_id == 4) { - tmem_allocator.allocate( - kTmemColumns, &smem.tmem_base); - __syncwarp(); - } - __syncthreads(); - cute::cluster_sync(); -#else - __trap(); -#endif - } else { - __syncthreads(); - } - - PipeState producer_state = is_producer - ? cutlass::make_producer_start_state() - : PipeState{}; - PipeState consumer_state; - - int cell_start = static_cast(blockIdx.x); - int cell_stride = static_cast(gridDim.x); - if constexpr (Compute == 100) { - cell_start /= Traits::ClusterM; - cell_stride /= Traits::ClusterM; - } - int kb_end = batch_kb_end >= 0 - ? batch_kb_end - : num_tokens / Traits::TileK; - - if (is_producer) { - mlp3_producer( - pipe, - producer_state, - smem, - tma_load_dyt, - tma_load_z, - expert_k_starts, - expert_k_ends, - num_experts, - hidden_dim, - intermediate_dim, - num_tokens, - num_m_tiles, - num_n_tiles, - outer_split, - cell_start, - cell_stride, - batch_kb_start, - kb_end, - k_split, - ring_kb); - } else if (is_consumer) { - if constexpr (Compute == 100) { - mlp3_consumer( - pipe, - consumer_state, - smem, - tma_reduce_da, - expert_k_starts, - expert_k_ends, - num_experts, - hidden_dim, - intermediate_dim, - total_n_rows, - num_m_tiles, - num_n_tiles, - outer_split, - cell_start, - cell_stride, - batch_kb_start, - kb_end, - k_split); - } else { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) - mlp3_consumer( - pipe, - consumer_state, - smem, - tma_reduce_da, - expert_k_starts, - expert_k_ends, - num_experts, - hidden_dim, - intermediate_dim, - total_n_rows, - num_m_tiles, - num_n_tiles, - outer_split, - cell_start, - cell_stride, - batch_kb_start, - kb_end, - k_split); -#else - __trap(); -#endif - } - } - - if constexpr (Compute == 100) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - cute::TMEM::Allocator2Sm tmem_allocator; - __syncthreads(); - cute::cluster_sync(); - constexpr int kTmemColumns = - Traits::AccStages * Traits::TileN; - if (warp_id == 4) { - tmem_allocator.release_allocation_lock(); - tmem_allocator.free( - smem.tmem_base, kTmemColumns); - } -#endif - } -} - -} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp3_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp3_sm100.cuh new file mode 100644 index 00000000..63f50add --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp3_sm100.cuh @@ -0,0 +1,1784 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 3: dA = dY^T @ Z (weight gradient, backward) +// ═══════════════════════════════════════════════════════════════════ +// +// TMA_REDUCE_ADD variant. Cooperative 2-WG consumer; the LARGER of +// (TileM, TileN) is the cooperation axis. Only two configs supported: +// (TileM=256, TileN=128) → M-split (Layout<_2,_1,_1>) +// (TileM=128, TileN=256) → N-split (Layout<_1,_2,_1>) ← default +// +// Per-WG output footprint = (WgTileM, WgTileN) = (128, 128) in both +// cases — each WG covers 128 rows × 128 cols per K-step via 2 wgmma +// m=64 issues, matching today's non-coop dual-WG per-WG work. +// +// Single fused DYT + Z TMA pipe (one mbarrier per stage, 2 TMA copies). +// Consumer computes a per-WG partial dA[m,n] in registers, casts to +// smem store_buf, then issues SM90_TMA_REDUCE_ADD into gmem — hardware +// atomically adds the partial. dA starts zero (torch::zeros / bwd host). +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for DYT / Z) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Chunk-fixed walk: each CTA owns ONE (chunk, lane) cell. We hold the +// SHARED (cooperative) operand chunk-constant and walk the cooperation +// (split) axis, maximizing L2 reuse of the operand both WGs consume: +// kMSplit=true (M-split, Z shared): chunk = (e, n_tile), walks m_tile; holds Z(n,*) +// kMSplit=false (N-split, dY^T shared): chunk = (e, m_tile), walks n_tile; holds dY^T(m,*) +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm100.cuh" +#include "tmem_load_op_sm100.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer specialization +// (Mlp3ConsumerImpl<100>). Header-only and safe to include under an sm_90a +// build — the tcgen05 PTX they emit is guarded by CUTE_ARCH_TCGEN05_* and only +// the Compute=100 body (itself gated on __CUDA_ARCH__ >= 1000) instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +// SM100 paired-CTA MLP3 support. Header-only, with the same guarding +// rationale as the other SM100 includes above. +#include +#include +#include + +namespace liger { + +using namespace cute; + +template +struct Mlp3UmmaTiledMmaSelector { + using type = void; +}; + +template +struct Mlp3UmmaTiledMmaSelector { + using type = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS{})); +}; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 256, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 64 +> +struct Mlp3Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + + // Cooperative 2-WG layout. Only two configs supported: + // (TileM=256, TileN=128) → kMSplit=true, WgLayout<_2,_1,_1> + // (TileM=128, TileN=256) → kMSplit=false, WgLayout<_1,_2,_1> + // Per-WG output footprint = (WgTileM=128, WgTileN=128) in both. Each + // WG covers 2 wgmma m=64 atoms per K-step (atom_count_M = 2). + static_assert((TileM == 256 && TileN == 128) || + (TileM == 128 && TileN == 256), + "Mlp3Traits supports only (TileM, TileN) = (256, 128) M-split " + "or (128, 256) N-split"); + + static constexpr bool kMSplit = (TileM > TileN); + static constexpr int AtomTileM = 64; + static constexpr int WgTileM = 128; + static constexpr int WgTileN = 128; + + static_assert(WgTileN % EpiChunkN == 0, + "EpiChunkN must divide WgTileN (=128)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + static constexpr int AtomN = WgTileN; + using GmmaAtom = typename GmmaSelectorMNMN::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtom = GMMA::Layout_MN_SW128_Atom; + + using SmemLayoutDYT_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, + Step<_2, _1>{})); + using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, + Step<_2, _1>{})); + using SmemLayoutDYT = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // ── SM100 (Compute=100) pipelines — single CTA, no cluster. ── + // MainloopPipelineUmma: UMMA-aware TMA pipeline whose consumer_release + // issues the UMMA-gated smem-buffer arrival. The Blackwell launcher must + // build it with num_consumers=1 (one umma_arrive per stage releases the + // buffer). Params/PipelineState/SharedStorage alias the Hopper + // PipelineTmaAsync's, so producer + Compute=100 consumer drive it unchanged. + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Double-buffered TMEM accumulators: MMA for the next (cell, walk) can fill + // the alternate stage while warps 4..11 drain/reduce-add the previous stage. + static constexpr int AccStages = 2; + static_assert(AccStages * TileN <= 512, + "SM100 MLP3 accumulator stages must fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + // SM100 UMMA TiledMMA — one 1SM tcgen05 atom spanning (TileM, TileN) with + // the accumulator in TMEM. mlp3 is the FIRST ported kernel with BOTH + // operands MN-major: A = dY^T (M=H contiguous) and B = Z (N=I contiguous), + // exactly matching SmemLayoutDYT / SmemLayoutZ (both Layout_MN_SW128_Atom + + // Step<_2,_1,_3>). make_umma_desc reads the MN-contiguous stride + // from each smem layout, so no transpose is needed (cf. mlp2_t's ). + using TiledMmaUmma = typename Mlp3UmmaTiledMmaSelector< + Element, ElementAccum, TileM, TileN>::type; + + static constexpr int TmaTransBytesDYT = + static_cast(size(SmemLayoutDYT_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesZ = + static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); + // Fused pipeline: DYT + Z share one mbarrier per stage. Both WGs + // consume the same Z via cooperative TiledMma slicing. + static constexpr int TmaTransBytes = TmaTransBytesDYT + TmaTransBytesZ; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; // 256 + static constexpr int NumThreads = 384; + + // Per-WG store buffer slot: (WgTileM, EpiChunkN). One slot per WG; it + // holds kAtomsPerWg = WgTileM/AtomTileM contiguous atom-rows after the + // acc→smem fold. + static constexpr int kAtomsPerWg = WgTileM / AtomTileM; // 2 + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + // TMA-reduce descriptor box = ONE wgmma atom-row (AtomTileM × EpiChunkN). + // At TileM=256 each WG owns 2 *interleaved* atoms (rows {0-63,128-191} / + // {64-127,192-255}) that map to non-contiguous gmem rows, so a WG issues + // kAtomsPerWg separate TMA stores per epi round (mirrors mlp1_fused.cuh, + // which uses an AtomTileM box but only ever has 1 atom/WG). + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; +}; + +template +using Mlp3MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Shared Memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp3Smem { + using Element = typename Traits::Element; + + static constexpr int smem_DYT_size = cosize_v; + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_DYT[smem_DYT_size]; + alignas(128) Element smem_Z[smem_Z_size]; + // Per-WG store buffer: WG1 uses [0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + typename Traits::MainloopPipeline::SharedStorage pipe_storage; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these; they + // keep the consumer signature identical across Compute values. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP3 shared memory (data only — pipe storage lives in the +// outer MlpFusedBwdSmem union; otherwise identical to Mlp3Smem). +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp3FusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_DYT_size = cosize_v; + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_DYT[smem_DYT_size]; + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element store_buf[2 * smem_store_size]; + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } +}; + + +// ═══════════════════════════════════════════════════════════════════ +// Pipe construction helper +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors mlp2/mlp5's make_pipe pattern: hides the role/transaction- +// byte setup so callers (mlp3_fwd, or any fused parent kernel) get a +// consistent Pipeline. The plain SM90 helper uses warp 0 as TMA producer and +// warps 4-11 as WGMMA consumers; the SM100 helper additionally enrolls warp 3 +// as the UMMA consumer of the TMA pipeline. + +template +__device__ __forceinline__ typename Traits::MainloopPipeline +mlp3_make_pipe(typename Traits::MainloopPipeline::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp3_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. The ctor arg pattern is identical to the Hopper pipe, so the +// arch-agnostic mlp3_producer drives it unchanged. +template +__device__ __forceinline__ typename Traits::MainloopPipelineUmma +mlp3_make_pipe_umma(typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Chunk-fixed cooperative producer/consumer +// ═══════════════════════════════════════════════════════════════════ +// +// Used by the standalone mlp3 kernel (mlp3.cu) and mlp_bwd Phase 2. +// Each CTA owns ONE (chunk, lane) cell. We hold the SHARED operand and +// walk the cooperation (split) axis: +// kMSplit=true (M-split): chunk = (e, n_tile), lane partitions +// num_m_tiles, walks m_tile. Z(n_tile, *) — the shared +// N-side operand — is the L2-held chunk-constant. +// kMSplit=false (N-split): chunk = (e, m_tile), lane partitions +// num_n_tiles, walks n_tile. dY^T(m_tile, *) — the +// shared M-side operand — is L2-held. +// +// Cell decode (chunk-major: cell_idx = chunk_idx × outer_split + lane): +// total_cells = total_chunks × outer_split +// cell → chunk_idx = cell_idx / outer_split +// lane = cell_idx % outer_split + +template +struct Mlp3ProducerImpl; + +template <> +struct Mlp3ProducerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadDYT const& tma_load_dyt, + TmaLoadZ const& tma_load_z, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + // Staging-ring depth in K-blocks. > 0 means batch_kb_*/the per-expert + // k-ranges are UNWRAPPED and the gmem K-tile coordinate wraps by % ring_kb + // (grouped grouped Phase-2 may straddle the ring). 0 = no wrap. + int ring_kb = 0) { + + auto sDYT = make_tensor(make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + + auto mDYT = tma_load_dyt.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(num_tokens))); + auto mZT = tma_load_z.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + + auto cta_tma_dyt = tma_load_dyt.get_slice(Int<0>{}); + auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); + + auto tDYTsDYT = cta_tma_dyt.partition_D(sDYT); + auto tZsZ = cta_tma_z.partition_D(sZ); + + int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); + // K-split: each (chunk, walk-lane) is replicated k_split times; the + // k_slice handles a 1/k_split sub-range of the expert's K-loop. Partial + // reductions accumulate into dA via the TMA_REDUCE_ADD epilogue. + int total_cells = total_chunks * outer_split * k_split; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e, m_chunk, n_chunk, walk_begin, walk_end; + if constexpr (Traits::kMSplit) { + // M-split: hold Z(n) shared. chunk = (e, n_tile), walks m_tile. + e = chunk_idx / num_n_tiles; + n_chunk = chunk_idx - e * num_n_tiles; + m_chunk = 0; + int m_per = num_m_tiles / outer_split; + walk_begin = lane * m_per; + walk_end = walk_begin + m_per; + } else { + // N-split: hold dY^T(m) shared. chunk = (e, m_tile), walks n_tile. + e = chunk_idx / num_m_tiles; + m_chunk = chunk_idx - e * num_m_tiles; + n_chunk = 0; + int n_per = num_n_tiles / outer_split; + walk_begin = lane * n_per; + walk_end = walk_begin + n_per; + } + + // Intersect this expert's global K range with the current + // batch's K-window. Global k_starts/ends were filled once per + // pass; per-batch we just clip. + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) continue; + + // K-split: clip to this cell's k_slice sub-range of [kb_lo, kb_hi). + { + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int s_lo = kb_lo + k_slice * k_per; + int s_hi = min(s_lo + k_per, kb_hi); + if (s_hi <= s_lo) continue; // empty slice (e.g. straggler) + kb_lo = s_lo; + kb_hi = s_hi; + } + + for (int w = walk_begin; w < walk_end; ++w) { + int m_tile = Traits::kMSplit ? w : m_chunk; + int n_tile = Traits::kMSplit ? n_chunk : w; + auto gDYT = local_tile(mDYT, + make_tile(Int{}, Int{}), + make_coord(m_tile, _)); + auto gZ = local_tile(mZT, + make_tile(Int{}, Int{}), + make_coord(n_tile, _)); + auto tDYTgDYT = cta_tma_dyt.partition_S(gDYT); + auto tZgZ = cta_tma_z.partition_S(gZ); + + for (int kb = kb_lo; kb < kb_hi; ++kb) { + int kb_g = (ring_kb > 0) ? (kb % ring_kb) : kb; // physical ring slot + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_dyt.with(*bar, 0), + tDYTgDYT(_, _, _, kb_g), tDYTsDYT(_, _, _, state.index())); + copy(tma_load_z.with(*bar, 0), + tZgZ(_, _, _, kb_g), tZsZ(_, _, _, state.index())); + } + ++state; + } + } + } +} +}; // Mlp3ProducerImpl<90> + +template +__device__ __forceinline__ void mlp3_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadDYT const& tma_load_dyt, + TmaLoadZ const& tma_load_z, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + int ring_kb = 0) { + Mlp3ProducerImpl::template run( + pipe, state, smem, tma_load_dyt, tma_load_z, + expert_k_starts, expert_k_ends, num_experts, + hidden_dim, intermediate_dim, num_tokens, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split, ring_kb); +} + +// ═══════════════════════════════════════════════════════════════════ +// Consumer — arch-split into Compute=90 (WGMMA) / Compute=100 (UMMA) +// ═══════════════════════════════════════════════════════════════════ +// Compute=90 → Hopper / WGMMA. Both WGs cooperate in ONE TiledMMA that +// partitions (TileM,TileN); each thread holds a register acc +// slice, scatters into store_buf, and the WG leader issues +// SM90_TMA_REDUCE_ADD. +// Compute=100 → Blackwell / UMMA. One 1SM tcgen05 atom (warp 3) fills ONE +// TMEM accumulator over the whole (TileM,TileN); both WGs pull +// their N-half from TMEM into the SAME store_buf slots, then the +// IDENTICAL REDUCE_ADD leader block stores them. + +template +struct Mlp3ConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper WGMMA. Verbatim from the original mlp3_consumer. +// ─────────────────────────────────────────────────────────────────── +template <> +struct Mlp3ConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDA const& tma_reduce_da, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int total_n_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split) { + + (void)Compute; + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + const int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sDYT = make_tensor(make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + + auto tCsDYT = thr_mma.partition_A(sDYT); + auto tCsZ = thr_mma.partition_B(sZ); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + // Identity coords for per-thread acc → store_buf mapping. + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + constexpr int store_slot_elems = Traits::WgTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + const bool is_my_wg_leader = (tid_in_wg == 0); + + auto mdA = [&]() { + if constexpr (Expert3D) { + return tma_reduce_da.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(num_experts))); + } else { + return tma_reduce_da.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(intermediate_dim))); + } + }(); + auto cta_tma_da = tma_reduce_da.get_slice(Int<0>{}); + + int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); + int total_cells = total_chunks * outer_split * k_split; + + constexpr int K_PIPE_MMAS = 1; + + // Carries across walk iters AND across cells for cross-tile K-loop / + // store overlap. Drained once at the end. + bool store_in_flight = false; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e, m_chunk, n_chunk, walk_begin, walk_end; + if constexpr (Traits::kMSplit) { + // M-split: hold Z(n) shared. chunk = (e, n_tile), walks m_tile. + e = chunk_idx / num_n_tiles; + n_chunk = chunk_idx - e * num_n_tiles; + m_chunk = 0; + int m_per = num_m_tiles / outer_split; + walk_begin = lane * m_per; + walk_end = walk_begin + m_per; + } else { + // N-split: hold dY^T(m) shared. chunk = (e, m_tile), walks n_tile. + e = chunk_idx / num_m_tiles; + m_chunk = chunk_idx - e * num_m_tiles; + n_chunk = 0; + int n_per = num_n_tiles / outer_split; + walk_begin = lane * n_per; + walk_end = walk_begin + n_per; + } + + // Intersect this expert's global K range with the current + // batch's K-window. + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) continue; + + // K-split: clip to this cell's k_slice sub-range (must match producer). + { + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int s_lo = kb_lo + k_slice * k_per; + int s_hi = min(s_lo + k_per, kb_hi); + if (s_hi <= s_lo) continue; // empty slice + kb_lo = s_lo; + kb_hi = s_hi; + } + + int k_count = kb_hi - kb_lo; + + for (int w = walk_begin; w < walk_end; ++w) { + int m_tile = Traits::kMSplit ? w : m_chunk; + int n_tile = Traits::kMSplit ? n_chunk : w; + + clear(acc); + auto state_release = state; + int prologue_count = (k_count < K_PIPE_MMAS) ? k_count : K_PIPE_MMAS; + + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsDYT(_, _, _, state.index()), + tCsZ(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + + for (int k = prologue_count; k < k_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsDYT(_, _, _, state.index()), + tCsZ(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + + ++state; + ++state_release; + } + + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // Each WG writes its (WgTileM × EpiChunkN) slot to store_buf + // and issues TMA stores. NumEpiRounds rounds per WG. + // kMSplit=true : M-split. WG_w owns WgTileM contiguous + // rows of TileM. With WgLayout_M=2 and + // atom_count_M = WgTileM/AtomTileM, the + // natural CUTE m-atom positions interleave + // between WGs — formula below maps m_loc + // back to a contiguous local row. + // kMSplit=false : N-split. WG_w owns full TileM, cols + // [w*WgTileN, (w+1)*WgTileN). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + // CUTE atom-iter interleaves WGs in M direction + // when atom_count_M > 1 (TileM=256). Unified + // formula folds 64-row WG strips into a + // contiguous WgTileM-row local range: + // m_atom_global = m_loc / 64 (atom row idx) + // m_atom_in_wg = m_atom_global / 2 + // m_local = m_atom_in_wg*64 + (m_loc % 64) + // = (m_loc / 128)*64 + (m_loc % 64) + // For tile_count_M=1 (TileM=128) the / 128 term + // is 0 so m_local = m_loc % 64. + m_local = (m_loc / 128) * 64 + (m_loc % 64); + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int da_m = Expert3D ? m_tile : e * num_m_tiles + m_tile; + // Store box = ONE atom-row (AtomTileM). Each WG issues + // kAtomsPerWg stores, one per atom it owns. The store_buf + // holds the WG's atoms in contiguous local rows + // [a*AtomTileM, +AtomTileM); their gmem atom-rows differ: + // M-split: the 2 atoms interleave → 4*da_m + my_wg + 2a + // (gmem M-grid is TileM/AtomTileM rows per tile). + // N-split: WG owns the full TileM (2 contiguous atoms) and + // its own N-half → 2*da_m + a, n offset by my_wg. + int n_tile_idx; + if constexpr (Traits::kMSplit) { + n_tile_idx = n_tile * Traits::NumEpiRounds + r; + } else { + n_tile_idx = n_tile * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + constexpr int kRowAtoms = Traits::TileM / Traits::AtomTileM; // 4 (M-split) / 2 (N-split) + CUTE_UNROLL + for (int a = 0; a < Traits::kAtomsPerWg; ++a) { + int m_atom_row = Traits::kMSplit + ? (kRowAtoms * da_m + my_wg + Traits::kAtomsPerWg * a) + : (kRowAtoms * da_m + a); + auto sStore_a = local_tile(sStore, + make_tile(Int{}, Int{}), + make_coord(a, 0)); + auto gdA = [&]() { + if constexpr (Expert3D) { + return local_tile(mdA, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx, e)); + } else { + return local_tile(mdA, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx)); + } + }(); + copy(tma_reduce_da, cta_tma_da.partition_S(sStore_a), + cta_tma_da.partition_D(gdA)); + } + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp3ConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — Compute defaults to 90 so existing call sites +// (mlp3_consumer(...)) are unchanged; pass mlp3_consumer +// for the Blackwell path. +// ─────────────────────────────────────────────────────────────────── +template +__device__ __forceinline__ void mlp3_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDA const& tma_reduce_da, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int total_n_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, + uint32_t* pair_init_phase = nullptr) { + if constexpr (Compute == 100) { + Mlp3ConsumerImpl::template run( + pipe, state, smem, tma_reduce_da, + expert_k_starts, expert_k_ends, num_experts, + hidden_dim, intermediate_dim, total_n_rows, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split, + pair_init_barrier, pair_init_phase); + } else { + Mlp3ConsumerImpl::template run( + pipe, state, smem, tma_reduce_da, + expert_k_starts, expert_k_ends, num_experts, + hidden_dim, intermediate_dim, total_n_rows, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split); + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Host launcher (raw-pointer; defined in mlp3.cu) +// ═══════════════════════════════════════════════════════════════════ + +// dA = dY^T @ Z, accumulated atomically into the [E, H, I] output via +// TMA_REDUCE_ADD. dA MUST be zero-initialized by the caller. +// +// dY : [num_tokens, hidden_dim] +// Z : [num_tokens, intermediate_dim] +// dA : [num_experts, hidden_dim, intermediate_dim] +// expert_for_k_block: [num_tokens / TileK] int32, TileK=64 +void mlp3_fwd_bf16_launch( + const cutlass::bfloat16_t* dY, + const cutlass::bfloat16_t* Z, + cutlass::bfloat16_t* dA, + const int* expert_for_k_block, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int device, + cudaStream_t stream); + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM (paired-CTA) MLP3 — basic ClusterM=2 path +// ═══════════════════════════════════════════════════════════════════ +// +// Joins two adjacent CTAs into one cta_group::2 UMMA atom covering a +// 256x256x64 joined tile, doubling the M extent a single CTA-pair can +// produce per wave versus the 1SM path above. Callers must launch with +// cudaLaunchKernelEx, clusterDim=(2,1,1), an even gridDim.x, and +// pair-aware SM100 2SM TMA load descriptors for dY^T / Z. The dA +// descriptor remains an ordinary (non-paired) TMA_REDUCE_ADD. +// +// TileK, Stages, and EpiChunkN stay trait parameters — the validated +// production candidate is TileK=64/Stages=5/EpiChunkN=64 (see +// 2sm.md / 2sm_stage.md), but that tuple is a caller-side choice, not +// hardcoded here. ClusterM is fixed at 2 (one paired-CTA "pair"): the +// ClusterM=4 (two-pair, N-split-across-pairs) and shared-operand TMA +// multicast variants explored alongside this trait did not clear the +// required performance bar (see the mlp3_multicast bench results) and +// are intentionally not ported, nor are the CompactEpilogue and +// EarlyTmemRelease epilogue experiments (see mlp3_s6_bounce / +// mlp3_s5_early_release) — both were rejected and are omitted here. +// +// Host-launcher contract (no launcher lives in this header — these are +// hard requirements on any .cu translation unit that selects Compute=100, +// e.g. an mlp_bwd.cuh fused entry point): +// 1. Grid: gridDim.x must be even (mlp3_fwd __traps if it is not); +// launch must use cudaLaunchKernelEx with cudaLaunchConfig_t:: +// cudaLaunchAttributeClusterDimension = {ClusterM, 1, 1}. +// 2. The kernel function must first be registered via +// cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) +// before the cudaLaunchKernelEx call — the driver rejects a 2-CTA +// cluster launch of a kernel that exceeds the "portable" static +// shared-memory budget without this opt-in. +// 3. The dY^T / Z TMA operands passed in as TmaLoadDYT / TmaLoadZ MUST +// be constructed with cute::make_tma_copy_A_sm100 / +// cute::make_tma_copy_B_sm100 (NOT the ordinary 1SM make_tma_copy) +// — only the _sm100 constructors produce the cta_group::2, +// pair-aware copy atoms that create_tma_multicast_mask / +// tma_partition below assume. The dA reduce-add operand stays an +// ordinary (non-paired) TMA_REDUCE_ADD, per above. +// 4. Dynamic shared memory for Mlp3Smem2Sm / Mlp3FusedSmem2Sm +// must be sized and opted in at runtime by the launcher — +// query cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemory +// PerBlockOptin, device), then cudaFuncSetAttribute(kernel, +// cudaFuncAttributeMaxDynamicSharedMemorySize, requested) with +// requested <= optin. This header intentionally has no compile- +// time static_assert on a fixed smem byte ceiling (only the TMEM +// column assert on AccStages*TileN below) — the ceiling is a +// runtime, device-dependent quantity and must not be hardcoded. +template < + typename Element_, + int TileM_ = 256, + int TileN_ = 256, + int TileK_ = 64, + int Stages_ = 3, + int EpiChunkN_ = 64, + int AccStages_ = 2, + int ClusterM_ = 2> +struct Mlp3Traits2Sm { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + static constexpr int AccStages = AccStages_; + static constexpr int ClusterM = ClusterM_; + static constexpr int NumPairs = ClusterM / 2; + static constexpr int CtaTileM = TileM / 2; + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = TileN / 2; + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + static constexpr int kAtomsPerCta = CtaTileM / AtomTileM; + + static_assert(TileM == 256, + "MLP3 2SM requires joined TileM=256; joined TileM=128 is invalid"); + static_assert(TileN == 128 || TileN == 256, + "MLP3 2SM supports TileN=128 or TileN=256"); + static_assert(TileK == 32 || TileK == 64, + "MLP3 2SM pilot supports TileK=32 or 64"); + static_assert(Stages >= 2 && Stages <= 12, + "MLP3 2SM supports two to twelve mainloop stages"); + static_assert(CtaTileM == 128, "each peer CTA must own 128 M rows"); + static_assert(WgTileN % EpiChunkN == 0, + "EpiChunkN must divide each warpgroup's N half"); + static_assert(kAtomsPerCta == 2, + "each CTA must reduce-add two 64-row atoms"); + static_assert(AccStages >= 2, "MLP3 2SM requires accumulator double buffering"); + static_assert(AccStages * TileN <= 512, + "MLP3 2SM accumulator stages must fit in 512 TMEM columns"); + static_assert(ClusterM == 2, + "MLP3 2SM only ports the validated single-pair path; ClusterM=4 " + "(two-pair N-split, evaluated together with shared-operand " + "multicast) was rejected -- see the mlp3_multicast bench results"); + + using TileShape = Shape, Int, Int>; + using ClusterShape = Shape, _1, _1>; + using AtomThrShape = Shape<_2, _1, _1>; + + using TiledMma2Sm = decltype(make_tiled_mma( + SM100_MMA_F16BF16_2x1SM_SS< + Element, Element, ElementAccum, TileM, TileN, + UMMA::Major::MN, UMMA::Major::MN>{})); + static_assert(size(typename TiledMma2Sm::AtomThrID{}) == 2, + "MLP3 2SM requires a two-CTA MMA atom"); + + using ClusterLayoutVMNK = decltype(tiled_divide( + make_layout(ClusterShape{}), + make_tile(typename TiledMma2Sm::AtomThrID{}))); + + using MmaShapeA_MK = decltype(partition_shape_A( + TiledMma2Sm{}, make_shape(Int{}, Int{}))); + using MmaShapeB_NK = decltype(partition_shape_B( + TiledMma2Sm{}, make_shape(Int{}, Int{}))); + + using SmemLayoutAtom = UMMA::Layout_MN_SW128_Atom; + using SmemLayoutDYT = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtom{}, append(MmaShapeA_MK{}, Int{}), + Step<_2, _1, _3>{})); + using SmemLayoutZ = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtom{}, append(MmaShapeB_NK{}, Int{}), + Step<_2, _1, _3>{})); + using SmemLayoutDYT_1 = + decltype(SmemLayoutDYT{}(_, _, _, Int<0>{})); + using SmemLayoutZ_1 = + decltype(SmemLayoutZ{}(_, _, _, Int<0>{})); + + static constexpr int TmaTransBytesDYT = + static_cast(cosize_v * sizeof(Element)); + static constexpr int TmaTransBytesZ = + static_cast(cosize_v * sizeof(Element)); + static constexpr int TmaTransBytes = + 2 * (TmaTransBytesDYT + TmaTransBytesZ); + + using MainloopPipelineUmma2Sm = cutlass::PipelineTmaUmmaAsync< + Stages, ClusterShape, AtomThrShape>; + using MainloopPipeline = MainloopPipelineUmma2Sm; + using MainloopPipelineUmma = MainloopPipelineUmma2Sm; + using PipelineState = typename MainloopPipelineUmma2Sm::PipelineState; + using AccumulatorPipeline2Sm = cutlass::PipelineUmmaAsync< + AccStages, AtomThrShape>; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; + + // One (CtaTileM, EpiChunkN) store slot per epilogue warpgroup; WG0 + // uses [0..S-1], WG1 uses [S..2S-1] (S = smem_store_size, see + // Mlp3Smem2Sm below). Mirrors the 1SM Mlp3Traits::SmemLayoutStoreSlot, + // sized CtaTileM (=128) instead of WgTileM. + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; +}; + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM Shared Memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp3Smem2Sm { + using Element = typename Traits::Element; + + static constexpr int smem_DYT_size = + cosize_v; + static constexpr int smem_Z_size = + cosize_v; + static constexpr int smem_store_size = + cosize_v; + + alignas(128) Element smem_DYT[smem_DYT_size]; + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element store_buf[2 * smem_store_size]; + alignas(16) typename Traits::MainloopPipelineUmma2Sm::SharedStorage pipe_storage; + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline2Sm::SharedStorage acc_pipe; + + CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } +}; + +// Fused MLP3 2SM shared memory (data only — the mainloop pipe's storage +// lives in the outer fused-kernel smem union; otherwise identical to +// Mlp3Smem2Sm). Mirrors the 1SM Mlp3FusedSmem / Mlp3Smem split. +template +struct Mlp3FusedSmem2Sm { + using Element = typename Traits::Element; + + static constexpr int smem_DYT_size = + cosize_v; + static constexpr int smem_Z_size = + cosize_v; + static constexpr int smem_store_size = + cosize_v; + + alignas(128) Element smem_DYT[smem_DYT_size]; + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element store_buf[2 * smem_store_size]; + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline2Sm::SharedStorage acc_pipe; + + CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM pipe construction helper +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors mlp3_make_pipe's role/transaction-byte setup, adjusted for the +// paired-CTA cluster: warp 0 is the TMA producer (leader CTA elects the +// arrival), warps 3-11 are UMMA/epilogue consumers. +template +__device__ __forceinline__ typename Traits::MainloopPipelineUmma2Sm +mlp3_make_pipe_umma_2sm( + typename Traits::MainloopPipelineUmma2Sm::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipelineUmma2Sm; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = warp_id == 0; + bool is_consumer = warp_id >= 3 && warp_id <= 11; + + typename Pipeline::Params params; + params.transaction_bytes = Traits::TmaTransBytes; + params.num_producers = 1; + params.num_consumers = 1; + params.initializing_warp = 0; + if (is_producer) { + params.role = Category::Producer; + params.is_leader = + threadIdx.x == 0 && + cute::block_rank_in_cluster() % 2 == 0; + } else if (is_consumer) { + params.role = Category::Consumer; + } else { + params.role = Category::NonParticipant; + } + return Pipeline( + storage, params, typename Traits::ClusterShape{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM producer / consumer +// ═══════════════════════════════════════════════════════════════════ +// +// Chunk-fixed walk, same principle as mlp3_producer/mlp3_consumer above: +// each CTA-pair owns ONE (chunk, lane) cell and walks m_tile × kb +// internally. Within a pair, the leader CTA (pair_rank==0) drives the +// paired UMMA and owns the accumulator producer role; both CTAs run the +// epilogue over their own CtaTileM=128 half of the joined 256-row tile. +template <> +struct Mlp3ProducerImpl<100> { +template < + typename Traits, + typename Pipeline, + typename SmemType, + typename TmaLoadDYT, + typename TmaLoadZ> +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadDYT const& tma_load_dyt, + TmaLoadZ const& tma_load_z, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + int ring_kb = 0) { + // Compute!=100 must stay 1SM: this body issues SM100-only paired-CTA + // TMA ops (pair-aware tma_partition / create_tma_multicast_mask on a + // TiledMma2Sm/MainloopPipelineUmma2Sm operand) and must never be live + // code for an sm_90a compilation pass of this same translation unit. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + auto sDYT = make_tensor( + make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); + auto sZ = make_tensor( + make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + + auto mDYT = tma_load_dyt.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(num_tokens))); + auto mZT = tma_load_z.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + + typename Traits::TiledMma2Sm tiled_mma; + int cta_rank = static_cast(cute::block_rank_in_cluster()); + int pair_rank = cta_rank % 2; + auto cta_mma = tiled_mma.get_slice(pair_rank); + auto pair_layout_vmnk = tiled_divide( + make_layout(typename Traits::ClusterShape{}), + make_tile(typename Traits::TiledMma2Sm::AtomThrID{})); + auto pair_coord_vmnk = + pair_layout_vmnk.get_flat_coord(pair_rank); + uint16_t mcast_mask_dyt = + create_tma_multicast_mask<2>(pair_layout_vmnk, pair_coord_vmnk); + uint16_t mcast_mask_z = + create_tma_multicast_mask<1>(pair_layout_vmnk, pair_coord_vmnk); + + int total_chunks = num_experts * num_n_tiles; + int total_cells = total_chunks * outer_split * k_split; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e = chunk_idx / num_n_tiles; + int n_tile = chunk_idx - e * num_n_tiles; + int walk_begin = lane * num_m_tiles / outer_split; + int walk_end = (lane + 1) * num_m_tiles / outer_split; + + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) + continue; + + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int slice_lo = kb_lo + k_slice * k_per; + int slice_hi = min(slice_lo + k_per, kb_hi); + if (slice_hi <= slice_lo) + continue; + kb_lo = slice_lo; + kb_hi = slice_hi; + + for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { + auto coord = make_coord(m_tile, n_tile, _); + auto gDYT = local_tile( + mDYT, typename Traits::TileShape{}, coord, + Step<_1, X, _1>{}); + auto gZ = local_tile( + mZT, typename Traits::TileShape{}, coord, + Step{}); + auto tCgDYT = cta_mma.partition_A(gDYT); + auto tCgZ = cta_mma.partition_B(gZ); + + auto [tDYTgDYT, tDYTsDYT] = tma_partition( + tma_load_dyt, + get<2>(pair_coord_vmnk), + make_layout(size<2>(pair_layout_vmnk)), + group_modes<0, 3>(sDYT), + group_modes<0, 3>(tCgDYT)); + auto [tZgZ, tZsZ] = tma_partition( + tma_load_z, + get<1>(pair_coord_vmnk), + make_layout(size<1>(pair_layout_vmnk)), + group_modes<0, 3>(sZ), + group_modes<0, 3>(tCgZ)); + + for (int kb = kb_lo; kb < kb_hi; ++kb) { + int kb_g = ring_kb > 0 ? kb % ring_kb : kb; + pipe.producer_acquire(state); + if (cute::elect_one_sync()) { + auto* barrier = pipe.producer_get_barrier(state); + copy( + tma_load_dyt.with( + *barrier, mcast_mask_dyt), + tDYTgDYT(_, kb_g), + tDYTsDYT(_, state.index())); + copy( + tma_load_z.with( + *barrier, mcast_mask_z), + tZgZ(_, kb_g), + tZsZ(_, state.index())); + } + ++state; + } + } + } + pipe.producer_tail(state); +#else + __trap(); +#endif +} +}; // Mlp3ProducerImpl<100> + +template <> +struct Mlp3ConsumerImpl<100> { +template < + typename Traits, + bool Expert3D = false, + typename Pipeline, + typename SmemType, + typename TmaReduceAddDA> +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDA const& tma_reduce_da, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int total_n_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + const cutlass::arch::ClusterBarrier* pair_init_barrier, + uint32_t* pair_init_phase) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + (void)hidden_dim; + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int CtaTileM = Traits::CtaTileM; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int NumEpiRounds = Traits::NumEpiRounds; + constexpr int kMmaEpiThreads = + Traits::ConsumerThreads + Traits::WarpSize; + + int cta_rank = static_cast(cute::block_rank_in_cluster()); + int pair_rank = cta_rank % 2; + bool is_leader_cta = pair_rank == 0; + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_mma_warp = warp_id == 3; + bool is_epilogue = warp_id >= 4 && warp_id <= 11; + int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; + int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + int tid_in_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + int wg_barrier_id = 1 + wg; + bool is_wg_leader = is_epilogue && tid_in_wg == 0; + + typename Traits::TiledMma2Sm tiled_mma; + auto cta_mma = tiled_mma.get_slice(pair_rank); + auto sDYT = make_tensor( + make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); + auto sZ = make_tensor( + make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto tCrDYT = cta_mma.make_fragment_A(sDYT); + auto tCrZ = cta_mma.make_fragment_B(sZ); + + auto cAccFull = make_identity_tensor( + make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + if (pair_init_barrier != nullptr) { + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + if (warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + pair_init_barrier->arrive(peer, !leader); + pair_init_barrier->wait(*pair_init_phase); + pair_init_barrier->arrive(peer, leader); + *pair_init_phase ^= 1; + } + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + } + + using AccPipe = typename Traits::AccumulatorPipeline2Sm; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp && is_leader_cta + ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; + acc_params.consumer_arv_count = 2; + acc_params.initializing_warp = 4; + AccPipe acc_pipe( + smem.acc_pipe, acc_params, typename Traits::ClusterShape{}); + auto acc_prod_state = + cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + if (pair_init_barrier != nullptr && warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + pair_init_barrier->arrive(peer, !leader); + pair_init_barrier->wait(*pair_init_phase); + pair_init_barrier->arrive(peer, leader); + *pair_init_phase ^= 1; + } + __syncwarp(); + } + if (pair_init_barrier != nullptr) + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + constexpr int store_slot_elems = CtaTileM * EpiChunkN; + Element* store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor( + make_smem_ptr(store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + auto mdA = tma_reduce_da.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(intermediate_dim))); + auto cta_tma_da = tma_reduce_da.get_slice(Int<0>{}); + + auto epi_tile = + make_tile(Int{}, Int{}); + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi = flat_divide(acc_mn, epi_tile); + auto t2r = make_tmem_copy( + TmemLoadOp{}, + tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tid_in_wg); + auto cChunk = make_identity_tensor( + make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); + + Layout tmem_warp_layout = + typename decltype(make_tmem_warp_partitioner( + tAcc_epi(_, _, _0{}, _0{})))::TiledLayout_TV{}; + constexpr bool predicate_tmem_load = + size(tmem_warp_layout) != cosize(tmem_warp_layout); + + int total_chunks = num_experts * num_n_tiles; + int total_cells = total_chunks * outer_split * k_split; + bool store_in_flight = false; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e = chunk_idx / num_n_tiles; + int n_tile = chunk_idx - e * num_n_tiles; + int walk_begin = lane * num_m_tiles / outer_split; + int walk_end = (lane + 1) * num_m_tiles / outer_split; + + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) + continue; + + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int slice_lo = kb_lo + k_slice * k_per; + int slice_hi = min(slice_lo + k_per, kb_hi); + if (slice_hi <= slice_lo) + continue; + kb_lo = slice_lo; + kb_hi = slice_hi; + + for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { + if (is_mma_warp && is_leader_cta) { + acc_pipe.producer_acquire(acc_prod_state); + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = + tmem_base + uint32_t(acc_stage * TileN); + + bool first = true; + for (int kb = kb_lo; kb < kb_hi; ++kb) { + pipe.consumer_wait(state); + CUTE_UNROLL + for (int ks = 0; ks < size<2>(tCrDYT); ++ks) { + tiled_mma.accumulate_ = first + ? UMMA::ScaleOut::Zero + : UMMA::ScaleOut::One; + first = false; + gemm( + tiled_mma, + tCrDYT(_, _, ks, state.index()), + tCrZ(_, _, ks, state.index()), + tCtAcc); + } + pipe.consumer_release(state); + ++state; + } + acc_pipe.producer_commit(acc_prod_state); + ++acc_prod_state; + } + + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = + tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = + tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = + flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc = + thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int round = 0; + round < NumEpiRounds; + ++round) { + int chunk = wg * NumEpiRounds + round; + auto tAccChunk = + tTR_tAcc(_, _, _, _0{}, chunk); + bool issue_tmem_load = true; + if constexpr (predicate_tmem_load) { + int subpart = + (tAccChunk.data().dp_ / 32) % 4; + issue_tmem_load = + tid_in_wg / Traits::WarpSize == subpart; + } + if (issue_tmem_load) { + copy(t2r, tAccChunk, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + } + + int da_m = e * num_m_tiles + m_tile; + int n_chunk = + n_tile * + (Traits::TileN / Traits::EpiChunkN) + + wg * Traits::NumEpiRounds + round; + + if (store_in_flight) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, + wg_barrier_id); + + if (issue_tmem_load) { + CUTE_UNROLL + for (int i = 0; + i < size(tTR_rAcc); + ++i) { + int m_local = + get<0>(tTR_cChunk(i)); + int n_local = + get<1>(tTR_cChunk(i)); + sStore(m_local, n_local) = + static_cast( + tTR_rAcc(i)); + } + } + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, + wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + CUTE_UNROLL + for (int atom = 0; + atom < Traits::kAtomsPerCta; + ++atom) { + int m_atom = + (Traits::TileM / + Traits::AtomTileM) * + da_m + + Traits::kAtomsPerCta * + pair_rank + + atom; + auto sStoreAtom = + local_tile( + sStore, + make_tile( + Int{}, + Int{}), + make_coord(atom, 0)); + auto gdA = local_tile( + mdA, + make_tile( + Int{}, + Int{}), + make_coord(m_atom, n_chunk)); + copy( + tma_reduce_da, + cta_tma_da.partition_S( + sStoreAtom), + cta_tma_da.partition_D(gdA)); + } + cute::tma_store_arrive(); + } + store_in_flight = true; + } + + // TMEM reads done → release the acc so the next + // (cell, walk) MMA may reuse it (one elected consumer + // thread arrives). Always released after all epilogue + // rounds' register→SMEM fold + TMA_REDUCE_ADD store + // work has been issued for this m_tile — matches the + // 1SM mlp3_consumer's release placement. (The donor's + // EarlyTmemRelease experiment moved this earlier for + // one accumulator stage's worth of extra overlap; it + // was rejected — see mlp3_s5_early_release — so is not + // ported.) + cutlass::arch::NamedBarrier::sync( + Traits::ConsumerThreads, 0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + } + + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); +#else + __trap(); +#endif +} +}; // Mlp3ConsumerImpl<100> + +// ═══════════════════════════════════════════════════════════════════ +// Single-launch driver, chunk-fixed and architecture-dispatched +// ═══════════════════════════════════════════════════════════════════ +// +// Compute=90 uses the original per-CTA Hopper pipeline and cell mapping. +// Compute=100 uses a paired-CTA pipeline, pair-granular cell mapping, a +// shared Allocator2Sm allocation, and cluster synchronization. +template < + typename Traits, + int Compute = 90, + bool Expert3D = false, + typename TmaLoadDYT, + typename TmaLoadZ, + typename TmaReduceAddDA> +__device__ __forceinline__ void mlp3_fwd( + cute::conditional_t< + Compute == 100, + Mlp3Smem2Sm, + Mlp3Smem>& smem, + TmaLoadDYT const& tma_load_dyt, + TmaLoadZ const& tma_load_z, + TmaReduceAddDA const& tma_reduce_da, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int total_n_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int batch_kb_start = 0, + int batch_kb_end = -1, + int k_split = 1, + int ring_kb = 0) { + static_assert(Compute == 90 || Compute == 100, + "mlp3_fwd supports Compute=90 or Compute=100"); + using PipeState = typename Traits::PipelineState; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = warp_id == 0; + constexpr int kFirstConsumerWarp = Compute == 100 ? 3 : 4; + bool is_consumer = + warp_id >= kFirstConsumerWarp && warp_id <= 11; + + cute::prefetch_tma_descriptor( + tma_load_dyt.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_load_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_reduce_da.get_tma_descriptor()); + + auto pipe = [&]() { + if constexpr (Compute == 100) + return mlp3_make_pipe_umma_2sm( + smem.pipe_storage); + else + return mlp3_make_pipe(smem.pipe_storage); + }(); + using Pipeline = decltype(pipe); + + if constexpr (Compute == 100) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + cute::TMEM::Allocator2Sm tmem_allocator; + if (gridDim.x % Traits::ClusterM != 0) + __trap(); + cute::cluster_sync(); + constexpr int kTmemColumns = + Traits::AccStages * Traits::TileN; + if (warp_id == 4) { + tmem_allocator.allocate( + kTmemColumns, &smem.tmem_base); + __syncwarp(); + } + __syncthreads(); + cute::cluster_sync(); +#else + __trap(); +#endif + } else { + __syncthreads(); + } + + PipeState producer_state = is_producer + ? cutlass::make_producer_start_state() + : PipeState{}; + PipeState consumer_state; + + int cell_start = static_cast(blockIdx.x); + int cell_stride = static_cast(gridDim.x); + if constexpr (Compute == 100) { + cell_start /= Traits::ClusterM; + cell_stride /= Traits::ClusterM; + } + int kb_end = batch_kb_end >= 0 + ? batch_kb_end + : num_tokens / Traits::TileK; + + if (is_producer) { + mlp3_producer( + pipe, + producer_state, + smem, + tma_load_dyt, + tma_load_z, + expert_k_starts, + expert_k_ends, + num_experts, + hidden_dim, + intermediate_dim, + num_tokens, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split, + ring_kb); + } else if (is_consumer) { + if constexpr (Compute == 100) { + mlp3_consumer( + pipe, + consumer_state, + smem, + tma_reduce_da, + expert_k_starts, + expert_k_ends, + num_experts, + hidden_dim, + intermediate_dim, + total_n_rows, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split); + } else { +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) + mlp3_consumer( + pipe, + consumer_state, + smem, + tma_reduce_da, + expert_k_starts, + expert_k_ends, + num_experts, + hidden_dim, + intermediate_dim, + total_n_rows, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split); +#else + __trap(); +#endif + } + } + + if constexpr (Compute == 100) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + cute::TMEM::Allocator2Sm tmem_allocator; + __syncthreads(); + cute::cluster_sync(); + constexpr int kTmemColumns = + Traits::AccStages * Traits::TileN; + if (warp_id == 4) { + tmem_allocator.release_allocation_lock(); + tmem_allocator.free( + smem.tmem_base, kTmemColumns); + } +#endif + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp3_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp3_sm90.cuh new file mode 100644 index 00000000..4ef32d1a --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp3_sm90.cuh @@ -0,0 +1,1784 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 3: dA = dY^T @ Z (weight gradient, backward) +// ═══════════════════════════════════════════════════════════════════ +// +// TMA_REDUCE_ADD variant. Cooperative 2-WG consumer; the LARGER of +// (TileM, TileN) is the cooperation axis. Only two configs supported: +// (TileM=256, TileN=128) → M-split (Layout<_2,_1,_1>) +// (TileM=128, TileN=256) → N-split (Layout<_1,_2,_1>) ← default +// +// Per-WG output footprint = (WgTileM, WgTileN) = (128, 128) in both +// cases — each WG covers 128 rows × 128 cols per K-step via 2 wgmma +// m=64 issues, matching today's non-coop dual-WG per-WG work. +// +// Single fused DYT + Z TMA pipe (one mbarrier per stage, 2 TMA copies). +// Consumer computes a per-WG partial dA[m,n] in registers, casts to +// smem store_buf, then issues SM90_TMA_REDUCE_ADD into gmem — hardware +// atomically adds the partial. dA starts zero (torch::zeros / bwd host). +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for DYT / Z) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Chunk-fixed walk: each CTA owns ONE (chunk, lane) cell. We hold the +// SHARED (cooperative) operand chunk-constant and walk the cooperation +// (split) axis, maximizing L2 reuse of the operand both WGs consume: +// kMSplit=true (M-split, Z shared): chunk = (e, n_tile), walks m_tile; holds Z(n,*) +// kMSplit=false (N-split, dY^T shared): chunk = (e, m_tile), walks n_tile; holds dY^T(m,*) +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm90.cuh" +#include "tmem_load_op_sm90.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer specialization +// (Mlp3ConsumerImpl<100>). Header-only and safe to include under an sm_90a +// build — the tcgen05 PTX they emit is guarded by CUTE_ARCH_TCGEN05_* and only +// the Compute=100 body (itself gated on __CUDA_ARCH__ >= 1000) instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +// SM100 paired-CTA MLP3 support. Header-only, with the same guarding +// rationale as the other SM100 includes above. +#include +#include +#include + +namespace liger { + +using namespace cute; + +template +struct Mlp3UmmaTiledMmaSelector { + using type = void; +}; + +template +struct Mlp3UmmaTiledMmaSelector { + using type = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS{})); +}; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 256, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 64 +> +struct Mlp3Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + + // Cooperative 2-WG layout. Only two configs supported: + // (TileM=256, TileN=128) → kMSplit=true, WgLayout<_2,_1,_1> + // (TileM=128, TileN=256) → kMSplit=false, WgLayout<_1,_2,_1> + // Per-WG output footprint = (WgTileM=128, WgTileN=128) in both. Each + // WG covers 2 wgmma m=64 atoms per K-step (atom_count_M = 2). + static_assert((TileM == 256 && TileN == 128) || + (TileM == 128 && TileN == 256), + "Mlp3Traits supports only (TileM, TileN) = (256, 128) M-split " + "or (128, 256) N-split"); + + static constexpr bool kMSplit = (TileM > TileN); + static constexpr int AtomTileM = 64; + static constexpr int WgTileM = 128; + static constexpr int WgTileN = 128; + + static_assert(WgTileN % EpiChunkN == 0, + "EpiChunkN must divide WgTileN (=128)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + static constexpr int AtomN = WgTileN; + using GmmaAtom = typename GmmaSelectorMNMN::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtom = GMMA::Layout_MN_SW128_Atom; + + using SmemLayoutDYT_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, + Step<_2, _1>{})); + using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, + Step<_2, _1>{})); + using SmemLayoutDYT = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // ── SM100 (Compute=100) pipelines — single CTA, no cluster. ── + // MainloopPipelineUmma: UMMA-aware TMA pipeline whose consumer_release + // issues the UMMA-gated smem-buffer arrival. The Blackwell launcher must + // build it with num_consumers=1 (one umma_arrive per stage releases the + // buffer). Params/PipelineState/SharedStorage alias the Hopper + // PipelineTmaAsync's, so producer + Compute=100 consumer drive it unchanged. + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Double-buffered TMEM accumulators: MMA for the next (cell, walk) can fill + // the alternate stage while warps 4..11 drain/reduce-add the previous stage. + static constexpr int AccStages = 2; + static_assert(AccStages * TileN <= 512, + "SM100 MLP3 accumulator stages must fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + // SM100 UMMA TiledMMA — one 1SM tcgen05 atom spanning (TileM, TileN) with + // the accumulator in TMEM. mlp3 is the FIRST ported kernel with BOTH + // operands MN-major: A = dY^T (M=H contiguous) and B = Z (N=I contiguous), + // exactly matching SmemLayoutDYT / SmemLayoutZ (both Layout_MN_SW128_Atom + + // Step<_2,_1,_3>). make_umma_desc reads the MN-contiguous stride + // from each smem layout, so no transpose is needed (cf. mlp2_t's ). + using TiledMmaUmma = typename Mlp3UmmaTiledMmaSelector< + Element, ElementAccum, TileM, TileN>::type; + + static constexpr int TmaTransBytesDYT = + static_cast(size(SmemLayoutDYT_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesZ = + static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); + // Fused pipeline: DYT + Z share one mbarrier per stage. Both WGs + // consume the same Z via cooperative TiledMma slicing. + static constexpr int TmaTransBytes = TmaTransBytesDYT + TmaTransBytesZ; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; // 256 + static constexpr int NumThreads = 384; + + // Per-WG store buffer slot: (WgTileM, EpiChunkN). One slot per WG; it + // holds kAtomsPerWg = WgTileM/AtomTileM contiguous atom-rows after the + // acc→smem fold. + static constexpr int kAtomsPerWg = WgTileM / AtomTileM; // 2 + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + // TMA-reduce descriptor box = ONE wgmma atom-row (AtomTileM × EpiChunkN). + // At TileM=256 each WG owns 2 *interleaved* atoms (rows {0-63,128-191} / + // {64-127,192-255}) that map to non-contiguous gmem rows, so a WG issues + // kAtomsPerWg separate TMA stores per epi round (mirrors mlp1_fused.cuh, + // which uses an AtomTileM box but only ever has 1 atom/WG). + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; +}; + +template +using Mlp3MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Shared Memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp3Smem { + using Element = typename Traits::Element; + + static constexpr int smem_DYT_size = cosize_v; + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_DYT[smem_DYT_size]; + alignas(128) Element smem_Z[smem_Z_size]; + // Per-WG store buffer: WG1 uses [0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + typename Traits::MainloopPipeline::SharedStorage pipe_storage; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these; they + // keep the consumer signature identical across Compute values. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP3 shared memory (data only — pipe storage lives in the +// outer MlpFusedBwdSmem union; otherwise identical to Mlp3Smem). +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp3FusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_DYT_size = cosize_v; + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_DYT[smem_DYT_size]; + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element store_buf[2 * smem_store_size]; + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } +}; + + +// ═══════════════════════════════════════════════════════════════════ +// Pipe construction helper +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors mlp2/mlp5's make_pipe pattern: hides the role/transaction- +// byte setup so callers (mlp3_fwd, or any fused parent kernel) get a +// consistent Pipeline. The plain SM90 helper uses warp 0 as TMA producer and +// warps 4-11 as WGMMA consumers; the SM100 helper additionally enrolls warp 3 +// as the UMMA consumer of the TMA pipeline. + +template +__device__ __forceinline__ typename Traits::MainloopPipeline +mlp3_make_pipe(typename Traits::MainloopPipeline::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp3_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. The ctor arg pattern is identical to the Hopper pipe, so the +// arch-agnostic mlp3_producer drives it unchanged. +template +__device__ __forceinline__ typename Traits::MainloopPipelineUmma +mlp3_make_pipe_umma(typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Chunk-fixed cooperative producer/consumer +// ═══════════════════════════════════════════════════════════════════ +// +// Used by the standalone mlp3 kernel (mlp3.cu) and mlp_bwd Phase 2. +// Each CTA owns ONE (chunk, lane) cell. We hold the SHARED operand and +// walk the cooperation (split) axis: +// kMSplit=true (M-split): chunk = (e, n_tile), lane partitions +// num_m_tiles, walks m_tile. Z(n_tile, *) — the shared +// N-side operand — is the L2-held chunk-constant. +// kMSplit=false (N-split): chunk = (e, m_tile), lane partitions +// num_n_tiles, walks n_tile. dY^T(m_tile, *) — the +// shared M-side operand — is L2-held. +// +// Cell decode (chunk-major: cell_idx = chunk_idx × outer_split + lane): +// total_cells = total_chunks × outer_split +// cell → chunk_idx = cell_idx / outer_split +// lane = cell_idx % outer_split + +template +struct Mlp3ProducerImpl; + +template <> +struct Mlp3ProducerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadDYT const& tma_load_dyt, + TmaLoadZ const& tma_load_z, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + // Staging-ring depth in K-blocks. > 0 means batch_kb_*/the per-expert + // k-ranges are UNWRAPPED and the gmem K-tile coordinate wraps by % ring_kb + // (grouped grouped Phase-2 may straddle the ring). 0 = no wrap. + int ring_kb = 0) { + + auto sDYT = make_tensor(make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + + auto mDYT = tma_load_dyt.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(num_tokens))); + auto mZT = tma_load_z.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + + auto cta_tma_dyt = tma_load_dyt.get_slice(Int<0>{}); + auto cta_tma_z = tma_load_z.get_slice(Int<0>{}); + + auto tDYTsDYT = cta_tma_dyt.partition_D(sDYT); + auto tZsZ = cta_tma_z.partition_D(sZ); + + int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); + // K-split: each (chunk, walk-lane) is replicated k_split times; the + // k_slice handles a 1/k_split sub-range of the expert's K-loop. Partial + // reductions accumulate into dA via the TMA_REDUCE_ADD epilogue. + int total_cells = total_chunks * outer_split * k_split; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e, m_chunk, n_chunk, walk_begin, walk_end; + if constexpr (Traits::kMSplit) { + // M-split: hold Z(n) shared. chunk = (e, n_tile), walks m_tile. + e = chunk_idx / num_n_tiles; + n_chunk = chunk_idx - e * num_n_tiles; + m_chunk = 0; + int m_per = num_m_tiles / outer_split; + walk_begin = lane * m_per; + walk_end = walk_begin + m_per; + } else { + // N-split: hold dY^T(m) shared. chunk = (e, m_tile), walks n_tile. + e = chunk_idx / num_m_tiles; + m_chunk = chunk_idx - e * num_m_tiles; + n_chunk = 0; + int n_per = num_n_tiles / outer_split; + walk_begin = lane * n_per; + walk_end = walk_begin + n_per; + } + + // Intersect this expert's global K range with the current + // batch's K-window. Global k_starts/ends were filled once per + // pass; per-batch we just clip. + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) continue; + + // K-split: clip to this cell's k_slice sub-range of [kb_lo, kb_hi). + { + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int s_lo = kb_lo + k_slice * k_per; + int s_hi = min(s_lo + k_per, kb_hi); + if (s_hi <= s_lo) continue; // empty slice (e.g. straggler) + kb_lo = s_lo; + kb_hi = s_hi; + } + + for (int w = walk_begin; w < walk_end; ++w) { + int m_tile = Traits::kMSplit ? w : m_chunk; + int n_tile = Traits::kMSplit ? n_chunk : w; + auto gDYT = local_tile(mDYT, + make_tile(Int{}, Int{}), + make_coord(m_tile, _)); + auto gZ = local_tile(mZT, + make_tile(Int{}, Int{}), + make_coord(n_tile, _)); + auto tDYTgDYT = cta_tma_dyt.partition_S(gDYT); + auto tZgZ = cta_tma_z.partition_S(gZ); + + for (int kb = kb_lo; kb < kb_hi; ++kb) { + int kb_g = (ring_kb > 0) ? (kb % ring_kb) : kb; // physical ring slot + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_dyt.with(*bar, 0), + tDYTgDYT(_, _, _, kb_g), tDYTsDYT(_, _, _, state.index())); + copy(tma_load_z.with(*bar, 0), + tZgZ(_, _, _, kb_g), tZsZ(_, _, _, state.index())); + } + ++state; + } + } + } +} +}; // Mlp3ProducerImpl<90> + +template +__device__ __forceinline__ void mlp3_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadDYT const& tma_load_dyt, + TmaLoadZ const& tma_load_z, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + int ring_kb = 0) { + Mlp3ProducerImpl::template run( + pipe, state, smem, tma_load_dyt, tma_load_z, + expert_k_starts, expert_k_ends, num_experts, + hidden_dim, intermediate_dim, num_tokens, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split, ring_kb); +} + +// ═══════════════════════════════════════════════════════════════════ +// Consumer — arch-split into Compute=90 (WGMMA) / Compute=100 (UMMA) +// ═══════════════════════════════════════════════════════════════════ +// Compute=90 → Hopper / WGMMA. Both WGs cooperate in ONE TiledMMA that +// partitions (TileM,TileN); each thread holds a register acc +// slice, scatters into store_buf, and the WG leader issues +// SM90_TMA_REDUCE_ADD. +// Compute=100 → Blackwell / UMMA. One 1SM tcgen05 atom (warp 3) fills ONE +// TMEM accumulator over the whole (TileM,TileN); both WGs pull +// their N-half from TMEM into the SAME store_buf slots, then the +// IDENTICAL REDUCE_ADD leader block stores them. + +template +struct Mlp3ConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper WGMMA. Verbatim from the original mlp3_consumer. +// ─────────────────────────────────────────────────────────────────── +template <> +struct Mlp3ConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDA const& tma_reduce_da, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int total_n_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split) { + + (void)Compute; + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + const int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sDYT = make_tensor(make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); + auto sZ = make_tensor(make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + + auto tCsDYT = thr_mma.partition_A(sDYT); + auto tCsZ = thr_mma.partition_B(sZ); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + // Identity coords for per-thread acc → store_buf mapping. + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + constexpr int store_slot_elems = Traits::WgTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + const bool is_my_wg_leader = (tid_in_wg == 0); + + auto mdA = [&]() { + if constexpr (Expert3D) { + return tma_reduce_da.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(num_experts))); + } else { + return tma_reduce_da.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(intermediate_dim))); + } + }(); + auto cta_tma_da = tma_reduce_da.get_slice(Int<0>{}); + + int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); + int total_cells = total_chunks * outer_split * k_split; + + constexpr int K_PIPE_MMAS = 1; + + // Carries across walk iters AND across cells for cross-tile K-loop / + // store overlap. Drained once at the end. + bool store_in_flight = false; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e, m_chunk, n_chunk, walk_begin, walk_end; + if constexpr (Traits::kMSplit) { + // M-split: hold Z(n) shared. chunk = (e, n_tile), walks m_tile. + e = chunk_idx / num_n_tiles; + n_chunk = chunk_idx - e * num_n_tiles; + m_chunk = 0; + int m_per = num_m_tiles / outer_split; + walk_begin = lane * m_per; + walk_end = walk_begin + m_per; + } else { + // N-split: hold dY^T(m) shared. chunk = (e, m_tile), walks n_tile. + e = chunk_idx / num_m_tiles; + m_chunk = chunk_idx - e * num_m_tiles; + n_chunk = 0; + int n_per = num_n_tiles / outer_split; + walk_begin = lane * n_per; + walk_end = walk_begin + n_per; + } + + // Intersect this expert's global K range with the current + // batch's K-window. + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) continue; + + // K-split: clip to this cell's k_slice sub-range (must match producer). + { + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int s_lo = kb_lo + k_slice * k_per; + int s_hi = min(s_lo + k_per, kb_hi); + if (s_hi <= s_lo) continue; // empty slice + kb_lo = s_lo; + kb_hi = s_hi; + } + + int k_count = kb_hi - kb_lo; + + for (int w = walk_begin; w < walk_end; ++w) { + int m_tile = Traits::kMSplit ? w : m_chunk; + int n_tile = Traits::kMSplit ? n_chunk : w; + + clear(acc); + auto state_release = state; + int prologue_count = (k_count < K_PIPE_MMAS) ? k_count : K_PIPE_MMAS; + + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsDYT(_, _, _, state.index()), + tCsZ(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + + for (int k = prologue_count; k < k_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsDYT(_, _, _, state.index()), + tCsZ(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + + ++state; + ++state_release; + } + + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // Each WG writes its (WgTileM × EpiChunkN) slot to store_buf + // and issues TMA stores. NumEpiRounds rounds per WG. + // kMSplit=true : M-split. WG_w owns WgTileM contiguous + // rows of TileM. With WgLayout_M=2 and + // atom_count_M = WgTileM/AtomTileM, the + // natural CUTE m-atom positions interleave + // between WGs — formula below maps m_loc + // back to a contiguous local row. + // kMSplit=false : N-split. WG_w owns full TileM, cols + // [w*WgTileN, (w+1)*WgTileN). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + // CUTE atom-iter interleaves WGs in M direction + // when atom_count_M > 1 (TileM=256). Unified + // formula folds 64-row WG strips into a + // contiguous WgTileM-row local range: + // m_atom_global = m_loc / 64 (atom row idx) + // m_atom_in_wg = m_atom_global / 2 + // m_local = m_atom_in_wg*64 + (m_loc % 64) + // = (m_loc / 128)*64 + (m_loc % 64) + // For tile_count_M=1 (TileM=128) the / 128 term + // is 0 so m_local = m_loc % 64. + m_local = (m_loc / 128) * 64 + (m_loc % 64); + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int da_m = Expert3D ? m_tile : e * num_m_tiles + m_tile; + // Store box = ONE atom-row (AtomTileM). Each WG issues + // kAtomsPerWg stores, one per atom it owns. The store_buf + // holds the WG's atoms in contiguous local rows + // [a*AtomTileM, +AtomTileM); their gmem atom-rows differ: + // M-split: the 2 atoms interleave → 4*da_m + my_wg + 2a + // (gmem M-grid is TileM/AtomTileM rows per tile). + // N-split: WG owns the full TileM (2 contiguous atoms) and + // its own N-half → 2*da_m + a, n offset by my_wg. + int n_tile_idx; + if constexpr (Traits::kMSplit) { + n_tile_idx = n_tile * Traits::NumEpiRounds + r; + } else { + n_tile_idx = n_tile * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + constexpr int kRowAtoms = Traits::TileM / Traits::AtomTileM; // 4 (M-split) / 2 (N-split) + CUTE_UNROLL + for (int a = 0; a < Traits::kAtomsPerWg; ++a) { + int m_atom_row = Traits::kMSplit + ? (kRowAtoms * da_m + my_wg + Traits::kAtomsPerWg * a) + : (kRowAtoms * da_m + a); + auto sStore_a = local_tile(sStore, + make_tile(Int{}, Int{}), + make_coord(a, 0)); + auto gdA = [&]() { + if constexpr (Expert3D) { + return local_tile(mdA, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx, e)); + } else { + return local_tile(mdA, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx)); + } + }(); + copy(tma_reduce_da, cta_tma_da.partition_S(sStore_a), + cta_tma_da.partition_D(gdA)); + } + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp3ConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — Compute defaults to 90 so existing call sites +// (mlp3_consumer(...)) are unchanged; pass mlp3_consumer +// for the Blackwell path. +// ─────────────────────────────────────────────────────────────────── +template +__device__ __forceinline__ void mlp3_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDA const& tma_reduce_da, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int total_n_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, + uint32_t* pair_init_phase = nullptr) { + if constexpr (Compute == 100) { + Mlp3ConsumerImpl::template run( + pipe, state, smem, tma_reduce_da, + expert_k_starts, expert_k_ends, num_experts, + hidden_dim, intermediate_dim, total_n_rows, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split, + pair_init_barrier, pair_init_phase); + } else { + Mlp3ConsumerImpl::template run( + pipe, state, smem, tma_reduce_da, + expert_k_starts, expert_k_ends, num_experts, + hidden_dim, intermediate_dim, total_n_rows, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split); + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Host launcher (raw-pointer; defined in mlp3.cu) +// ═══════════════════════════════════════════════════════════════════ + +// dA = dY^T @ Z, accumulated atomically into the [E, H, I] output via +// TMA_REDUCE_ADD. dA MUST be zero-initialized by the caller. +// +// dY : [num_tokens, hidden_dim] +// Z : [num_tokens, intermediate_dim] +// dA : [num_experts, hidden_dim, intermediate_dim] +// expert_for_k_block: [num_tokens / TileK] int32, TileK=64 +void mlp3_fwd_bf16_launch( + const cutlass::bfloat16_t* dY, + const cutlass::bfloat16_t* Z, + cutlass::bfloat16_t* dA, + const int* expert_for_k_block, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int device, + cudaStream_t stream); + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM (paired-CTA) MLP3 — basic ClusterM=2 path +// ═══════════════════════════════════════════════════════════════════ +// +// Joins two adjacent CTAs into one cta_group::2 UMMA atom covering a +// 256x256x64 joined tile, doubling the M extent a single CTA-pair can +// produce per wave versus the 1SM path above. Callers must launch with +// cudaLaunchKernelEx, clusterDim=(2,1,1), an even gridDim.x, and +// pair-aware SM100 2SM TMA load descriptors for dY^T / Z. The dA +// descriptor remains an ordinary (non-paired) TMA_REDUCE_ADD. +// +// TileK, Stages, and EpiChunkN stay trait parameters — the validated +// production candidate is TileK=64/Stages=5/EpiChunkN=64 (see +// 2sm.md / 2sm_stage.md), but that tuple is a caller-side choice, not +// hardcoded here. ClusterM is fixed at 2 (one paired-CTA "pair"): the +// ClusterM=4 (two-pair, N-split-across-pairs) and shared-operand TMA +// multicast variants explored alongside this trait did not clear the +// required performance bar (see the mlp3_multicast bench results) and +// are intentionally not ported, nor are the CompactEpilogue and +// EarlyTmemRelease epilogue experiments (see mlp3_s6_bounce / +// mlp3_s5_early_release) — both were rejected and are omitted here. +// +// Host-launcher contract (no launcher lives in this header — these are +// hard requirements on any .cu translation unit that selects Compute=100, +// e.g. an mlp_bwd.cuh fused entry point): +// 1. Grid: gridDim.x must be even (mlp3_fwd __traps if it is not); +// launch must use cudaLaunchKernelEx with cudaLaunchConfig_t:: +// cudaLaunchAttributeClusterDimension = {ClusterM, 1, 1}. +// 2. The kernel function must first be registered via +// cudaFuncSetAttribute(kernel, cudaFuncAttributeNonPortableClusterSizeAllowed, 1) +// before the cudaLaunchKernelEx call — the driver rejects a 2-CTA +// cluster launch of a kernel that exceeds the "portable" static +// shared-memory budget without this opt-in. +// 3. The dY^T / Z TMA operands passed in as TmaLoadDYT / TmaLoadZ MUST +// be constructed with cute::make_tma_copy_A_sm100 / +// cute::make_tma_copy_B_sm100 (NOT the ordinary 1SM make_tma_copy) +// — only the _sm100 constructors produce the cta_group::2, +// pair-aware copy atoms that create_tma_multicast_mask / +// tma_partition below assume. The dA reduce-add operand stays an +// ordinary (non-paired) TMA_REDUCE_ADD, per above. +// 4. Dynamic shared memory for Mlp3Smem2Sm / Mlp3FusedSmem2Sm +// must be sized and opted in at runtime by the launcher — +// query cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemory +// PerBlockOptin, device), then cudaFuncSetAttribute(kernel, +// cudaFuncAttributeMaxDynamicSharedMemorySize, requested) with +// requested <= optin. This header intentionally has no compile- +// time static_assert on a fixed smem byte ceiling (only the TMEM +// column assert on AccStages*TileN below) — the ceiling is a +// runtime, device-dependent quantity and must not be hardcoded. +template < + typename Element_, + int TileM_ = 256, + int TileN_ = 256, + int TileK_ = 64, + int Stages_ = 3, + int EpiChunkN_ = 64, + int AccStages_ = 2, + int ClusterM_ = 2> +struct Mlp3Traits2Sm { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + static constexpr int AccStages = AccStages_; + static constexpr int ClusterM = ClusterM_; + static constexpr int NumPairs = ClusterM / 2; + static constexpr int CtaTileM = TileM / 2; + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = TileN / 2; + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + static constexpr int kAtomsPerCta = CtaTileM / AtomTileM; + + static_assert(TileM == 256, + "MLP3 2SM requires joined TileM=256; joined TileM=128 is invalid"); + static_assert(TileN == 128 || TileN == 256, + "MLP3 2SM supports TileN=128 or TileN=256"); + static_assert(TileK == 32 || TileK == 64, + "MLP3 2SM pilot supports TileK=32 or 64"); + static_assert(Stages >= 2 && Stages <= 12, + "MLP3 2SM supports two to twelve mainloop stages"); + static_assert(CtaTileM == 128, "each peer CTA must own 128 M rows"); + static_assert(WgTileN % EpiChunkN == 0, + "EpiChunkN must divide each warpgroup's N half"); + static_assert(kAtomsPerCta == 2, + "each CTA must reduce-add two 64-row atoms"); + static_assert(AccStages >= 2, "MLP3 2SM requires accumulator double buffering"); + static_assert(AccStages * TileN <= 512, + "MLP3 2SM accumulator stages must fit in 512 TMEM columns"); + static_assert(ClusterM == 2, + "MLP3 2SM only ports the validated single-pair path; ClusterM=4 " + "(two-pair N-split, evaluated together with shared-operand " + "multicast) was rejected -- see the mlp3_multicast bench results"); + + using TileShape = Shape, Int, Int>; + using ClusterShape = Shape, _1, _1>; + using AtomThrShape = Shape<_2, _1, _1>; + + using TiledMma2Sm = decltype(make_tiled_mma( + SM100_MMA_F16BF16_2x1SM_SS< + Element, Element, ElementAccum, TileM, TileN, + UMMA::Major::MN, UMMA::Major::MN>{})); + static_assert(size(typename TiledMma2Sm::AtomThrID{}) == 2, + "MLP3 2SM requires a two-CTA MMA atom"); + + using ClusterLayoutVMNK = decltype(tiled_divide( + make_layout(ClusterShape{}), + make_tile(typename TiledMma2Sm::AtomThrID{}))); + + using MmaShapeA_MK = decltype(partition_shape_A( + TiledMma2Sm{}, make_shape(Int{}, Int{}))); + using MmaShapeB_NK = decltype(partition_shape_B( + TiledMma2Sm{}, make_shape(Int{}, Int{}))); + + using SmemLayoutAtom = UMMA::Layout_MN_SW128_Atom; + using SmemLayoutDYT = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtom{}, append(MmaShapeA_MK{}, Int{}), + Step<_2, _1, _3>{})); + using SmemLayoutZ = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtom{}, append(MmaShapeB_NK{}, Int{}), + Step<_2, _1, _3>{})); + using SmemLayoutDYT_1 = + decltype(SmemLayoutDYT{}(_, _, _, Int<0>{})); + using SmemLayoutZ_1 = + decltype(SmemLayoutZ{}(_, _, _, Int<0>{})); + + static constexpr int TmaTransBytesDYT = + static_cast(cosize_v * sizeof(Element)); + static constexpr int TmaTransBytesZ = + static_cast(cosize_v * sizeof(Element)); + static constexpr int TmaTransBytes = + 2 * (TmaTransBytesDYT + TmaTransBytesZ); + + using MainloopPipelineUmma2Sm = cutlass::PipelineTmaUmmaAsync< + Stages, ClusterShape, AtomThrShape>; + using MainloopPipeline = MainloopPipelineUmma2Sm; + using MainloopPipelineUmma = MainloopPipelineUmma2Sm; + using PipelineState = typename MainloopPipelineUmma2Sm::PipelineState; + using AccumulatorPipeline2Sm = cutlass::PipelineUmmaAsync< + AccStages, AtomThrShape>; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; + + // One (CtaTileM, EpiChunkN) store slot per epilogue warpgroup; WG0 + // uses [0..S-1], WG1 uses [S..2S-1] (S = smem_store_size, see + // Mlp3Smem2Sm below). Mirrors the 1SM Mlp3Traits::SmemLayoutStoreSlot, + // sized CtaTileM (=128) instead of WgTileM. + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; +}; + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM Shared Memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp3Smem2Sm { + using Element = typename Traits::Element; + + static constexpr int smem_DYT_size = + cosize_v; + static constexpr int smem_Z_size = + cosize_v; + static constexpr int smem_store_size = + cosize_v; + + alignas(128) Element smem_DYT[smem_DYT_size]; + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element store_buf[2 * smem_store_size]; + alignas(16) typename Traits::MainloopPipelineUmma2Sm::SharedStorage pipe_storage; + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline2Sm::SharedStorage acc_pipe; + + CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } +}; + +// Fused MLP3 2SM shared memory (data only — the mainloop pipe's storage +// lives in the outer fused-kernel smem union; otherwise identical to +// Mlp3Smem2Sm). Mirrors the 1SM Mlp3FusedSmem / Mlp3Smem split. +template +struct Mlp3FusedSmem2Sm { + using Element = typename Traits::Element; + + static constexpr int smem_DYT_size = + cosize_v; + static constexpr int smem_Z_size = + cosize_v; + static constexpr int smem_store_size = + cosize_v; + + alignas(128) Element smem_DYT[smem_DYT_size]; + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element store_buf[2 * smem_store_size]; + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline2Sm::SharedStorage acc_pipe; + + CUTE_DEVICE Element* DYT_data() { return &smem_DYT[0]; } + CUTE_DEVICE Element* Z_data() { return &smem_Z[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM pipe construction helper +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors mlp3_make_pipe's role/transaction-byte setup, adjusted for the +// paired-CTA cluster: warp 0 is the TMA producer (leader CTA elects the +// arrival), warps 3-11 are UMMA/epilogue consumers. +template +__device__ __forceinline__ typename Traits::MainloopPipelineUmma2Sm +mlp3_make_pipe_umma_2sm( + typename Traits::MainloopPipelineUmma2Sm::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipelineUmma2Sm; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = warp_id == 0; + bool is_consumer = warp_id >= 3 && warp_id <= 11; + + typename Pipeline::Params params; + params.transaction_bytes = Traits::TmaTransBytes; + params.num_producers = 1; + params.num_consumers = 1; + params.initializing_warp = 0; + if (is_producer) { + params.role = Category::Producer; + params.is_leader = + threadIdx.x == 0 && + cute::block_rank_in_cluster() % 2 == 0; + } else if (is_consumer) { + params.role = Category::Consumer; + } else { + params.role = Category::NonParticipant; + } + return Pipeline( + storage, params, typename Traits::ClusterShape{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM producer / consumer +// ═══════════════════════════════════════════════════════════════════ +// +// Chunk-fixed walk, same principle as mlp3_producer/mlp3_consumer above: +// each CTA-pair owns ONE (chunk, lane) cell and walks m_tile × kb +// internally. Within a pair, the leader CTA (pair_rank==0) drives the +// paired UMMA and owns the accumulator producer role; both CTAs run the +// epilogue over their own CtaTileM=128 half of the joined 256-row tile. +template <> +struct Mlp3ProducerImpl<100> { +template < + typename Traits, + typename Pipeline, + typename SmemType, + typename TmaLoadDYT, + typename TmaLoadZ> +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadDYT const& tma_load_dyt, + TmaLoadZ const& tma_load_z, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + int ring_kb = 0) { + // Compute!=100 must stay 1SM: this body issues SM100-only paired-CTA + // TMA ops (pair-aware tma_partition / create_tma_multicast_mask on a + // TiledMma2Sm/MainloopPipelineUmma2Sm operand) and must never be live + // code for an sm_90a compilation pass of this same translation unit. +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + auto sDYT = make_tensor( + make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); + auto sZ = make_tensor( + make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + + auto mDYT = tma_load_dyt.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(num_tokens))); + auto mZT = tma_load_z.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + + typename Traits::TiledMma2Sm tiled_mma; + int cta_rank = static_cast(cute::block_rank_in_cluster()); + int pair_rank = cta_rank % 2; + auto cta_mma = tiled_mma.get_slice(pair_rank); + auto pair_layout_vmnk = tiled_divide( + make_layout(typename Traits::ClusterShape{}), + make_tile(typename Traits::TiledMma2Sm::AtomThrID{})); + auto pair_coord_vmnk = + pair_layout_vmnk.get_flat_coord(pair_rank); + uint16_t mcast_mask_dyt = + create_tma_multicast_mask<2>(pair_layout_vmnk, pair_coord_vmnk); + uint16_t mcast_mask_z = + create_tma_multicast_mask<1>(pair_layout_vmnk, pair_coord_vmnk); + + int total_chunks = num_experts * num_n_tiles; + int total_cells = total_chunks * outer_split * k_split; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e = chunk_idx / num_n_tiles; + int n_tile = chunk_idx - e * num_n_tiles; + int walk_begin = lane * num_m_tiles / outer_split; + int walk_end = (lane + 1) * num_m_tiles / outer_split; + + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) + continue; + + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int slice_lo = kb_lo + k_slice * k_per; + int slice_hi = min(slice_lo + k_per, kb_hi); + if (slice_hi <= slice_lo) + continue; + kb_lo = slice_lo; + kb_hi = slice_hi; + + for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { + auto coord = make_coord(m_tile, n_tile, _); + auto gDYT = local_tile( + mDYT, typename Traits::TileShape{}, coord, + Step<_1, X, _1>{}); + auto gZ = local_tile( + mZT, typename Traits::TileShape{}, coord, + Step{}); + auto tCgDYT = cta_mma.partition_A(gDYT); + auto tCgZ = cta_mma.partition_B(gZ); + + auto [tDYTgDYT, tDYTsDYT] = tma_partition( + tma_load_dyt, + get<2>(pair_coord_vmnk), + make_layout(size<2>(pair_layout_vmnk)), + group_modes<0, 3>(sDYT), + group_modes<0, 3>(tCgDYT)); + auto [tZgZ, tZsZ] = tma_partition( + tma_load_z, + get<1>(pair_coord_vmnk), + make_layout(size<1>(pair_layout_vmnk)), + group_modes<0, 3>(sZ), + group_modes<0, 3>(tCgZ)); + + for (int kb = kb_lo; kb < kb_hi; ++kb) { + int kb_g = ring_kb > 0 ? kb % ring_kb : kb; + pipe.producer_acquire(state); + if (cute::elect_one_sync()) { + auto* barrier = pipe.producer_get_barrier(state); + copy( + tma_load_dyt.with( + *barrier, mcast_mask_dyt), + tDYTgDYT(_, kb_g), + tDYTsDYT(_, state.index())); + copy( + tma_load_z.with( + *barrier, mcast_mask_z), + tZgZ(_, kb_g), + tZsZ(_, state.index())); + } + ++state; + } + } + } + pipe.producer_tail(state); +#else + __trap(); +#endif +} +}; // Mlp3ProducerImpl<100> + +template <> +struct Mlp3ConsumerImpl<100> { +template < + typename Traits, + bool Expert3D = false, + typename Pipeline, + typename SmemType, + typename TmaReduceAddDA> +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDA const& tma_reduce_da, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int total_n_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + const cutlass::arch::ClusterBarrier* pair_init_barrier, + uint32_t* pair_init_phase) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + (void)hidden_dim; + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int CtaTileM = Traits::CtaTileM; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int NumEpiRounds = Traits::NumEpiRounds; + constexpr int kMmaEpiThreads = + Traits::ConsumerThreads + Traits::WarpSize; + + int cta_rank = static_cast(cute::block_rank_in_cluster()); + int pair_rank = cta_rank % 2; + bool is_leader_cta = pair_rank == 0; + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_mma_warp = warp_id == 3; + bool is_epilogue = warp_id >= 4 && warp_id <= 11; + int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; + int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + int tid_in_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + int wg_barrier_id = 1 + wg; + bool is_wg_leader = is_epilogue && tid_in_wg == 0; + + typename Traits::TiledMma2Sm tiled_mma; + auto cta_mma = tiled_mma.get_slice(pair_rank); + auto sDYT = make_tensor( + make_smem_ptr(smem.DYT_data()), typename Traits::SmemLayoutDYT{}); + auto sZ = make_tensor( + make_smem_ptr(smem.Z_data()), typename Traits::SmemLayoutZ{}); + auto tCrDYT = cta_mma.make_fragment_A(sDYT); + auto tCrZ = cta_mma.make_fragment_B(sZ); + + auto cAccFull = make_identity_tensor( + make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + if (pair_init_barrier != nullptr) { + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + if (warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + pair_init_barrier->arrive(peer, !leader); + pair_init_barrier->wait(*pair_init_phase); + pair_init_barrier->arrive(peer, leader); + *pair_init_phase ^= 1; + } + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + } + + using AccPipe = typename Traits::AccumulatorPipeline2Sm; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp && is_leader_cta + ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; + acc_params.consumer_arv_count = 2; + acc_params.initializing_warp = 4; + AccPipe acc_pipe( + smem.acc_pipe, acc_params, typename Traits::ClusterShape{}); + auto acc_prod_state = + cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + if (pair_init_barrier != nullptr && warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + pair_init_barrier->arrive(peer, !leader); + pair_init_barrier->wait(*pair_init_phase); + pair_init_barrier->arrive(peer, leader); + *pair_init_phase ^= 1; + } + __syncwarp(); + } + if (pair_init_barrier != nullptr) + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + constexpr int store_slot_elems = CtaTileM * EpiChunkN; + Element* store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor( + make_smem_ptr(store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + auto mdA = tma_reduce_da.get_tma_tensor(make_shape( + static_cast(total_n_rows), + static_cast(intermediate_dim))); + auto cta_tma_da = tma_reduce_da.get_slice(Int<0>{}); + + auto epi_tile = + make_tile(Int{}, Int{}); + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi = flat_divide(acc_mn, epi_tile); + auto t2r = make_tmem_copy( + TmemLoadOp{}, + tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tid_in_wg); + auto cChunk = make_identity_tensor( + make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); + + Layout tmem_warp_layout = + typename decltype(make_tmem_warp_partitioner( + tAcc_epi(_, _, _0{}, _0{})))::TiledLayout_TV{}; + constexpr bool predicate_tmem_load = + size(tmem_warp_layout) != cosize(tmem_warp_layout); + + int total_chunks = num_experts * num_n_tiles; + int total_cells = total_chunks * outer_split * k_split; + bool store_in_flight = false; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e = chunk_idx / num_n_tiles; + int n_tile = chunk_idx - e * num_n_tiles; + int walk_begin = lane * num_m_tiles / outer_split; + int walk_end = (lane + 1) * num_m_tiles / outer_split; + + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) + continue; + + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int slice_lo = kb_lo + k_slice * k_per; + int slice_hi = min(slice_lo + k_per, kb_hi); + if (slice_hi <= slice_lo) + continue; + kb_lo = slice_lo; + kb_hi = slice_hi; + + for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { + if (is_mma_warp && is_leader_cta) { + acc_pipe.producer_acquire(acc_prod_state); + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = + tmem_base + uint32_t(acc_stage * TileN); + + bool first = true; + for (int kb = kb_lo; kb < kb_hi; ++kb) { + pipe.consumer_wait(state); + CUTE_UNROLL + for (int ks = 0; ks < size<2>(tCrDYT); ++ks) { + tiled_mma.accumulate_ = first + ? UMMA::ScaleOut::Zero + : UMMA::ScaleOut::One; + first = false; + gemm( + tiled_mma, + tCrDYT(_, _, ks, state.index()), + tCrZ(_, _, ks, state.index()), + tCtAcc); + } + pipe.consumer_release(state); + ++state; + } + acc_pipe.producer_commit(acc_prod_state); + ++acc_prod_state; + } + + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = + tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = + tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = + flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc = + thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int round = 0; + round < NumEpiRounds; + ++round) { + int chunk = wg * NumEpiRounds + round; + auto tAccChunk = + tTR_tAcc(_, _, _, _0{}, chunk); + bool issue_tmem_load = true; + if constexpr (predicate_tmem_load) { + int subpart = + (tAccChunk.data().dp_ / 32) % 4; + issue_tmem_load = + tid_in_wg / Traits::WarpSize == subpart; + } + if (issue_tmem_load) { + copy(t2r, tAccChunk, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + } + + int da_m = e * num_m_tiles + m_tile; + int n_chunk = + n_tile * + (Traits::TileN / Traits::EpiChunkN) + + wg * Traits::NumEpiRounds + round; + + if (store_in_flight) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, + wg_barrier_id); + + if (issue_tmem_load) { + CUTE_UNROLL + for (int i = 0; + i < size(tTR_rAcc); + ++i) { + int m_local = + get<0>(tTR_cChunk(i)); + int n_local = + get<1>(tTR_cChunk(i)); + sStore(m_local, n_local) = + static_cast( + tTR_rAcc(i)); + } + } + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, + wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + CUTE_UNROLL + for (int atom = 0; + atom < Traits::kAtomsPerCta; + ++atom) { + int m_atom = + (Traits::TileM / + Traits::AtomTileM) * + da_m + + Traits::kAtomsPerCta * + pair_rank + + atom; + auto sStoreAtom = + local_tile( + sStore, + make_tile( + Int{}, + Int{}), + make_coord(atom, 0)); + auto gdA = local_tile( + mdA, + make_tile( + Int{}, + Int{}), + make_coord(m_atom, n_chunk)); + copy( + tma_reduce_da, + cta_tma_da.partition_S( + sStoreAtom), + cta_tma_da.partition_D(gdA)); + } + cute::tma_store_arrive(); + } + store_in_flight = true; + } + + // TMEM reads done → release the acc so the next + // (cell, walk) MMA may reuse it (one elected consumer + // thread arrives). Always released after all epilogue + // rounds' register→SMEM fold + TMA_REDUCE_ADD store + // work has been issued for this m_tile — matches the + // 1SM mlp3_consumer's release placement. (The donor's + // EarlyTmemRelease experiment moved this earlier for + // one accumulator stage's worth of extra overlap; it + // was rejected — see mlp3_s5_early_release — so is not + // ported.) + cutlass::arch::NamedBarrier::sync( + Traits::ConsumerThreads, 0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + } + + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); +#else + __trap(); +#endif +} +}; // Mlp3ConsumerImpl<100> + +// ═══════════════════════════════════════════════════════════════════ +// Single-launch driver, chunk-fixed and architecture-dispatched +// ═══════════════════════════════════════════════════════════════════ +// +// Compute=90 uses the original per-CTA Hopper pipeline and cell mapping. +// Compute=100 uses a paired-CTA pipeline, pair-granular cell mapping, a +// shared Allocator2Sm allocation, and cluster synchronization. +template < + typename Traits, + int Compute = 90, + bool Expert3D = false, + typename TmaLoadDYT, + typename TmaLoadZ, + typename TmaReduceAddDA> +__device__ __forceinline__ void mlp3_fwd( + cute::conditional_t< + Compute == 100, + Mlp3Smem2Sm, + Mlp3Smem>& smem, + TmaLoadDYT const& tma_load_dyt, + TmaLoadZ const& tma_load_z, + TmaReduceAddDA const& tma_reduce_da, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int total_n_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int batch_kb_start = 0, + int batch_kb_end = -1, + int k_split = 1, + int ring_kb = 0) { + static_assert(Compute == 90 || Compute == 100, + "mlp3_fwd supports Compute=90 or Compute=100"); + using PipeState = typename Traits::PipelineState; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = warp_id == 0; + constexpr int kFirstConsumerWarp = Compute == 100 ? 3 : 4; + bool is_consumer = + warp_id >= kFirstConsumerWarp && warp_id <= 11; + + cute::prefetch_tma_descriptor( + tma_load_dyt.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_load_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_reduce_da.get_tma_descriptor()); + + auto pipe = [&]() { + if constexpr (Compute == 100) + return mlp3_make_pipe_umma_2sm( + smem.pipe_storage); + else + return mlp3_make_pipe(smem.pipe_storage); + }(); + using Pipeline = decltype(pipe); + + if constexpr (Compute == 100) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + cute::TMEM::Allocator2Sm tmem_allocator; + if (gridDim.x % Traits::ClusterM != 0) + __trap(); + cute::cluster_sync(); + constexpr int kTmemColumns = + Traits::AccStages * Traits::TileN; + if (warp_id == 4) { + tmem_allocator.allocate( + kTmemColumns, &smem.tmem_base); + __syncwarp(); + } + __syncthreads(); + cute::cluster_sync(); +#else + __trap(); +#endif + } else { + __syncthreads(); + } + + PipeState producer_state = is_producer + ? cutlass::make_producer_start_state() + : PipeState{}; + PipeState consumer_state; + + int cell_start = static_cast(blockIdx.x); + int cell_stride = static_cast(gridDim.x); + if constexpr (Compute == 100) { + cell_start /= Traits::ClusterM; + cell_stride /= Traits::ClusterM; + } + int kb_end = batch_kb_end >= 0 + ? batch_kb_end + : num_tokens / Traits::TileK; + + if (is_producer) { + mlp3_producer( + pipe, + producer_state, + smem, + tma_load_dyt, + tma_load_z, + expert_k_starts, + expert_k_ends, + num_experts, + hidden_dim, + intermediate_dim, + num_tokens, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split, + ring_kb); + } else if (is_consumer) { + if constexpr (Compute == 100) { + mlp3_consumer( + pipe, + consumer_state, + smem, + tma_reduce_da, + expert_k_starts, + expert_k_ends, + num_experts, + hidden_dim, + intermediate_dim, + total_n_rows, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split); + } else { +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) + mlp3_consumer( + pipe, + consumer_state, + smem, + tma_reduce_da, + expert_k_starts, + expert_k_ends, + num_experts, + hidden_dim, + intermediate_dim, + total_n_rows, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split); +#else + __trap(); +#endif + } + } + + if constexpr (Compute == 100) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + cute::TMEM::Allocator2Sm tmem_allocator; + __syncthreads(); + cute::cluster_sync(); + constexpr int kTmemColumns = + Traits::AccStages * Traits::TileN; + if (warp_id == 4) { + tmem_allocator.release_allocation_lock(); + tmem_allocator.free( + smem.tmem_base, kTmemColumns); + } +#endif + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp4.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp4.cuh index b14867dd..6928325e 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp4.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp4.cuh @@ -1,1839 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Phase 4: dB = dU^T @ X, dC = dV^T @ X (weight grads, bwd) -// ═══════════════════════════════════════════════════════════════════ -// -// TMA_REDUCE_ADD variant. Cooperative 2-WG consumer, mirroring mlp3.cuh. -// The two WGs cooperate on ONE output tile at a time, so dB and dC are -// produced SEQUENTIALLY (two phases per cell). This lets a SINGLE A-side -// smem buffer be reused for dU^T (phase dB) then dV^T (phase dC) — vs the -// old design's two simultaneous A buffers — so the footprint matches mlp3 -// and fits the 228 KiB cap at 256-wide tiles. -// -// Only two configs supported (mirrors Mlp3Traits): -// (TileM=256, TileN=128) → M-split (Layout<_2,_1,_1>) ← default -// (TileM=128, TileN=256) → N-split (Layout<_1,_2,_1>) -// -// GEMM orientation: dB = dU^T @ X, so C[M=I, N=H], A[M=I,K=T] = dU^T/dV^T, -// B[N=H,K=T] = X. TileM tiles the intermediate (I) dim, TileN the hidden -// (H) dim. Per-WG output footprint = (WgTileM, WgTileN) = (128, 128). -// -// Single fused A + X TMA pipe (one mbarrier per stage, 2 TMA copies). The -// shared operand X is read by both dB and dC phases and by both WGs. -// -// CTA layout (384 threads = 12 warps): -// WG0 warp 0 : Producer (TMA loads for A / X) -// WG0 warps 1-2 : Idle / reserved for fused-kernel communication -// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 -// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 -// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 -// -// Chunk-fixed walk (same principle as mlp3): hold the SHARED operand -// chunk-constant and walk the cooperation (split) axis: -// kMSplit=true (M-split, X shared): chunk = (e, n_tile), walks m_tile; holds X(n,*) -// kMSplit=false (N-split, dU^T/dV^T shared): chunk = (e, m_tile), walks n_tile; holds A(m,*) -// In the M-split default, X(n_tile) is held across BOTH the m-walk and the -// dB/dC phases — re-read for phase dC hits L2. -// ═══════════════════════════════════════════════════════════════════ +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -#include "models.cuh" -#include "tmem_load_op.cuh" - -#include -#include -#include -#include -#include - -// ── SM100 (Compute=100 / Blackwell UMMA) headers ──────────────────── -// Header-only; safe to include under an sm_90a build. The tcgen05 PTX they -// emit is guarded by CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself -// gated on __CUDA_ARCH__ >= 1000) ever instantiates it. Mirrors mlp5.cuh / -// mlp1_fused.cuh. -#include -#include -#include -#include -#include -#include // PipelineTmaUmmaAsync, PipelineUmmaAsync - -// ── SM100 2-SM (paired-CTA cluster) headers ───────────────────────── -// Needed by the paired MLP4 path below: cluster synchronization/rank queries -// for the (2,1,1) cluster and the 2SM TMA copy traits used by tma_partition(). -#include -#include - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// Traits -// ═══════════════════════════════════════════════════════════════════ - -template < - typename Element_, - int TileM_ = 256, - int TileN_ = 128, - int TileK_ = 64, - int Stages_ = 4, - int EpiChunkN_ = 64 -> -struct Mlp4Traits { - using Element = Element_; - using ElementAccum = float; - - static constexpr int TileM = TileM_; - static constexpr int TileN = TileN_; - static constexpr int TileK = TileK_; - static constexpr int Stages = Stages_; - static constexpr int EpiChunkN = EpiChunkN_; - // Cooperative 2-WG layout. Only two configs supported: - // (TileM=256, TileN=128) → kMSplit=true, WgLayout<_2,_1,_1> - // (TileM=128, TileN=256) → kMSplit=false, WgLayout<_1,_2,_1> - // Per-WG output footprint = (WgTileM=128, WgTileN=128) in both. Each - // WG covers 2 wgmma m=64 atoms per K-step (atom_count_M = 2). - static_assert((TileM == 256 && TileN == 128) || - (TileM == 128 && TileN == 256), - "Mlp4Traits supports only (TileM, TileN) = (256, 128) M-split " - "or (128, 256) N-split"); - - static constexpr bool kMSplit = (TileM > TileN); - static constexpr int AtomTileM = 64; - static constexpr int WgTileM = 128; - static constexpr int WgTileN = 128; - - static_assert(WgTileN % EpiChunkN == 0, - "EpiChunkN must divide WgTileN (=128)"); - static constexpr int NumEpiRounds = WgTileN / EpiChunkN; - - static constexpr int AtomN = WgTileN; - using GmmaAtom = typename GmmaSelectorMNMN::Atom; - using WgLayout = cute::conditional_t< - kMSplit, - Layout>, - Layout>>; - using TiledMma = TiledMMA, WgLayout>; - - using SmemLayoutAtom = GMMA::Layout_MN_SW128_Atom; - - // A-side = dU^T / dV^T (M-side, TileM); X = shared B-side (N-side, TileN). - using SmemLayoutA_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, - Step<_2, _1>{})); - using SmemLayoutX_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, - Step<_2, _1>{})); - using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, - Step<_2, _1, _3>{})); - using SmemLayoutX = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, - Step<_2, _1, _3>{})); - - using MainloopPipeline = cutlass::PipelineTmaAsync; - using PipelineState = cutlass::PipelineState; - - // ── SM100 (Compute=100 / Blackwell UMMA) ──────────────────────── - // The mainloop pipeline is the UMMA-aware TMA pipeline (its consumer_release - // issues the UMMA-gated smem-buffer arrival). Params/PipelineState/ - // SharedStorage alias the Hopper PipelineTmaAsync's, so smem.pipe_storage - // (typed MainloopPipeline::SharedStorage) drives BOTH — the standalone - // launcher passes it to mlp4_make_pipe_umma unchanged. num_consumers=1 on - // Blackwell (one umma_arrive per stage releases the buffer). - using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< - Stages, cute::Shape, - cute::Shape>; - // Per-WG UMMA atom: SM100_MMA_F16BF16_SS with M=WgTileM=128, N=WgTileN=128, - // BOTH operands MN-major (A = dU^T/dV^T is M-major in smem; X is N-major). - // Warp 3 issues TWO of these per k-step into two independent - // (128,128) TMEM accumulators (one per consumer WG) — see - // Mlp4ConsumerImpl<100>. UMMA requires M ∈ {64,128}, so the 256-wide tile - // cannot use one atom; the per-WG (128,128) atom is used for both configs. - using TiledMmaUmma = decltype(make_tiled_mma( - SM100_MMA_F16BF16_SS{})); - // Double-buffered TMEM accumulators: each stage owns one (WgTileM,WgTileN) - // accumulator per epilogue WG. The two dB/dC phases still clear - // independently — this only lets MMA for the next (phase,walk) run while - // the epilogue drains the previous accumulator stage. - static constexpr int AccStages = 2; - static_assert(AccStages * (2 * WgTileN) <= 512, - "SM100 MLP4 needs AccStages x two per-WG accumulators to fit in 512 TMEM columns"); - using AccumulatorPipeline = cutlass::PipelineUmmaAsync; - - static constexpr int TmaTransBytesA = - static_cast(size(SmemLayoutA_1{}) * sizeof(Element)); - static constexpr int TmaTransBytesX = - static_cast(size(SmemLayoutX_1{}) * sizeof(Element)); - // Fused pipeline: X + ONE A operand (dU^T or dV^T, per phase) share one - // mbarrier per stage. Both WGs consume the same X via cooperative TiledMma. - static constexpr int TmaTransBytes = TmaTransBytesA + TmaTransBytesX; - - static constexpr int WarpSize = 32; - static constexpr int WarpGroupSize = 128; - static constexpr int NumConsumers = 2; - static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; // 256 - static constexpr int NumThreads = 384; - - // Per-WG store buffer slot: (WgTileM, EpiChunkN). One slot per WG; it - // holds kAtomsPerWg = WgTileM/AtomTileM contiguous atom-rows after the - // acc→smem fold. - static constexpr int kAtomsPerWg = WgTileM / AtomTileM; // 2 - using SmemLayoutStoreSlot = Layout< - Shape, Int>, - Stride, _1>>; - // TMA-reduce descriptor box = ONE wgmma atom-row (AtomTileM × EpiChunkN). - // At TileM=256 each WG owns 2 *interleaved* atoms that map to non-contiguous - // gmem rows, so a WG issues kAtomsPerWg separate TMA stores per epi round - // (mirrors mlp1_fused.cuh / mlp3.cuh). - using SmemLayoutStore = Layout< - Shape, Int>, - Stride, _1>>; -}; - -template -using Mlp4MainloopPipelineFor = cute::conditional_t< - Compute == 100, - typename Traits::MainloopPipelineUmma, - typename Traits::MainloopPipeline>; - -// ═══════════════════════════════════════════════════════════════════ -// Shared Memory -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp4Smem { - using Element = typename Traits::Element; - - static constexpr int smem_X_size = cosize_v; - static constexpr int smem_A_size = cosize_v; - static constexpr int smem_store_size = cosize_v; - - // Single A buffer, reused for dU^T (phase dB) then dV^T (phase dC). - alignas(128) Element smem_X[smem_X_size]; - alignas(128) Element smem_A[smem_A_size]; - // Per-WG store buffer: WG1 uses [0..S-1], WG2 uses [S..2S-1]. - alignas(128) Element store_buf[2 * smem_store_size]; - - typename Traits::MainloopPipeline::SharedStorage pipe_storage; - - // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM - // base address, plus the accumulator pipeline's barrier storage (UMMA→ - // epilogue handoff). The Hopper (Compute=90) path never touches these; - // SM100 launchers allocate/free TMEM outside the consumer and publish the - // base here; the accumulator pipeline storage still lives in this smem. - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; - - CUTE_DEVICE Element* X_data() { return &smem_X[0]; } - CUTE_DEVICE Element* A_data() { return &smem_A[0]; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// Fused MLP4 shared memory (data only — pipe storage lives in the -// outer MlpFusedBwdSmem union; otherwise identical to Mlp4Smem). -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp4FusedSmem { - using Element = typename Traits::Element; - - static constexpr int smem_X_size = cosize_v; - static constexpr int smem_A_size = cosize_v; - static constexpr int smem_store_size = cosize_v; - - alignas(128) Element smem_X[smem_X_size]; - alignas(128) Element smem_A[smem_A_size]; - alignas(128) Element store_buf[2 * smem_store_size]; - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; - - CUTE_DEVICE Element* X_data() { return &smem_X[0]; } - CUTE_DEVICE Element* A_data() { return &smem_A[0]; } -}; - - -// ═══════════════════════════════════════════════════════════════════ -// Pipe construction helper (mirrors mlp3_make_pipe) -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ typename Traits::MainloopPipeline -mlp4_make_pipe(typename Traits::MainloopPipeline::SharedStorage& storage) { - using Pipeline = typename Traits::MainloopPipeline; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 4 && warp_id <= 11); - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = Traits::ConsumerThreads; - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// SM100 (Compute=100) mainloop pipe. Mirrors mlp4_make_pipe but builds the -// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues -// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer -// is released by a single umma_arrive per stage (the MMA warp), not per -// consumer thread. Params/PipelineState/SharedStorage alias the Hopper -// PipelineTmaAsync's, and the ctor arg pattern is identical, so the (shared) -// mlp4_producer and the Compute=100 consumer drive it unchanged. -template -__device__ __forceinline__ typename Traits::MainloopPipelineUmma -mlp4_make_pipe_umma(typename Traits::MainloopPipelineUmma::SharedStorage& storage) { - using Pipeline = typename Traits::MainloopPipelineUmma; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// ═══════════════════════════════════════════════════════════════════ -// Chunk-fixed cooperative producer/consumer — sequential dB / dC -// ═══════════════════════════════════════════════════════════════════ -// -// Used by the standalone mlp4 kernel (mlp4.cu) and mlp_bwd Phase 2. -// Each CTA owns ONE (chunk, lane) cell and processes it in TWO phases -// (phase 0 = dB via dU^T, phase 1 = dC via dV^T). We hold the SHARED -// operand and walk the cooperation (split) axis: -// kMSplit=true (M-split): chunk = (e, n_tile), lane partitions -// num_m_tiles, walks m_tile. X(n_tile, *) — the shared -// N-side operand — is the L2-held chunk-constant, -// reused across the m-walk AND both phases. -// kMSplit=false (N-split): chunk = (e, m_tile), lane partitions -// num_n_tiles, walks n_tile. A(m_tile, *) (dU^T then -// dV^T) is L2-held; X is re-streamed per phase. -// -// Cell decode (chunk-major: cell_idx = chunk_idx × outer_split + lane): -// total_cells = total_chunks × outer_split -// cell → chunk_idx = cell_idx / outer_split -// lane = cell_idx % outer_split - -template -struct Mlp4ProducerImpl; - -template <> -struct Mlp4ProducerImpl<90> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaLoadX const& tma_load_x, - TmaLoadA const& tma_load_dut, - TmaLoadA const& tma_load_dvt, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int intermediate_dim, - int hidden_dim, - int num_tokens, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - // Staging-ring depth in K-blocks; > 0 wraps the gmem K-tile coordinate - // (% ring_kb) for the grouped Phase-2 that may straddle the ring. See - // mlp3_producer. 0 = no wrap. - int ring_kb = 0) { - - auto sA = make_tensor(make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); - auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - - // int64_t casts: H·T, I·T can exceed 2^31 elems at production shapes. - auto mX = tma_load_x.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(num_tokens))); - auto mdUT = tma_load_dut.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(num_tokens))); - auto mdVT = tma_load_dvt.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(num_tokens))); - - auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); - auto cta_tma_dut = tma_load_dut.get_slice(Int<0>{}); - auto cta_tma_dvt = tma_load_dvt.get_slice(Int<0>{}); - - auto tXsX = cta_tma_x.partition_D(sX); - auto tAsA = cta_tma_dut.partition_D(sA); // dU^T/dV^T share A layout - - int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); - // K-split: each (chunk, walk-lane) replicated k_split times; k_slice owns a - // 1/k_split sub-range of the K-loop. Partials sum via TMA_REDUCE_ADD. - int total_cells = total_chunks * outer_split * k_split; - - for (int cell_idx = cell_start; - cell_idx < total_cells; - cell_idx += cell_stride) { - - int k_slice = cell_idx % k_split; - int cell_om = cell_idx / k_split; - int chunk_idx = cell_om / outer_split; - int lane = cell_om - chunk_idx * outer_split; - - int e, m_chunk, n_chunk, walk_begin, walk_end; - if constexpr (Traits::kMSplit) { - // M-split: hold X(n) shared. chunk = (e, n_tile), walks m_tile. - e = chunk_idx / num_n_tiles; - n_chunk = chunk_idx - e * num_n_tiles; - m_chunk = 0; - int m_per = num_m_tiles / outer_split; - walk_begin = lane * m_per; - walk_end = walk_begin + m_per; - } else { - // N-split: hold A(m) shared. chunk = (e, m_tile), walks n_tile. - e = chunk_idx / num_m_tiles; - m_chunk = chunk_idx - e * num_m_tiles; - n_chunk = 0; - int n_per = num_n_tiles / outer_split; - walk_begin = lane * n_per; - walk_end = walk_begin + n_per; - } - - // Intersect this expert's global K range with the current - // batch's K-window. - int kb_lo = max(expert_k_starts[e], batch_kb_start); - int kb_hi = min(expert_k_ends[e], batch_kb_end); - if (kb_hi <= kb_lo) continue; - - // K-split: clip to this cell's k_slice sub-range of [kb_lo, kb_hi). - { - int k_total = kb_hi - kb_lo; - int k_per = (k_total + k_split - 1) / k_split; - int s_lo = kb_lo + k_slice * k_per; - int s_hi = min(s_lo + k_per, kb_hi); - if (s_hi <= s_lo) continue; // empty slice - kb_lo = s_lo; - kb_hi = s_hi; - } - - // Two phases per cell: phase 0 = dB (dU^T), phase 1 = dC (dV^T). - // dU^T/dV^T share descriptor layout; pair the tma object, its slice, - // and its gmem tensor consistently per phase. - CUTE_UNROLL - for (int phase = 0; phase < 2; ++phase) { - auto mA = (phase == 0) ? mdUT : mdVT; - const auto& tma_A = (phase == 0) ? tma_load_dut : tma_load_dvt; - const auto& cta_tma_A = (phase == 0) ? cta_tma_dut : cta_tma_dvt; - - for (int w = walk_begin; w < walk_end; ++w) { - int m_tile = Traits::kMSplit ? w : m_chunk; - int n_tile = Traits::kMSplit ? n_chunk : w; - auto gX = local_tile(mX, - make_tile(Int{}, Int{}), - make_coord(n_tile, _)); - auto gA = local_tile(mA, - make_tile(Int{}, Int{}), - make_coord(m_tile, _)); - auto tXgX = cta_tma_x.partition_S(gX); - auto tAgA = cta_tma_A.partition_S(gA); - - for (int kb = kb_lo; kb < kb_hi; ++kb) { - int kb_g = (ring_kb > 0) ? (kb % ring_kb) : kb; // physical ring slot - pipe.producer_acquire(state); - if (threadIdx.x == 0) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_x.with(*bar, 0), - tXgX(_, _, _, kb_g), tXsX(_, _, _, state.index())); - copy(tma_A.with(*bar, 0), - tAgA(_, _, _, kb_g), tAsA(_, _, _, state.index())); - } - ++state; - } - } - } - } -} -}; // Mlp4ProducerImpl<90> - -template -__device__ __forceinline__ void mlp4_producer( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaLoadX const& tma_load_x, - TmaLoadA const& tma_load_dut, - TmaLoadA const& tma_load_dvt, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int intermediate_dim, - int hidden_dim, - int num_tokens, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - int ring_kb = 0) { - Mlp4ProducerImpl::template run( - pipe, state, smem, tma_load_x, tma_load_dut, tma_load_dvt, - expert_k_starts, expert_k_ends, num_experts, - intermediate_dim, hidden_dim, num_tokens, - num_m_tiles, num_n_tiles, outer_split, - cell_start, cell_stride, - batch_kb_start, batch_kb_end, k_split, ring_kb); -} - -// ═══════════════════════════════════════════════════════════════════ -// Consumer — architecture-specialized on `int Compute` -// ═══════════════════════════════════════════════════════════════════ -// -// Primary template is undefined; one full specialization per supported compute -// capability provides a `run(...)` member (a function-template on the Traits/ -// Pipeline/Smem/Tma types, since function templates can't be partially -// specialized). The free function `mlp4_consumer` below forwards to the right -// specialization — call sites (mlp_bwd.cuh's mlp4_consumer(...)) stay -// unchanged (Compute defaults to 90). -// -// Compute=90 → Hopper / WGMMA (cooperative 2-WG, register acc, M-atom -// interleave + (m_loc/128)*64+(m_loc%64) store remap) -// Compute=100 → Blackwell / UMMA (warp 3 issues 2 per-WG (128,128) UMMAs -// into double-buffered TMEM stages; warps 4-11 run the -// epilogue, with independent per-phase clear and contiguous -// per-WG store — no remap) - -template -struct Mlp4ConsumerImpl; - -// ─────────────────────────────────────────────────────────────────── -// Compute=90 — Hopper / WGMMA. (Verbatim from the original mlp4_consumer.) -// ─────────────────────────────────────────────────────────────────── -template <> -struct Mlp4ConsumerImpl<90> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaReduceAddDB const& tma_reduce_db, - TmaReduceAddDC const& tma_reduce_dc, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int intermediate_dim, - int hidden_dim, - int total_m_rows, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split) { - - (void)Compute; - using Element = typename Traits::Element; - typename Traits::TiledMma tiled_mma; - - const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 - const int my_barrier_id = 1 + my_wg; // 1 or 2 - const int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 - const int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 - auto thr_mma = tiled_mma.get_slice(tid_in_mma); - - auto sA = make_tensor(make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); - auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - - auto tCsA = thr_mma.partition_A(sA); - auto tCsX = thr_mma.partition_B(sX); - - auto acc = partition_fragment_C(tiled_mma, - Shape, Int>{}); - - // Identity coords for per-thread acc → store_buf mapping. - auto cC = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCcC = thr_mma.partition_C(cC); - - constexpr int store_slot_elems = Traits::WgTileM * Traits::EpiChunkN; - Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - const bool is_my_wg_leader = (tid_in_wg == 0); - - auto mdB = [&]() { - if constexpr (Expert3D) { - return tma_reduce_db.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(num_experts))); - } else { - return tma_reduce_db.get_tma_tensor(make_shape( - static_cast(total_m_rows), - static_cast(hidden_dim))); - } - }(); - auto mdC = [&]() { - if constexpr (Expert3D) { - return tma_reduce_dc.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(num_experts))); - } else { - return tma_reduce_dc.get_tma_tensor(make_shape( - static_cast(total_m_rows), - static_cast(hidden_dim))); - } - }(); - auto cta_tma_db = tma_reduce_db.get_slice(Int<0>{}); - auto cta_tma_dc = tma_reduce_dc.get_slice(Int<0>{}); - - int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); - int total_cells = total_chunks * outer_split * k_split; - - constexpr int K_PIPE_MMAS = 1; - - // Carries across phases, walk iters AND cells for cross-tile K-loop / - // store overlap. Drained once at the end. - bool store_in_flight = false; - - for (int cell_idx = cell_start; - cell_idx < total_cells; - cell_idx += cell_stride) { - - int k_slice = cell_idx % k_split; - int cell_om = cell_idx / k_split; - int chunk_idx = cell_om / outer_split; - int lane = cell_om - chunk_idx * outer_split; - - int e, m_chunk, n_chunk, walk_begin, walk_end; - if constexpr (Traits::kMSplit) { - // M-split: hold X(n) shared. chunk = (e, n_tile), walks m_tile. - e = chunk_idx / num_n_tiles; - n_chunk = chunk_idx - e * num_n_tiles; - m_chunk = 0; - int m_per = num_m_tiles / outer_split; - walk_begin = lane * m_per; - walk_end = walk_begin + m_per; - } else { - // N-split: hold A(m) shared. chunk = (e, m_tile), walks n_tile. - e = chunk_idx / num_m_tiles; - m_chunk = chunk_idx - e * num_m_tiles; - n_chunk = 0; - int n_per = num_n_tiles / outer_split; - walk_begin = lane * n_per; - walk_end = walk_begin + n_per; - } - - // Intersect this expert's global K range with the current - // batch's K-window. - int kb_lo = max(expert_k_starts[e], batch_kb_start); - int kb_hi = min(expert_k_ends[e], batch_kb_end); - if (kb_hi <= kb_lo) continue; - - // K-split: clip to this cell's k_slice sub-range (must match producer). - { - int k_total = kb_hi - kb_lo; - int k_per = (k_total + k_split - 1) / k_split; - int s_lo = kb_lo + k_slice * k_per; - int s_hi = min(s_lo + k_per, kb_hi); - if (s_hi <= s_lo) continue; // empty slice - kb_lo = s_lo; - kb_hi = s_hi; - } - - int k_count = kb_hi - kb_lo; - - // Two phases per cell: phase 0 = dB, phase 1 = dC. - CUTE_UNROLL - for (int phase = 0; phase < 2; ++phase) { - - for (int w = walk_begin; w < walk_end; ++w) { - int m_tile = Traits::kMSplit ? w : m_chunk; - int n_tile = Traits::kMSplit ? n_chunk : w; - - clear(acc); - auto state_release = state; - int prologue_count = (k_count < K_PIPE_MMAS) ? k_count : K_PIPE_MMAS; - - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsA(_, _, _, state.index()), - tCsX(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - ++state; - } - - for (int k = prologue_count; k < k_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsA(_, _, _, state.index()), - tCsX(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - - warpgroup_wait(); - warpgroup_fence_operand(acc); - pipe.consumer_release(state_release); - - ++state; - ++state_release; - } - - warpgroup_wait<0>(); - warpgroup_fence_operand(acc); - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_release(state_release); - ++state_release; - } - - // ── Epilogue ────────────────────────────────────── - // Each WG writes its (WgTileM × EpiChunkN) slot to store_buf - // and TMA_REDUCE_ADDs it into dB (phase 0) or dC (phase 1). - // Store indexing mirrors mlp3_consumer (M-split / N-split). - CUTE_UNROLL - for (int r = 0; r < Traits::NumEpiRounds; ++r) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - int chunk_start = r * Traits::EpiChunkN; - CUTE_UNROLL - for (int i = 0; i < size(acc); ++i) { - auto coord = tCcC(i); - int m_loc = get<0>(coord); - int n_loc = get<1>(coord); - int m_local, n_local; - if constexpr (Traits::kMSplit) { - m_local = (m_loc / 128) * 64 + (m_loc % 64); - n_local = n_loc; - } else { - m_local = m_loc; - n_local = n_loc - my_wg * Traits::WgTileN; - } - if (n_local >= chunk_start && - n_local < chunk_start + Traits::EpiChunkN) { - int chunk_n = n_local - chunk_start; - sStore(m_local, chunk_n) = static_cast(acc(i)); - } - } - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - if (is_my_wg_leader) { - cute::tma_store_fence(); - int out_m = Expert3D ? m_tile : e * num_m_tiles + m_tile; - // Store box = ONE atom-row (AtomTileM); kAtomsPerWg stores - // per WG. M-split: the 2 atoms interleave in gmem - // (4*out_m + my_wg + 2a). N-split: WG owns the full TileM - // (2 contiguous atoms) and its own N-half (2*out_m + a). - int n_tile_idx; - if constexpr (Traits::kMSplit) { - n_tile_idx = n_tile * Traits::NumEpiRounds + r; - } else { - n_tile_idx = n_tile * (Traits::TileN / Traits::EpiChunkN) - + my_wg * Traits::NumEpiRounds + r; - } - constexpr int kRowAtoms = Traits::TileM / Traits::AtomTileM; // 4 (M-split) / 2 (N-split) - CUTE_UNROLL - for (int a = 0; a < Traits::kAtomsPerWg; ++a) { - int m_atom_row = Traits::kMSplit - ? (kRowAtoms * out_m + my_wg + Traits::kAtomsPerWg * a) - : (kRowAtoms * out_m + a); - auto sStore_a = local_tile(sStore, - make_tile(Int{}, Int{}), - make_coord(a, 0)); - if (phase == 0) { - auto gdB = [&]() { - if constexpr (Expert3D) { - return local_tile(mdB, - make_tile(Int{}, Int{}), - make_coord(m_atom_row, n_tile_idx, e)); - } else { - return local_tile(mdB, - make_tile(Int{}, Int{}), - make_coord(m_atom_row, n_tile_idx)); - } - }(); - copy(tma_reduce_db, cta_tma_db.partition_S(sStore_a), - cta_tma_db.partition_D(gdB)); - } else { - auto gdC = [&]() { - if constexpr (Expert3D) { - return local_tile(mdC, - make_tile(Int{}, Int{}), - make_coord(m_atom_row, n_tile_idx, e)); - } else { - return local_tile(mdC, - make_tile(Int{}, Int{}), - make_coord(m_atom_row, n_tile_idx)); - } - }(); - copy(tma_reduce_dc, cta_tma_dc.partition_S(sStore_a), - cta_tma_dc.partition_D(gdC)); - } - } - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - } - } - if (store_in_flight) - cute::tma_store_wait<0>(); -} -}; // struct Mlp4ConsumerImpl<90> - -// ─────────────────────────────────────────────────────────────────── -// Free-function forwarder. Existing call sites (mlp4_consumer(...)) -// are unchanged (Compute defaults to 90); pass mlp4_consumer(...) -// to select the Blackwell path. -// ─────────────────────────────────────────────────────────────────── -template -__device__ __forceinline__ void mlp4_consumer( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaReduceAddDB const& tma_reduce_db, - TmaReduceAddDC const& tma_reduce_dc, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int intermediate_dim, - int hidden_dim, - int total_m_rows, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, - uint32_t* pair_init_phase = nullptr) { - if constexpr (Compute == 100) { - Mlp4ConsumerImpl::template run( - pipe, state, smem, tma_reduce_db, tma_reduce_dc, - expert_k_starts, expert_k_ends, num_experts, - intermediate_dim, hidden_dim, total_m_rows, - num_m_tiles, num_n_tiles, outer_split, - cell_start, cell_stride, batch_kb_start, batch_kb_end, k_split, - pair_init_barrier, pair_init_phase); - } else { - Mlp4ConsumerImpl::template run( - pipe, state, smem, tma_reduce_db, tma_reduce_dc, - expert_k_starts, expert_k_ends, num_experts, - intermediate_dim, hidden_dim, total_m_rows, - num_m_tiles, num_n_tiles, outer_split, - cell_start, cell_stride, batch_kb_start, batch_kb_end, k_split); - } -} - -// ═══════════════════════════════════════════════════════════════════ -// SM100 2-SM (paired-CTA) MLP4 — Mlp4Traits2Sm, pipe/producer/consumer -// ═══════════════════════════════════════════════════════════════════ -// -// A cluster-(2,1,1) SM100 UMMA path: two CTAs jointly compute one -// (TileM=256, TileN=256) tile via SM100_MMA_F16BF16_2x1SM_SS, each CTA -// owning CtaTileM=128 M-rows (kAtomsPerCta=2 wgmma-style 64-row atoms). -// Reuses the generic Mlp4Smem / Mlp4FusedSmem above — -// Mlp4Traits2Sm supplies the same nested type names (SmemLayoutA, -// SmemLayoutX, SmemLayoutStoreSlot, MainloopPipeline, AccumulatorPipeline) -// so no separate 2SM smem struct is needed. -// -// Scope: only the BASIC ClusterM=2 (single paired-CTA) configuration -// required by the fused backward kernel is ported here. Intentionally -// EXCLUDED (upstream "topology experiment" scope, not ported): the -// multi-pair (ClusterM=4) path, shared-operand multicast -// (MulticastShared_), N-split, the compact-epilogue variant, and early -// TMEM release — the latter three never existed for MLP4 in the donor -// tree (MLP3-only experiments) and are noted here only for the audit -// record. Stages and EpiChunkN remain free template parameters (no -// S5/E64 production hardcoding); callers pick tuned values at the -// trait-instantiation site (e.g. moe_bwd.cu). -// -// Cross-file invariants a caller/launcher MUST uphold (none of this is -// enforceable from inside this header-only file; documented for the -// next integrator of mlp_bwd.cuh / moe_bwd.cu / mlp4.cu): -// * grid.x must be a MULTIPLE OF ClusterM (=2) before launch; the runtime -// `__trap()` check in mlp4_fwd is only a defensive backstop. -// * Compute != 100 must stay on the 1SM traits/path. Mlp4Traits2Sm's -// TiledMma2Sm / ClusterShape / MainloopPipelineUmma2Sm / -// AccumulatorPipeline2Sm members do not exist on Mlp4Traits, so mixing -// an architecture with the wrong trait remains a compile-time error. -// * Launch must use cudaLaunchKernelEx with a cluster dimension of -// (ClusterM,1,1), and cudaFuncSetAttribute(..., -// cudaFuncAttributeNonPortableClusterSizeAllowed, 1) must be set on -// the kernel beforehand. -// * The dynamic shared-memory allocation for Mlp4Smem> -// must be sized/capped against the device's *dynamic* opt-in limit -// (cudaDeviceGetAttribute(..., cudaDevAttrMaxSharedMemoryPerBlockOptin, -// ...)), not the static 48 KiB default, via -// cudaFuncAttributeMaxDynamicSharedMemorySize. -// * X / dU^T / dV^T TMA descriptors must be built with -// make_tma_copy_A_sm100 / make_tma_copy_B_sm100 (SM100_TMA_2SM_LOAD -// copy op) — the plain SM90 TMA copy-atom builders are not -// pair-aware and would not multicast/replicate operands across the -// two peer CTAs. dB / dC stay TMA REDUCE_ADD outputs. -// * The producer/consumer specializations assume TMEM is already allocated -// and smem.tmem_base populated. The standalone mlp4_fwd driver owns that -// allocation and rendezvouses peer CTAs with cute::cluster_sync() (safe -// because both CTAs run symmetric control flow). A fused caller sharing -// one TMEM region across warp groups with asymmetric control flow must not -// reuse a bare cluster_sync() for that rendezvous; it must pair -// cute::TMEM::Allocator2Sm's -// allocate()/free() with its own dedicated -// cutlass::arch::ClusterBarrier rendezvous scoped to just the -// participating warps, the way donor's mlp_bwd.cuh -// MlpFusedBwdSmem::tmem_pair_barrier does. That smem field and its -// init/arrive/wait call sites live in mlp_bwd.cuh, outside this -// file's scope, and are not present in this target repo yet. - -template < - typename Element_, - int TileM_ = 256, - int TileN_ = 256, - int TileK_ = 64, - int Stages_ = 3, - int EpiChunkN_ = 128, - int AccStages_ = 2, - int ClusterM_ = 2 -> -struct Mlp4Traits2Sm { - using Element = Element_; - using ElementAccum = float; - - static constexpr int TileM = TileM_; - static constexpr int TileN = TileN_; - static constexpr int TileK = TileK_; - static constexpr int Stages = Stages_; - static constexpr int EpiChunkN = EpiChunkN_; - static constexpr int AccStages = AccStages_; - static constexpr int ClusterM = ClusterM_; - static constexpr int NumPairs = ClusterM / 2; - static constexpr int CtaTileM = TileM / 2; - static constexpr int AtomTileM = 64; - static constexpr int WgTileN = TileN / 2; - static constexpr int NumEpiRounds = WgTileN / EpiChunkN; - static constexpr int kAtomsPerCta = CtaTileM / AtomTileM; - - static_assert(TileM == 256, "MLP4 2SM requires joined TileM=256"); - static_assert(TileN == 256, "MLP4 2SM requires joined TileN=256"); - static_assert(TileK == 32 || TileK == 64, - "MLP4 2SM supports TileK=32 or 64"); - static_assert(Stages >= 2 && Stages <= 12, - "MLP4 2SM supports two to twelve mainloop stages"); - static_assert(CtaTileM == 128, "each peer CTA must own 128 M rows"); - static_assert(WgTileN % EpiChunkN == 0, - "EpiChunkN must divide each warpgroup's N half"); - static_assert(kAtomsPerCta == 2, - "each CTA must reduce-add two 64-row atoms"); - static_assert(AccStages == 2, - "MLP4 2SM fixes accumulator double buffering"); - static_assert(AccStages * TileN <= 512, - "MLP4 2SM accumulator stages must fit in 512 TMEM columns"); - static_assert(ClusterM == 2, - "MLP4 2SM ports only the basic single paired-CTA cluster; " - "ClusterM=4 pair-of-pairs multicast is an excluded experiment"); - - using TileShape = Shape, Int, Int>; - using ClusterShape = Shape, _1, _1>; - using AtomThrShape = Shape<_2, _1, _1>; - - using TiledMma2Sm = decltype(make_tiled_mma( - SM100_MMA_F16BF16_2x1SM_SS< - Element, Element, ElementAccum, TileM, TileN, - UMMA::Major::MN, UMMA::Major::MN>{})); - static_assert(size(typename TiledMma2Sm::AtomThrID{}) == 2, - "MLP4 2SM requires a two-CTA MMA atom"); - - using MmaShapeA_MK = decltype(partition_shape_A( - TiledMma2Sm{}, make_shape(Int{}, Int{}))); - using MmaShapeB_NK = decltype(partition_shape_B( - TiledMma2Sm{}, make_shape(Int{}, Int{}))); - - using SmemLayoutAtom = UMMA::Layout_MN_SW128_Atom; - using SmemLayoutA = decltype(UMMA::tile_to_mma_shape( - SmemLayoutAtom{}, append(MmaShapeA_MK{}, Int{}), - Step<_2, _1, _3>{})); - using SmemLayoutX = decltype(UMMA::tile_to_mma_shape( - SmemLayoutAtom{}, append(MmaShapeB_NK{}, Int{}), - Step<_2, _1, _3>{})); - using SmemLayoutA_1 = decltype(SmemLayoutA{}(_, _, _, Int<0>{})); - using SmemLayoutX_1 = decltype(SmemLayoutX{}(_, _, _, Int<0>{})); - - static constexpr int TmaTransBytesA = - static_cast(cosize_v * sizeof(Element)); - static constexpr int TmaTransBytesX = - static_cast(cosize_v * sizeof(Element)); - // Both peer CTAs issue each 2SM load, and both contributions arrive on - // the cluster-scoped full barrier. - static constexpr int TmaTransBytes = - 2 * (TmaTransBytesA + TmaTransBytesX); - - using MainloopPipelineUmma2Sm = cutlass::PipelineTmaUmmaAsync< - Stages, ClusterShape, AtomThrShape>; - // Aliased to the names the generic Mlp4Smem/Mlp4FusedSmem - // (and Mlp4MainloopPipelineFor) look up. - using MainloopPipelineUmma = MainloopPipelineUmma2Sm; - using MainloopPipeline = MainloopPipelineUmma2Sm; - using PipelineState = typename MainloopPipelineUmma2Sm::PipelineState; - using AccumulatorPipeline2Sm = cutlass::PipelineUmmaAsync< - AccStages, AtomThrShape>; - using AccumulatorPipeline = AccumulatorPipeline2Sm; - - static constexpr int WarpSize = 32; - static constexpr int WarpGroupSize = 128; - static constexpr int NumConsumers = 2; - static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; - static constexpr int NumThreads = 384; - - using SmemLayoutStoreSlot = Layout< - Shape, Int>, - Stride, _1>>; - using SmemLayoutStore = Layout< - Shape, Int>, - Stride, _1>>; -}; - -// Compatibility alias: Mlp4Traits2Sm's nested type names line up -// field-for-field with the generic Mlp4Smem above (see the -// comment atop this section), so no bespoke 2SM smem struct is defined — -// this alias just lets callers/tooling written against the name -// "Mlp4Smem2Sm" (e.g. bench_opt/bench_mlp4_2sm.cu) resolve to the -// same generic struct without any source changes on their end. -template -using Mlp4Smem2Sm = Mlp4Smem; - -// SM100 2SM mainloop pipe. Mirrors mlp4_make_pipe_umma; num_consumers=1 -// (one umma_arrive per stage releases the buffer). is_leader restricts the -// TMA-issuing thread to lane 0 of the pair's rank-0 CTA — both peer CTAs' -// smem buffers arrive on the shared cluster barrier via the 2SM TMA -// multicast, so only one CTA needs to issue the copy. -template -__device__ __forceinline__ typename Traits::MainloopPipelineUmma2Sm -mlp4_make_pipe_umma_2sm( - typename Traits::MainloopPipelineUmma2Sm::SharedStorage& storage) { - using Pipeline = typename Traits::MainloopPipelineUmma2Sm; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = warp_id == 0; - bool is_consumer = warp_id >= 3 && warp_id <= 11; - - typename Pipeline::Params params; - params.transaction_bytes = Traits::TmaTransBytes; - params.num_producers = 1; - params.num_consumers = 1; - params.initializing_warp = 0; - if (is_producer) { - params.role = Category::Producer; - params.is_leader = - threadIdx.x == 0 && cute::block_rank_in_cluster() % 2 == 0; - } else if (is_consumer) { - params.role = Category::Consumer; - } else { - params.role = Category::NonParticipant; - } - return Pipeline( - storage, params, typename Traits::ClusterShape{}, - cute::true_type{}, cute::true_type{}); -} - -// ═══════════════════════════════════════════════════════════════════ -// 2SM producer/consumer — chunk-fixed, sequential dB/dC (same cell -// decode as mlp4_producer/mlp4_consumer above), ClusterM=2 only: the -// walk axis (num_m_tiles) is partitioned by `lane` alone (no cluster/pair -// multicast fan-out across multiple pairs). -// ═══════════════════════════════════════════════════════════════════ - -template <> -struct Mlp4ProducerImpl<100> { -template < - typename Traits, - typename Pipeline, - typename TmaLoadX, - typename TmaLoadA, - typename SmemType> -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaLoadX const& tma_load_x, - TmaLoadA const& tma_load_dut, - TmaLoadA const& tma_load_dvt, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int intermediate_dim, - int hidden_dim, - int num_tokens, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - int ring_kb = 0) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - auto sA = make_tensor( - make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); - auto sX = make_tensor( - make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - - auto mX = tma_load_x.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(num_tokens))); - auto mdUT = tma_load_dut.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(num_tokens))); - auto mdVT = tma_load_dvt.get_tma_tensor(make_shape( - static_cast(intermediate_dim), - static_cast(num_tokens))); - - typename Traits::TiledMma2Sm tiled_mma; - int cta_rank = static_cast(cute::block_rank_in_cluster()); - int pair_rank = cta_rank % 2; - auto cta_mma = tiled_mma.get_slice(pair_rank); - // ClusterM=2 → the pair IS the whole cluster. - auto pair_layout_vmnk = tiled_divide( - make_layout(typename Traits::ClusterShape{}), - make_tile(typename Traits::TiledMma2Sm::AtomThrID{})); - auto pair_coord_vmnk = pair_layout_vmnk.get_flat_coord(pair_rank); - uint16_t mcast_mask_a = - create_tma_multicast_mask<2>(pair_layout_vmnk, pair_coord_vmnk); - uint16_t mcast_mask_x = - create_tma_multicast_mask<1>(pair_layout_vmnk, pair_coord_vmnk); - - int total_chunks = num_experts * num_n_tiles; - int total_cells = total_chunks * outer_split * k_split; - - for (int cell_idx = cell_start; - cell_idx < total_cells; - cell_idx += cell_stride) { - int k_slice = cell_idx % k_split; - int cell_om = cell_idx / k_split; - int chunk_idx = cell_om / outer_split; - int lane = cell_om - chunk_idx * outer_split; - - int e = chunk_idx / num_n_tiles; - int n_tile = chunk_idx - e * num_n_tiles; - // Balanced split: multiply BEFORE dividing so any remainder - // (num_m_tiles % outer_split != 0) is distributed across the - // trailing lanes instead of being silently dropped — e.g. - // num_m_tiles=3, outer_split=2 gives lane0=[0,1), lane1=[1,3), - // covering all 3 tiles. `lane * (num_m_tiles/outer_split)` (a - // naive floor-then-multiply) would leave the tail tile - // unvisited by any lane. ClusterM=2 needs no further - // valid_m_tile guard: walk_end for the last lane is always - // exactly num_m_tiles. - int walk_begin = lane * num_m_tiles / outer_split; - int walk_end = (lane + 1) * num_m_tiles / outer_split; - - int kb_lo = max(expert_k_starts[e], batch_kb_start); - int kb_hi = min(expert_k_ends[e], batch_kb_end); - if (kb_hi <= kb_lo) - continue; - - int k_total = kb_hi - kb_lo; - int k_per = (k_total + k_split - 1) / k_split; - int slice_lo = kb_lo + k_slice * k_per; - int slice_hi = min(slice_lo + k_per, kb_hi); - if (slice_hi <= slice_lo) - continue; - kb_lo = slice_lo; - kb_hi = slice_hi; - - CUTE_UNROLL - for (int phase = 0; phase < 2; ++phase) { - auto mA = phase == 0 ? mdUT : mdVT; - const auto& tma_a = - phase == 0 ? tma_load_dut : tma_load_dvt; - - for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { - auto coord = make_coord(m_tile, n_tile, _); - auto gA = local_tile( - mA, typename Traits::TileShape{}, coord, - Step<_1, X, _1>{}); - auto gX = local_tile( - mX, typename Traits::TileShape{}, coord, - Step{}); - auto tCgA = cta_mma.partition_A(gA); - auto tCgX = cta_mma.partition_B(gX); - - auto [tAgA, tAsA] = tma_partition( - tma_a, - get<2>(pair_coord_vmnk), - make_layout(size<2>(pair_layout_vmnk)), - group_modes<0, 3>(sA), - group_modes<0, 3>(tCgA)); - auto [tXgX, tXsX] = tma_partition( - tma_load_x, - get<1>(pair_coord_vmnk), - make_layout(size<1>(pair_layout_vmnk)), - group_modes<0, 3>(sX), - group_modes<0, 3>(tCgX)); - - for (int kb = kb_lo; kb < kb_hi; ++kb) { - int kb_g = ring_kb > 0 ? kb % ring_kb : kb; - pipe.producer_acquire(state); - if (cute::elect_one_sync()) { - auto* barrier = - pipe.producer_get_barrier(state); - copy( - tma_a.with(*barrier, mcast_mask_a), - tAgA(_, kb_g), - tAsA(_, state.index())); - copy( - tma_load_x.with( - *barrier, mcast_mask_x), - tXgX(_, kb_g), - tXsX(_, state.index())); - } - ++state; - } - } - } - } - pipe.producer_tail(state); +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp4_sm90.cuh" #else - __trap(); +#include "mlp4_sm100.cuh" #endif -} -}; // Mlp4ProducerImpl<100> - -template <> -struct Mlp4ConsumerImpl<100> { -template < - typename Traits, - bool Expert3D = false, - typename Pipeline, - typename TmaReduceAddDB, - typename TmaReduceAddDC, - typename SmemType> -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - SmemType& smem, - TmaReduceAddDB const& tma_reduce_db, - TmaReduceAddDC const& tma_reduce_dc, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int intermediate_dim, - int hidden_dim, - int total_m_rows, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int cell_start, - int cell_stride, - int batch_kb_start, - int batch_kb_end, - int k_split, - const cutlass::arch::ClusterBarrier* pair_init_barrier, - uint32_t* pair_init_phase) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - (void)intermediate_dim; - using Element = typename Traits::Element; - constexpr int TileM = Traits::TileM; - constexpr int TileN = Traits::TileN; - constexpr int CtaTileM = Traits::CtaTileM; - constexpr int EpiChunkN = Traits::EpiChunkN; - constexpr int NumEpiRounds = Traits::NumEpiRounds; - constexpr int kMmaEpiThreads = - Traits::ConsumerThreads + Traits::WarpSize; - - int cta_rank = static_cast(cute::block_rank_in_cluster()); - int pair_rank = cta_rank % 2; - bool is_leader_cta = pair_rank == 0; - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_mma_warp = warp_id == 3; - bool is_epilogue = warp_id >= 4 && warp_id <= 11; - int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; - int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; - int tid_in_wg = is_epilogue - ? tid_in_epi % Traits::WarpGroupSize - : 0; - int wg_barrier_id = 1 + wg; - bool is_wg_leader = is_epilogue && tid_in_wg == 0; - - typename Traits::TiledMma2Sm tiled_mma; - auto cta_mma = tiled_mma.get_slice(pair_rank); - auto sA = make_tensor( - make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); - auto sX = make_tensor( - make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); - auto tCrA = cta_mma.make_fragment_A(sA); - auto tCrX = cta_mma.make_fragment_B(sX); - - auto cAccFull = make_identity_tensor( - make_shape(Int{}, Int{})); - auto tCgC = cta_mma.partition_C(cAccFull); - auto tCtAcc = cta_mma.make_fragment_C(tCgC); - - if (pair_init_barrier != nullptr) { - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); - if (warp_id == 3) { - if (cute::elect_one_sync()) { - uint32_t rank = cute::block_rank_in_cluster(); - uint32_t peer = rank ^ 1; - bool leader = (rank % 2) == 0; - pair_init_barrier->arrive(peer, !leader); - pair_init_barrier->wait(*pair_init_phase); - pair_init_barrier->arrive(peer, leader); - *pair_init_phase ^= 1; - } - __syncwarp(); - } - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); - } - - using AccPipe = typename Traits::AccumulatorPipeline2Sm; - typename AccPipe::Params acc_params; - acc_params.role = is_mma_warp && is_leader_cta - ? AccPipe::ThreadCategory::Producer - : AccPipe::ThreadCategory::Consumer; - acc_params.producer_arv_count = 1; - acc_params.consumer_arv_count = 2; - acc_params.initializing_warp = 4; - AccPipe acc_pipe( - smem.acc_pipe, acc_params, typename Traits::ClusterShape{}); - auto acc_prod_state = - cutlass::make_producer_start_state(); - typename AccPipe::PipelineState acc_cons_state; - - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); - if (pair_init_barrier != nullptr && warp_id == 3) { - if (cute::elect_one_sync()) { - uint32_t rank = cute::block_rank_in_cluster(); - uint32_t peer = rank ^ 1; - bool leader = (rank % 2) == 0; - pair_init_barrier->arrive(peer, !leader); - pair_init_barrier->wait(*pair_init_phase); - pair_init_barrier->arrive(peer, leader); - *pair_init_phase ^= 1; - } - __syncwarp(); - } - if (pair_init_barrier != nullptr) - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); - uint32_t tmem_base = smem.tmem_base; - tCtAcc.data() = tmem_base; - - constexpr int store_slot_elems = CtaTileM * EpiChunkN; - Element* store_ptr = smem.store_buf + wg * store_slot_elems; - auto sStore = make_tensor( - make_smem_ptr(store_ptr), - typename Traits::SmemLayoutStoreSlot{}); - - auto mdB = tma_reduce_db.get_tma_tensor(make_shape( - static_cast(total_m_rows), - static_cast(hidden_dim))); - auto mdC = tma_reduce_dc.get_tma_tensor(make_shape( - static_cast(total_m_rows), - static_cast(hidden_dim))); - auto cta_tma_db = tma_reduce_db.get_slice(Int<0>{}); - auto cta_tma_dc = tma_reduce_dc.get_slice(Int<0>{}); - - auto epi_tile = - make_tile(Int{}, Int{}); - auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); - auto tAcc_epi = flat_divide(acc_mn, epi_tile); - auto t2r = make_tmem_copy( - TmemLoadOp{}, - tAcc_epi(_, _, _0{}, _0{})); - auto thr_t2r = t2r.get_slice(tid_in_wg); - auto cChunk = make_identity_tensor( - make_shape(Int{}, Int{})); - auto tTR_cChunk = thr_t2r.partition_D(cChunk); - auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); - - Layout tmem_warp_layout = - typename decltype(make_tmem_warp_partitioner( - tAcc_epi(_, _, _0{}, _0{})))::TiledLayout_TV{}; - constexpr bool predicate_tmem_load = - size(tmem_warp_layout) != cosize(tmem_warp_layout); - - int total_chunks = num_experts * num_n_tiles; - int total_cells = total_chunks * outer_split * k_split; - bool store_in_flight = false; - - for (int cell_idx = cell_start; - cell_idx < total_cells; - cell_idx += cell_stride) { - int k_slice = cell_idx % k_split; - int cell_om = cell_idx / k_split; - int chunk_idx = cell_om / outer_split; - int lane = cell_om - chunk_idx * outer_split; - - int e = chunk_idx / num_n_tiles; - int n_tile = chunk_idx - e * num_n_tiles; - // Match the producer's balanced split: multiply-before-divide - // avoids dropping the tail - // M-tile(s) when num_m_tiles % outer_split != 0. - int walk_begin = lane * num_m_tiles / outer_split; - int walk_end = (lane + 1) * num_m_tiles / outer_split; - - int kb_lo = max(expert_k_starts[e], batch_kb_start); - int kb_hi = min(expert_k_ends[e], batch_kb_end); - if (kb_hi <= kb_lo) - continue; - - int k_total = kb_hi - kb_lo; - int k_per = (k_total + k_split - 1) / k_split; - int slice_lo = kb_lo + k_slice * k_per; - int slice_hi = min(slice_lo + k_per, kb_hi); - if (slice_hi <= slice_lo) - continue; - kb_lo = slice_lo; - kb_hi = slice_hi; - - CUTE_UNROLL - for (int phase = 0; phase < 2; ++phase) { - for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { - if (is_mma_warp && is_leader_cta) { - acc_pipe.producer_acquire(acc_prod_state); - int acc_stage = acc_prod_state.index(); - tCtAcc.data() = - tmem_base + uint32_t(acc_stage * TileN); - - bool first_mma = true; - for (int kb = kb_lo; kb < kb_hi; ++kb) { - pipe.consumer_wait(state); - CUTE_UNROLL - for (int ks = 0; - ks < size<2>(tCrA); - ++ks) { - tiled_mma.accumulate_ = first_mma - ? UMMA::ScaleOut::Zero - : UMMA::ScaleOut::One; - first_mma = false; - gemm( - tiled_mma, - tCrA(_, _, ks, state.index()), - tCrX(_, _, ks, state.index()), - tCtAcc); - } - pipe.consumer_release(state); - ++state; - } - acc_pipe.producer_commit(acc_prod_state); - ++acc_prod_state; - } - - if (is_epilogue) { - acc_pipe.consumer_wait(acc_cons_state); - int acc_stage = acc_cons_state.index(); - tCtAcc.data() = - tmem_base + uint32_t(acc_stage * TileN); - auto acc_mn_stage = - tCtAcc(make_coord(_, _), _0{}, _0{}); - auto tAcc_epi_stage = - flat_divide(acc_mn_stage, epi_tile); - auto tTR_tAcc = - thr_t2r.partition_S(tAcc_epi_stage); - - CUTE_UNROLL - for (int round = 0; - round < NumEpiRounds; - ++round) { - int chunk = wg * NumEpiRounds + round; - auto tAccChunk = - tTR_tAcc(_, _, _, _0{}, chunk); - bool issue_tmem_load = true; - if constexpr (predicate_tmem_load) { - int subpart = - (tAccChunk.data().dp_ / 32) % 4; - issue_tmem_load = - tid_in_wg / Traits::WarpSize == - subpart; - } - if (issue_tmem_load) { - copy(t2r, tAccChunk, tTR_rAcc); - cutlass::arch::fence_view_async_tmem_load(); - } - - if (store_in_flight) - cute::tma_store_wait<0>(); - cutlass::arch::NamedBarrier::sync( - Traits::WarpGroupSize, - wg_barrier_id); - - if (issue_tmem_load) { - CUTE_UNROLL - for (int i = 0; - i < size(tTR_rAcc); - ++i) { - int m_local = - get<0>(tTR_cChunk(i)); - int n_local = - get<1>(tTR_cChunk(i)); - sStore(m_local, n_local) = - static_cast( - tTR_rAcc(i)); - } - } - cutlass::arch::NamedBarrier::sync( - Traits::WarpGroupSize, - wg_barrier_id); - - if (is_wg_leader) { - cute::tma_store_fence(); - int out_m = - e * num_m_tiles + m_tile; - int n_chunk = - n_tile * - (TileN / EpiChunkN) + - wg * NumEpiRounds + - round; - CUTE_UNROLL - for (int atom = 0; - atom < - Traits::kAtomsPerCta; - ++atom) { - int m_atom = - (TileM / - Traits::AtomTileM) * - out_m + - Traits::kAtomsPerCta * - pair_rank + - atom; - auto sStoreAtom = - local_tile( - sStore, - make_tile( - Int< - Traits:: - AtomTileM>{}, - Int< - EpiChunkN>{}), - make_coord( - atom, 0)); - if (phase == 0) { - auto gdB = local_tile( - mdB, - make_tile( - Int< - Traits:: - AtomTileM>{}, - Int< - EpiChunkN>{}), - make_coord( - m_atom, - n_chunk)); - copy( - tma_reduce_db, - cta_tma_db.partition_S( - sStoreAtom), - cta_tma_db.partition_D( - gdB)); - } else { - auto gdC = local_tile( - mdC, - make_tile( - Int< - Traits:: - AtomTileM>{}, - Int< - EpiChunkN>{}), - make_coord( - m_atom, - n_chunk)); - copy( - tma_reduce_dc, - cta_tma_dc.partition_S( - sStoreAtom), - cta_tma_dc.partition_D( - gdC)); - } - } - cute::tma_store_arrive(); - } - store_in_flight = true; - } - - cutlass::arch::NamedBarrier::sync( - Traits::ConsumerThreads, 0); - // One elected release from each CTA satisfies the - // cluster-scoped accumulator pipeline's count of two. - if (tid_in_epi == 0) - acc_pipe.consumer_release( - acc_cons_state); - ++acc_cons_state; - } - } - } - } - - if (is_epilogue && store_in_flight) - cute::tma_store_wait<0>(); - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); -#else - __trap(); -#endif -} -}; // Mlp4ConsumerImpl<100> - -// ═══════════════════════════════════════════════════════════════════ -// Single-launch driver, chunk-fixed and architecture-dispatched -// ═══════════════════════════════════════════════════════════════════ -// -// Compute=90 uses the original per-CTA Hopper pipeline and cell mapping. -// Compute=100 uses a paired-CTA pipeline, pair-granular cell mapping, a -// shared Allocator2Sm allocation, and cluster synchronization. -template < - typename Traits, - int Compute = 90, - bool Expert3D = false, - typename TmaLoadX, - typename TmaLoadA, - typename TmaReduceAddDB, - typename TmaReduceAddDC> -__device__ __forceinline__ void mlp4_fwd( - Mlp4Smem& smem, - TmaLoadX const& tma_load_x, - TmaLoadA const& tma_load_dut, - TmaLoadA const& tma_load_dvt, - TmaReduceAddDB const& tma_reduce_db, - TmaReduceAddDC const& tma_reduce_dc, - const int* expert_k_starts, - const int* expert_k_ends, - int num_experts, - int intermediate_dim, - int hidden_dim, - int num_tokens, - int total_m_rows, - int num_m_tiles, - int num_n_tiles, - int outer_split, - int batch_kb_start = 0, - int batch_kb_end = -1, - int k_split = 1, - int ring_kb = 0) { - static_assert(Compute == 90 || Compute == 100, - "mlp4_fwd supports Compute=90 or Compute=100"); - using PipeState = typename Traits::PipelineState; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = warp_id == 0; - constexpr int kFirstConsumerWarp = Compute == 100 ? 3 : 4; - bool is_consumer = - warp_id >= kFirstConsumerWarp && warp_id <= 11; - - cute::prefetch_tma_descriptor( - tma_load_x.get_tma_descriptor()); - cute::prefetch_tma_descriptor( - tma_load_dut.get_tma_descriptor()); - cute::prefetch_tma_descriptor( - tma_load_dvt.get_tma_descriptor()); - cute::prefetch_tma_descriptor( - tma_reduce_db.get_tma_descriptor()); - cute::prefetch_tma_descriptor( - tma_reduce_dc.get_tma_descriptor()); - - auto pipe = [&]() { - if constexpr (Compute == 100) - return mlp4_make_pipe_umma_2sm( - smem.pipe_storage); - else - return mlp4_make_pipe(smem.pipe_storage); - }(); - using Pipeline = decltype(pipe); - - if constexpr (Compute == 100) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - cute::TMEM::Allocator2Sm tmem_allocator; - if (gridDim.x % Traits::ClusterM != 0) - __trap(); - cute::cluster_sync(); - constexpr int kTmemColumns = - Traits::AccStages * Traits::TileN; - if (warp_id == 4) { - tmem_allocator.allocate( - kTmemColumns, &smem.tmem_base); - __syncwarp(); - } - __syncthreads(); - cute::cluster_sync(); -#else - __trap(); -#endif - } else { - __syncthreads(); - } - - PipeState producer_state = is_producer - ? cutlass::make_producer_start_state() - : PipeState{}; - PipeState consumer_state; - - int cell_start = static_cast(blockIdx.x); - int cell_stride = static_cast(gridDim.x); - if constexpr (Compute == 100) { - cell_start /= Traits::ClusterM; - cell_stride /= Traits::ClusterM; - } - int kb_end = batch_kb_end >= 0 - ? batch_kb_end - : num_tokens / Traits::TileK; - - if (is_producer) { - mlp4_producer( - pipe, - producer_state, - smem, - tma_load_x, - tma_load_dut, - tma_load_dvt, - expert_k_starts, - expert_k_ends, - num_experts, - intermediate_dim, - hidden_dim, - num_tokens, - num_m_tiles, - num_n_tiles, - outer_split, - cell_start, - cell_stride, - batch_kb_start, - kb_end, - k_split, - ring_kb); - } else if (is_consumer) { - if constexpr (Compute == 100) { - mlp4_consumer( - pipe, - consumer_state, - smem, - tma_reduce_db, - tma_reduce_dc, - expert_k_starts, - expert_k_ends, - num_experts, - intermediate_dim, - hidden_dim, - total_m_rows, - num_m_tiles, - num_n_tiles, - outer_split, - cell_start, - cell_stride, - batch_kb_start, - kb_end, - k_split); - } else { -#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) - mlp4_consumer( - pipe, - consumer_state, - smem, - tma_reduce_db, - tma_reduce_dc, - expert_k_starts, - expert_k_ends, - num_experts, - intermediate_dim, - hidden_dim, - total_m_rows, - num_m_tiles, - num_n_tiles, - outer_split, - cell_start, - cell_stride, - batch_kb_start, - kb_end, - k_split); -#else - __trap(); -#endif - } - } - - if constexpr (Compute == 100) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - cute::TMEM::Allocator2Sm tmem_allocator; - __syncthreads(); - cute::cluster_sync(); - constexpr int kTmemColumns = - Traits::AccStages * Traits::TileN; - if (warp_id == 4) { - tmem_allocator.release_allocation_lock(); - tmem_allocator.free( - smem.tmem_base, kTmemColumns); - } -#endif - } -} - -// ═══════════════════════════════════════════════════════════════════ -// Host launcher (raw-pointer; defined in mlp4.cu) -// ═══════════════════════════════════════════════════════════════════ - -// dB = dU^T @ X, dC = dV^T @ X. Both outputs accumulated atomically -// via TMA_REDUCE_ADD. dB and dC MUST be zero-initialized by caller. -// -// dU : [num_tokens, intermediate_dim] -// dV : [num_tokens, intermediate_dim] -// X : [num_tokens, hidden_dim] -// dB, dC : [num_experts, intermediate_dim, hidden_dim] -// expert_for_k_block: [num_tokens / TileK] int32, TileK=64 -void mlp4_fwd_bf16_launch( - const cutlass::bfloat16_t* dU, - const cutlass::bfloat16_t* dV, - const cutlass::bfloat16_t* X, - cutlass::bfloat16_t* dB, - cutlass::bfloat16_t* dC, - const int* expert_for_k_block, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int num_experts, - int device, - cudaStream_t stream); - -} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp4_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp4_sm100.cuh new file mode 100644 index 00000000..640ebed6 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp4_sm100.cuh @@ -0,0 +1,1839 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 4: dB = dU^T @ X, dC = dV^T @ X (weight grads, bwd) +// ═══════════════════════════════════════════════════════════════════ +// +// TMA_REDUCE_ADD variant. Cooperative 2-WG consumer, mirroring mlp3.cuh. +// The two WGs cooperate on ONE output tile at a time, so dB and dC are +// produced SEQUENTIALLY (two phases per cell). This lets a SINGLE A-side +// smem buffer be reused for dU^T (phase dB) then dV^T (phase dC) — vs the +// old design's two simultaneous A buffers — so the footprint matches mlp3 +// and fits the 228 KiB cap at 256-wide tiles. +// +// Only two configs supported (mirrors Mlp3Traits): +// (TileM=256, TileN=128) → M-split (Layout<_2,_1,_1>) ← default +// (TileM=128, TileN=256) → N-split (Layout<_1,_2,_1>) +// +// GEMM orientation: dB = dU^T @ X, so C[M=I, N=H], A[M=I,K=T] = dU^T/dV^T, +// B[N=H,K=T] = X. TileM tiles the intermediate (I) dim, TileN the hidden +// (H) dim. Per-WG output footprint = (WgTileM, WgTileN) = (128, 128). +// +// Single fused A + X TMA pipe (one mbarrier per stage, 2 TMA copies). The +// shared operand X is read by both dB and dC phases and by both WGs. +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for A / X) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Chunk-fixed walk (same principle as mlp3): hold the SHARED operand +// chunk-constant and walk the cooperation (split) axis: +// kMSplit=true (M-split, X shared): chunk = (e, n_tile), walks m_tile; holds X(n,*) +// kMSplit=false (N-split, dU^T/dV^T shared): chunk = (e, m_tile), walks n_tile; holds A(m,*) +// In the M-split default, X(n_tile) is held across BOTH the m-walk and the +// dB/dC phases — re-read for phase dC hits L2. +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm100.cuh" +#include "tmem_load_op_sm100.cuh" + +#include +#include +#include +#include +#include + +// ── SM100 (Compute=100 / Blackwell UMMA) headers ──────────────────── +// Header-only; safe to include under an sm_90a build. The tcgen05 PTX they +// emit is guarded by CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself +// gated on __CUDA_ARCH__ >= 1000) ever instantiates it. Mirrors mlp5.cuh / +// mlp1_fused.cuh. +#include +#include +#include +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +// ── SM100 2-SM (paired-CTA cluster) headers ───────────────────────── +// Needed by the paired MLP4 path below: cluster synchronization/rank queries +// for the (2,1,1) cluster and the 2SM TMA copy traits used by tma_partition(). +#include +#include + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 256, + int TileN_ = 128, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 64 +> +struct Mlp4Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + // Cooperative 2-WG layout. Only two configs supported: + // (TileM=256, TileN=128) → kMSplit=true, WgLayout<_2,_1,_1> + // (TileM=128, TileN=256) → kMSplit=false, WgLayout<_1,_2,_1> + // Per-WG output footprint = (WgTileM=128, WgTileN=128) in both. Each + // WG covers 2 wgmma m=64 atoms per K-step (atom_count_M = 2). + static_assert((TileM == 256 && TileN == 128) || + (TileM == 128 && TileN == 256), + "Mlp4Traits supports only (TileM, TileN) = (256, 128) M-split " + "or (128, 256) N-split"); + + static constexpr bool kMSplit = (TileM > TileN); + static constexpr int AtomTileM = 64; + static constexpr int WgTileM = 128; + static constexpr int WgTileN = 128; + + static_assert(WgTileN % EpiChunkN == 0, + "EpiChunkN must divide WgTileN (=128)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + static constexpr int AtomN = WgTileN; + using GmmaAtom = typename GmmaSelectorMNMN::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtom = GMMA::Layout_MN_SW128_Atom; + + // A-side = dU^T / dV^T (M-side, TileM); X = shared B-side (N-side, TileN). + using SmemLayoutA_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, + Step<_2, _1>{})); + using SmemLayoutX_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, + Step<_2, _1>{})); + using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + using SmemLayoutX = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // ── SM100 (Compute=100 / Blackwell UMMA) ──────────────────────── + // The mainloop pipeline is the UMMA-aware TMA pipeline (its consumer_release + // issues the UMMA-gated smem-buffer arrival). Params/PipelineState/ + // SharedStorage alias the Hopper PipelineTmaAsync's, so smem.pipe_storage + // (typed MainloopPipeline::SharedStorage) drives BOTH — the standalone + // launcher passes it to mlp4_make_pipe_umma unchanged. num_consumers=1 on + // Blackwell (one umma_arrive per stage releases the buffer). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Per-WG UMMA atom: SM100_MMA_F16BF16_SS with M=WgTileM=128, N=WgTileN=128, + // BOTH operands MN-major (A = dU^T/dV^T is M-major in smem; X is N-major). + // Warp 3 issues TWO of these per k-step into two independent + // (128,128) TMEM accumulators (one per consumer WG) — see + // Mlp4ConsumerImpl<100>. UMMA requires M ∈ {64,128}, so the 256-wide tile + // cannot use one atom; the per-WG (128,128) atom is used for both configs. + using TiledMmaUmma = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS{})); + // Double-buffered TMEM accumulators: each stage owns one (WgTileM,WgTileN) + // accumulator per epilogue WG. The two dB/dC phases still clear + // independently — this only lets MMA for the next (phase,walk) run while + // the epilogue drains the previous accumulator stage. + static constexpr int AccStages = 2; + static_assert(AccStages * (2 * WgTileN) <= 512, + "SM100 MLP4 needs AccStages x two per-WG accumulators to fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + static constexpr int TmaTransBytesA = + static_cast(size(SmemLayoutA_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesX = + static_cast(size(SmemLayoutX_1{}) * sizeof(Element)); + // Fused pipeline: X + ONE A operand (dU^T or dV^T, per phase) share one + // mbarrier per stage. Both WGs consume the same X via cooperative TiledMma. + static constexpr int TmaTransBytes = TmaTransBytesA + TmaTransBytesX; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; // 256 + static constexpr int NumThreads = 384; + + // Per-WG store buffer slot: (WgTileM, EpiChunkN). One slot per WG; it + // holds kAtomsPerWg = WgTileM/AtomTileM contiguous atom-rows after the + // acc→smem fold. + static constexpr int kAtomsPerWg = WgTileM / AtomTileM; // 2 + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + // TMA-reduce descriptor box = ONE wgmma atom-row (AtomTileM × EpiChunkN). + // At TileM=256 each WG owns 2 *interleaved* atoms that map to non-contiguous + // gmem rows, so a WG issues kAtomsPerWg separate TMA stores per epi round + // (mirrors mlp1_fused.cuh / mlp3.cuh). + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; +}; + +template +using Mlp4MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Shared Memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp4Smem { + using Element = typename Traits::Element; + + static constexpr int smem_X_size = cosize_v; + static constexpr int smem_A_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + // Single A buffer, reused for dU^T (phase dB) then dV^T (phase dC). + alignas(128) Element smem_X[smem_X_size]; + alignas(128) Element smem_A[smem_A_size]; + // Per-WG store buffer: WG1 uses [0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + typename Traits::MainloopPipeline::SharedStorage pipe_storage; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these; + // SM100 launchers allocate/free TMEM outside the consumer and publish the + // base here; the accumulator pipeline storage still lives in this smem. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* X_data() { return &smem_X[0]; } + CUTE_DEVICE Element* A_data() { return &smem_A[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP4 shared memory (data only — pipe storage lives in the +// outer MlpFusedBwdSmem union; otherwise identical to Mlp4Smem). +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp4FusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_X_size = cosize_v; + static constexpr int smem_A_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_X[smem_X_size]; + alignas(128) Element smem_A[smem_A_size]; + alignas(128) Element store_buf[2 * smem_store_size]; + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* X_data() { return &smem_X[0]; } + CUTE_DEVICE Element* A_data() { return &smem_A[0]; } +}; + + +// ═══════════════════════════════════════════════════════════════════ +// Pipe construction helper (mirrors mlp3_make_pipe) +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ typename Traits::MainloopPipeline +mlp4_make_pipe(typename Traits::MainloopPipeline::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp4_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the (shared) +// mlp4_producer and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ typename Traits::MainloopPipelineUmma +mlp4_make_pipe_umma(typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Chunk-fixed cooperative producer/consumer — sequential dB / dC +// ═══════════════════════════════════════════════════════════════════ +// +// Used by the standalone mlp4 kernel (mlp4.cu) and mlp_bwd Phase 2. +// Each CTA owns ONE (chunk, lane) cell and processes it in TWO phases +// (phase 0 = dB via dU^T, phase 1 = dC via dV^T). We hold the SHARED +// operand and walk the cooperation (split) axis: +// kMSplit=true (M-split): chunk = (e, n_tile), lane partitions +// num_m_tiles, walks m_tile. X(n_tile, *) — the shared +// N-side operand — is the L2-held chunk-constant, +// reused across the m-walk AND both phases. +// kMSplit=false (N-split): chunk = (e, m_tile), lane partitions +// num_n_tiles, walks n_tile. A(m_tile, *) (dU^T then +// dV^T) is L2-held; X is re-streamed per phase. +// +// Cell decode (chunk-major: cell_idx = chunk_idx × outer_split + lane): +// total_cells = total_chunks × outer_split +// cell → chunk_idx = cell_idx / outer_split +// lane = cell_idx % outer_split + +template +struct Mlp4ProducerImpl; + +template <> +struct Mlp4ProducerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadX const& tma_load_x, + TmaLoadA const& tma_load_dut, + TmaLoadA const& tma_load_dvt, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + // Staging-ring depth in K-blocks; > 0 wraps the gmem K-tile coordinate + // (% ring_kb) for the grouped Phase-2 that may straddle the ring. See + // mlp3_producer. 0 = no wrap. + int ring_kb = 0) { + + auto sA = make_tensor(make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + + // int64_t casts: H·T, I·T can exceed 2^31 elems at production shapes. + auto mX = tma_load_x.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(num_tokens))); + auto mdUT = tma_load_dut.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + auto mdVT = tma_load_dvt.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + + auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); + auto cta_tma_dut = tma_load_dut.get_slice(Int<0>{}); + auto cta_tma_dvt = tma_load_dvt.get_slice(Int<0>{}); + + auto tXsX = cta_tma_x.partition_D(sX); + auto tAsA = cta_tma_dut.partition_D(sA); // dU^T/dV^T share A layout + + int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); + // K-split: each (chunk, walk-lane) replicated k_split times; k_slice owns a + // 1/k_split sub-range of the K-loop. Partials sum via TMA_REDUCE_ADD. + int total_cells = total_chunks * outer_split * k_split; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e, m_chunk, n_chunk, walk_begin, walk_end; + if constexpr (Traits::kMSplit) { + // M-split: hold X(n) shared. chunk = (e, n_tile), walks m_tile. + e = chunk_idx / num_n_tiles; + n_chunk = chunk_idx - e * num_n_tiles; + m_chunk = 0; + int m_per = num_m_tiles / outer_split; + walk_begin = lane * m_per; + walk_end = walk_begin + m_per; + } else { + // N-split: hold A(m) shared. chunk = (e, m_tile), walks n_tile. + e = chunk_idx / num_m_tiles; + m_chunk = chunk_idx - e * num_m_tiles; + n_chunk = 0; + int n_per = num_n_tiles / outer_split; + walk_begin = lane * n_per; + walk_end = walk_begin + n_per; + } + + // Intersect this expert's global K range with the current + // batch's K-window. + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) continue; + + // K-split: clip to this cell's k_slice sub-range of [kb_lo, kb_hi). + { + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int s_lo = kb_lo + k_slice * k_per; + int s_hi = min(s_lo + k_per, kb_hi); + if (s_hi <= s_lo) continue; // empty slice + kb_lo = s_lo; + kb_hi = s_hi; + } + + // Two phases per cell: phase 0 = dB (dU^T), phase 1 = dC (dV^T). + // dU^T/dV^T share descriptor layout; pair the tma object, its slice, + // and its gmem tensor consistently per phase. + CUTE_UNROLL + for (int phase = 0; phase < 2; ++phase) { + auto mA = (phase == 0) ? mdUT : mdVT; + const auto& tma_A = (phase == 0) ? tma_load_dut : tma_load_dvt; + const auto& cta_tma_A = (phase == 0) ? cta_tma_dut : cta_tma_dvt; + + for (int w = walk_begin; w < walk_end; ++w) { + int m_tile = Traits::kMSplit ? w : m_chunk; + int n_tile = Traits::kMSplit ? n_chunk : w; + auto gX = local_tile(mX, + make_tile(Int{}, Int{}), + make_coord(n_tile, _)); + auto gA = local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(m_tile, _)); + auto tXgX = cta_tma_x.partition_S(gX); + auto tAgA = cta_tma_A.partition_S(gA); + + for (int kb = kb_lo; kb < kb_hi; ++kb) { + int kb_g = (ring_kb > 0) ? (kb % ring_kb) : kb; // physical ring slot + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_x.with(*bar, 0), + tXgX(_, _, _, kb_g), tXsX(_, _, _, state.index())); + copy(tma_A.with(*bar, 0), + tAgA(_, _, _, kb_g), tAsA(_, _, _, state.index())); + } + ++state; + } + } + } + } +} +}; // Mlp4ProducerImpl<90> + +template +__device__ __forceinline__ void mlp4_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadX const& tma_load_x, + TmaLoadA const& tma_load_dut, + TmaLoadA const& tma_load_dvt, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + int ring_kb = 0) { + Mlp4ProducerImpl::template run( + pipe, state, smem, tma_load_x, tma_load_dut, tma_load_dvt, + expert_k_starts, expert_k_ends, num_experts, + intermediate_dim, hidden_dim, num_tokens, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split, ring_kb); +} + +// ═══════════════════════════════════════════════════════════════════ +// Consumer — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported compute +// capability provides a `run(...)` member (a function-template on the Traits/ +// Pipeline/Smem/Tma types, since function templates can't be partially +// specialized). The free function `mlp4_consumer` below forwards to the right +// specialization — call sites (mlp_bwd.cuh's mlp4_consumer(...)) stay +// unchanged (Compute defaults to 90). +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, register acc, M-atom +// interleave + (m_loc/128)*64+(m_loc%64) store remap) +// Compute=100 → Blackwell / UMMA (warp 3 issues 2 per-WG (128,128) UMMAs +// into double-buffered TMEM stages; warps 4-11 run the +// epilogue, with independent per-phase clear and contiguous +// per-WG store — no remap) + +template +struct Mlp4ConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper / WGMMA. (Verbatim from the original mlp4_consumer.) +// ─────────────────────────────────────────────────────────────────── +template <> +struct Mlp4ConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDB const& tma_reduce_db, + TmaReduceAddDC const& tma_reduce_dc, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int total_m_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split) { + + (void)Compute; + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + const int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sA = make_tensor(make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + + auto tCsA = thr_mma.partition_A(sA); + auto tCsX = thr_mma.partition_B(sX); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + // Identity coords for per-thread acc → store_buf mapping. + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + constexpr int store_slot_elems = Traits::WgTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + const bool is_my_wg_leader = (tid_in_wg == 0); + + auto mdB = [&]() { + if constexpr (Expert3D) { + return tma_reduce_db.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_reduce_db.get_tma_tensor(make_shape( + static_cast(total_m_rows), + static_cast(hidden_dim))); + } + }(); + auto mdC = [&]() { + if constexpr (Expert3D) { + return tma_reduce_dc.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_reduce_dc.get_tma_tensor(make_shape( + static_cast(total_m_rows), + static_cast(hidden_dim))); + } + }(); + auto cta_tma_db = tma_reduce_db.get_slice(Int<0>{}); + auto cta_tma_dc = tma_reduce_dc.get_slice(Int<0>{}); + + int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); + int total_cells = total_chunks * outer_split * k_split; + + constexpr int K_PIPE_MMAS = 1; + + // Carries across phases, walk iters AND cells for cross-tile K-loop / + // store overlap. Drained once at the end. + bool store_in_flight = false; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e, m_chunk, n_chunk, walk_begin, walk_end; + if constexpr (Traits::kMSplit) { + // M-split: hold X(n) shared. chunk = (e, n_tile), walks m_tile. + e = chunk_idx / num_n_tiles; + n_chunk = chunk_idx - e * num_n_tiles; + m_chunk = 0; + int m_per = num_m_tiles / outer_split; + walk_begin = lane * m_per; + walk_end = walk_begin + m_per; + } else { + // N-split: hold A(m) shared. chunk = (e, m_tile), walks n_tile. + e = chunk_idx / num_m_tiles; + m_chunk = chunk_idx - e * num_m_tiles; + n_chunk = 0; + int n_per = num_n_tiles / outer_split; + walk_begin = lane * n_per; + walk_end = walk_begin + n_per; + } + + // Intersect this expert's global K range with the current + // batch's K-window. + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) continue; + + // K-split: clip to this cell's k_slice sub-range (must match producer). + { + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int s_lo = kb_lo + k_slice * k_per; + int s_hi = min(s_lo + k_per, kb_hi); + if (s_hi <= s_lo) continue; // empty slice + kb_lo = s_lo; + kb_hi = s_hi; + } + + int k_count = kb_hi - kb_lo; + + // Two phases per cell: phase 0 = dB, phase 1 = dC. + CUTE_UNROLL + for (int phase = 0; phase < 2; ++phase) { + + for (int w = walk_begin; w < walk_end; ++w) { + int m_tile = Traits::kMSplit ? w : m_chunk; + int n_tile = Traits::kMSplit ? n_chunk : w; + + clear(acc); + auto state_release = state; + int prologue_count = (k_count < K_PIPE_MMAS) ? k_count : K_PIPE_MMAS; + + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsA(_, _, _, state.index()), + tCsX(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + + for (int k = prologue_count; k < k_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsA(_, _, _, state.index()), + tCsX(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + + ++state; + ++state_release; + } + + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // Each WG writes its (WgTileM × EpiChunkN) slot to store_buf + // and TMA_REDUCE_ADDs it into dB (phase 0) or dC (phase 1). + // Store indexing mirrors mlp3_consumer (M-split / N-split). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = (m_loc / 128) * 64 + (m_loc % 64); + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int out_m = Expert3D ? m_tile : e * num_m_tiles + m_tile; + // Store box = ONE atom-row (AtomTileM); kAtomsPerWg stores + // per WG. M-split: the 2 atoms interleave in gmem + // (4*out_m + my_wg + 2a). N-split: WG owns the full TileM + // (2 contiguous atoms) and its own N-half (2*out_m + a). + int n_tile_idx; + if constexpr (Traits::kMSplit) { + n_tile_idx = n_tile * Traits::NumEpiRounds + r; + } else { + n_tile_idx = n_tile * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + constexpr int kRowAtoms = Traits::TileM / Traits::AtomTileM; // 4 (M-split) / 2 (N-split) + CUTE_UNROLL + for (int a = 0; a < Traits::kAtomsPerWg; ++a) { + int m_atom_row = Traits::kMSplit + ? (kRowAtoms * out_m + my_wg + Traits::kAtomsPerWg * a) + : (kRowAtoms * out_m + a); + auto sStore_a = local_tile(sStore, + make_tile(Int{}, Int{}), + make_coord(a, 0)); + if (phase == 0) { + auto gdB = [&]() { + if constexpr (Expert3D) { + return local_tile(mdB, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx, e)); + } else { + return local_tile(mdB, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx)); + } + }(); + copy(tma_reduce_db, cta_tma_db.partition_S(sStore_a), + cta_tma_db.partition_D(gdB)); + } else { + auto gdC = [&]() { + if constexpr (Expert3D) { + return local_tile(mdC, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx, e)); + } else { + return local_tile(mdC, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx)); + } + }(); + copy(tma_reduce_dc, cta_tma_dc.partition_S(sStore_a), + cta_tma_dc.partition_D(gdC)); + } + } + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // struct Mlp4ConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Free-function forwarder. Existing call sites (mlp4_consumer(...)) +// are unchanged (Compute defaults to 90); pass mlp4_consumer(...) +// to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── +template +__device__ __forceinline__ void mlp4_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDB const& tma_reduce_db, + TmaReduceAddDC const& tma_reduce_dc, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int total_m_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, + uint32_t* pair_init_phase = nullptr) { + if constexpr (Compute == 100) { + Mlp4ConsumerImpl::template run( + pipe, state, smem, tma_reduce_db, tma_reduce_dc, + expert_k_starts, expert_k_ends, num_experts, + intermediate_dim, hidden_dim, total_m_rows, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, batch_kb_start, batch_kb_end, k_split, + pair_init_barrier, pair_init_phase); + } else { + Mlp4ConsumerImpl::template run( + pipe, state, smem, tma_reduce_db, tma_reduce_dc, + expert_k_starts, expert_k_ends, num_experts, + intermediate_dim, hidden_dim, total_m_rows, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, batch_kb_start, batch_kb_end, k_split); + } +} + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM (paired-CTA) MLP4 — Mlp4Traits2Sm, pipe/producer/consumer +// ═══════════════════════════════════════════════════════════════════ +// +// A cluster-(2,1,1) SM100 UMMA path: two CTAs jointly compute one +// (TileM=256, TileN=256) tile via SM100_MMA_F16BF16_2x1SM_SS, each CTA +// owning CtaTileM=128 M-rows (kAtomsPerCta=2 wgmma-style 64-row atoms). +// Reuses the generic Mlp4Smem / Mlp4FusedSmem above — +// Mlp4Traits2Sm supplies the same nested type names (SmemLayoutA, +// SmemLayoutX, SmemLayoutStoreSlot, MainloopPipeline, AccumulatorPipeline) +// so no separate 2SM smem struct is needed. +// +// Scope: only the BASIC ClusterM=2 (single paired-CTA) configuration +// required by the fused backward kernel is ported here. Intentionally +// EXCLUDED (upstream "topology experiment" scope, not ported): the +// multi-pair (ClusterM=4) path, shared-operand multicast +// (MulticastShared_), N-split, the compact-epilogue variant, and early +// TMEM release — the latter three never existed for MLP4 in the donor +// tree (MLP3-only experiments) and are noted here only for the audit +// record. Stages and EpiChunkN remain free template parameters (no +// S5/E64 production hardcoding); callers pick tuned values at the +// trait-instantiation site (e.g. moe_bwd.cu). +// +// Cross-file invariants a caller/launcher MUST uphold (none of this is +// enforceable from inside this header-only file; documented for the +// next integrator of mlp_bwd.cuh / moe_bwd.cu / mlp4.cu): +// * grid.x must be a MULTIPLE OF ClusterM (=2) before launch; the runtime +// `__trap()` check in mlp4_fwd is only a defensive backstop. +// * Compute != 100 must stay on the 1SM traits/path. Mlp4Traits2Sm's +// TiledMma2Sm / ClusterShape / MainloopPipelineUmma2Sm / +// AccumulatorPipeline2Sm members do not exist on Mlp4Traits, so mixing +// an architecture with the wrong trait remains a compile-time error. +// * Launch must use cudaLaunchKernelEx with a cluster dimension of +// (ClusterM,1,1), and cudaFuncSetAttribute(..., +// cudaFuncAttributeNonPortableClusterSizeAllowed, 1) must be set on +// the kernel beforehand. +// * The dynamic shared-memory allocation for Mlp4Smem> +// must be sized/capped against the device's *dynamic* opt-in limit +// (cudaDeviceGetAttribute(..., cudaDevAttrMaxSharedMemoryPerBlockOptin, +// ...)), not the static 48 KiB default, via +// cudaFuncAttributeMaxDynamicSharedMemorySize. +// * X / dU^T / dV^T TMA descriptors must be built with +// make_tma_copy_A_sm100 / make_tma_copy_B_sm100 (SM100_TMA_2SM_LOAD +// copy op) — the plain SM90 TMA copy-atom builders are not +// pair-aware and would not multicast/replicate operands across the +// two peer CTAs. dB / dC stay TMA REDUCE_ADD outputs. +// * The producer/consumer specializations assume TMEM is already allocated +// and smem.tmem_base populated. The standalone mlp4_fwd driver owns that +// allocation and rendezvouses peer CTAs with cute::cluster_sync() (safe +// because both CTAs run symmetric control flow). A fused caller sharing +// one TMEM region across warp groups with asymmetric control flow must not +// reuse a bare cluster_sync() for that rendezvous; it must pair +// cute::TMEM::Allocator2Sm's +// allocate()/free() with its own dedicated +// cutlass::arch::ClusterBarrier rendezvous scoped to just the +// participating warps, the way donor's mlp_bwd.cuh +// MlpFusedBwdSmem::tmem_pair_barrier does. That smem field and its +// init/arrive/wait call sites live in mlp_bwd.cuh, outside this +// file's scope, and are not present in this target repo yet. + +template < + typename Element_, + int TileM_ = 256, + int TileN_ = 256, + int TileK_ = 64, + int Stages_ = 3, + int EpiChunkN_ = 128, + int AccStages_ = 2, + int ClusterM_ = 2 +> +struct Mlp4Traits2Sm { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + static constexpr int AccStages = AccStages_; + static constexpr int ClusterM = ClusterM_; + static constexpr int NumPairs = ClusterM / 2; + static constexpr int CtaTileM = TileM / 2; + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = TileN / 2; + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + static constexpr int kAtomsPerCta = CtaTileM / AtomTileM; + + static_assert(TileM == 256, "MLP4 2SM requires joined TileM=256"); + static_assert(TileN == 256, "MLP4 2SM requires joined TileN=256"); + static_assert(TileK == 32 || TileK == 64, + "MLP4 2SM supports TileK=32 or 64"); + static_assert(Stages >= 2 && Stages <= 12, + "MLP4 2SM supports two to twelve mainloop stages"); + static_assert(CtaTileM == 128, "each peer CTA must own 128 M rows"); + static_assert(WgTileN % EpiChunkN == 0, + "EpiChunkN must divide each warpgroup's N half"); + static_assert(kAtomsPerCta == 2, + "each CTA must reduce-add two 64-row atoms"); + static_assert(AccStages == 2, + "MLP4 2SM fixes accumulator double buffering"); + static_assert(AccStages * TileN <= 512, + "MLP4 2SM accumulator stages must fit in 512 TMEM columns"); + static_assert(ClusterM == 2, + "MLP4 2SM ports only the basic single paired-CTA cluster; " + "ClusterM=4 pair-of-pairs multicast is an excluded experiment"); + + using TileShape = Shape, Int, Int>; + using ClusterShape = Shape, _1, _1>; + using AtomThrShape = Shape<_2, _1, _1>; + + using TiledMma2Sm = decltype(make_tiled_mma( + SM100_MMA_F16BF16_2x1SM_SS< + Element, Element, ElementAccum, TileM, TileN, + UMMA::Major::MN, UMMA::Major::MN>{})); + static_assert(size(typename TiledMma2Sm::AtomThrID{}) == 2, + "MLP4 2SM requires a two-CTA MMA atom"); + + using MmaShapeA_MK = decltype(partition_shape_A( + TiledMma2Sm{}, make_shape(Int{}, Int{}))); + using MmaShapeB_NK = decltype(partition_shape_B( + TiledMma2Sm{}, make_shape(Int{}, Int{}))); + + using SmemLayoutAtom = UMMA::Layout_MN_SW128_Atom; + using SmemLayoutA = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtom{}, append(MmaShapeA_MK{}, Int{}), + Step<_2, _1, _3>{})); + using SmemLayoutX = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtom{}, append(MmaShapeB_NK{}, Int{}), + Step<_2, _1, _3>{})); + using SmemLayoutA_1 = decltype(SmemLayoutA{}(_, _, _, Int<0>{})); + using SmemLayoutX_1 = decltype(SmemLayoutX{}(_, _, _, Int<0>{})); + + static constexpr int TmaTransBytesA = + static_cast(cosize_v * sizeof(Element)); + static constexpr int TmaTransBytesX = + static_cast(cosize_v * sizeof(Element)); + // Both peer CTAs issue each 2SM load, and both contributions arrive on + // the cluster-scoped full barrier. + static constexpr int TmaTransBytes = + 2 * (TmaTransBytesA + TmaTransBytesX); + + using MainloopPipelineUmma2Sm = cutlass::PipelineTmaUmmaAsync< + Stages, ClusterShape, AtomThrShape>; + // Aliased to the names the generic Mlp4Smem/Mlp4FusedSmem + // (and Mlp4MainloopPipelineFor) look up. + using MainloopPipelineUmma = MainloopPipelineUmma2Sm; + using MainloopPipeline = MainloopPipelineUmma2Sm; + using PipelineState = typename MainloopPipelineUmma2Sm::PipelineState; + using AccumulatorPipeline2Sm = cutlass::PipelineUmmaAsync< + AccStages, AtomThrShape>; + using AccumulatorPipeline = AccumulatorPipeline2Sm; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; + + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; +}; + +// Compatibility alias: Mlp4Traits2Sm's nested type names line up +// field-for-field with the generic Mlp4Smem above (see the +// comment atop this section), so no bespoke 2SM smem struct is defined — +// this alias just lets callers/tooling written against the name +// "Mlp4Smem2Sm" (e.g. bench_opt/bench_mlp4_2sm.cu) resolve to the +// same generic struct without any source changes on their end. +template +using Mlp4Smem2Sm = Mlp4Smem; + +// SM100 2SM mainloop pipe. Mirrors mlp4_make_pipe_umma; num_consumers=1 +// (one umma_arrive per stage releases the buffer). is_leader restricts the +// TMA-issuing thread to lane 0 of the pair's rank-0 CTA — both peer CTAs' +// smem buffers arrive on the shared cluster barrier via the 2SM TMA +// multicast, so only one CTA needs to issue the copy. +template +__device__ __forceinline__ typename Traits::MainloopPipelineUmma2Sm +mlp4_make_pipe_umma_2sm( + typename Traits::MainloopPipelineUmma2Sm::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipelineUmma2Sm; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = warp_id == 0; + bool is_consumer = warp_id >= 3 && warp_id <= 11; + + typename Pipeline::Params params; + params.transaction_bytes = Traits::TmaTransBytes; + params.num_producers = 1; + params.num_consumers = 1; + params.initializing_warp = 0; + if (is_producer) { + params.role = Category::Producer; + params.is_leader = + threadIdx.x == 0 && cute::block_rank_in_cluster() % 2 == 0; + } else if (is_consumer) { + params.role = Category::Consumer; + } else { + params.role = Category::NonParticipant; + } + return Pipeline( + storage, params, typename Traits::ClusterShape{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// 2SM producer/consumer — chunk-fixed, sequential dB/dC (same cell +// decode as mlp4_producer/mlp4_consumer above), ClusterM=2 only: the +// walk axis (num_m_tiles) is partitioned by `lane` alone (no cluster/pair +// multicast fan-out across multiple pairs). +// ═══════════════════════════════════════════════════════════════════ + +template <> +struct Mlp4ProducerImpl<100> { +template < + typename Traits, + typename Pipeline, + typename TmaLoadX, + typename TmaLoadA, + typename SmemType> +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadX const& tma_load_x, + TmaLoadA const& tma_load_dut, + TmaLoadA const& tma_load_dvt, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + int ring_kb = 0) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + auto sA = make_tensor( + make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); + auto sX = make_tensor( + make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + + auto mX = tma_load_x.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(num_tokens))); + auto mdUT = tma_load_dut.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + auto mdVT = tma_load_dvt.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + + typename Traits::TiledMma2Sm tiled_mma; + int cta_rank = static_cast(cute::block_rank_in_cluster()); + int pair_rank = cta_rank % 2; + auto cta_mma = tiled_mma.get_slice(pair_rank); + // ClusterM=2 → the pair IS the whole cluster. + auto pair_layout_vmnk = tiled_divide( + make_layout(typename Traits::ClusterShape{}), + make_tile(typename Traits::TiledMma2Sm::AtomThrID{})); + auto pair_coord_vmnk = pair_layout_vmnk.get_flat_coord(pair_rank); + uint16_t mcast_mask_a = + create_tma_multicast_mask<2>(pair_layout_vmnk, pair_coord_vmnk); + uint16_t mcast_mask_x = + create_tma_multicast_mask<1>(pair_layout_vmnk, pair_coord_vmnk); + + int total_chunks = num_experts * num_n_tiles; + int total_cells = total_chunks * outer_split * k_split; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e = chunk_idx / num_n_tiles; + int n_tile = chunk_idx - e * num_n_tiles; + // Balanced split: multiply BEFORE dividing so any remainder + // (num_m_tiles % outer_split != 0) is distributed across the + // trailing lanes instead of being silently dropped — e.g. + // num_m_tiles=3, outer_split=2 gives lane0=[0,1), lane1=[1,3), + // covering all 3 tiles. `lane * (num_m_tiles/outer_split)` (a + // naive floor-then-multiply) would leave the tail tile + // unvisited by any lane. ClusterM=2 needs no further + // valid_m_tile guard: walk_end for the last lane is always + // exactly num_m_tiles. + int walk_begin = lane * num_m_tiles / outer_split; + int walk_end = (lane + 1) * num_m_tiles / outer_split; + + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) + continue; + + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int slice_lo = kb_lo + k_slice * k_per; + int slice_hi = min(slice_lo + k_per, kb_hi); + if (slice_hi <= slice_lo) + continue; + kb_lo = slice_lo; + kb_hi = slice_hi; + + CUTE_UNROLL + for (int phase = 0; phase < 2; ++phase) { + auto mA = phase == 0 ? mdUT : mdVT; + const auto& tma_a = + phase == 0 ? tma_load_dut : tma_load_dvt; + + for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { + auto coord = make_coord(m_tile, n_tile, _); + auto gA = local_tile( + mA, typename Traits::TileShape{}, coord, + Step<_1, X, _1>{}); + auto gX = local_tile( + mX, typename Traits::TileShape{}, coord, + Step{}); + auto tCgA = cta_mma.partition_A(gA); + auto tCgX = cta_mma.partition_B(gX); + + auto [tAgA, tAsA] = tma_partition( + tma_a, + get<2>(pair_coord_vmnk), + make_layout(size<2>(pair_layout_vmnk)), + group_modes<0, 3>(sA), + group_modes<0, 3>(tCgA)); + auto [tXgX, tXsX] = tma_partition( + tma_load_x, + get<1>(pair_coord_vmnk), + make_layout(size<1>(pair_layout_vmnk)), + group_modes<0, 3>(sX), + group_modes<0, 3>(tCgX)); + + for (int kb = kb_lo; kb < kb_hi; ++kb) { + int kb_g = ring_kb > 0 ? kb % ring_kb : kb; + pipe.producer_acquire(state); + if (cute::elect_one_sync()) { + auto* barrier = + pipe.producer_get_barrier(state); + copy( + tma_a.with(*barrier, mcast_mask_a), + tAgA(_, kb_g), + tAsA(_, state.index())); + copy( + tma_load_x.with( + *barrier, mcast_mask_x), + tXgX(_, kb_g), + tXsX(_, state.index())); + } + ++state; + } + } + } + } + pipe.producer_tail(state); +#else + __trap(); +#endif +} +}; // Mlp4ProducerImpl<100> + +template <> +struct Mlp4ConsumerImpl<100> { +template < + typename Traits, + bool Expert3D = false, + typename Pipeline, + typename TmaReduceAddDB, + typename TmaReduceAddDC, + typename SmemType> +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDB const& tma_reduce_db, + TmaReduceAddDC const& tma_reduce_dc, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int total_m_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + const cutlass::arch::ClusterBarrier* pair_init_barrier, + uint32_t* pair_init_phase) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + (void)intermediate_dim; + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int CtaTileM = Traits::CtaTileM; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int NumEpiRounds = Traits::NumEpiRounds; + constexpr int kMmaEpiThreads = + Traits::ConsumerThreads + Traits::WarpSize; + + int cta_rank = static_cast(cute::block_rank_in_cluster()); + int pair_rank = cta_rank % 2; + bool is_leader_cta = pair_rank == 0; + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_mma_warp = warp_id == 3; + bool is_epilogue = warp_id >= 4 && warp_id <= 11; + int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; + int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + int tid_in_wg = is_epilogue + ? tid_in_epi % Traits::WarpGroupSize + : 0; + int wg_barrier_id = 1 + wg; + bool is_wg_leader = is_epilogue && tid_in_wg == 0; + + typename Traits::TiledMma2Sm tiled_mma; + auto cta_mma = tiled_mma.get_slice(pair_rank); + auto sA = make_tensor( + make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); + auto sX = make_tensor( + make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto tCrA = cta_mma.make_fragment_A(sA); + auto tCrX = cta_mma.make_fragment_B(sX); + + auto cAccFull = make_identity_tensor( + make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + if (pair_init_barrier != nullptr) { + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + if (warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + pair_init_barrier->arrive(peer, !leader); + pair_init_barrier->wait(*pair_init_phase); + pair_init_barrier->arrive(peer, leader); + *pair_init_phase ^= 1; + } + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + } + + using AccPipe = typename Traits::AccumulatorPipeline2Sm; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp && is_leader_cta + ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; + acc_params.consumer_arv_count = 2; + acc_params.initializing_warp = 4; + AccPipe acc_pipe( + smem.acc_pipe, acc_params, typename Traits::ClusterShape{}); + auto acc_prod_state = + cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + if (pair_init_barrier != nullptr && warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + pair_init_barrier->arrive(peer, !leader); + pair_init_barrier->wait(*pair_init_phase); + pair_init_barrier->arrive(peer, leader); + *pair_init_phase ^= 1; + } + __syncwarp(); + } + if (pair_init_barrier != nullptr) + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + constexpr int store_slot_elems = CtaTileM * EpiChunkN; + Element* store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor( + make_smem_ptr(store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + auto mdB = tma_reduce_db.get_tma_tensor(make_shape( + static_cast(total_m_rows), + static_cast(hidden_dim))); + auto mdC = tma_reduce_dc.get_tma_tensor(make_shape( + static_cast(total_m_rows), + static_cast(hidden_dim))); + auto cta_tma_db = tma_reduce_db.get_slice(Int<0>{}); + auto cta_tma_dc = tma_reduce_dc.get_slice(Int<0>{}); + + auto epi_tile = + make_tile(Int{}, Int{}); + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi = flat_divide(acc_mn, epi_tile); + auto t2r = make_tmem_copy( + TmemLoadOp{}, + tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tid_in_wg); + auto cChunk = make_identity_tensor( + make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); + + Layout tmem_warp_layout = + typename decltype(make_tmem_warp_partitioner( + tAcc_epi(_, _, _0{}, _0{})))::TiledLayout_TV{}; + constexpr bool predicate_tmem_load = + size(tmem_warp_layout) != cosize(tmem_warp_layout); + + int total_chunks = num_experts * num_n_tiles; + int total_cells = total_chunks * outer_split * k_split; + bool store_in_flight = false; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e = chunk_idx / num_n_tiles; + int n_tile = chunk_idx - e * num_n_tiles; + // Match the producer's balanced split: multiply-before-divide + // avoids dropping the tail + // M-tile(s) when num_m_tiles % outer_split != 0. + int walk_begin = lane * num_m_tiles / outer_split; + int walk_end = (lane + 1) * num_m_tiles / outer_split; + + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) + continue; + + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int slice_lo = kb_lo + k_slice * k_per; + int slice_hi = min(slice_lo + k_per, kb_hi); + if (slice_hi <= slice_lo) + continue; + kb_lo = slice_lo; + kb_hi = slice_hi; + + CUTE_UNROLL + for (int phase = 0; phase < 2; ++phase) { + for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { + if (is_mma_warp && is_leader_cta) { + acc_pipe.producer_acquire(acc_prod_state); + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = + tmem_base + uint32_t(acc_stage * TileN); + + bool first_mma = true; + for (int kb = kb_lo; kb < kb_hi; ++kb) { + pipe.consumer_wait(state); + CUTE_UNROLL + for (int ks = 0; + ks < size<2>(tCrA); + ++ks) { + tiled_mma.accumulate_ = first_mma + ? UMMA::ScaleOut::Zero + : UMMA::ScaleOut::One; + first_mma = false; + gemm( + tiled_mma, + tCrA(_, _, ks, state.index()), + tCrX(_, _, ks, state.index()), + tCtAcc); + } + pipe.consumer_release(state); + ++state; + } + acc_pipe.producer_commit(acc_prod_state); + ++acc_prod_state; + } + + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = + tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = + tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = + flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc = + thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int round = 0; + round < NumEpiRounds; + ++round) { + int chunk = wg * NumEpiRounds + round; + auto tAccChunk = + tTR_tAcc(_, _, _, _0{}, chunk); + bool issue_tmem_load = true; + if constexpr (predicate_tmem_load) { + int subpart = + (tAccChunk.data().dp_ / 32) % 4; + issue_tmem_load = + tid_in_wg / Traits::WarpSize == + subpart; + } + if (issue_tmem_load) { + copy(t2r, tAccChunk, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + } + + if (store_in_flight) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, + wg_barrier_id); + + if (issue_tmem_load) { + CUTE_UNROLL + for (int i = 0; + i < size(tTR_rAcc); + ++i) { + int m_local = + get<0>(tTR_cChunk(i)); + int n_local = + get<1>(tTR_cChunk(i)); + sStore(m_local, n_local) = + static_cast( + tTR_rAcc(i)); + } + } + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, + wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int out_m = + e * num_m_tiles + m_tile; + int n_chunk = + n_tile * + (TileN / EpiChunkN) + + wg * NumEpiRounds + + round; + CUTE_UNROLL + for (int atom = 0; + atom < + Traits::kAtomsPerCta; + ++atom) { + int m_atom = + (TileM / + Traits::AtomTileM) * + out_m + + Traits::kAtomsPerCta * + pair_rank + + atom; + auto sStoreAtom = + local_tile( + sStore, + make_tile( + Int< + Traits:: + AtomTileM>{}, + Int< + EpiChunkN>{}), + make_coord( + atom, 0)); + if (phase == 0) { + auto gdB = local_tile( + mdB, + make_tile( + Int< + Traits:: + AtomTileM>{}, + Int< + EpiChunkN>{}), + make_coord( + m_atom, + n_chunk)); + copy( + tma_reduce_db, + cta_tma_db.partition_S( + sStoreAtom), + cta_tma_db.partition_D( + gdB)); + } else { + auto gdC = local_tile( + mdC, + make_tile( + Int< + Traits:: + AtomTileM>{}, + Int< + EpiChunkN>{}), + make_coord( + m_atom, + n_chunk)); + copy( + tma_reduce_dc, + cta_tma_dc.partition_S( + sStoreAtom), + cta_tma_dc.partition_D( + gdC)); + } + } + cute::tma_store_arrive(); + } + store_in_flight = true; + } + + cutlass::arch::NamedBarrier::sync( + Traits::ConsumerThreads, 0); + // One elected release from each CTA satisfies the + // cluster-scoped accumulator pipeline's count of two. + if (tid_in_epi == 0) + acc_pipe.consumer_release( + acc_cons_state); + ++acc_cons_state; + } + } + } + } + + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); +#else + __trap(); +#endif +} +}; // Mlp4ConsumerImpl<100> + +// ═══════════════════════════════════════════════════════════════════ +// Single-launch driver, chunk-fixed and architecture-dispatched +// ═══════════════════════════════════════════════════════════════════ +// +// Compute=90 uses the original per-CTA Hopper pipeline and cell mapping. +// Compute=100 uses a paired-CTA pipeline, pair-granular cell mapping, a +// shared Allocator2Sm allocation, and cluster synchronization. +template < + typename Traits, + int Compute = 90, + bool Expert3D = false, + typename TmaLoadX, + typename TmaLoadA, + typename TmaReduceAddDB, + typename TmaReduceAddDC> +__device__ __forceinline__ void mlp4_fwd( + Mlp4Smem& smem, + TmaLoadX const& tma_load_x, + TmaLoadA const& tma_load_dut, + TmaLoadA const& tma_load_dvt, + TmaReduceAddDB const& tma_reduce_db, + TmaReduceAddDC const& tma_reduce_dc, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int num_tokens, + int total_m_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int batch_kb_start = 0, + int batch_kb_end = -1, + int k_split = 1, + int ring_kb = 0) { + static_assert(Compute == 90 || Compute == 100, + "mlp4_fwd supports Compute=90 or Compute=100"); + using PipeState = typename Traits::PipelineState; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = warp_id == 0; + constexpr int kFirstConsumerWarp = Compute == 100 ? 3 : 4; + bool is_consumer = + warp_id >= kFirstConsumerWarp && warp_id <= 11; + + cute::prefetch_tma_descriptor( + tma_load_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_load_dut.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_load_dvt.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_reduce_db.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_reduce_dc.get_tma_descriptor()); + + auto pipe = [&]() { + if constexpr (Compute == 100) + return mlp4_make_pipe_umma_2sm( + smem.pipe_storage); + else + return mlp4_make_pipe(smem.pipe_storage); + }(); + using Pipeline = decltype(pipe); + + if constexpr (Compute == 100) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + cute::TMEM::Allocator2Sm tmem_allocator; + if (gridDim.x % Traits::ClusterM != 0) + __trap(); + cute::cluster_sync(); + constexpr int kTmemColumns = + Traits::AccStages * Traits::TileN; + if (warp_id == 4) { + tmem_allocator.allocate( + kTmemColumns, &smem.tmem_base); + __syncwarp(); + } + __syncthreads(); + cute::cluster_sync(); +#else + __trap(); +#endif + } else { + __syncthreads(); + } + + PipeState producer_state = is_producer + ? cutlass::make_producer_start_state() + : PipeState{}; + PipeState consumer_state; + + int cell_start = static_cast(blockIdx.x); + int cell_stride = static_cast(gridDim.x); + if constexpr (Compute == 100) { + cell_start /= Traits::ClusterM; + cell_stride /= Traits::ClusterM; + } + int kb_end = batch_kb_end >= 0 + ? batch_kb_end + : num_tokens / Traits::TileK; + + if (is_producer) { + mlp4_producer( + pipe, + producer_state, + smem, + tma_load_x, + tma_load_dut, + tma_load_dvt, + expert_k_starts, + expert_k_ends, + num_experts, + intermediate_dim, + hidden_dim, + num_tokens, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split, + ring_kb); + } else if (is_consumer) { + if constexpr (Compute == 100) { + mlp4_consumer( + pipe, + consumer_state, + smem, + tma_reduce_db, + tma_reduce_dc, + expert_k_starts, + expert_k_ends, + num_experts, + intermediate_dim, + hidden_dim, + total_m_rows, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split); + } else { +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) + mlp4_consumer( + pipe, + consumer_state, + smem, + tma_reduce_db, + tma_reduce_dc, + expert_k_starts, + expert_k_ends, + num_experts, + intermediate_dim, + hidden_dim, + total_m_rows, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split); +#else + __trap(); +#endif + } + } + + if constexpr (Compute == 100) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + cute::TMEM::Allocator2Sm tmem_allocator; + __syncthreads(); + cute::cluster_sync(); + constexpr int kTmemColumns = + Traits::AccStages * Traits::TileN; + if (warp_id == 4) { + tmem_allocator.release_allocation_lock(); + tmem_allocator.free( + smem.tmem_base, kTmemColumns); + } +#endif + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Host launcher (raw-pointer; defined in mlp4.cu) +// ═══════════════════════════════════════════════════════════════════ + +// dB = dU^T @ X, dC = dV^T @ X. Both outputs accumulated atomically +// via TMA_REDUCE_ADD. dB and dC MUST be zero-initialized by caller. +// +// dU : [num_tokens, intermediate_dim] +// dV : [num_tokens, intermediate_dim] +// X : [num_tokens, hidden_dim] +// dB, dC : [num_experts, intermediate_dim, hidden_dim] +// expert_for_k_block: [num_tokens / TileK] int32, TileK=64 +void mlp4_fwd_bf16_launch( + const cutlass::bfloat16_t* dU, + const cutlass::bfloat16_t* dV, + const cutlass::bfloat16_t* X, + cutlass::bfloat16_t* dB, + cutlass::bfloat16_t* dC, + const int* expert_for_k_block, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int device, + cudaStream_t stream); + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp4_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp4_sm90.cuh new file mode 100644 index 00000000..fb0e28f4 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp4_sm90.cuh @@ -0,0 +1,1839 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 4: dB = dU^T @ X, dC = dV^T @ X (weight grads, bwd) +// ═══════════════════════════════════════════════════════════════════ +// +// TMA_REDUCE_ADD variant. Cooperative 2-WG consumer, mirroring mlp3.cuh. +// The two WGs cooperate on ONE output tile at a time, so dB and dC are +// produced SEQUENTIALLY (two phases per cell). This lets a SINGLE A-side +// smem buffer be reused for dU^T (phase dB) then dV^T (phase dC) — vs the +// old design's two simultaneous A buffers — so the footprint matches mlp3 +// and fits the 228 KiB cap at 256-wide tiles. +// +// Only two configs supported (mirrors Mlp3Traits): +// (TileM=256, TileN=128) → M-split (Layout<_2,_1,_1>) ← default +// (TileM=128, TileN=256) → N-split (Layout<_1,_2,_1>) +// +// GEMM orientation: dB = dU^T @ X, so C[M=I, N=H], A[M=I,K=T] = dU^T/dV^T, +// B[N=H,K=T] = X. TileM tiles the intermediate (I) dim, TileN the hidden +// (H) dim. Per-WG output footprint = (WgTileM, WgTileN) = (128, 128). +// +// Single fused A + X TMA pipe (one mbarrier per stage, 2 TMA copies). The +// shared operand X is read by both dB and dC phases and by both WGs. +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for A / X) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Chunk-fixed walk (same principle as mlp3): hold the SHARED operand +// chunk-constant and walk the cooperation (split) axis: +// kMSplit=true (M-split, X shared): chunk = (e, n_tile), walks m_tile; holds X(n,*) +// kMSplit=false (N-split, dU^T/dV^T shared): chunk = (e, m_tile), walks n_tile; holds A(m,*) +// In the M-split default, X(n_tile) is held across BOTH the m-walk and the +// dB/dC phases — re-read for phase dC hits L2. +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm90.cuh" +#include "tmem_load_op_sm90.cuh" + +#include +#include +#include +#include +#include + +// ── SM100 (Compute=100 / Blackwell UMMA) headers ──────────────────── +// Header-only; safe to include under an sm_90a build. The tcgen05 PTX they +// emit is guarded by CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself +// gated on __CUDA_ARCH__ >= 1000) ever instantiates it. Mirrors mlp5.cuh / +// mlp1_fused.cuh. +#include +#include +#include +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +// ── SM100 2-SM (paired-CTA cluster) headers ───────────────────────── +// Needed by the paired MLP4 path below: cluster synchronization/rank queries +// for the (2,1,1) cluster and the 2SM TMA copy traits used by tma_partition(). +#include +#include + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 256, + int TileN_ = 128, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 64 +> +struct Mlp4Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + // Cooperative 2-WG layout. Only two configs supported: + // (TileM=256, TileN=128) → kMSplit=true, WgLayout<_2,_1,_1> + // (TileM=128, TileN=256) → kMSplit=false, WgLayout<_1,_2,_1> + // Per-WG output footprint = (WgTileM=128, WgTileN=128) in both. Each + // WG covers 2 wgmma m=64 atoms per K-step (atom_count_M = 2). + static_assert((TileM == 256 && TileN == 128) || + (TileM == 128 && TileN == 256), + "Mlp4Traits supports only (TileM, TileN) = (256, 128) M-split " + "or (128, 256) N-split"); + + static constexpr bool kMSplit = (TileM > TileN); + static constexpr int AtomTileM = 64; + static constexpr int WgTileM = 128; + static constexpr int WgTileN = 128; + + static_assert(WgTileN % EpiChunkN == 0, + "EpiChunkN must divide WgTileN (=128)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + static constexpr int AtomN = WgTileN; + using GmmaAtom = typename GmmaSelectorMNMN::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtom = GMMA::Layout_MN_SW128_Atom; + + // A-side = dU^T / dV^T (M-side, TileM); X = shared B-side (N-side, TileN). + using SmemLayoutA_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, + Step<_2, _1>{})); + using SmemLayoutX_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{}, + Step<_2, _1>{})); + using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + using SmemLayoutX = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{}, + Step<_2, _1, _3>{})); + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // ── SM100 (Compute=100 / Blackwell UMMA) ──────────────────────── + // The mainloop pipeline is the UMMA-aware TMA pipeline (its consumer_release + // issues the UMMA-gated smem-buffer arrival). Params/PipelineState/ + // SharedStorage alias the Hopper PipelineTmaAsync's, so smem.pipe_storage + // (typed MainloopPipeline::SharedStorage) drives BOTH — the standalone + // launcher passes it to mlp4_make_pipe_umma unchanged. num_consumers=1 on + // Blackwell (one umma_arrive per stage releases the buffer). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Per-WG UMMA atom: SM100_MMA_F16BF16_SS with M=WgTileM=128, N=WgTileN=128, + // BOTH operands MN-major (A = dU^T/dV^T is M-major in smem; X is N-major). + // Warp 3 issues TWO of these per k-step into two independent + // (128,128) TMEM accumulators (one per consumer WG) — see + // Mlp4ConsumerImpl<100>. UMMA requires M ∈ {64,128}, so the 256-wide tile + // cannot use one atom; the per-WG (128,128) atom is used for both configs. + using TiledMmaUmma = decltype(make_tiled_mma( + SM100_MMA_F16BF16_SS{})); + // Double-buffered TMEM accumulators: each stage owns one (WgTileM,WgTileN) + // accumulator per epilogue WG. The two dB/dC phases still clear + // independently — this only lets MMA for the next (phase,walk) run while + // the epilogue drains the previous accumulator stage. + static constexpr int AccStages = 2; + static_assert(AccStages * (2 * WgTileN) <= 512, + "SM100 MLP4 needs AccStages x two per-WG accumulators to fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + static constexpr int TmaTransBytesA = + static_cast(size(SmemLayoutA_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesX = + static_cast(size(SmemLayoutX_1{}) * sizeof(Element)); + // Fused pipeline: X + ONE A operand (dU^T or dV^T, per phase) share one + // mbarrier per stage. Both WGs consume the same X via cooperative TiledMma. + static constexpr int TmaTransBytes = TmaTransBytesA + TmaTransBytesX; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; // 256 + static constexpr int NumThreads = 384; + + // Per-WG store buffer slot: (WgTileM, EpiChunkN). One slot per WG; it + // holds kAtomsPerWg = WgTileM/AtomTileM contiguous atom-rows after the + // acc→smem fold. + static constexpr int kAtomsPerWg = WgTileM / AtomTileM; // 2 + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + // TMA-reduce descriptor box = ONE wgmma atom-row (AtomTileM × EpiChunkN). + // At TileM=256 each WG owns 2 *interleaved* atoms that map to non-contiguous + // gmem rows, so a WG issues kAtomsPerWg separate TMA stores per epi round + // (mirrors mlp1_fused.cuh / mlp3.cuh). + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; +}; + +template +using Mlp4MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Shared Memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp4Smem { + using Element = typename Traits::Element; + + static constexpr int smem_X_size = cosize_v; + static constexpr int smem_A_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + // Single A buffer, reused for dU^T (phase dB) then dV^T (phase dC). + alignas(128) Element smem_X[smem_X_size]; + alignas(128) Element smem_A[smem_A_size]; + // Per-WG store buffer: WG1 uses [0..S-1], WG2 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + typename Traits::MainloopPipeline::SharedStorage pipe_storage; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these; + // SM100 launchers allocate/free TMEM outside the consumer and publish the + // base here; the accumulator pipeline storage still lives in this smem. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* X_data() { return &smem_X[0]; } + CUTE_DEVICE Element* A_data() { return &smem_A[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP4 shared memory (data only — pipe storage lives in the +// outer MlpFusedBwdSmem union; otherwise identical to Mlp4Smem). +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp4FusedSmem { + using Element = typename Traits::Element; + + static constexpr int smem_X_size = cosize_v; + static constexpr int smem_A_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_X[smem_X_size]; + alignas(128) Element smem_A[smem_A_size]; + alignas(128) Element store_buf[2 * smem_store_size]; + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; + + CUTE_DEVICE Element* X_data() { return &smem_X[0]; } + CUTE_DEVICE Element* A_data() { return &smem_A[0]; } +}; + + +// ═══════════════════════════════════════════════════════════════════ +// Pipe construction helper (mirrors mlp3_make_pipe) +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ typename Traits::MainloopPipeline +mlp4_make_pipe(typename Traits::MainloopPipeline::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp4_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the (shared) +// mlp4_producer and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ typename Traits::MainloopPipelineUmma +mlp4_make_pipe_umma(typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// Chunk-fixed cooperative producer/consumer — sequential dB / dC +// ═══════════════════════════════════════════════════════════════════ +// +// Used by the standalone mlp4 kernel (mlp4.cu) and mlp_bwd Phase 2. +// Each CTA owns ONE (chunk, lane) cell and processes it in TWO phases +// (phase 0 = dB via dU^T, phase 1 = dC via dV^T). We hold the SHARED +// operand and walk the cooperation (split) axis: +// kMSplit=true (M-split): chunk = (e, n_tile), lane partitions +// num_m_tiles, walks m_tile. X(n_tile, *) — the shared +// N-side operand — is the L2-held chunk-constant, +// reused across the m-walk AND both phases. +// kMSplit=false (N-split): chunk = (e, m_tile), lane partitions +// num_n_tiles, walks n_tile. A(m_tile, *) (dU^T then +// dV^T) is L2-held; X is re-streamed per phase. +// +// Cell decode (chunk-major: cell_idx = chunk_idx × outer_split + lane): +// total_cells = total_chunks × outer_split +// cell → chunk_idx = cell_idx / outer_split +// lane = cell_idx % outer_split + +template +struct Mlp4ProducerImpl; + +template <> +struct Mlp4ProducerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadX const& tma_load_x, + TmaLoadA const& tma_load_dut, + TmaLoadA const& tma_load_dvt, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + // Staging-ring depth in K-blocks; > 0 wraps the gmem K-tile coordinate + // (% ring_kb) for the grouped Phase-2 that may straddle the ring. See + // mlp3_producer. 0 = no wrap. + int ring_kb = 0) { + + auto sA = make_tensor(make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + + // int64_t casts: H·T, I·T can exceed 2^31 elems at production shapes. + auto mX = tma_load_x.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(num_tokens))); + auto mdUT = tma_load_dut.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + auto mdVT = tma_load_dvt.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + + auto cta_tma_x = tma_load_x.get_slice(Int<0>{}); + auto cta_tma_dut = tma_load_dut.get_slice(Int<0>{}); + auto cta_tma_dvt = tma_load_dvt.get_slice(Int<0>{}); + + auto tXsX = cta_tma_x.partition_D(sX); + auto tAsA = cta_tma_dut.partition_D(sA); // dU^T/dV^T share A layout + + int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); + // K-split: each (chunk, walk-lane) replicated k_split times; k_slice owns a + // 1/k_split sub-range of the K-loop. Partials sum via TMA_REDUCE_ADD. + int total_cells = total_chunks * outer_split * k_split; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e, m_chunk, n_chunk, walk_begin, walk_end; + if constexpr (Traits::kMSplit) { + // M-split: hold X(n) shared. chunk = (e, n_tile), walks m_tile. + e = chunk_idx / num_n_tiles; + n_chunk = chunk_idx - e * num_n_tiles; + m_chunk = 0; + int m_per = num_m_tiles / outer_split; + walk_begin = lane * m_per; + walk_end = walk_begin + m_per; + } else { + // N-split: hold A(m) shared. chunk = (e, m_tile), walks n_tile. + e = chunk_idx / num_m_tiles; + m_chunk = chunk_idx - e * num_m_tiles; + n_chunk = 0; + int n_per = num_n_tiles / outer_split; + walk_begin = lane * n_per; + walk_end = walk_begin + n_per; + } + + // Intersect this expert's global K range with the current + // batch's K-window. + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) continue; + + // K-split: clip to this cell's k_slice sub-range of [kb_lo, kb_hi). + { + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int s_lo = kb_lo + k_slice * k_per; + int s_hi = min(s_lo + k_per, kb_hi); + if (s_hi <= s_lo) continue; // empty slice + kb_lo = s_lo; + kb_hi = s_hi; + } + + // Two phases per cell: phase 0 = dB (dU^T), phase 1 = dC (dV^T). + // dU^T/dV^T share descriptor layout; pair the tma object, its slice, + // and its gmem tensor consistently per phase. + CUTE_UNROLL + for (int phase = 0; phase < 2; ++phase) { + auto mA = (phase == 0) ? mdUT : mdVT; + const auto& tma_A = (phase == 0) ? tma_load_dut : tma_load_dvt; + const auto& cta_tma_A = (phase == 0) ? cta_tma_dut : cta_tma_dvt; + + for (int w = walk_begin; w < walk_end; ++w) { + int m_tile = Traits::kMSplit ? w : m_chunk; + int n_tile = Traits::kMSplit ? n_chunk : w; + auto gX = local_tile(mX, + make_tile(Int{}, Int{}), + make_coord(n_tile, _)); + auto gA = local_tile(mA, + make_tile(Int{}, Int{}), + make_coord(m_tile, _)); + auto tXgX = cta_tma_x.partition_S(gX); + auto tAgA = cta_tma_A.partition_S(gA); + + for (int kb = kb_lo; kb < kb_hi; ++kb) { + int kb_g = (ring_kb > 0) ? (kb % ring_kb) : kb; // physical ring slot + pipe.producer_acquire(state); + if (threadIdx.x == 0) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_x.with(*bar, 0), + tXgX(_, _, _, kb_g), tXsX(_, _, _, state.index())); + copy(tma_A.with(*bar, 0), + tAgA(_, _, _, kb_g), tAsA(_, _, _, state.index())); + } + ++state; + } + } + } + } +} +}; // Mlp4ProducerImpl<90> + +template +__device__ __forceinline__ void mlp4_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadX const& tma_load_x, + TmaLoadA const& tma_load_dut, + TmaLoadA const& tma_load_dvt, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + int ring_kb = 0) { + Mlp4ProducerImpl::template run( + pipe, state, smem, tma_load_x, tma_load_dut, tma_load_dvt, + expert_k_starts, expert_k_ends, num_experts, + intermediate_dim, hidden_dim, num_tokens, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split, ring_kb); +} + +// ═══════════════════════════════════════════════════════════════════ +// Consumer — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported compute +// capability provides a `run(...)` member (a function-template on the Traits/ +// Pipeline/Smem/Tma types, since function templates can't be partially +// specialized). The free function `mlp4_consumer` below forwards to the right +// specialization — call sites (mlp_bwd.cuh's mlp4_consumer(...)) stay +// unchanged (Compute defaults to 90). +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, register acc, M-atom +// interleave + (m_loc/128)*64+(m_loc%64) store remap) +// Compute=100 → Blackwell / UMMA (warp 3 issues 2 per-WG (128,128) UMMAs +// into double-buffered TMEM stages; warps 4-11 run the +// epilogue, with independent per-phase clear and contiguous +// per-WG store — no remap) + +template +struct Mlp4ConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper / WGMMA. (Verbatim from the original mlp4_consumer.) +// ─────────────────────────────────────────────────────────────────── +template <> +struct Mlp4ConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDB const& tma_reduce_db, + TmaReduceAddDC const& tma_reduce_dc, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int total_m_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split) { + + (void)Compute; + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; // 0 or 1 + const int my_barrier_id = 1 + my_wg; // 1 or 2 + const int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; // 0..255 + const int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; // 0..127 + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sA = make_tensor(make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); + auto sX = make_tensor(make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + + auto tCsA = thr_mma.partition_A(sA); + auto tCsX = thr_mma.partition_B(sX); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + // Identity coords for per-thread acc → store_buf mapping. + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + constexpr int store_slot_elems = Traits::WgTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + const bool is_my_wg_leader = (tid_in_wg == 0); + + auto mdB = [&]() { + if constexpr (Expert3D) { + return tma_reduce_db.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_reduce_db.get_tma_tensor(make_shape( + static_cast(total_m_rows), + static_cast(hidden_dim))); + } + }(); + auto mdC = [&]() { + if constexpr (Expert3D) { + return tma_reduce_dc.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(num_experts))); + } else { + return tma_reduce_dc.get_tma_tensor(make_shape( + static_cast(total_m_rows), + static_cast(hidden_dim))); + } + }(); + auto cta_tma_db = tma_reduce_db.get_slice(Int<0>{}); + auto cta_tma_dc = tma_reduce_dc.get_slice(Int<0>{}); + + int total_chunks = num_experts * (Traits::kMSplit ? num_n_tiles : num_m_tiles); + int total_cells = total_chunks * outer_split * k_split; + + constexpr int K_PIPE_MMAS = 1; + + // Carries across phases, walk iters AND cells for cross-tile K-loop / + // store overlap. Drained once at the end. + bool store_in_flight = false; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e, m_chunk, n_chunk, walk_begin, walk_end; + if constexpr (Traits::kMSplit) { + // M-split: hold X(n) shared. chunk = (e, n_tile), walks m_tile. + e = chunk_idx / num_n_tiles; + n_chunk = chunk_idx - e * num_n_tiles; + m_chunk = 0; + int m_per = num_m_tiles / outer_split; + walk_begin = lane * m_per; + walk_end = walk_begin + m_per; + } else { + // N-split: hold A(m) shared. chunk = (e, m_tile), walks n_tile. + e = chunk_idx / num_m_tiles; + m_chunk = chunk_idx - e * num_m_tiles; + n_chunk = 0; + int n_per = num_n_tiles / outer_split; + walk_begin = lane * n_per; + walk_end = walk_begin + n_per; + } + + // Intersect this expert's global K range with the current + // batch's K-window. + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) continue; + + // K-split: clip to this cell's k_slice sub-range (must match producer). + { + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int s_lo = kb_lo + k_slice * k_per; + int s_hi = min(s_lo + k_per, kb_hi); + if (s_hi <= s_lo) continue; // empty slice + kb_lo = s_lo; + kb_hi = s_hi; + } + + int k_count = kb_hi - kb_lo; + + // Two phases per cell: phase 0 = dB, phase 1 = dC. + CUTE_UNROLL + for (int phase = 0; phase < 2; ++phase) { + + for (int w = walk_begin; w < walk_end; ++w) { + int m_tile = Traits::kMSplit ? w : m_chunk; + int n_tile = Traits::kMSplit ? n_chunk : w; + + clear(acc); + auto state_release = state; + int prologue_count = (k_count < K_PIPE_MMAS) ? k_count : K_PIPE_MMAS; + + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsA(_, _, _, state.index()), + tCsX(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + + for (int k = prologue_count; k < k_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsA(_, _, _, state.index()), + tCsX(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + + ++state; + ++state_release; + } + + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + // Each WG writes its (WgTileM × EpiChunkN) slot to store_buf + // and TMA_REDUCE_ADDs it into dB (phase 0) or dC (phase 1). + // Store indexing mirrors mlp3_consumer (M-split / N-split). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = (m_loc / 128) * 64 + (m_loc % 64); + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int out_m = Expert3D ? m_tile : e * num_m_tiles + m_tile; + // Store box = ONE atom-row (AtomTileM); kAtomsPerWg stores + // per WG. M-split: the 2 atoms interleave in gmem + // (4*out_m + my_wg + 2a). N-split: WG owns the full TileM + // (2 contiguous atoms) and its own N-half (2*out_m + a). + int n_tile_idx; + if constexpr (Traits::kMSplit) { + n_tile_idx = n_tile * Traits::NumEpiRounds + r; + } else { + n_tile_idx = n_tile * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + constexpr int kRowAtoms = Traits::TileM / Traits::AtomTileM; // 4 (M-split) / 2 (N-split) + CUTE_UNROLL + for (int a = 0; a < Traits::kAtomsPerWg; ++a) { + int m_atom_row = Traits::kMSplit + ? (kRowAtoms * out_m + my_wg + Traits::kAtomsPerWg * a) + : (kRowAtoms * out_m + a); + auto sStore_a = local_tile(sStore, + make_tile(Int{}, Int{}), + make_coord(a, 0)); + if (phase == 0) { + auto gdB = [&]() { + if constexpr (Expert3D) { + return local_tile(mdB, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx, e)); + } else { + return local_tile(mdB, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx)); + } + }(); + copy(tma_reduce_db, cta_tma_db.partition_S(sStore_a), + cta_tma_db.partition_D(gdB)); + } else { + auto gdC = [&]() { + if constexpr (Expert3D) { + return local_tile(mdC, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx, e)); + } else { + return local_tile(mdC, + make_tile(Int{}, Int{}), + make_coord(m_atom_row, n_tile_idx)); + } + }(); + copy(tma_reduce_dc, cta_tma_dc.partition_S(sStore_a), + cta_tma_dc.partition_D(gdC)); + } + } + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // struct Mlp4ConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Free-function forwarder. Existing call sites (mlp4_consumer(...)) +// are unchanged (Compute defaults to 90); pass mlp4_consumer(...) +// to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── +template +__device__ __forceinline__ void mlp4_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDB const& tma_reduce_db, + TmaReduceAddDC const& tma_reduce_dc, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int total_m_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, + uint32_t* pair_init_phase = nullptr) { + if constexpr (Compute == 100) { + Mlp4ConsumerImpl::template run( + pipe, state, smem, tma_reduce_db, tma_reduce_dc, + expert_k_starts, expert_k_ends, num_experts, + intermediate_dim, hidden_dim, total_m_rows, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, batch_kb_start, batch_kb_end, k_split, + pair_init_barrier, pair_init_phase); + } else { + Mlp4ConsumerImpl::template run( + pipe, state, smem, tma_reduce_db, tma_reduce_dc, + expert_k_starts, expert_k_ends, num_experts, + intermediate_dim, hidden_dim, total_m_rows, + num_m_tiles, num_n_tiles, outer_split, + cell_start, cell_stride, batch_kb_start, batch_kb_end, k_split); + } +} + +// ═══════════════════════════════════════════════════════════════════ +// SM100 2-SM (paired-CTA) MLP4 — Mlp4Traits2Sm, pipe/producer/consumer +// ═══════════════════════════════════════════════════════════════════ +// +// A cluster-(2,1,1) SM100 UMMA path: two CTAs jointly compute one +// (TileM=256, TileN=256) tile via SM100_MMA_F16BF16_2x1SM_SS, each CTA +// owning CtaTileM=128 M-rows (kAtomsPerCta=2 wgmma-style 64-row atoms). +// Reuses the generic Mlp4Smem / Mlp4FusedSmem above — +// Mlp4Traits2Sm supplies the same nested type names (SmemLayoutA, +// SmemLayoutX, SmemLayoutStoreSlot, MainloopPipeline, AccumulatorPipeline) +// so no separate 2SM smem struct is needed. +// +// Scope: only the BASIC ClusterM=2 (single paired-CTA) configuration +// required by the fused backward kernel is ported here. Intentionally +// EXCLUDED (upstream "topology experiment" scope, not ported): the +// multi-pair (ClusterM=4) path, shared-operand multicast +// (MulticastShared_), N-split, the compact-epilogue variant, and early +// TMEM release — the latter three never existed for MLP4 in the donor +// tree (MLP3-only experiments) and are noted here only for the audit +// record. Stages and EpiChunkN remain free template parameters (no +// S5/E64 production hardcoding); callers pick tuned values at the +// trait-instantiation site (e.g. moe_bwd.cu). +// +// Cross-file invariants a caller/launcher MUST uphold (none of this is +// enforceable from inside this header-only file; documented for the +// next integrator of mlp_bwd.cuh / moe_bwd.cu / mlp4.cu): +// * grid.x must be a MULTIPLE OF ClusterM (=2) before launch; the runtime +// `__trap()` check in mlp4_fwd is only a defensive backstop. +// * Compute != 100 must stay on the 1SM traits/path. Mlp4Traits2Sm's +// TiledMma2Sm / ClusterShape / MainloopPipelineUmma2Sm / +// AccumulatorPipeline2Sm members do not exist on Mlp4Traits, so mixing +// an architecture with the wrong trait remains a compile-time error. +// * Launch must use cudaLaunchKernelEx with a cluster dimension of +// (ClusterM,1,1), and cudaFuncSetAttribute(..., +// cudaFuncAttributeNonPortableClusterSizeAllowed, 1) must be set on +// the kernel beforehand. +// * The dynamic shared-memory allocation for Mlp4Smem> +// must be sized/capped against the device's *dynamic* opt-in limit +// (cudaDeviceGetAttribute(..., cudaDevAttrMaxSharedMemoryPerBlockOptin, +// ...)), not the static 48 KiB default, via +// cudaFuncAttributeMaxDynamicSharedMemorySize. +// * X / dU^T / dV^T TMA descriptors must be built with +// make_tma_copy_A_sm100 / make_tma_copy_B_sm100 (SM100_TMA_2SM_LOAD +// copy op) — the plain SM90 TMA copy-atom builders are not +// pair-aware and would not multicast/replicate operands across the +// two peer CTAs. dB / dC stay TMA REDUCE_ADD outputs. +// * The producer/consumer specializations assume TMEM is already allocated +// and smem.tmem_base populated. The standalone mlp4_fwd driver owns that +// allocation and rendezvouses peer CTAs with cute::cluster_sync() (safe +// because both CTAs run symmetric control flow). A fused caller sharing +// one TMEM region across warp groups with asymmetric control flow must not +// reuse a bare cluster_sync() for that rendezvous; it must pair +// cute::TMEM::Allocator2Sm's +// allocate()/free() with its own dedicated +// cutlass::arch::ClusterBarrier rendezvous scoped to just the +// participating warps, the way donor's mlp_bwd.cuh +// MlpFusedBwdSmem::tmem_pair_barrier does. That smem field and its +// init/arrive/wait call sites live in mlp_bwd.cuh, outside this +// file's scope, and are not present in this target repo yet. + +template < + typename Element_, + int TileM_ = 256, + int TileN_ = 256, + int TileK_ = 64, + int Stages_ = 3, + int EpiChunkN_ = 128, + int AccStages_ = 2, + int ClusterM_ = 2 +> +struct Mlp4Traits2Sm { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + static constexpr int AccStages = AccStages_; + static constexpr int ClusterM = ClusterM_; + static constexpr int NumPairs = ClusterM / 2; + static constexpr int CtaTileM = TileM / 2; + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = TileN / 2; + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + static constexpr int kAtomsPerCta = CtaTileM / AtomTileM; + + static_assert(TileM == 256, "MLP4 2SM requires joined TileM=256"); + static_assert(TileN == 256, "MLP4 2SM requires joined TileN=256"); + static_assert(TileK == 32 || TileK == 64, + "MLP4 2SM supports TileK=32 or 64"); + static_assert(Stages >= 2 && Stages <= 12, + "MLP4 2SM supports two to twelve mainloop stages"); + static_assert(CtaTileM == 128, "each peer CTA must own 128 M rows"); + static_assert(WgTileN % EpiChunkN == 0, + "EpiChunkN must divide each warpgroup's N half"); + static_assert(kAtomsPerCta == 2, + "each CTA must reduce-add two 64-row atoms"); + static_assert(AccStages == 2, + "MLP4 2SM fixes accumulator double buffering"); + static_assert(AccStages * TileN <= 512, + "MLP4 2SM accumulator stages must fit in 512 TMEM columns"); + static_assert(ClusterM == 2, + "MLP4 2SM ports only the basic single paired-CTA cluster; " + "ClusterM=4 pair-of-pairs multicast is an excluded experiment"); + + using TileShape = Shape, Int, Int>; + using ClusterShape = Shape, _1, _1>; + using AtomThrShape = Shape<_2, _1, _1>; + + using TiledMma2Sm = decltype(make_tiled_mma( + SM100_MMA_F16BF16_2x1SM_SS< + Element, Element, ElementAccum, TileM, TileN, + UMMA::Major::MN, UMMA::Major::MN>{})); + static_assert(size(typename TiledMma2Sm::AtomThrID{}) == 2, + "MLP4 2SM requires a two-CTA MMA atom"); + + using MmaShapeA_MK = decltype(partition_shape_A( + TiledMma2Sm{}, make_shape(Int{}, Int{}))); + using MmaShapeB_NK = decltype(partition_shape_B( + TiledMma2Sm{}, make_shape(Int{}, Int{}))); + + using SmemLayoutAtom = UMMA::Layout_MN_SW128_Atom; + using SmemLayoutA = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtom{}, append(MmaShapeA_MK{}, Int{}), + Step<_2, _1, _3>{})); + using SmemLayoutX = decltype(UMMA::tile_to_mma_shape( + SmemLayoutAtom{}, append(MmaShapeB_NK{}, Int{}), + Step<_2, _1, _3>{})); + using SmemLayoutA_1 = decltype(SmemLayoutA{}(_, _, _, Int<0>{})); + using SmemLayoutX_1 = decltype(SmemLayoutX{}(_, _, _, Int<0>{})); + + static constexpr int TmaTransBytesA = + static_cast(cosize_v * sizeof(Element)); + static constexpr int TmaTransBytesX = + static_cast(cosize_v * sizeof(Element)); + // Both peer CTAs issue each 2SM load, and both contributions arrive on + // the cluster-scoped full barrier. + static constexpr int TmaTransBytes = + 2 * (TmaTransBytesA + TmaTransBytesX); + + using MainloopPipelineUmma2Sm = cutlass::PipelineTmaUmmaAsync< + Stages, ClusterShape, AtomThrShape>; + // Aliased to the names the generic Mlp4Smem/Mlp4FusedSmem + // (and Mlp4MainloopPipelineFor) look up. + using MainloopPipelineUmma = MainloopPipelineUmma2Sm; + using MainloopPipeline = MainloopPipelineUmma2Sm; + using PipelineState = typename MainloopPipelineUmma2Sm::PipelineState; + using AccumulatorPipeline2Sm = cutlass::PipelineUmmaAsync< + AccStages, AtomThrShape>; + using AccumulatorPipeline = AccumulatorPipeline2Sm; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; + + using SmemLayoutStoreSlot = Layout< + Shape, Int>, + Stride, _1>>; + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; +}; + +// Compatibility alias: Mlp4Traits2Sm's nested type names line up +// field-for-field with the generic Mlp4Smem above (see the +// comment atop this section), so no bespoke 2SM smem struct is defined — +// this alias just lets callers/tooling written against the name +// "Mlp4Smem2Sm" (e.g. bench_opt/bench_mlp4_2sm.cu) resolve to the +// same generic struct without any source changes on their end. +template +using Mlp4Smem2Sm = Mlp4Smem; + +// SM100 2SM mainloop pipe. Mirrors mlp4_make_pipe_umma; num_consumers=1 +// (one umma_arrive per stage releases the buffer). is_leader restricts the +// TMA-issuing thread to lane 0 of the pair's rank-0 CTA — both peer CTAs' +// smem buffers arrive on the shared cluster barrier via the 2SM TMA +// multicast, so only one CTA needs to issue the copy. +template +__device__ __forceinline__ typename Traits::MainloopPipelineUmma2Sm +mlp4_make_pipe_umma_2sm( + typename Traits::MainloopPipelineUmma2Sm::SharedStorage& storage) { + using Pipeline = typename Traits::MainloopPipelineUmma2Sm; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = warp_id == 0; + bool is_consumer = warp_id >= 3 && warp_id <= 11; + + typename Pipeline::Params params; + params.transaction_bytes = Traits::TmaTransBytes; + params.num_producers = 1; + params.num_consumers = 1; + params.initializing_warp = 0; + if (is_producer) { + params.role = Category::Producer; + params.is_leader = + threadIdx.x == 0 && cute::block_rank_in_cluster() % 2 == 0; + } else if (is_consumer) { + params.role = Category::Consumer; + } else { + params.role = Category::NonParticipant; + } + return Pipeline( + storage, params, typename Traits::ClusterShape{}, + cute::true_type{}, cute::true_type{}); +} + +// ═══════════════════════════════════════════════════════════════════ +// 2SM producer/consumer — chunk-fixed, sequential dB/dC (same cell +// decode as mlp4_producer/mlp4_consumer above), ClusterM=2 only: the +// walk axis (num_m_tiles) is partitioned by `lane` alone (no cluster/pair +// multicast fan-out across multiple pairs). +// ═══════════════════════════════════════════════════════════════════ + +template <> +struct Mlp4ProducerImpl<100> { +template < + typename Traits, + typename Pipeline, + typename TmaLoadX, + typename TmaLoadA, + typename SmemType> +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaLoadX const& tma_load_x, + TmaLoadA const& tma_load_dut, + TmaLoadA const& tma_load_dvt, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int num_tokens, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + int ring_kb = 0) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + auto sA = make_tensor( + make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); + auto sX = make_tensor( + make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + + auto mX = tma_load_x.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(num_tokens))); + auto mdUT = tma_load_dut.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + auto mdVT = tma_load_dvt.get_tma_tensor(make_shape( + static_cast(intermediate_dim), + static_cast(num_tokens))); + + typename Traits::TiledMma2Sm tiled_mma; + int cta_rank = static_cast(cute::block_rank_in_cluster()); + int pair_rank = cta_rank % 2; + auto cta_mma = tiled_mma.get_slice(pair_rank); + // ClusterM=2 → the pair IS the whole cluster. + auto pair_layout_vmnk = tiled_divide( + make_layout(typename Traits::ClusterShape{}), + make_tile(typename Traits::TiledMma2Sm::AtomThrID{})); + auto pair_coord_vmnk = pair_layout_vmnk.get_flat_coord(pair_rank); + uint16_t mcast_mask_a = + create_tma_multicast_mask<2>(pair_layout_vmnk, pair_coord_vmnk); + uint16_t mcast_mask_x = + create_tma_multicast_mask<1>(pair_layout_vmnk, pair_coord_vmnk); + + int total_chunks = num_experts * num_n_tiles; + int total_cells = total_chunks * outer_split * k_split; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e = chunk_idx / num_n_tiles; + int n_tile = chunk_idx - e * num_n_tiles; + // Balanced split: multiply BEFORE dividing so any remainder + // (num_m_tiles % outer_split != 0) is distributed across the + // trailing lanes instead of being silently dropped — e.g. + // num_m_tiles=3, outer_split=2 gives lane0=[0,1), lane1=[1,3), + // covering all 3 tiles. `lane * (num_m_tiles/outer_split)` (a + // naive floor-then-multiply) would leave the tail tile + // unvisited by any lane. ClusterM=2 needs no further + // valid_m_tile guard: walk_end for the last lane is always + // exactly num_m_tiles. + int walk_begin = lane * num_m_tiles / outer_split; + int walk_end = (lane + 1) * num_m_tiles / outer_split; + + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) + continue; + + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int slice_lo = kb_lo + k_slice * k_per; + int slice_hi = min(slice_lo + k_per, kb_hi); + if (slice_hi <= slice_lo) + continue; + kb_lo = slice_lo; + kb_hi = slice_hi; + + CUTE_UNROLL + for (int phase = 0; phase < 2; ++phase) { + auto mA = phase == 0 ? mdUT : mdVT; + const auto& tma_a = + phase == 0 ? tma_load_dut : tma_load_dvt; + + for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { + auto coord = make_coord(m_tile, n_tile, _); + auto gA = local_tile( + mA, typename Traits::TileShape{}, coord, + Step<_1, X, _1>{}); + auto gX = local_tile( + mX, typename Traits::TileShape{}, coord, + Step{}); + auto tCgA = cta_mma.partition_A(gA); + auto tCgX = cta_mma.partition_B(gX); + + auto [tAgA, tAsA] = tma_partition( + tma_a, + get<2>(pair_coord_vmnk), + make_layout(size<2>(pair_layout_vmnk)), + group_modes<0, 3>(sA), + group_modes<0, 3>(tCgA)); + auto [tXgX, tXsX] = tma_partition( + tma_load_x, + get<1>(pair_coord_vmnk), + make_layout(size<1>(pair_layout_vmnk)), + group_modes<0, 3>(sX), + group_modes<0, 3>(tCgX)); + + for (int kb = kb_lo; kb < kb_hi; ++kb) { + int kb_g = ring_kb > 0 ? kb % ring_kb : kb; + pipe.producer_acquire(state); + if (cute::elect_one_sync()) { + auto* barrier = + pipe.producer_get_barrier(state); + copy( + tma_a.with(*barrier, mcast_mask_a), + tAgA(_, kb_g), + tAsA(_, state.index())); + copy( + tma_load_x.with( + *barrier, mcast_mask_x), + tXgX(_, kb_g), + tXsX(_, state.index())); + } + ++state; + } + } + } + } + pipe.producer_tail(state); +#else + __trap(); +#endif +} +}; // Mlp4ProducerImpl<100> + +template <> +struct Mlp4ConsumerImpl<100> { +template < + typename Traits, + bool Expert3D = false, + typename Pipeline, + typename TmaReduceAddDB, + typename TmaReduceAddDC, + typename SmemType> +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + SmemType& smem, + TmaReduceAddDB const& tma_reduce_db, + TmaReduceAddDC const& tma_reduce_dc, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int total_m_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int cell_start, + int cell_stride, + int batch_kb_start, + int batch_kb_end, + int k_split, + const cutlass::arch::ClusterBarrier* pair_init_barrier, + uint32_t* pair_init_phase) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + (void)intermediate_dim; + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int CtaTileM = Traits::CtaTileM; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int NumEpiRounds = Traits::NumEpiRounds; + constexpr int kMmaEpiThreads = + Traits::ConsumerThreads + Traits::WarpSize; + + int cta_rank = static_cast(cute::block_rank_in_cluster()); + int pair_rank = cta_rank % 2; + bool is_leader_cta = pair_rank == 0; + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_mma_warp = warp_id == 3; + bool is_epilogue = warp_id >= 4 && warp_id <= 11; + int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; + int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + int tid_in_wg = is_epilogue + ? tid_in_epi % Traits::WarpGroupSize + : 0; + int wg_barrier_id = 1 + wg; + bool is_wg_leader = is_epilogue && tid_in_wg == 0; + + typename Traits::TiledMma2Sm tiled_mma; + auto cta_mma = tiled_mma.get_slice(pair_rank); + auto sA = make_tensor( + make_smem_ptr(smem.A_data()), typename Traits::SmemLayoutA{}); + auto sX = make_tensor( + make_smem_ptr(smem.X_data()), typename Traits::SmemLayoutX{}); + auto tCrA = cta_mma.make_fragment_A(sA); + auto tCrX = cta_mma.make_fragment_B(sX); + + auto cAccFull = make_identity_tensor( + make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + if (pair_init_barrier != nullptr) { + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + if (warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + pair_init_barrier->arrive(peer, !leader); + pair_init_barrier->wait(*pair_init_phase); + pair_init_barrier->arrive(peer, leader); + *pair_init_phase ^= 1; + } + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + } + + using AccPipe = typename Traits::AccumulatorPipeline2Sm; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp && is_leader_cta + ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; + acc_params.consumer_arv_count = 2; + acc_params.initializing_warp = 4; + AccPipe acc_pipe( + smem.acc_pipe, acc_params, typename Traits::ClusterShape{}); + auto acc_prod_state = + cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + if (pair_init_barrier != nullptr && warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + pair_init_barrier->arrive(peer, !leader); + pair_init_barrier->wait(*pair_init_phase); + pair_init_barrier->arrive(peer, leader); + *pair_init_phase ^= 1; + } + __syncwarp(); + } + if (pair_init_barrier != nullptr) + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); + uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + constexpr int store_slot_elems = CtaTileM * EpiChunkN; + Element* store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor( + make_smem_ptr(store_ptr), + typename Traits::SmemLayoutStoreSlot{}); + + auto mdB = tma_reduce_db.get_tma_tensor(make_shape( + static_cast(total_m_rows), + static_cast(hidden_dim))); + auto mdC = tma_reduce_dc.get_tma_tensor(make_shape( + static_cast(total_m_rows), + static_cast(hidden_dim))); + auto cta_tma_db = tma_reduce_db.get_slice(Int<0>{}); + auto cta_tma_dc = tma_reduce_dc.get_slice(Int<0>{}); + + auto epi_tile = + make_tile(Int{}, Int{}); + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi = flat_divide(acc_mn, epi_tile); + auto t2r = make_tmem_copy( + TmemLoadOp{}, + tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tid_in_wg); + auto cChunk = make_identity_tensor( + make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); + + Layout tmem_warp_layout = + typename decltype(make_tmem_warp_partitioner( + tAcc_epi(_, _, _0{}, _0{})))::TiledLayout_TV{}; + constexpr bool predicate_tmem_load = + size(tmem_warp_layout) != cosize(tmem_warp_layout); + + int total_chunks = num_experts * num_n_tiles; + int total_cells = total_chunks * outer_split * k_split; + bool store_in_flight = false; + + for (int cell_idx = cell_start; + cell_idx < total_cells; + cell_idx += cell_stride) { + int k_slice = cell_idx % k_split; + int cell_om = cell_idx / k_split; + int chunk_idx = cell_om / outer_split; + int lane = cell_om - chunk_idx * outer_split; + + int e = chunk_idx / num_n_tiles; + int n_tile = chunk_idx - e * num_n_tiles; + // Match the producer's balanced split: multiply-before-divide + // avoids dropping the tail + // M-tile(s) when num_m_tiles % outer_split != 0. + int walk_begin = lane * num_m_tiles / outer_split; + int walk_end = (lane + 1) * num_m_tiles / outer_split; + + int kb_lo = max(expert_k_starts[e], batch_kb_start); + int kb_hi = min(expert_k_ends[e], batch_kb_end); + if (kb_hi <= kb_lo) + continue; + + int k_total = kb_hi - kb_lo; + int k_per = (k_total + k_split - 1) / k_split; + int slice_lo = kb_lo + k_slice * k_per; + int slice_hi = min(slice_lo + k_per, kb_hi); + if (slice_hi <= slice_lo) + continue; + kb_lo = slice_lo; + kb_hi = slice_hi; + + CUTE_UNROLL + for (int phase = 0; phase < 2; ++phase) { + for (int m_tile = walk_begin; m_tile < walk_end; ++m_tile) { + if (is_mma_warp && is_leader_cta) { + acc_pipe.producer_acquire(acc_prod_state); + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = + tmem_base + uint32_t(acc_stage * TileN); + + bool first_mma = true; + for (int kb = kb_lo; kb < kb_hi; ++kb) { + pipe.consumer_wait(state); + CUTE_UNROLL + for (int ks = 0; + ks < size<2>(tCrA); + ++ks) { + tiled_mma.accumulate_ = first_mma + ? UMMA::ScaleOut::Zero + : UMMA::ScaleOut::One; + first_mma = false; + gemm( + tiled_mma, + tCrA(_, _, ks, state.index()), + tCrX(_, _, ks, state.index()), + tCtAcc); + } + pipe.consumer_release(state); + ++state; + } + acc_pipe.producer_commit(acc_prod_state); + ++acc_prod_state; + } + + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = + tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = + tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = + flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc = + thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int round = 0; + round < NumEpiRounds; + ++round) { + int chunk = wg * NumEpiRounds + round; + auto tAccChunk = + tTR_tAcc(_, _, _, _0{}, chunk); + bool issue_tmem_load = true; + if constexpr (predicate_tmem_load) { + int subpart = + (tAccChunk.data().dp_ / 32) % 4; + issue_tmem_load = + tid_in_wg / Traits::WarpSize == + subpart; + } + if (issue_tmem_load) { + copy(t2r, tAccChunk, tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + } + + if (store_in_flight) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, + wg_barrier_id); + + if (issue_tmem_load) { + CUTE_UNROLL + for (int i = 0; + i < size(tTR_rAcc); + ++i) { + int m_local = + get<0>(tTR_cChunk(i)); + int n_local = + get<1>(tTR_cChunk(i)); + sStore(m_local, n_local) = + static_cast( + tTR_rAcc(i)); + } + } + cutlass::arch::NamedBarrier::sync( + Traits::WarpGroupSize, + wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int out_m = + e * num_m_tiles + m_tile; + int n_chunk = + n_tile * + (TileN / EpiChunkN) + + wg * NumEpiRounds + + round; + CUTE_UNROLL + for (int atom = 0; + atom < + Traits::kAtomsPerCta; + ++atom) { + int m_atom = + (TileM / + Traits::AtomTileM) * + out_m + + Traits::kAtomsPerCta * + pair_rank + + atom; + auto sStoreAtom = + local_tile( + sStore, + make_tile( + Int< + Traits:: + AtomTileM>{}, + Int< + EpiChunkN>{}), + make_coord( + atom, 0)); + if (phase == 0) { + auto gdB = local_tile( + mdB, + make_tile( + Int< + Traits:: + AtomTileM>{}, + Int< + EpiChunkN>{}), + make_coord( + m_atom, + n_chunk)); + copy( + tma_reduce_db, + cta_tma_db.partition_S( + sStoreAtom), + cta_tma_db.partition_D( + gdB)); + } else { + auto gdC = local_tile( + mdC, + make_tile( + Int< + Traits:: + AtomTileM>{}, + Int< + EpiChunkN>{}), + make_coord( + m_atom, + n_chunk)); + copy( + tma_reduce_dc, + cta_tma_dc.partition_S( + sStoreAtom), + cta_tma_dc.partition_D( + gdC)); + } + } + cute::tma_store_arrive(); + } + store_in_flight = true; + } + + cutlass::arch::NamedBarrier::sync( + Traits::ConsumerThreads, 0); + // One elected release from each CTA satisfies the + // cluster-scoped accumulator pipeline's count of two. + if (tid_in_epi == 0) + acc_pipe.consumer_release( + acc_cons_state); + ++acc_cons_state; + } + } + } + } + + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, 3); +#else + __trap(); +#endif +} +}; // Mlp4ConsumerImpl<100> + +// ═══════════════════════════════════════════════════════════════════ +// Single-launch driver, chunk-fixed and architecture-dispatched +// ═══════════════════════════════════════════════════════════════════ +// +// Compute=90 uses the original per-CTA Hopper pipeline and cell mapping. +// Compute=100 uses a paired-CTA pipeline, pair-granular cell mapping, a +// shared Allocator2Sm allocation, and cluster synchronization. +template < + typename Traits, + int Compute = 90, + bool Expert3D = false, + typename TmaLoadX, + typename TmaLoadA, + typename TmaReduceAddDB, + typename TmaReduceAddDC> +__device__ __forceinline__ void mlp4_fwd( + Mlp4Smem& smem, + TmaLoadX const& tma_load_x, + TmaLoadA const& tma_load_dut, + TmaLoadA const& tma_load_dvt, + TmaReduceAddDB const& tma_reduce_db, + TmaReduceAddDC const& tma_reduce_dc, + const int* expert_k_starts, + const int* expert_k_ends, + int num_experts, + int intermediate_dim, + int hidden_dim, + int num_tokens, + int total_m_rows, + int num_m_tiles, + int num_n_tiles, + int outer_split, + int batch_kb_start = 0, + int batch_kb_end = -1, + int k_split = 1, + int ring_kb = 0) { + static_assert(Compute == 90 || Compute == 100, + "mlp4_fwd supports Compute=90 or Compute=100"); + using PipeState = typename Traits::PipelineState; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = warp_id == 0; + constexpr int kFirstConsumerWarp = Compute == 100 ? 3 : 4; + bool is_consumer = + warp_id >= kFirstConsumerWarp && warp_id <= 11; + + cute::prefetch_tma_descriptor( + tma_load_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_load_dut.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_load_dvt.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_reduce_db.get_tma_descriptor()); + cute::prefetch_tma_descriptor( + tma_reduce_dc.get_tma_descriptor()); + + auto pipe = [&]() { + if constexpr (Compute == 100) + return mlp4_make_pipe_umma_2sm( + smem.pipe_storage); + else + return mlp4_make_pipe(smem.pipe_storage); + }(); + using Pipeline = decltype(pipe); + + if constexpr (Compute == 100) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + cute::TMEM::Allocator2Sm tmem_allocator; + if (gridDim.x % Traits::ClusterM != 0) + __trap(); + cute::cluster_sync(); + constexpr int kTmemColumns = + Traits::AccStages * Traits::TileN; + if (warp_id == 4) { + tmem_allocator.allocate( + kTmemColumns, &smem.tmem_base); + __syncwarp(); + } + __syncthreads(); + cute::cluster_sync(); +#else + __trap(); +#endif + } else { + __syncthreads(); + } + + PipeState producer_state = is_producer + ? cutlass::make_producer_start_state() + : PipeState{}; + PipeState consumer_state; + + int cell_start = static_cast(blockIdx.x); + int cell_stride = static_cast(gridDim.x); + if constexpr (Compute == 100) { + cell_start /= Traits::ClusterM; + cell_stride /= Traits::ClusterM; + } + int kb_end = batch_kb_end >= 0 + ? batch_kb_end + : num_tokens / Traits::TileK; + + if (is_producer) { + mlp4_producer( + pipe, + producer_state, + smem, + tma_load_x, + tma_load_dut, + tma_load_dvt, + expert_k_starts, + expert_k_ends, + num_experts, + intermediate_dim, + hidden_dim, + num_tokens, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split, + ring_kb); + } else if (is_consumer) { + if constexpr (Compute == 100) { + mlp4_consumer( + pipe, + consumer_state, + smem, + tma_reduce_db, + tma_reduce_dc, + expert_k_starts, + expert_k_ends, + num_experts, + intermediate_dim, + hidden_dim, + total_m_rows, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split); + } else { +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) + mlp4_consumer( + pipe, + consumer_state, + smem, + tma_reduce_db, + tma_reduce_dc, + expert_k_starts, + expert_k_ends, + num_experts, + intermediate_dim, + hidden_dim, + total_m_rows, + num_m_tiles, + num_n_tiles, + outer_split, + cell_start, + cell_stride, + batch_kb_start, + kb_end, + k_split); +#else + __trap(); +#endif + } + } + + if constexpr (Compute == 100) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + cute::TMEM::Allocator2Sm tmem_allocator; + __syncthreads(); + cute::cluster_sync(); + constexpr int kTmemColumns = + Traits::AccStages * Traits::TileN; + if (warp_id == 4) { + tmem_allocator.release_allocation_lock(); + tmem_allocator.free( + smem.tmem_base, kTmemColumns); + } +#endif + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Host launcher (raw-pointer; defined in mlp4.cu) +// ═══════════════════════════════════════════════════════════════════ + +// dB = dU^T @ X, dC = dV^T @ X. Both outputs accumulated atomically +// via TMA_REDUCE_ADD. dB and dC MUST be zero-initialized by caller. +// +// dU : [num_tokens, intermediate_dim] +// dV : [num_tokens, intermediate_dim] +// X : [num_tokens, hidden_dim] +// dB, dC : [num_experts, intermediate_dim, hidden_dim] +// expert_for_k_block: [num_tokens / TileK] int32, TileK=64 +void mlp4_fwd_bf16_launch( + const cutlass::bfloat16_t* dU, + const cutlass::bfloat16_t* dV, + const cutlass::bfloat16_t* X, + cutlass::bfloat16_t* dB, + cutlass::bfloat16_t* dC, + const int* expert_for_k_block, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int device, + cudaStream_t stream); + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp5.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp5.cuh index b4f2b919..cb2708e1 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp5.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp5.cuh @@ -1,599 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Phase 5: dX = dU @ B + dV @ C (input gradient, backward) -// NESTED-LOOP scheduler — the default. -// ═══════════════════════════════════════════════════════════════════ -// -// Cooperative 2-WG consumer; split axis selected by TileM: -// TileM=128 → Layout<_2,_1,_1>, split along M (each WG owns 64×TileN) -// TileM=64 → Layout<_1,_2,_1>, split along N (each WG owns 64×TileN/2) -// ONE accumulator per WG covering (AtomTileM=64, AtomN). Each WG writes -// its strip independently — no cross-WG NamedBarrier sync in the epilogue. -// -// Fused mainloop runs a single continuous (2 × num_k_tiles) k-loop: -// Phase 1 (k=0..num_k_tiles-1): acc += dU @ B -// Phase 2 (k=num_k_tiles..2K-1): acc += dV @ C -// -// Scheduler: -// - 2D grid (num_sms / NSplit, NSplit) — host-side launcher MUST -// match (see mlp5.cu). -// - nested `for (m = blockIdx.x; ...) for (n = blockIdx.y; ...)` -// - amortizes expert lookup + dU/dV tile-coord setup once per m -// - cross-CTA dU/dV multicast in L2 between (m, 0) and (m, 1) in the -// same m-column -// -// Alternative flat work_idx scheduler is preserved in `mlp5_flat.cuh` -// (which includes this header for the shared building blocks). -// -// CUTLASS-canonical patterns: -// - K_PIPE_MMAS = 1 -// - __launch_bounds__(NumThreads, 1) — applied by the caller kernel -// -// CTA layout (384 threads = 12 warps): -// WG0 warp 0 : Producer (TMA loads for dU/dV/B/C) -// WG0 warps 1-2 : Idle / reserved for fused-kernel communication -// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 -// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 -// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 -// -// Per `expert_for_m_block`: one int per TileM=128 m-block. -// -// Smem at default config (TileM=128, TileN=256, TileK=64, Stages=4, -// EpiChunkN=64): -// Z (TileM × TileK × S × 2) = 64 KiB -// W (TileN × TileK × S × 2) = 128 KiB -// store_buf (2 WGs × AtomTileM × EpiChunkN × 2) = 16 KiB -// pipe + barriers : ~1 KiB -// ───────────────────────────────────────────── -// Total : ~209 KiB (≤ 228 KiB cap) -// ═══════════════════════════════════════════════════════════════════ - -#include "models.cuh" -#include "tmem_load_op.cuh" - -#include -#include -#include -#include -#include - -// SM100 (Blackwell / UMMA) support for the Compute=100 consumer -// specialization (mlp5_fused.cuh). These headers are header-only and safe to -// include under an sm_90a build — the tcgen05 PTX they emit is guarded by -// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on -// __CUDA_ARCH__ >= 1000) ever instantiates it. -#include -#include - -#include -#include -#include // PipelineTmaUmmaAsync, PipelineUmmaAsync - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// Traits -// ═══════════════════════════════════════════════════════════════════ - -template < - typename Element_, - int TileM_ = 128, - int TileN_ = 256, - int TileK_ = 64, - int Stages_ = 4, - int EpiChunkN_ = 64 -> -struct Mlp5Traits { - using Element = Element_; - using ElementAccum = float; - - static constexpr int TileM = TileM_; - static constexpr int TileN = TileN_; - static constexpr int TileK = TileK_; - static constexpr int Stages = Stages_; - static constexpr int EpiChunkN = EpiChunkN_; - - static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); - - // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate - // along M (each owns 64 rows × TileN). With TileM=64 the WGs cooperate - // along N (each owns 64 rows × TileN/2). AtomTileM is the per-WG M - // extent — always 64. Mirrors mlp2_fused.cuh's kMSplit/N-split pattern. - static constexpr bool kMSplit = (TileM == 128); - static constexpr int AtomTileM = 64; - static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); - - static_assert(TileN <= 256, "TileN must fit TMA box limit"); - static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); - static_assert(kMSplit || (TileN % 2 == 0), - "TileN must be even when TileM=64 (cooperative N-split halves)"); - static constexpr int NumEpiRounds = WgTileN / EpiChunkN; - - // Cooperative 2-WG layout: M-split for TileM=128, N-split for TileM=64. - // AtomN is the per-WG WGMMA atom N — combined across both WGs it must - // cover exactly TileN cols (= sW's N extent). - static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); - using GmmaAtom = typename GmmaSelectorKMN::Atom; - using WgLayout = cute::conditional_t< - kMSplit, - Layout>, - Layout>>; - using TiledMma = TiledMMA, WgLayout>; - - using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom; - using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom; - - // One-stage layouts (used for TMA descriptors). - using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{}, - Shape, Int>{})); - using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtomW{}, - Shape, Int>{}, Step<_2, _1>{})); - - // Multi-stage smem layouts. - using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{}, - Shape, Int, Int>{})); - using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{}, - Shape, Int, Int>{}, Step<_2, _1, _3>{})); - - // Per-WG store buffer: (AtomTileM, EpiChunkN). One slot per WG. - using SmemLayoutStore = Layout< - Shape, Int>, - Stride, _1>>; - // Alias for callers that build TMA descriptors using the single-slot - // 2D layout (e.g., moe_bwd's mlp_bwd.cu/moe_bwd.cu). - using SmemLayoutStore_1 = SmemLayoutStore; - - using MainloopPipeline = cutlass::PipelineTmaAsync; - using PipelineState = cutlass::PipelineState; - - // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop - // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the - // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the single - // TMEM accumulator from the MMA warp to the epilogue warps. See - // Mlp5FusedConsumerImpl<100> in mlp5_fused.cuh. The Blackwell launcher must - // build the mainloop pipe with num_consumers=1 (one umma_arrive per stage). - using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< - Stages, cute::Shape, - cute::Shape>; - // Double-buffered TMEM accumulators: the single accumulator per stage still - // carries the cross-phase sum dU·B + dV·C across the whole 2·num_k_tiles - // mainloop, while the alternate stage lets the next n-tile's MMA overlap - // the previous n-tile's epilogue. - static constexpr int AccStages = 2; - static_assert(AccStages * TileN <= 512, - "SM100 MLP5 accumulator stages must fit in 512 TMEM columns"); - using AccumulatorPipeline = cutlass::PipelineUmmaAsync; - - static constexpr int TmaTransBytesZ = - static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); - static constexpr int TmaTransBytesW = - static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); - static constexpr int TmaTransBytes = TmaTransBytesZ + TmaTransBytesW; - - static constexpr int WarpSize = 32; - static constexpr int WarpGroupSize = 128; - static constexpr int NumConsumers = 2; - static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; - static constexpr int NumThreads = 384; -}; - -template -using Mlp5MainloopPipelineFor = cute::conditional_t< - Compute == 100, - typename Traits::MainloopPipelineUmma, - typename Traits::MainloopPipeline>; - -// ═══════════════════════════════════════════════════════════════════ -// Shared memory layout -// ═══════════════════════════════════════════════════════════════════ - -template -struct Mlp5Smem { - using Element = typename Traits::Element; - - static constexpr int smem_Z_size = cosize_v; - static constexpr int smem_W_size = cosize_v; - static constexpr int smem_store_size = cosize_v; - - alignas(128) Element smem_Z[smem_Z_size]; - alignas(128) Element smem_W[smem_W_size]; - // Per-WG store buffer: WG0 uses store_buf[0..S-1], WG1 uses [S..2S-1]. - alignas(128) Element store_buf[2 * smem_store_size]; - - typename Traits::MainloopPipeline::SharedStorage pipe_storage; - - // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM - // base address, plus the accumulator pipeline's barrier storage (UMMA→ - // epilogue handoff). The Hopper (Compute=90) path never touches these. - alignas(16) uint32_t tmem_base; - alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; -}; - -// ═══════════════════════════════════════════════════════════════════ -// Pipeline construction -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ auto mlp5_make_pipe( - typename Traits::MainloopPipeline::SharedStorage& storage) { - - using Pipeline = typename Traits::MainloopPipeline; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 4 && warp_id <= 11); - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - // Both consumer WGs (cooperative) consume each pipe slot. - pp.num_consumers = Traits::ConsumerThreads; - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// SM100 (Compute=100) mainloop pipe. Mirrors mlp5_make_pipe but builds the -// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues -// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer -// is released by a single umma_arrive per stage (the MMA warp), not per -// consumer thread. Params/PipelineState/SharedStorage alias the Hopper -// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer -// (mlp5_fused_producer) and the Compute=100 consumer drive it unchanged. -template -__device__ __forceinline__ auto mlp5_make_pipe_umma( - typename Traits::MainloopPipelineUmma::SharedStorage& storage) { - - using Pipeline = typename Traits::MainloopPipelineUmma; - using Category = typename Pipeline::ThreadCategory; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer - - typename Pipeline::Params pp; - pp.transaction_bytes = Traits::TmaTransBytes; - pp.num_producers = 1; - pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer - if (is_producer) { - pp.role = Category::Producer; - pp.is_leader = (threadIdx.x == 0); - } else if (is_consumer) { - pp.role = Category::Consumer; - } else { - pp.role = Category::NonParticipant; - } - return Pipeline(storage, pp, Shape<_1, _1, _1>{}, - cute::true_type{}, cute::true_type{}); -} - -// ── Producer (nested) ─────────────────────────────────────────────── - -template -__device__ __forceinline__ void mlp5_producer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp5Smem& smem, - TmaLoadZ const& tma_load_du, - TmaLoadZ const& tma_load_dv, - TmaLoadW const& tma_load_b, - TmaLoadW const& tma_load_c, - const int* expert_for_m_block, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int total_k_cols, - int num_m_tiles, - int num_n_tiles, - int num_k_tiles) { - - auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); - - auto mDU = tma_load_du.get_tma_tensor(make_shape(num_tokens, intermediate_dim)); - auto mDV = tma_load_dv.get_tma_tensor(make_shape(num_tokens, intermediate_dim)); - auto mB = tma_load_b.get_tma_tensor(make_shape(hidden_dim, total_k_cols)); - auto mC = tma_load_c.get_tma_tensor(make_shape(hidden_dim, total_k_cols)); - - auto cta_tma_du = tma_load_du.get_slice(Int<0>{}); - auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); - - auto tZsZ = cta_tma_du.partition_D(sZ); - auto tWsW = cta_tma_b.partition_D(sW); - - auto gB_all = local_tile(mB, - make_tile(Int{}, Int{}), - make_coord(_, _)); - auto gC_all = local_tile(mC, - make_tile(Int{}, Int{}), - make_coord(_, _)); - auto tBgB_all = cta_tma_b.partition_S(gB_all); - auto tCgC_all = cta_tma_b.partition_S(gC_all); - - bool is_leader = (threadIdx.x == 0); - - int n_start = (int)blockIdx.y; - int n_stride = (int)gridDim.y; - - for (int m = (int)blockIdx.x; m < num_m_tiles; m += (int)gridDim.x) { - int expert = expert_for_m_block[m]; - int expert_k_offset = expert * num_k_tiles; - - auto gDU = local_tile(mDU, - make_tile(Int{}, Int{}), - make_coord(m, _)); - auto tZgDU = cta_tma_du.partition_S(gDU); - auto gDV = local_tile(mDV, - make_tile(Int{}, Int{}), - make_coord(m, _)); - auto tZgDV = cta_tma_du.partition_S(gDV); - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - // Phase 1: dU @ B - for (int k = 0; k < num_k_tiles; ++k) { - pipe.producer_acquire(state); - if (is_leader) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_du.with(*bar, 0), - tZgDU(_, _, _, k), tZsZ(_, _, _, state.index())); - copy(tma_load_b.with(*bar, 0), - tBgB_all(_, _, _, n, expert_k_offset + k), - tWsW(_, _, _, state.index())); - } - ++state; - } - // Phase 2: dV @ C - for (int k = 0; k < num_k_tiles; ++k) { - pipe.producer_acquire(state); - if (is_leader) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_dv.with(*bar, 0), - tZgDV(_, _, _, k), tZsZ(_, _, _, state.index())); - copy(tma_load_c.with(*bar, 0), - tCgC_all(_, _, _, n, expert_k_offset + k), - tWsW(_, _, _, state.index())); - } - ++state; - } - } - } -} - -// ── Consumer (nested) ─────────────────────────────────────────────── - -template -__device__ __forceinline__ void mlp5_consumer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp5Smem& smem, - TmaStoreDX const& tma_store_dx, - int hidden_dim, - int num_m_tiles, - int num_n_tiles, - int num_k_tiles) { - - (void)Compute; - using Element = typename Traits::Element; - typename Traits::TiledMma tiled_mma; - int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; - int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; - auto thr_mma = tiled_mma.get_slice(tid_in_mma); - - auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); - - auto tCsZ = thr_mma.partition_A(sZ); - auto tCsW = thr_mma.partition_B(sW); - - auto acc = partition_fragment_C(tiled_mma, - Shape, Int>{}); - - auto cC = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCcC = thr_mma.partition_C(cC); - - const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; - const int my_barrier_id = 1 + my_wg; - const bool is_my_wg_leader = (tid_in_wg == 0); - - constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; - Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStore{}); - - auto mDX = tma_store_dx.get_tma_tensor( - make_shape(num_m_tiles * Traits::TileM, hidden_dim)); - auto cta_tma_dx = tma_store_dx.get_slice(Int<0>{}); - - int total_k = 2 * num_k_tiles; - constexpr int K_PIPE_MMAS = 1; - bool store_in_flight = false; - - int n_start = (int)blockIdx.y; - int n_stride = (int)gridDim.y; - - for (int m = (int)blockIdx.x; m < num_m_tiles; m += (int)gridDim.x) { - for (int n = n_start; n < num_n_tiles; n += n_stride) { - - clear(acc); - auto state_release = state; - int prologue_count = (total_k < K_PIPE_MMAS) ? total_k : K_PIPE_MMAS; - - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - ++state; - } - for (int k = prologue_count; k < total_k; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - - warpgroup_wait(); - warpgroup_fence_operand(acc); - pipe.consumer_release(state_release); - ++state; - ++state_release; - } - warpgroup_wait<0>(); - warpgroup_fence_operand(acc); - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_release(state_release); - ++state_release; - } - - // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. - // TileM=64 (N-split): WG_w owns full M (=64), cols [w*WgTileN, ...). - CUTE_UNROLL - for (int r = 0; r < Traits::NumEpiRounds; ++r) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - int chunk_start = r * Traits::EpiChunkN; - CUTE_UNROLL - for (int i = 0; i < size(acc); ++i) { - auto coord = tCcC(i); - int m_loc = get<0>(coord); - int n_loc = get<1>(coord); - int m_local, n_local; - if constexpr (Traits::kMSplit) { - m_local = m_loc - my_wg * Traits::AtomTileM; - n_local = n_loc; - } else { - m_local = m_loc; - n_local = n_loc - my_wg * Traits::WgTileN; - } - if (n_local >= chunk_start && - n_local < chunk_start + Traits::EpiChunkN) { - int chunk_n = n_local - chunk_start; - sStore(m_local, chunk_n) = static_cast(acc(i)); - } - } - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - if (is_my_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx, n_tile_idx; - if constexpr (Traits::kMSplit) { - // TileM=128: 2 row tiles per logical M tile. - m_tile_idx = 2 * m + my_wg; - n_tile_idx = n * Traits::NumEpiRounds + r; - } else { - // TileM=64: 1 row tile per logical M tile; WG_w handles - // its N-half. - m_tile_idx = m; - n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) - + my_wg * Traits::NumEpiRounds + r; - } - auto gDX = local_tile(mDX, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_dx, cta_tma_dx.partition_S(sStore), - cta_tma_dx.partition_D(gDX)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - } - if (store_in_flight) - cute::tma_store_wait<0>(); -} - -// ── Driver ────────────────────────────────────────────────────────── - -template -__device__ __forceinline__ void mlp5_fwd( - Mlp5Smem& smem, - TmaLoadZ const& tma_load_du, - TmaLoadZ const& tma_load_dv, - TmaLoadW const& tma_load_b, - TmaLoadW const& tma_load_c, - TmaStoreDX const& tma_store_dx, - const int* expert_for_m_block, - int num_tokens, - int intermediate_dim, - int hidden_dim, - int total_k_cols, - int num_m_tiles, - int num_n_tiles, - int num_k_tiles) { - - using Pipeline = typename Traits::MainloopPipeline; - using PipeState = typename Traits::PipelineState; - - int warp_id = threadIdx.x / Traits::WarpSize; - bool is_producer = (warp_id == 0); - bool is_consumer = (warp_id >= 4 && warp_id <= 11); - - auto pipe = mlp5_make_pipe(smem.pipe_storage); - - __syncthreads(); - - PipeState state = is_producer ? - cutlass::make_producer_start_state() : PipeState{}; - PipeState cons_state; - - if (is_producer) { - mlp5_producer(pipe, state, smem, - tma_load_du, tma_load_dv, tma_load_b, tma_load_c, - expert_for_m_block, - num_tokens, hidden_dim, intermediate_dim, total_k_cols, - num_m_tiles, num_n_tiles, num_k_tiles); - } - if (is_consumer) { - mlp5_consumer(pipe, cons_state, smem, - tma_store_dx, hidden_dim, - num_m_tiles, num_n_tiles, num_k_tiles); - } - - __syncthreads(); -} - -// ═══════════════════════════════════════════════════════════════════ -// Host launcher (raw-pointer; defined in mlp5.cu) -// ═══════════════════════════════════════════════════════════════════ -// -// dX = dU @ B + dV @ C. Output dX is fully written; no zero-init needed. -// -// dU : [num_tokens, intermediate_dim] -// dV : [num_tokens, intermediate_dim] -// B, C : [num_experts, intermediate_dim, hidden_dim] -// dX : [num_tokens, hidden_dim] (output) -// expert_for_m_block: [num_tokens / TileM] int32 (TileM is Traits::TileM) -void mlp5_fwd_bf16_launch( - const cutlass::bfloat16_t* dU, - const cutlass::bfloat16_t* dV, - const cutlass::bfloat16_t* B, - const cutlass::bfloat16_t* C, - cutlass::bfloat16_t* dX, - const int* expert_for_m_block, - int num_tokens, - int intermediate_dim, - int hidden_dim, - int num_experts, - int device, - cudaStream_t stream); - -} // namespace liger +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif + +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp5_sm90.cuh" +#else +#include "mlp5_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp5_fused.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp5_fused.cuh index cf7c2588..4390d07d 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp5_fused.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp5_fused.cuh @@ -1,629 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Phase 5 (single-tile API for fused use) -// ═══════════════════════════════════════════════════════════════════ -// -// Same kernel body as mlp5.cuh's nested-loop default, exposed as a -// single-(m)-tile producer + consumer so a parent kernel can drive the -// outer m loop itself (e.g. mlp_bwd.cuh composing Phase 1d). -// -// Mirrors how mlp1_fused.cuh exposes single-tile entry points: callers -// own the Pipeline construction (use -// mlp5_make_pipe from mlp5.cuh) and the Mlp5Smem allocation; this -// header just provides the producer/consumer body parameterized for -// ONE m-tile. -// -// Params unique to fused use: -// - m : the m-tile this invocation processes -// - expert_k_offset : expert · num_k_tiles (caller resolves the -// expert lookup once per m) -// - split_idx : the column index within an N-split (caller -// picks; analogous to blockIdx.y in the -// standalone driver) -// - num_splits : total columns (analogous to gridDim.y) -// -// Output store goes to a global dX tensor view of shape -// (num_m_tiles · TileM, hidden_dim) (same as the standalone consumer). -// ═══════════════════════════════════════════════════════════════════ - -#include "mlp5.cuh" // Mlp5Traits, Mlp5Smem, mlp5_make_pipe - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// Producer (single m-tile) — fills 2·num_k_tiles slots per (m, n) -// across n ∈ {split_idx, split_idx+num_splits, ...}. -// Phase 1 (k=0..num_k_tiles-1): Z=dU, W=B -// Phase 2 (k=num_k_tiles..2K-1): Z=dV, W=C -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void mlp5_fused_producer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp5Smem& smem, - TmaLoadZ const& tma_load_du, - TmaLoadZ const& tma_load_dv, - TmaLoadW const& tma_load_b, - TmaLoadW const& tma_load_c, - int m, - int expert_or_k_offset, - int num_tokens, - int hidden_dim, - int intermediate_dim, - int num_experts, - int total_k_cols, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { - - auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); - - // int64_t cast for production shapes where T · I or E · I · H can - // exceed 2^31 elements (same rationale as in mlp1_fused_act.cuh and - // the int64 commit on the fwd path). - auto mDU = tma_load_du.get_tma_tensor(make_shape( - static_cast(num_tokens), - static_cast(intermediate_dim))); - auto mDV = tma_load_dv.get_tma_tensor(make_shape( - static_cast(num_tokens), - static_cast(intermediate_dim))); - auto mB = [&]() { - if constexpr (Expert3D) { - return tma_load_b.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(intermediate_dim), - static_cast(num_experts))); - } else { - return tma_load_b.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(total_k_cols))); - } - }(); - auto mC = [&]() { - if constexpr (Expert3D) { - return tma_load_c.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(intermediate_dim), - static_cast(num_experts))); - } else { - return tma_load_c.get_tma_tensor(make_shape( - static_cast(hidden_dim), - static_cast(total_k_cols))); - } - }(); - - auto cta_tma_du = tma_load_du.get_slice(Int<0>{}); - auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); - - auto tZsZ = cta_tma_du.partition_D(sZ); - auto tWsW = cta_tma_b.partition_D(sW); - - auto gB_all = [&]() { - if constexpr (Expert3D) { - return local_tile(mB, - make_tile(Int{}, Int{}), - make_coord(_, _, expert_or_k_offset)); - } else { - return local_tile(mB, - make_tile(Int{}, Int{}), - make_coord(_, _)); - } - }(); - auto gC_all = [&]() { - if constexpr (Expert3D) { - return local_tile(mC, - make_tile(Int{}, Int{}), - make_coord(_, _, expert_or_k_offset)); - } else { - return local_tile(mC, - make_tile(Int{}, Int{}), - make_coord(_, _)); - } - }(); - auto tBgB_all = cta_tma_b.partition_S(gB_all); - auto tCgC_all = cta_tma_b.partition_S(gC_all); - - auto gDU = local_tile(mDU, - make_tile(Int{}, Int{}), - make_coord(m, _)); - auto tZgDU = cta_tma_du.partition_S(gDU); - auto gDV = local_tile(mDV, - make_tile(Int{}, Int{}), - make_coord(m, _)); - auto tZgDV = cta_tma_du.partition_S(gDV); - - const bool is_leader = (threadIdx.x == 0); - - for (int n = split_idx; n < num_n_tiles; n += num_splits) { - // Phase 1: dU @ B - for (int k = 0; k < num_k_tiles; ++k) { - pipe.producer_acquire(state); - if (is_leader) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_du.with(*bar, 0), - tZgDU(_, _, _, k), tZsZ(_, _, _, state.index())); - if constexpr (Expert3D) { - copy(tma_load_b.with(*bar, 0), - tBgB_all(_, _, _, n, k), tWsW(_, _, _, state.index())); - } else { - copy(tma_load_b.with(*bar, 0), - tBgB_all(_, _, _, n, expert_or_k_offset + k), - tWsW(_, _, _, state.index())); - } - } - ++state; - } - // Phase 2: dV @ C - for (int k = 0; k < num_k_tiles; ++k) { - pipe.producer_acquire(state); - if (is_leader) { - auto* bar = pipe.producer_get_barrier(state); - copy(tma_load_dv.with(*bar, 0), - tZgDV(_, _, _, k), tZsZ(_, _, _, state.index())); - if constexpr (Expert3D) { - copy(tma_load_c.with(*bar, 0), - tCgC_all(_, _, _, n, k), tWsW(_, _, _, state.index())); - } else { - copy(tma_load_c.with(*bar, 0), - tCgC_all(_, _, _, n, expert_or_k_offset + k), - tWsW(_, _, _, state.index())); - } - } - ++state; - } - } -} - -// ═══════════════════════════════════════════════════════════════════ -// Consumer (single m-tile) — architecture-specialized on `int Compute` -// ═══════════════════════════════════════════════════════════════════ -// -// Primary template is undefined; one full specialization per supported -// compute capability provides a `run(...)` member (a function-template on -// Traits/Pipeline/TmaStoreDX, since function templates can't be partially -// specialized). The free function `mlp5_fused_consumer` below forwards to -// the right specialization — call sites stay in the existing style. -// -// Compute=90 → Hopper / WGMMA (cooperative 2-WG, single register acc, -// cross-phase sum accumulated across 2·num_k_tiles MMAs) -// Compute=100 → Blackwell / UMMA (single-warp issue, ONE TMEM accumulator, -// cross-phase accumulate via the per-instruction ScaleOut bit) - -template -struct Mlp5FusedConsumerImpl; - -// ─────────────────────────────────────────────────────────────────── -// Compute=90 — Hopper. The two WGs split (TileM,TileN); the single register -// accumulator carries dU·B (phase 1, k=0..K-1) + dV·C (phase 2, k=K..2K-1) -// via cute::gemm's running accumulation over the whole 2·num_k_tiles k-loop. -// (Verbatim from the original mlp5_fused_consumer.) -// ─────────────────────────────────────────────────────────────────── - -template <> -struct Mlp5FusedConsumerImpl<90> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp5Smem& smem, - TmaStoreDX const& tma_store_dx, - int m, - int hidden_dim, - int num_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { - - using Element = typename Traits::Element; - typename Traits::TiledMma tiled_mma; - int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; - int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; - auto thr_mma = tiled_mma.get_slice(tid_in_mma); - - auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); - - auto tCsZ = thr_mma.partition_A(sZ); - auto tCsW = thr_mma.partition_B(sW); - - auto acc = partition_fragment_C(tiled_mma, - Shape, Int>{}); - - auto cC = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCcC = thr_mma.partition_C(cC); - - const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; - const int my_barrier_id = 1 + my_wg; - const bool is_my_wg_leader = (tid_in_wg == 0); - - constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; - Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStore{}); - - auto mDX = tma_store_dx.get_tma_tensor(make_shape( - static_cast(num_m_tiles) * Traits::TileM, - static_cast(hidden_dim))); - auto cta_tma_dx = tma_store_dx.get_slice(Int<0>{}); - - int total_k = 2 * num_k_tiles; - constexpr int K_PIPE_MMAS = 1; - bool store_in_flight = false; - - for (int n = split_idx; n < num_n_tiles; n += num_splits) { - - clear(acc); - auto state_release = state; - int prologue_count = (total_k < K_PIPE_MMAS) ? total_k : K_PIPE_MMAS; - - // Prologue - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - ++state; - } - // Steady state - for (int k = prologue_count; k < total_k; ++k) { - pipe.consumer_wait(state); - warpgroup_fence_operand(acc); - warpgroup_arrive(); - gemm(tiled_mma, tCsZ(_, _, _, state.index()), - tCsW(_, _, _, state.index()), acc); - warpgroup_commit_batch(); - - warpgroup_wait(); - warpgroup_fence_operand(acc); - pipe.consumer_release(state_release); - ++state; - ++state_release; - } - // Drain - warpgroup_wait<0>(); - warpgroup_fence_operand(acc); - for (int k = 0; k < prologue_count; ++k) { - pipe.consumer_release(state_release); - ++state_release; - } - - // ── Epilogue ────────────────────────────────────── - CUTE_UNROLL - for (int r = 0; r < Traits::NumEpiRounds; ++r) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - int chunk_start = r * Traits::EpiChunkN; - CUTE_UNROLL - for (int i = 0; i < size(acc); ++i) { - auto coord = tCcC(i); - int m_loc = get<0>(coord); - int n_loc = get<1>(coord); - int m_local, n_local; - if constexpr (Traits::kMSplit) { - m_local = m_loc - my_wg * Traits::AtomTileM; - n_local = n_loc; - } else { - m_local = m_loc; - n_local = n_loc - my_wg * Traits::WgTileN; - } - if (n_local >= chunk_start && - n_local < chunk_start + Traits::EpiChunkN) { - int chunk_n = n_local - chunk_start; - sStore(m_local, chunk_n) = static_cast(acc(i)); - } - } - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); - - if (is_my_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx, n_tile_idx; - if constexpr (Traits::kMSplit) { - m_tile_idx = 2 * m + my_wg; - n_tile_idx = n * Traits::NumEpiRounds + r; - } else { - m_tile_idx = m; - n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) - + my_wg * Traits::NumEpiRounds + r; - } - auto gDX = local_tile(mDX, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_dx, cta_tma_dx.partition_S(sStore), - cta_tma_dx.partition_D(gDX)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - if (store_in_flight) - cute::tma_store_wait<0>(); -} -}; // Mlp5FusedConsumerImpl<90> - -// ─────────────────────────────────────────────────────────────────── -// Compute=100 — Blackwell / UMMA. -// -// One 1SM UMMA atom with M=TileM covers the whole tile; the single -// accumulator dX = dU·B + dV·C lives in TMEM. Operand A (dU/dV) is K-major; -// operand B (B/C) is MN-major — reusing mlp2_t's column-major weight view -// (SmemLayoutAtomW = Layout_MN_SW128_Atom, already in Mlp5Traits) with the -// UMMA atom's b_major = UMMA::Major::MN (the descriptor builder reads the -// MN-major stride from the smem layout, so no transpose is needed). -// -// Cross-phase accumulate (the crux): the producer fills 2·num_k_tiles pipe -// stages — B in phase 1 (k -struct Mlp5FusedConsumerImpl<100> { -template -static __device__ __forceinline__ void run( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp5Smem& smem, - TmaStoreDX const& tma_store_dx, - int m, - int hidden_dim, - int num_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { -#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - using Element = typename Traits::Element; - constexpr int TileM = Traits::TileM; - constexpr int TileN = Traits::TileN; - constexpr int EpiChunkN = Traits::EpiChunkN; - constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile - static_assert(TileN % 2 == 0, - "Blackwell consumer splits TileN across the two consumer warpgroups"); - constexpr int WgN = TileN / 2; // per-warpgroup N width - static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); - constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per warpgroup - constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) - - // ── Thread identity ───────────────────────────────────── - // Warp 3 is the dedicated, epilogue-free UMMA producer. Warps 4..11 are - // the two aligned epilogue WGs. - const int warp_id = threadIdx.x / Traits::WarpSize; - const bool is_mma_warp = (warp_id == 3); - const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); - const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 - const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; - const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; - const bool is_wg_leader = is_epilogue && (tid_wg == 0); - const int wg_barrier_id = 1 + wg; // 1 or 2 - constexpr int kEpilogueThreads = Traits::ConsumerThreads; - constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; - static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); - static_assert(kEpilogueThreads == 8 * Traits::WarpSize); - static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); - - // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, SS. - // Operand A = dU/dV → K-major; operand B = B/C → MN-major (the - // column-major weight view, matching Mlp5Traits::SmemLayoutW). ── - auto tiled_mma = make_tiled_mma( - SM100_MMA_F16BF16_SS{}); - auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 - - auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); - auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); - - // ── One TMEM accumulator: dX = dU·B + dV·C ── - auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); - auto tCgC = cta_mma.partition_C(cAccFull); - auto tCtAcc = cta_mma.make_fragment_C(tCgC); - - // TMEM is allocated by the outer fused/standalone launcher once per CTA. - - // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers - // (warps 4..11, both WGs). AccStages stages let MMA for the next n-tile - // overlap the previous n-tile's epilogue. - using AccPipe = typename Traits::AccumulatorPipeline; - typename AccPipe::Params acc_params; - acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer - : AccPipe::ThreadCategory::Consumer; - acc_params.producer_arv_count = 1; // one umma_arrive per commit - acc_params.consumer_arv_count = 1; // one elected releaser - acc_params.initializing_warp = 4; // warp 4 inits the acc barriers - AccPipe acc_pipe(smem.acc_pipe, acc_params, - cute::Shape{}); - auto acc_prod_state = cutlass::make_producer_start_state(); - typename AccPipe::PipelineState acc_cons_state; - - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); - const uint32_t tmem_base = smem.tmem_base; - tCtAcc.data() = tmem_base; - - // ── Per-WG store slot (AtomTileM × EpiChunkN) — same smem as Hopper ── - constexpr int store_slot_elems = AtomTileM * EpiChunkN; - Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; - auto sStore = make_tensor(make_smem_ptr(my_store_ptr), - typename Traits::SmemLayoutStore{}); - - // dX view: [num_m_tiles·TileM, hidden_dim] row-major. int64_t cast keeps - // CUTE's layout math in 64-bit for large output buffers. - auto mDX = tma_store_dx.get_tma_tensor(make_shape( - static_cast(num_m_tiles) * TileM, - static_cast(hidden_dim))); - auto cta_tma_dx = tma_store_dx.get_slice(Int<0>{}); - - // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local tid). - // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) modes, - // which is what make_tmem_copy's cotiled builder requires. The register - // fragment is sized from partition_D (the DEST coords), NOT partition_S, - // so it excludes the warp-collective datapath-lane dim. ── - auto epi_tile = make_tile(Int{}, Int{}); - auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) - auto tAcc_epi = flat_divide(acc_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) - auto t2r = make_tmem_copy(TmemLoadOp{}, tAcc_epi(_, _, _0{}, _0{})); - auto thr_t2r = t2r.get_slice(tid_wg); - auto tTR_tAcc = thr_t2r.partition_S(tAcc_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) - auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); - auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) - auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); // f32 regs - - int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; - bool store_in_flight = false; - - const int total_k = 2 * num_k_tiles; // phase 1 (dU·B) + phase 2 (dV·C) - - for (int n = n_start; n < num_n_tiles; n += n_stride) { - - // ── Mainloop (warp 3 only): one continuous k-loop over ALL - // 2·num_k_tiles stages into the single TMEM accumulator. The - // accumulate bit is Zero on the very first MMA (clears the acc) and - // One thereafter — including across the phase-1→phase-2 boundary, so - // dX = dU·B + dV·C accumulates without dropping the dU·B term. ── - if (is_mma_warp) { - acc_pipe.producer_acquire(acc_prod_state); // TMEM acc free (prev epilogue done) - int acc_stage = acc_prod_state.index(); - tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); - for (int k = 0; k < total_k; ++k) { - pipe.consumer_wait(state); - auto tCsZ = cta_mma.partition_A(sZ(_, _, state.index())); - auto tCsW = cta_mma.partition_B(sW(_, _, state.index())); - auto tCrZ = cta_mma.make_fragment_A(tCsZ); - auto tCrW = cta_mma.make_fragment_B(tCsW); - CUTE_UNROLL - for (int kb = 0; kb < size<2>(tCrZ); ++kb) { - // First MMA of the whole 2K loop clears TMEM; every other - // MMA (incl. the phase boundary) accumulates. Do NOT reset. - tiled_mma.accumulate_ = (k == 0 && kb == 0) - ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; - gemm(tiled_mma, tCrZ(_, _, kb), tCrW(_, _, kb), tCtAcc); - } - pipe.consumer_release(state); // UMMA-gated smem-buffer release - ++state; - } - acc_pipe.producer_commit(acc_prod_state); // signal acc ready (umma_arrive) - ++acc_prod_state; - } - - // ── Epilogue: wait for the accumulator, then this WG processes its - // n-chunks [wg·NChunksHalf, +NChunksHalf). ── - if (is_epilogue) { - acc_pipe.consumer_wait(acc_cons_state); - int acc_stage = acc_cons_state.index(); - tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); - auto acc_mn_stage = tCtAcc(make_coord(_, _), _0{}, _0{}); - auto tAcc_epi_stage = flat_divide(acc_mn_stage, epi_tile); - auto tTR_tAcc_stage = thr_t2r.partition_S(tAcc_epi_stage); - - CUTE_UNROLL - for (int r = 0; r < NChunksHalf; ++r) { - int chunk = wg * NChunksHalf + r; // absolute n-chunk index - - // TMEM → registers (this chunk, full TileM rows). - copy(t2r, tTR_tAcc_stage(_, _, _, _0{}, chunk), tTR_rAcc); - cutlass::arch::fence_view_async_tmem_load(); - - // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). - CUTE_UNROLL - for (int ms = 0; ms < MSub; ++ms) { - if (store_in_flight) - cute::tma_store_wait<0>(); - - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - CUTE_UNROLL - for (int i = 0; i < size(tTR_rAcc); ++i) { - int m_row = get<0>(tTR_cChunk(i)); // 0..TileM - int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN - if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) - sStore(m_row - ms * AtomTileM, n_col) = - static_cast(tTR_rAcc(i)); - } - cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); - - if (is_wg_leader) { - cute::tma_store_fence(); - int m_tile_idx = MSub * m + ms; - int n_tile_idx = n * (TileN / EpiChunkN) + chunk; - auto gDX = local_tile(mDX, - make_tile(Int{}, Int{}), - make_coord(m_tile_idx, n_tile_idx)); - copy(tma_store_dx, cta_tma_dx.partition_S(sStore), - cta_tma_dx.partition_D(gDX)); - cute::tma_store_arrive(); - } - store_in_flight = true; - } - } - - // All epilogue TMEM reads are done; release the accumulator so the next - // n-tile's MMA may reuse it (one elected consumer thread arrives). - cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); - if (tid_in_epi == 0) - acc_pipe.consumer_release(acc_cons_state); - ++acc_cons_state; - } - } - if (is_epilogue && store_in_flight) - cute::tma_store_wait<0>(); - - // The TMEM allocation is owned by the launcher (freed once per CTA after the - // m-loop). Do NOT relinquish/free here — this consumer runs once per m-tile - // and freeing per-tile then re-allocating next tile trips the tcgen05 - // "phase invalid during alloc" guardrail. Keep this final 288-thread - // MMA+epilogue barrier so the next m-tile's accumulator-pipeline init is - // isolated from this tile's in-flight handshake. - cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); -#else - // The Compute=100 specialization is never dispatched on non-SM100 targets - // (host picks Compute=90 there). Trap if it is ever reached. - __trap(); +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 #endif -} -}; // Mlp5FusedConsumerImpl<100> -// ─────────────────────────────────────────────────────────────────── -// Forwarder — keeps the existing call style. `Compute` defaults to 90 so -// current call sites (mlp5_fused_consumer(...)) are unchanged; pass -// mlp5_fused_consumer(...) to select the Blackwell path. -// ─────────────────────────────────────────────────────────────────── - -template -__device__ __forceinline__ void mlp5_fused_consumer( - Pipeline& pipe, - typename Traits::PipelineState& state, - Mlp5Smem& smem, - TmaStoreDX const& tma_store_dx, - int m, - int hidden_dim, - int num_m_tiles, - int num_n_tiles, - int num_k_tiles, - int split_idx, - int num_splits) { - Mlp5FusedConsumerImpl::template run( - pipe, state, smem, tma_store_dx, m, hidden_dim, - num_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits); -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp5_fused_sm90.cuh" +#else +#include "mlp5_fused_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm100.cuh new file mode 100644 index 00000000..6f603752 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm100.cuh @@ -0,0 +1,629 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 5 (single-tile API for fused use) +// ═══════════════════════════════════════════════════════════════════ +// +// Same kernel body as mlp5.cuh's nested-loop default, exposed as a +// single-(m)-tile producer + consumer so a parent kernel can drive the +// outer m loop itself (e.g. mlp_bwd.cuh composing Phase 1d). +// +// Mirrors how mlp1_fused.cuh exposes single-tile entry points: callers +// own the Pipeline construction (use +// mlp5_make_pipe from mlp5.cuh) and the Mlp5Smem allocation; this +// header just provides the producer/consumer body parameterized for +// ONE m-tile. +// +// Params unique to fused use: +// - m : the m-tile this invocation processes +// - expert_k_offset : expert · num_k_tiles (caller resolves the +// expert lookup once per m) +// - split_idx : the column index within an N-split (caller +// picks; analogous to blockIdx.y in the +// standalone driver) +// - num_splits : total columns (analogous to gridDim.y) +// +// Output store goes to a global dX tensor view of shape +// (num_m_tiles · TileM, hidden_dim) (same as the standalone consumer). +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp5_sm100.cuh" // Mlp5Traits, Mlp5Smem, mlp5_make_pipe + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Producer (single m-tile) — fills 2·num_k_tiles slots per (m, n) +// across n ∈ {split_idx, split_idx+num_splits, ...}. +// Phase 1 (k=0..num_k_tiles-1): Z=dU, W=B +// Phase 2 (k=num_k_tiles..2K-1): Z=dV, W=C +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp5_fused_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaLoadZ const& tma_load_du, + TmaLoadZ const& tma_load_dv, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + int m, + int expert_or_k_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_k_cols, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + // int64_t cast for production shapes where T · I or E · I · H can + // exceed 2^31 elements (same rationale as in mlp1_fused_act.cuh and + // the int64 commit on the fwd path). + auto mDU = tma_load_du.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(intermediate_dim))); + auto mDV = tma_load_dv.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(intermediate_dim))); + auto mB = [&]() { + if constexpr (Expert3D) { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(num_experts))); + } else { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(total_k_cols))); + } + }(); + auto mC = [&]() { + if constexpr (Expert3D) { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(num_experts))); + } else { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(total_k_cols))); + } + }(); + + auto cta_tma_du = tma_load_du.get_slice(Int<0>{}); + auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_du.partition_D(sZ); + auto tWsW = cta_tma_b.partition_D(sW); + + auto gB_all = [&]() { + if constexpr (Expert3D) { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(_, _, expert_or_k_offset)); + } else { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(_, _)); + } + }(); + auto gC_all = [&]() { + if constexpr (Expert3D) { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(_, _, expert_or_k_offset)); + } else { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(_, _)); + } + }(); + auto tBgB_all = cta_tma_b.partition_S(gB_all); + auto tCgC_all = cta_tma_b.partition_S(gC_all); + + auto gDU = local_tile(mDU, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgDU = cta_tma_du.partition_S(gDU); + auto gDV = local_tile(mDV, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgDV = cta_tma_du.partition_S(gDV); + + const bool is_leader = (threadIdx.x == 0); + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + // Phase 1: dU @ B + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (is_leader) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_du.with(*bar, 0), + tZgDU(_, _, _, k), tZsZ(_, _, _, state.index())); + if constexpr (Expert3D) { + copy(tma_load_b.with(*bar, 0), + tBgB_all(_, _, _, n, k), tWsW(_, _, _, state.index())); + } else { + copy(tma_load_b.with(*bar, 0), + tBgB_all(_, _, _, n, expert_or_k_offset + k), + tWsW(_, _, _, state.index())); + } + } + ++state; + } + // Phase 2: dV @ C + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (is_leader) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_dv.with(*bar, 0), + tZgDV(_, _, _, k), tZsZ(_, _, _, state.index())); + if constexpr (Expert3D) { + copy(tma_load_c.with(*bar, 0), + tCgC_all(_, _, _, n, k), tWsW(_, _, _, state.index())); + } else { + copy(tma_load_c.with(*bar, 0), + tCgC_all(_, _, _, n, expert_or_k_offset + k), + tWsW(_, _, _, state.index())); + } + } + ++state; + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Consumer (single m-tile) — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported +// compute capability provides a `run(...)` member (a function-template on +// Traits/Pipeline/TmaStoreDX, since function templates can't be partially +// specialized). The free function `mlp5_fused_consumer` below forwards to +// the right specialization — call sites stay in the existing style. +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, single register acc, +// cross-phase sum accumulated across 2·num_k_tiles MMAs) +// Compute=100 → Blackwell / UMMA (single-warp issue, ONE TMEM accumulator, +// cross-phase accumulate via the per-instruction ScaleOut bit) + +template +struct Mlp5FusedConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. The two WGs split (TileM,TileN); the single register +// accumulator carries dU·B (phase 1, k=0..K-1) + dV·C (phase 2, k=K..2K-1) +// via cute::gemm's running accumulation over the whole 2·num_k_tiles k-loop. +// (Verbatim from the original mlp5_fused_consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp5FusedConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaStoreDX const& tma_store_dx, + int m, + int hidden_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; + const int my_barrier_id = 1 + my_wg; + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStore{}); + + auto mDX = tma_store_dx.get_tma_tensor(make_shape( + static_cast(num_m_tiles) * Traits::TileM, + static_cast(hidden_dim))); + auto cta_tma_dx = tma_store_dx.get_slice(Int<0>{}); + + int total_k = 2 * num_k_tiles; + constexpr int K_PIPE_MMAS = 1; + bool store_in_flight = false; + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + + clear(acc); + auto state_release = state; + int prologue_count = (total_k < K_PIPE_MMAS) ? total_k : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < total_k; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + m_tile_idx = 2 * m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + m_tile_idx = m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gDX = local_tile(mDX, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_dx, cta_tma_dx.partition_S(sStore), + cta_tma_dx.partition_D(gDX)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp5FusedConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. +// +// One 1SM UMMA atom with M=TileM covers the whole tile; the single +// accumulator dX = dU·B + dV·C lives in TMEM. Operand A (dU/dV) is K-major; +// operand B (B/C) is MN-major — reusing mlp2_t's column-major weight view +// (SmemLayoutAtomW = Layout_MN_SW128_Atom, already in Mlp5Traits) with the +// UMMA atom's b_major = UMMA::Major::MN (the descriptor builder reads the +// MN-major stride from the smem layout, so no transpose is needed). +// +// Cross-phase accumulate (the crux): the producer fills 2·num_k_tiles pipe +// stages — B in phase 1 (k +struct Mlp5FusedConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaStoreDX const& tma_store_dx, + int m, + int hidden_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two consumer warpgroups"); + constexpr int WgN = TileN / 2; // per-warpgroup N width + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per warpgroup + constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) + + // ── Thread identity ───────────────────────────────────── + // Warp 3 is the dedicated, epilogue-free UMMA producer. Warps 4..11 are + // the two aligned epilogue WGs. + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; // 1 or 2 + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, SS. + // Operand A = dU/dV → K-major; operand B = B/C → MN-major (the + // column-major weight view, matching Mlp5Traits::SmemLayoutW). ── + auto tiled_mma = make_tiled_mma( + SM100_MMA_F16BF16_SS{}); + auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + // ── One TMEM accumulator: dX = dU·B + dV·C ── + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages stages let MMA for the next n-tile + // overlap the previous n-tile's epilogue. + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; // one umma_arrive per commit + acc_params.consumer_arv_count = 1; // one elected releaser + acc_params.initializing_warp = 4; // warp 4 inits the acc barriers + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + // ── Per-WG store slot (AtomTileM × EpiChunkN) — same smem as Hopper ── + constexpr int store_slot_elems = AtomTileM * EpiChunkN; + Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStore{}); + + // dX view: [num_m_tiles·TileM, hidden_dim] row-major. int64_t cast keeps + // CUTE's layout math in 64-bit for large output buffers. + auto mDX = tma_store_dx.get_tma_tensor(make_shape( + static_cast(num_m_tiles) * TileM, + static_cast(hidden_dim))); + auto cta_tma_dx = tma_store_dx.get_slice(Int<0>{}); + + // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local tid). + // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) modes, + // which is what make_tmem_copy's cotiled builder requires. The register + // fragment is sized from partition_D (the DEST coords), NOT partition_S, + // so it excludes the warp-collective datapath-lane dim. ── + auto epi_tile = make_tile(Int{}, Int{}); + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto tAcc_epi = flat_divide(acc_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto t2r = make_tmem_copy(TmemLoadOp{}, tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tid_wg); + auto tTR_tAcc = thr_t2r.partition_S(tAcc_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); // f32 regs + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + const int total_k = 2 * num_k_tiles; // phase 1 (dU·B) + phase 2 (dV·C) + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + // ── Mainloop (warp 3 only): one continuous k-loop over ALL + // 2·num_k_tiles stages into the single TMEM accumulator. The + // accumulate bit is Zero on the very first MMA (clears the acc) and + // One thereafter — including across the phase-1→phase-2 boundary, so + // dX = dU·B + dV·C accumulates without dropping the dU·B term. ── + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); // TMEM acc free (prev epilogue done) + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + for (int k = 0; k < total_k; ++k) { + pipe.consumer_wait(state); + auto tCsZ = cta_mma.partition_A(sZ(_, _, state.index())); + auto tCsW = cta_mma.partition_B(sW(_, _, state.index())); + auto tCrZ = cta_mma.make_fragment_A(tCsZ); + auto tCrW = cta_mma.make_fragment_B(tCsW); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrZ); ++kb) { + // First MMA of the whole 2K loop clears TMEM; every other + // MMA (incl. the phase boundary) accumulates. Do NOT reset. + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrZ(_, _, kb), tCrW(_, _, kb), tCtAcc); + } + pipe.consumer_release(state); // UMMA-gated smem-buffer release + ++state; + } + acc_pipe.producer_commit(acc_prod_state); // signal acc ready (umma_arrive) + ++acc_prod_state; + } + + // ── Epilogue: wait for the accumulator, then this WG processes its + // n-chunks [wg·NChunksHalf, +NChunksHalf). ── + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc_stage = thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; // absolute n-chunk index + + // TMEM → registers (this chunk, full TileM rows). + copy(t2r, tTR_tAcc_stage(_, _, _, _0{}, chunk), tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + + // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rAcc); ++i) { + int m_row = get<0>(tTR_cChunk(i)); // 0..TileM + int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN + if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) + sStore(m_row - ms * AtomTileM, n_col) = + static_cast(tTR_rAcc(i)); + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx = MSub * m + ms; + int n_tile_idx = n * (TileN / EpiChunkN) + chunk; + auto gDX = local_tile(mDX, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_dx, cta_tma_dx.partition_S(sStore), + cta_tma_dx.partition_D(gDX)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // All epilogue TMEM reads are done; release the accumulator so the next + // n-tile's MMA may reuse it (one elected consumer thread arrives). + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + // The TMEM allocation is owned by the launcher (freed once per CTA after the + // m-loop). Do NOT relinquish/free here — this consumer runs once per m-tile + // and freeing per-tile then re-allocating next tile trips the tcgen05 + // "phase invalid during alloc" guardrail. Keep this final 288-thread + // MMA+epilogue barrier so the next m-tile's accumulator-pipeline init is + // isolated from this tile's in-flight handshake. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + // The Compute=100 specialization is never dispatched on non-SM100 targets + // (host picks Compute=90 there). Trap if it is ever reached. + __trap(); +#endif +} +}; // Mlp5FusedConsumerImpl<100> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — keeps the existing call style. `Compute` defaults to 90 so +// current call sites (mlp5_fused_consumer(...)) are unchanged; pass +// mlp5_fused_consumer(...) to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp5_fused_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaStoreDX const& tma_store_dx, + int m, + int hidden_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + Mlp5FusedConsumerImpl::template run( + pipe, state, smem, tma_store_dx, m, hidden_dim, + num_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm90.cuh new file mode 100644 index 00000000..1f130fe5 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm90.cuh @@ -0,0 +1,629 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 5 (single-tile API for fused use) +// ═══════════════════════════════════════════════════════════════════ +// +// Same kernel body as mlp5.cuh's nested-loop default, exposed as a +// single-(m)-tile producer + consumer so a parent kernel can drive the +// outer m loop itself (e.g. mlp_bwd.cuh composing Phase 1d). +// +// Mirrors how mlp1_fused.cuh exposes single-tile entry points: callers +// own the Pipeline construction (use +// mlp5_make_pipe from mlp5.cuh) and the Mlp5Smem allocation; this +// header just provides the producer/consumer body parameterized for +// ONE m-tile. +// +// Params unique to fused use: +// - m : the m-tile this invocation processes +// - expert_k_offset : expert · num_k_tiles (caller resolves the +// expert lookup once per m) +// - split_idx : the column index within an N-split (caller +// picks; analogous to blockIdx.y in the +// standalone driver) +// - num_splits : total columns (analogous to gridDim.y) +// +// Output store goes to a global dX tensor view of shape +// (num_m_tiles · TileM, hidden_dim) (same as the standalone consumer). +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp5_sm90.cuh" // Mlp5Traits, Mlp5Smem, mlp5_make_pipe + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Producer (single m-tile) — fills 2·num_k_tiles slots per (m, n) +// across n ∈ {split_idx, split_idx+num_splits, ...}. +// Phase 1 (k=0..num_k_tiles-1): Z=dU, W=B +// Phase 2 (k=num_k_tiles..2K-1): Z=dV, W=C +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp5_fused_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaLoadZ const& tma_load_du, + TmaLoadZ const& tma_load_dv, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + int m, + int expert_or_k_offset, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int num_experts, + int total_k_cols, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + // int64_t cast for production shapes where T · I or E · I · H can + // exceed 2^31 elements (same rationale as in mlp1_fused_act.cuh and + // the int64 commit on the fwd path). + auto mDU = tma_load_du.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(intermediate_dim))); + auto mDV = tma_load_dv.get_tma_tensor(make_shape( + static_cast(num_tokens), + static_cast(intermediate_dim))); + auto mB = [&]() { + if constexpr (Expert3D) { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(num_experts))); + } else { + return tma_load_b.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(total_k_cols))); + } + }(); + auto mC = [&]() { + if constexpr (Expert3D) { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(num_experts))); + } else { + return tma_load_c.get_tma_tensor(make_shape( + static_cast(hidden_dim), + static_cast(total_k_cols))); + } + }(); + + auto cta_tma_du = tma_load_du.get_slice(Int<0>{}); + auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_du.partition_D(sZ); + auto tWsW = cta_tma_b.partition_D(sW); + + auto gB_all = [&]() { + if constexpr (Expert3D) { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(_, _, expert_or_k_offset)); + } else { + return local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(_, _)); + } + }(); + auto gC_all = [&]() { + if constexpr (Expert3D) { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(_, _, expert_or_k_offset)); + } else { + return local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(_, _)); + } + }(); + auto tBgB_all = cta_tma_b.partition_S(gB_all); + auto tCgC_all = cta_tma_b.partition_S(gC_all); + + auto gDU = local_tile(mDU, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgDU = cta_tma_du.partition_S(gDU); + auto gDV = local_tile(mDV, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgDV = cta_tma_du.partition_S(gDV); + + const bool is_leader = (threadIdx.x == 0); + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + // Phase 1: dU @ B + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (is_leader) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_du.with(*bar, 0), + tZgDU(_, _, _, k), tZsZ(_, _, _, state.index())); + if constexpr (Expert3D) { + copy(tma_load_b.with(*bar, 0), + tBgB_all(_, _, _, n, k), tWsW(_, _, _, state.index())); + } else { + copy(tma_load_b.with(*bar, 0), + tBgB_all(_, _, _, n, expert_or_k_offset + k), + tWsW(_, _, _, state.index())); + } + } + ++state; + } + // Phase 2: dV @ C + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (is_leader) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_dv.with(*bar, 0), + tZgDV(_, _, _, k), tZsZ(_, _, _, state.index())); + if constexpr (Expert3D) { + copy(tma_load_c.with(*bar, 0), + tCgC_all(_, _, _, n, k), tWsW(_, _, _, state.index())); + } else { + copy(tma_load_c.with(*bar, 0), + tCgC_all(_, _, _, n, expert_or_k_offset + k), + tWsW(_, _, _, state.index())); + } + } + ++state; + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Consumer (single m-tile) — architecture-specialized on `int Compute` +// ═══════════════════════════════════════════════════════════════════ +// +// Primary template is undefined; one full specialization per supported +// compute capability provides a `run(...)` member (a function-template on +// Traits/Pipeline/TmaStoreDX, since function templates can't be partially +// specialized). The free function `mlp5_fused_consumer` below forwards to +// the right specialization — call sites stay in the existing style. +// +// Compute=90 → Hopper / WGMMA (cooperative 2-WG, single register acc, +// cross-phase sum accumulated across 2·num_k_tiles MMAs) +// Compute=100 → Blackwell / UMMA (single-warp issue, ONE TMEM accumulator, +// cross-phase accumulate via the per-instruction ScaleOut bit) + +template +struct Mlp5FusedConsumerImpl; + +// ─────────────────────────────────────────────────────────────────── +// Compute=90 — Hopper. The two WGs split (TileM,TileN); the single register +// accumulator carries dU·B (phase 1, k=0..K-1) + dV·C (phase 2, k=K..2K-1) +// via cute::gemm's running accumulation over the whole 2·num_k_tiles k-loop. +// (Verbatim from the original mlp5_fused_consumer.) +// ─────────────────────────────────────────────────────────────────── + +template <> +struct Mlp5FusedConsumerImpl<90> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaStoreDX const& tma_store_dx, + int m, + int hidden_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; + const int my_barrier_id = 1 + my_wg; + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStore{}); + + auto mDX = tma_store_dx.get_tma_tensor(make_shape( + static_cast(num_m_tiles) * Traits::TileM, + static_cast(hidden_dim))); + auto cta_tma_dx = tma_store_dx.get_slice(Int<0>{}); + + int total_k = 2 * num_k_tiles; + constexpr int K_PIPE_MMAS = 1; + bool store_in_flight = false; + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + + clear(acc); + auto state_release = state; + int prologue_count = (total_k < K_PIPE_MMAS) ? total_k : K_PIPE_MMAS; + + // Prologue + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + // Steady state + for (int k = prologue_count; k < total_k; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + // Drain + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // ── Epilogue ────────────────────────────────────── + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + m_tile_idx = 2 * m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + m_tile_idx = m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gDX = local_tile(mDX, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_dx, cta_tma_dx.partition_S(sStore), + cta_tma_dx.partition_D(gDX)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} +}; // Mlp5FusedConsumerImpl<90> + +// ─────────────────────────────────────────────────────────────────── +// Compute=100 — Blackwell / UMMA. +// +// One 1SM UMMA atom with M=TileM covers the whole tile; the single +// accumulator dX = dU·B + dV·C lives in TMEM. Operand A (dU/dV) is K-major; +// operand B (B/C) is MN-major — reusing mlp2_t's column-major weight view +// (SmemLayoutAtomW = Layout_MN_SW128_Atom, already in Mlp5Traits) with the +// UMMA atom's b_major = UMMA::Major::MN (the descriptor builder reads the +// MN-major stride from the smem layout, so no transpose is needed). +// +// Cross-phase accumulate (the crux): the producer fills 2·num_k_tiles pipe +// stages — B in phase 1 (k +struct Mlp5FusedConsumerImpl<100> { +template +static __device__ __forceinline__ void run( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaStoreDX const& tma_store_dx, + int m, + int hidden_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using Element = typename Traits::Element; + constexpr int TileM = Traits::TileM; + constexpr int TileN = Traits::TileN; + constexpr int EpiChunkN = Traits::EpiChunkN; + constexpr int AtomTileM = Traits::AtomTileM; // 64-row TMA store tile + static_assert(TileN % 2 == 0, + "Blackwell consumer splits TileN across the two consumer warpgroups"); + constexpr int WgN = TileN / 2; // per-warpgroup N width + static_assert(WgN % EpiChunkN == 0, "EpiChunkN must divide TileN/2"); + constexpr int NChunksHalf = WgN / EpiChunkN; // n-chunks per warpgroup + constexpr int MSub = TileM / AtomTileM; // 1 (TileM=64) or 2 (TileM=128) + + // ── Thread identity ───────────────────────────────────── + // Warp 3 is the dedicated, epilogue-free UMMA producer. Warps 4..11 are + // the two aligned epilogue WGs. + const int warp_id = threadIdx.x / Traits::WarpSize; + const bool is_mma_warp = (warp_id == 3); + const bool is_epilogue = (warp_id >= 4 && warp_id <= 11); + const int tid_in_epi = threadIdx.x - Traits::WarpGroupSize; // warps 4..11 -> 0..255 + const int wg = is_epilogue ? tid_in_epi / Traits::WarpGroupSize : 0; + const int tid_wg = is_epilogue ? tid_in_epi % Traits::WarpGroupSize : 0; + const bool is_wg_leader = is_epilogue && (tid_wg == 0); + const int wg_barrier_id = 1 + wg; // 1 or 2 + constexpr int kEpilogueThreads = Traits::ConsumerThreads; + constexpr int kMmaEpiThreads = kEpilogueThreads + Traits::WarpSize; + static_assert(Traits::WarpGroupSize == 4 * Traits::WarpSize); + static_assert(kEpilogueThreads == 8 * Traits::WarpSize); + static_assert(kMmaEpiThreads == 9 * Traits::WarpSize); + + // ── TiledMMA: single 1SM UMMA atom, M=TileM, N=TileN, SS. + // Operand A = dU/dV → K-major; operand B = B/C → MN-major (the + // column-major weight view, matching Mlp5Traits::SmemLayoutW). ── + auto tiled_mma = make_tiled_mma( + SM100_MMA_F16BF16_SS{}); + auto cta_mma = tiled_mma.get_slice(0); // 1SM → single CTA, peer-coord 0 + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + // ── One TMEM accumulator: dX = dU·B + dV·C ── + auto cAccFull = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCgC = cta_mma.partition_C(cAccFull); + auto tCtAcc = cta_mma.make_fragment_C(tCgC); + + // TMEM is allocated by the outer fused/standalone launcher once per CTA. + + // ── Accumulator pipeline: UMMA producer (warp 3) → epilogue consumers + // (warps 4..11, both WGs). AccStages stages let MMA for the next n-tile + // overlap the previous n-tile's epilogue. + using AccPipe = typename Traits::AccumulatorPipeline; + typename AccPipe::Params acc_params; + acc_params.role = is_mma_warp ? AccPipe::ThreadCategory::Producer + : AccPipe::ThreadCategory::Consumer; + acc_params.producer_arv_count = 1; // one umma_arrive per commit + acc_params.consumer_arv_count = 1; // one elected releaser + acc_params.initializing_warp = 4; // warp 4 inits the acc barriers + AccPipe acc_pipe(smem.acc_pipe, acc_params, + cute::Shape{}); + auto acc_prod_state = cutlass::make_producer_start_state(); + typename AccPipe::PipelineState acc_cons_state; + + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); + const uint32_t tmem_base = smem.tmem_base; + tCtAcc.data() = tmem_base; + + // ── Per-WG store slot (AtomTileM × EpiChunkN) — same smem as Hopper ── + constexpr int store_slot_elems = AtomTileM * EpiChunkN; + Element* my_store_ptr = smem.store_buf + wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStore{}); + + // dX view: [num_m_tiles·TileM, hidden_dim] row-major. int64_t cast keeps + // CUTE's layout math in 64-bit for large output buffers. + auto mDX = tma_store_dx.get_tma_tensor(make_shape( + static_cast(num_m_tiles) * TileM, + static_cast(hidden_dim))); + auto cta_tma_dx = tma_store_dx.get_slice(Int<0>{}); + + // ── Epilogue TMEM→reg copy plumbing (built once; sliced per WG-local tid). + // flat_divide (not zipped_divide) keeps the epi tile as flat (M,N) modes, + // which is what make_tmem_copy's cotiled builder requires. The register + // fragment is sized from partition_D (the DEST coords), NOT partition_S, + // so it excludes the warp-collective datapath-lane dim. ── + auto epi_tile = make_tile(Int{}, Int{}); + auto acc_mn = tCtAcc(make_coord(_, _), _0{}, _0{}); // (TileM,TileN) + auto tAcc_epi = flat_divide(acc_mn, epi_tile); // (TileM,EpiChunkN,1,TileN/EpiChunkN) + auto t2r = make_tmem_copy(TmemLoadOp{}, tAcc_epi(_, _, _0{}, _0{})); + auto thr_t2r = t2r.get_slice(tid_wg); + auto tTR_tAcc = thr_t2r.partition_S(tAcc_epi); // (Cpy,Cpy_M,Cpy_N,1,nTiles) + auto cChunk = make_identity_tensor(make_shape(Int{}, Int{})); + auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N) + auto tTR_rAcc = make_tensor(shape(tTR_cChunk)); // f32 regs + + int n_start = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_stride = (num_splits >= 0) ? num_splits : (int)gridDim.y; + bool store_in_flight = false; + + const int total_k = 2 * num_k_tiles; // phase 1 (dU·B) + phase 2 (dV·C) + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + // ── Mainloop (warp 3 only): one continuous k-loop over ALL + // 2·num_k_tiles stages into the single TMEM accumulator. The + // accumulate bit is Zero on the very first MMA (clears the acc) and + // One thereafter — including across the phase-1→phase-2 boundary, so + // dX = dU·B + dV·C accumulates without dropping the dU·B term. ── + if (is_mma_warp) { + acc_pipe.producer_acquire(acc_prod_state); // TMEM acc free (prev epilogue done) + int acc_stage = acc_prod_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + for (int k = 0; k < total_k; ++k) { + pipe.consumer_wait(state); + auto tCsZ = cta_mma.partition_A(sZ(_, _, state.index())); + auto tCsW = cta_mma.partition_B(sW(_, _, state.index())); + auto tCrZ = cta_mma.make_fragment_A(tCsZ); + auto tCrW = cta_mma.make_fragment_B(tCsW); + CUTE_UNROLL + for (int kb = 0; kb < size<2>(tCrZ); ++kb) { + // First MMA of the whole 2K loop clears TMEM; every other + // MMA (incl. the phase boundary) accumulates. Do NOT reset. + tiled_mma.accumulate_ = (k == 0 && kb == 0) + ? UMMA::ScaleOut::Zero : UMMA::ScaleOut::One; + gemm(tiled_mma, tCrZ(_, _, kb), tCrW(_, _, kb), tCtAcc); + } + pipe.consumer_release(state); // UMMA-gated smem-buffer release + ++state; + } + acc_pipe.producer_commit(acc_prod_state); // signal acc ready (umma_arrive) + ++acc_prod_state; + } + + // ── Epilogue: wait for the accumulator, then this WG processes its + // n-chunks [wg·NChunksHalf, +NChunksHalf). ── + if (is_epilogue) { + acc_pipe.consumer_wait(acc_cons_state); + int acc_stage = acc_cons_state.index(); + tCtAcc.data() = tmem_base + uint32_t(acc_stage * TileN); + auto acc_mn_stage = tCtAcc(make_coord(_, _), _0{}, _0{}); + auto tAcc_epi_stage = flat_divide(acc_mn_stage, epi_tile); + auto tTR_tAcc_stage = thr_t2r.partition_S(tAcc_epi_stage); + + CUTE_UNROLL + for (int r = 0; r < NChunksHalf; ++r) { + int chunk = wg * NChunksHalf + r; // absolute n-chunk index + + // TMEM → registers (this chunk, full TileM rows). + copy(t2r, tTR_tAcc_stage(_, _, _, _0{}, chunk), tTR_rAcc); + cutlass::arch::fence_view_async_tmem_load(); + + // Store as MSub × (AtomTileM=64)-row TMA tiles (1 for TileM=64, 2 for 128). + CUTE_UNROLL + for (int ms = 0; ms < MSub; ++ms) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + CUTE_UNROLL + for (int i = 0; i < size(tTR_rAcc); ++i) { + int m_row = get<0>(tTR_cChunk(i)); // 0..TileM + int n_col = get<1>(tTR_cChunk(i)); // 0..EpiChunkN + if (m_row >= ms * AtomTileM && m_row < (ms + 1) * AtomTileM) + sStore(m_row - ms * AtomTileM, n_col) = + static_cast(tTR_rAcc(i)); + } + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, wg_barrier_id); + + if (is_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx = MSub * m + ms; + int n_tile_idx = n * (TileN / EpiChunkN) + chunk; + auto gDX = local_tile(mDX, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_dx, cta_tma_dx.partition_S(sStore), + cta_tma_dx.partition_D(gDX)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + + // All epilogue TMEM reads are done; release the accumulator so the next + // n-tile's MMA may reuse it (one elected consumer thread arrives). + cutlass::arch::NamedBarrier::sync(Traits::ConsumerThreads, /*id=*/0); + if (tid_in_epi == 0) + acc_pipe.consumer_release(acc_cons_state); + ++acc_cons_state; + } + } + if (is_epilogue && store_in_flight) + cute::tma_store_wait<0>(); + + // The TMEM allocation is owned by the launcher (freed once per CTA after the + // m-loop). Do NOT relinquish/free here — this consumer runs once per m-tile + // and freeing per-tile then re-allocating next tile trips the tcgen05 + // "phase invalid during alloc" guardrail. Keep this final 288-thread + // MMA+epilogue barrier so the next m-tile's accumulator-pipeline init is + // isolated from this tile's in-flight handshake. + cutlass::arch::NamedBarrier::sync(kMmaEpiThreads, /*id=*/3); +#else + // The Compute=100 specialization is never dispatched on non-SM100 targets + // (host picks Compute=90 there). Trap if it is ever reached. + __trap(); +#endif +} +}; // Mlp5FusedConsumerImpl<100> + +// ─────────────────────────────────────────────────────────────────── +// Forwarder — keeps the existing call style. `Compute` defaults to 90 so +// current call sites (mlp5_fused_consumer(...)) are unchanged; pass +// mlp5_fused_consumer(...) to select the Blackwell path. +// ─────────────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp5_fused_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaStoreDX const& tma_store_dx, + int m, + int hidden_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles, + int split_idx, + int num_splits) { + Mlp5FusedConsumerImpl::template run( + pipe, state, smem, tma_store_dx, m, hidden_dim, + num_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp5_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp5_sm100.cuh new file mode 100644 index 00000000..83ad0188 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp5_sm100.cuh @@ -0,0 +1,599 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 5: dX = dU @ B + dV @ C (input gradient, backward) +// NESTED-LOOP scheduler — the default. +// ═══════════════════════════════════════════════════════════════════ +// +// Cooperative 2-WG consumer; split axis selected by TileM: +// TileM=128 → Layout<_2,_1,_1>, split along M (each WG owns 64×TileN) +// TileM=64 → Layout<_1,_2,_1>, split along N (each WG owns 64×TileN/2) +// ONE accumulator per WG covering (AtomTileM=64, AtomN). Each WG writes +// its strip independently — no cross-WG NamedBarrier sync in the epilogue. +// +// Fused mainloop runs a single continuous (2 × num_k_tiles) k-loop: +// Phase 1 (k=0..num_k_tiles-1): acc += dU @ B +// Phase 2 (k=num_k_tiles..2K-1): acc += dV @ C +// +// Scheduler: +// - 2D grid (num_sms / NSplit, NSplit) — host-side launcher MUST +// match (see mlp5.cu). +// - nested `for (m = blockIdx.x; ...) for (n = blockIdx.y; ...)` +// - amortizes expert lookup + dU/dV tile-coord setup once per m +// - cross-CTA dU/dV multicast in L2 between (m, 0) and (m, 1) in the +// same m-column +// +// Alternative flat work_idx scheduler is preserved in `mlp5_flat.cuh` +// (which includes this header for the shared building blocks). +// +// CUTLASS-canonical patterns: +// - K_PIPE_MMAS = 1 +// - __launch_bounds__(NumThreads, 1) — applied by the caller kernel +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for dU/dV/B/C) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Per `expert_for_m_block`: one int per TileM=128 m-block. +// +// Smem at default config (TileM=128, TileN=256, TileK=64, Stages=4, +// EpiChunkN=64): +// Z (TileM × TileK × S × 2) = 64 KiB +// W (TileN × TileK × S × 2) = 128 KiB +// store_buf (2 WGs × AtomTileM × EpiChunkN × 2) = 16 KiB +// pipe + barriers : ~1 KiB +// ───────────────────────────────────────────── +// Total : ~209 KiB (≤ 228 KiB cap) +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm100.cuh" +#include "tmem_load_op_sm100.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer +// specialization (mlp5_fused.cuh). These headers are header-only and safe to +// include under an sm_90a build — the tcgen05 PTX they emit is guarded by +// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on +// __CUDA_ARCH__ >= 1000) ever instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 256, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 64 +> +struct Mlp5Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + + static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); + + // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate + // along M (each owns 64 rows × TileN). With TileM=64 the WGs cooperate + // along N (each owns 64 rows × TileN/2). AtomTileM is the per-WG M + // extent — always 64. Mirrors mlp2_fused.cuh's kMSplit/N-split pattern. + static constexpr bool kMSplit = (TileM == 128); + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); + + static_assert(TileN <= 256, "TileN must fit TMA box limit"); + static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); + static_assert(kMSplit || (TileN % 2 == 0), + "TileN must be even when TileM=64 (cooperative N-split halves)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + // Cooperative 2-WG layout: M-split for TileM=128, N-split for TileM=64. + // AtomN is the per-WG WGMMA atom N — combined across both WGs it must + // cover exactly TileN cols (= sW's N extent). + static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); + using GmmaAtom = typename GmmaSelectorKMN::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom; + using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom; + + // One-stage layouts (used for TMA descriptors). + using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{}, + Shape, Int>{})); + using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtomW{}, + Shape, Int>{}, Step<_2, _1>{})); + + // Multi-stage smem layouts. + using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{}, + Shape, Int, Int>{})); + using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{}, + Shape, Int, Int>{}, Step<_2, _1, _3>{})); + + // Per-WG store buffer: (AtomTileM, EpiChunkN). One slot per WG. + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; + // Alias for callers that build TMA descriptors using the single-slot + // 2D layout (e.g., moe_bwd's mlp_bwd.cu/moe_bwd.cu). + using SmemLayoutStore_1 = SmemLayoutStore; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop + // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the + // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the single + // TMEM accumulator from the MMA warp to the epilogue warps. See + // Mlp5FusedConsumerImpl<100> in mlp5_fused.cuh. The Blackwell launcher must + // build the mainloop pipe with num_consumers=1 (one umma_arrive per stage). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Double-buffered TMEM accumulators: the single accumulator per stage still + // carries the cross-phase sum dU·B + dV·C across the whole 2·num_k_tiles + // mainloop, while the alternate stage lets the next n-tile's MMA overlap + // the previous n-tile's epilogue. + static constexpr int AccStages = 2; + static_assert(AccStages * TileN <= 512, + "SM100 MLP5 accumulator stages must fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + static constexpr int TmaTransBytesZ = + static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesW = + static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); + static constexpr int TmaTransBytes = TmaTransBytesZ + TmaTransBytesW; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; +}; + +template +using Mlp5MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Shared memory layout +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp5Smem { + using Element = typename Traits::Element; + + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element smem_W[smem_W_size]; + // Per-WG store buffer: WG0 uses store_buf[0..S-1], WG1 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + typename Traits::MainloopPipeline::SharedStorage pipe_storage; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; +}; + +// ═══════════════════════════════════════════════════════════════════ +// Pipeline construction +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ auto mlp5_make_pipe( + typename Traits::MainloopPipeline::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + // Both consumer WGs (cooperative) consume each pipe slot. + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp5_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer +// (mlp5_fused_producer) and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ auto mlp5_make_pipe_umma( + typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ── Producer (nested) ─────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp5_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaLoadZ const& tma_load_du, + TmaLoadZ const& tma_load_dv, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + const int* expert_for_m_block, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int total_k_cols, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles) { + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + auto mDU = tma_load_du.get_tma_tensor(make_shape(num_tokens, intermediate_dim)); + auto mDV = tma_load_dv.get_tma_tensor(make_shape(num_tokens, intermediate_dim)); + auto mB = tma_load_b.get_tma_tensor(make_shape(hidden_dim, total_k_cols)); + auto mC = tma_load_c.get_tma_tensor(make_shape(hidden_dim, total_k_cols)); + + auto cta_tma_du = tma_load_du.get_slice(Int<0>{}); + auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_du.partition_D(sZ); + auto tWsW = cta_tma_b.partition_D(sW); + + auto gB_all = local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(_, _)); + auto gC_all = local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(_, _)); + auto tBgB_all = cta_tma_b.partition_S(gB_all); + auto tCgC_all = cta_tma_b.partition_S(gC_all); + + bool is_leader = (threadIdx.x == 0); + + int n_start = (int)blockIdx.y; + int n_stride = (int)gridDim.y; + + for (int m = (int)blockIdx.x; m < num_m_tiles; m += (int)gridDim.x) { + int expert = expert_for_m_block[m]; + int expert_k_offset = expert * num_k_tiles; + + auto gDU = local_tile(mDU, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgDU = cta_tma_du.partition_S(gDU); + auto gDV = local_tile(mDV, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgDV = cta_tma_du.partition_S(gDV); + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + // Phase 1: dU @ B + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (is_leader) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_du.with(*bar, 0), + tZgDU(_, _, _, k), tZsZ(_, _, _, state.index())); + copy(tma_load_b.with(*bar, 0), + tBgB_all(_, _, _, n, expert_k_offset + k), + tWsW(_, _, _, state.index())); + } + ++state; + } + // Phase 2: dV @ C + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (is_leader) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_dv.with(*bar, 0), + tZgDV(_, _, _, k), tZsZ(_, _, _, state.index())); + copy(tma_load_c.with(*bar, 0), + tCgC_all(_, _, _, n, expert_k_offset + k), + tWsW(_, _, _, state.index())); + } + ++state; + } + } + } +} + +// ── Consumer (nested) ─────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp5_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaStoreDX const& tma_store_dx, + int hidden_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles) { + + (void)Compute; + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; + const int my_barrier_id = 1 + my_wg; + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStore{}); + + auto mDX = tma_store_dx.get_tma_tensor( + make_shape(num_m_tiles * Traits::TileM, hidden_dim)); + auto cta_tma_dx = tma_store_dx.get_slice(Int<0>{}); + + int total_k = 2 * num_k_tiles; + constexpr int K_PIPE_MMAS = 1; + bool store_in_flight = false; + + int n_start = (int)blockIdx.y; + int n_stride = (int)gridDim.y; + + for (int m = (int)blockIdx.x; m < num_m_tiles; m += (int)gridDim.x) { + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc); + auto state_release = state; + int prologue_count = (total_k < K_PIPE_MMAS) ? total_k : K_PIPE_MMAS; + + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + for (int k = prologue_count; k < total_k; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. + // TileM=64 (N-split): WG_w owns full M (=64), cols [w*WgTileN, ...). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: 2 row tiles per logical M tile. + m_tile_idx = 2 * m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M tile; WG_w handles + // its N-half. + m_tile_idx = m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gDX = local_tile(mDX, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_dx, cta_tma_dx.partition_S(sStore), + cta_tma_dx.partition_D(gDX)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} + +// ── Driver ────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp5_fwd( + Mlp5Smem& smem, + TmaLoadZ const& tma_load_du, + TmaLoadZ const& tma_load_dv, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + TmaStoreDX const& tma_store_dx, + const int* expert_for_m_block, + int num_tokens, + int intermediate_dim, + int hidden_dim, + int total_k_cols, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles) { + + using Pipeline = typename Traits::MainloopPipeline; + using PipeState = typename Traits::PipelineState; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + auto pipe = mlp5_make_pipe(smem.pipe_storage); + + __syncthreads(); + + PipeState state = is_producer ? + cutlass::make_producer_start_state() : PipeState{}; + PipeState cons_state; + + if (is_producer) { + mlp5_producer(pipe, state, smem, + tma_load_du, tma_load_dv, tma_load_b, tma_load_c, + expert_for_m_block, + num_tokens, hidden_dim, intermediate_dim, total_k_cols, + num_m_tiles, num_n_tiles, num_k_tiles); + } + if (is_consumer) { + mlp5_consumer(pipe, cons_state, smem, + tma_store_dx, hidden_dim, + num_m_tiles, num_n_tiles, num_k_tiles); + } + + __syncthreads(); +} + +// ═══════════════════════════════════════════════════════════════════ +// Host launcher (raw-pointer; defined in mlp5.cu) +// ═══════════════════════════════════════════════════════════════════ +// +// dX = dU @ B + dV @ C. Output dX is fully written; no zero-init needed. +// +// dU : [num_tokens, intermediate_dim] +// dV : [num_tokens, intermediate_dim] +// B, C : [num_experts, intermediate_dim, hidden_dim] +// dX : [num_tokens, hidden_dim] (output) +// expert_for_m_block: [num_tokens / TileM] int32 (TileM is Traits::TileM) +void mlp5_fwd_bf16_launch( + const cutlass::bfloat16_t* dU, + const cutlass::bfloat16_t* dV, + const cutlass::bfloat16_t* B, + const cutlass::bfloat16_t* C, + cutlass::bfloat16_t* dX, + const int* expert_for_m_block, + int num_tokens, + int intermediate_dim, + int hidden_dim, + int num_experts, + int device, + cudaStream_t stream); + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp5_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp5_sm90.cuh new file mode 100644 index 00000000..c619bb97 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp5_sm90.cuh @@ -0,0 +1,599 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Phase 5: dX = dU @ B + dV @ C (input gradient, backward) +// NESTED-LOOP scheduler — the default. +// ═══════════════════════════════════════════════════════════════════ +// +// Cooperative 2-WG consumer; split axis selected by TileM: +// TileM=128 → Layout<_2,_1,_1>, split along M (each WG owns 64×TileN) +// TileM=64 → Layout<_1,_2,_1>, split along N (each WG owns 64×TileN/2) +// ONE accumulator per WG covering (AtomTileM=64, AtomN). Each WG writes +// its strip independently — no cross-WG NamedBarrier sync in the epilogue. +// +// Fused mainloop runs a single continuous (2 × num_k_tiles) k-loop: +// Phase 1 (k=0..num_k_tiles-1): acc += dU @ B +// Phase 2 (k=num_k_tiles..2K-1): acc += dV @ C +// +// Scheduler: +// - 2D grid (num_sms / NSplit, NSplit) — host-side launcher MUST +// match (see mlp5.cu). +// - nested `for (m = blockIdx.x; ...) for (n = blockIdx.y; ...)` +// - amortizes expert lookup + dU/dV tile-coord setup once per m +// - cross-CTA dU/dV multicast in L2 between (m, 0) and (m, 1) in the +// same m-column +// +// Alternative flat work_idx scheduler is preserved in `mlp5_flat.cuh` +// (which includes this header for the shared building blocks). +// +// CUTLASS-canonical patterns: +// - K_PIPE_MMAS = 1 +// - __launch_bounds__(NumThreads, 1) — applied by the caller kernel +// +// CTA layout (384 threads = 12 warps): +// WG0 warp 0 : Producer (TMA loads for dU/dV/B/C) +// WG0 warps 1-2 : Idle / reserved for fused-kernel communication +// WG0 warp 3 : Idle on SM90; dedicated UMMA producer on SM100 +// WG1 (warps 4-7) : WGMMA consumer on SM90; epilogue WG0 on SM100 +// WG2 (warps 8-11) : WGMMA consumer on SM90; epilogue WG1 on SM100 +// +// Per `expert_for_m_block`: one int per TileM=128 m-block. +// +// Smem at default config (TileM=128, TileN=256, TileK=64, Stages=4, +// EpiChunkN=64): +// Z (TileM × TileK × S × 2) = 64 KiB +// W (TileN × TileK × S × 2) = 128 KiB +// store_buf (2 WGs × AtomTileM × EpiChunkN × 2) = 16 KiB +// pipe + barriers : ~1 KiB +// ───────────────────────────────────────────── +// Total : ~209 KiB (≤ 228 KiB cap) +// ═══════════════════════════════════════════════════════════════════ + +#include "models_sm90.cuh" +#include "tmem_load_op_sm90.cuh" + +#include +#include +#include +#include +#include + +// SM100 (Blackwell / UMMA) support for the Compute=100 consumer +// specialization (mlp5_fused.cuh). These headers are header-only and safe to +// include under an sm_90a build — the tcgen05 PTX they emit is guarded by +// CUTE_ARCH_TCGEN05_* and only the Compute=100 body (itself gated on +// __CUDA_ARCH__ >= 1000) ever instantiates it. +#include +#include + +#include +#include +#include // PipelineTmaUmmaAsync, PipelineUmmaAsync + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Traits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileM_ = 128, + int TileN_ = 256, + int TileK_ = 64, + int Stages_ = 4, + int EpiChunkN_ = 64 +> +struct Mlp5Traits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = TileM_; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int EpiChunkN = EpiChunkN_; + + static_assert(TileM == 64 || TileM == 128, "TileM must be 64 or 128"); + + // Hopper WGMMA atom has m=64. With TileM=128 the two WGs cooperate + // along M (each owns 64 rows × TileN). With TileM=64 the WGs cooperate + // along N (each owns 64 rows × TileN/2). AtomTileM is the per-WG M + // extent — always 64. Mirrors mlp2_fused.cuh's kMSplit/N-split pattern. + static constexpr bool kMSplit = (TileM == 128); + static constexpr int AtomTileM = 64; + static constexpr int WgTileN = kMSplit ? TileN : (TileN / 2); + + static_assert(TileN <= 256, "TileN must fit TMA box limit"); + static_assert(WgTileN % EpiChunkN == 0, "EpiChunkN must divide WgTileN"); + static_assert(kMSplit || (TileN % 2 == 0), + "TileN must be even when TileM=64 (cooperative N-split halves)"); + static constexpr int NumEpiRounds = WgTileN / EpiChunkN; + + // Cooperative 2-WG layout: M-split for TileM=128, N-split for TileM=64. + // AtomN is the per-WG WGMMA atom N — combined across both WGs it must + // cover exactly TileN cols (= sW's N extent). + static constexpr int AtomN = kMSplit ? TileN : (TileN / 2); + using GmmaAtom = typename GmmaSelectorKMN::Atom; + using WgLayout = cute::conditional_t< + kMSplit, + Layout>, + Layout>>; + using TiledMma = TiledMMA, WgLayout>; + + using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom; + using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom; + + // One-stage layouts (used for TMA descriptors). + using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{}, + Shape, Int>{})); + using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtomW{}, + Shape, Int>{}, Step<_2, _1>{})); + + // Multi-stage smem layouts. + using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{}, + Shape, Int, Int>{})); + using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{}, + Shape, Int, Int>{}, Step<_2, _1, _3>{})); + + // Per-WG store buffer: (AtomTileM, EpiChunkN). One slot per WG. + using SmemLayoutStore = Layout< + Shape, Int>, + Stride, _1>>; + // Alias for callers that build TMA descriptors using the single-slot + // 2D layout (e.g., moe_bwd's mlp_bwd.cu/moe_bwd.cu). + using SmemLayoutStore_1 = SmemLayoutStore; + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + // SM100 (Compute=100) pipelines — single CTA, no cluster. The mainloop + // pipeline is the UMMA-aware TMA pipeline (its consumer_release issues the + // UMMA-gated smem-buffer arrival); the accumulator pipeline hands the single + // TMEM accumulator from the MMA warp to the epilogue warps. See + // Mlp5FusedConsumerImpl<100> in mlp5_fused.cuh. The Blackwell launcher must + // build the mainloop pipe with num_consumers=1 (one umma_arrive per stage). + using MainloopPipelineUmma = cutlass::PipelineTmaUmmaAsync< + Stages, cute::Shape, + cute::Shape>; + // Double-buffered TMEM accumulators: the single accumulator per stage still + // carries the cross-phase sum dU·B + dV·C across the whole 2·num_k_tiles + // mainloop, while the alternate stage lets the next n-tile's MMA overlap + // the previous n-tile's epilogue. + static constexpr int AccStages = 2; + static_assert(AccStages * TileN <= 512, + "SM100 MLP5 accumulator stages must fit in 512 TMEM columns"); + using AccumulatorPipeline = cutlass::PipelineUmmaAsync; + + static constexpr int TmaTransBytesZ = + static_cast(size(SmemLayoutZ_1{}) * sizeof(Element)); + static constexpr int TmaTransBytesW = + static_cast(size(SmemLayoutW_1{}) * sizeof(Element)); + static constexpr int TmaTransBytes = TmaTransBytesZ + TmaTransBytesW; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumConsumers = 2; + static constexpr int ConsumerThreads = NumConsumers * WarpGroupSize; + static constexpr int NumThreads = 384; +}; + +template +using Mlp5MainloopPipelineFor = cute::conditional_t< + Compute == 100, + typename Traits::MainloopPipelineUmma, + typename Traits::MainloopPipeline>; + +// ═══════════════════════════════════════════════════════════════════ +// Shared memory layout +// ═══════════════════════════════════════════════════════════════════ + +template +struct Mlp5Smem { + using Element = typename Traits::Element; + + static constexpr int smem_Z_size = cosize_v; + static constexpr int smem_W_size = cosize_v; + static constexpr int smem_store_size = cosize_v; + + alignas(128) Element smem_Z[smem_Z_size]; + alignas(128) Element smem_W[smem_W_size]; + // Per-WG store buffer: WG0 uses store_buf[0..S-1], WG1 uses [S..2S-1]. + alignas(128) Element store_buf[2 * smem_store_size]; + + typename Traits::MainloopPipeline::SharedStorage pipe_storage; + + // SM100 (Compute=100) only: landing slot for tcgen05.alloc's granted TMEM + // base address, plus the accumulator pipeline's barrier storage (UMMA→ + // epilogue handoff). The Hopper (Compute=90) path never touches these. + alignas(16) uint32_t tmem_base; + alignas(16) typename Traits::AccumulatorPipeline::SharedStorage acc_pipe; +}; + +// ═══════════════════════════════════════════════════════════════════ +// Pipeline construction +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ auto mlp5_make_pipe( + typename Traits::MainloopPipeline::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipeline; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + // Both consumer WGs (cooperative) consume each pipe slot. + pp.num_consumers = Traits::ConsumerThreads; + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// SM100 (Compute=100) mainloop pipe. Mirrors mlp5_make_pipe but builds the +// UMMA-aware TMA pipeline (PipelineTmaUmmaAsync) whose consumer_release issues +// the UMMA-gated smem-buffer arrival. num_consumers=1: on Blackwell the buffer +// is released by a single umma_arrive per stage (the MMA warp), not per +// consumer thread. Params/PipelineState/SharedStorage alias the Hopper +// PipelineTmaAsync's, and the ctor arg pattern is identical, so the producer +// (mlp5_fused_producer) and the Compute=100 consumer drive it unchanged. +template +__device__ __forceinline__ auto mlp5_make_pipe_umma( + typename Traits::MainloopPipelineUmma::SharedStorage& storage) { + + using Pipeline = typename Traits::MainloopPipelineUmma; + using Category = typename Pipeline::ThreadCategory; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 3 && warp_id <= 11); // warp 3 = UMMA producer + + typename Pipeline::Params pp; + pp.transaction_bytes = Traits::TmaTransBytes; + pp.num_producers = 1; + pp.num_consumers = 1; // UMMA: one umma_arrive per stage releases the buffer + if (is_producer) { + pp.role = Category::Producer; + pp.is_leader = (threadIdx.x == 0); + } else if (is_consumer) { + pp.role = Category::Consumer; + } else { + pp.role = Category::NonParticipant; + } + return Pipeline(storage, pp, Shape<_1, _1, _1>{}, + cute::true_type{}, cute::true_type{}); +} + +// ── Producer (nested) ─────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp5_producer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaLoadZ const& tma_load_du, + TmaLoadZ const& tma_load_dv, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + const int* expert_for_m_block, + int num_tokens, + int hidden_dim, + int intermediate_dim, + int total_k_cols, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles) { + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + auto mDU = tma_load_du.get_tma_tensor(make_shape(num_tokens, intermediate_dim)); + auto mDV = tma_load_dv.get_tma_tensor(make_shape(num_tokens, intermediate_dim)); + auto mB = tma_load_b.get_tma_tensor(make_shape(hidden_dim, total_k_cols)); + auto mC = tma_load_c.get_tma_tensor(make_shape(hidden_dim, total_k_cols)); + + auto cta_tma_du = tma_load_du.get_slice(Int<0>{}); + auto cta_tma_b = tma_load_b.get_slice(Int<0>{}); + + auto tZsZ = cta_tma_du.partition_D(sZ); + auto tWsW = cta_tma_b.partition_D(sW); + + auto gB_all = local_tile(mB, + make_tile(Int{}, Int{}), + make_coord(_, _)); + auto gC_all = local_tile(mC, + make_tile(Int{}, Int{}), + make_coord(_, _)); + auto tBgB_all = cta_tma_b.partition_S(gB_all); + auto tCgC_all = cta_tma_b.partition_S(gC_all); + + bool is_leader = (threadIdx.x == 0); + + int n_start = (int)blockIdx.y; + int n_stride = (int)gridDim.y; + + for (int m = (int)blockIdx.x; m < num_m_tiles; m += (int)gridDim.x) { + int expert = expert_for_m_block[m]; + int expert_k_offset = expert * num_k_tiles; + + auto gDU = local_tile(mDU, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgDU = cta_tma_du.partition_S(gDU); + auto gDV = local_tile(mDV, + make_tile(Int{}, Int{}), + make_coord(m, _)); + auto tZgDV = cta_tma_du.partition_S(gDV); + + for (int n = n_start; n < num_n_tiles; n += n_stride) { + // Phase 1: dU @ B + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (is_leader) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_du.with(*bar, 0), + tZgDU(_, _, _, k), tZsZ(_, _, _, state.index())); + copy(tma_load_b.with(*bar, 0), + tBgB_all(_, _, _, n, expert_k_offset + k), + tWsW(_, _, _, state.index())); + } + ++state; + } + // Phase 2: dV @ C + for (int k = 0; k < num_k_tiles; ++k) { + pipe.producer_acquire(state); + if (is_leader) { + auto* bar = pipe.producer_get_barrier(state); + copy(tma_load_dv.with(*bar, 0), + tZgDV(_, _, _, k), tZsZ(_, _, _, state.index())); + copy(tma_load_c.with(*bar, 0), + tCgC_all(_, _, _, n, expert_k_offset + k), + tWsW(_, _, _, state.index())); + } + ++state; + } + } + } +} + +// ── Consumer (nested) ─────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp5_consumer( + Pipeline& pipe, + typename Traits::PipelineState& state, + Mlp5Smem& smem, + TmaStoreDX const& tma_store_dx, + int hidden_dim, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles) { + + (void)Compute; + using Element = typename Traits::Element; + typename Traits::TiledMma tiled_mma; + int tid_in_mma = threadIdx.x - Traits::WarpGroupSize; + int tid_in_wg = tid_in_mma % Traits::WarpGroupSize; + auto thr_mma = tiled_mma.get_slice(tid_in_mma); + + auto sZ = make_tensor(make_smem_ptr(smem.smem_Z), typename Traits::SmemLayoutZ{}); + auto sW = make_tensor(make_smem_ptr(smem.smem_W), typename Traits::SmemLayoutW{}); + + auto tCsZ = thr_mma.partition_A(sZ); + auto tCsW = thr_mma.partition_B(sW); + + auto acc = partition_fragment_C(tiled_mma, + Shape, Int>{}); + + auto cC = make_identity_tensor(make_shape(Int{}, Int{})); + auto tCcC = thr_mma.partition_C(cC); + + const int my_wg = (threadIdx.x / Traits::WarpGroupSize) - 1; + const int my_barrier_id = 1 + my_wg; + const bool is_my_wg_leader = (tid_in_wg == 0); + + constexpr int store_slot_elems = Traits::AtomTileM * Traits::EpiChunkN; + Element* my_store_ptr = smem.store_buf + my_wg * store_slot_elems; + auto sStore = make_tensor(make_smem_ptr(my_store_ptr), + typename Traits::SmemLayoutStore{}); + + auto mDX = tma_store_dx.get_tma_tensor( + make_shape(num_m_tiles * Traits::TileM, hidden_dim)); + auto cta_tma_dx = tma_store_dx.get_slice(Int<0>{}); + + int total_k = 2 * num_k_tiles; + constexpr int K_PIPE_MMAS = 1; + bool store_in_flight = false; + + int n_start = (int)blockIdx.y; + int n_stride = (int)gridDim.y; + + for (int m = (int)blockIdx.x; m < num_m_tiles; m += (int)gridDim.x) { + for (int n = n_start; n < num_n_tiles; n += n_stride) { + + clear(acc); + auto state_release = state; + int prologue_count = (total_k < K_PIPE_MMAS) ? total_k : K_PIPE_MMAS; + + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + ++state; + } + for (int k = prologue_count; k < total_k; ++k) { + pipe.consumer_wait(state); + warpgroup_fence_operand(acc); + warpgroup_arrive(); + gemm(tiled_mma, tCsZ(_, _, _, state.index()), + tCsW(_, _, _, state.index()), acc); + warpgroup_commit_batch(); + + warpgroup_wait(); + warpgroup_fence_operand(acc); + pipe.consumer_release(state_release); + ++state; + ++state_release; + } + warpgroup_wait<0>(); + warpgroup_fence_operand(acc); + for (int k = 0; k < prologue_count; ++k) { + pipe.consumer_release(state_release); + ++state_release; + } + + // TileM=128 (M-split): WG_w owns rows [w*64, (w+1)*64), full N. + // TileM=64 (N-split): WG_w owns full M (=64), cols [w*WgTileN, ...). + CUTE_UNROLL + for (int r = 0; r < Traits::NumEpiRounds; ++r) { + if (store_in_flight) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + int chunk_start = r * Traits::EpiChunkN; + CUTE_UNROLL + for (int i = 0; i < size(acc); ++i) { + auto coord = tCcC(i); + int m_loc = get<0>(coord); + int n_loc = get<1>(coord); + int m_local, n_local; + if constexpr (Traits::kMSplit) { + m_local = m_loc - my_wg * Traits::AtomTileM; + n_local = n_loc; + } else { + m_local = m_loc; + n_local = n_loc - my_wg * Traits::WgTileN; + } + if (n_local >= chunk_start && + n_local < chunk_start + Traits::EpiChunkN) { + int chunk_n = n_local - chunk_start; + sStore(m_local, chunk_n) = static_cast(acc(i)); + } + } + + cutlass::arch::NamedBarrier::sync(Traits::WarpGroupSize, my_barrier_id); + + if (is_my_wg_leader) { + cute::tma_store_fence(); + int m_tile_idx, n_tile_idx; + if constexpr (Traits::kMSplit) { + // TileM=128: 2 row tiles per logical M tile. + m_tile_idx = 2 * m + my_wg; + n_tile_idx = n * Traits::NumEpiRounds + r; + } else { + // TileM=64: 1 row tile per logical M tile; WG_w handles + // its N-half. + m_tile_idx = m; + n_tile_idx = n * (Traits::TileN / Traits::EpiChunkN) + + my_wg * Traits::NumEpiRounds + r; + } + auto gDX = local_tile(mDX, + make_tile(Int{}, Int{}), + make_coord(m_tile_idx, n_tile_idx)); + copy(tma_store_dx, cta_tma_dx.partition_S(sStore), + cta_tma_dx.partition_D(gDX)); + cute::tma_store_arrive(); + } + store_in_flight = true; + } + } + } + if (store_in_flight) + cute::tma_store_wait<0>(); +} + +// ── Driver ────────────────────────────────────────────────────────── + +template +__device__ __forceinline__ void mlp5_fwd( + Mlp5Smem& smem, + TmaLoadZ const& tma_load_du, + TmaLoadZ const& tma_load_dv, + TmaLoadW const& tma_load_b, + TmaLoadW const& tma_load_c, + TmaStoreDX const& tma_store_dx, + const int* expert_for_m_block, + int num_tokens, + int intermediate_dim, + int hidden_dim, + int total_k_cols, + int num_m_tiles, + int num_n_tiles, + int num_k_tiles) { + + using Pipeline = typename Traits::MainloopPipeline; + using PipeState = typename Traits::PipelineState; + + int warp_id = threadIdx.x / Traits::WarpSize; + bool is_producer = (warp_id == 0); + bool is_consumer = (warp_id >= 4 && warp_id <= 11); + + auto pipe = mlp5_make_pipe(smem.pipe_storage); + + __syncthreads(); + + PipeState state = is_producer ? + cutlass::make_producer_start_state() : PipeState{}; + PipeState cons_state; + + if (is_producer) { + mlp5_producer(pipe, state, smem, + tma_load_du, tma_load_dv, tma_load_b, tma_load_c, + expert_for_m_block, + num_tokens, hidden_dim, intermediate_dim, total_k_cols, + num_m_tiles, num_n_tiles, num_k_tiles); + } + if (is_consumer) { + mlp5_consumer(pipe, cons_state, smem, + tma_store_dx, hidden_dim, + num_m_tiles, num_n_tiles, num_k_tiles); + } + + __syncthreads(); +} + +// ═══════════════════════════════════════════════════════════════════ +// Host launcher (raw-pointer; defined in mlp5.cu) +// ═══════════════════════════════════════════════════════════════════ +// +// dX = dU @ B + dV @ C. Output dX is fully written; no zero-init needed. +// +// dU : [num_tokens, intermediate_dim] +// dV : [num_tokens, intermediate_dim] +// B, C : [num_experts, intermediate_dim, hidden_dim] +// dX : [num_tokens, hidden_dim] (output) +// expert_for_m_block: [num_tokens / TileM] int32 (TileM is Traits::TileM) +void mlp5_fwd_bf16_launch( + const cutlass::bfloat16_t* dU, + const cutlass::bfloat16_t* dV, + const cutlass::bfloat16_t* B, + const cutlass::bfloat16_t* C, + cutlass::bfloat16_t* dX, + const int* expert_for_m_block, + int num_tokens, + int intermediate_dim, + int hidden_dim, + int num_experts, + int device, + cudaStream_t stream); + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_bwd.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_bwd.cuh index db10238e..7adcc7b2 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp_bwd.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_bwd.cuh @@ -1,1885 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// Fused MLP Backward — batch-interleaved Phase 1 + Phase 2 -// ═══════════════════════════════════════════════════════════════════ -// -// Mirrors forward mlp.cuh: pipeline construction + tile/batch loop. -// Called once for local tiles, once for remote tiles. -// -// Per batch (gridDim.x tiles): -// Phase 1: one tile per CTA (a→b→c→d sub-phases) -// Global barrier -// Phase 2: all CTAs do weight grads for this batch's K-blocks -// Global barrier -// -// ═══════════════════════════════════════════════════════════════════ - -#include "mlp.cuh" // kMlpBarrierId, mlp1 pipe helpers -#include "mlp1_fused_act.cuh" // bwd Phase 1a — TileM=128 cooperative-M-split -#include "mlp2_t_fused.cuh" // bwd Phase 1b' — single-m-tile entry over mlp2_t.cuh -#include "mlp5.cuh" -#include "mlp5_fused.cuh" // mlp5_fused_producer/consumer (single-m-tile entry) -#include "mlp3.cuh" // Mlp3Traits, Mlp3FusedSmem -#include "mlp4.cuh" // Mlp4Traits, Mlp4FusedSmem -#include "silu_bwd_fused.cuh" // silu_bwd_pair_tile — separate elementwise phase -#include "cta_barrier.cuh" // target-based CtaCounterBarrier - -#include // INT_MAX for Phase 2 per-batch range scan - -namespace liger { - -using namespace cute; - -// MLP-warps barrier count for the bwd kernel. Hopper uses the historical -// 9-warp partition (w0 + w4-11). Blackwell mirrors the forward path: comm -// uses only warps 1-2 and warp 3 joins MLP as the dedicated UMMA producer. -template -static constexpr int kMlpBwdBarrierThreadsFor = (Compute == 100 ? 10 : 9) * 32; -static_assert(kMlpBwdBarrierThreadsFor<90> == 288); -static_assert(kMlpBwdBarrierThreadsFor<100> == 320); - -// Compute == 100 (Blackwell/SM100) always means the Phase-2 GEMMs (mlp3 AND -// mlp4) run the SM100 ClusterM=2 paired-CTA path; Compute == 90 (Hopper/SM90) -// always means plain 1SM. There is no per-phase or per-trait choice anymore -// (see mlp3.cuh / mlp4.cuh's unified mlp{3,4}_producer / -// mlp{3,4}_consumer), so `Compute == 100` alone drives the -// TMEM allocator choice (Allocator2Sm vs Allocator1Sm) and the extra -// tmem_pair_barrier below: both mlp3 and mlp4 share ONE TMEM allocation per -// CTA-pair, so the pair must rendezvous around the allocate/free calls. - -// ── Compute == 100 (paired-CTA) caller contract ── -// -// This header only builds the __device__ Phase-2 GEMM bodies; it does not -// launch the kernel or build TMA descriptors, so the following invariants -// must hold at the call site (host launcher, e.g. moe_bwd.cu) whenever -// Compute == 100 is selected — violating any of them is a silent-corruption -// or launch-failure risk, not something this file can assert at compile time: -// • Grid must be launched EVEN (gridDim.x·gridDim.y divisible by -// Traits{3,4}::ClusterM) — the flat cell-walk above (cell_start = -// flat_id / ClusterM, cell_stride = total_phase2_ctas / ClusterM) -// assumes every CTA has a live cluster peer. -// • Launch via cudaLaunchKernelEx with clusterDim = (ClusterM, 1, 1) (or -// the ClusterShape Traits{3,4} defines) and the -// cudaLaunchAttributeClusterDimension / -// NonPortableClusterSizeAllowed attribute set (ClusterM > 8 or any -// non-default shape needs the "non-portable" opt-in). -// • The dYT/Z (mlp3) and X/dU^T/dV^T (mlp4) TMA descriptors passed in -// (TmaLoadDYT3/TmaLoadZT3/... below) must be built with -// make_tma_copy_A_sm100 / make_tma_copy_B_sm100 (2-CTA-aware multicast -// descriptors) — a plain SM90 make_tma_copy descriptor does not carry -// the multicast-mask semantics mlp{3,4}_producer's Compute == 100 -// `copy(tma.with(barrier, mcast_mask), ...)` calls rely on. -// • MlpFusedBwdSmem's static size (larger for the Compute == 100 -// Mlp3FusedSmem2Sm variant) must fit within the dynamically-queried -// cudaDevAttrMaxSharedMemoryPerBlockOptin limit, and the kernel's -// cudaFuncAttributeMaxDynamicSharedMemorySize must be opted in to that -// size before launch — this header assumes the smem struct it defines -// already fits whatever the host allocated. -// See Compute == 100's use below (TMEM Allocator2Sm + tmem_pair_barrier) and -// the Compute == 90/100 static_asserts in mlp_fused_bwd_run_batches / -// mlp_fused_bwd / mlp_fused_bwd_dual for the invariants this file DOES -// enforce. - -// ═══════════════════════════════════════════════════════════════════ -// Phase-1 sub-batch grouping (mlp1/2/5 small chunks, mlp3/4 one big batch) -// ═══════════════════════════════════════════════════════════════════ -// -// Remote BWD groups kBwdSubBatch Phase-1 sub-batches before running ONE -// Phase-2 (mlp3/mlp4) over all of them at once. Phase 1 (mlp1/2/5) prefers -// SMALL per-batch chunks (less CTA barrier-stall); Phase 2 (weight grads, -// dA/dB/dC) prefers a LONG K-reduction so the cooperative cell-walk fills -// the grid (deeper outer_split / k_split). Grouping serves both. -// -// Buffer survival: the SubBatch sub-batches keep their X/dY comm-staging slots -// un-released until the grouped Phase 2 consumes them, and their Z/dU/dV scratch -// slots must not be overwritten. This needs the staging ring to be DEEPER than -// the group span — remote_num_m_tiles > SubBatch·grid_x — so the group's live -// slots fit with comm prefetch headroom (host-checked in moe_bwd.cu). No ring/ -// group ALIGNMENT is required: the grouped Phase-2 window wraps the ring exactly -// as the original per-batch loop does (a window that runs past the ring end is -// only correct when it does not actually straddle — i.e. global_total ≤ ring or -// grid_x·SubBatch divides the ring — same property the un-grouped loop has). -// SubBatch == 1 reduces the grouped loop EXACTLY to the original per-batch loop. -// -// The sub-batch count is a pure TEMPLATE ARGUMENT (SubBatch) on moe_fused_bwd / -// mlp_fused_bwd_dual, sourced from MoeBwdConfig::kSubBatch (whose default is a -// literal — see moe_bwd.cu). Distinct SubBatch instantiations coexist in one -// binary; there is no preprocessor knob. -static constexpr int kBwdSubBatchMax = 8; // upper bound for the per-group save arrays - -// ═══════════════════════════════════════════════════════════════════ -// Fused backward shared memory — union of all phases -// ═══════════════════════════════════════════════════════════════════ - -template -struct MlpFusedBwdSmem { - union { - Mlp1FusedActSmem mlp1; // TileM=128 cooperative-M-split - Mlp2TFusedSmem mlp2t; // TileM=128 cooperative-M-split (mlp2_t.cuh) - Mlp5Smem mlp5; // TileM=128 cooperative-M-split - // mlp4's fused smem is traits-conditional internally (see mlp4.cuh), - // so it needs no separate 2SM specialization here. mlp3 instead uses - // a wholly distinct 2SM smem layout (Mlp3FusedSmem2Sm — paired-CTA - // accumulator pipeline in place of mlp3's normal AccumulatorPipeline), - // selected purely by Compute == 100 (Blackwell always pairs mlp3's - // Phase-2 GEMM across CTAs; Compute == 90 never does). - cute::conditional_t< - Compute == 100, - Mlp3FusedSmem2Sm, - Mlp3FusedSmem> mlp3; - Mlp4FusedSmem mlp4; - }; - - // Phase 2 per-expert K-range scratch lives in gmem (MlpBwdBufs: - // expert_k_starts / expert_k_ends). Allocated to - // exactly `experts_per_pe` ints — no compile-time cap. - - // Phase 1a, 1b', 1d: single fused pipe each (1c silu_bwd is - // elementwise gmem→gmem, no pipe). - typename Mlp1MainloopPipelineFor::SharedStorage pa_pipe; - typename Mlp2TMainloopPipelineFor::SharedStorage pb_pipe; - typename Mlp5MainloopPipelineFor::SharedStorage pd_pipe; - - // Phase 2: single fused pipe each for mlp3 and mlp4. Outputs (dA, - // dB, dC) are accumulated via TMA_REDUCE_ADD to gmem — no prev-pipe - // needed. - typename Mlp3MainloopPipelineFor::SharedStorage p3_pipe; - typename Mlp4MainloopPipelineFor::SharedStorage p4_pipe; - alignas(16) uint32_t tmem_base; - // Rendezvous barrier for the paired-CTA TMEM allocation window - // (Compute == 100 only — one CTA-pair shares one TMEM allocation, so - // both peers must arrive before either frees it). Collapses to a plain - // uint32_t placeholder (no ClusterBarrier construction/arrive/wait) on - // Compute == 90, which never pairs CTAs. - alignas(16) cute::conditional_t< - Compute == 100, - cutlass::arch::ClusterBarrier, - uint32_t> tmem_pair_barrier; - alignas(16) cute::conditional_t< - Compute == 100, - cutlass::arch::ClusterBarrier, - uint32_t> tmem_init_barrier; - alignas(16) cute::conditional_t< - Compute == 100, - cutlass::arch::ClusterBarrier, - uint32_t> acc_pair_barrier; -}; - -__device__ __forceinline__ void mlp_bwd_pair_barrier_sync( - const cutlass::arch::ClusterBarrier& barrier, - int warp_id) { - if (warp_id == 3) { - if (cute::elect_one_sync()) { - uint32_t rank = cute::block_rank_in_cluster(); - uint32_t peer = rank ^ 1; - bool leader = (rank % 2) == 0; - barrier.arrive(peer, !leader); - barrier.wait(0); - barrier.arrive(peer, leader); - } - __syncwarp(); - } -} - -template -struct MlpBwdTmemColumns { - static constexpr int kMlp1Cols = Traits1::AccStages * (2 * Traits1::TileN); - static constexpr int kMlp2TCols = Traits2T::AccStages * Traits2T::TileN; - static constexpr int kMlp5Cols = Traits5::AccStages * Traits5::TileN; - static constexpr int kMlp3Cols = Traits3::AccStages * Traits3::TileN; - static constexpr int kMlp4Cols = Traits4::AccStages * (2 * Traits4::WgTileN); - - static constexpr int kPhase15Cols = (kMlp1Cols > kMlp5Cols) - ? kMlp1Cols : kMlp5Cols; - static constexpr int kPhase2T3Cols = (kMlp2TCols > kMlp3Cols) - ? kMlp2TCols : kMlp3Cols; - static constexpr int kPhase1235Cols = (kPhase15Cols > kPhase2T3Cols) - ? kPhase15Cols : kPhase2T3Cols; - static constexpr int value = (kPhase1235Cols > kMlp4Cols) - ? kPhase1235Cols : kMlp4Cols; - static_assert(value <= 512, "BWD SM100 TMEM allocation must fit one CTA's 512 columns"); -}; - -// ═══════════════════════════════════════════════════════════════════ -// MlpBwdDims — invariant backward dimensions (__grid_constant__) -// ═══════════════════════════════════════════════════════════════════ - -struct MlpBwdDims { - int hidden_dim; - int intermediate_dim; - - // Flat 1-D launch geometry (host-computed). The grid is launched as - // dim3(num_blocks) with num_blocks = num_sms, so gridDim is no longer - // the logical grid. Phase 1 + comm partition over the GEMM-active subset: - // n_gemm = floor_NS = largest multiple of NSplit ≤ num_sms - // grid_x = n_gemm / NSplit (Phase 1 column count) - // CTAs with flat_id ≥ n_gemm fall through Phase 1/comm but still run - // Phase 2 (grid-stride over all num_blocks CTAs) and every barrier. - int grid_x; - int n_gemm; - - // Phase 1a: recompute - int total_n_rows_1; - int num_n_tiles_1; - int num_k_tiles_1; - - // Phase 1b: dZ = dY @ A - int total_k_cols_2t; - int num_n_tiles_2t; - int num_k_tiles_2t; - - // Phase 1d: dX = dU@B + dV@C - int total_k_cols_5; - int num_n_tiles_5; - int num_k_tiles_5; - - // Phase 2: mlp3 (dA) - int num_m_tiles_3; - int num_n_tiles_3; - int total_n_rows_3; - - // Phase 2: mlp4 (dB, dC) - int num_m_tiles_4; - int num_n_tiles_4; - int total_m_rows_4; - - int* phase_counter; - - // Scalar params (in __grid_constant__ to avoid register pressure) - int num_tokens; - int num_m_tiles; // Phase 1 output buffer height (in m-tiles). - // Local: total tile capacity (= max_total_slots/TileM). - // Remote: gridDim.x * CommNumStages (the staging ring). - // Phase 2 wraps bk_start modulo this when computing - // the slot index into z_buf/du_buf/dv_buf — for the - // local pass that's a no-op; for the remote pass it - // brings batches ≥ NumStages back into the live ring. - int total_m_tiles_in_pass; // Phase 2 K-block iteration bound (in m-tiles). - // Local: equals num_m_tiles (existing semantics — Phase 2 - // walks the full padded buffer and skips fake K-blocks - // via the expert_for_k_block sentinel). - // Remote: equals smem.comm.global_total — strictly less - // than num_m_tiles when per_cta_tiles > NumStages, so - // batches that would walk past the ring still run Phase 2 - // (with bk_start wrapped) instead of being silently - // clamped to bk_tiles=0. - int num_batches; - int local_expert_start; - int num_pes; - int experts_per_pe; // Phase 3/4 use this to skip k-blocks whose expert - // falls outside [local_expert_start, - // local_expert_start + experts_per_pe). Required - // because expert_for_k_block (built post-sort) holds - // GLOBAL expert IDs spanning all PEs, but dA/dB/dC - // on this PE only have slots for the local subset. - - // Phase 2 virtual-grid block count along the non-NSplit2 axis: - // phase2_num_blocks = (gridDim.x * gridDim.y) / NSplit2. - // Host-known (depends only on launch grid + template constant NSplit2), - // so storing here lets every batch's mlp3/mlp4 call read it from grid - // constant memory instead of recomputing `gridDim.x * gridDim.y / - // NSplit2` per batch. - int phase2_num_blocks; - - // Remote-pass overrides for num_tokens / num_m_tiles. Both are host- - // known (gridDim.x * CommNumStages and that × Traits1::TileM); kept here - // so moe_fused_bwd's remote_dims build can just assign from grid constant - // instead of computing on device. The remote-only fields stay = 0 for - // the local pass (the local pass doesn't read them). - int remote_num_m_tiles; - int remote_num_tokens; - - // 1 when the bwd TMA-GET comm path is active (NVLink P2P + divisible - // hidden_dim + bounce smem fits). Selects the comm warp layout AND the - // remote MLP iterator's producer/consumer counts (1·NC get / 2·NC put - // vs 2·NC get / 1·NC put). Host-set; 0 keeps the original getmem layout. - int tma_get_enabled; - - // (efkb_stride_3 / efkb_stride_4 moved to template parameters on - // mlp_fused_bwd — see EfkbStride3/EfkbStride4 below. They were always - // compile-time constants per call site (local: 1, remote: - // Traits1::TileM/Traits{3,4}::TileK), so passing them as runtime ints - // just kept Traits1::TileM and the divides live in registers across the - // whole batch loop. Hoisting to templates lets the compiler constant- - // fold the `/1` (local) into a no-op and the remote ratios into a shift. -}; - -// ═══════════════════════════════════════════════════════════════════ -// MlpBwdBufs — buffer pointers (__grid_constant__) -// ═══════════════════════════════════════════════════════════════════ - -template -struct MlpBwdBufs { - // du_buf holds U' (silu-bwd coefficient) after mlp1, then dU after silu_bwd. - // dv_buf holds V' (silu-bwd coefficient) after mlp1, then dV after silu_bwd. - // dz_buf holds dZ between mlp2_t (Phase 1b') and silu_bwd (Phase 1c). - Element* du_buf; - Element* dv_buf; - Element* dz_buf; - - // Expert id per K-block, read by Phase 2's mlp4 and mlp3 respectively. - // They hold the same logical content — for K-block k, the local expert - // id of the tile that K-block belongs to — but the source array differs - // because each Phase-2 sub-pass has a different freshness requirement. - // - // expert_for_k_block_mlp4 is read by mlp4 (which reads X). In the - // remote pass it points at - // CommBuffersBwd::tile_expert_ids_x — comm - // won't overwrite it until MLP's release_src, - // which is deferred past mlp4. - // - // expert_for_k_block_mlp3 is read by mlp3 (which reads dY). In the - // remote pass it points at - // CommBuffersBwd::tile_expert_ids_dy — comm - // won't overwrite it until MLP's release_dy, - // which is deferred past mlp3. - // - // In the local pass and the standalone bwd kernel, both pointers alias - // the same post-sort expert_for_k_block built by build_expert_for_k_block - // — there is no comm race to defend against, and a single source is - // sufficient. - const int* expert_for_k_block_mlp4; - const int* expert_for_k_block_mlp3; - const int* valid_rows_mlp4; - - // Per-pass per-expert K-range scratch for Phase 2. Filled once at - // the start of each pass (in the bwd preamble) and consumed by both - // mlp3 and mlp4 batches — they require Traits3::TileK == Traits4:: - // TileK (the codebase already enforces this since a single - // expert_for_k_block indexes both), so one shared scan suffices. - // Sized exactly `experts_per_pe` ints each. - int* expert_k_starts; - int* expert_k_ends; - - Element const* gmem_da; - Element const* gmem_db; - Element const* gmem_dc; - int* barrier_counter; -}; - -// ═══════════════════════════════════════════════════════════════════ -// MLP-warps-only barriers (BWD) -// ═══════════════════════════════════════════════════════════════════ -// -// Both barriers use the target-based CtaCounterBarrier (see -// cta_barrier.cuh). The earlier modulo variants were susceptible to a -// CTA-skew deadlock where a fast CTA could push the counter past the -// `% stride == 0` window before slow CTAs polled. -// -// • MlpBwdGlobalBarrier: cross-CTA, stride = gridDim.x * gridDim.y. -// Counter is a single int in device memory. -// • MlpBwdXBarrier: per-blockIdx.x, stride = gridDim.y. -// Counter is a gmem array of length gridDim.x. -template -using MlpBwdCtaBarrierT = - CtaCounterBarrier, kMlpBarrierId>; - -// ═══════════════════════════════════════════════════════════════════ -// Per-phase inline runners -// ═══════════════════════════════════════════════════════════════════ -// -// Each phase of the bwd batch loop (1a / 1b' / 1d / 2-mlp4 / 2-mlp3) is -// pulled out into its own __forceinline__ device function. They take the -// per-phase PipelineState by REFERENCE — the caller (`mlp_fused_bwd`) -// holds the persistent state at function scope, and the inline body sees -// it as a direct register reference once inlined. Using actual functions -// instead of capture-all-by-reference lambdas avoids forcing the captured -// state into stack/local memory (CUDA closures pin captured vars to -// addressable storage, even when __forceinline__ would otherwise let -// them stay in registers). -// -// State persistence: the in-smem pipeline barriers carry a phase counter -// that advances each iteration and persists across phase invocations. To -// keep correctness with construction-inside-the-phase, the outer scope -// persists only a monotonic `uint32_t count` per phase (1 reg vs the -// 3-reg full PipelineState). The phase function: -// 1. Constructs a fresh PipelineState locally (warp-role-aware initial -// phase: producer starts phase=1, consumer starts phase=0). -// 2. Calls `advance(count)` to align it with the actual smem-barrier -// phase as left by the previous invocation. -// 3. Runs the phase, naturally advancing local state. -// 4. Writes `count = state.count()` back through the reference. -// This shrinks the cross-phase live set from 5×3 = 15 regs to 5×1 = 5 regs -// without re-introducing the desync bug we hit when state was lambda-local -// (T4096_D4096_I14336_E16_K4: dA off by 10% on expert 7 when num_iter % -// (2*Stages) != 0). - -template -__device__ __forceinline__ cutlass::PipelineState -mlp_bwd_resume_state(uint32_t count, bool is_producer) { - auto s = is_producer - ? cutlass::make_producer_start_state() - : cutlass::PipelineState{}; - s.advance(count); - return s; -} - -template -__device__ __forceinline__ void mlp_bwd_run_phase_1a( - Mlp1MainloopPipelineFor& pa_pipe, - uint32_t& pa_count, - Mlp1FusedActSmem& smem_mlp1, - TmaLoadX const& tma_load_x, - TmaLoadW1 const& tma_load_b_fwd, - TmaLoadW1 const& tma_load_c_fwd, - TmaStoreU const& tma_store_du, - TmaStoreV const& tma_store_dv, - TmaStoreZ const& tma_store_z, - int warp_id, int m, int expert, - const MlpBwdDims& dims, - uint32_t tmem_base, - int split_idx = -1, int num_splits = -1) { - int ns = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int nc = (num_splits >= 0) ? num_splits : (int)gridDim.y; - auto s = mlp_bwd_resume_state>( - pa_count, warp_id == 0); - if (warp_id == 0) { - mlp1_fused_act_producer(pa_pipe, s, - smem_mlp1, tma_load_x, tma_load_b_fwd, tma_load_c_fwd, - m, (Compute == 90) ? expert : expert * dims.num_n_tiles_1, - dims.num_tokens, dims.hidden_dim, dims.intermediate_dim, - dims.experts_per_pe, dims.total_n_rows_1, - dims.num_n_tiles_1, dims.num_k_tiles_1, - ns, nc); - if constexpr (Compute == 100) - pa_pipe.producer_tail(s); - } - if constexpr (Compute == 100) - if (warp_id == 3 && threadIdx.x % Traits1::WarpSize == 0) - smem_mlp1.tmem_base = tmem_base; - constexpr int kFirstMlp12ConsumerWarp = (Compute == 100) ? 3 : 4; - if (warp_id >= kFirstMlp12ConsumerWarp) - mlp1_fused_act_consumer(pa_pipe, s, - smem_mlp1, tma_store_du, tma_store_dv, tma_store_z, - m, dims.intermediate_dim, dims.num_m_tiles, - dims.num_n_tiles_1, dims.num_k_tiles_1, - ns, nc); - pa_count = s.count(); -} - -template -__device__ __forceinline__ void mlp_bwd_run_phase_1b( - Mlp2TMainloopPipelineFor& pb_pipe, - uint32_t& pb_count, - Mlp2TFusedSmem& smem_mlp2t, - TmaLoadDY const& tma_load_dy, - TmaLoadW2T const& tma_load_a_col, - TmaStoreDZ const& tma_store_dz, - int warp_id, int m, int expert, - const MlpBwdDims& dims, - uint32_t tmem_base, - int split_idx = -1, int num_splits = -1) { - int ns = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int nc = (num_splits >= 0) ? num_splits : (int)gridDim.y; - auto s = mlp_bwd_resume_state>( - pb_count, warp_id == 0); - if (warp_id == 0) { - mlp2_t_fused_producer(pb_pipe, s, - smem_mlp2t, tma_load_dy, tma_load_a_col, - m, expert, - dims.num_tokens, dims.hidden_dim, dims.intermediate_dim, - dims.experts_per_pe, dims.total_k_cols_2t, - dims.num_n_tiles_2t, dims.num_k_tiles_2t, - ns, nc); - if constexpr (Compute == 100) - pb_pipe.producer_tail(s); - } - if constexpr (Compute == 100) - if (warp_id == 3 && threadIdx.x % Traits2T::WarpSize == 0) - smem_mlp2t.tmem_base = tmem_base; - constexpr int kFirstMlp12ConsumerWarp = (Compute == 100) ? 3 : 4; - if (warp_id >= kFirstMlp12ConsumerWarp) - mlp2_t_fused_consumer(pb_pipe, s, - smem_mlp2t, tma_store_dz, - m, dims.intermediate_dim, dims.num_m_tiles, - dims.num_n_tiles_2t, dims.num_k_tiles_2t, - ns, nc); - pb_count = s.count(); -} - -template -__device__ __forceinline__ void mlp_bwd_run_phase_1d( - Mlp5MainloopPipelineFor& pd_pipe, - uint32_t& pd_count, - Mlp5Smem& smem_mlp5, - TmaLoadDU const& tma_load_du, - TmaLoadDV const& tma_load_dv, - TmaLoadB const& tma_load_b_col, - TmaLoadC const& tma_load_c_col, - TmaStoreDX const& tma_store_dx, - int warp_id, int m, int expert, - const MlpBwdDims& dims, - uint32_t tmem_base, - int split_idx = -1, int num_splits = -1) { - int ns = (split_idx >= 0) ? split_idx : (int)blockIdx.y; - int nc = (num_splits >= 0) ? num_splits : (int)gridDim.y; - auto s = mlp_bwd_resume_state>( - pd_count, warp_id == 0); - if (warp_id == 0) { - mlp5_fused_producer(pd_pipe, s, - smem_mlp5, - tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, - m, (Compute == 90) ? expert : expert * dims.num_k_tiles_5, - dims.num_tokens, dims.hidden_dim, dims.intermediate_dim, - dims.experts_per_pe, dims.total_k_cols_5, - dims.num_n_tiles_5, dims.num_k_tiles_5, - ns, nc); - if constexpr (Compute == 100) - pd_pipe.producer_tail(s); - } - if constexpr (Compute == 100) - if (warp_id == 3 && threadIdx.x % Traits5::WarpSize == 0) - smem_mlp5.tmem_base = tmem_base; - constexpr int kFirstMlp5ConsumerWarp = (Compute == 100) ? 3 : 4; - if (warp_id >= kFirstMlp5ConsumerWarp) - mlp5_fused_consumer(pd_pipe, s, - smem_mlp5, tma_store_dx, - m, dims.hidden_dim, - dims.num_m_tiles, dims.num_n_tiles_5, dims.num_k_tiles_5, - ns, nc); - pd_count = s.count(); -} - -// ═══════════════════════════════════════════════════════════════════ -// Phase 2 per-batch build-ranges scan -// ═══════════════════════════════════════════════════════════════════ -// -// CTA 0 warp 0 fills (k_start, k_end) per local expert for the current batch. -// Empty experts stay at (INT_MAX, 0); the chunk loop skips them via -// `k_end <= k_start`. The caller follows this with a grid-wide barrier before -// any CTA consumes the ranges. Restricting reset+scan to one warp avoids a -// race where a late CTA reset erased ranges another CTA had already updated. -// -// Called once per (batch, phase). mlp3 and mlp4 use different -// expert_for_k_block arrays and strides, so each phase scans its own. - -// ring_kb (K-blocks in the staging ring): when > 0 the window [k_offset, -// k_offset+num_k_in_batch) is in UNWRAPPED coordinates, and the physical ring -// slot for the expert-id lookup is (k_offset+kb) % ring_kb — while the recorded -// k_start/k_end ranges stay UNWRAPPED so each expert's range is contiguous even -// when the grouped window straddles the ring wrap. ring_kb == 0 (default) means -// no wrap (the un-grouped / local-pass callers, whose window already fits). -template -__device__ __forceinline__ void mlp_bwd_phase_2_build_ranges( - int* k_starts, - int* k_ends, - const int* expert_for_k_block, - const int* valid_rows_for_tile, - int k_offset, - int num_k_in_batch, - int experts_per_pe, - int local_expert_start, - int efkb_stride, - int tile_k, - int warp_id, - int ring_kb = 0) { - constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; - if (blockIdx.x == 0 && blockIdx.y == 0 && warp_id == 0) { - int lane = threadIdx.x; - for (int e = lane; e < experts_per_pe; e += 32) { - k_starts[e] = INT_MAX; - k_ends[e] = 0; - } - __syncwarp(); - for (int kb = lane; kb < num_k_in_batch; kb += 32) { - int kb_abs = k_offset + kb; // unwrapped - int kb_phys = (ring_kb > 0) ? (kb_abs % ring_kb) : kb_abs; // ring slot - int tile = kb_phys / efkb_stride; - if (valid_rows_for_tile != nullptr) { - int row_in_tile = (kb_phys - tile * efkb_stride) * tile_k; - if (row_in_tile >= valid_rows_for_tile[tile]) - continue; - } - int eg = expert_for_k_block[tile]; - int el = eg - local_expert_start; - if (el >= 0 && el < experts_per_pe) { - atomicMin(&k_starts[el], kb_abs); - atomicMax(&k_ends[el], kb_abs + 1); - } - } - } - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); -} - -// ── mlp4 Phase-2 runner ── -// -// Unified over Compute: Compute == 90 (Hopper) drives the adaptive kMSplit -// ternary walk over plain per-CTA cells; Compute == 100 (Blackwell) always -// drives the SM100 ClusterM=2 paired-CTA path (fixed M-split — the paired -// UMMA atom only supports the shared-N/split-M cooperative layout — and -// cell = CTA-PAIR, not CTA). Calls the unified mlp4_producer / mlp4_consumer (mlp4.cuh), which do their own -// internal Compute == 100 dispatch to the paired-CTA multicast TMA path. -// Reuses the plain Mlp4FusedSmem smem type (see MlpFusedBwdSmem — -// mlp4.cuh's Mlp4FusedSmem is traits-compatible internally, no separate -// Compute-conditional smem needed here). -template -__device__ __forceinline__ void mlp_bwd_run_phase_2_mlp4( - Mlp4MainloopPipelineFor& p4_pipe, - uint32_t& p4_count, - Mlp4FusedSmem& smem_mlp4, - const int* k_starts, - const int* k_ends, - TmaLoadXT4 const& tma_load_xt4, - TmaLoaddUT4 const& tma_load_dut4, - TmaLoaddVT4 const& tma_load_dvt4, - TmaReduceDB const& tma_reduce_db, - TmaReduceDC const& tma_reduce_dc, - int warp_id, int bk_start, int bk_tiles, - int phase2_div, int phase2_mod, - const MlpBwdDims& dims, - uint32_t tmem_base, - int ring_kb = 0, - const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, - uint32_t* pair_init_phase = nullptr) { - static_assert(Compute == 90 || Compute == 100, - "mlp_bwd_run_phase_2_mlp4: Compute must be 90 (Hopper 1SM) or " - "100 (Blackwell ClusterM=2 paired-CTA) — only Compute == 100 " - "executes the paired-CTA (ClusterM/NumPairs) code below."); - constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; - auto s = mlp_bwd_resume_state>( - p4_count, warp_id == 0); - - int experts_per_pe = dims.experts_per_pe; - - // Batch's K-window in mlp4's TileK_4 grain. Producer/consumer - // intersect this with the global per-expert ranges in `k_starts`/ - // `k_ends` (filled once per pass in the bwd preamble). - int batch_kb_start = (bk_start * Traits1::TileM) / Traits4::TileK; - int batch_kb_end = batch_kb_start - + (bk_tiles * Traits1::TileM) / Traits4::TileK; - - // Flat cell index over all Phase 2 CTAs. cell_start = phase2_div·NSplit2 - // + phase2_mod = flat_id ∈ [0,num_blocks) (the div/mod from the caller - // recombine to the flat id), so coverage is complete and disjoint for - // ANY num_blocks — no NSplit2 divisibility. - int flat_id = phase2_div * NSplit2 + phase2_mod; - // Phase 2 uses ALL launched CTAs (flat grid-stride over cells), so the - // cell stride is the full launched count gridDim.x·gridDim.y = num_blocks - // on Compute == 90; Compute == 100 pairs CTAs (ClusterM=2), so the stride - // is in CTA-PAIRS instead. - int total_phase2_ctas = (int)gridDim.x * (int)gridDim.y; - // Size outer_split from the experts LIVE in this batch's K-window, NOT the - // global expert set. Sizing from experts_per_pe under-splits when only ~1 - // expert is live per batch (large T) — the dead experts inflate the sizing - // input, so outer_split stays small and the lone live expert's work spreads - // over too few CTAs. k_starts/k_ends are per-expert ranges in TileK grain. - int live_experts_4 = 0; - for (int e = 0; e < experts_per_pe; ++e) - if (k_starts[e] < batch_kb_end && k_ends[e] > batch_kb_start) live_experts_4++; - if (live_experts_4 < 1) live_experts_4 = 1; - - int cell_start, cell_stride, outer_split_4, k_split_4; - if constexpr (Compute == 100) { - // SM100 ClusterM=2 paired-CTA path: fixed M-split (no kMSplit - // ternary — the paired-CTA MMA atom only supports the shared-N/ - // split-M cooperative layout), cell = CTA-PAIR. - int num_clusters = total_phase2_ctas / Traits4::ClusterM; - int total_chunks = live_experts_4 * dims.num_n_tiles_4; - int walk_extent = dims.num_m_tiles_4; - outer_split_4 = 1; - { - int target_cells = 2 * num_clusters; - while (outer_split_4 * 2 * Traits4::NumPairs <= walk_extent - && total_chunks * outer_split_4 < target_cells) - outer_split_4 *= 2; - while (walk_extent % outer_split_4 != 0) - outer_split_4 /= 2; - if (outer_split_4 < 1) outer_split_4 = 1; - } - // K-split (KS): split each cell's K-loop across k_split_4 CTA-pairs - // when the live cells don't fill the grid. TMA_REDUCE_ADD sums - // partials. - k_split_4 = 1; - { - int active_cells = total_chunks * outer_split_4; - int batch_k_tiles = batch_kb_end - batch_kb_start; - constexpr int kKMin = 8; - int max_by_depth = batch_k_tiles / kKMin; - if (max_by_depth < 1) max_by_depth = 1; - if (active_cells > 0 && active_cells < num_clusters) { - k_split_4 = (num_clusters + active_cells - 1) / active_cells; - if (k_split_4 > max_by_depth) k_split_4 = max_by_depth; - if (k_split_4 < 1) k_split_4 = 1; - } - } - cell_start = flat_id / Traits4::ClusterM; - cell_stride = num_clusters; - } else { - // SM90 (Hopper) 1SM path. mlp4 holds the SHARED operand and walks - // the cooperation (split) axis (see mlp4.cuh): - // kMSplit=true (M-split): chunk = (e, n_tile), lane partitions num_m_tiles_4 - // kMSplit=false (N-split): chunk = (e, m_tile), lane partitions num_n_tiles_4 - // Adaptive split: largest pow2 divisor of walk_extent so total_cells - // covers all Phase 2 CTAs. - int total_chunks_4 = Traits4::kMSplit - ? (live_experts_4 * dims.num_n_tiles_4) - : (live_experts_4 * dims.num_m_tiles_4); - int walk_extent_4 = Traits4::kMSplit ? dims.num_m_tiles_4 : dims.num_n_tiles_4; - outer_split_4 = 1; - { - int target_cells = 2 * total_phase2_ctas; - while (outer_split_4 * 2 <= walk_extent_4 - && total_chunks_4 * outer_split_4 < target_cells) { - outer_split_4 *= 2; - } - while (walk_extent_4 % outer_split_4 != 0) outer_split_4 /= 2; - if (outer_split_4 < 1) outer_split_4 = 1; - } - // K-split (KS): split each cell's K-loop across k_split_4 CTAs when - // the live cells don't fill the grid. TMA_REDUCE_ADD sums partials. - k_split_4 = 1; - { - int active_cells = total_chunks_4 * outer_split_4; - int batch_k_tiles = batch_kb_end - batch_kb_start; - constexpr int kKMin = 8; - int max_by_depth = batch_k_tiles / kKMin; - if (max_by_depth < 1) max_by_depth = 1; - if (active_cells > 0 && active_cells < total_phase2_ctas) { - k_split_4 = (total_phase2_ctas + active_cells - 1) / active_cells; - if (k_split_4 > max_by_depth) k_split_4 = max_by_depth; - if (k_split_4 < 1) k_split_4 = 1; - } - } - cell_start = flat_id; - cell_stride = total_phase2_ctas; - } - - if (warp_id == 0) - mlp4_producer( - p4_pipe, s, - smem_mlp4, tma_load_xt4, tma_load_dut4, tma_load_dvt4, - k_starts, k_ends, experts_per_pe, - dims.intermediate_dim, dims.hidden_dim, dims.num_tokens, - dims.num_m_tiles_4, dims.num_n_tiles_4, outer_split_4, - cell_start, cell_stride, - batch_kb_start, batch_kb_end, k_split_4, ring_kb); - if constexpr (Compute == 100) - if (warp_id == 3 && threadIdx.x % Traits4::WarpSize == 0) - smem_mlp4.tmem_base = tmem_base; - constexpr int kFirstMlp4ConsumerWarp = (Compute == 100) ? 3 : 4; - if (warp_id >= kFirstMlp4ConsumerWarp) - mlp4_consumer( - p4_pipe, s, - smem_mlp4, tma_reduce_db, tma_reduce_dc, - k_starts, k_ends, experts_per_pe, - dims.intermediate_dim, dims.hidden_dim, dims.total_m_rows_4, - dims.num_m_tiles_4, dims.num_n_tiles_4, outer_split_4, - cell_start, cell_stride, - batch_kb_start, batch_kb_end, k_split_4, - pair_init_barrier, pair_init_phase); - - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - p4_count = s.count(); -} - -// ── mlp3 Phase-2 runner ── -// -// Unified over Compute: Compute == 90 (Hopper) drives the adaptive kMSplit -// ternary walk over plain per-CTA cells; Compute == 100 (Blackwell) always -// drives the SM100 ClusterM=2 paired-CTA path (fixed M-split — the paired -// UMMA atom only supports the shared-N/split-M cooperative layout — and -// cell = CTA-PAIR, not CTA), and uses the distinct Mlp3FusedSmem2Sm smem -// layout (paired-CTA accumulator pipeline; see MlpFusedBwdSmem). Calls the -// unified mlp3_producer / mlp3_consumer -// (mlp3.cuh), which do their own internal Compute == 100 dispatch to the -// paired-CTA multicast TMA path. -template -__device__ __forceinline__ void mlp_bwd_run_phase_2_mlp3( - Mlp3MainloopPipelineFor& p3_pipe, - uint32_t& p3_count, - cute::conditional_t, Mlp3FusedSmem>& smem_mlp3, - const int* k_starts, - const int* k_ends, - TmaLoadDYT3 const& tma_load_dyt3, - TmaLoadZT3 const& tma_load_zt3, - TmaReduceDA const& tma_reduce_da, - int warp_id, int bk_start, int bk_tiles, - int phase2_div, int phase2_mod, - const MlpBwdDims& dims, - uint32_t tmem_base, - int ring_kb = 0, - const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, - uint32_t* pair_init_phase = nullptr) { - static_assert(Compute == 90 || Compute == 100, - "mlp_bwd_run_phase_2_mlp3: Compute must be 90 (Hopper 1SM) or " - "100 (Blackwell ClusterM=2 paired-CTA) — only Compute == 100 " - "executes the paired-CTA (ClusterM/NumPairs) code below."); - constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; - auto s = mlp_bwd_resume_state>( - p3_count, warp_id == 0); - - int experts_per_pe = dims.experts_per_pe; - - // Batch's K-window in mlp3's TileK_3 grain (see mlp4 helper above). - int batch_kb_start = (bk_start * Traits1::TileM) / Traits3::TileK; - int batch_kb_end = batch_kb_start - + (bk_tiles * Traits1::TileM) / Traits3::TileK; - - // Flat cell index over all Phase 2 CTAs. cell_start = phase2_div·NSplit2 - // + phase2_mod = flat_id ∈ [0,num_blocks) (the div/mod from the caller - // recombine to the flat id), so coverage is complete and disjoint for - // ANY num_blocks — no NSplit2 divisibility. - int flat_id = phase2_div * NSplit2 + phase2_mod; - // All launched CTAs (gridDim.x·gridDim.y = num_blocks) — see mlp4 helper. - int total_phase2_ctas = (int)gridDim.x * (int)gridDim.y; - // Size outer_split from the experts LIVE in this batch's K-window (see the - // mlp4 helper above for the rationale) — not the global experts_per_pe. - int live_experts_3 = 0; - for (int e = 0; e < experts_per_pe; ++e) - if (k_starts[e] < batch_kb_end && k_ends[e] > batch_kb_start) live_experts_3++; - if (live_experts_3 < 1) live_experts_3 = 1; - - int cell_start, cell_stride, outer_split_3, k_split_3; - if constexpr (Compute == 100) { - // SM100 ClusterM=2 paired-CTA path: fixed M-split, cell = CTA-PAIR. - int total_phase2_pairs = total_phase2_ctas / Traits3::ClusterM; - int total_chunks = live_experts_3 * dims.num_n_tiles_3; - outer_split_3 = 1; - { - int target_cells = 2 * total_phase2_pairs; - while (outer_split_3 * 2 * Traits3::NumPairs <= dims.num_m_tiles_3 - && total_chunks * outer_split_3 < target_cells) - outer_split_3 *= 2; - while (dims.num_m_tiles_3 % (outer_split_3 * Traits3::NumPairs) != 0) - outer_split_3 /= 2; - outer_split_3 = max(outer_split_3, 1); - } - // K-split (KS): split each cell's K-loop across k_split_3 CTA-pairs - // when the live cells don't fill the grid. TMA_REDUCE_ADD sums - // partials. - k_split_3 = 1; - { - int active_cells = total_chunks * outer_split_3; - int batch_k_tiles = batch_kb_end - batch_kb_start; - constexpr int kKMin = 8; - int max_by_depth = max(batch_k_tiles / kKMin, 1); - if (active_cells > 0 && active_cells < total_phase2_pairs) { - k_split_3 = (total_phase2_pairs + active_cells - 1) / active_cells; - k_split_3 = min(k_split_3, max_by_depth); - } - } - cell_start = flat_id / Traits3::ClusterM; - cell_stride = total_phase2_pairs; - } else { - // SM90 (Hopper) 1SM path. mlp3 holds the SHARED operand and walks - // the cooperation (split) axis: - // kMSplit=true : chunk = (e, n_tile), lane partitions num_m_tiles - // kMSplit=false : chunk = (e, m_tile), lane partitions num_n_tiles - // Adaptive split: pick largest pow2 divisor of walk_extent so - // total_cells covers all Phase 2 CTAs. - int total_chunks_3 = Traits3::kMSplit - ? (live_experts_3 * dims.num_n_tiles_3) - : (live_experts_3 * dims.num_m_tiles_3); - int walk_extent_3 = Traits3::kMSplit ? dims.num_m_tiles_3 : dims.num_n_tiles_3; - outer_split_3 = 1; - { - int target_cells = 2 * total_phase2_ctas; - while (outer_split_3 * 2 <= walk_extent_3 - && total_chunks_3 * outer_split_3 < target_cells) { - outer_split_3 *= 2; - } - while (walk_extent_3 % outer_split_3 != 0) outer_split_3 /= 2; - if (outer_split_3 < 1) outer_split_3 = 1; - } - // K-split (KS): when the live cells (total_chunks_3 × outer_split_3) - // don't fill the grid, split each cell's K-loop across k_split_3 - // CTAs and let the TMA_REDUCE_ADD epilogue sum the partials. - // Bounded so each slice keeps >= kKMin K-tiles (else the Stages-deep - // pipeline can't amortize its fill). - k_split_3 = 1; - { - int active_cells = total_chunks_3 * outer_split_3; - int batch_k_tiles = batch_kb_end - batch_kb_start; - constexpr int kKMin = 8; - int max_by_depth = batch_k_tiles / kKMin; - if (max_by_depth < 1) max_by_depth = 1; - if (active_cells > 0 && active_cells < total_phase2_ctas) { - k_split_3 = (total_phase2_ctas + active_cells - 1) / active_cells; - if (k_split_3 > max_by_depth) k_split_3 = max_by_depth; - if (k_split_3 < 1) k_split_3 = 1; - } - } - cell_start = flat_id; - cell_stride = total_phase2_ctas; - } - - if (warp_id == 0) - mlp3_producer( - p3_pipe, s, - smem_mlp3, tma_load_dyt3, tma_load_zt3, - k_starts, k_ends, experts_per_pe, - dims.hidden_dim, dims.intermediate_dim, dims.num_tokens, - dims.num_m_tiles_3, dims.num_n_tiles_3, outer_split_3, - cell_start, cell_stride, - batch_kb_start, batch_kb_end, k_split_3, ring_kb); - if constexpr (Compute == 100) - if (warp_id == 3 && threadIdx.x % Traits3::WarpSize == 0) - smem_mlp3.tmem_base = tmem_base; - constexpr int kFirstMlp3ConsumerWarp = (Compute == 100) ? 3 : 4; - if (warp_id >= kFirstMlp3ConsumerWarp) - mlp3_consumer( - p3_pipe, s, - smem_mlp3, tma_reduce_da, - k_starts, k_ends, experts_per_pe, - dims.hidden_dim, dims.intermediate_dim, dims.total_n_rows_3, - dims.num_m_tiles_3, dims.num_n_tiles_3, outer_split_3, - cell_start, cell_stride, - batch_kb_start, batch_kb_end, k_split_3, - pair_init_barrier, pair_init_phase); - - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - p3_count = s.count(); -} - -// ═══════════════════════════════════════════════════════════════════ -// mlp_fused_bwd_run_batches — batch loop body extracted -// ═══════════════════════════════════════════════════════════════════ -// -// Inner per-pass batch loop. Caller (mlp_fused_bwd or mlp_fused_bwd_dual) -// owns pipeline construction + the persistent per-phase counts. This lets -// the dual entry share ONE set of pipelines / counts across the local -// and remote passes — they sit at the dual-function's stack scope and -// PTXAS can analyze both batch loops as sequential code in one CFG, -// reusing spill slots between them. -// -// Templated on TileIter / TMA types / EfkbStride{3,4} — gets inlined into -// the caller's body. When called from mlp_fused_bwd_dual twice, the two -// inlined copies appear as two adjacent code regions in the same outer -// function (vs two separate inlined instances of the entire mlp_fused_bwd, -// which is the current 2-call shape). - -template -__device__ __forceinline__ void mlp_fused_bwd_run_batches( - MlpFusedBwdSmem& smem, - Mlp1MainloopPipelineFor& pa_pipe, - Mlp2TMainloopPipelineFor& pb_pipe, - Mlp5MainloopPipelineFor& pd_pipe, - Mlp3MainloopPipelineFor& p3_pipe, - Mlp4MainloopPipelineFor& p4_pipe, - uint32_t& pa_count, uint32_t& pb_count, uint32_t& pd_count, - uint32_t& p3_count, uint32_t& p4_count, - TileIter& iter, - TmaLoadX const& tma_load_x, TmaLoadW1 const& tma_load_b_fwd, - TmaLoadW1 const& tma_load_c_fwd, TmaStoreZ const& tma_store_z, - TmaStoreDU const& tma_store_du, TmaStoreDV const& tma_store_dv, - TmaLoadDY const& tma_load_dy, TmaLoadW2T const& tma_load_a_col, - TmaStoreDZ const& tma_store_dz, - TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, - TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, - TmaStoreDX const& tma_store_dx, - TmaLoadDYT3 const& tma_load_dyt3, TmaLoadZT3 const& tma_load_zt3, - TmaReduceDA const& tma_reduce_da, - TmaLoadXT4 const& tma_load_xt4, - TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, - TmaReduceDB const& tma_reduce_db, TmaReduceDC const& tma_reduce_dc, - int efkb_stride_3, int efkb_stride_4, - const MlpBwdDims& dims, - const MlpBwdBufs& bufs, - MlpBwdCtaBarrierT& global_barrier, - MlpBwdCtaBarrierT& x_barrier, - int warp_id) { - - using Element = typename Traits1::Element; - constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; - - // Phase 2 per-expert K-range scan is done PER BATCH over the current - // batch's K-window (below), NOT once per pass. mlp3 and mlp4 share TileK, - // so a single per-batch scan (over the X-pipe expert-id array) feeds both. - // - // A once-per-pass global scan over [0, num_tokens/TileK) is only correct - // when each physical slot maps to a SINGLE expert for the whole pass — - // true for the local pass (fully-materialized buffer) and for a remote - // ring that never wraps (global_total ≤ L). When the remote staging ring - // wraps to a 2nd+ lap (global_total > L = MC·CommNumStages), a physical - // slot is REUSED for a different expert each lap, so a global min/max - // range mis-attributes the reused slots — corrupting the experts whose - // tiles straddle the lap-reuse boundary (the ring's first/last experts). - // Scanning the LIVE tile_expert_ids per batch (which holds this batch's - // residency, sync-guaranteed by Phase 1's acquire) attributes every slot - // correctly regardless of lap. This restores the pre-#102 per-batch scan - // semantics under #102's cooperative cell-walk mlp3/mlp4. - static_assert(Traits3::TileK == Traits4::TileK, - "Phase 2 shared k_starts/ends require Traits3::TileK == Traits4::TileK"); - // The Compute == 100 ClusterM=2 paired-CTA Phase-2 path issues tcgen05 - // 2x1SM UMMA and cross-CTA cluster-scope mbarriers that only exist on - // SM100. Compute is the SOLE selector (see mlp_bwd_run_phase_2_mlp3 / - // mlp_bwd_run_phase_2_mlp4 above) — catch a bogus Compute value here at - // compile time rather than as a cryptic pipeline/SharedStorage type - // mismatch further down. - static_assert(Compute == 90 || Compute == 100, - "mlp_fused_bwd_run_batches: Compute must be 90 (SM90/Hopper, 1SM " - "Phase-2) or 100 (SM100/Blackwell, ClusterM=2 paired-CTA Phase-2)."); - - for (int batch = 0; batch < dims.num_batches; ++batch) { - - int saved_m = -1, saved_expert = -1; - bool has_tile = false; - - // Phase 1: recompute + silu coefs, dZ via dY @ A, silu_bwd - if (iter.has_next()) { - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - - iter.acquire_src(); - auto tile = iter.next(); - int expert = tile.expert - dims.local_expert_start; - int m = tile.y_m; - - saved_m = m; - saved_expert = expert; - has_tile = true; - - mlp_bwd_run_phase_1a( - pa_pipe, pa_count, smem.mlp1, - tma_load_x, tma_load_b_fwd, tma_load_c_fwd, - tma_store_du, tma_store_dv, tma_store_z, - warp_id, m, expert, dims, smem.tmem_base); - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - - iter.acquire_dy(); - mlp_bwd_run_phase_1b( - pb_pipe, pb_count, smem.mlp2t, - tma_load_dy, tma_load_a_col, tma_store_dz, - warp_id, m, expert, dims, smem.tmem_base); - x_barrier.wait(); - x_barrier.wait(); - - silu_bwd_pair_tile( - bufs.dz_buf, bufs.du_buf, bufs.dv_buf, - bufs.du_buf, bufs.dv_buf, - m, Traits1::TileM, Traits2T::TileN, dims.intermediate_dim, - blockIdx.y, gridDim.y); - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - } - - // Global barrier: Z, dU, dV all visible - global_barrier.wait(); - - // Phase 1d: dX = dU@B + dV@C - if (has_tile) { - iter.acquire_dst(); - mlp_bwd_run_phase_1d( - pd_pipe, pd_count, smem.mlp5, - tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, - tma_store_dx, - warp_id, saved_m, saved_expert, dims, smem.tmem_base); - __threadfence(); // system-scope dX flush for the NIC put is done by release_dst - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - iter.release_dst(threadIdx.x); - } - - // Phase 2: weight grads — mlp4 first, then mlp3 - int phase2_flat_cta = (int)blockIdx.x + (int)blockIdx.y * (int)gridDim.x; - int phase2_div = phase2_flat_cta / NSplit2; - int phase2_mod = phase2_flat_cta % NSplit2; - - // Per-batch K-window = the tiles Phase 1 delivered THIS batch. Phase 1 - // advances each of grid_x columns by one tile per batch, so batch b - // delivers global tiles [b·grid_x, (b+1)·grid_x). Phase 2 must reduce - // exactly those — striding by grid_x (= n_gemm/NSplit), NOT gridDim.x. - // Under the flat 1-D launch gridDim.x = num_blocks (= num_sms) ≠ grid_x, - // so using gridDim.x here over-strides and over-spans: in the remote - // pass it reads slots from other laps of the circular staging ring - // (mis-attributed tile_expert_ids → corrupt dA/dB/dC at the ring-wrap - // boundary experts). The local pass tolerated it (fully-materialized - // non-circular buffer, harmless no-op tail batches), which is why 1-PE - // passed and multi-PE failed. num_batches is already sized ceil(total / - // grid_x), so grid_x is the matching stride. - int bk_start_unwrapped = batch * dims.grid_x; - int bk_tiles = min(dims.grid_x, - dims.total_m_tiles_in_pass - bk_start_unwrapped); - int bk_start = bk_start_unwrapped % dims.num_m_tiles; - - // Per-batch per-expert K-range scan over THIS batch's window of the - // (possibly wrapped) ring. Reads the LIVE X-pipe expert-id array - // (tile_expert_ids_x in the remote pass; release_src is deferred past - // mlp4, so it stays valid here AND for mlp3 since the ranges are cached - // into k_starts/ends before the scan's barrier). bk_start is already - // wrapped mod num_m_tiles, so the window indexes physical ring slots — - // matching the kb coordinates the mlp3/mlp4 producers feed to TMA. - if (bk_tiles > 0) { - int batch_kb_off = (bk_start * Traits1::TileM) / Traits4::TileK; - int batch_kb_cnt = (bk_tiles * Traits1::TileM) / Traits4::TileK; - mlp_bwd_phase_2_build_ranges( - bufs.expert_k_starts, bufs.expert_k_ends, - bufs.expert_for_k_block_mlp4, bufs.valid_rows_mlp4, - batch_kb_off, batch_kb_cnt, - dims.experts_per_pe, dims.local_expert_start, - efkb_stride_4, Traits4::TileK, warp_id); - } - - // Cross-CTA barrier after CTA 0 builds the shared ranges. MLP4 must not - // read them before the scan completes, or producer/consumer tile counts - // can diverge and deadlock the pipeline. - if (bk_tiles > 0) - global_barrier.wait(); - - if (bk_tiles > 0) { - mlp_bwd_run_phase_2_mlp4< - Traits1, Traits4, NSplit2, Compute>( - p4_pipe, p4_count, smem.mlp4, - bufs.expert_k_starts, bufs.expert_k_ends, - tma_load_xt4, tma_load_dut4, tma_load_dvt4, - tma_reduce_db, tma_reduce_dc, - warp_id, bk_start, bk_tiles, - phase2_div, phase2_mod, - dims, smem.tmem_base); - } - - global_barrier.wait(); - if (has_tile) - iter.release_src(threadIdx.x); - - if (bk_tiles > 0) { - mlp_bwd_run_phase_2_mlp3< - Traits1, Traits3, NSplit2, Compute>( - p3_pipe, p3_count, smem.mlp3, - bufs.expert_k_starts, bufs.expert_k_ends, - tma_load_dyt3, tma_load_zt3, tma_reduce_da, - warp_id, bk_start, bk_tiles, - phase2_div, phase2_mod, - dims, smem.tmem_base); - } - - global_barrier.wait(); - if (has_tile) - iter.release_dy(threadIdx.x); - } -} - -// ═══════════════════════════════════════════════════════════════════ -// mlp_fused_bwd — pipeline construction + batch loop (single pass) -// ═══════════════════════════════════════════════════════════════════ -// -// Generic over TileIter. Called once for local, once for remote. -// Mirrors forward mlp_fused_fwd: prologue → tile/batch loop. - -// Per-phase N-split parameters (autotuning surface). -// Phase 1's effective NSplits (1a, 1b', 1c, 1d) are tied to gridDim.y by -// construction: those phases iterate n-tiles via blockIdx.y/gridDim.y. To -// retune them, change the launch grid shape at the host level. -// Phase 2 (mlp3, mlp4) uses NSplit2 with virtual (vx, vy) remapping — it -// can differ from gridDim.y as long as gridDim.x * gridDim.y is divisible -// by NSplit2. -template -__device__ __forceinline__ void mlp_fused_bwd( - MlpFusedBwdSmem& smem, - TileIter& iter, - // Phase 1 TMA - TmaLoadX const& tma_load_x, TmaLoadW1 const& tma_load_b_fwd, - TmaLoadW1 const& tma_load_c_fwd, TmaStoreZ const& tma_store_z, - TmaStoreDU const& tma_store_du, TmaStoreDV const& tma_store_dv, - TmaLoadDY const& tma_load_dy, TmaLoadW2T const& tma_load_a_col, - TmaStoreDZ const& tma_store_dz, - TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, - TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, - TmaStoreDX const& tma_store_dx, - // Phase 2 TMA - TmaLoadDYT3 const& tma_load_dyt3, TmaLoadZT3 const& tma_load_zt3, - TmaReduceDA const& tma_reduce_da, - TmaLoadXT4 const& tma_load_xt4, - TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, - TmaReduceDB const& tma_reduce_db, TmaReduceDC const& tma_reduce_dc, - // Dims + buffers - const MlpBwdDims& dims, - const MlpBwdBufs& bufs, - // Target-based barriers — constructed ONCE by the caller in - // moe_fused_bwd and threaded through here so target counters - // accumulate monotonically across local/cross-phase/remote calls. - MlpBwdCtaBarrierT& global_barrier, - MlpBwdCtaBarrierT& x_barrier) { - - using Element = typename Traits1::Element; - int warp_id = threadIdx.x / Traits1::WarpSize; - constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; - // See mlp_fused_bwd_run_batches — Compute is the SOLE selector of the - // Compute == 100 ClusterM=2 paired-CTA Phase-2 path; the paired-CTA - // UMMA + cluster mbarrier rendezvous below is SM100-only. - static_assert(Compute == 90 || Compute == 100, - "mlp_fused_bwd: Compute must be 90 (SM90/Hopper, 1SM Phase-2) or " - "100 (SM100/Blackwell, ClusterM=2 paired-CTA Phase-2)."); - - // ── Prefetch TMA descriptors ── - cute::prefetch_tma_descriptor(tma_load_x.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_b_fwd.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_c_fwd.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_z.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_du.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_dv.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dy.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_a_col.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_dz.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_du.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dv.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_b_col.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_c_col.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_dx.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dyt3.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_zt3.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_reduce_da.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_xt4.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dut4.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dvt4.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_reduce_db.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_reduce_dc.get_tma_descriptor()); - - // ── Prologue: construct pipelines ── - // Pipeline constructor requires __syncthreads (all 384 threads). - // Must happen before comm warps exit. All five phases (1a/1b/1d, - // 2-mlp3, 2-mlp4) use a single cooperative fused pipe each; phase - // 1b additionally uses a tiny scale-load aux pipe (Stages=1) to - // carry the silu-bwd coefficients U'/V' produced by mlp1. - auto pa_pipe = [&]() { - if constexpr (Compute == 100) - return mlp1_make_pipe_umma(smem.pa_pipe); - else - return mlp1_make_pipe(smem.pa_pipe); - }(); - auto pb_pipe = [&]() { - if constexpr (Compute == 100) - return mlp2_t_make_pipe_umma(smem.pb_pipe); - else - return mlp2_t_make_pipe(smem.pb_pipe); - }(); - auto pd_pipe = [&]() { - if constexpr (Compute == 100) - return mlp5_make_pipe_umma(smem.pd_pipe); - else - return mlp5_make_pipe(smem.pd_pipe); - }(); - - // Single fused pipe for mlp3 carrying dYT + Z per acquire. On SM100, - // warp 3 issues UMMA and epilogue WGs 4-11 drain the same output tile. - // Compute == 100 always drives the SM100 ClusterM=2 paired-CTA pipe - // (mlp3_make_pipe_umma_2sm); Compute == 90 drives the existing Hopper - // pipe (mlp3_make_pipe). - auto p3_pipe = [&]() { - if constexpr (Compute == 100) - return mlp3_make_pipe_umma_2sm(smem.p3_pipe); - else - return mlp3_make_pipe(smem.p3_pipe); - }(); - // Single fused pipe for mlp4 carrying X + dU^T + dV^T per acquire. On - // SM100, warp 3 issues the two UMMAs and epilogue WGs 4-11 reduce-add - // dB/dC. Compute == 100 always drives the SM100 ClusterM=2 paired-CTA - // pipe (mlp4_make_pipe_umma_2sm); Compute == 90 drives the existing - // Hopper pipe (mlp4_make_pipe). - auto p4_pipe = [&]() { - if constexpr (Compute == 100) - return mlp4_make_pipe_umma_2sm(smem.p4_pipe); - else - return mlp4_make_pipe(smem.p4_pipe); - }(); - - // Comm warps exit. Hopper keeps the historical MLP partition (warps 0,4-11) - // and drops warp 3 immediately. Blackwell mirrors the forward path: comm - // uses only warps 1-2 and warp 3 joins MLP as the dedicated UMMA producer. - if constexpr (Compute == 90) { - if (warp_id == 3) return; - } - if (warp_id == 1 || warp_id == 2) return; - - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - cute::conditional_t< - Compute == 100, - cute::TMEM::Allocator2Sm, - cute::TMEM::Allocator1Sm> tmem_alloc{}; - if constexpr (Compute == 100) { - constexpr int kTmemColumns = - MlpBwdTmemColumns::value; - if (warp_id == 3) { - tmem_alloc.allocate(kTmemColumns, &smem.tmem_base); - if constexpr (Compute == 100) { - if (cute::elect_one_sync()) { - smem.tmem_pair_barrier.init(1); - smem.tmem_init_barrier.init(1); - smem.acc_pair_barrier.init(1); - } - cutlass::arch::fence_barrier_init(); - } - __syncwarp(); - } - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - mlp_bwd_pair_barrier_sync(smem.tmem_init_barrier, warp_id); - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - } - - // Keep phase-role tests at their call sites. Hoisting producer/MMA/epilogue - // booleans across every batch keeps them live through all phase mainloops, - // increasing register pressure in a kernel that already spills. - - // ── Persistent per-phase pipeline counts ───────────────────── - // - // We persist only the monotonic iteration count per phase (1 uint32_t) - // rather than the full PipelineState (3 ints: index_, phase_, count_). - // Each phase function reconstructs the live state on entry by calling - // `make_producer_start_state().advance(count)` (or the consumer - // equivalent) — this is correct because `advance(N)` deterministically - // computes the (index, phase) tuple that ++state would arrive at after - // N iterations, and that tuple matches the in-smem barrier phase the - // previous invocation left behind. - // - // PipelineState itself must NOT be persistent across phase invocations - // if we want construction inside the phase — but persisting only the - // count (5 regs across 5 phases vs 15 regs for full states) trims the - // cross-phase live set while keeping correctness. We hit a desync bug - // on T4096_D4096_I14336_E16_K4 (dA off 10% on expert 7) when state was - // reset fresh per batch without preserving the count. - uint32_t pa_count = 0; - uint32_t pb_count = 0; - uint32_t pd_count = 0; - uint32_t p3_count = 0; - uint32_t p4_count = 0; - - mlp_fused_bwd_run_batches< - Traits1, Traits2T, Traits3, Traits4, Traits5, - NSplit2, Compute>( - smem, pa_pipe, pb_pipe, pd_pipe, p3_pipe, p4_pipe, - pa_count, pb_count, pd_count, p3_count, p4_count, - iter, - tma_load_x, tma_load_b_fwd, tma_load_c_fwd, - tma_store_z, tma_store_du, tma_store_dv, - tma_load_dy, tma_load_a_col, tma_store_dz, - tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, - tma_load_dyt3, tma_load_zt3, tma_reduce_da, - tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, - EfkbStride3, EfkbStride4, - dims, bufs, global_barrier, x_barrier, warp_id); - - if constexpr (Compute == 100) { - constexpr int kTmemColumns = - MlpBwdTmemColumns::value; - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - if (warp_id == 3) { - tmem_alloc.release_allocation_lock(); - mlp_bwd_pair_barrier_sync(smem.tmem_pair_barrier, warp_id); - tmem_alloc.free(smem.tmem_base, kTmemColumns); - } - } -} - -// ═══════════════════════════════════════════════════════════════════ -// mlp_fused_bwd_dual — single-instance local + remote passes -// ═══════════════════════════════════════════════════════════════════ -// -// Builds pipelines + per-phase counts ONCE, runs the local batch loop, -// crosses the global barrier, and runs the remote batch loop in the -// SAME function body. The aim is to let PTXAS analyze both batch loops -// as sequential code in one CFG, reusing spill slots between them -// (lifetimes don't overlap — separated by global_barrier.wait()). This -// is the alternative to the current pattern in moe_fused_bwd of two -// inlined `mlp_fused_bwd` calls, each constructing its own pipelines -// and re-initializing the smem barriers. -// -// Pipeline barriers in smem are constructed ONCE here. The persistent -// `*_count` state is shared across passes and consumed by the count→ -// PipelineState resume logic inside each phase helper (see -// mlp_bwd_resume_state) — this matches the smem barriers' phase -// counters that advance across both passes. -template 1 → Phase-1 (mlp1a/mlp2t/ - // silu/mlp5) steps SubTiles GemmTileM-row sub-tiles per comm slot; - // Phase-2 sees num_m_tiles·SubTiles GemmTileM token-tiles with - // tile_expert_ids replicated SubTiles× (the K-block count is invariant). - // 1 → unchanged single-tile path. - int SubTiles = 1, - int Compute = 90, - typename FusedIter, - // Shared TMA (Phase 1 weights, Phase 1d B/C, Phase 2 reduce + Z/U/V) - typename TmaLoadW1, typename TmaStoreZ, typename TmaStoreDU, - typename TmaStoreDV, typename TmaLoadW2T, typename TmaStoreDZ, - typename TmaLoadDU, typename TmaLoadDV, - typename TmaLoadB, typename TmaLoadC, - typename TmaLoadZT3, typename TmaReduceDA, - typename TmaLoaddUT4, typename TmaLoaddVT4, - typename TmaReduceDB, typename TmaReduceDC, - // Per-role TMA types — local and remote instances share the - // same C++ type (same cute Layout: dynamic shape, _1 col stride, - // identical smem layout). Only the runtime gmem pointer differs. - typename TmaLoadX, typename TmaLoadDY, typename TmaStoreDX, - typename TmaLoadDYT3, typename TmaLoadXT4> -__device__ __forceinline__ void mlp_fused_bwd_dual( - MlpFusedBwdSmem& smem, - FusedIter& iter, - bool remote_active, - int efkb_stride_4_remote, - // Phase 1 TMA — weights are shared across tiles; X/dY/dX (the *_remote - // instances below) come from the comm staging ring, where EVERY tile, - // local or remote, is staged. - TmaLoadW1 const& tma_load_b_fwd, TmaLoadW1 const& tma_load_c_fwd, - TmaStoreZ const& tma_store_z, - TmaStoreDU const& tma_store_du, TmaStoreDV const& tma_store_dv, - TmaLoadW2T const& tma_load_a_col, TmaStoreDZ const& tma_store_dz, - TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, - TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, - // Phase 2 TMA (shared) - TmaLoadZT3 const& tma_load_zt3, TmaReduceDA const& tma_reduce_da, - TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, - TmaReduceDB const& tma_reduce_db, TmaReduceDC const& tma_reduce_dc, - // Staging-ring data TMAs: X/dY/dX (Phase 1) + dYᵀ/Xᵀ (Phase 2). - TmaLoadX const& tma_load_x_remote, - TmaLoadDY const& tma_load_dy_remote, - TmaStoreDX const& tma_store_dx_remote, - TmaLoadDYT3 const& tma_load_dyt3_remote, - TmaLoadXT4 const& tma_load_xt4_remote, - const MlpBwdDims& remote_dims, - const MlpBwdBufs& remote_bufs, - MlpBwdCtaBarrierT& global_barrier, - MlpBwdCtaBarrierT& x_barrier, - // Flat 1-D launch coordinates (computed once in moe_fused_bwd): - // flat_id = blockIdx.x - // col = flat_id / NSplit (Phase 1 logical column) - // grid_x = n_gemm / NSplit (Phase 1 column count) - // split = flat_id % NSplit (N-split index) - // num_splits = NSplit - // gemm_active = flat_id < n_gemm → this CTA runs Phase 1 + comm; - // gap CTAs skip Phase 1 but still run Phase 2 + barriers. - int flat_id, int col, int grid_x, int split, int num_splits, - bool gemm_active) { - - static_assert(SubBatch >= 1 && SubBatch <= kBwdSubBatchMax, - "SubBatch (MoeBwdConfig::kSubBatch) must be in [1, kBwdSubBatchMax]"); - // See mlp_fused_bwd_run_batches — Compute is the SOLE selector of the - // Compute == 100 ClusterM=2 paired-CTA Phase-2 path; the paired-CTA - // UMMA + cluster mbarrier rendezvous below is SM100-only. - static_assert(Compute == 90 || Compute == 100, - "mlp_fused_bwd_dual: Compute must be 90 (SM90/Hopper, 1SM Phase-2) " - "or 100 (SM100/Blackwell, ClusterM=2 paired-CTA Phase-2)."); - - int warp_id = threadIdx.x / Traits1::WarpSize; - constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; - - // Prefetch the TMA descriptors this pass uses: shared weights/intermediates - // + the staging-ring data descriptors. - cute::prefetch_tma_descriptor(tma_load_b_fwd.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_c_fwd.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_z.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_du.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_dv.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_a_col.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_dz.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_du.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dv.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_b_col.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_c_col.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_zt3.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_reduce_da.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dut4.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dvt4.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_reduce_db.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_reduce_dc.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_x_remote.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dy_remote.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_store_dx_remote.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_dyt3_remote.get_tma_descriptor()); - cute::prefetch_tma_descriptor(tma_load_xt4_remote.get_tma_descriptor()); - - // Build pipelines ONCE (smem barriers initialized in one place — the - // remote pass picks up smem-barrier phase from where local left off). - auto pa_pipe = [&]() { - if constexpr (Compute == 100) - return mlp1_make_pipe_umma(smem.pa_pipe); - else - return mlp1_make_pipe(smem.pa_pipe); - }(); - auto pb_pipe = [&]() { - if constexpr (Compute == 100) - return mlp2_t_make_pipe_umma(smem.pb_pipe); - else - return mlp2_t_make_pipe(smem.pb_pipe); - }(); - auto pd_pipe = [&]() { - if constexpr (Compute == 100) - return mlp5_make_pipe_umma(smem.pd_pipe); - else - return mlp5_make_pipe(smem.pd_pipe); - }(); - - // Single fused pipe for mlp3 carrying dYT + Z per acquire. Compute == - // 100 always drives the SM100 ClusterM=2 paired-CTA pipe - // (mlp3_make_pipe_umma_2sm); Compute == 90 drives the existing Hopper - // pipe (mlp3_make_pipe). - auto p3_pipe = [&]() { - if constexpr (Compute == 100) - return mlp3_make_pipe_umma_2sm(smem.p3_pipe); - else - return mlp3_make_pipe(smem.p3_pipe); - }(); - // Single fused pipe for mlp4 carrying X + dU^T + dV^T per acquire. - // Compute == 100 always drives the SM100 ClusterM=2 paired-CTA pipe - // (mlp4_make_pipe_umma_2sm); Compute == 90 drives the existing Hopper - // pipe (mlp4_make_pipe). - auto p4_pipe = [&]() { - if constexpr (Compute == 100) - return mlp4_make_pipe_umma_2sm(smem.p4_pipe); - else - return mlp4_make_pipe(smem.p4_pipe); - }(); - - if constexpr (Compute == 90) { - if (warp_id == 3) return; - } - if (warp_id == 1 || warp_id == 2) return; - - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - cute::conditional_t< - Compute == 100, - cute::TMEM::Allocator2Sm, - cute::TMEM::Allocator1Sm> tmem_alloc{}; - if constexpr (Compute == 100) { - constexpr int kTmemColumns = - MlpBwdTmemColumns::value; - if (warp_id == 3) { - tmem_alloc.allocate(kTmemColumns, &smem.tmem_base); - if constexpr (Compute == 100) { - if (cute::elect_one_sync()) { - smem.tmem_pair_barrier.init(1); - smem.tmem_init_barrier.init(1); - smem.acc_pair_barrier.init(1); - } - cutlass::arch::fence_barrier_init(); - } - __syncwarp(); - } - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - mlp_bwd_pair_barrier_sync(smem.tmem_init_barrier, warp_id); - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - } - - // Counts persist across both passes — smem barriers continue advancing - // from local into remote, and PipelineState::advance(count) re-derives - // the right (index, phase) for the resume. - uint32_t pa_count = 0; - uint32_t pb_count = 0; - uint32_t pd_count = 0; - uint32_t p3_count = 0; - uint32_t p4_count = 0; - uint32_t acc_pair_phase = 0; - - using Element = typename Traits1::Element; - - // Phase 2 flat-cell coords are CTA-invariant — hoist out of the loops. - int phase2_div = flat_id / NSplit2; - int phase2_mod = flat_id % NSplit2; - - // Shared k_starts/k_ends require equal TileK (one scan feeds mlp3+mlp4). - static_assert(Traits3::TileK == Traits4::TileK, - "Phase 2 shared k_starts/ends require Traits3::TileK == Traits4::TileK"); - - // Nothing to do if this PE has no tiles (remote_active is grid-uniform, so - // every CTA agrees and the early-out keeps the barriers below well-defined). - if (!remote_active) { - global_barrier.wait(); - if constexpr (Compute == 100) { - constexpr int kTmemColumns = - MlpBwdTmemColumns::value; - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - if (warp_id == 3) { - tmem_alloc.release_allocation_lock(); - mlp_bwd_pair_barrier_sync( - smem.tmem_pair_barrier, warp_id); - tmem_alloc.free(smem.tmem_base, kTmemColumns); - } - } - return; - } - - // One cross-CTA sync so all MLP CTAs enter the grouped loop together: - // SM90 synchronizes w0+w4-11 (288 threads); SM100 synchronizes - // w0+w3-11 (320 threads). Comm warps 1-2 have already returned. - global_barrier.wait(); - - // ═══════════════════════════════════════════════════════════════ - // REMOTE GROUPED PASS — kBwdSubBatch Phase-1 sub-batches, then ONE - // Phase-2 (mlp4+mlp3) over their combined kBwdSubBatch·grid_x K-window. - // ═══════════════════════════════════════════════════════════════ - // Phase 1 (mlp1/2/5) runs in SMALL per-sub-batch chunks (less CTA - // barrier-stall); Phase 2 (weight grads dA/dB/dC) runs ONCE over the whole - // group so its cooperative cell-walk has a long K-reduction to fill the grid - // (deeper outer_split / k_split). Each sub-batch keeps its X/dY staging slot - // un-released until the grouped Phase 2 has read X^T (mlp4) and dY^T (mlp3); - // the saved slots are freed in a batch after each Phase-2 sub-pass. The host - // guarantees CommNumStages is a multiple of kBwdSubBatch and ≥ it, so the - // group window never straddles the ring wrap (per-expert k-ranges stay - // contiguous) and comm has prefetch headroom. kBwdSubBatch == 1 reduces this - // to the original per-batch loop. - const int remote_batches = remote_dims.num_batches; - - // Adaptive group size: a group of `g` sub-batches holds g·grid_x live staging - // slots that must map to DISTINCT ring slots, so g·grid_x ≤ ring. Cap the - // group by how many grid_x-windows the ring holds — deep rings get the full - // SubBatch; shallow rings degrade to fewer (down to 1 = original per-batch - // loop). Grid-uniform (num_m_tiles, grid_x are host-set), so every CTA agrees. - const int eff_subbatch = - max(1, min((int)SubBatch, remote_dims.num_m_tiles / grid_x)); - - // Phase-1 GEMM consumers + Phase-2 see the token axis in GemmTileM - // (=Traits1::TileM) tiles: num_m_tiles (Phase-1 buffer height) and num_tokens - // (M extent) scale by SubTiles = CommTileM/GemmTileM. The comm loop control - // (eff_subbatch, group_start, bk_tiles clamp, slot releases) stays in - // communication-slot units via remote_dims. At SubTiles=1 the two are identical. - // (Other dims fields — N/K extents, expert counts — are M-axis-independent.) - MlpBwdDims gemm_dims = remote_dims; - gemm_dims.num_m_tiles *= SubTiles; - gemm_dims.num_tokens *= SubTiles; - int first_sub = 0; - int last_sub = SubTiles; - int gemm_split = split; - int gemm_splits = num_splits; - constexpr int kM64PerGemmTile = Traits1::TileM / 64; - - for (int g0 = 0; g0 < remote_batches; g0 += eff_subbatch) { - const int gcount = min(eff_subbatch, remote_batches - g0); - - // Per-sub-batch staging slots, freed together after Phase 2. cur_slot in - // the iterator only tracks the LAST acquire, so releasing the earlier - // sub-batches needs their saved slot (release indexes *_consumed[slot]). - int saved_slot[SubBatch]; - bool saved_live[SubBatch]; - - // ── Phase 1: gcount sub-batches (mlp1a / mlp2t / silu_bwd / mlp5) ── - for (int s = 0; s < gcount; ++s) { - int saved_m = -1, saved_expert = -1; - int saved_valid_m64_subtiles = 0; - bool has_tile = false; - bool compute_tile = false; - - // Gated by gemm_active — gap CTAs (flat_id ≥ n_gemm) skip Phase 1 and - // the per-column x_barrier, but still hit every global barrier below. - if (gemm_active && iter.has_next()) { - iter.acquire_src(); - cutlass::arch::NamedBarrier::sync( - kMlpBwdBarrierThreads, kMlpBarrierId); - auto tile = iter.next(); - int expert = tile.expert - remote_dims.local_expert_start; - int m = tile.y_m; - - saved_m = m; - saved_expert = expert; - saved_valid_m64_subtiles = tile.valid_m64_subtiles; - has_tile = true; - compute_tile = - first_sub * kM64PerGemmTile < - saved_valid_m64_subtiles; - - // One CommTileM-row slot feeds SubTiles GemmTileM-row sub-tiles. - // gemm coordinate of sub-tile `sub` = m·SubTiles + sub (the - // Phase-1 kernels multiply by Traits1::TileM=GemmTileM → rows - // m·CommTileM + sub·GemmTileM). gemm_dims carries the SubTiles-scaled - // num_m_tiles / num_tokens so the buffer tensor shapes match. - // Acquire/release stays once per communication tile. - if (compute_tile) { - for (int sub = first_sub; sub < last_sub; ++sub) { - mlp_bwd_run_phase_1a( - pa_pipe, pa_count, smem.mlp1, - tma_load_x_remote, tma_load_b_fwd, tma_load_c_fwd, - tma_store_du, tma_store_dv, tma_store_z, - warp_id, m * SubTiles + sub, expert, gemm_dims, - smem.tmem_base, gemm_split, gemm_splits); - } - } - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - x_barrier.wait(); - - iter.acquire_dy(); - if (compute_tile) { - for (int sub = first_sub; sub < last_sub; ++sub) { - mlp_bwd_run_phase_1b( - pb_pipe, pb_count, smem.mlp2t, - tma_load_dy_remote, tma_load_a_col, tma_store_dz, - warp_id, m * SubTiles + sub, expert, gemm_dims, - smem.tmem_base, gemm_split, gemm_splits); - } - } - - x_barrier.wait(); - - if (compute_tile) { - for (int sub = first_sub; sub < last_sub; ++sub) { - silu_bwd_pair_tile( - remote_bufs.dz_buf, remote_bufs.du_buf, remote_bufs.dv_buf, - remote_bufs.du_buf, remote_bufs.dv_buf, - m * SubTiles + sub, Traits1::TileM, Traits2T::TileN, - remote_dims.intermediate_dim, gemm_split, gemm_splits); - } - } - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - } - - // Save this sub-batch's staging slot (== tile.y_m == iter cur_slot) - // for the deferred grouped release. -1 / false when this CTA had no - // tile (gap CTA or column exhausted). - saved_slot[s] = saved_m; - saved_live[s] = has_tile; - - // Global barrier: this sub-batch's Z, dU, dV visible to all CTAs. - global_barrier.wait(); - - // Phase 1d: dX = dU@B + dV@C (per sub-batch — Phase 2 never reads dX, - // so it is emitted and its dst slot released immediately). - if (has_tile) { - iter.acquire_dst(); - if (compute_tile) { - for (int sub = first_sub; sub < last_sub; ++sub) { - mlp_bwd_run_phase_1d( - pd_pipe, pd_count, smem.mlp5, - tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, - tma_store_dx_remote, - warp_id, saved_m * SubTiles + sub, saved_expert, gemm_dims, - smem.tmem_base, gemm_split, gemm_splits); - } - } - __threadfence(); // dX flush for the NIC put (release_dst signals) - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - iter.release_dst(threadIdx.x); - } - } - - // ── Phase 2: ONE pass over the group's gcount·grid_x K-window ── - // Phase 1 advances each of grid_x columns by one tile per sub-batch, so - // the group delivered global tiles [g0·grid_x, (g0+gcount)·grid_x). - // bk_start is UNWRAPPED (group_start) so the group's K-window stays - // contiguous even when it straddles the staging-ring wrap; build_ranges - // and the producers wrap the physical slot by % ring_kb internally - // (ring_kb > 0 enables it). The group's gcount·grid_x live slots are - // guaranteed distinct (host check: ring > SubBatch·grid_x). - // group_start / bk_tiles_comm are in communication-slot units. Phase 2 - // works in GemmTileM units, so scale the contiguous parent window by - // SubTiles; range construction skips padded K-blocks inside it. - int group_start = g0 * grid_x; - int bk_tiles_comm = min( - gcount * grid_x, - remote_dims.total_m_tiles_in_pass - group_start); - int bk_start = group_start * SubTiles; - int bk_tiles = bk_tiles_comm * SubTiles; - int ring_kb = (gemm_dims.num_m_tiles * Traits1::TileM) / Traits4::TileK; - const cutlass::arch::ClusterBarrier* acc_pair_barrier = nullptr; - if constexpr (Compute == 100) - acc_pair_barrier = &smem.acc_pair_barrier; - - // Per-group per-expert K-range scan over the group's window. Reads the - // LIVE X-pipe expert ids — none of the group's slots have been released - // yet, so all gcount·grid_x are valid. Feeds mlp3+mlp4 (unwrapped ranges). - if (bk_tiles > 0) { - int batch_kb_off = (bk_start * Traits1::TileM) / Traits4::TileK; - int batch_kb_cnt = (bk_tiles * Traits1::TileM) / Traits4::TileK; - mlp_bwd_phase_2_build_ranges( - remote_bufs.expert_k_starts, remote_bufs.expert_k_ends, - remote_bufs.expert_for_k_block_mlp4, - remote_bufs.valid_rows_mlp4, - batch_kb_off, batch_kb_cnt, - remote_dims.experts_per_pe, remote_dims.local_expert_start, - efkb_stride_4_remote, Traits4::TileK, warp_id, ring_kb); - // Cross-CTA barrier: CTA 0 owns reset+scan, so no producer may read - // the shared ranges until that work completes. - global_barrier.wait(); - } - - if (bk_tiles > 0) { - mlp_bwd_run_phase_2_mlp4< - Traits1, Traits4, NSplit2, Compute>( - p4_pipe, p4_count, smem.mlp4, - remote_bufs.expert_k_starts, remote_bufs.expert_k_ends, - tma_load_xt4_remote, tma_load_dut4, tma_load_dvt4, - tma_reduce_db, tma_reduce_dc, - warp_id, bk_start, bk_tiles, - phase2_div, phase2_mod, - gemm_dims, smem.tmem_base, ring_kb, - acc_pair_barrier, &acc_pair_phase); - } - - global_barrier.wait(); - // mlp4 done reading X^T — free every X slot the group held. - { - __threadfence(); - for (int s = 0; s < gcount; ++s) - if (saved_live[s]) iter.release_src_slot(saved_slot[s], threadIdx.x); - } - - if (bk_tiles > 0) { - mlp_bwd_run_phase_2_mlp3< - Traits1, Traits3, NSplit2, Compute>( - p3_pipe, p3_count, smem.mlp3, - remote_bufs.expert_k_starts, remote_bufs.expert_k_ends, - tma_load_dyt3_remote, tma_load_zt3, tma_reduce_da, - warp_id, bk_start, bk_tiles, - phase2_div, phase2_mod, - gemm_dims, smem.tmem_base, ring_kb, - acc_pair_barrier, &acc_pair_phase); - } - - global_barrier.wait(); - // mlp3 done reading dY^T — free every dY slot the group held. - { - __threadfence(); - for (int s = 0; s < gcount; ++s) - if (saved_live[s]) iter.release_dy_slot(saved_slot[s], threadIdx.x); - } - } - - if constexpr (Compute == 100) { - constexpr int kTmemColumns = - MlpBwdTmemColumns::value; - cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); - if (warp_id == 3) { - tmem_alloc.release_allocation_lock(); - mlp_bwd_pair_barrier_sync(smem.tmem_pair_barrier, warp_id); - tmem_alloc.free(smem.tmem_base, kTmemColumns); - } - } -} - - -} // namespace liger +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif + +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp_bwd_sm90.cuh" +#else +#include "mlp_bwd_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm100.cuh new file mode 100644 index 00000000..3bd6458f --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm100.cuh @@ -0,0 +1,1885 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP Backward — batch-interleaved Phase 1 + Phase 2 +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors forward mlp.cuh: pipeline construction + tile/batch loop. +// Called once for local tiles, once for remote tiles. +// +// Per batch (gridDim.x tiles): +// Phase 1: one tile per CTA (a→b→c→d sub-phases) +// Global barrier +// Phase 2: all CTAs do weight grads for this batch's K-blocks +// Global barrier +// +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp_sm100.cuh" // kMlpBarrierId, mlp1 pipe helpers +#include "mlp1_fused_act_sm100.cuh" // bwd Phase 1a — TileM=128 cooperative-M-split +#include "mlp2_t_fused_sm100.cuh" // bwd Phase 1b' — single-m-tile entry over mlp2_t.cuh +#include "mlp5_sm100.cuh" +#include "mlp5_fused_sm100.cuh" // mlp5_fused_producer/consumer (single-m-tile entry) +#include "mlp3_sm100.cuh" // Mlp3Traits, Mlp3FusedSmem +#include "mlp4_sm100.cuh" // Mlp4Traits, Mlp4FusedSmem +#include "silu_bwd_fused_sm100.cuh" // silu_bwd_pair_tile — separate elementwise phase +#include "cta_barrier_sm100.cuh" // target-based CtaCounterBarrier + +#include // INT_MAX for Phase 2 per-batch range scan + +namespace liger { + +using namespace cute; + +// MLP-warps barrier count for the bwd kernel. Hopper uses the historical +// 9-warp partition (w0 + w4-11). Blackwell mirrors the forward path: comm +// uses only warps 1-2 and warp 3 joins MLP as the dedicated UMMA producer. +template +static constexpr int kMlpBwdBarrierThreadsFor = (Compute == 100 ? 10 : 9) * 32; +static_assert(kMlpBwdBarrierThreadsFor<90> == 288); +static_assert(kMlpBwdBarrierThreadsFor<100> == 320); + +// Compute == 100 (Blackwell/SM100) always means the Phase-2 GEMMs (mlp3 AND +// mlp4) run the SM100 ClusterM=2 paired-CTA path; Compute == 90 (Hopper/SM90) +// always means plain 1SM. There is no per-phase or per-trait choice anymore +// (see mlp3.cuh / mlp4.cuh's unified mlp{3,4}_producer / +// mlp{3,4}_consumer), so `Compute == 100` alone drives the +// TMEM allocator choice (Allocator2Sm vs Allocator1Sm) and the extra +// tmem_pair_barrier below: both mlp3 and mlp4 share ONE TMEM allocation per +// CTA-pair, so the pair must rendezvous around the allocate/free calls. + +// ── Compute == 100 (paired-CTA) caller contract ── +// +// This header only builds the __device__ Phase-2 GEMM bodies; it does not +// launch the kernel or build TMA descriptors, so the following invariants +// must hold at the call site (host launcher, e.g. moe_bwd.cu) whenever +// Compute == 100 is selected — violating any of them is a silent-corruption +// or launch-failure risk, not something this file can assert at compile time: +// • Grid must be launched EVEN (gridDim.x·gridDim.y divisible by +// Traits{3,4}::ClusterM) — the flat cell-walk above (cell_start = +// flat_id / ClusterM, cell_stride = total_phase2_ctas / ClusterM) +// assumes every CTA has a live cluster peer. +// • Launch via cudaLaunchKernelEx with clusterDim = (ClusterM, 1, 1) (or +// the ClusterShape Traits{3,4} defines) and the +// cudaLaunchAttributeClusterDimension / +// NonPortableClusterSizeAllowed attribute set (ClusterM > 8 or any +// non-default shape needs the "non-portable" opt-in). +// • The dYT/Z (mlp3) and X/dU^T/dV^T (mlp4) TMA descriptors passed in +// (TmaLoadDYT3/TmaLoadZT3/... below) must be built with +// make_tma_copy_A_sm100 / make_tma_copy_B_sm100 (2-CTA-aware multicast +// descriptors) — a plain SM90 make_tma_copy descriptor does not carry +// the multicast-mask semantics mlp{3,4}_producer's Compute == 100 +// `copy(tma.with(barrier, mcast_mask), ...)` calls rely on. +// • MlpFusedBwdSmem's static size (larger for the Compute == 100 +// Mlp3FusedSmem2Sm variant) must fit within the dynamically-queried +// cudaDevAttrMaxSharedMemoryPerBlockOptin limit, and the kernel's +// cudaFuncAttributeMaxDynamicSharedMemorySize must be opted in to that +// size before launch — this header assumes the smem struct it defines +// already fits whatever the host allocated. +// See Compute == 100's use below (TMEM Allocator2Sm + tmem_pair_barrier) and +// the Compute == 90/100 static_asserts in mlp_fused_bwd_run_batches / +// mlp_fused_bwd / mlp_fused_bwd_dual for the invariants this file DOES +// enforce. + +// ═══════════════════════════════════════════════════════════════════ +// Phase-1 sub-batch grouping (mlp1/2/5 small chunks, mlp3/4 one big batch) +// ═══════════════════════════════════════════════════════════════════ +// +// Remote BWD groups kBwdSubBatch Phase-1 sub-batches before running ONE +// Phase-2 (mlp3/mlp4) over all of them at once. Phase 1 (mlp1/2/5) prefers +// SMALL per-batch chunks (less CTA barrier-stall); Phase 2 (weight grads, +// dA/dB/dC) prefers a LONG K-reduction so the cooperative cell-walk fills +// the grid (deeper outer_split / k_split). Grouping serves both. +// +// Buffer survival: the SubBatch sub-batches keep their X/dY comm-staging slots +// un-released until the grouped Phase 2 consumes them, and their Z/dU/dV scratch +// slots must not be overwritten. This needs the staging ring to be DEEPER than +// the group span — remote_num_m_tiles > SubBatch·grid_x — so the group's live +// slots fit with comm prefetch headroom (host-checked in moe_bwd.cu). No ring/ +// group ALIGNMENT is required: the grouped Phase-2 window wraps the ring exactly +// as the original per-batch loop does (a window that runs past the ring end is +// only correct when it does not actually straddle — i.e. global_total ≤ ring or +// grid_x·SubBatch divides the ring — same property the un-grouped loop has). +// SubBatch == 1 reduces the grouped loop EXACTLY to the original per-batch loop. +// +// The sub-batch count is a pure TEMPLATE ARGUMENT (SubBatch) on moe_fused_bwd / +// mlp_fused_bwd_dual, sourced from MoeBwdConfig::kSubBatch (whose default is a +// literal — see moe_bwd.cu). Distinct SubBatch instantiations coexist in one +// binary; there is no preprocessor knob. +static constexpr int kBwdSubBatchMax = 8; // upper bound for the per-group save arrays + +// ═══════════════════════════════════════════════════════════════════ +// Fused backward shared memory — union of all phases +// ═══════════════════════════════════════════════════════════════════ + +template +struct MlpFusedBwdSmem { + union { + Mlp1FusedActSmem mlp1; // TileM=128 cooperative-M-split + Mlp2TFusedSmem mlp2t; // TileM=128 cooperative-M-split (mlp2_t.cuh) + Mlp5Smem mlp5; // TileM=128 cooperative-M-split + // mlp4's fused smem is traits-conditional internally (see mlp4.cuh), + // so it needs no separate 2SM specialization here. mlp3 instead uses + // a wholly distinct 2SM smem layout (Mlp3FusedSmem2Sm — paired-CTA + // accumulator pipeline in place of mlp3's normal AccumulatorPipeline), + // selected purely by Compute == 100 (Blackwell always pairs mlp3's + // Phase-2 GEMM across CTAs; Compute == 90 never does). + cute::conditional_t< + Compute == 100, + Mlp3FusedSmem2Sm, + Mlp3FusedSmem> mlp3; + Mlp4FusedSmem mlp4; + }; + + // Phase 2 per-expert K-range scratch lives in gmem (MlpBwdBufs: + // expert_k_starts / expert_k_ends). Allocated to + // exactly `experts_per_pe` ints — no compile-time cap. + + // Phase 1a, 1b', 1d: single fused pipe each (1c silu_bwd is + // elementwise gmem→gmem, no pipe). + typename Mlp1MainloopPipelineFor::SharedStorage pa_pipe; + typename Mlp2TMainloopPipelineFor::SharedStorage pb_pipe; + typename Mlp5MainloopPipelineFor::SharedStorage pd_pipe; + + // Phase 2: single fused pipe each for mlp3 and mlp4. Outputs (dA, + // dB, dC) are accumulated via TMA_REDUCE_ADD to gmem — no prev-pipe + // needed. + typename Mlp3MainloopPipelineFor::SharedStorage p3_pipe; + typename Mlp4MainloopPipelineFor::SharedStorage p4_pipe; + alignas(16) uint32_t tmem_base; + // Rendezvous barrier for the paired-CTA TMEM allocation window + // (Compute == 100 only — one CTA-pair shares one TMEM allocation, so + // both peers must arrive before either frees it). Collapses to a plain + // uint32_t placeholder (no ClusterBarrier construction/arrive/wait) on + // Compute == 90, which never pairs CTAs. + alignas(16) cute::conditional_t< + Compute == 100, + cutlass::arch::ClusterBarrier, + uint32_t> tmem_pair_barrier; + alignas(16) cute::conditional_t< + Compute == 100, + cutlass::arch::ClusterBarrier, + uint32_t> tmem_init_barrier; + alignas(16) cute::conditional_t< + Compute == 100, + cutlass::arch::ClusterBarrier, + uint32_t> acc_pair_barrier; +}; + +__device__ __forceinline__ void mlp_bwd_pair_barrier_sync( + const cutlass::arch::ClusterBarrier& barrier, + int warp_id) { + if (warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + barrier.arrive(peer, !leader); + barrier.wait(0); + barrier.arrive(peer, leader); + } + __syncwarp(); + } +} + +template +struct MlpBwdTmemColumns { + static constexpr int kMlp1Cols = Traits1::AccStages * (2 * Traits1::TileN); + static constexpr int kMlp2TCols = Traits2T::AccStages * Traits2T::TileN; + static constexpr int kMlp5Cols = Traits5::AccStages * Traits5::TileN; + static constexpr int kMlp3Cols = Traits3::AccStages * Traits3::TileN; + static constexpr int kMlp4Cols = Traits4::AccStages * (2 * Traits4::WgTileN); + + static constexpr int kPhase15Cols = (kMlp1Cols > kMlp5Cols) + ? kMlp1Cols : kMlp5Cols; + static constexpr int kPhase2T3Cols = (kMlp2TCols > kMlp3Cols) + ? kMlp2TCols : kMlp3Cols; + static constexpr int kPhase1235Cols = (kPhase15Cols > kPhase2T3Cols) + ? kPhase15Cols : kPhase2T3Cols; + static constexpr int value = (kPhase1235Cols > kMlp4Cols) + ? kPhase1235Cols : kMlp4Cols; + static_assert(value <= 512, "BWD SM100 TMEM allocation must fit one CTA's 512 columns"); +}; + +// ═══════════════════════════════════════════════════════════════════ +// MlpBwdDims — invariant backward dimensions (__grid_constant__) +// ═══════════════════════════════════════════════════════════════════ + +struct MlpBwdDims { + int hidden_dim; + int intermediate_dim; + + // Flat 1-D launch geometry (host-computed). The grid is launched as + // dim3(num_blocks) with num_blocks = num_sms, so gridDim is no longer + // the logical grid. Phase 1 + comm partition over the GEMM-active subset: + // n_gemm = floor_NS = largest multiple of NSplit ≤ num_sms + // grid_x = n_gemm / NSplit (Phase 1 column count) + // CTAs with flat_id ≥ n_gemm fall through Phase 1/comm but still run + // Phase 2 (grid-stride over all num_blocks CTAs) and every barrier. + int grid_x; + int n_gemm; + + // Phase 1a: recompute + int total_n_rows_1; + int num_n_tiles_1; + int num_k_tiles_1; + + // Phase 1b: dZ = dY @ A + int total_k_cols_2t; + int num_n_tiles_2t; + int num_k_tiles_2t; + + // Phase 1d: dX = dU@B + dV@C + int total_k_cols_5; + int num_n_tiles_5; + int num_k_tiles_5; + + // Phase 2: mlp3 (dA) + int num_m_tiles_3; + int num_n_tiles_3; + int total_n_rows_3; + + // Phase 2: mlp4 (dB, dC) + int num_m_tiles_4; + int num_n_tiles_4; + int total_m_rows_4; + + int* phase_counter; + + // Scalar params (in __grid_constant__ to avoid register pressure) + int num_tokens; + int num_m_tiles; // Phase 1 output buffer height (in m-tiles). + // Local: total tile capacity (= max_total_slots/TileM). + // Remote: gridDim.x * CommNumStages (the staging ring). + // Phase 2 wraps bk_start modulo this when computing + // the slot index into z_buf/du_buf/dv_buf — for the + // local pass that's a no-op; for the remote pass it + // brings batches ≥ NumStages back into the live ring. + int total_m_tiles_in_pass; // Phase 2 K-block iteration bound (in m-tiles). + // Local: equals num_m_tiles (existing semantics — Phase 2 + // walks the full padded buffer and skips fake K-blocks + // via the expert_for_k_block sentinel). + // Remote: equals smem.comm.global_total — strictly less + // than num_m_tiles when per_cta_tiles > NumStages, so + // batches that would walk past the ring still run Phase 2 + // (with bk_start wrapped) instead of being silently + // clamped to bk_tiles=0. + int num_batches; + int local_expert_start; + int num_pes; + int experts_per_pe; // Phase 3/4 use this to skip k-blocks whose expert + // falls outside [local_expert_start, + // local_expert_start + experts_per_pe). Required + // because expert_for_k_block (built post-sort) holds + // GLOBAL expert IDs spanning all PEs, but dA/dB/dC + // on this PE only have slots for the local subset. + + // Phase 2 virtual-grid block count along the non-NSplit2 axis: + // phase2_num_blocks = (gridDim.x * gridDim.y) / NSplit2. + // Host-known (depends only on launch grid + template constant NSplit2), + // so storing here lets every batch's mlp3/mlp4 call read it from grid + // constant memory instead of recomputing `gridDim.x * gridDim.y / + // NSplit2` per batch. + int phase2_num_blocks; + + // Remote-pass overrides for num_tokens / num_m_tiles. Both are host- + // known (gridDim.x * CommNumStages and that × Traits1::TileM); kept here + // so moe_fused_bwd's remote_dims build can just assign from grid constant + // instead of computing on device. The remote-only fields stay = 0 for + // the local pass (the local pass doesn't read them). + int remote_num_m_tiles; + int remote_num_tokens; + + // 1 when the bwd TMA-GET comm path is active (NVLink P2P + divisible + // hidden_dim + bounce smem fits). Selects the comm warp layout AND the + // remote MLP iterator's producer/consumer counts (1·NC get / 2·NC put + // vs 2·NC get / 1·NC put). Host-set; 0 keeps the original getmem layout. + int tma_get_enabled; + + // (efkb_stride_3 / efkb_stride_4 moved to template parameters on + // mlp_fused_bwd — see EfkbStride3/EfkbStride4 below. They were always + // compile-time constants per call site (local: 1, remote: + // Traits1::TileM/Traits{3,4}::TileK), so passing them as runtime ints + // just kept Traits1::TileM and the divides live in registers across the + // whole batch loop. Hoisting to templates lets the compiler constant- + // fold the `/1` (local) into a no-op and the remote ratios into a shift. +}; + +// ═══════════════════════════════════════════════════════════════════ +// MlpBwdBufs — buffer pointers (__grid_constant__) +// ═══════════════════════════════════════════════════════════════════ + +template +struct MlpBwdBufs { + // du_buf holds U' (silu-bwd coefficient) after mlp1, then dU after silu_bwd. + // dv_buf holds V' (silu-bwd coefficient) after mlp1, then dV after silu_bwd. + // dz_buf holds dZ between mlp2_t (Phase 1b') and silu_bwd (Phase 1c). + Element* du_buf; + Element* dv_buf; + Element* dz_buf; + + // Expert id per K-block, read by Phase 2's mlp4 and mlp3 respectively. + // They hold the same logical content — for K-block k, the local expert + // id of the tile that K-block belongs to — but the source array differs + // because each Phase-2 sub-pass has a different freshness requirement. + // + // expert_for_k_block_mlp4 is read by mlp4 (which reads X). In the + // remote pass it points at + // CommBuffersBwd::tile_expert_ids_x — comm + // won't overwrite it until MLP's release_src, + // which is deferred past mlp4. + // + // expert_for_k_block_mlp3 is read by mlp3 (which reads dY). In the + // remote pass it points at + // CommBuffersBwd::tile_expert_ids_dy — comm + // won't overwrite it until MLP's release_dy, + // which is deferred past mlp3. + // + // In the local pass and the standalone bwd kernel, both pointers alias + // the same post-sort expert_for_k_block built by build_expert_for_k_block + // — there is no comm race to defend against, and a single source is + // sufficient. + const int* expert_for_k_block_mlp4; + const int* expert_for_k_block_mlp3; + const int* valid_rows_mlp4; + + // Per-pass per-expert K-range scratch for Phase 2. Filled once at + // the start of each pass (in the bwd preamble) and consumed by both + // mlp3 and mlp4 batches — they require Traits3::TileK == Traits4:: + // TileK (the codebase already enforces this since a single + // expert_for_k_block indexes both), so one shared scan suffices. + // Sized exactly `experts_per_pe` ints each. + int* expert_k_starts; + int* expert_k_ends; + + Element const* gmem_da; + Element const* gmem_db; + Element const* gmem_dc; + int* barrier_counter; +}; + +// ═══════════════════════════════════════════════════════════════════ +// MLP-warps-only barriers (BWD) +// ═══════════════════════════════════════════════════════════════════ +// +// Both barriers use the target-based CtaCounterBarrier (see +// cta_barrier.cuh). The earlier modulo variants were susceptible to a +// CTA-skew deadlock where a fast CTA could push the counter past the +// `% stride == 0` window before slow CTAs polled. +// +// • MlpBwdGlobalBarrier: cross-CTA, stride = gridDim.x * gridDim.y. +// Counter is a single int in device memory. +// • MlpBwdXBarrier: per-blockIdx.x, stride = gridDim.y. +// Counter is a gmem array of length gridDim.x. +template +using MlpBwdCtaBarrierT = + CtaCounterBarrier, kMlpBarrierId>; + +// ═══════════════════════════════════════════════════════════════════ +// Per-phase inline runners +// ═══════════════════════════════════════════════════════════════════ +// +// Each phase of the bwd batch loop (1a / 1b' / 1d / 2-mlp4 / 2-mlp3) is +// pulled out into its own __forceinline__ device function. They take the +// per-phase PipelineState by REFERENCE — the caller (`mlp_fused_bwd`) +// holds the persistent state at function scope, and the inline body sees +// it as a direct register reference once inlined. Using actual functions +// instead of capture-all-by-reference lambdas avoids forcing the captured +// state into stack/local memory (CUDA closures pin captured vars to +// addressable storage, even when __forceinline__ would otherwise let +// them stay in registers). +// +// State persistence: the in-smem pipeline barriers carry a phase counter +// that advances each iteration and persists across phase invocations. To +// keep correctness with construction-inside-the-phase, the outer scope +// persists only a monotonic `uint32_t count` per phase (1 reg vs the +// 3-reg full PipelineState). The phase function: +// 1. Constructs a fresh PipelineState locally (warp-role-aware initial +// phase: producer starts phase=1, consumer starts phase=0). +// 2. Calls `advance(count)` to align it with the actual smem-barrier +// phase as left by the previous invocation. +// 3. Runs the phase, naturally advancing local state. +// 4. Writes `count = state.count()` back through the reference. +// This shrinks the cross-phase live set from 5×3 = 15 regs to 5×1 = 5 regs +// without re-introducing the desync bug we hit when state was lambda-local +// (T4096_D4096_I14336_E16_K4: dA off by 10% on expert 7 when num_iter % +// (2*Stages) != 0). + +template +__device__ __forceinline__ cutlass::PipelineState +mlp_bwd_resume_state(uint32_t count, bool is_producer) { + auto s = is_producer + ? cutlass::make_producer_start_state() + : cutlass::PipelineState{}; + s.advance(count); + return s; +} + +template +__device__ __forceinline__ void mlp_bwd_run_phase_1a( + Mlp1MainloopPipelineFor& pa_pipe, + uint32_t& pa_count, + Mlp1FusedActSmem& smem_mlp1, + TmaLoadX const& tma_load_x, + TmaLoadW1 const& tma_load_b_fwd, + TmaLoadW1 const& tma_load_c_fwd, + TmaStoreU const& tma_store_du, + TmaStoreV const& tma_store_dv, + TmaStoreZ const& tma_store_z, + int warp_id, int m, int expert, + const MlpBwdDims& dims, + uint32_t tmem_base, + int split_idx = -1, int num_splits = -1) { + int ns = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int nc = (num_splits >= 0) ? num_splits : (int)gridDim.y; + auto s = mlp_bwd_resume_state>( + pa_count, warp_id == 0); + if (warp_id == 0) { + mlp1_fused_act_producer(pa_pipe, s, + smem_mlp1, tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + m, (Compute == 90) ? expert : expert * dims.num_n_tiles_1, + dims.num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.experts_per_pe, dims.total_n_rows_1, + dims.num_n_tiles_1, dims.num_k_tiles_1, + ns, nc); + if constexpr (Compute == 100) + pa_pipe.producer_tail(s); + } + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits1::WarpSize == 0) + smem_mlp1.tmem_base = tmem_base; + constexpr int kFirstMlp12ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp12ConsumerWarp) + mlp1_fused_act_consumer(pa_pipe, s, + smem_mlp1, tma_store_du, tma_store_dv, tma_store_z, + m, dims.intermediate_dim, dims.num_m_tiles, + dims.num_n_tiles_1, dims.num_k_tiles_1, + ns, nc); + pa_count = s.count(); +} + +template +__device__ __forceinline__ void mlp_bwd_run_phase_1b( + Mlp2TMainloopPipelineFor& pb_pipe, + uint32_t& pb_count, + Mlp2TFusedSmem& smem_mlp2t, + TmaLoadDY const& tma_load_dy, + TmaLoadW2T const& tma_load_a_col, + TmaStoreDZ const& tma_store_dz, + int warp_id, int m, int expert, + const MlpBwdDims& dims, + uint32_t tmem_base, + int split_idx = -1, int num_splits = -1) { + int ns = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int nc = (num_splits >= 0) ? num_splits : (int)gridDim.y; + auto s = mlp_bwd_resume_state>( + pb_count, warp_id == 0); + if (warp_id == 0) { + mlp2_t_fused_producer(pb_pipe, s, + smem_mlp2t, tma_load_dy, tma_load_a_col, + m, expert, + dims.num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.experts_per_pe, dims.total_k_cols_2t, + dims.num_n_tiles_2t, dims.num_k_tiles_2t, + ns, nc); + if constexpr (Compute == 100) + pb_pipe.producer_tail(s); + } + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits2T::WarpSize == 0) + smem_mlp2t.tmem_base = tmem_base; + constexpr int kFirstMlp12ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp12ConsumerWarp) + mlp2_t_fused_consumer(pb_pipe, s, + smem_mlp2t, tma_store_dz, + m, dims.intermediate_dim, dims.num_m_tiles, + dims.num_n_tiles_2t, dims.num_k_tiles_2t, + ns, nc); + pb_count = s.count(); +} + +template +__device__ __forceinline__ void mlp_bwd_run_phase_1d( + Mlp5MainloopPipelineFor& pd_pipe, + uint32_t& pd_count, + Mlp5Smem& smem_mlp5, + TmaLoadDU const& tma_load_du, + TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, + TmaLoadC const& tma_load_c_col, + TmaStoreDX const& tma_store_dx, + int warp_id, int m, int expert, + const MlpBwdDims& dims, + uint32_t tmem_base, + int split_idx = -1, int num_splits = -1) { + int ns = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int nc = (num_splits >= 0) ? num_splits : (int)gridDim.y; + auto s = mlp_bwd_resume_state>( + pd_count, warp_id == 0); + if (warp_id == 0) { + mlp5_fused_producer(pd_pipe, s, + smem_mlp5, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, + m, (Compute == 90) ? expert : expert * dims.num_k_tiles_5, + dims.num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.experts_per_pe, dims.total_k_cols_5, + dims.num_n_tiles_5, dims.num_k_tiles_5, + ns, nc); + if constexpr (Compute == 100) + pd_pipe.producer_tail(s); + } + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits5::WarpSize == 0) + smem_mlp5.tmem_base = tmem_base; + constexpr int kFirstMlp5ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp5ConsumerWarp) + mlp5_fused_consumer(pd_pipe, s, + smem_mlp5, tma_store_dx, + m, dims.hidden_dim, + dims.num_m_tiles, dims.num_n_tiles_5, dims.num_k_tiles_5, + ns, nc); + pd_count = s.count(); +} + +// ═══════════════════════════════════════════════════════════════════ +// Phase 2 per-batch build-ranges scan +// ═══════════════════════════════════════════════════════════════════ +// +// CTA 0 warp 0 fills (k_start, k_end) per local expert for the current batch. +// Empty experts stay at (INT_MAX, 0); the chunk loop skips them via +// `k_end <= k_start`. The caller follows this with a grid-wide barrier before +// any CTA consumes the ranges. Restricting reset+scan to one warp avoids a +// race where a late CTA reset erased ranges another CTA had already updated. +// +// Called once per (batch, phase). mlp3 and mlp4 use different +// expert_for_k_block arrays and strides, so each phase scans its own. + +// ring_kb (K-blocks in the staging ring): when > 0 the window [k_offset, +// k_offset+num_k_in_batch) is in UNWRAPPED coordinates, and the physical ring +// slot for the expert-id lookup is (k_offset+kb) % ring_kb — while the recorded +// k_start/k_end ranges stay UNWRAPPED so each expert's range is contiguous even +// when the grouped window straddles the ring wrap. ring_kb == 0 (default) means +// no wrap (the un-grouped / local-pass callers, whose window already fits). +template +__device__ __forceinline__ void mlp_bwd_phase_2_build_ranges( + int* k_starts, + int* k_ends, + const int* expert_for_k_block, + const int* valid_rows_for_tile, + int k_offset, + int num_k_in_batch, + int experts_per_pe, + int local_expert_start, + int efkb_stride, + int tile_k, + int warp_id, + int ring_kb = 0) { + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + if (blockIdx.x == 0 && blockIdx.y == 0 && warp_id == 0) { + int lane = threadIdx.x; + for (int e = lane; e < experts_per_pe; e += 32) { + k_starts[e] = INT_MAX; + k_ends[e] = 0; + } + __syncwarp(); + for (int kb = lane; kb < num_k_in_batch; kb += 32) { + int kb_abs = k_offset + kb; // unwrapped + int kb_phys = (ring_kb > 0) ? (kb_abs % ring_kb) : kb_abs; // ring slot + int tile = kb_phys / efkb_stride; + if (valid_rows_for_tile != nullptr) { + int row_in_tile = (kb_phys - tile * efkb_stride) * tile_k; + if (row_in_tile >= valid_rows_for_tile[tile]) + continue; + } + int eg = expert_for_k_block[tile]; + int el = eg - local_expert_start; + if (el >= 0 && el < experts_per_pe) { + atomicMin(&k_starts[el], kb_abs); + atomicMax(&k_ends[el], kb_abs + 1); + } + } + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); +} + +// ── mlp4 Phase-2 runner ── +// +// Unified over Compute: Compute == 90 (Hopper) drives the adaptive kMSplit +// ternary walk over plain per-CTA cells; Compute == 100 (Blackwell) always +// drives the SM100 ClusterM=2 paired-CTA path (fixed M-split — the paired +// UMMA atom only supports the shared-N/split-M cooperative layout — and +// cell = CTA-PAIR, not CTA). Calls the unified mlp4_producer / mlp4_consumer (mlp4.cuh), which do their own +// internal Compute == 100 dispatch to the paired-CTA multicast TMA path. +// Reuses the plain Mlp4FusedSmem smem type (see MlpFusedBwdSmem — +// mlp4.cuh's Mlp4FusedSmem is traits-compatible internally, no separate +// Compute-conditional smem needed here). +template +__device__ __forceinline__ void mlp_bwd_run_phase_2_mlp4( + Mlp4MainloopPipelineFor& p4_pipe, + uint32_t& p4_count, + Mlp4FusedSmem& smem_mlp4, + const int* k_starts, + const int* k_ends, + TmaLoadXT4 const& tma_load_xt4, + TmaLoaddUT4 const& tma_load_dut4, + TmaLoaddVT4 const& tma_load_dvt4, + TmaReduceDB const& tma_reduce_db, + TmaReduceDC const& tma_reduce_dc, + int warp_id, int bk_start, int bk_tiles, + int phase2_div, int phase2_mod, + const MlpBwdDims& dims, + uint32_t tmem_base, + int ring_kb = 0, + const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, + uint32_t* pair_init_phase = nullptr) { + static_assert(Compute == 90 || Compute == 100, + "mlp_bwd_run_phase_2_mlp4: Compute must be 90 (Hopper 1SM) or " + "100 (Blackwell ClusterM=2 paired-CTA) — only Compute == 100 " + "executes the paired-CTA (ClusterM/NumPairs) code below."); + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + auto s = mlp_bwd_resume_state>( + p4_count, warp_id == 0); + + int experts_per_pe = dims.experts_per_pe; + + // Batch's K-window in mlp4's TileK_4 grain. Producer/consumer + // intersect this with the global per-expert ranges in `k_starts`/ + // `k_ends` (filled once per pass in the bwd preamble). + int batch_kb_start = (bk_start * Traits1::TileM) / Traits4::TileK; + int batch_kb_end = batch_kb_start + + (bk_tiles * Traits1::TileM) / Traits4::TileK; + + // Flat cell index over all Phase 2 CTAs. cell_start = phase2_div·NSplit2 + // + phase2_mod = flat_id ∈ [0,num_blocks) (the div/mod from the caller + // recombine to the flat id), so coverage is complete and disjoint for + // ANY num_blocks — no NSplit2 divisibility. + int flat_id = phase2_div * NSplit2 + phase2_mod; + // Phase 2 uses ALL launched CTAs (flat grid-stride over cells), so the + // cell stride is the full launched count gridDim.x·gridDim.y = num_blocks + // on Compute == 90; Compute == 100 pairs CTAs (ClusterM=2), so the stride + // is in CTA-PAIRS instead. + int total_phase2_ctas = (int)gridDim.x * (int)gridDim.y; + // Size outer_split from the experts LIVE in this batch's K-window, NOT the + // global expert set. Sizing from experts_per_pe under-splits when only ~1 + // expert is live per batch (large T) — the dead experts inflate the sizing + // input, so outer_split stays small and the lone live expert's work spreads + // over too few CTAs. k_starts/k_ends are per-expert ranges in TileK grain. + int live_experts_4 = 0; + for (int e = 0; e < experts_per_pe; ++e) + if (k_starts[e] < batch_kb_end && k_ends[e] > batch_kb_start) live_experts_4++; + if (live_experts_4 < 1) live_experts_4 = 1; + + int cell_start, cell_stride, outer_split_4, k_split_4; + if constexpr (Compute == 100) { + // SM100 ClusterM=2 paired-CTA path: fixed M-split (no kMSplit + // ternary — the paired-CTA MMA atom only supports the shared-N/ + // split-M cooperative layout), cell = CTA-PAIR. + int num_clusters = total_phase2_ctas / Traits4::ClusterM; + int total_chunks = live_experts_4 * dims.num_n_tiles_4; + int walk_extent = dims.num_m_tiles_4; + outer_split_4 = 1; + { + int target_cells = 2 * num_clusters; + while (outer_split_4 * 2 * Traits4::NumPairs <= walk_extent + && total_chunks * outer_split_4 < target_cells) + outer_split_4 *= 2; + while (walk_extent % outer_split_4 != 0) + outer_split_4 /= 2; + if (outer_split_4 < 1) outer_split_4 = 1; + } + // K-split (KS): split each cell's K-loop across k_split_4 CTA-pairs + // when the live cells don't fill the grid. TMA_REDUCE_ADD sums + // partials. + k_split_4 = 1; + { + int active_cells = total_chunks * outer_split_4; + int batch_k_tiles = batch_kb_end - batch_kb_start; + constexpr int kKMin = 8; + int max_by_depth = batch_k_tiles / kKMin; + if (max_by_depth < 1) max_by_depth = 1; + if (active_cells > 0 && active_cells < num_clusters) { + k_split_4 = (num_clusters + active_cells - 1) / active_cells; + if (k_split_4 > max_by_depth) k_split_4 = max_by_depth; + if (k_split_4 < 1) k_split_4 = 1; + } + } + cell_start = flat_id / Traits4::ClusterM; + cell_stride = num_clusters; + } else { + // SM90 (Hopper) 1SM path. mlp4 holds the SHARED operand and walks + // the cooperation (split) axis (see mlp4.cuh): + // kMSplit=true (M-split): chunk = (e, n_tile), lane partitions num_m_tiles_4 + // kMSplit=false (N-split): chunk = (e, m_tile), lane partitions num_n_tiles_4 + // Adaptive split: largest pow2 divisor of walk_extent so total_cells + // covers all Phase 2 CTAs. + int total_chunks_4 = Traits4::kMSplit + ? (live_experts_4 * dims.num_n_tiles_4) + : (live_experts_4 * dims.num_m_tiles_4); + int walk_extent_4 = Traits4::kMSplit ? dims.num_m_tiles_4 : dims.num_n_tiles_4; + outer_split_4 = 1; + { + int target_cells = 2 * total_phase2_ctas; + while (outer_split_4 * 2 <= walk_extent_4 + && total_chunks_4 * outer_split_4 < target_cells) { + outer_split_4 *= 2; + } + while (walk_extent_4 % outer_split_4 != 0) outer_split_4 /= 2; + if (outer_split_4 < 1) outer_split_4 = 1; + } + // K-split (KS): split each cell's K-loop across k_split_4 CTAs when + // the live cells don't fill the grid. TMA_REDUCE_ADD sums partials. + k_split_4 = 1; + { + int active_cells = total_chunks_4 * outer_split_4; + int batch_k_tiles = batch_kb_end - batch_kb_start; + constexpr int kKMin = 8; + int max_by_depth = batch_k_tiles / kKMin; + if (max_by_depth < 1) max_by_depth = 1; + if (active_cells > 0 && active_cells < total_phase2_ctas) { + k_split_4 = (total_phase2_ctas + active_cells - 1) / active_cells; + if (k_split_4 > max_by_depth) k_split_4 = max_by_depth; + if (k_split_4 < 1) k_split_4 = 1; + } + } + cell_start = flat_id; + cell_stride = total_phase2_ctas; + } + + if (warp_id == 0) + mlp4_producer( + p4_pipe, s, + smem_mlp4, tma_load_xt4, tma_load_dut4, tma_load_dvt4, + k_starts, k_ends, experts_per_pe, + dims.intermediate_dim, dims.hidden_dim, dims.num_tokens, + dims.num_m_tiles_4, dims.num_n_tiles_4, outer_split_4, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split_4, ring_kb); + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits4::WarpSize == 0) + smem_mlp4.tmem_base = tmem_base; + constexpr int kFirstMlp4ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp4ConsumerWarp) + mlp4_consumer( + p4_pipe, s, + smem_mlp4, tma_reduce_db, tma_reduce_dc, + k_starts, k_ends, experts_per_pe, + dims.intermediate_dim, dims.hidden_dim, dims.total_m_rows_4, + dims.num_m_tiles_4, dims.num_n_tiles_4, outer_split_4, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split_4, + pair_init_barrier, pair_init_phase); + + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + p4_count = s.count(); +} + +// ── mlp3 Phase-2 runner ── +// +// Unified over Compute: Compute == 90 (Hopper) drives the adaptive kMSplit +// ternary walk over plain per-CTA cells; Compute == 100 (Blackwell) always +// drives the SM100 ClusterM=2 paired-CTA path (fixed M-split — the paired +// UMMA atom only supports the shared-N/split-M cooperative layout — and +// cell = CTA-PAIR, not CTA), and uses the distinct Mlp3FusedSmem2Sm smem +// layout (paired-CTA accumulator pipeline; see MlpFusedBwdSmem). Calls the +// unified mlp3_producer / mlp3_consumer +// (mlp3.cuh), which do their own internal Compute == 100 dispatch to the +// paired-CTA multicast TMA path. +template +__device__ __forceinline__ void mlp_bwd_run_phase_2_mlp3( + Mlp3MainloopPipelineFor& p3_pipe, + uint32_t& p3_count, + cute::conditional_t, Mlp3FusedSmem>& smem_mlp3, + const int* k_starts, + const int* k_ends, + TmaLoadDYT3 const& tma_load_dyt3, + TmaLoadZT3 const& tma_load_zt3, + TmaReduceDA const& tma_reduce_da, + int warp_id, int bk_start, int bk_tiles, + int phase2_div, int phase2_mod, + const MlpBwdDims& dims, + uint32_t tmem_base, + int ring_kb = 0, + const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, + uint32_t* pair_init_phase = nullptr) { + static_assert(Compute == 90 || Compute == 100, + "mlp_bwd_run_phase_2_mlp3: Compute must be 90 (Hopper 1SM) or " + "100 (Blackwell ClusterM=2 paired-CTA) — only Compute == 100 " + "executes the paired-CTA (ClusterM/NumPairs) code below."); + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + auto s = mlp_bwd_resume_state>( + p3_count, warp_id == 0); + + int experts_per_pe = dims.experts_per_pe; + + // Batch's K-window in mlp3's TileK_3 grain (see mlp4 helper above). + int batch_kb_start = (bk_start * Traits1::TileM) / Traits3::TileK; + int batch_kb_end = batch_kb_start + + (bk_tiles * Traits1::TileM) / Traits3::TileK; + + // Flat cell index over all Phase 2 CTAs. cell_start = phase2_div·NSplit2 + // + phase2_mod = flat_id ∈ [0,num_blocks) (the div/mod from the caller + // recombine to the flat id), so coverage is complete and disjoint for + // ANY num_blocks — no NSplit2 divisibility. + int flat_id = phase2_div * NSplit2 + phase2_mod; + // All launched CTAs (gridDim.x·gridDim.y = num_blocks) — see mlp4 helper. + int total_phase2_ctas = (int)gridDim.x * (int)gridDim.y; + // Size outer_split from the experts LIVE in this batch's K-window (see the + // mlp4 helper above for the rationale) — not the global experts_per_pe. + int live_experts_3 = 0; + for (int e = 0; e < experts_per_pe; ++e) + if (k_starts[e] < batch_kb_end && k_ends[e] > batch_kb_start) live_experts_3++; + if (live_experts_3 < 1) live_experts_3 = 1; + + int cell_start, cell_stride, outer_split_3, k_split_3; + if constexpr (Compute == 100) { + // SM100 ClusterM=2 paired-CTA path: fixed M-split, cell = CTA-PAIR. + int total_phase2_pairs = total_phase2_ctas / Traits3::ClusterM; + int total_chunks = live_experts_3 * dims.num_n_tiles_3; + outer_split_3 = 1; + { + int target_cells = 2 * total_phase2_pairs; + while (outer_split_3 * 2 * Traits3::NumPairs <= dims.num_m_tiles_3 + && total_chunks * outer_split_3 < target_cells) + outer_split_3 *= 2; + while (dims.num_m_tiles_3 % (outer_split_3 * Traits3::NumPairs) != 0) + outer_split_3 /= 2; + outer_split_3 = max(outer_split_3, 1); + } + // K-split (KS): split each cell's K-loop across k_split_3 CTA-pairs + // when the live cells don't fill the grid. TMA_REDUCE_ADD sums + // partials. + k_split_3 = 1; + { + int active_cells = total_chunks * outer_split_3; + int batch_k_tiles = batch_kb_end - batch_kb_start; + constexpr int kKMin = 8; + int max_by_depth = max(batch_k_tiles / kKMin, 1); + if (active_cells > 0 && active_cells < total_phase2_pairs) { + k_split_3 = (total_phase2_pairs + active_cells - 1) / active_cells; + k_split_3 = min(k_split_3, max_by_depth); + } + } + cell_start = flat_id / Traits3::ClusterM; + cell_stride = total_phase2_pairs; + } else { + // SM90 (Hopper) 1SM path. mlp3 holds the SHARED operand and walks + // the cooperation (split) axis: + // kMSplit=true : chunk = (e, n_tile), lane partitions num_m_tiles + // kMSplit=false : chunk = (e, m_tile), lane partitions num_n_tiles + // Adaptive split: pick largest pow2 divisor of walk_extent so + // total_cells covers all Phase 2 CTAs. + int total_chunks_3 = Traits3::kMSplit + ? (live_experts_3 * dims.num_n_tiles_3) + : (live_experts_3 * dims.num_m_tiles_3); + int walk_extent_3 = Traits3::kMSplit ? dims.num_m_tiles_3 : dims.num_n_tiles_3; + outer_split_3 = 1; + { + int target_cells = 2 * total_phase2_ctas; + while (outer_split_3 * 2 <= walk_extent_3 + && total_chunks_3 * outer_split_3 < target_cells) { + outer_split_3 *= 2; + } + while (walk_extent_3 % outer_split_3 != 0) outer_split_3 /= 2; + if (outer_split_3 < 1) outer_split_3 = 1; + } + // K-split (KS): when the live cells (total_chunks_3 × outer_split_3) + // don't fill the grid, split each cell's K-loop across k_split_3 + // CTAs and let the TMA_REDUCE_ADD epilogue sum the partials. + // Bounded so each slice keeps >= kKMin K-tiles (else the Stages-deep + // pipeline can't amortize its fill). + k_split_3 = 1; + { + int active_cells = total_chunks_3 * outer_split_3; + int batch_k_tiles = batch_kb_end - batch_kb_start; + constexpr int kKMin = 8; + int max_by_depth = batch_k_tiles / kKMin; + if (max_by_depth < 1) max_by_depth = 1; + if (active_cells > 0 && active_cells < total_phase2_ctas) { + k_split_3 = (total_phase2_ctas + active_cells - 1) / active_cells; + if (k_split_3 > max_by_depth) k_split_3 = max_by_depth; + if (k_split_3 < 1) k_split_3 = 1; + } + } + cell_start = flat_id; + cell_stride = total_phase2_ctas; + } + + if (warp_id == 0) + mlp3_producer( + p3_pipe, s, + smem_mlp3, tma_load_dyt3, tma_load_zt3, + k_starts, k_ends, experts_per_pe, + dims.hidden_dim, dims.intermediate_dim, dims.num_tokens, + dims.num_m_tiles_3, dims.num_n_tiles_3, outer_split_3, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split_3, ring_kb); + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits3::WarpSize == 0) + smem_mlp3.tmem_base = tmem_base; + constexpr int kFirstMlp3ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp3ConsumerWarp) + mlp3_consumer( + p3_pipe, s, + smem_mlp3, tma_reduce_da, + k_starts, k_ends, experts_per_pe, + dims.hidden_dim, dims.intermediate_dim, dims.total_n_rows_3, + dims.num_m_tiles_3, dims.num_n_tiles_3, outer_split_3, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split_3, + pair_init_barrier, pair_init_phase); + + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + p3_count = s.count(); +} + +// ═══════════════════════════════════════════════════════════════════ +// mlp_fused_bwd_run_batches — batch loop body extracted +// ═══════════════════════════════════════════════════════════════════ +// +// Inner per-pass batch loop. Caller (mlp_fused_bwd or mlp_fused_bwd_dual) +// owns pipeline construction + the persistent per-phase counts. This lets +// the dual entry share ONE set of pipelines / counts across the local +// and remote passes — they sit at the dual-function's stack scope and +// PTXAS can analyze both batch loops as sequential code in one CFG, +// reusing spill slots between them. +// +// Templated on TileIter / TMA types / EfkbStride{3,4} — gets inlined into +// the caller's body. When called from mlp_fused_bwd_dual twice, the two +// inlined copies appear as two adjacent code regions in the same outer +// function (vs two separate inlined instances of the entire mlp_fused_bwd, +// which is the current 2-call shape). + +template +__device__ __forceinline__ void mlp_fused_bwd_run_batches( + MlpFusedBwdSmem& smem, + Mlp1MainloopPipelineFor& pa_pipe, + Mlp2TMainloopPipelineFor& pb_pipe, + Mlp5MainloopPipelineFor& pd_pipe, + Mlp3MainloopPipelineFor& p3_pipe, + Mlp4MainloopPipelineFor& p4_pipe, + uint32_t& pa_count, uint32_t& pb_count, uint32_t& pd_count, + uint32_t& p3_count, uint32_t& p4_count, + TileIter& iter, + TmaLoadX const& tma_load_x, TmaLoadW1 const& tma_load_b_fwd, + TmaLoadW1 const& tma_load_c_fwd, TmaStoreZ const& tma_store_z, + TmaStoreDU const& tma_store_du, TmaStoreDV const& tma_store_dv, + TmaLoadDY const& tma_load_dy, TmaLoadW2T const& tma_load_a_col, + TmaStoreDZ const& tma_store_dz, + TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, + TmaStoreDX const& tma_store_dx, + TmaLoadDYT3 const& tma_load_dyt3, TmaLoadZT3 const& tma_load_zt3, + TmaReduceDA const& tma_reduce_da, + TmaLoadXT4 const& tma_load_xt4, + TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, + TmaReduceDB const& tma_reduce_db, TmaReduceDC const& tma_reduce_dc, + int efkb_stride_3, int efkb_stride_4, + const MlpBwdDims& dims, + const MlpBwdBufs& bufs, + MlpBwdCtaBarrierT& global_barrier, + MlpBwdCtaBarrierT& x_barrier, + int warp_id) { + + using Element = typename Traits1::Element; + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + + // Phase 2 per-expert K-range scan is done PER BATCH over the current + // batch's K-window (below), NOT once per pass. mlp3 and mlp4 share TileK, + // so a single per-batch scan (over the X-pipe expert-id array) feeds both. + // + // A once-per-pass global scan over [0, num_tokens/TileK) is only correct + // when each physical slot maps to a SINGLE expert for the whole pass — + // true for the local pass (fully-materialized buffer) and for a remote + // ring that never wraps (global_total ≤ L). When the remote staging ring + // wraps to a 2nd+ lap (global_total > L = MC·CommNumStages), a physical + // slot is REUSED for a different expert each lap, so a global min/max + // range mis-attributes the reused slots — corrupting the experts whose + // tiles straddle the lap-reuse boundary (the ring's first/last experts). + // Scanning the LIVE tile_expert_ids per batch (which holds this batch's + // residency, sync-guaranteed by Phase 1's acquire) attributes every slot + // correctly regardless of lap. This restores the pre-#102 per-batch scan + // semantics under #102's cooperative cell-walk mlp3/mlp4. + static_assert(Traits3::TileK == Traits4::TileK, + "Phase 2 shared k_starts/ends require Traits3::TileK == Traits4::TileK"); + // The Compute == 100 ClusterM=2 paired-CTA Phase-2 path issues tcgen05 + // 2x1SM UMMA and cross-CTA cluster-scope mbarriers that only exist on + // SM100. Compute is the SOLE selector (see mlp_bwd_run_phase_2_mlp3 / + // mlp_bwd_run_phase_2_mlp4 above) — catch a bogus Compute value here at + // compile time rather than as a cryptic pipeline/SharedStorage type + // mismatch further down. + static_assert(Compute == 90 || Compute == 100, + "mlp_fused_bwd_run_batches: Compute must be 90 (SM90/Hopper, 1SM " + "Phase-2) or 100 (SM100/Blackwell, ClusterM=2 paired-CTA Phase-2)."); + + for (int batch = 0; batch < dims.num_batches; ++batch) { + + int saved_m = -1, saved_expert = -1; + bool has_tile = false; + + // Phase 1: recompute + silu coefs, dZ via dY @ A, silu_bwd + if (iter.has_next()) { + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + + iter.acquire_src(); + auto tile = iter.next(); + int expert = tile.expert - dims.local_expert_start; + int m = tile.y_m; + + saved_m = m; + saved_expert = expert; + has_tile = true; + + mlp_bwd_run_phase_1a( + pa_pipe, pa_count, smem.mlp1, + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_du, tma_store_dv, tma_store_z, + warp_id, m, expert, dims, smem.tmem_base); + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + + iter.acquire_dy(); + mlp_bwd_run_phase_1b( + pb_pipe, pb_count, smem.mlp2t, + tma_load_dy, tma_load_a_col, tma_store_dz, + warp_id, m, expert, dims, smem.tmem_base); + x_barrier.wait(); + x_barrier.wait(); + + silu_bwd_pair_tile( + bufs.dz_buf, bufs.du_buf, bufs.dv_buf, + bufs.du_buf, bufs.dv_buf, + m, Traits1::TileM, Traits2T::TileN, dims.intermediate_dim, + blockIdx.y, gridDim.y); + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + } + + // Global barrier: Z, dU, dV all visible + global_barrier.wait(); + + // Phase 1d: dX = dU@B + dV@C + if (has_tile) { + iter.acquire_dst(); + mlp_bwd_run_phase_1d( + pd_pipe, pd_count, smem.mlp5, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, + tma_store_dx, + warp_id, saved_m, saved_expert, dims, smem.tmem_base); + __threadfence(); // system-scope dX flush for the NIC put is done by release_dst + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + iter.release_dst(threadIdx.x); + } + + // Phase 2: weight grads — mlp4 first, then mlp3 + int phase2_flat_cta = (int)blockIdx.x + (int)blockIdx.y * (int)gridDim.x; + int phase2_div = phase2_flat_cta / NSplit2; + int phase2_mod = phase2_flat_cta % NSplit2; + + // Per-batch K-window = the tiles Phase 1 delivered THIS batch. Phase 1 + // advances each of grid_x columns by one tile per batch, so batch b + // delivers global tiles [b·grid_x, (b+1)·grid_x). Phase 2 must reduce + // exactly those — striding by grid_x (= n_gemm/NSplit), NOT gridDim.x. + // Under the flat 1-D launch gridDim.x = num_blocks (= num_sms) ≠ grid_x, + // so using gridDim.x here over-strides and over-spans: in the remote + // pass it reads slots from other laps of the circular staging ring + // (mis-attributed tile_expert_ids → corrupt dA/dB/dC at the ring-wrap + // boundary experts). The local pass tolerated it (fully-materialized + // non-circular buffer, harmless no-op tail batches), which is why 1-PE + // passed and multi-PE failed. num_batches is already sized ceil(total / + // grid_x), so grid_x is the matching stride. + int bk_start_unwrapped = batch * dims.grid_x; + int bk_tiles = min(dims.grid_x, + dims.total_m_tiles_in_pass - bk_start_unwrapped); + int bk_start = bk_start_unwrapped % dims.num_m_tiles; + + // Per-batch per-expert K-range scan over THIS batch's window of the + // (possibly wrapped) ring. Reads the LIVE X-pipe expert-id array + // (tile_expert_ids_x in the remote pass; release_src is deferred past + // mlp4, so it stays valid here AND for mlp3 since the ranges are cached + // into k_starts/ends before the scan's barrier). bk_start is already + // wrapped mod num_m_tiles, so the window indexes physical ring slots — + // matching the kb coordinates the mlp3/mlp4 producers feed to TMA. + if (bk_tiles > 0) { + int batch_kb_off = (bk_start * Traits1::TileM) / Traits4::TileK; + int batch_kb_cnt = (bk_tiles * Traits1::TileM) / Traits4::TileK; + mlp_bwd_phase_2_build_ranges( + bufs.expert_k_starts, bufs.expert_k_ends, + bufs.expert_for_k_block_mlp4, bufs.valid_rows_mlp4, + batch_kb_off, batch_kb_cnt, + dims.experts_per_pe, dims.local_expert_start, + efkb_stride_4, Traits4::TileK, warp_id); + } + + // Cross-CTA barrier after CTA 0 builds the shared ranges. MLP4 must not + // read them before the scan completes, or producer/consumer tile counts + // can diverge and deadlock the pipeline. + if (bk_tiles > 0) + global_barrier.wait(); + + if (bk_tiles > 0) { + mlp_bwd_run_phase_2_mlp4< + Traits1, Traits4, NSplit2, Compute>( + p4_pipe, p4_count, smem.mlp4, + bufs.expert_k_starts, bufs.expert_k_ends, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, + tma_reduce_db, tma_reduce_dc, + warp_id, bk_start, bk_tiles, + phase2_div, phase2_mod, + dims, smem.tmem_base); + } + + global_barrier.wait(); + if (has_tile) + iter.release_src(threadIdx.x); + + if (bk_tiles > 0) { + mlp_bwd_run_phase_2_mlp3< + Traits1, Traits3, NSplit2, Compute>( + p3_pipe, p3_count, smem.mlp3, + bufs.expert_k_starts, bufs.expert_k_ends, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + warp_id, bk_start, bk_tiles, + phase2_div, phase2_mod, + dims, smem.tmem_base); + } + + global_barrier.wait(); + if (has_tile) + iter.release_dy(threadIdx.x); + } +} + +// ═══════════════════════════════════════════════════════════════════ +// mlp_fused_bwd — pipeline construction + batch loop (single pass) +// ═══════════════════════════════════════════════════════════════════ +// +// Generic over TileIter. Called once for local, once for remote. +// Mirrors forward mlp_fused_fwd: prologue → tile/batch loop. + +// Per-phase N-split parameters (autotuning surface). +// Phase 1's effective NSplits (1a, 1b', 1c, 1d) are tied to gridDim.y by +// construction: those phases iterate n-tiles via blockIdx.y/gridDim.y. To +// retune them, change the launch grid shape at the host level. +// Phase 2 (mlp3, mlp4) uses NSplit2 with virtual (vx, vy) remapping — it +// can differ from gridDim.y as long as gridDim.x * gridDim.y is divisible +// by NSplit2. +template +__device__ __forceinline__ void mlp_fused_bwd( + MlpFusedBwdSmem& smem, + TileIter& iter, + // Phase 1 TMA + TmaLoadX const& tma_load_x, TmaLoadW1 const& tma_load_b_fwd, + TmaLoadW1 const& tma_load_c_fwd, TmaStoreZ const& tma_store_z, + TmaStoreDU const& tma_store_du, TmaStoreDV const& tma_store_dv, + TmaLoadDY const& tma_load_dy, TmaLoadW2T const& tma_load_a_col, + TmaStoreDZ const& tma_store_dz, + TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, + TmaStoreDX const& tma_store_dx, + // Phase 2 TMA + TmaLoadDYT3 const& tma_load_dyt3, TmaLoadZT3 const& tma_load_zt3, + TmaReduceDA const& tma_reduce_da, + TmaLoadXT4 const& tma_load_xt4, + TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, + TmaReduceDB const& tma_reduce_db, TmaReduceDC const& tma_reduce_dc, + // Dims + buffers + const MlpBwdDims& dims, + const MlpBwdBufs& bufs, + // Target-based barriers — constructed ONCE by the caller in + // moe_fused_bwd and threaded through here so target counters + // accumulate monotonically across local/cross-phase/remote calls. + MlpBwdCtaBarrierT& global_barrier, + MlpBwdCtaBarrierT& x_barrier) { + + using Element = typename Traits1::Element; + int warp_id = threadIdx.x / Traits1::WarpSize; + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + // See mlp_fused_bwd_run_batches — Compute is the SOLE selector of the + // Compute == 100 ClusterM=2 paired-CTA Phase-2 path; the paired-CTA + // UMMA + cluster mbarrier rendezvous below is SM100-only. + static_assert(Compute == 90 || Compute == 100, + "mlp_fused_bwd: Compute must be 90 (SM90/Hopper, 1SM Phase-2) or " + "100 (SM100/Blackwell, ClusterM=2 paired-CTA Phase-2)."); + + // ── Prefetch TMA descriptors ── + cute::prefetch_tma_descriptor(tma_load_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_b_fwd.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c_fwd.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_du.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dv.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dy.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_a_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dz.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_du.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dv.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_b_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dx.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dyt3.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_zt3.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_da.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_xt4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dut4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dvt4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_db.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_dc.get_tma_descriptor()); + + // ── Prologue: construct pipelines ── + // Pipeline constructor requires __syncthreads (all 384 threads). + // Must happen before comm warps exit. All five phases (1a/1b/1d, + // 2-mlp3, 2-mlp4) use a single cooperative fused pipe each; phase + // 1b additionally uses a tiny scale-load aux pipe (Stages=1) to + // carry the silu-bwd coefficients U'/V' produced by mlp1. + auto pa_pipe = [&]() { + if constexpr (Compute == 100) + return mlp1_make_pipe_umma(smem.pa_pipe); + else + return mlp1_make_pipe(smem.pa_pipe); + }(); + auto pb_pipe = [&]() { + if constexpr (Compute == 100) + return mlp2_t_make_pipe_umma(smem.pb_pipe); + else + return mlp2_t_make_pipe(smem.pb_pipe); + }(); + auto pd_pipe = [&]() { + if constexpr (Compute == 100) + return mlp5_make_pipe_umma(smem.pd_pipe); + else + return mlp5_make_pipe(smem.pd_pipe); + }(); + + // Single fused pipe for mlp3 carrying dYT + Z per acquire. On SM100, + // warp 3 issues UMMA and epilogue WGs 4-11 drain the same output tile. + // Compute == 100 always drives the SM100 ClusterM=2 paired-CTA pipe + // (mlp3_make_pipe_umma_2sm); Compute == 90 drives the existing Hopper + // pipe (mlp3_make_pipe). + auto p3_pipe = [&]() { + if constexpr (Compute == 100) + return mlp3_make_pipe_umma_2sm(smem.p3_pipe); + else + return mlp3_make_pipe(smem.p3_pipe); + }(); + // Single fused pipe for mlp4 carrying X + dU^T + dV^T per acquire. On + // SM100, warp 3 issues the two UMMAs and epilogue WGs 4-11 reduce-add + // dB/dC. Compute == 100 always drives the SM100 ClusterM=2 paired-CTA + // pipe (mlp4_make_pipe_umma_2sm); Compute == 90 drives the existing + // Hopper pipe (mlp4_make_pipe). + auto p4_pipe = [&]() { + if constexpr (Compute == 100) + return mlp4_make_pipe_umma_2sm(smem.p4_pipe); + else + return mlp4_make_pipe(smem.p4_pipe); + }(); + + // Comm warps exit. Hopper keeps the historical MLP partition (warps 0,4-11) + // and drops warp 3 immediately. Blackwell mirrors the forward path: comm + // uses only warps 1-2 and warp 3 joins MLP as the dedicated UMMA producer. + if constexpr (Compute == 90) { + if (warp_id == 3) return; + } + if (warp_id == 1 || warp_id == 2) return; + + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + cute::conditional_t< + Compute == 100, + cute::TMEM::Allocator2Sm, + cute::TMEM::Allocator1Sm> tmem_alloc{}; + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + if (warp_id == 3) { + tmem_alloc.allocate(kTmemColumns, &smem.tmem_base); + if constexpr (Compute == 100) { + if (cute::elect_one_sync()) { + smem.tmem_pair_barrier.init(1); + smem.tmem_init_barrier.init(1); + smem.acc_pair_barrier.init(1); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + mlp_bwd_pair_barrier_sync(smem.tmem_init_barrier, warp_id); + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + } + + // Keep phase-role tests at their call sites. Hoisting producer/MMA/epilogue + // booleans across every batch keeps them live through all phase mainloops, + // increasing register pressure in a kernel that already spills. + + // ── Persistent per-phase pipeline counts ───────────────────── + // + // We persist only the monotonic iteration count per phase (1 uint32_t) + // rather than the full PipelineState (3 ints: index_, phase_, count_). + // Each phase function reconstructs the live state on entry by calling + // `make_producer_start_state().advance(count)` (or the consumer + // equivalent) — this is correct because `advance(N)` deterministically + // computes the (index, phase) tuple that ++state would arrive at after + // N iterations, and that tuple matches the in-smem barrier phase the + // previous invocation left behind. + // + // PipelineState itself must NOT be persistent across phase invocations + // if we want construction inside the phase — but persisting only the + // count (5 regs across 5 phases vs 15 regs for full states) trims the + // cross-phase live set while keeping correctness. We hit a desync bug + // on T4096_D4096_I14336_E16_K4 (dA off 10% on expert 7) when state was + // reset fresh per batch without preserving the count. + uint32_t pa_count = 0; + uint32_t pb_count = 0; + uint32_t pd_count = 0; + uint32_t p3_count = 0; + uint32_t p4_count = 0; + + mlp_fused_bwd_run_batches< + Traits1, Traits2T, Traits3, Traits4, Traits5, + NSplit2, Compute>( + smem, pa_pipe, pb_pipe, pd_pipe, p3_pipe, p4_pipe, + pa_count, pb_count, pd_count, p3_count, p4_count, + iter, + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_du, tma_store_dv, + tma_load_dy, tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, + EfkbStride3, EfkbStride4, + dims, bufs, global_barrier, x_barrier, warp_id); + + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + if (warp_id == 3) { + tmem_alloc.release_allocation_lock(); + mlp_bwd_pair_barrier_sync(smem.tmem_pair_barrier, warp_id); + tmem_alloc.free(smem.tmem_base, kTmemColumns); + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// mlp_fused_bwd_dual — single-instance local + remote passes +// ═══════════════════════════════════════════════════════════════════ +// +// Builds pipelines + per-phase counts ONCE, runs the local batch loop, +// crosses the global barrier, and runs the remote batch loop in the +// SAME function body. The aim is to let PTXAS analyze both batch loops +// as sequential code in one CFG, reusing spill slots between them +// (lifetimes don't overlap — separated by global_barrier.wait()). This +// is the alternative to the current pattern in moe_fused_bwd of two +// inlined `mlp_fused_bwd` calls, each constructing its own pipelines +// and re-initializing the smem barriers. +// +// Pipeline barriers in smem are constructed ONCE here. The persistent +// `*_count` state is shared across passes and consumed by the count→ +// PipelineState resume logic inside each phase helper (see +// mlp_bwd_resume_state) — this matches the smem barriers' phase +// counters that advance across both passes. +template 1 → Phase-1 (mlp1a/mlp2t/ + // silu/mlp5) steps SubTiles GemmTileM-row sub-tiles per comm slot; + // Phase-2 sees num_m_tiles·SubTiles GemmTileM token-tiles with + // tile_expert_ids replicated SubTiles× (the K-block count is invariant). + // 1 → unchanged single-tile path. + int SubTiles = 1, + int Compute = 90, + typename FusedIter, + // Shared TMA (Phase 1 weights, Phase 1d B/C, Phase 2 reduce + Z/U/V) + typename TmaLoadW1, typename TmaStoreZ, typename TmaStoreDU, + typename TmaStoreDV, typename TmaLoadW2T, typename TmaStoreDZ, + typename TmaLoadDU, typename TmaLoadDV, + typename TmaLoadB, typename TmaLoadC, + typename TmaLoadZT3, typename TmaReduceDA, + typename TmaLoaddUT4, typename TmaLoaddVT4, + typename TmaReduceDB, typename TmaReduceDC, + // Per-role TMA types — local and remote instances share the + // same C++ type (same cute Layout: dynamic shape, _1 col stride, + // identical smem layout). Only the runtime gmem pointer differs. + typename TmaLoadX, typename TmaLoadDY, typename TmaStoreDX, + typename TmaLoadDYT3, typename TmaLoadXT4> +__device__ __forceinline__ void mlp_fused_bwd_dual( + MlpFusedBwdSmem& smem, + FusedIter& iter, + bool remote_active, + int efkb_stride_4_remote, + // Phase 1 TMA — weights are shared across tiles; X/dY/dX (the *_remote + // instances below) come from the comm staging ring, where EVERY tile, + // local or remote, is staged. + TmaLoadW1 const& tma_load_b_fwd, TmaLoadW1 const& tma_load_c_fwd, + TmaStoreZ const& tma_store_z, + TmaStoreDU const& tma_store_du, TmaStoreDV const& tma_store_dv, + TmaLoadW2T const& tma_load_a_col, TmaStoreDZ const& tma_store_dz, + TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, + // Phase 2 TMA (shared) + TmaLoadZT3 const& tma_load_zt3, TmaReduceDA const& tma_reduce_da, + TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, + TmaReduceDB const& tma_reduce_db, TmaReduceDC const& tma_reduce_dc, + // Staging-ring data TMAs: X/dY/dX (Phase 1) + dYᵀ/Xᵀ (Phase 2). + TmaLoadX const& tma_load_x_remote, + TmaLoadDY const& tma_load_dy_remote, + TmaStoreDX const& tma_store_dx_remote, + TmaLoadDYT3 const& tma_load_dyt3_remote, + TmaLoadXT4 const& tma_load_xt4_remote, + const MlpBwdDims& remote_dims, + const MlpBwdBufs& remote_bufs, + MlpBwdCtaBarrierT& global_barrier, + MlpBwdCtaBarrierT& x_barrier, + // Flat 1-D launch coordinates (computed once in moe_fused_bwd): + // flat_id = blockIdx.x + // col = flat_id / NSplit (Phase 1 logical column) + // grid_x = n_gemm / NSplit (Phase 1 column count) + // split = flat_id % NSplit (N-split index) + // num_splits = NSplit + // gemm_active = flat_id < n_gemm → this CTA runs Phase 1 + comm; + // gap CTAs skip Phase 1 but still run Phase 2 + barriers. + int flat_id, int col, int grid_x, int split, int num_splits, + bool gemm_active) { + + static_assert(SubBatch >= 1 && SubBatch <= kBwdSubBatchMax, + "SubBatch (MoeBwdConfig::kSubBatch) must be in [1, kBwdSubBatchMax]"); + // See mlp_fused_bwd_run_batches — Compute is the SOLE selector of the + // Compute == 100 ClusterM=2 paired-CTA Phase-2 path; the paired-CTA + // UMMA + cluster mbarrier rendezvous below is SM100-only. + static_assert(Compute == 90 || Compute == 100, + "mlp_fused_bwd_dual: Compute must be 90 (SM90/Hopper, 1SM Phase-2) " + "or 100 (SM100/Blackwell, ClusterM=2 paired-CTA Phase-2)."); + + int warp_id = threadIdx.x / Traits1::WarpSize; + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + + // Prefetch the TMA descriptors this pass uses: shared weights/intermediates + // + the staging-ring data descriptors. + cute::prefetch_tma_descriptor(tma_load_b_fwd.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c_fwd.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_du.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dv.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_a_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dz.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_du.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dv.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_b_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_zt3.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_da.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dut4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dvt4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_db.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_dc.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_x_remote.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dy_remote.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dx_remote.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dyt3_remote.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_xt4_remote.get_tma_descriptor()); + + // Build pipelines ONCE (smem barriers initialized in one place — the + // remote pass picks up smem-barrier phase from where local left off). + auto pa_pipe = [&]() { + if constexpr (Compute == 100) + return mlp1_make_pipe_umma(smem.pa_pipe); + else + return mlp1_make_pipe(smem.pa_pipe); + }(); + auto pb_pipe = [&]() { + if constexpr (Compute == 100) + return mlp2_t_make_pipe_umma(smem.pb_pipe); + else + return mlp2_t_make_pipe(smem.pb_pipe); + }(); + auto pd_pipe = [&]() { + if constexpr (Compute == 100) + return mlp5_make_pipe_umma(smem.pd_pipe); + else + return mlp5_make_pipe(smem.pd_pipe); + }(); + + // Single fused pipe for mlp3 carrying dYT + Z per acquire. Compute == + // 100 always drives the SM100 ClusterM=2 paired-CTA pipe + // (mlp3_make_pipe_umma_2sm); Compute == 90 drives the existing Hopper + // pipe (mlp3_make_pipe). + auto p3_pipe = [&]() { + if constexpr (Compute == 100) + return mlp3_make_pipe_umma_2sm(smem.p3_pipe); + else + return mlp3_make_pipe(smem.p3_pipe); + }(); + // Single fused pipe for mlp4 carrying X + dU^T + dV^T per acquire. + // Compute == 100 always drives the SM100 ClusterM=2 paired-CTA pipe + // (mlp4_make_pipe_umma_2sm); Compute == 90 drives the existing Hopper + // pipe (mlp4_make_pipe). + auto p4_pipe = [&]() { + if constexpr (Compute == 100) + return mlp4_make_pipe_umma_2sm(smem.p4_pipe); + else + return mlp4_make_pipe(smem.p4_pipe); + }(); + + if constexpr (Compute == 90) { + if (warp_id == 3) return; + } + if (warp_id == 1 || warp_id == 2) return; + + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + cute::conditional_t< + Compute == 100, + cute::TMEM::Allocator2Sm, + cute::TMEM::Allocator1Sm> tmem_alloc{}; + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + if (warp_id == 3) { + tmem_alloc.allocate(kTmemColumns, &smem.tmem_base); + if constexpr (Compute == 100) { + if (cute::elect_one_sync()) { + smem.tmem_pair_barrier.init(1); + smem.tmem_init_barrier.init(1); + smem.acc_pair_barrier.init(1); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + mlp_bwd_pair_barrier_sync(smem.tmem_init_barrier, warp_id); + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + } + + // Counts persist across both passes — smem barriers continue advancing + // from local into remote, and PipelineState::advance(count) re-derives + // the right (index, phase) for the resume. + uint32_t pa_count = 0; + uint32_t pb_count = 0; + uint32_t pd_count = 0; + uint32_t p3_count = 0; + uint32_t p4_count = 0; + uint32_t acc_pair_phase = 0; + + using Element = typename Traits1::Element; + + // Phase 2 flat-cell coords are CTA-invariant — hoist out of the loops. + int phase2_div = flat_id / NSplit2; + int phase2_mod = flat_id % NSplit2; + + // Shared k_starts/k_ends require equal TileK (one scan feeds mlp3+mlp4). + static_assert(Traits3::TileK == Traits4::TileK, + "Phase 2 shared k_starts/ends require Traits3::TileK == Traits4::TileK"); + + // Nothing to do if this PE has no tiles (remote_active is grid-uniform, so + // every CTA agrees and the early-out keeps the barriers below well-defined). + if (!remote_active) { + global_barrier.wait(); + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + if (warp_id == 3) { + tmem_alloc.release_allocation_lock(); + mlp_bwd_pair_barrier_sync( + smem.tmem_pair_barrier, warp_id); + tmem_alloc.free(smem.tmem_base, kTmemColumns); + } + } + return; + } + + // One cross-CTA sync so all MLP CTAs enter the grouped loop together: + // SM90 synchronizes w0+w4-11 (288 threads); SM100 synchronizes + // w0+w3-11 (320 threads). Comm warps 1-2 have already returned. + global_barrier.wait(); + + // ═══════════════════════════════════════════════════════════════ + // REMOTE GROUPED PASS — kBwdSubBatch Phase-1 sub-batches, then ONE + // Phase-2 (mlp4+mlp3) over their combined kBwdSubBatch·grid_x K-window. + // ═══════════════════════════════════════════════════════════════ + // Phase 1 (mlp1/2/5) runs in SMALL per-sub-batch chunks (less CTA + // barrier-stall); Phase 2 (weight grads dA/dB/dC) runs ONCE over the whole + // group so its cooperative cell-walk has a long K-reduction to fill the grid + // (deeper outer_split / k_split). Each sub-batch keeps its X/dY staging slot + // un-released until the grouped Phase 2 has read X^T (mlp4) and dY^T (mlp3); + // the saved slots are freed in a batch after each Phase-2 sub-pass. The host + // guarantees CommNumStages is a multiple of kBwdSubBatch and ≥ it, so the + // group window never straddles the ring wrap (per-expert k-ranges stay + // contiguous) and comm has prefetch headroom. kBwdSubBatch == 1 reduces this + // to the original per-batch loop. + const int remote_batches = remote_dims.num_batches; + + // Adaptive group size: a group of `g` sub-batches holds g·grid_x live staging + // slots that must map to DISTINCT ring slots, so g·grid_x ≤ ring. Cap the + // group by how many grid_x-windows the ring holds — deep rings get the full + // SubBatch; shallow rings degrade to fewer (down to 1 = original per-batch + // loop). Grid-uniform (num_m_tiles, grid_x are host-set), so every CTA agrees. + const int eff_subbatch = + max(1, min((int)SubBatch, remote_dims.num_m_tiles / grid_x)); + + // Phase-1 GEMM consumers + Phase-2 see the token axis in GemmTileM + // (=Traits1::TileM) tiles: num_m_tiles (Phase-1 buffer height) and num_tokens + // (M extent) scale by SubTiles = CommTileM/GemmTileM. The comm loop control + // (eff_subbatch, group_start, bk_tiles clamp, slot releases) stays in + // communication-slot units via remote_dims. At SubTiles=1 the two are identical. + // (Other dims fields — N/K extents, expert counts — are M-axis-independent.) + MlpBwdDims gemm_dims = remote_dims; + gemm_dims.num_m_tiles *= SubTiles; + gemm_dims.num_tokens *= SubTiles; + int first_sub = 0; + int last_sub = SubTiles; + int gemm_split = split; + int gemm_splits = num_splits; + constexpr int kM64PerGemmTile = Traits1::TileM / 64; + + for (int g0 = 0; g0 < remote_batches; g0 += eff_subbatch) { + const int gcount = min(eff_subbatch, remote_batches - g0); + + // Per-sub-batch staging slots, freed together after Phase 2. cur_slot in + // the iterator only tracks the LAST acquire, so releasing the earlier + // sub-batches needs their saved slot (release indexes *_consumed[slot]). + int saved_slot[SubBatch]; + bool saved_live[SubBatch]; + + // ── Phase 1: gcount sub-batches (mlp1a / mlp2t / silu_bwd / mlp5) ── + for (int s = 0; s < gcount; ++s) { + int saved_m = -1, saved_expert = -1; + int saved_valid_m64_subtiles = 0; + bool has_tile = false; + bool compute_tile = false; + + // Gated by gemm_active — gap CTAs (flat_id ≥ n_gemm) skip Phase 1 and + // the per-column x_barrier, but still hit every global barrier below. + if (gemm_active && iter.has_next()) { + iter.acquire_src(); + cutlass::arch::NamedBarrier::sync( + kMlpBwdBarrierThreads, kMlpBarrierId); + auto tile = iter.next(); + int expert = tile.expert - remote_dims.local_expert_start; + int m = tile.y_m; + + saved_m = m; + saved_expert = expert; + saved_valid_m64_subtiles = tile.valid_m64_subtiles; + has_tile = true; + compute_tile = + first_sub * kM64PerGemmTile < + saved_valid_m64_subtiles; + + // One CommTileM-row slot feeds SubTiles GemmTileM-row sub-tiles. + // gemm coordinate of sub-tile `sub` = m·SubTiles + sub (the + // Phase-1 kernels multiply by Traits1::TileM=GemmTileM → rows + // m·CommTileM + sub·GemmTileM). gemm_dims carries the SubTiles-scaled + // num_m_tiles / num_tokens so the buffer tensor shapes match. + // Acquire/release stays once per communication tile. + if (compute_tile) { + for (int sub = first_sub; sub < last_sub; ++sub) { + mlp_bwd_run_phase_1a( + pa_pipe, pa_count, smem.mlp1, + tma_load_x_remote, tma_load_b_fwd, tma_load_c_fwd, + tma_store_du, tma_store_dv, tma_store_z, + warp_id, m * SubTiles + sub, expert, gemm_dims, + smem.tmem_base, gemm_split, gemm_splits); + } + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + x_barrier.wait(); + + iter.acquire_dy(); + if (compute_tile) { + for (int sub = first_sub; sub < last_sub; ++sub) { + mlp_bwd_run_phase_1b( + pb_pipe, pb_count, smem.mlp2t, + tma_load_dy_remote, tma_load_a_col, tma_store_dz, + warp_id, m * SubTiles + sub, expert, gemm_dims, + smem.tmem_base, gemm_split, gemm_splits); + } + } + + x_barrier.wait(); + + if (compute_tile) { + for (int sub = first_sub; sub < last_sub; ++sub) { + silu_bwd_pair_tile( + remote_bufs.dz_buf, remote_bufs.du_buf, remote_bufs.dv_buf, + remote_bufs.du_buf, remote_bufs.dv_buf, + m * SubTiles + sub, Traits1::TileM, Traits2T::TileN, + remote_dims.intermediate_dim, gemm_split, gemm_splits); + } + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + } + + // Save this sub-batch's staging slot (== tile.y_m == iter cur_slot) + // for the deferred grouped release. -1 / false when this CTA had no + // tile (gap CTA or column exhausted). + saved_slot[s] = saved_m; + saved_live[s] = has_tile; + + // Global barrier: this sub-batch's Z, dU, dV visible to all CTAs. + global_barrier.wait(); + + // Phase 1d: dX = dU@B + dV@C (per sub-batch — Phase 2 never reads dX, + // so it is emitted and its dst slot released immediately). + if (has_tile) { + iter.acquire_dst(); + if (compute_tile) { + for (int sub = first_sub; sub < last_sub; ++sub) { + mlp_bwd_run_phase_1d( + pd_pipe, pd_count, smem.mlp5, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, + tma_store_dx_remote, + warp_id, saved_m * SubTiles + sub, saved_expert, gemm_dims, + smem.tmem_base, gemm_split, gemm_splits); + } + } + __threadfence(); // dX flush for the NIC put (release_dst signals) + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + iter.release_dst(threadIdx.x); + } + } + + // ── Phase 2: ONE pass over the group's gcount·grid_x K-window ── + // Phase 1 advances each of grid_x columns by one tile per sub-batch, so + // the group delivered global tiles [g0·grid_x, (g0+gcount)·grid_x). + // bk_start is UNWRAPPED (group_start) so the group's K-window stays + // contiguous even when it straddles the staging-ring wrap; build_ranges + // and the producers wrap the physical slot by % ring_kb internally + // (ring_kb > 0 enables it). The group's gcount·grid_x live slots are + // guaranteed distinct (host check: ring > SubBatch·grid_x). + // group_start / bk_tiles_comm are in communication-slot units. Phase 2 + // works in GemmTileM units, so scale the contiguous parent window by + // SubTiles; range construction skips padded K-blocks inside it. + int group_start = g0 * grid_x; + int bk_tiles_comm = min( + gcount * grid_x, + remote_dims.total_m_tiles_in_pass - group_start); + int bk_start = group_start * SubTiles; + int bk_tiles = bk_tiles_comm * SubTiles; + int ring_kb = (gemm_dims.num_m_tiles * Traits1::TileM) / Traits4::TileK; + const cutlass::arch::ClusterBarrier* acc_pair_barrier = nullptr; + if constexpr (Compute == 100) + acc_pair_barrier = &smem.acc_pair_barrier; + + // Per-group per-expert K-range scan over the group's window. Reads the + // LIVE X-pipe expert ids — none of the group's slots have been released + // yet, so all gcount·grid_x are valid. Feeds mlp3+mlp4 (unwrapped ranges). + if (bk_tiles > 0) { + int batch_kb_off = (bk_start * Traits1::TileM) / Traits4::TileK; + int batch_kb_cnt = (bk_tiles * Traits1::TileM) / Traits4::TileK; + mlp_bwd_phase_2_build_ranges( + remote_bufs.expert_k_starts, remote_bufs.expert_k_ends, + remote_bufs.expert_for_k_block_mlp4, + remote_bufs.valid_rows_mlp4, + batch_kb_off, batch_kb_cnt, + remote_dims.experts_per_pe, remote_dims.local_expert_start, + efkb_stride_4_remote, Traits4::TileK, warp_id, ring_kb); + // Cross-CTA barrier: CTA 0 owns reset+scan, so no producer may read + // the shared ranges until that work completes. + global_barrier.wait(); + } + + if (bk_tiles > 0) { + mlp_bwd_run_phase_2_mlp4< + Traits1, Traits4, NSplit2, Compute>( + p4_pipe, p4_count, smem.mlp4, + remote_bufs.expert_k_starts, remote_bufs.expert_k_ends, + tma_load_xt4_remote, tma_load_dut4, tma_load_dvt4, + tma_reduce_db, tma_reduce_dc, + warp_id, bk_start, bk_tiles, + phase2_div, phase2_mod, + gemm_dims, smem.tmem_base, ring_kb, + acc_pair_barrier, &acc_pair_phase); + } + + global_barrier.wait(); + // mlp4 done reading X^T — free every X slot the group held. + { + __threadfence(); + for (int s = 0; s < gcount; ++s) + if (saved_live[s]) iter.release_src_slot(saved_slot[s], threadIdx.x); + } + + if (bk_tiles > 0) { + mlp_bwd_run_phase_2_mlp3< + Traits1, Traits3, NSplit2, Compute>( + p3_pipe, p3_count, smem.mlp3, + remote_bufs.expert_k_starts, remote_bufs.expert_k_ends, + tma_load_dyt3_remote, tma_load_zt3, tma_reduce_da, + warp_id, bk_start, bk_tiles, + phase2_div, phase2_mod, + gemm_dims, smem.tmem_base, ring_kb, + acc_pair_barrier, &acc_pair_phase); + } + + global_barrier.wait(); + // mlp3 done reading dY^T — free every dY slot the group held. + { + __threadfence(); + for (int s = 0; s < gcount; ++s) + if (saved_live[s]) iter.release_dy_slot(saved_slot[s], threadIdx.x); + } + } + + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + if (warp_id == 3) { + tmem_alloc.release_allocation_lock(); + mlp_bwd_pair_barrier_sync(smem.tmem_pair_barrier, warp_id); + tmem_alloc.free(smem.tmem_base, kTmemColumns); + } + } +} + + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm90.cuh new file mode 100644 index 00000000..d9fb1eb4 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm90.cuh @@ -0,0 +1,1885 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP Backward — batch-interleaved Phase 1 + Phase 2 +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors forward mlp.cuh: pipeline construction + tile/batch loop. +// Called once for local tiles, once for remote tiles. +// +// Per batch (gridDim.x tiles): +// Phase 1: one tile per CTA (a→b→c→d sub-phases) +// Global barrier +// Phase 2: all CTAs do weight grads for this batch's K-blocks +// Global barrier +// +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp_sm90.cuh" // kMlpBarrierId, mlp1 pipe helpers +#include "mlp1_fused_act_sm90.cuh" // bwd Phase 1a — TileM=128 cooperative-M-split +#include "mlp2_t_fused_sm90.cuh" // bwd Phase 1b' — single-m-tile entry over mlp2_t.cuh +#include "mlp5_sm90.cuh" +#include "mlp5_fused_sm90.cuh" // mlp5_fused_producer/consumer (single-m-tile entry) +#include "mlp3_sm90.cuh" // Mlp3Traits, Mlp3FusedSmem +#include "mlp4_sm90.cuh" // Mlp4Traits, Mlp4FusedSmem +#include "silu_bwd_fused_sm90.cuh" // silu_bwd_pair_tile — separate elementwise phase +#include "cta_barrier_sm90.cuh" // target-based CtaCounterBarrier + +#include // INT_MAX for Phase 2 per-batch range scan + +namespace liger { + +using namespace cute; + +// MLP-warps barrier count for the bwd kernel. Hopper uses the historical +// 9-warp partition (w0 + w4-11). Blackwell mirrors the forward path: comm +// uses only warps 1-2 and warp 3 joins MLP as the dedicated UMMA producer. +template +static constexpr int kMlpBwdBarrierThreadsFor = (Compute == 100 ? 10 : 9) * 32; +static_assert(kMlpBwdBarrierThreadsFor<90> == 288); +static_assert(kMlpBwdBarrierThreadsFor<100> == 320); + +// Compute == 100 (Blackwell/SM100) always means the Phase-2 GEMMs (mlp3 AND +// mlp4) run the SM100 ClusterM=2 paired-CTA path; Compute == 90 (Hopper/SM90) +// always means plain 1SM. There is no per-phase or per-trait choice anymore +// (see mlp3.cuh / mlp4.cuh's unified mlp{3,4}_producer / +// mlp{3,4}_consumer), so `Compute == 100` alone drives the +// TMEM allocator choice (Allocator2Sm vs Allocator1Sm) and the extra +// tmem_pair_barrier below: both mlp3 and mlp4 share ONE TMEM allocation per +// CTA-pair, so the pair must rendezvous around the allocate/free calls. + +// ── Compute == 100 (paired-CTA) caller contract ── +// +// This header only builds the __device__ Phase-2 GEMM bodies; it does not +// launch the kernel or build TMA descriptors, so the following invariants +// must hold at the call site (host launcher, e.g. moe_bwd.cu) whenever +// Compute == 100 is selected — violating any of them is a silent-corruption +// or launch-failure risk, not something this file can assert at compile time: +// • Grid must be launched EVEN (gridDim.x·gridDim.y divisible by +// Traits{3,4}::ClusterM) — the flat cell-walk above (cell_start = +// flat_id / ClusterM, cell_stride = total_phase2_ctas / ClusterM) +// assumes every CTA has a live cluster peer. +// • Launch via cudaLaunchKernelEx with clusterDim = (ClusterM, 1, 1) (or +// the ClusterShape Traits{3,4} defines) and the +// cudaLaunchAttributeClusterDimension / +// NonPortableClusterSizeAllowed attribute set (ClusterM > 8 or any +// non-default shape needs the "non-portable" opt-in). +// • The dYT/Z (mlp3) and X/dU^T/dV^T (mlp4) TMA descriptors passed in +// (TmaLoadDYT3/TmaLoadZT3/... below) must be built with +// make_tma_copy_A_sm100 / make_tma_copy_B_sm100 (2-CTA-aware multicast +// descriptors) — a plain SM90 make_tma_copy descriptor does not carry +// the multicast-mask semantics mlp{3,4}_producer's Compute == 100 +// `copy(tma.with(barrier, mcast_mask), ...)` calls rely on. +// • MlpFusedBwdSmem's static size (larger for the Compute == 100 +// Mlp3FusedSmem2Sm variant) must fit within the dynamically-queried +// cudaDevAttrMaxSharedMemoryPerBlockOptin limit, and the kernel's +// cudaFuncAttributeMaxDynamicSharedMemorySize must be opted in to that +// size before launch — this header assumes the smem struct it defines +// already fits whatever the host allocated. +// See Compute == 100's use below (TMEM Allocator2Sm + tmem_pair_barrier) and +// the Compute == 90/100 static_asserts in mlp_fused_bwd_run_batches / +// mlp_fused_bwd / mlp_fused_bwd_dual for the invariants this file DOES +// enforce. + +// ═══════════════════════════════════════════════════════════════════ +// Phase-1 sub-batch grouping (mlp1/2/5 small chunks, mlp3/4 one big batch) +// ═══════════════════════════════════════════════════════════════════ +// +// Remote BWD groups kBwdSubBatch Phase-1 sub-batches before running ONE +// Phase-2 (mlp3/mlp4) over all of them at once. Phase 1 (mlp1/2/5) prefers +// SMALL per-batch chunks (less CTA barrier-stall); Phase 2 (weight grads, +// dA/dB/dC) prefers a LONG K-reduction so the cooperative cell-walk fills +// the grid (deeper outer_split / k_split). Grouping serves both. +// +// Buffer survival: the SubBatch sub-batches keep their X/dY comm-staging slots +// un-released until the grouped Phase 2 consumes them, and their Z/dU/dV scratch +// slots must not be overwritten. This needs the staging ring to be DEEPER than +// the group span — remote_num_m_tiles > SubBatch·grid_x — so the group's live +// slots fit with comm prefetch headroom (host-checked in moe_bwd.cu). No ring/ +// group ALIGNMENT is required: the grouped Phase-2 window wraps the ring exactly +// as the original per-batch loop does (a window that runs past the ring end is +// only correct when it does not actually straddle — i.e. global_total ≤ ring or +// grid_x·SubBatch divides the ring — same property the un-grouped loop has). +// SubBatch == 1 reduces the grouped loop EXACTLY to the original per-batch loop. +// +// The sub-batch count is a pure TEMPLATE ARGUMENT (SubBatch) on moe_fused_bwd / +// mlp_fused_bwd_dual, sourced from MoeBwdConfig::kSubBatch (whose default is a +// literal — see moe_bwd.cu). Distinct SubBatch instantiations coexist in one +// binary; there is no preprocessor knob. +static constexpr int kBwdSubBatchMax = 8; // upper bound for the per-group save arrays + +// ═══════════════════════════════════════════════════════════════════ +// Fused backward shared memory — union of all phases +// ═══════════════════════════════════════════════════════════════════ + +template +struct MlpFusedBwdSmem { + union { + Mlp1FusedActSmem mlp1; // TileM=128 cooperative-M-split + Mlp2TFusedSmem mlp2t; // TileM=128 cooperative-M-split (mlp2_t.cuh) + Mlp5Smem mlp5; // TileM=128 cooperative-M-split + // mlp4's fused smem is traits-conditional internally (see mlp4.cuh), + // so it needs no separate 2SM specialization here. mlp3 instead uses + // a wholly distinct 2SM smem layout (Mlp3FusedSmem2Sm — paired-CTA + // accumulator pipeline in place of mlp3's normal AccumulatorPipeline), + // selected purely by Compute == 100 (Blackwell always pairs mlp3's + // Phase-2 GEMM across CTAs; Compute == 90 never does). + cute::conditional_t< + Compute == 100, + Mlp3FusedSmem2Sm, + Mlp3FusedSmem> mlp3; + Mlp4FusedSmem mlp4; + }; + + // Phase 2 per-expert K-range scratch lives in gmem (MlpBwdBufs: + // expert_k_starts / expert_k_ends). Allocated to + // exactly `experts_per_pe` ints — no compile-time cap. + + // Phase 1a, 1b', 1d: single fused pipe each (1c silu_bwd is + // elementwise gmem→gmem, no pipe). + typename Mlp1MainloopPipelineFor::SharedStorage pa_pipe; + typename Mlp2TMainloopPipelineFor::SharedStorage pb_pipe; + typename Mlp5MainloopPipelineFor::SharedStorage pd_pipe; + + // Phase 2: single fused pipe each for mlp3 and mlp4. Outputs (dA, + // dB, dC) are accumulated via TMA_REDUCE_ADD to gmem — no prev-pipe + // needed. + typename Mlp3MainloopPipelineFor::SharedStorage p3_pipe; + typename Mlp4MainloopPipelineFor::SharedStorage p4_pipe; + alignas(16) uint32_t tmem_base; + // Rendezvous barrier for the paired-CTA TMEM allocation window + // (Compute == 100 only — one CTA-pair shares one TMEM allocation, so + // both peers must arrive before either frees it). Collapses to a plain + // uint32_t placeholder (no ClusterBarrier construction/arrive/wait) on + // Compute == 90, which never pairs CTAs. + alignas(16) cute::conditional_t< + Compute == 100, + cutlass::arch::ClusterBarrier, + uint32_t> tmem_pair_barrier; + alignas(16) cute::conditional_t< + Compute == 100, + cutlass::arch::ClusterBarrier, + uint32_t> tmem_init_barrier; + alignas(16) cute::conditional_t< + Compute == 100, + cutlass::arch::ClusterBarrier, + uint32_t> acc_pair_barrier; +}; + +__device__ __forceinline__ void mlp_bwd_pair_barrier_sync( + const cutlass::arch::ClusterBarrier& barrier, + int warp_id) { + if (warp_id == 3) { + if (cute::elect_one_sync()) { + uint32_t rank = cute::block_rank_in_cluster(); + uint32_t peer = rank ^ 1; + bool leader = (rank % 2) == 0; + barrier.arrive(peer, !leader); + barrier.wait(0); + barrier.arrive(peer, leader); + } + __syncwarp(); + } +} + +template +struct MlpBwdTmemColumns { + static constexpr int kMlp1Cols = Traits1::AccStages * (2 * Traits1::TileN); + static constexpr int kMlp2TCols = Traits2T::AccStages * Traits2T::TileN; + static constexpr int kMlp5Cols = Traits5::AccStages * Traits5::TileN; + static constexpr int kMlp3Cols = Traits3::AccStages * Traits3::TileN; + static constexpr int kMlp4Cols = Traits4::AccStages * (2 * Traits4::WgTileN); + + static constexpr int kPhase15Cols = (kMlp1Cols > kMlp5Cols) + ? kMlp1Cols : kMlp5Cols; + static constexpr int kPhase2T3Cols = (kMlp2TCols > kMlp3Cols) + ? kMlp2TCols : kMlp3Cols; + static constexpr int kPhase1235Cols = (kPhase15Cols > kPhase2T3Cols) + ? kPhase15Cols : kPhase2T3Cols; + static constexpr int value = (kPhase1235Cols > kMlp4Cols) + ? kPhase1235Cols : kMlp4Cols; + static_assert(value <= 512, "BWD SM100 TMEM allocation must fit one CTA's 512 columns"); +}; + +// ═══════════════════════════════════════════════════════════════════ +// MlpBwdDims — invariant backward dimensions (__grid_constant__) +// ═══════════════════════════════════════════════════════════════════ + +struct MlpBwdDims { + int hidden_dim; + int intermediate_dim; + + // Flat 1-D launch geometry (host-computed). The grid is launched as + // dim3(num_blocks) with num_blocks = num_sms, so gridDim is no longer + // the logical grid. Phase 1 + comm partition over the GEMM-active subset: + // n_gemm = floor_NS = largest multiple of NSplit ≤ num_sms + // grid_x = n_gemm / NSplit (Phase 1 column count) + // CTAs with flat_id ≥ n_gemm fall through Phase 1/comm but still run + // Phase 2 (grid-stride over all num_blocks CTAs) and every barrier. + int grid_x; + int n_gemm; + + // Phase 1a: recompute + int total_n_rows_1; + int num_n_tiles_1; + int num_k_tiles_1; + + // Phase 1b: dZ = dY @ A + int total_k_cols_2t; + int num_n_tiles_2t; + int num_k_tiles_2t; + + // Phase 1d: dX = dU@B + dV@C + int total_k_cols_5; + int num_n_tiles_5; + int num_k_tiles_5; + + // Phase 2: mlp3 (dA) + int num_m_tiles_3; + int num_n_tiles_3; + int total_n_rows_3; + + // Phase 2: mlp4 (dB, dC) + int num_m_tiles_4; + int num_n_tiles_4; + int total_m_rows_4; + + int* phase_counter; + + // Scalar params (in __grid_constant__ to avoid register pressure) + int num_tokens; + int num_m_tiles; // Phase 1 output buffer height (in m-tiles). + // Local: total tile capacity (= max_total_slots/TileM). + // Remote: gridDim.x * CommNumStages (the staging ring). + // Phase 2 wraps bk_start modulo this when computing + // the slot index into z_buf/du_buf/dv_buf — for the + // local pass that's a no-op; for the remote pass it + // brings batches ≥ NumStages back into the live ring. + int total_m_tiles_in_pass; // Phase 2 K-block iteration bound (in m-tiles). + // Local: equals num_m_tiles (existing semantics — Phase 2 + // walks the full padded buffer and skips fake K-blocks + // via the expert_for_k_block sentinel). + // Remote: equals smem.comm.global_total — strictly less + // than num_m_tiles when per_cta_tiles > NumStages, so + // batches that would walk past the ring still run Phase 2 + // (with bk_start wrapped) instead of being silently + // clamped to bk_tiles=0. + int num_batches; + int local_expert_start; + int num_pes; + int experts_per_pe; // Phase 3/4 use this to skip k-blocks whose expert + // falls outside [local_expert_start, + // local_expert_start + experts_per_pe). Required + // because expert_for_k_block (built post-sort) holds + // GLOBAL expert IDs spanning all PEs, but dA/dB/dC + // on this PE only have slots for the local subset. + + // Phase 2 virtual-grid block count along the non-NSplit2 axis: + // phase2_num_blocks = (gridDim.x * gridDim.y) / NSplit2. + // Host-known (depends only on launch grid + template constant NSplit2), + // so storing here lets every batch's mlp3/mlp4 call read it from grid + // constant memory instead of recomputing `gridDim.x * gridDim.y / + // NSplit2` per batch. + int phase2_num_blocks; + + // Remote-pass overrides for num_tokens / num_m_tiles. Both are host- + // known (gridDim.x * CommNumStages and that × Traits1::TileM); kept here + // so moe_fused_bwd's remote_dims build can just assign from grid constant + // instead of computing on device. The remote-only fields stay = 0 for + // the local pass (the local pass doesn't read them). + int remote_num_m_tiles; + int remote_num_tokens; + + // 1 when the bwd TMA-GET comm path is active (NVLink P2P + divisible + // hidden_dim + bounce smem fits). Selects the comm warp layout AND the + // remote MLP iterator's producer/consumer counts (1·NC get / 2·NC put + // vs 2·NC get / 1·NC put). Host-set; 0 keeps the original getmem layout. + int tma_get_enabled; + + // (efkb_stride_3 / efkb_stride_4 moved to template parameters on + // mlp_fused_bwd — see EfkbStride3/EfkbStride4 below. They were always + // compile-time constants per call site (local: 1, remote: + // Traits1::TileM/Traits{3,4}::TileK), so passing them as runtime ints + // just kept Traits1::TileM and the divides live in registers across the + // whole batch loop. Hoisting to templates lets the compiler constant- + // fold the `/1` (local) into a no-op and the remote ratios into a shift. +}; + +// ═══════════════════════════════════════════════════════════════════ +// MlpBwdBufs — buffer pointers (__grid_constant__) +// ═══════════════════════════════════════════════════════════════════ + +template +struct MlpBwdBufs { + // du_buf holds U' (silu-bwd coefficient) after mlp1, then dU after silu_bwd. + // dv_buf holds V' (silu-bwd coefficient) after mlp1, then dV after silu_bwd. + // dz_buf holds dZ between mlp2_t (Phase 1b') and silu_bwd (Phase 1c). + Element* du_buf; + Element* dv_buf; + Element* dz_buf; + + // Expert id per K-block, read by Phase 2's mlp4 and mlp3 respectively. + // They hold the same logical content — for K-block k, the local expert + // id of the tile that K-block belongs to — but the source array differs + // because each Phase-2 sub-pass has a different freshness requirement. + // + // expert_for_k_block_mlp4 is read by mlp4 (which reads X). In the + // remote pass it points at + // CommBuffersBwd::tile_expert_ids_x — comm + // won't overwrite it until MLP's release_src, + // which is deferred past mlp4. + // + // expert_for_k_block_mlp3 is read by mlp3 (which reads dY). In the + // remote pass it points at + // CommBuffersBwd::tile_expert_ids_dy — comm + // won't overwrite it until MLP's release_dy, + // which is deferred past mlp3. + // + // In the local pass and the standalone bwd kernel, both pointers alias + // the same post-sort expert_for_k_block built by build_expert_for_k_block + // — there is no comm race to defend against, and a single source is + // sufficient. + const int* expert_for_k_block_mlp4; + const int* expert_for_k_block_mlp3; + const int* valid_rows_mlp4; + + // Per-pass per-expert K-range scratch for Phase 2. Filled once at + // the start of each pass (in the bwd preamble) and consumed by both + // mlp3 and mlp4 batches — they require Traits3::TileK == Traits4:: + // TileK (the codebase already enforces this since a single + // expert_for_k_block indexes both), so one shared scan suffices. + // Sized exactly `experts_per_pe` ints each. + int* expert_k_starts; + int* expert_k_ends; + + Element const* gmem_da; + Element const* gmem_db; + Element const* gmem_dc; + int* barrier_counter; +}; + +// ═══════════════════════════════════════════════════════════════════ +// MLP-warps-only barriers (BWD) +// ═══════════════════════════════════════════════════════════════════ +// +// Both barriers use the target-based CtaCounterBarrier (see +// cta_barrier.cuh). The earlier modulo variants were susceptible to a +// CTA-skew deadlock where a fast CTA could push the counter past the +// `% stride == 0` window before slow CTAs polled. +// +// • MlpBwdGlobalBarrier: cross-CTA, stride = gridDim.x * gridDim.y. +// Counter is a single int in device memory. +// • MlpBwdXBarrier: per-blockIdx.x, stride = gridDim.y. +// Counter is a gmem array of length gridDim.x. +template +using MlpBwdCtaBarrierT = + CtaCounterBarrier, kMlpBarrierId>; + +// ═══════════════════════════════════════════════════════════════════ +// Per-phase inline runners +// ═══════════════════════════════════════════════════════════════════ +// +// Each phase of the bwd batch loop (1a / 1b' / 1d / 2-mlp4 / 2-mlp3) is +// pulled out into its own __forceinline__ device function. They take the +// per-phase PipelineState by REFERENCE — the caller (`mlp_fused_bwd`) +// holds the persistent state at function scope, and the inline body sees +// it as a direct register reference once inlined. Using actual functions +// instead of capture-all-by-reference lambdas avoids forcing the captured +// state into stack/local memory (CUDA closures pin captured vars to +// addressable storage, even when __forceinline__ would otherwise let +// them stay in registers). +// +// State persistence: the in-smem pipeline barriers carry a phase counter +// that advances each iteration and persists across phase invocations. To +// keep correctness with construction-inside-the-phase, the outer scope +// persists only a monotonic `uint32_t count` per phase (1 reg vs the +// 3-reg full PipelineState). The phase function: +// 1. Constructs a fresh PipelineState locally (warp-role-aware initial +// phase: producer starts phase=1, consumer starts phase=0). +// 2. Calls `advance(count)` to align it with the actual smem-barrier +// phase as left by the previous invocation. +// 3. Runs the phase, naturally advancing local state. +// 4. Writes `count = state.count()` back through the reference. +// This shrinks the cross-phase live set from 5×3 = 15 regs to 5×1 = 5 regs +// without re-introducing the desync bug we hit when state was lambda-local +// (T4096_D4096_I14336_E16_K4: dA off by 10% on expert 7 when num_iter % +// (2*Stages) != 0). + +template +__device__ __forceinline__ cutlass::PipelineState +mlp_bwd_resume_state(uint32_t count, bool is_producer) { + auto s = is_producer + ? cutlass::make_producer_start_state() + : cutlass::PipelineState{}; + s.advance(count); + return s; +} + +template +__device__ __forceinline__ void mlp_bwd_run_phase_1a( + Mlp1MainloopPipelineFor& pa_pipe, + uint32_t& pa_count, + Mlp1FusedActSmem& smem_mlp1, + TmaLoadX const& tma_load_x, + TmaLoadW1 const& tma_load_b_fwd, + TmaLoadW1 const& tma_load_c_fwd, + TmaStoreU const& tma_store_du, + TmaStoreV const& tma_store_dv, + TmaStoreZ const& tma_store_z, + int warp_id, int m, int expert, + const MlpBwdDims& dims, + uint32_t tmem_base, + int split_idx = -1, int num_splits = -1) { + int ns = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int nc = (num_splits >= 0) ? num_splits : (int)gridDim.y; + auto s = mlp_bwd_resume_state>( + pa_count, warp_id == 0); + if (warp_id == 0) { + mlp1_fused_act_producer(pa_pipe, s, + smem_mlp1, tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + m, (Compute == 90) ? expert : expert * dims.num_n_tiles_1, + dims.num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.experts_per_pe, dims.total_n_rows_1, + dims.num_n_tiles_1, dims.num_k_tiles_1, + ns, nc); + if constexpr (Compute == 100) + pa_pipe.producer_tail(s); + } + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits1::WarpSize == 0) + smem_mlp1.tmem_base = tmem_base; + constexpr int kFirstMlp12ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp12ConsumerWarp) + mlp1_fused_act_consumer(pa_pipe, s, + smem_mlp1, tma_store_du, tma_store_dv, tma_store_z, + m, dims.intermediate_dim, dims.num_m_tiles, + dims.num_n_tiles_1, dims.num_k_tiles_1, + ns, nc); + pa_count = s.count(); +} + +template +__device__ __forceinline__ void mlp_bwd_run_phase_1b( + Mlp2TMainloopPipelineFor& pb_pipe, + uint32_t& pb_count, + Mlp2TFusedSmem& smem_mlp2t, + TmaLoadDY const& tma_load_dy, + TmaLoadW2T const& tma_load_a_col, + TmaStoreDZ const& tma_store_dz, + int warp_id, int m, int expert, + const MlpBwdDims& dims, + uint32_t tmem_base, + int split_idx = -1, int num_splits = -1) { + int ns = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int nc = (num_splits >= 0) ? num_splits : (int)gridDim.y; + auto s = mlp_bwd_resume_state>( + pb_count, warp_id == 0); + if (warp_id == 0) { + mlp2_t_fused_producer(pb_pipe, s, + smem_mlp2t, tma_load_dy, tma_load_a_col, + m, expert, + dims.num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.experts_per_pe, dims.total_k_cols_2t, + dims.num_n_tiles_2t, dims.num_k_tiles_2t, + ns, nc); + if constexpr (Compute == 100) + pb_pipe.producer_tail(s); + } + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits2T::WarpSize == 0) + smem_mlp2t.tmem_base = tmem_base; + constexpr int kFirstMlp12ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp12ConsumerWarp) + mlp2_t_fused_consumer(pb_pipe, s, + smem_mlp2t, tma_store_dz, + m, dims.intermediate_dim, dims.num_m_tiles, + dims.num_n_tiles_2t, dims.num_k_tiles_2t, + ns, nc); + pb_count = s.count(); +} + +template +__device__ __forceinline__ void mlp_bwd_run_phase_1d( + Mlp5MainloopPipelineFor& pd_pipe, + uint32_t& pd_count, + Mlp5Smem& smem_mlp5, + TmaLoadDU const& tma_load_du, + TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, + TmaLoadC const& tma_load_c_col, + TmaStoreDX const& tma_store_dx, + int warp_id, int m, int expert, + const MlpBwdDims& dims, + uint32_t tmem_base, + int split_idx = -1, int num_splits = -1) { + int ns = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int nc = (num_splits >= 0) ? num_splits : (int)gridDim.y; + auto s = mlp_bwd_resume_state>( + pd_count, warp_id == 0); + if (warp_id == 0) { + mlp5_fused_producer(pd_pipe, s, + smem_mlp5, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, + m, (Compute == 90) ? expert : expert * dims.num_k_tiles_5, + dims.num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.experts_per_pe, dims.total_k_cols_5, + dims.num_n_tiles_5, dims.num_k_tiles_5, + ns, nc); + if constexpr (Compute == 100) + pd_pipe.producer_tail(s); + } + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits5::WarpSize == 0) + smem_mlp5.tmem_base = tmem_base; + constexpr int kFirstMlp5ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp5ConsumerWarp) + mlp5_fused_consumer(pd_pipe, s, + smem_mlp5, tma_store_dx, + m, dims.hidden_dim, + dims.num_m_tiles, dims.num_n_tiles_5, dims.num_k_tiles_5, + ns, nc); + pd_count = s.count(); +} + +// ═══════════════════════════════════════════════════════════════════ +// Phase 2 per-batch build-ranges scan +// ═══════════════════════════════════════════════════════════════════ +// +// CTA 0 warp 0 fills (k_start, k_end) per local expert for the current batch. +// Empty experts stay at (INT_MAX, 0); the chunk loop skips them via +// `k_end <= k_start`. The caller follows this with a grid-wide barrier before +// any CTA consumes the ranges. Restricting reset+scan to one warp avoids a +// race where a late CTA reset erased ranges another CTA had already updated. +// +// Called once per (batch, phase). mlp3 and mlp4 use different +// expert_for_k_block arrays and strides, so each phase scans its own. + +// ring_kb (K-blocks in the staging ring): when > 0 the window [k_offset, +// k_offset+num_k_in_batch) is in UNWRAPPED coordinates, and the physical ring +// slot for the expert-id lookup is (k_offset+kb) % ring_kb — while the recorded +// k_start/k_end ranges stay UNWRAPPED so each expert's range is contiguous even +// when the grouped window straddles the ring wrap. ring_kb == 0 (default) means +// no wrap (the un-grouped / local-pass callers, whose window already fits). +template +__device__ __forceinline__ void mlp_bwd_phase_2_build_ranges( + int* k_starts, + int* k_ends, + const int* expert_for_k_block, + const int* valid_rows_for_tile, + int k_offset, + int num_k_in_batch, + int experts_per_pe, + int local_expert_start, + int efkb_stride, + int tile_k, + int warp_id, + int ring_kb = 0) { + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + if (blockIdx.x == 0 && blockIdx.y == 0 && warp_id == 0) { + int lane = threadIdx.x; + for (int e = lane; e < experts_per_pe; e += 32) { + k_starts[e] = INT_MAX; + k_ends[e] = 0; + } + __syncwarp(); + for (int kb = lane; kb < num_k_in_batch; kb += 32) { + int kb_abs = k_offset + kb; // unwrapped + int kb_phys = (ring_kb > 0) ? (kb_abs % ring_kb) : kb_abs; // ring slot + int tile = kb_phys / efkb_stride; + if (valid_rows_for_tile != nullptr) { + int row_in_tile = (kb_phys - tile * efkb_stride) * tile_k; + if (row_in_tile >= valid_rows_for_tile[tile]) + continue; + } + int eg = expert_for_k_block[tile]; + int el = eg - local_expert_start; + if (el >= 0 && el < experts_per_pe) { + atomicMin(&k_starts[el], kb_abs); + atomicMax(&k_ends[el], kb_abs + 1); + } + } + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); +} + +// ── mlp4 Phase-2 runner ── +// +// Unified over Compute: Compute == 90 (Hopper) drives the adaptive kMSplit +// ternary walk over plain per-CTA cells; Compute == 100 (Blackwell) always +// drives the SM100 ClusterM=2 paired-CTA path (fixed M-split — the paired +// UMMA atom only supports the shared-N/split-M cooperative layout — and +// cell = CTA-PAIR, not CTA). Calls the unified mlp4_producer / mlp4_consumer (mlp4.cuh), which do their own +// internal Compute == 100 dispatch to the paired-CTA multicast TMA path. +// Reuses the plain Mlp4FusedSmem smem type (see MlpFusedBwdSmem — +// mlp4.cuh's Mlp4FusedSmem is traits-compatible internally, no separate +// Compute-conditional smem needed here). +template +__device__ __forceinline__ void mlp_bwd_run_phase_2_mlp4( + Mlp4MainloopPipelineFor& p4_pipe, + uint32_t& p4_count, + Mlp4FusedSmem& smem_mlp4, + const int* k_starts, + const int* k_ends, + TmaLoadXT4 const& tma_load_xt4, + TmaLoaddUT4 const& tma_load_dut4, + TmaLoaddVT4 const& tma_load_dvt4, + TmaReduceDB const& tma_reduce_db, + TmaReduceDC const& tma_reduce_dc, + int warp_id, int bk_start, int bk_tiles, + int phase2_div, int phase2_mod, + const MlpBwdDims& dims, + uint32_t tmem_base, + int ring_kb = 0, + const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, + uint32_t* pair_init_phase = nullptr) { + static_assert(Compute == 90 || Compute == 100, + "mlp_bwd_run_phase_2_mlp4: Compute must be 90 (Hopper 1SM) or " + "100 (Blackwell ClusterM=2 paired-CTA) — only Compute == 100 " + "executes the paired-CTA (ClusterM/NumPairs) code below."); + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + auto s = mlp_bwd_resume_state>( + p4_count, warp_id == 0); + + int experts_per_pe = dims.experts_per_pe; + + // Batch's K-window in mlp4's TileK_4 grain. Producer/consumer + // intersect this with the global per-expert ranges in `k_starts`/ + // `k_ends` (filled once per pass in the bwd preamble). + int batch_kb_start = (bk_start * Traits1::TileM) / Traits4::TileK; + int batch_kb_end = batch_kb_start + + (bk_tiles * Traits1::TileM) / Traits4::TileK; + + // Flat cell index over all Phase 2 CTAs. cell_start = phase2_div·NSplit2 + // + phase2_mod = flat_id ∈ [0,num_blocks) (the div/mod from the caller + // recombine to the flat id), so coverage is complete and disjoint for + // ANY num_blocks — no NSplit2 divisibility. + int flat_id = phase2_div * NSplit2 + phase2_mod; + // Phase 2 uses ALL launched CTAs (flat grid-stride over cells), so the + // cell stride is the full launched count gridDim.x·gridDim.y = num_blocks + // on Compute == 90; Compute == 100 pairs CTAs (ClusterM=2), so the stride + // is in CTA-PAIRS instead. + int total_phase2_ctas = (int)gridDim.x * (int)gridDim.y; + // Size outer_split from the experts LIVE in this batch's K-window, NOT the + // global expert set. Sizing from experts_per_pe under-splits when only ~1 + // expert is live per batch (large T) — the dead experts inflate the sizing + // input, so outer_split stays small and the lone live expert's work spreads + // over too few CTAs. k_starts/k_ends are per-expert ranges in TileK grain. + int live_experts_4 = 0; + for (int e = 0; e < experts_per_pe; ++e) + if (k_starts[e] < batch_kb_end && k_ends[e] > batch_kb_start) live_experts_4++; + if (live_experts_4 < 1) live_experts_4 = 1; + + int cell_start, cell_stride, outer_split_4, k_split_4; + if constexpr (Compute == 100) { + // SM100 ClusterM=2 paired-CTA path: fixed M-split (no kMSplit + // ternary — the paired-CTA MMA atom only supports the shared-N/ + // split-M cooperative layout), cell = CTA-PAIR. + int num_clusters = total_phase2_ctas / Traits4::ClusterM; + int total_chunks = live_experts_4 * dims.num_n_tiles_4; + int walk_extent = dims.num_m_tiles_4; + outer_split_4 = 1; + { + int target_cells = 2 * num_clusters; + while (outer_split_4 * 2 * Traits4::NumPairs <= walk_extent + && total_chunks * outer_split_4 < target_cells) + outer_split_4 *= 2; + while (walk_extent % outer_split_4 != 0) + outer_split_4 /= 2; + if (outer_split_4 < 1) outer_split_4 = 1; + } + // K-split (KS): split each cell's K-loop across k_split_4 CTA-pairs + // when the live cells don't fill the grid. TMA_REDUCE_ADD sums + // partials. + k_split_4 = 1; + { + int active_cells = total_chunks * outer_split_4; + int batch_k_tiles = batch_kb_end - batch_kb_start; + constexpr int kKMin = 8; + int max_by_depth = batch_k_tiles / kKMin; + if (max_by_depth < 1) max_by_depth = 1; + if (active_cells > 0 && active_cells < num_clusters) { + k_split_4 = (num_clusters + active_cells - 1) / active_cells; + if (k_split_4 > max_by_depth) k_split_4 = max_by_depth; + if (k_split_4 < 1) k_split_4 = 1; + } + } + cell_start = flat_id / Traits4::ClusterM; + cell_stride = num_clusters; + } else { + // SM90 (Hopper) 1SM path. mlp4 holds the SHARED operand and walks + // the cooperation (split) axis (see mlp4.cuh): + // kMSplit=true (M-split): chunk = (e, n_tile), lane partitions num_m_tiles_4 + // kMSplit=false (N-split): chunk = (e, m_tile), lane partitions num_n_tiles_4 + // Adaptive split: largest pow2 divisor of walk_extent so total_cells + // covers all Phase 2 CTAs. + int total_chunks_4 = Traits4::kMSplit + ? (live_experts_4 * dims.num_n_tiles_4) + : (live_experts_4 * dims.num_m_tiles_4); + int walk_extent_4 = Traits4::kMSplit ? dims.num_m_tiles_4 : dims.num_n_tiles_4; + outer_split_4 = 1; + { + int target_cells = 2 * total_phase2_ctas; + while (outer_split_4 * 2 <= walk_extent_4 + && total_chunks_4 * outer_split_4 < target_cells) { + outer_split_4 *= 2; + } + while (walk_extent_4 % outer_split_4 != 0) outer_split_4 /= 2; + if (outer_split_4 < 1) outer_split_4 = 1; + } + // K-split (KS): split each cell's K-loop across k_split_4 CTAs when + // the live cells don't fill the grid. TMA_REDUCE_ADD sums partials. + k_split_4 = 1; + { + int active_cells = total_chunks_4 * outer_split_4; + int batch_k_tiles = batch_kb_end - batch_kb_start; + constexpr int kKMin = 8; + int max_by_depth = batch_k_tiles / kKMin; + if (max_by_depth < 1) max_by_depth = 1; + if (active_cells > 0 && active_cells < total_phase2_ctas) { + k_split_4 = (total_phase2_ctas + active_cells - 1) / active_cells; + if (k_split_4 > max_by_depth) k_split_4 = max_by_depth; + if (k_split_4 < 1) k_split_4 = 1; + } + } + cell_start = flat_id; + cell_stride = total_phase2_ctas; + } + + if (warp_id == 0) + mlp4_producer( + p4_pipe, s, + smem_mlp4, tma_load_xt4, tma_load_dut4, tma_load_dvt4, + k_starts, k_ends, experts_per_pe, + dims.intermediate_dim, dims.hidden_dim, dims.num_tokens, + dims.num_m_tiles_4, dims.num_n_tiles_4, outer_split_4, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split_4, ring_kb); + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits4::WarpSize == 0) + smem_mlp4.tmem_base = tmem_base; + constexpr int kFirstMlp4ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp4ConsumerWarp) + mlp4_consumer( + p4_pipe, s, + smem_mlp4, tma_reduce_db, tma_reduce_dc, + k_starts, k_ends, experts_per_pe, + dims.intermediate_dim, dims.hidden_dim, dims.total_m_rows_4, + dims.num_m_tiles_4, dims.num_n_tiles_4, outer_split_4, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split_4, + pair_init_barrier, pair_init_phase); + + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + p4_count = s.count(); +} + +// ── mlp3 Phase-2 runner ── +// +// Unified over Compute: Compute == 90 (Hopper) drives the adaptive kMSplit +// ternary walk over plain per-CTA cells; Compute == 100 (Blackwell) always +// drives the SM100 ClusterM=2 paired-CTA path (fixed M-split — the paired +// UMMA atom only supports the shared-N/split-M cooperative layout — and +// cell = CTA-PAIR, not CTA), and uses the distinct Mlp3FusedSmem2Sm smem +// layout (paired-CTA accumulator pipeline; see MlpFusedBwdSmem). Calls the +// unified mlp3_producer / mlp3_consumer +// (mlp3.cuh), which do their own internal Compute == 100 dispatch to the +// paired-CTA multicast TMA path. +template +__device__ __forceinline__ void mlp_bwd_run_phase_2_mlp3( + Mlp3MainloopPipelineFor& p3_pipe, + uint32_t& p3_count, + cute::conditional_t, Mlp3FusedSmem>& smem_mlp3, + const int* k_starts, + const int* k_ends, + TmaLoadDYT3 const& tma_load_dyt3, + TmaLoadZT3 const& tma_load_zt3, + TmaReduceDA const& tma_reduce_da, + int warp_id, int bk_start, int bk_tiles, + int phase2_div, int phase2_mod, + const MlpBwdDims& dims, + uint32_t tmem_base, + int ring_kb = 0, + const cutlass::arch::ClusterBarrier* pair_init_barrier = nullptr, + uint32_t* pair_init_phase = nullptr) { + static_assert(Compute == 90 || Compute == 100, + "mlp_bwd_run_phase_2_mlp3: Compute must be 90 (Hopper 1SM) or " + "100 (Blackwell ClusterM=2 paired-CTA) — only Compute == 100 " + "executes the paired-CTA (ClusterM/NumPairs) code below."); + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + auto s = mlp_bwd_resume_state>( + p3_count, warp_id == 0); + + int experts_per_pe = dims.experts_per_pe; + + // Batch's K-window in mlp3's TileK_3 grain (see mlp4 helper above). + int batch_kb_start = (bk_start * Traits1::TileM) / Traits3::TileK; + int batch_kb_end = batch_kb_start + + (bk_tiles * Traits1::TileM) / Traits3::TileK; + + // Flat cell index over all Phase 2 CTAs. cell_start = phase2_div·NSplit2 + // + phase2_mod = flat_id ∈ [0,num_blocks) (the div/mod from the caller + // recombine to the flat id), so coverage is complete and disjoint for + // ANY num_blocks — no NSplit2 divisibility. + int flat_id = phase2_div * NSplit2 + phase2_mod; + // All launched CTAs (gridDim.x·gridDim.y = num_blocks) — see mlp4 helper. + int total_phase2_ctas = (int)gridDim.x * (int)gridDim.y; + // Size outer_split from the experts LIVE in this batch's K-window (see the + // mlp4 helper above for the rationale) — not the global experts_per_pe. + int live_experts_3 = 0; + for (int e = 0; e < experts_per_pe; ++e) + if (k_starts[e] < batch_kb_end && k_ends[e] > batch_kb_start) live_experts_3++; + if (live_experts_3 < 1) live_experts_3 = 1; + + int cell_start, cell_stride, outer_split_3, k_split_3; + if constexpr (Compute == 100) { + // SM100 ClusterM=2 paired-CTA path: fixed M-split, cell = CTA-PAIR. + int total_phase2_pairs = total_phase2_ctas / Traits3::ClusterM; + int total_chunks = live_experts_3 * dims.num_n_tiles_3; + outer_split_3 = 1; + { + int target_cells = 2 * total_phase2_pairs; + while (outer_split_3 * 2 * Traits3::NumPairs <= dims.num_m_tiles_3 + && total_chunks * outer_split_3 < target_cells) + outer_split_3 *= 2; + while (dims.num_m_tiles_3 % (outer_split_3 * Traits3::NumPairs) != 0) + outer_split_3 /= 2; + outer_split_3 = max(outer_split_3, 1); + } + // K-split (KS): split each cell's K-loop across k_split_3 CTA-pairs + // when the live cells don't fill the grid. TMA_REDUCE_ADD sums + // partials. + k_split_3 = 1; + { + int active_cells = total_chunks * outer_split_3; + int batch_k_tiles = batch_kb_end - batch_kb_start; + constexpr int kKMin = 8; + int max_by_depth = max(batch_k_tiles / kKMin, 1); + if (active_cells > 0 && active_cells < total_phase2_pairs) { + k_split_3 = (total_phase2_pairs + active_cells - 1) / active_cells; + k_split_3 = min(k_split_3, max_by_depth); + } + } + cell_start = flat_id / Traits3::ClusterM; + cell_stride = total_phase2_pairs; + } else { + // SM90 (Hopper) 1SM path. mlp3 holds the SHARED operand and walks + // the cooperation (split) axis: + // kMSplit=true : chunk = (e, n_tile), lane partitions num_m_tiles + // kMSplit=false : chunk = (e, m_tile), lane partitions num_n_tiles + // Adaptive split: pick largest pow2 divisor of walk_extent so + // total_cells covers all Phase 2 CTAs. + int total_chunks_3 = Traits3::kMSplit + ? (live_experts_3 * dims.num_n_tiles_3) + : (live_experts_3 * dims.num_m_tiles_3); + int walk_extent_3 = Traits3::kMSplit ? dims.num_m_tiles_3 : dims.num_n_tiles_3; + outer_split_3 = 1; + { + int target_cells = 2 * total_phase2_ctas; + while (outer_split_3 * 2 <= walk_extent_3 + && total_chunks_3 * outer_split_3 < target_cells) { + outer_split_3 *= 2; + } + while (walk_extent_3 % outer_split_3 != 0) outer_split_3 /= 2; + if (outer_split_3 < 1) outer_split_3 = 1; + } + // K-split (KS): when the live cells (total_chunks_3 × outer_split_3) + // don't fill the grid, split each cell's K-loop across k_split_3 + // CTAs and let the TMA_REDUCE_ADD epilogue sum the partials. + // Bounded so each slice keeps >= kKMin K-tiles (else the Stages-deep + // pipeline can't amortize its fill). + k_split_3 = 1; + { + int active_cells = total_chunks_3 * outer_split_3; + int batch_k_tiles = batch_kb_end - batch_kb_start; + constexpr int kKMin = 8; + int max_by_depth = batch_k_tiles / kKMin; + if (max_by_depth < 1) max_by_depth = 1; + if (active_cells > 0 && active_cells < total_phase2_ctas) { + k_split_3 = (total_phase2_ctas + active_cells - 1) / active_cells; + if (k_split_3 > max_by_depth) k_split_3 = max_by_depth; + if (k_split_3 < 1) k_split_3 = 1; + } + } + cell_start = flat_id; + cell_stride = total_phase2_ctas; + } + + if (warp_id == 0) + mlp3_producer( + p3_pipe, s, + smem_mlp3, tma_load_dyt3, tma_load_zt3, + k_starts, k_ends, experts_per_pe, + dims.hidden_dim, dims.intermediate_dim, dims.num_tokens, + dims.num_m_tiles_3, dims.num_n_tiles_3, outer_split_3, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split_3, ring_kb); + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits3::WarpSize == 0) + smem_mlp3.tmem_base = tmem_base; + constexpr int kFirstMlp3ConsumerWarp = (Compute == 100) ? 3 : 4; + if (warp_id >= kFirstMlp3ConsumerWarp) + mlp3_consumer( + p3_pipe, s, + smem_mlp3, tma_reduce_da, + k_starts, k_ends, experts_per_pe, + dims.hidden_dim, dims.intermediate_dim, dims.total_n_rows_3, + dims.num_m_tiles_3, dims.num_n_tiles_3, outer_split_3, + cell_start, cell_stride, + batch_kb_start, batch_kb_end, k_split_3, + pair_init_barrier, pair_init_phase); + + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + p3_count = s.count(); +} + +// ═══════════════════════════════════════════════════════════════════ +// mlp_fused_bwd_run_batches — batch loop body extracted +// ═══════════════════════════════════════════════════════════════════ +// +// Inner per-pass batch loop. Caller (mlp_fused_bwd or mlp_fused_bwd_dual) +// owns pipeline construction + the persistent per-phase counts. This lets +// the dual entry share ONE set of pipelines / counts across the local +// and remote passes — they sit at the dual-function's stack scope and +// PTXAS can analyze both batch loops as sequential code in one CFG, +// reusing spill slots between them. +// +// Templated on TileIter / TMA types / EfkbStride{3,4} — gets inlined into +// the caller's body. When called from mlp_fused_bwd_dual twice, the two +// inlined copies appear as two adjacent code regions in the same outer +// function (vs two separate inlined instances of the entire mlp_fused_bwd, +// which is the current 2-call shape). + +template +__device__ __forceinline__ void mlp_fused_bwd_run_batches( + MlpFusedBwdSmem& smem, + Mlp1MainloopPipelineFor& pa_pipe, + Mlp2TMainloopPipelineFor& pb_pipe, + Mlp5MainloopPipelineFor& pd_pipe, + Mlp3MainloopPipelineFor& p3_pipe, + Mlp4MainloopPipelineFor& p4_pipe, + uint32_t& pa_count, uint32_t& pb_count, uint32_t& pd_count, + uint32_t& p3_count, uint32_t& p4_count, + TileIter& iter, + TmaLoadX const& tma_load_x, TmaLoadW1 const& tma_load_b_fwd, + TmaLoadW1 const& tma_load_c_fwd, TmaStoreZ const& tma_store_z, + TmaStoreDU const& tma_store_du, TmaStoreDV const& tma_store_dv, + TmaLoadDY const& tma_load_dy, TmaLoadW2T const& tma_load_a_col, + TmaStoreDZ const& tma_store_dz, + TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, + TmaStoreDX const& tma_store_dx, + TmaLoadDYT3 const& tma_load_dyt3, TmaLoadZT3 const& tma_load_zt3, + TmaReduceDA const& tma_reduce_da, + TmaLoadXT4 const& tma_load_xt4, + TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, + TmaReduceDB const& tma_reduce_db, TmaReduceDC const& tma_reduce_dc, + int efkb_stride_3, int efkb_stride_4, + const MlpBwdDims& dims, + const MlpBwdBufs& bufs, + MlpBwdCtaBarrierT& global_barrier, + MlpBwdCtaBarrierT& x_barrier, + int warp_id) { + + using Element = typename Traits1::Element; + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + + // Phase 2 per-expert K-range scan is done PER BATCH over the current + // batch's K-window (below), NOT once per pass. mlp3 and mlp4 share TileK, + // so a single per-batch scan (over the X-pipe expert-id array) feeds both. + // + // A once-per-pass global scan over [0, num_tokens/TileK) is only correct + // when each physical slot maps to a SINGLE expert for the whole pass — + // true for the local pass (fully-materialized buffer) and for a remote + // ring that never wraps (global_total ≤ L). When the remote staging ring + // wraps to a 2nd+ lap (global_total > L = MC·CommNumStages), a physical + // slot is REUSED for a different expert each lap, so a global min/max + // range mis-attributes the reused slots — corrupting the experts whose + // tiles straddle the lap-reuse boundary (the ring's first/last experts). + // Scanning the LIVE tile_expert_ids per batch (which holds this batch's + // residency, sync-guaranteed by Phase 1's acquire) attributes every slot + // correctly regardless of lap. This restores the pre-#102 per-batch scan + // semantics under #102's cooperative cell-walk mlp3/mlp4. + static_assert(Traits3::TileK == Traits4::TileK, + "Phase 2 shared k_starts/ends require Traits3::TileK == Traits4::TileK"); + // The Compute == 100 ClusterM=2 paired-CTA Phase-2 path issues tcgen05 + // 2x1SM UMMA and cross-CTA cluster-scope mbarriers that only exist on + // SM100. Compute is the SOLE selector (see mlp_bwd_run_phase_2_mlp3 / + // mlp_bwd_run_phase_2_mlp4 above) — catch a bogus Compute value here at + // compile time rather than as a cryptic pipeline/SharedStorage type + // mismatch further down. + static_assert(Compute == 90 || Compute == 100, + "mlp_fused_bwd_run_batches: Compute must be 90 (SM90/Hopper, 1SM " + "Phase-2) or 100 (SM100/Blackwell, ClusterM=2 paired-CTA Phase-2)."); + + for (int batch = 0; batch < dims.num_batches; ++batch) { + + int saved_m = -1, saved_expert = -1; + bool has_tile = false; + + // Phase 1: recompute + silu coefs, dZ via dY @ A, silu_bwd + if (iter.has_next()) { + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + + iter.acquire_src(); + auto tile = iter.next(); + int expert = tile.expert - dims.local_expert_start; + int m = tile.y_m; + + saved_m = m; + saved_expert = expert; + has_tile = true; + + mlp_bwd_run_phase_1a( + pa_pipe, pa_count, smem.mlp1, + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_du, tma_store_dv, tma_store_z, + warp_id, m, expert, dims, smem.tmem_base); + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + + iter.acquire_dy(); + mlp_bwd_run_phase_1b( + pb_pipe, pb_count, smem.mlp2t, + tma_load_dy, tma_load_a_col, tma_store_dz, + warp_id, m, expert, dims, smem.tmem_base); + x_barrier.wait(); + x_barrier.wait(); + + silu_bwd_pair_tile( + bufs.dz_buf, bufs.du_buf, bufs.dv_buf, + bufs.du_buf, bufs.dv_buf, + m, Traits1::TileM, Traits2T::TileN, dims.intermediate_dim, + blockIdx.y, gridDim.y); + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + } + + // Global barrier: Z, dU, dV all visible + global_barrier.wait(); + + // Phase 1d: dX = dU@B + dV@C + if (has_tile) { + iter.acquire_dst(); + mlp_bwd_run_phase_1d( + pd_pipe, pd_count, smem.mlp5, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, + tma_store_dx, + warp_id, saved_m, saved_expert, dims, smem.tmem_base); + __threadfence(); // system-scope dX flush for the NIC put is done by release_dst + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + iter.release_dst(threadIdx.x); + } + + // Phase 2: weight grads — mlp4 first, then mlp3 + int phase2_flat_cta = (int)blockIdx.x + (int)blockIdx.y * (int)gridDim.x; + int phase2_div = phase2_flat_cta / NSplit2; + int phase2_mod = phase2_flat_cta % NSplit2; + + // Per-batch K-window = the tiles Phase 1 delivered THIS batch. Phase 1 + // advances each of grid_x columns by one tile per batch, so batch b + // delivers global tiles [b·grid_x, (b+1)·grid_x). Phase 2 must reduce + // exactly those — striding by grid_x (= n_gemm/NSplit), NOT gridDim.x. + // Under the flat 1-D launch gridDim.x = num_blocks (= num_sms) ≠ grid_x, + // so using gridDim.x here over-strides and over-spans: in the remote + // pass it reads slots from other laps of the circular staging ring + // (mis-attributed tile_expert_ids → corrupt dA/dB/dC at the ring-wrap + // boundary experts). The local pass tolerated it (fully-materialized + // non-circular buffer, harmless no-op tail batches), which is why 1-PE + // passed and multi-PE failed. num_batches is already sized ceil(total / + // grid_x), so grid_x is the matching stride. + int bk_start_unwrapped = batch * dims.grid_x; + int bk_tiles = min(dims.grid_x, + dims.total_m_tiles_in_pass - bk_start_unwrapped); + int bk_start = bk_start_unwrapped % dims.num_m_tiles; + + // Per-batch per-expert K-range scan over THIS batch's window of the + // (possibly wrapped) ring. Reads the LIVE X-pipe expert-id array + // (tile_expert_ids_x in the remote pass; release_src is deferred past + // mlp4, so it stays valid here AND for mlp3 since the ranges are cached + // into k_starts/ends before the scan's barrier). bk_start is already + // wrapped mod num_m_tiles, so the window indexes physical ring slots — + // matching the kb coordinates the mlp3/mlp4 producers feed to TMA. + if (bk_tiles > 0) { + int batch_kb_off = (bk_start * Traits1::TileM) / Traits4::TileK; + int batch_kb_cnt = (bk_tiles * Traits1::TileM) / Traits4::TileK; + mlp_bwd_phase_2_build_ranges( + bufs.expert_k_starts, bufs.expert_k_ends, + bufs.expert_for_k_block_mlp4, bufs.valid_rows_mlp4, + batch_kb_off, batch_kb_cnt, + dims.experts_per_pe, dims.local_expert_start, + efkb_stride_4, Traits4::TileK, warp_id); + } + + // Cross-CTA barrier after CTA 0 builds the shared ranges. MLP4 must not + // read them before the scan completes, or producer/consumer tile counts + // can diverge and deadlock the pipeline. + if (bk_tiles > 0) + global_barrier.wait(); + + if (bk_tiles > 0) { + mlp_bwd_run_phase_2_mlp4< + Traits1, Traits4, NSplit2, Compute>( + p4_pipe, p4_count, smem.mlp4, + bufs.expert_k_starts, bufs.expert_k_ends, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, + tma_reduce_db, tma_reduce_dc, + warp_id, bk_start, bk_tiles, + phase2_div, phase2_mod, + dims, smem.tmem_base); + } + + global_barrier.wait(); + if (has_tile) + iter.release_src(threadIdx.x); + + if (bk_tiles > 0) { + mlp_bwd_run_phase_2_mlp3< + Traits1, Traits3, NSplit2, Compute>( + p3_pipe, p3_count, smem.mlp3, + bufs.expert_k_starts, bufs.expert_k_ends, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + warp_id, bk_start, bk_tiles, + phase2_div, phase2_mod, + dims, smem.tmem_base); + } + + global_barrier.wait(); + if (has_tile) + iter.release_dy(threadIdx.x); + } +} + +// ═══════════════════════════════════════════════════════════════════ +// mlp_fused_bwd — pipeline construction + batch loop (single pass) +// ═══════════════════════════════════════════════════════════════════ +// +// Generic over TileIter. Called once for local, once for remote. +// Mirrors forward mlp_fused_fwd: prologue → tile/batch loop. + +// Per-phase N-split parameters (autotuning surface). +// Phase 1's effective NSplits (1a, 1b', 1c, 1d) are tied to gridDim.y by +// construction: those phases iterate n-tiles via blockIdx.y/gridDim.y. To +// retune them, change the launch grid shape at the host level. +// Phase 2 (mlp3, mlp4) uses NSplit2 with virtual (vx, vy) remapping — it +// can differ from gridDim.y as long as gridDim.x * gridDim.y is divisible +// by NSplit2. +template +__device__ __forceinline__ void mlp_fused_bwd( + MlpFusedBwdSmem& smem, + TileIter& iter, + // Phase 1 TMA + TmaLoadX const& tma_load_x, TmaLoadW1 const& tma_load_b_fwd, + TmaLoadW1 const& tma_load_c_fwd, TmaStoreZ const& tma_store_z, + TmaStoreDU const& tma_store_du, TmaStoreDV const& tma_store_dv, + TmaLoadDY const& tma_load_dy, TmaLoadW2T const& tma_load_a_col, + TmaStoreDZ const& tma_store_dz, + TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, + TmaStoreDX const& tma_store_dx, + // Phase 2 TMA + TmaLoadDYT3 const& tma_load_dyt3, TmaLoadZT3 const& tma_load_zt3, + TmaReduceDA const& tma_reduce_da, + TmaLoadXT4 const& tma_load_xt4, + TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, + TmaReduceDB const& tma_reduce_db, TmaReduceDC const& tma_reduce_dc, + // Dims + buffers + const MlpBwdDims& dims, + const MlpBwdBufs& bufs, + // Target-based barriers — constructed ONCE by the caller in + // moe_fused_bwd and threaded through here so target counters + // accumulate monotonically across local/cross-phase/remote calls. + MlpBwdCtaBarrierT& global_barrier, + MlpBwdCtaBarrierT& x_barrier) { + + using Element = typename Traits1::Element; + int warp_id = threadIdx.x / Traits1::WarpSize; + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + // See mlp_fused_bwd_run_batches — Compute is the SOLE selector of the + // Compute == 100 ClusterM=2 paired-CTA Phase-2 path; the paired-CTA + // UMMA + cluster mbarrier rendezvous below is SM100-only. + static_assert(Compute == 90 || Compute == 100, + "mlp_fused_bwd: Compute must be 90 (SM90/Hopper, 1SM Phase-2) or " + "100 (SM100/Blackwell, ClusterM=2 paired-CTA Phase-2)."); + + // ── Prefetch TMA descriptors ── + cute::prefetch_tma_descriptor(tma_load_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_b_fwd.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c_fwd.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_du.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dv.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dy.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_a_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dz.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_du.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dv.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_b_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dx.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dyt3.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_zt3.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_da.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_xt4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dut4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dvt4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_db.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_dc.get_tma_descriptor()); + + // ── Prologue: construct pipelines ── + // Pipeline constructor requires __syncthreads (all 384 threads). + // Must happen before comm warps exit. All five phases (1a/1b/1d, + // 2-mlp3, 2-mlp4) use a single cooperative fused pipe each; phase + // 1b additionally uses a tiny scale-load aux pipe (Stages=1) to + // carry the silu-bwd coefficients U'/V' produced by mlp1. + auto pa_pipe = [&]() { + if constexpr (Compute == 100) + return mlp1_make_pipe_umma(smem.pa_pipe); + else + return mlp1_make_pipe(smem.pa_pipe); + }(); + auto pb_pipe = [&]() { + if constexpr (Compute == 100) + return mlp2_t_make_pipe_umma(smem.pb_pipe); + else + return mlp2_t_make_pipe(smem.pb_pipe); + }(); + auto pd_pipe = [&]() { + if constexpr (Compute == 100) + return mlp5_make_pipe_umma(smem.pd_pipe); + else + return mlp5_make_pipe(smem.pd_pipe); + }(); + + // Single fused pipe for mlp3 carrying dYT + Z per acquire. On SM100, + // warp 3 issues UMMA and epilogue WGs 4-11 drain the same output tile. + // Compute == 100 always drives the SM100 ClusterM=2 paired-CTA pipe + // (mlp3_make_pipe_umma_2sm); Compute == 90 drives the existing Hopper + // pipe (mlp3_make_pipe). + auto p3_pipe = [&]() { + if constexpr (Compute == 100) + return mlp3_make_pipe_umma_2sm(smem.p3_pipe); + else + return mlp3_make_pipe(smem.p3_pipe); + }(); + // Single fused pipe for mlp4 carrying X + dU^T + dV^T per acquire. On + // SM100, warp 3 issues the two UMMAs and epilogue WGs 4-11 reduce-add + // dB/dC. Compute == 100 always drives the SM100 ClusterM=2 paired-CTA + // pipe (mlp4_make_pipe_umma_2sm); Compute == 90 drives the existing + // Hopper pipe (mlp4_make_pipe). + auto p4_pipe = [&]() { + if constexpr (Compute == 100) + return mlp4_make_pipe_umma_2sm(smem.p4_pipe); + else + return mlp4_make_pipe(smem.p4_pipe); + }(); + + // Comm warps exit. Hopper keeps the historical MLP partition (warps 0,4-11) + // and drops warp 3 immediately. Blackwell mirrors the forward path: comm + // uses only warps 1-2 and warp 3 joins MLP as the dedicated UMMA producer. + if constexpr (Compute == 90) { + if (warp_id == 3) return; + } + if (warp_id == 1 || warp_id == 2) return; + + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + cute::conditional_t< + Compute == 100, + cute::TMEM::Allocator2Sm, + cute::TMEM::Allocator1Sm> tmem_alloc{}; + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + if (warp_id == 3) { + tmem_alloc.allocate(kTmemColumns, &smem.tmem_base); + if constexpr (Compute == 100) { + if (cute::elect_one_sync()) { + smem.tmem_pair_barrier.init(1); + smem.tmem_init_barrier.init(1); + smem.acc_pair_barrier.init(1); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + mlp_bwd_pair_barrier_sync(smem.tmem_init_barrier, warp_id); + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + } + + // Keep phase-role tests at their call sites. Hoisting producer/MMA/epilogue + // booleans across every batch keeps them live through all phase mainloops, + // increasing register pressure in a kernel that already spills. + + // ── Persistent per-phase pipeline counts ───────────────────── + // + // We persist only the monotonic iteration count per phase (1 uint32_t) + // rather than the full PipelineState (3 ints: index_, phase_, count_). + // Each phase function reconstructs the live state on entry by calling + // `make_producer_start_state().advance(count)` (or the consumer + // equivalent) — this is correct because `advance(N)` deterministically + // computes the (index, phase) tuple that ++state would arrive at after + // N iterations, and that tuple matches the in-smem barrier phase the + // previous invocation left behind. + // + // PipelineState itself must NOT be persistent across phase invocations + // if we want construction inside the phase — but persisting only the + // count (5 regs across 5 phases vs 15 regs for full states) trims the + // cross-phase live set while keeping correctness. We hit a desync bug + // on T4096_D4096_I14336_E16_K4 (dA off 10% on expert 7) when state was + // reset fresh per batch without preserving the count. + uint32_t pa_count = 0; + uint32_t pb_count = 0; + uint32_t pd_count = 0; + uint32_t p3_count = 0; + uint32_t p4_count = 0; + + mlp_fused_bwd_run_batches< + Traits1, Traits2T, Traits3, Traits4, Traits5, + NSplit2, Compute>( + smem, pa_pipe, pb_pipe, pd_pipe, p3_pipe, p4_pipe, + pa_count, pb_count, pd_count, p3_count, p4_count, + iter, + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_du, tma_store_dv, + tma_load_dy, tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, + EfkbStride3, EfkbStride4, + dims, bufs, global_barrier, x_barrier, warp_id); + + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + if (warp_id == 3) { + tmem_alloc.release_allocation_lock(); + mlp_bwd_pair_barrier_sync(smem.tmem_pair_barrier, warp_id); + tmem_alloc.free(smem.tmem_base, kTmemColumns); + } + } +} + +// ═══════════════════════════════════════════════════════════════════ +// mlp_fused_bwd_dual — single-instance local + remote passes +// ═══════════════════════════════════════════════════════════════════ +// +// Builds pipelines + per-phase counts ONCE, runs the local batch loop, +// crosses the global barrier, and runs the remote batch loop in the +// SAME function body. The aim is to let PTXAS analyze both batch loops +// as sequential code in one CFG, reusing spill slots between them +// (lifetimes don't overlap — separated by global_barrier.wait()). This +// is the alternative to the current pattern in moe_fused_bwd of two +// inlined `mlp_fused_bwd` calls, each constructing its own pipelines +// and re-initializing the smem barriers. +// +// Pipeline barriers in smem are constructed ONCE here. The persistent +// `*_count` state is shared across passes and consumed by the count→ +// PipelineState resume logic inside each phase helper (see +// mlp_bwd_resume_state) — this matches the smem barriers' phase +// counters that advance across both passes. +template 1 → Phase-1 (mlp1a/mlp2t/ + // silu/mlp5) steps SubTiles GemmTileM-row sub-tiles per comm slot; + // Phase-2 sees num_m_tiles·SubTiles GemmTileM token-tiles with + // tile_expert_ids replicated SubTiles× (the K-block count is invariant). + // 1 → unchanged single-tile path. + int SubTiles = 1, + int Compute = 90, + typename FusedIter, + // Shared TMA (Phase 1 weights, Phase 1d B/C, Phase 2 reduce + Z/U/V) + typename TmaLoadW1, typename TmaStoreZ, typename TmaStoreDU, + typename TmaStoreDV, typename TmaLoadW2T, typename TmaStoreDZ, + typename TmaLoadDU, typename TmaLoadDV, + typename TmaLoadB, typename TmaLoadC, + typename TmaLoadZT3, typename TmaReduceDA, + typename TmaLoaddUT4, typename TmaLoaddVT4, + typename TmaReduceDB, typename TmaReduceDC, + // Per-role TMA types — local and remote instances share the + // same C++ type (same cute Layout: dynamic shape, _1 col stride, + // identical smem layout). Only the runtime gmem pointer differs. + typename TmaLoadX, typename TmaLoadDY, typename TmaStoreDX, + typename TmaLoadDYT3, typename TmaLoadXT4> +__device__ __forceinline__ void mlp_fused_bwd_dual( + MlpFusedBwdSmem& smem, + FusedIter& iter, + bool remote_active, + int efkb_stride_4_remote, + // Phase 1 TMA — weights are shared across tiles; X/dY/dX (the *_remote + // instances below) come from the comm staging ring, where EVERY tile, + // local or remote, is staged. + TmaLoadW1 const& tma_load_b_fwd, TmaLoadW1 const& tma_load_c_fwd, + TmaStoreZ const& tma_store_z, + TmaStoreDU const& tma_store_du, TmaStoreDV const& tma_store_dv, + TmaLoadW2T const& tma_load_a_col, TmaStoreDZ const& tma_store_dz, + TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, + // Phase 2 TMA (shared) + TmaLoadZT3 const& tma_load_zt3, TmaReduceDA const& tma_reduce_da, + TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, + TmaReduceDB const& tma_reduce_db, TmaReduceDC const& tma_reduce_dc, + // Staging-ring data TMAs: X/dY/dX (Phase 1) + dYᵀ/Xᵀ (Phase 2). + TmaLoadX const& tma_load_x_remote, + TmaLoadDY const& tma_load_dy_remote, + TmaStoreDX const& tma_store_dx_remote, + TmaLoadDYT3 const& tma_load_dyt3_remote, + TmaLoadXT4 const& tma_load_xt4_remote, + const MlpBwdDims& remote_dims, + const MlpBwdBufs& remote_bufs, + MlpBwdCtaBarrierT& global_barrier, + MlpBwdCtaBarrierT& x_barrier, + // Flat 1-D launch coordinates (computed once in moe_fused_bwd): + // flat_id = blockIdx.x + // col = flat_id / NSplit (Phase 1 logical column) + // grid_x = n_gemm / NSplit (Phase 1 column count) + // split = flat_id % NSplit (N-split index) + // num_splits = NSplit + // gemm_active = flat_id < n_gemm → this CTA runs Phase 1 + comm; + // gap CTAs skip Phase 1 but still run Phase 2 + barriers. + int flat_id, int col, int grid_x, int split, int num_splits, + bool gemm_active) { + + static_assert(SubBatch >= 1 && SubBatch <= kBwdSubBatchMax, + "SubBatch (MoeBwdConfig::kSubBatch) must be in [1, kBwdSubBatchMax]"); + // See mlp_fused_bwd_run_batches — Compute is the SOLE selector of the + // Compute == 100 ClusterM=2 paired-CTA Phase-2 path; the paired-CTA + // UMMA + cluster mbarrier rendezvous below is SM100-only. + static_assert(Compute == 90 || Compute == 100, + "mlp_fused_bwd_dual: Compute must be 90 (SM90/Hopper, 1SM Phase-2) " + "or 100 (SM100/Blackwell, ClusterM=2 paired-CTA Phase-2)."); + + int warp_id = threadIdx.x / Traits1::WarpSize; + constexpr int kMlpBwdBarrierThreads = kMlpBwdBarrierThreadsFor; + + // Prefetch the TMA descriptors this pass uses: shared weights/intermediates + // + the staging-ring data descriptors. + cute::prefetch_tma_descriptor(tma_load_b_fwd.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c_fwd.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_du.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dv.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_a_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dz.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_du.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dv.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_b_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c_col.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_zt3.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_da.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dut4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dvt4.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_db.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_reduce_dc.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_x_remote.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dy_remote.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_dx_remote.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_dyt3_remote.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_xt4_remote.get_tma_descriptor()); + + // Build pipelines ONCE (smem barriers initialized in one place — the + // remote pass picks up smem-barrier phase from where local left off). + auto pa_pipe = [&]() { + if constexpr (Compute == 100) + return mlp1_make_pipe_umma(smem.pa_pipe); + else + return mlp1_make_pipe(smem.pa_pipe); + }(); + auto pb_pipe = [&]() { + if constexpr (Compute == 100) + return mlp2_t_make_pipe_umma(smem.pb_pipe); + else + return mlp2_t_make_pipe(smem.pb_pipe); + }(); + auto pd_pipe = [&]() { + if constexpr (Compute == 100) + return mlp5_make_pipe_umma(smem.pd_pipe); + else + return mlp5_make_pipe(smem.pd_pipe); + }(); + + // Single fused pipe for mlp3 carrying dYT + Z per acquire. Compute == + // 100 always drives the SM100 ClusterM=2 paired-CTA pipe + // (mlp3_make_pipe_umma_2sm); Compute == 90 drives the existing Hopper + // pipe (mlp3_make_pipe). + auto p3_pipe = [&]() { + if constexpr (Compute == 100) + return mlp3_make_pipe_umma_2sm(smem.p3_pipe); + else + return mlp3_make_pipe(smem.p3_pipe); + }(); + // Single fused pipe for mlp4 carrying X + dU^T + dV^T per acquire. + // Compute == 100 always drives the SM100 ClusterM=2 paired-CTA pipe + // (mlp4_make_pipe_umma_2sm); Compute == 90 drives the existing Hopper + // pipe (mlp4_make_pipe). + auto p4_pipe = [&]() { + if constexpr (Compute == 100) + return mlp4_make_pipe_umma_2sm(smem.p4_pipe); + else + return mlp4_make_pipe(smem.p4_pipe); + }(); + + if constexpr (Compute == 90) { + if (warp_id == 3) return; + } + if (warp_id == 1 || warp_id == 2) return; + + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + cute::conditional_t< + Compute == 100, + cute::TMEM::Allocator2Sm, + cute::TMEM::Allocator1Sm> tmem_alloc{}; + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + if (warp_id == 3) { + tmem_alloc.allocate(kTmemColumns, &smem.tmem_base); + if constexpr (Compute == 100) { + if (cute::elect_one_sync()) { + smem.tmem_pair_barrier.init(1); + smem.tmem_init_barrier.init(1); + smem.acc_pair_barrier.init(1); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + mlp_bwd_pair_barrier_sync(smem.tmem_init_barrier, warp_id); + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + } + + // Counts persist across both passes — smem barriers continue advancing + // from local into remote, and PipelineState::advance(count) re-derives + // the right (index, phase) for the resume. + uint32_t pa_count = 0; + uint32_t pb_count = 0; + uint32_t pd_count = 0; + uint32_t p3_count = 0; + uint32_t p4_count = 0; + uint32_t acc_pair_phase = 0; + + using Element = typename Traits1::Element; + + // Phase 2 flat-cell coords are CTA-invariant — hoist out of the loops. + int phase2_div = flat_id / NSplit2; + int phase2_mod = flat_id % NSplit2; + + // Shared k_starts/k_ends require equal TileK (one scan feeds mlp3+mlp4). + static_assert(Traits3::TileK == Traits4::TileK, + "Phase 2 shared k_starts/ends require Traits3::TileK == Traits4::TileK"); + + // Nothing to do if this PE has no tiles (remote_active is grid-uniform, so + // every CTA agrees and the early-out keeps the barriers below well-defined). + if (!remote_active) { + global_barrier.wait(); + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + if (warp_id == 3) { + tmem_alloc.release_allocation_lock(); + mlp_bwd_pair_barrier_sync( + smem.tmem_pair_barrier, warp_id); + tmem_alloc.free(smem.tmem_base, kTmemColumns); + } + } + return; + } + + // One cross-CTA sync so all MLP CTAs enter the grouped loop together: + // SM90 synchronizes w0+w4-11 (288 threads); SM100 synchronizes + // w0+w3-11 (320 threads). Comm warps 1-2 have already returned. + global_barrier.wait(); + + // ═══════════════════════════════════════════════════════════════ + // REMOTE GROUPED PASS — kBwdSubBatch Phase-1 sub-batches, then ONE + // Phase-2 (mlp4+mlp3) over their combined kBwdSubBatch·grid_x K-window. + // ═══════════════════════════════════════════════════════════════ + // Phase 1 (mlp1/2/5) runs in SMALL per-sub-batch chunks (less CTA + // barrier-stall); Phase 2 (weight grads dA/dB/dC) runs ONCE over the whole + // group so its cooperative cell-walk has a long K-reduction to fill the grid + // (deeper outer_split / k_split). Each sub-batch keeps its X/dY staging slot + // un-released until the grouped Phase 2 has read X^T (mlp4) and dY^T (mlp3); + // the saved slots are freed in a batch after each Phase-2 sub-pass. The host + // guarantees CommNumStages is a multiple of kBwdSubBatch and ≥ it, so the + // group window never straddles the ring wrap (per-expert k-ranges stay + // contiguous) and comm has prefetch headroom. kBwdSubBatch == 1 reduces this + // to the original per-batch loop. + const int remote_batches = remote_dims.num_batches; + + // Adaptive group size: a group of `g` sub-batches holds g·grid_x live staging + // slots that must map to DISTINCT ring slots, so g·grid_x ≤ ring. Cap the + // group by how many grid_x-windows the ring holds — deep rings get the full + // SubBatch; shallow rings degrade to fewer (down to 1 = original per-batch + // loop). Grid-uniform (num_m_tiles, grid_x are host-set), so every CTA agrees. + const int eff_subbatch = + max(1, min((int)SubBatch, remote_dims.num_m_tiles / grid_x)); + + // Phase-1 GEMM consumers + Phase-2 see the token axis in GemmTileM + // (=Traits1::TileM) tiles: num_m_tiles (Phase-1 buffer height) and num_tokens + // (M extent) scale by SubTiles = CommTileM/GemmTileM. The comm loop control + // (eff_subbatch, group_start, bk_tiles clamp, slot releases) stays in + // communication-slot units via remote_dims. At SubTiles=1 the two are identical. + // (Other dims fields — N/K extents, expert counts — are M-axis-independent.) + MlpBwdDims gemm_dims = remote_dims; + gemm_dims.num_m_tiles *= SubTiles; + gemm_dims.num_tokens *= SubTiles; + int first_sub = 0; + int last_sub = SubTiles; + int gemm_split = split; + int gemm_splits = num_splits; + constexpr int kM64PerGemmTile = Traits1::TileM / 64; + + for (int g0 = 0; g0 < remote_batches; g0 += eff_subbatch) { + const int gcount = min(eff_subbatch, remote_batches - g0); + + // Per-sub-batch staging slots, freed together after Phase 2. cur_slot in + // the iterator only tracks the LAST acquire, so releasing the earlier + // sub-batches needs their saved slot (release indexes *_consumed[slot]). + int saved_slot[SubBatch]; + bool saved_live[SubBatch]; + + // ── Phase 1: gcount sub-batches (mlp1a / mlp2t / silu_bwd / mlp5) ── + for (int s = 0; s < gcount; ++s) { + int saved_m = -1, saved_expert = -1; + int saved_valid_m64_subtiles = 0; + bool has_tile = false; + bool compute_tile = false; + + // Gated by gemm_active — gap CTAs (flat_id ≥ n_gemm) skip Phase 1 and + // the per-column x_barrier, but still hit every global barrier below. + if (gemm_active && iter.has_next()) { + iter.acquire_src(); + cutlass::arch::NamedBarrier::sync( + kMlpBwdBarrierThreads, kMlpBarrierId); + auto tile = iter.next(); + int expert = tile.expert - remote_dims.local_expert_start; + int m = tile.y_m; + + saved_m = m; + saved_expert = expert; + saved_valid_m64_subtiles = tile.valid_m64_subtiles; + has_tile = true; + compute_tile = + first_sub * kM64PerGemmTile < + saved_valid_m64_subtiles; + + // One CommTileM-row slot feeds SubTiles GemmTileM-row sub-tiles. + // gemm coordinate of sub-tile `sub` = m·SubTiles + sub (the + // Phase-1 kernels multiply by Traits1::TileM=GemmTileM → rows + // m·CommTileM + sub·GemmTileM). gemm_dims carries the SubTiles-scaled + // num_m_tiles / num_tokens so the buffer tensor shapes match. + // Acquire/release stays once per communication tile. + if (compute_tile) { + for (int sub = first_sub; sub < last_sub; ++sub) { + mlp_bwd_run_phase_1a( + pa_pipe, pa_count, smem.mlp1, + tma_load_x_remote, tma_load_b_fwd, tma_load_c_fwd, + tma_store_du, tma_store_dv, tma_store_z, + warp_id, m * SubTiles + sub, expert, gemm_dims, + smem.tmem_base, gemm_split, gemm_splits); + } + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + x_barrier.wait(); + + iter.acquire_dy(); + if (compute_tile) { + for (int sub = first_sub; sub < last_sub; ++sub) { + mlp_bwd_run_phase_1b( + pb_pipe, pb_count, smem.mlp2t, + tma_load_dy_remote, tma_load_a_col, tma_store_dz, + warp_id, m * SubTiles + sub, expert, gemm_dims, + smem.tmem_base, gemm_split, gemm_splits); + } + } + + x_barrier.wait(); + + if (compute_tile) { + for (int sub = first_sub; sub < last_sub; ++sub) { + silu_bwd_pair_tile( + remote_bufs.dz_buf, remote_bufs.du_buf, remote_bufs.dv_buf, + remote_bufs.du_buf, remote_bufs.dv_buf, + m * SubTiles + sub, Traits1::TileM, Traits2T::TileN, + remote_dims.intermediate_dim, gemm_split, gemm_splits); + } + } + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + } + + // Save this sub-batch's staging slot (== tile.y_m == iter cur_slot) + // for the deferred grouped release. -1 / false when this CTA had no + // tile (gap CTA or column exhausted). + saved_slot[s] = saved_m; + saved_live[s] = has_tile; + + // Global barrier: this sub-batch's Z, dU, dV visible to all CTAs. + global_barrier.wait(); + + // Phase 1d: dX = dU@B + dV@C (per sub-batch — Phase 2 never reads dX, + // so it is emitted and its dst slot released immediately). + if (has_tile) { + iter.acquire_dst(); + if (compute_tile) { + for (int sub = first_sub; sub < last_sub; ++sub) { + mlp_bwd_run_phase_1d( + pd_pipe, pd_count, smem.mlp5, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, + tma_store_dx_remote, + warp_id, saved_m * SubTiles + sub, saved_expert, gemm_dims, + smem.tmem_base, gemm_split, gemm_splits); + } + } + __threadfence(); // dX flush for the NIC put (release_dst signals) + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + iter.release_dst(threadIdx.x); + } + } + + // ── Phase 2: ONE pass over the group's gcount·grid_x K-window ── + // Phase 1 advances each of grid_x columns by one tile per sub-batch, so + // the group delivered global tiles [g0·grid_x, (g0+gcount)·grid_x). + // bk_start is UNWRAPPED (group_start) so the group's K-window stays + // contiguous even when it straddles the staging-ring wrap; build_ranges + // and the producers wrap the physical slot by % ring_kb internally + // (ring_kb > 0 enables it). The group's gcount·grid_x live slots are + // guaranteed distinct (host check: ring > SubBatch·grid_x). + // group_start / bk_tiles_comm are in communication-slot units. Phase 2 + // works in GemmTileM units, so scale the contiguous parent window by + // SubTiles; range construction skips padded K-blocks inside it. + int group_start = g0 * grid_x; + int bk_tiles_comm = min( + gcount * grid_x, + remote_dims.total_m_tiles_in_pass - group_start); + int bk_start = group_start * SubTiles; + int bk_tiles = bk_tiles_comm * SubTiles; + int ring_kb = (gemm_dims.num_m_tiles * Traits1::TileM) / Traits4::TileK; + const cutlass::arch::ClusterBarrier* acc_pair_barrier = nullptr; + if constexpr (Compute == 100) + acc_pair_barrier = &smem.acc_pair_barrier; + + // Per-group per-expert K-range scan over the group's window. Reads the + // LIVE X-pipe expert ids — none of the group's slots have been released + // yet, so all gcount·grid_x are valid. Feeds mlp3+mlp4 (unwrapped ranges). + if (bk_tiles > 0) { + int batch_kb_off = (bk_start * Traits1::TileM) / Traits4::TileK; + int batch_kb_cnt = (bk_tiles * Traits1::TileM) / Traits4::TileK; + mlp_bwd_phase_2_build_ranges( + remote_bufs.expert_k_starts, remote_bufs.expert_k_ends, + remote_bufs.expert_for_k_block_mlp4, + remote_bufs.valid_rows_mlp4, + batch_kb_off, batch_kb_cnt, + remote_dims.experts_per_pe, remote_dims.local_expert_start, + efkb_stride_4_remote, Traits4::TileK, warp_id, ring_kb); + // Cross-CTA barrier: CTA 0 owns reset+scan, so no producer may read + // the shared ranges until that work completes. + global_barrier.wait(); + } + + if (bk_tiles > 0) { + mlp_bwd_run_phase_2_mlp4< + Traits1, Traits4, NSplit2, Compute>( + p4_pipe, p4_count, smem.mlp4, + remote_bufs.expert_k_starts, remote_bufs.expert_k_ends, + tma_load_xt4_remote, tma_load_dut4, tma_load_dvt4, + tma_reduce_db, tma_reduce_dc, + warp_id, bk_start, bk_tiles, + phase2_div, phase2_mod, + gemm_dims, smem.tmem_base, ring_kb, + acc_pair_barrier, &acc_pair_phase); + } + + global_barrier.wait(); + // mlp4 done reading X^T — free every X slot the group held. + { + __threadfence(); + for (int s = 0; s < gcount; ++s) + if (saved_live[s]) iter.release_src_slot(saved_slot[s], threadIdx.x); + } + + if (bk_tiles > 0) { + mlp_bwd_run_phase_2_mlp3< + Traits1, Traits3, NSplit2, Compute>( + p3_pipe, p3_count, smem.mlp3, + remote_bufs.expert_k_starts, remote_bufs.expert_k_ends, + tma_load_dyt3_remote, tma_load_zt3, tma_reduce_da, + warp_id, bk_start, bk_tiles, + phase2_div, phase2_mod, + gemm_dims, smem.tmem_base, ring_kb, + acc_pair_barrier, &acc_pair_phase); + } + + global_barrier.wait(); + // mlp3 done reading dY^T — free every dY slot the group held. + { + __threadfence(); + for (int s = 0; s < gcount; ++s) + if (saved_live[s]) iter.release_dy_slot(saved_slot[s], threadIdx.x); + } + } + + if constexpr (Compute == 100) { + constexpr int kTmemColumns = + MlpBwdTmemColumns::value; + cutlass::arch::NamedBarrier::sync(kMlpBwdBarrierThreads, kMlpBarrierId); + if (warp_id == 3) { + tmem_alloc.release_allocation_lock(); + mlp_bwd_pair_barrier_sync(smem.tmem_pair_barrier, warp_id); + tmem_alloc.free(smem.tmem_base, kTmemColumns); + } + } +} + + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_comms.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_comms.cuh index e30bb495..f7d480d5 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp_comms.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_comms.cuh @@ -1,1343 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// Fused MoE MLP with NVSHMEM communication -// ═══════════════════════════════════════════════════════════════════ -// -// Combines nvshmem get/put with MLP compute in a single persistent -// kernel. Warps 1-3 of WG0 handle nvshmem communication while the -// rest of the CTA runs the MLP1 pipeline. -// -// Warp roles (384 threads = 12 warps): -// -// WG0 (warps 0–3): -// warp 0 TMA producer (mlp1 pipelines) -// warp 1 NVSHMEM get warp 0 (fetches X into src_staging) -// warp 2 NVSHMEM get warp 1 (fetches X into src_staging) -// warp 3 NVSHMEM put warp (sends Y from dst_staging) -// -// WG1 (warps 4–7): Consumer A (mlp1 WGMMA) -// WG2 (warps 8–11): Consumer B (mlp1 WGMMA) -// -// Staging areas (per blockIdx.x, shared by both blockIdx.y): -// src_staging[NumStages][TileM][hidden_dim] -// dst_staging[NumStages][TileM][intermediate_dim] -// -// Schedule: the two get warps and the put warp each run an -// independent flat loop over the same tile sequence. The NumStages -// staging ring bounds how far the get warps can lead the put warp. -// -// NVSHMEM get throughput is roughly half of put; we counteract that -// by issuing 2× as many gets per tile. Each tile is chunked across -// 2·NC warps on the get side (2 warps per CTA × NC CTAs) but only NC -// warps on the put side. Get chunk = tile_bytes / (2·NC); put chunk -// = tile_bytes / NC. -// -// Transport-dependent chunking: -// * NVLink P2P (nvshmem_ptr != nullptr): all 2·NC get warps / NC put -// warps cooperate per tile, each copying its chunk with coalesced -// int4 loads/stores — many warps maximise NVLink throughput. -// * IB (nvshmem_ptr == nullptr): the per-warp chunking is counter- -// productive (small RDMA messages, per-tile rendezvous). Instead a -// SINGLE warp issues one whole-tile getmem/putmem, and the warp that -// issues rotates round-robin across the 2·NC (get) / NC (put) warps -// keyed by the tile sequence index, so consecutive tiles are issued -// by different warps. Their blocking transfers then overlap (bounded -// by the NumStages staging ring). The non-issuing warps do no I/O and -// just rubber-stamp the StagePipe so the producer/consumer counters -// stay intact. -// -// ═══════════════════════════════════════════════════════════════════ - -#include // CUtensorMap (raw TMA descriptor) -#include -#include -#include // typed TMA / mbarrier intrinsics (replaces inline PTX) -#include -#include -#include -#include "nvshmem_helpers.cuh" // g_dest_table - -#include "cta_barrier.cuh" // SyncThreadsCtaCounterBarrier - -namespace liger { - -// ═══════════════════════════════════════════════════════════════════ -// StagePipe — bidirectional acquire/release for a staging buffer -// ═══════════════════════════════════════════════════════════════════ -// -// Wraps a pair of atomic counters (ready[S], consumed[S]) shared -// between a producer and a consumer. Each side calls acquire/release: -// -// Producer: acquire() waits on consumed[] (slot free), then -// release() signals ready[] (data written). -// Consumer: acquire() waits on ready[] (data available), then -// release() signals consumed[] (slot freed). -// -// Internally tracks stage index and monotonic thresholds. -// is_leader is set at init — only the leader thread polls atomics. - -template -struct StagePipe { - int* ready; // producer signals, consumer waits - int* consumed; // consumer signals, producer waits - int stride; // step in slots between consecutive logical stages - bool is_leader; - int num_producers; - int num_consumers; - - int prod_count; - int cons_count; - int prod_expected[NumStages]; // threshold for consumed[] (producer waits on) - int cons_expected[NumStages]; // threshold for ready[] (consumer waits on) - - // Stride lets the pipe address into a flat slot list with non-unit - // stage spacing (e.g. comm pipes use stride = MC, MLP pipes use - // stride = gridDim.x). Default stride = 1 preserves the original - // contiguous-ring layout used by the bwd path. - __device__ void init(int* ready_, int* consumed_, bool is_leader_, - int stride_ = 1) { - ready = ready_; - consumed = consumed_; - stride = stride_; - is_leader = is_leader_; - num_producers = NumProducers; - num_consumers = NumConsumers; - prod_count = 0; - cons_count = 0; - for (int i = 0; i < NumStages; ++i) { - prod_expected[i] = 0; - cons_expected[i] = 0; - } - } - - __device__ void init(int* ready_, int* consumed_, bool is_leader_, - int stride_, int num_producers_, int num_consumers_) { - init(ready_, consumed_, is_leader_, stride_); - num_producers = num_producers_; - num_consumers = num_consumers_; - } - - // ── Producer side ─────────────────────────────────────── - - __device__ int producer_stage() const { - return prod_count % NumStages; - } - - // Wait until this slot is free (consumer has released it). - // - // Acquire-side memory ordering: only the leader does the atomic poll - // (cheaper than 32 lanes spinning), but ALL lanes need to see the - // counterparty's writes after this returns. __syncwarp alone does NOT - // establish global-memory ordering — non-leader lanes may still hit - // stale L1/L2. We add a __threadfence after the leader's poll AND after - // __syncwarp so every lane has a release/acquire pair. - __device__ void producer_acquire() { - int s = producer_stage(); - if (prod_count >= NumStages) { - if (is_leader) { - cuda::atomic_ref r(consumed[s * stride]); - while (r.load(cuda::memory_order_relaxed) < prod_expected[s]) {} - } - __syncwarp(); - __threadfence(); - } - prod_expected[s] += num_consumers; - } - - // Signal that data has been written to this slot. - __device__ void producer_release(int lane) { - int s = producer_stage(); - // Device-scope: this is the get-side publish (NIC-written src → MLP) or - // a slot-free signal — both GPU-internal w.r.t. the consumer. The NIC's - // inbound-write visibility is NOT fixable by a fence (needs GDRCopy), so - // system scope here only costs bandwidth without buying correctness. The - // one publish that genuinely crosses to the NIC (dst → put warp) keeps - // system scope in RemoteIter::release_dst. - __threadfence(); - if (lane == 0) { - atomicAdd(&ready[s * stride], 1); - } - prod_count++; - } - - // ── Consumer side ─────────────────────────────────────── - - __device__ int consumer_stage() const { - return cons_count % NumStages; - } - - // Wait until this slot has data (producer has released it). - // - // Acquire-side memory ordering: see comment on producer_acquire above. - // __syncwarp alone does NOT establish global-memory visibility for - // non-leader lanes — they may load stale data from L1/L2 even after - // the leader has confirmed the producer's release. The __threadfence - // after __syncwarp gives every lane in the warp a fresh view of the - // producer's writes (HBM coherence) before the warp-collective load - // in the subsequent put. - __device__ void consumer_acquire() { - int s = consumer_stage(); - if (is_leader) { - cuda::atomic_ref r(ready[s * stride]); - while (r.load(cuda::memory_order_relaxed) < cons_expected[s] + num_producers) {} - } - __syncwarp(); - // Device-scope: the put warp observes dst_staging (Y), already flushed - // to HBM by the MLP's system-scope release_dst. Observe is GPU-internal. - __threadfence(); - cons_expected[s] += num_producers; - } - - // Signal that this slot has been consumed and is free. - __device__ void consumer_release(int lane) { - int s = consumer_stage(); - __threadfence(); - if (lane == 0) { - atomicAdd(&consumed[s * stride], 1); - } - cons_count++; - } -}; - -// ═══════════════════════════════════════════════════════════════════ -// CommSmem — per-CTA shared memory (smem-local, not shared across CTAs) -// ═══════════════════════════════════════════════════════════════════ - -struct CommSmem { - // Total number of remote tiles across all CTAs. Initialized to -1. - // Comm warps write the actual count; consumers wait for >= 0. - int total_tiles; - - // Per-CTA tile count (this blockIdx.x's share of total_tiles). - int per_cta_tiles; - - // Grid-wide remote tile count (sum across all CTAs/columns). Same - // value on every CTA in the grid; used by moe_fused_bwd to gate the - // remote phase consistently — the inner mlp_fused_bwd's cross-CTA - // mlp_global_barrier requires every CTA to make the same enter/skip - // decision, so per-CTA per_cta_tiles can't be the gate. - int global_total; - int runtime_nsplit; // N-split CTAs per GEMM M-row - int runtime_n_gemm; - int runtime_grid_x; - - // Pre-fetched remote expert offsets, assigned to dynamic smem at runtime. - // Layout: [num_pes][experts_per_pe + 1] - int* remote_offsets; - // Exact routed-row counts matching remote_offsets. - // Layout: [num_pes][experts_per_pe] - int* remote_counts; -}; - -__device__ __forceinline__ int select_runtime_nsplit_from_tiles( - int total_tiles, int num_n_tiles_1, int num_n_tiles_2, - int num_k_tiles_1, int num_k_tiles_2, - int num_blocks, int static_nsplit, int m_subtiles = 1) { - if (total_tiles <= 0) return 2 * m_subtiles; - constexpr int base_candidates[5] = {2, 4, 6, 8, 16}; - int best_base_ns = static_nsplit; - long long best_cost = -1; - long long best_waste = -1; - int best_n_gemm = -1; - long long phase1_weight = 2LL * num_k_tiles_1; // gate + up - long long phase2_weight = (long long)num_k_tiles_2; - #pragma unroll - for (int i = 0; i < 5; ++i) { - int base_ns = base_candidates[i]; - if (base_ns * m_subtiles > num_blocks) continue; - int n_gemm = (num_blocks / base_ns) * base_ns; - int ms = n_gemm / base_ns; - if (ms <= 0) continue; - int m_waves = (total_tiles + ms - 1) / ms; - int n1_waves = (num_n_tiles_1 + base_ns - 1) / base_ns; - int n2_waves = (num_n_tiles_2 + base_ns - 1) / base_ns; - - long long phase_cost = phase1_weight * n1_waves - + phase2_weight * n2_waves; - long long cost = (long long)m_waves * phase_cost; - long long phase1_waste = - (long long)m_waves * ms * n1_waves * base_ns - - (long long)total_tiles * num_n_tiles_1; - long long phase2_waste = - (long long)m_waves * ms * n2_waves * base_ns - - (long long)total_tiles * num_n_tiles_2; - long long waste = phase1_weight * phase1_waste - + phase2_weight * phase2_waste; - if (best_cost < 0 || - cost < best_cost || - (cost == best_cost && waste < best_waste) || - (cost == best_cost && waste == best_waste && n_gemm > best_n_gemm)) { - best_cost = cost; - best_waste = waste; - best_n_gemm = n_gemm; - best_base_ns = base_ns; - } - } - return best_base_ns * m_subtiles; -} - -// ═══════════════════════════════════════════════════════════════════ -// PeSync — device-level phase synchronization -// ═══════════════════════════════════════════════════════════════════ - -struct PeSync { - int* cta_counter; // [1] device memory — shared monotonic counter - int num_pes; - int my_pe; - nvshmem_team_t team; - - // ── barrier: cross-CTA global memory fence ────────── - // Defers to a persistent SyncThreadsCtaCounterBarrier owned by the - // caller — see cta_barrier.cuh for why the barrier object must be - // constructed ONCE per kernel and reused (target accumulates). - // PeSync cannot own the barrier itself because it lives inside a - // __grid_constant__ MoeParams; the caller builds the barrier on - // the kernel stack and threads it into every signal/wait call. - __device__ void barrier(SyncThreadsCtaCounterBarrier& b) const { - b.wait(); - } - - __device__ void barrier_all(SyncThreadsCtaCounterBarrier& b) const { - b.wait(); - int cta_id = blockIdx.x + blockIdx.y * gridDim.x; - if (cta_id == 0) { - if (num_pes > 1 && threadIdx.x == 0) - nvshmem_barrier(team); - __syncthreads(); - } - b.wait(); - } -}; - -// ═══════════════════════════════════════════════════════════════════ -// CommBuffers — pointers passed to the kernel -// ═══════════════════════════════════════════════════════════════════ - -struct CommBuffers { - // Staging areas: per-blockIdx.x, shared by both blockIdx.y CTAs. - // Layout: [gridDim.x][NumStages][TileM][dim] - void* src_staging; // [gridDim.x * S * TileM * hidden_dim] - void* dst_staging; // [gridDim.x * S * TileM * intermediate_dim] - - // Stage pipe signals in device memory (HBM), per-blockIdx.x. - // Layout: [gridDim.x][NumStages] — shared by both blockIdx.y CTAs. - int* src_ready; // nvshmem → consumer: tile data written - int* src_consumed; // consumer → nvshmem: tile data read - int* dst_ready; // consumer → nvshmem: tile data written - int* dst_consumed; // nvshmem → consumer: put completed, slot free - - const int* expert_offsets; // [E + 1] local (from radix sort, TileM-aligned, symmetric) - const int* expert_counts; // [E] exact routed rows per expert - const int* sorted_token_ids; // [total_slots] local sorted token IDs - - void* local_tokens; // [T, hidden_dim] local input (symmetric) - void* local_output; // [total_slots, intermediate_dim] local output (symmetric) - - // Per-staging-slot expert IDs in HBM. Written by do_get, read by MLP. - // Layout: [gridDim.x][NumStages] - int* tile_expert_ids; - - // All PEs' expert offsets, gathered into local symmetric memory. - // Layout: [num_pes][num_experts + 1] — symmetric. - int* all_expert_offsets; - // All PEs' exact expert counts. Layout: [num_pes][num_experts]. - int* all_expert_counts; - - // Cross-PE synchronization for offset broadcast/consume. - PeSync pe_sync; - - int hidden_dim; - int num_experts; // total global experts - int experts_per_pe; // num_experts / num_pes (evenly split) - int all_expert_offsets_stride; // max_num_experts + 1 (fixed across configs) - // GEMM/comm-active CTA count = floor_NS (largest multiple of NSplit <= - // num_sms). The flat-grid launch may run more CTAs than this (grid-stride - // phases use them); GEMM and comm gate to the first n_gemm CTAs so the - // comm↔MLP staging ring geometry (grid_x = n_gemm/NSplit, MC = n_gemm/NC) - // stays consistent. Replaces the old gridDim.x·gridDim.y product. - int n_gemm; - - // Convenience accessors for fields stored in pe_sync. - __host__ __device__ __forceinline__ int my_pe() const { return pe_sync.my_pe; } - __host__ __device__ __forceinline__ int num_pes() const { return pe_sync.num_pes; } - __host__ __device__ __forceinline__ nvshmem_team_t team() const { return pe_sync.team; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// TileInfo — describes one tile to get or put -// ═══════════════════════════════════════════════════════════════════ - -struct TileInfo { - int pe; // remote PE - int token_offset; // offset into remote PE's token buffer (in tokens) - int expert; // local expert index (for weight selection) - int valid_rows; // valid prefix rows in this communication tile - int m_subtile; // GemmTileM row within this communication tile - int parent_tile; // flat communication-tile index -}; - -// ═══════════════════════════════════════════════════════════════════ -// TileIterator — enumerates tiles: expert-outer, PE-inner, strided -// ═══════════════════════════════════════════════════════════════════ -// -// Tiles within each (expert, pe) pair are strided by gridDim.x. -// Each blockIdx.x starts at its own offset and advances by gridDim.x. - -template -struct TileIterator { - static constexpr int TileM = TileM_; - - const int* remote_offsets; // smem: [num_pes][experts_per_pe + 1] - int offsets_stride; // experts_per_pe + 1 - const int* remote_counts; // smem: [num_pes][experts_per_pe] (nullable) - int counts_stride; // experts_per_pe - int experts_per_pe; - int num_pes; - int my_pe; - - int stride; // distance between consecutive tiles - - // Current position in the iteration. - int le; // local expert index - int pe_idx; // 1..num_pes-1 (skips my_pe) - int tile; // tile within current (expert, pe) pair - int remote_tiles; // total tiles for current (expert, pe) pair - int remote_start; // start offset for current (expert, pe) - int remote_count; // exact valid rows for current (expert, pe) - - int total_tiles; // this iterator's tile count - - __device__ void init(int start, int str) { - stride = str; - le = 0; - pe_idx = 0; - total_tiles = 0; - - // Count total tiles across all (expert, pe) pairs and - // compute this iterator's share. - // NOTE: p starts at 0 → the LOCAL PE (my_pe, via g_dest_table slot - // offset 0) is enumerated first for each expert, alongside the remote - // PEs (p=1..num_pes-1). Unified path: local experts are staged by the - // comm get warp from local symmetric memory just like remote tiles, so - // the separate local MLP pass is gone (see moe.cuh moe_fused_fwd). - int lane = threadIdx.x & 31; - int local_total = 0; - int n_experts = experts_per_pe * num_pes; - for (int idx = lane; idx < n_experts; idx += 32) { - int e = idx / num_pes; - int p = idx - e * num_pes; - int pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + p) % num_pes]; - int rs = remote_offsets[pe * offsets_stride + e]; - int re = remote_offsets[pe * offsets_stride + e + 1]; - local_total += (re - rs) / TileM; - } - #pragma unroll - for (int delta = 16; delta > 0; delta >>= 1) { - local_total += __shfl_down_sync(0xffffffff, local_total, delta); - } - int global_total = __shfl_sync(0xffffffff, local_total, 0); - if (start < global_total) - total_tiles = (global_total - 1 - start) / stride + 1; - - // Find the starting position: walk the flat tile space - // to find where start lands. - int skip = start; - for (int e = 0; e < experts_per_pe; ++e) { - for (int p = 0; p < num_pes; ++p) { - int pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + p) % num_pes]; - int rs = remote_offsets[pe * offsets_stride + e]; - int re = remote_offsets[pe * offsets_stride + e + 1]; - int ntiles = (re - rs) / TileM; - if (skip < ntiles) { - le = e; - pe_idx = p; - tile = skip; - remote_start = rs; - remote_tiles = ntiles; - remote_count = remote_counts != nullptr - ? remote_counts[pe * counts_stride + e] - : ntiles * TileM; - return; - } - skip -= ntiles; - } - } - // No tiles for this iterator. - le = experts_per_pe; - } - - __device__ bool has_next() const { - return le < experts_per_pe; - } - - __device__ TileInfo next() { - int pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; - TileInfo info; - info.pe = pe; - info.token_offset = remote_start + tile * TileM; - info.expert = le; - info.valid_rows = max(0, min(TileM, remote_count - tile * TileM)); - info.m_subtile = 0; - info.parent_tile = -1; - - // Advance by stride in the flat tile space. - // Carry the remainder across (expert, pe) boundaries. - tile += stride; - while (tile >= remote_tiles && le < experts_per_pe) { - int remaining = tile - remote_tiles; - pe_idx++; - if (pe_idx >= num_pes) { - pe_idx = 0; - le++; - } - if (le >= experts_per_pe) break; - int next_pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; - remote_start = remote_offsets[next_pe * offsets_stride + le]; - int remote_end = remote_offsets[next_pe * offsets_stride + le + 1]; - remote_tiles = (remote_end - remote_start) / TileM; - remote_count = remote_counts != nullptr - ? remote_counts[next_pe * counts_stride + le] - : remote_tiles * TileM; - tile = remaining; - } - return info; - } -}; - -template -struct GemmTileIterator { - static constexpr int TileM = TileM_; - static constexpr int MSubTiles = MSubTiles_; - static constexpr int GemmTileM = TileM / MSubTiles; - - const int* remote_offsets; - int offsets_stride; - const int* remote_counts; - int counts_stride; - int experts_per_pe; - int num_pes; - int my_pe; - int stride; - - int le; - int pe_idx; - int tile; - int remote_tiles; - int remote_parent_tiles; - int remote_start; - int remote_count; - int parent_prefix; - int total_tiles; - - __device__ void init(int start, int str) { - stride = str; - le = 0; - pe_idx = 0; - total_tiles = 0; - - int lane = threadIdx.x & 31; - int local_total = 0; - int n_experts = experts_per_pe * num_pes; - for (int idx = lane; idx < n_experts; idx += 32) { - int e = idx / num_pes; - int p = idx - e * num_pes; - int pe = liger::g_dest_table[ - (liger::comm_slot_of(my_pe) + p) % num_pes]; - int count = remote_counts[pe * counts_stride + e]; - local_total += (count + GemmTileM - 1) / GemmTileM; - } - #pragma unroll - for (int delta = 16; delta > 0; delta >>= 1) - local_total += __shfl_down_sync(0xffffffff, local_total, delta); - int global_total = __shfl_sync(0xffffffff, local_total, 0); - if (start < global_total) - total_tiles = (global_total - 1 - start) / stride + 1; - - int skip = start; - parent_prefix = 0; - for (int e = 0; e < experts_per_pe; ++e) { - for (int p = 0; p < num_pes; ++p) { - int pe = liger::g_dest_table[ - (liger::comm_slot_of(my_pe) + p) % num_pes]; - int rs = remote_offsets[pe * offsets_stride + e]; - int re = remote_offsets[pe * offsets_stride + e + 1]; - int count = remote_counts[pe * counts_stride + e]; - int ntiles = (count + GemmTileM - 1) / GemmTileM; - int nparents = (re - rs) / TileM; - if (skip < ntiles) { - le = e; - pe_idx = p; - tile = skip; - remote_tiles = ntiles; - remote_parent_tiles = nparents; - remote_start = rs; - remote_count = count; - return; - } - skip -= ntiles; - parent_prefix += nparents; - } - } - le = experts_per_pe; - } - - __device__ bool has_next() const { - return le < experts_per_pe; - } - - __device__ int current_parent_tile() const { - return parent_prefix + tile / MSubTiles; - } - - __device__ int current_m_subtile() const { - return tile % MSubTiles; - } - - __device__ TileInfo next() { - int pe = liger::g_dest_table[ - (liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; - int parent_in_expert = tile / MSubTiles; - TileInfo info; - info.pe = pe; - info.token_offset = remote_start + parent_in_expert * TileM; - info.expert = le; - info.valid_rows = max( - 0, min(TileM, remote_count - parent_in_expert * TileM)); - info.m_subtile = tile % MSubTiles; - info.parent_tile = parent_prefix + parent_in_expert; - - tile += stride; - while (tile >= remote_tiles && le < experts_per_pe) { - int remaining = tile - remote_tiles; - parent_prefix += remote_parent_tiles; - pe_idx++; - if (pe_idx >= num_pes) { - pe_idx = 0; - le++; - } - if (le >= experts_per_pe) - break; - int next_pe = liger::g_dest_table[ - (liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; - remote_start = - remote_offsets[next_pe * offsets_stride + le]; - int remote_end = - remote_offsets[next_pe * offsets_stride + le + 1]; - remote_count = - remote_counts[next_pe * counts_stride + le]; - remote_tiles = - (remote_count + GemmTileM - 1) / GemmTileM; - remote_parent_tiles = - (remote_end - remote_start) / TileM; - tile = remaining; - } - return info; - } -}; - -// ═══════════════════════════════════════════════════════════════════ -// NVSHMEM Comm device functions — dedicated get/put warps -// ═══════════════════════════════════════════════════════════════════ -// -// Per CTA: warp 2 is the getter, warp 3 is the putter. -// -// The physical grid is launched as (gridDim.x, N_SPLIT) to match the -// GEMM's cooperative-WG layout. Comms re-indexes the flat CTA id into -// a (MC, NC) logical grid where NC must divide N_SPLIT: -// -// cta_id = blockIdx.x * gridDim.y + blockIdx.y -// xc = cta_id / NC in [0, MC) -// yc = cta_id % NC in [0, NC) — chunk index -// MC = gridDim.x * (gridDim.y / NC) -// -// NC CTAs (same xc, different yc) cooperate on one tile and chunk its -// bytes NC ways. MC groups stride through the remote-tile list, so -// more tiles are in flight per kernel launch when NC < N_SPLIT. -// -// A CTA's comm tile is chosen by xc; its GEMM tile is chosen by -// blockIdx.x. These can diverge when NC < N_SPLIT (the whole point of -// the re-indexing) — the staging ring is keyed by xc so MLP and comm -// rendezvous through the ring, not through shared CTA coordinates. -// -// The two roles run independent flat loops over all tiles. No 1F1B — -// the staging ring (NumStages) bounds how far the getter can lead -// the putter via the shared src/dst StagePipe signals. - -// Prologue + bwd use warps 2-3 only (kept for backward compatibility -// with the bwd path that has not been migrated to the 2-getter layout). -static constexpr int kCommWarpStart = 2; -static constexpr int kCommWarpEnd = 2; -static constexpr int kCommWarpsPerCta = kCommWarpEnd - kCommWarpStart + 1; // 2 -static constexpr int kCommThreadStart = kCommWarpStart * 32; - -// Fwd comm layout: 2 get warps + 1 put warp per CTA. -// warp 1 → get warp 0 -// warp 2 → get warp 1 (also prologue leader) -// warp 3 → put warp -static constexpr int kCommGetWarp = 1; -static constexpr int kCommPutWarp = 2; -static constexpr int kNumGetWarpsPerCta = 2; // shared w/ bwd — do NOT change - -// FWD-only active get-warp count. With the TMA GET (DMA engine, issue-bound -// not thread-bound), ONE warp saturates the per-SM TMA unit, so warp 2 is -// freed (idle in steady state; still helps the epilogue last-tile put drain). -// kNumGetWarpsPerCta stays 2 for the bwd path and for the 2-warp epilogue -// drain (kPutW); only the steady-state FWD producer count drops to this. -static constexpr int kNumGetWarpsFwd = 1; - -// FWD-only put-warp count. Keep one active put warp on Hopper; the freed warp -// enters the MLP-side branch and exits there for Compute != 100. -static constexpr int kNumPutWarpsFwd = 1; - -// Range of warps participating in fwd comm (used by callers to filter -// MLP vs comm warps). -static constexpr int kCommFwdWarpStart = kCommGetWarp; -static constexpr int kCommFwdWarpEnd = kCommPutWarp; - -// Get side has kNumGetWarpsPerCta × NC warps per tile (chunk = -// tile_bytes / (kNumGetWarpsPerCta × NC)); put side has NC warps per -// tile (chunk = tile_bytes / NC). - -// ═══════════════════════════════════════════════════════════════════ -// GET-via-TMA: fetch X over NVLink with the TMA engine, not warp LD/ST. -// ═══════════════════════════════════════════════════════════════════ -// -// The dedicated get warp drives a TMA load (peer X → smem) followed by a -// TMA store (smem → local src_staging HBM). TMA can't go global→global on -// Hopper, so the bounce through smem is mandatory; the store TARGET is -// LOCAL HBM, so completion (cp.async.bulk.wait_group) + a device-scope -// __threadfence in producer_release makes it visible to the MLP reader — -// the safe-completion path (unlike a store to a PEER, which failed before). -// -// Validated in microbench: ~370 GB/s, in- and out-of-kernel correct. -// -// kGetKChunk is the contiguous column box width (elements). Each of the -// kNumGetWarpsPerCta·NC warps owns a TileM/(kNumGetWarpsPerCta·NC) row -// band of the tile and loops over hidden_dim/kGetKChunk column boxes. -// Box = [kGetBoxRows, kGetKChunk]; descriptors built host-side in moe.cu. -static constexpr int kGetKChunk = 256; - -// Per-peer source descriptors + one local-staging dst descriptor, built on -// the host (nvshmem_ptr(x_sorted, peer)) and passed __grid_constant__. -// FWD-only: NOT referenced by bwd — kept out of CommBuffers so the shared -// struct's ABI is unchanged and moe_bwd.cu need not recompile. -// -// Multi-host: only SAME-HOST peers are P2P-reachable, so src_x_desc is indexed -// by intra-host local rank (pe % gpus_per_node), NOT by global team PE — the -// array therefore only ever needs gpus_per_node (≤ kGetMaxPes) entries -// regardless of total PE count. Remote (IB) peers fall back to getmem and never -// touch this array. Assumes ranks are host-contiguous and host-aligned to -// gpus_per_node (true for a WORLD-contiguous team); asserted host-side. -static constexpr int kGetMaxPes = 8; // upper bound on GPUs per host -struct GetTmaDescs { - CUtensorMap src_x_desc[kGetMaxPes]; // index by (TileInfo::pe % gpus_per_node) - CUtensorMap dst_staging_desc; // local src_staging - int enabled; // 1 = NVLink P2P TMA path available; 0 = getmem - int gpus_per_node; // local host PE count — host/local-rank divisor - int my_pe; // this PE in team space (for same-host test) -}; - -// Same-host predicate (pure arithmetic): true iff peer `pe` shares this PE's -// host. Used both to gate the TMA get path and to pick the release_dst fence -// scope (device for same-host NVLink, system for cross-host IB). -__host__ __device__ __forceinline__ bool get_pe_is_local( - const GetTmaDescs* descs, int pe) { - if (descs == nullptr || descs->gpus_per_node <= 0) return true; // single-host/test - return (pe / descs->gpus_per_node) == (descs->my_pe / descs->gpus_per_node); -} - -// Per-peer Y-store descriptors for the direct GEMM→peer TMA store (FWD-only). -// desc[r] is a CUTE TMA-store object built host-side over -// nvshmem_ptr(local_output, peer) for the same-host peer at intra-host local -// rank r (self included, rank = my_pe % gpus_per_node). When a tile is_local, -// MLP2 stores Y straight into the destination peer's symmetric local_output via -// desc[MlpTileInfo::peer_rank] instead of dst_staging — the put warp then only -// rubber-stamps. The direct store is gated per-tile on is_local (same-host → -// P2P guaranteed), so no runtime enable flag is needed. Templated on the CUTE -// store type so the same object the wrapper passes as tma_store_y populates it. -template -struct PeerYStoreDescs { - int enabled = 1; - TmaStoreY desc[kGetMaxPes]; // index by (TileInfo::pe % gpus_per_node) -}; - -// Smem bounce: per get-warp single box + its mbarrier, 128-aligned so the -// TMA dst smem address is aligned. Stride between the two get warps' regions. -template -struct GetBounce { - static constexpr int kWarpsPerTile = kNumGetWarpsFwd * NC; - static constexpr int kBoxRows = TileM / kWarpsPerTile; - static constexpr int kBoxElems = kBoxRows * kGetKChunk; - static constexpr int kBoxBytes = kBoxElems * (int)sizeof(Element); - // [box (128-aligned)] [mbar:8B], next warp's region 128-aligned. - static constexpr int kPerWarpStride = ((kBoxBytes + 8) + 127) & ~127; - static constexpr int kTotalBytes = kPerWarpStride * kNumGetWarpsFwd; -}; - -template -__device__ __forceinline__ void do_get( - StagePipe& src_pipe, - const CommBuffers& bufs, - const TileInfo& info, - Element* src_staging, // base of flat staging buffer - int src_tile_elems, - int xc, // start slot in flat list - int MC, // stride between consecutive comm stages - int chunk_idx, // 0..kNumGetWarpsPerCta * NC - 1 - int m_subtile, - int row_chunk_idx, - bool row_empty, - int& ib_seq, // IB-tile counter (round-robin turn; advanced here) - int lane, - const GetTmaDescs* descs, // per-peer src + local dst TMA descriptors (nullptr → getmem) - char* bounce_warp) { // this warp's GetBounce sub-region (128-aligned) - - using Bounce = GetBounce; - constexpr int kWarpsPerTile = Bounce::kWarpsPerTile; - constexpr int kBoxRows = Bounce::kBoxRows; - - src_pipe.producer_acquire(); - - // No per-tile signal wait — wait_dispatch() already confirmed - // all dispatch writes on all PEs are complete. - - int slot = xc + src_pipe.producer_stage() * MC; - if (row_empty) { - src_pipe.producer_release(lane); - return; - } - - // Locality of this tile's source PE (same-host NVLink vs cross-host IB). - // Pure arithmetic; gates the TMA get path below. The MLP's release_dst - // recomputes the SAME predicate independently (embedded TileIterator), so no - // per-slot flag is published here. - bool is_local = get_pe_is_local(descs, info.pe); - - // TMA path: peer X → smem → local src_staging, via the TMA engine. - // Gated on same-host NVLink P2P (is_local), TMA availability (descs->enabled), - // and hidden_dim divisibility; cross-host IB, odd shapes, or a null descs - // (e.g. the standalone test) fall back to getmem. - bool tma_ok = descs && descs->enabled && is_local - && (bufs.hidden_dim % kGetKChunk == 0); - if (tma_ok) { - // Typed TMA / mbarrier intrinsics from . These compile to the - // same SASS as the hand-written PTX they replace: the peer-X→smem load - // is the exact `cp.async.bulk.tensor.2d.shared::cluster.global.tile` - // instruction, and the smem→staging store the `.global.shared::cta.tile` - // one (`.tile` is the default tiled mode → identical encoding). - namespace ptx = cuda::ptx; - Element* sbuf = reinterpret_cast(bounce_warp); - uint64_t* mbar = reinterpret_cast(bounce_warp + Bounce::kBoxBytes); - - if (lane == 0) - ptx::mbarrier_init(mbar, 1u); - __syncwarp(); - // Make the barrier init visible to the async (TMA) proxy before use. - ptx::fence_proxy_async(ptx::space_shared); - - // Tile coords: each warp owns its kBoxRows-tall row band. - int y_src = info.token_offset + chunk_idx * kBoxRows; // peer X row - int y_dst = slot * TileM + chunk_idx * kBoxRows; // staging row - // Same-host (tma_ok ⇒ is_local) → index by intra-host local rank so the - // descriptor array stays gpus_per_node-sized regardless of total PEs. - const CUtensorMap* src_map = &descs->src_x_desc[info.pe % descs->gpus_per_node]; - const CUtensorMap* dst_map = &descs->dst_staging_desc; - - const int n_k = bufs.hidden_dim / kGetKChunk; - const uint32_t txbytes = (uint32_t)Bounce::kBoxBytes; - uint32_t phase = 0; - for (int k = 0; k < n_k; ++k) { - int x = k * kGetKChunk; - if (lane == 0) { - // Arm the barrier for the incoming box, then issue the - // peer-X → smem TMA load it tracks. - (void)ptx::mbarrier_arrive_expect_tx( - ptx::sem_release, ptx::scope_cta, ptx::space_shared, - mbar, txbytes); - int32_t load_crd[2] = {x, y_src}; - ptx::cp_async_bulk_tensor( - ptx::space_cluster, ptx::space_global, - sbuf, src_map, load_crd, mbar); - // Spin until the load completes (single-buffer bounce). - while (!ptx::mbarrier_try_wait_parity(mbar, phase)) {} - // smem now holds the loaded box; push it to local staging. - int32_t store_crd[2] = {x, y_dst}; - ptx::cp_async_bulk_tensor( - ptx::space_global, ptx::space_shared, - dst_map, store_crd, sbuf); - ptx::cp_async_bulk_commit_group(); - // Single-buffer: drain the store before the next load overwrites smem. - ptx::cp_async_bulk_wait_group(ptx::n32_t<0>{}); - } - phase ^= 1; - __syncwarp(); - } - - if (bufs.tile_expert_ids && chunk_idx == 0 && lane == 0) - bufs.tile_expert_ids[slot] = info.expert; - - // producer_release's __threadfence (device scope) orders the completed - // local TMA stores before the ready signal — MLP reader sees fresh X. - __syncwarp(); - src_pipe.producer_release(lane); - return; - } - - // ── Fallback: warp getmem (IB, or hidden_dim not divisible by kGetKChunk) ── - int chunk = src_tile_elems / kWarpsPerTile; - int offset = chunk_idx * chunk; - auto* local_tokens = static_cast(bufs.local_tokens); - Element* remote_base = local_tokens + info.token_offset * bufs.hidden_dim; - Element* local_base = src_staging + slot * src_tile_elems; - int global_pe = nvshmem_team_translate_pe(bufs.team(), info.pe, NVSHMEM_TEAM_WORLD); - bool is_ib = descs != nullptr && !is_local; - if (is_ib) { - constexpr int kRowNC = NC / MSubTiles; - constexpr int kRowsPerSubtile = TileM / MSubTiles; - int turn = ib_seq % kRowNC; - if (row_chunk_idx == turn) { - int row_start = m_subtile * kRowsPerSubtile; - int valid_rows = max( - 0, min(kRowsPerSubtile, info.valid_rows - row_start)); - int row_offset = row_start * bufs.hidden_dim; - nvshmemx_getmem_warp( - local_base + row_offset, remote_base + row_offset, - (size_t)valid_rows * bufs.hidden_dim * sizeof(Element), - global_pe); - } - ib_seq++; - } else { - nvshmemx_getmem_warp(local_base + offset, remote_base + offset, - (size_t)chunk * sizeof(Element), global_pe); - } - if (bufs.tile_expert_ids && chunk_idx == 0 && lane == 0) { - bufs.tile_expert_ids[slot] = info.expert; - } - src_pipe.producer_release(lane); -} - -template -__device__ __forceinline__ void do_put( - StagePipe& dst_pipe, - const CommBuffers& bufs, - const TileInfo& info, - Element* dst_staging, - int dst_tile_elems, - int xc, // start slot in flat list - int MC, // stride between consecutive comm stages - int yc, // 0..NC-1 — CTA's chunk index within the tile - int m_subtile, - int row_chunk_idx, - bool row_empty, - int& ib_seq, // IB-tile counter (round-robin turn; advanced here) - int lane, - int put_local_idx, // 0..kNumPutWarpsFwd-1 — put warp within the CTA - const GetTmaDescs* descs) { // null → standalone test (real put); else skip local - - // Two put warps per CTA: split the tile kNumPutWarpsFwd·NC ways, this - // warp owns slice (yc·kNumPutWarpsFwd + put_local_idx). - constexpr int kWarpsPerTile = kNumPutWarpsFwd * NC; - int chunk = dst_tile_elems / kWarpsPerTile; - int offset = (yc * kNumPutWarpsFwd + put_local_idx) * chunk; - int my_bytes = chunk * sizeof(Element); - - dst_pipe.consumer_acquire(); - if (row_empty) { - dst_pipe.consumer_release(lane); - return; - } - - // Direct-store: for same-host (NVLink) tiles the GEMM has already TMA-stored - // Y straight into the peer's symmetric local_output, so this put warp must - // NOT copy dst_staging over it (that would clobber the correct Y with a stale - // staging slot). It only RUBBER-STAMPS the dst pipe so the StagePipe counters - // stay consistent and the GEMM's acquire_dst never stalls. Cross-host (IB) - // tiles still need the real putmem. descs==nullptr is the standalone comm - // test (no direct store) → always do the real put. - bool skip_put = (descs != nullptr) && get_pe_is_local(descs, info.pe); - if (!skip_put) { - auto* local_output = static_cast(bufs.local_output); - int slot = xc + dst_pipe.consumer_stage() * MC; - Element* local_base = dst_staging + slot * dst_tile_elems; - Element* remote_base = local_output + info.token_offset * bufs.hidden_dim; - int global_pe = nvshmem_team_translate_pe(bufs.team(), info.pe, NVSHMEM_TEAM_WORLD); - bool is_ib = descs != nullptr; - if (is_ib) { - constexpr int kRowNC = NC / MSubTiles; - constexpr int kRowsPerSubtile = TileM / MSubTiles; - int turn = ib_seq % (kNumPutWarpsFwd * kRowNC); - int row_warp_idx = - row_chunk_idx * kNumPutWarpsFwd + put_local_idx; - if (row_warp_idx == turn) { - int row_start = m_subtile * kRowsPerSubtile; - int valid_rows = max( - 0, min(kRowsPerSubtile, info.valid_rows - row_start)); - int row_offset = row_start * bufs.hidden_dim; - nvshmemx_putmem_warp( - remote_base + row_offset, local_base + row_offset, - (size_t)valid_rows * bufs.hidden_dim * sizeof(Element), - global_pe); - } - ib_seq++; - } else { - nvshmemx_putmem_warp(remote_base + offset, local_base + offset, - my_bytes, global_pe); - } - } - - dst_pipe.consumer_release(lane); -} - -// Chunk-only put for the cooperative LAST-tile drain (epilogue get-warp put -// swap). The dedicated put warp owns the pipe (acquire/release); each of the -// PutW comm warps just writes its share — NO pipe ops here, so no extra -// dst_ready/dst_consumed signalling. P2P splits the tile PutW*NC ways; IB -// sends the whole tile from rank 0 only. dst_tile_elems must be divisible by -// PutW*NC (TileM*D divisible by 2*NC holds for the supported shapes). -template -__device__ __forceinline__ void do_put_chunk( - const CommBuffers& bufs, const TileInfo& info, - Element* dst_staging, int dst_tile_elems, - int slot, int yc, int put_rank, int lane) { - constexpr int kWarpsPerTile = PutW * NC; - int put_idx = yc * PutW + put_rank; - int chunk = dst_tile_elems / kWarpsPerTile; - int offset = put_idx * chunk; - Element* local_base = dst_staging + (size_t)slot * dst_tile_elems; - Element* remote_base = static_cast(bufs.local_output) - + (size_t)info.token_offset * bufs.hidden_dim; - int global_pe = nvshmem_team_translate_pe(bufs.team(), info.pe, NVSHMEM_TEAM_WORLD); - // nbi: the last tile's put has no downstream consumer waiting on it, so it - // need not block — completion is guaranteed by the nvshmem_quiet() each - // comm warp issues after comm_main returns (moe.cu). Single per-chunk put, - // no transport branch (NVLink-only build) → one call site. - nvshmemx_putmem_nbi_warp(remote_base + offset, local_base + offset, - (size_t)chunk * sizeof(Element), global_pe); -} - -// ═══════════════════════════════════════════════════════════════════ -// broadcast_expert_offsets — called after sort, before MLP phase. -// ═══════════════════════════════════════════════════════════════════ - -__device__ __forceinline__ void broadcast_expert_offsets( - const CommBuffers& bufs) { - - int cta_id = blockIdx.x + blockIdx.y * gridDim.x; - int num_experts = bufs.num_experts; - int num_pes = bufs.num_pes(); - int my_pe = bufs.my_pe(); - int offsets_stride = num_experts + 1; - const int* my_offsets = bufs.expert_offsets; - - // CTA 0: copy into own local all_expert_offsets[my_pe][:]. - if (cta_id == 0 && threadIdx.x == 0) { - int* local_dst = bufs.all_expert_offsets + my_pe * offsets_stride; - for (int i = 0; i <= num_experts; ++i) - local_dst[i] = my_offsets[i]; - int* local_counts_dst = - bufs.all_expert_counts + my_pe * num_experts; - for (int i = 0; i < num_experts; ++i) - local_counts_dst[i] = bufs.expert_counts[i]; - } - - // CTA i puts offsets to PE i (if valid remote PE). - // No signal here — signal_dispatch() after dispatch covers it. - // NVSHMEM ordering: when dispatch signal arrives, this put is visible. - int target_pe = cta_id; - if (target_pe >= num_pes || target_pe == my_pe) return; - - int global_target = nvshmem_team_translate_pe(bufs.team(), target_pe, NVSHMEM_TEAM_WORLD); - int* remote_dst = bufs.all_expert_offsets + my_pe * offsets_stride; - nvshmemx_int_put_nbi_block( - remote_dst, my_offsets, - offsets_stride, - global_target); - int* remote_counts_dst = - bufs.all_expert_counts + my_pe * num_experts; - nvshmemx_int_put_nbi_block( - remote_counts_dst, bufs.expert_counts, - num_experts, - global_target); -} - -// ═══════════════════════════════════════════════════════════════════ -// Prologue: waits for remote offsets, computes total_tiles. -// ═══════════════════════════════════════════════════════════════════ -// -// Called by all threads — comm warps do the work, others pass through. -// Must be followed by __syncthreads() before nvshmem_comm_main. - -template -__device__ __forceinline__ void nvshmem_comm_prologue( - CommSmem& smem, - const CommBuffers& bufs, - int static_nsplit) { - - int warp_id = threadIdx.x / 32; - if (warp_id < kCommWarpStart || warp_id > kCommWarpEnd) - return; - - int lane = threadIdx.x % 32; - int local_warp = warp_id - kCommWarpStart; - bool is_leader = (local_warp == 0 && lane == 0); - - int experts_per_pe = bufs.experts_per_pe; - int num_pes = bufs.num_pes(); - int my_pe = bufs.my_pe(); - - if (bufs.all_expert_offsets == nullptr) { - // Standalone test (no nvshmem offsets) — no tiles. - if (is_leader) { - smem.total_tiles = 0; - smem.per_cta_tiles = 0; - smem.global_total = 0; - smem.runtime_nsplit = static_nsplit; - smem.runtime_n_gemm = bufs.n_gemm; - smem.runtime_grid_x = bufs.n_gemm / static_nsplit; - } - __syncwarp(); - return; - } - // NOTE: num_pes == 1 is NO LONGER short-circuited. The unified TileIterator - // enumerates the LOCAL PE (p=0), so a single-GPU run stages its local tiles - // through the comm get/put path exactly like the multi-GPU remote path — - // this is the apples-to-apples isolation of the per-tile staging cost. - - // wait_dispatch() was called before entering the MLP phase, - // guaranteeing all remote PEs' offsets are visible. - - // Extract this PE's local expert offsets from each remote PE's full array. - // all_expert_offsets layout: [num_pes][num_experts + 1]. - // smem layout: [num_pes][experts_per_pe + 1] — only this PE's experts. - int local_e_start = my_pe * experts_per_pe; - int src_stride = bufs.num_experts + 1; - int dst_stride = experts_per_pe + 1; - { - int tid = threadIdx.x - kCommThreadStart; - for (int pe = 0; pe < num_pes; ++pe) { - for (int i = tid; i <= experts_per_pe; i += (kCommWarpsPerCta * 32)) - smem.remote_offsets[pe * dst_stride + i] = - bufs.all_expert_offsets[pe * src_stride + local_e_start + i]; - for (int i = tid; i < experts_per_pe; i += (kCommWarpsPerCta * 32)) - smem.remote_counts[pe * experts_per_pe + i] = - bufs.all_expert_counts[ - pe * bufs.num_experts + local_e_start + i]; - } - } - __syncwarp(); - - // Re-index the flat cta_id to the comm (xc, yc) grid. Comm walks - // tiles owned by (xc, yc) with stride MC. MLP walks tiles for this - // blockIdx.x column with stride gridDim.x — a different start AND - // a different stride, so the two counts must be computed separately - // (uneven distribution can make them differ by more than the ratio - // of the strides per CTA). - // Flat-grid launch: blockIdx.x is the flat cta_id; the logical comm/MLP - // grid comes from n_gemm + compile-time NC/N_SPLIT, not gridDim. - // comm: xc = cta_id / NC, MC = n_gemm / NC - // MLP: col = cta_id / N_SPLIT, grid_x = n_gemm / N_SPLIT - int cta_id = blockIdx.x; - // Comm grid is NC-aligned and DECOUPLED from the GEMM's NS-aligned n_gemm: - // recompute the comm count n_comm from the launched grid (NOT n_gemm). The - // comm side uses all NC-complete CTAs; the leftover remainder (xc >= MC) - // gets comm_count == 0 and self-gates. The MLP count below keeps n_gemm. - int n_comm = (gridDim.x / NC) * NC; - int xc = cta_id / NC; - int MC = n_comm / NC; - - TileIterator iter; - iter.remote_offsets = smem.remote_offsets; - iter.offsets_stride = dst_stride; - iter.remote_counts = smem.remote_counts; - iter.counts_stride = experts_per_pe; - iter.experts_per_pe = experts_per_pe; - iter.num_pes = num_pes; - iter.my_pe = my_pe; - - // Comm count: this CTA's per-(xc, yc) workload. - iter.init(xc, MC); - int comm_count = iter.total_tiles; - - // Grid-wide total tile count (start=0, stride=1 → all tiles once). - iter.init(0, 1); - int global_count = iter.total_tiles; - - // Default static-NS MLP count; moe_fused_kernel overwrites this after it - // selects a runtime NS using MLP N-tile divisibility. - int static_grid_x = bufs.n_gemm / static_nsplit; - iter.init(cta_id / static_nsplit, static_grid_x); - int mlp_count = iter.total_tiles; - - if (is_leader) { - smem.total_tiles = comm_count; // read by comm_main - smem.per_cta_tiles = mlp_count; // read by MLP iterator - smem.global_total = global_count; - smem.runtime_nsplit = static_nsplit; - smem.runtime_n_gemm = bufs.n_gemm; - smem.runtime_grid_x = static_grid_x; - } - __syncwarp(); -} - -// Main loop: warps 1, 2 (getters) and warp 3 (putter) each run an -// independent flat loop over the same tile sequence. Called after -// __syncthreads(). -// -// CTA re-indexing: -// cta_id = blockIdx.x * gridDim.y + blockIdx.y -// xc = cta_id / NC (tile-stream index in [0, MC)) -// yc = cta_id % NC (CTA's slot within its xc group) -// MC = (gridDim.x * gridDim.y) / NC (multiply before divide so MC is -// exact when NC ∤ gridDim.y) -// -// Each xc group strides through the remote-tile list by MC; all NC -// CTAs with the same xc cooperate on the same tile. Get side splits -// each tile kNumGetWarpsPerCta × NC ways, with -// chunk_idx = yc * kNumGetWarpsPerCta + local_get_idx. Put side -// splits NC ways with chunk_idx = yc. -// -// Signal arrays and staging are a flat list of L = MC * NumStages -// slots. Comm CTA at xc walks slots {xc, xc+MC, …, xc+(NumStages-1)·MC} -// (stride MC, NumStages entries — full column coverage of L). MLP at -// blockIdx.x walks slots {blockIdx.x, blockIdx.x+gridDim.x, …} (stride -// gridDim.x, L/gridDim.x = NumStages·N_SPLIT/NC entries — looping back -// across the full L). Both sides cover all L slots, and per-column -// subsets are disjoint across blockIdx.x (gridDim.x | L, since -// L = gridDim.x · NumStages · N_SPLIT/NC and NC | NumStages·N_SPLIT). -// -// The flat list bounds how far the getters can lead the putter: the -// getter blocks in producer_acquire once its NumStages-slot subset is -// pending MLP consumption, and the putter blocks in consumer_acquire -// until MLP releases a dst slot. -template -__device__ __forceinline__ void nvshmem_comm_main( - CommSmem& smem, - const CommBuffers& bufs, - const GetTmaDescs* get_descs, // per-peer X load + local staging store (nullptr → getmem) - char* get_bounce) { // base of this CTA's GetBounce smem region - - // Comm-side pipe depth is NumStages — one comm CTA cycle covers - // the full MC × NumStages staging area, independent of N_SPLIT and - // NC. The MLP side is ticket-based (RemoteMlpTileIterator) and - // shares no per-stage_id state, so no divisibility relationship is - // required between NumStages, N_SPLIT and NC. - constexpr int K = NumStages; - - int warp_id = threadIdx.x / 32; - if (warp_id != kCommGetWarp && warp_id != kCommPutWarp) - return; - - int lane = threadIdx.x % 32; - // Flat-grid launch: blockIdx.x is the flat cta_id. (xc, yc) comm grid and - // MC derive from n_gemm + compile-time NC, not gridDim. - int cta_id = blockIdx.x; - // n_comm: NC-aligned comm grid, decoupled from n_gemm, recomputed from the - // launched grid to match the prologue's comm_count. - int n_comm = (gridDim.x / NC) * NC; - int xc = cta_id / NC; - int yc = cta_id % NC; - int MC = n_comm / NC; - static_assert(NC % MSubTiles == 0, - "communication CTAs must divide evenly across GEMM M-subtiles"); - constexpr int kRowNC = NC / MSubTiles; - int m_subtile = yc / kRowNC; - int row_yc = yc % kRowNC; - if (cta_id >= n_comm) return; - - // Per-warp leader: each comm warp must independently spin in - // StagePipe acquire(). __syncwarp() is warp-local. - bool is_leader = (lane == 0); - - int total = smem.total_tiles; - if (total == 0) return; - - int tile_elems = TileM * bufs.hidden_dim; - - // ── Tile iterator: MC-strided through all remote tiles, starting at xc ── - // All three roles visit the same tile sequence in the same order. - TileIterator iter; - iter.remote_offsets = smem.remote_offsets; - iter.offsets_stride = bufs.experts_per_pe + 1; - iter.remote_counts = smem.remote_counts; - iter.counts_stride = bufs.experts_per_pe; - iter.experts_per_pe = bufs.experts_per_pe; - iter.num_pes = bufs.num_pes(); - iter.my_pe = bufs.my_pe(); - iter.init(xc, MC); - - int my_total = iter.total_tiles; - if (my_total == 0) return; - int runtime_nsplit = smem.runtime_nsplit; - - auto* src_staging_base = static_cast(bufs.src_staging); - auto* dst_staging_base = static_cast(bufs.dst_staging); - - // Warp roles (FWD, TMA GET): warp 1 = single TMA get warp; warp 3 = - // single put warp. Warp 2 is deliberately not a comm warp in this variant. - if (warp_id == kCommPutWarp) { - // One put warp per CTA. put_local_idx is 0; it puts its - // 1/(kNumPutWarpsFwd·NC) slice and independently drives the pipe. - // NumProducers = N_SPLIT (GEMM CTAs produce Y), NumConsumers = - // kNumPutWarpsFwd·NC. - const int put_local_idx = 0; - StagePipe dst_pipe; - dst_pipe.init( - bufs.dst_ready + xc * MSubTiles + m_subtile, - bufs.dst_consumed + xc * MSubTiles + m_subtile, - is_leader, MC * MSubTiles, - runtime_nsplit, kNumPutWarpsFwd * kRowNC); - - int ib_seq = 0; // advances only on IB tiles (inside do_put) - for (int i = 0; i < my_total; ++i) { - TileInfo info = iter.next(); - constexpr int kRowsPerSubtile = TileM / MSubTiles; - bool row_empty = - info.valid_rows <= m_subtile * kRowsPerSubtile; - if (row_empty && row_yc == 0 && lane == 0) { - int s = dst_pipe.consumer_stage(); - atomicAdd( - &dst_pipe.ready[s * dst_pipe.stride], - runtime_nsplit); - } - do_put( - dst_pipe, bufs, info, - dst_staging_base, tile_elems, xc, MC, yc, - m_subtile, row_yc, row_empty, ib_seq, lane, - put_local_idx, get_descs); - } - } else if (warp_id == kCommGetWarp) { // single TMA get warp - const int chunk_idx = yc; // kNumGetWarpsFwd == 1 → yc·1 + 0 - - // NumProducers = kNumGetWarpsFwd·NC (= NC: one get warp × NC CTAs), - // NumConsumers = N_SPLIT (all GEMM CTAs consume cooperatively). - StagePipe src_pipe; - src_pipe.init( - bufs.src_ready + xc * MSubTiles + m_subtile, - bufs.src_consumed + xc * MSubTiles + m_subtile, - is_leader, MC * MSubTiles, - kNumGetWarpsFwd * kRowNC, runtime_nsplit); - char* bounce_warp = get_bounce; // single warp → slice 0 - - int ib_seq = 0; // advances only on IB tiles (inside do_get) - for (int i = 0; i < my_total; ++i) { - TileInfo info = iter.next(); - constexpr int kRowsPerSubtile = TileM / MSubTiles; - bool row_empty = - info.valid_rows <= m_subtile * kRowsPerSubtile; - int s = src_pipe.producer_stage(); - do_get( - src_pipe, bufs, info, - src_staging_base, tile_elems, xc, MC, chunk_idx, - m_subtile, row_yc, row_empty, ib_seq, lane, - get_descs, bounce_warp); - if (row_empty && row_yc == 0 && lane == 0) { - atomicAdd( - &src_pipe.consumed[s * src_pipe.stride], - runtime_nsplit); - } - } - } -} - -} // namespace liger +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif + +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp_comms_sm90.cuh" +#else +#include "mlp_comms_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd.cuh index 4aef2359..75ce1ef9 100644 --- a/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd.cuh @@ -1,772 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// Backward MoE communication: gets X + dY, puts dX -// ═══════════════════════════════════════════════════════════════════ -// -// Warp roles: -// warp 1 → X/dY get warp -// warp 2 → dX put warp -// warp 3 → MLP-side SM100 UMMA producer (exits early on SM90) -// -// Grid is launched 2D as (gridDim.x, NSplit) where NSplit is the -// MLP-side cooperative count. Comms re-indexes the flat CTA id into -// a (MC, NC) logical grid: -// -// cta_id = blockIdx.x * gridDim.y + blockIdx.y -// xc = cta_id / NC in [0, MC) -// yc = cta_id % NC in [0, NC) -// MC = (gridDim.x * gridDim.y) / NC (multiply before divide so -// MC is exact when NC ∤ NSplit) -// -// NC CTAs (same xc, different yc) cooperate on one tile per get/put. -// MC > gridDim.x when NC < NSplit so excess comm CTAs run ahead and -// load the next tiles while the MLP is still on the current one. -// -// ── Tile-level locking (same as fwd) ─────────────────────────────── -// -// All three pipes (X, dY, dX) use the per-tile StagePipe over a flat -// [L = MC * NumStages]-sized signal array. Get pipes have -// NumProducers = kNumGetWarpsPerCta·NC, NumConsumers = NSplit. Put -// pipe has NumProducers = NSplit, NumConsumers = NC. Comm strides MC -// across stages (NumStages entries — full ring cycle, independent of -// NSplit and NC); MLP is ticket-based (RemoteMlpTileIteratorBwd) and -// keys on the absolute slot index = (m_base + idx·gridDim.x) mod L. -// Both sides cover the same flat L slots. -// -// ═══════════════════════════════════════════════════════════════════ - -#include "mlp_comms.cuh" - -namespace liger { - -// ═══════════════════════════════════════════════════════════════════ -// Bwd comm warp constants -// ═══════════════════════════════════════════════════════════════════ -// -// Main loop runs on warps 1-2 only. kCommBwdGetWarp1 remains for legacy -// prologue compatibility with shared helpers; it is not a main-loop role. - -static constexpr int kCommBwdGetWarp0 = 1; -static constexpr int kCommBwdGetWarp1 = 2; // prologue compatibility only -static constexpr int kCommBwdPutWarp = 2; -static constexpr int kCommBwdWarpStart = kCommBwdGetWarp0; -static constexpr int kCommBwdWarpEnd = kCommBwdPutWarp; - -// ── TMA-GET (NVLink) warp rebalance ───────────────────────────────── -// Mirrors the fwd #111 layout: one warp saturates the TMA GET path, so warp 1 -// gets X+dY and warp 2 puts dX. Warp 3 enters the MLP-side branch and exits -// there for Compute != 100. The MLP-side iterator picks producer/consumer -// counts from the same runtime tma_enabled flag. -static constexpr int kNumGetWarpsBwd = 1; // TMA mode: 1 get warp (X+dY) -static constexpr int kNumPutWarpsBwd = 1; // TMA mode: 1 put warp (dX) - -// ═══════════════════════════════════════════════════════════════════ -// GetTmaDescsBwd — raw CUtensorMap descriptors for the bwd TMA GET -// ═══════════════════════════════════════════════════════════════════ -// -// Per-same-host-peer source maps over nvshmem_ptr(x_sorted, peer) and -// nvshmem_ptr(dy_sorted, peer); one local dst map each over x_staging / -// dy_staging. Built host-side in moe_bwd.cu, passed __grid_constant__. -// Same as fwd GetTmaDescs: local peers index by pe % gpus_per_node; remote -// peers fall back to IB getmem in do_get_bwd_tma. -struct GetTmaDescsBwd { - CUtensorMap src_x_desc[kGetMaxPes]; // index by (TileInfo::pe % gpus_per_node) - CUtensorMap src_dy_desc[kGetMaxPes]; - CUtensorMap dst_x_staging_desc; // local x_staging - CUtensorMap dst_dy_staging_desc; // local dy_staging - int enabled; // 1 = local-peer TMA path available - int gpus_per_node; - int my_pe; -}; - -__host__ __device__ __forceinline__ bool get_pe_is_local_bwd( - const GetTmaDescsBwd* descs, int pe) { - if (descs == nullptr || descs->gpus_per_node <= 0) return true; - return (pe / descs->gpus_per_node) == (descs->my_pe / descs->gpus_per_node); -} - -// Box height (rows) of a single backward TMA transfer. SM100 mirrors -// forward's full [32×256] row-band box; SM90 uses a smaller 8-row box. -static constexpr int kGetBoxRowsBwdSm90 = 8; -static constexpr int kGetBoxRowsBwdSm100 = 32; - -// X and dY reuse one data box sequentially but keep distinct mbarriers. -// Reinitializing one mbarrier across the two TMA streams races the async proxy; -// the box itself is safe to reuse after cp.async.bulk.wait_group<0>. -template -struct GetBounceBwd { - static constexpr int kBoxRows = BoxRows; - static constexpr int kBoxElems = kBoxRows * kGetKChunk; - static constexpr int kBoxBytes = kBoxElems * (int)sizeof(Element); - static constexpr int kMbarOffset = kBoxBytes; - static constexpr int kMbarStride = sizeof(uint64_t); - static constexpr int kTotalBytes = - (kBoxBytes + 2 * kMbarStride + 127) & ~127; -}; - -// do_get_bwd_tma is defined below, after CommBuffersBwd (it needs the -// complete struct type for bufs.hidden_dim). - -// ═══════════════════════════════════════════════════════════════════ -// CommBuffersBwd — backward-specific buffer pointers -// ═══════════════════════════════════════════════════════════════════ - -struct CommBuffersBwd { - // Staging: flat slot list [L][TileM][hidden_dim] where - // L = gridDim.x * NumStages. Indexed by slot ∈ [0, L). - void* x_staging; // gets: X from remote PE - void* dy_staging; // gets: dY from remote PE - void* dx_staging; // puts: dX back to remote PE - - // Per-tile signals ([L] entries each), one counter pair per slot — - // same granularity as fwd src/dst pipes. - int* x_src_ready; - int* x_src_consumed; - int* dy_src_ready; - int* dy_src_consumed; - int* dst_ready; - int* dst_consumed; - - // Per-slot expert IDs ([L] each), one array per get pipe. - // - // They hold the SAME value for any given live slot — info.expert at the - // time the comm warp wrote that slot — but each is gated by its own - // pipe's StagePipe release window: - // - // tile_expert_ids_x is overwritten between MLP's release_src and - // comm's next-lap producer_release on the X pipe. - // Phase 2 mlp4 reads this (mlp4 reads X) and is - // safe as long as release_src is deferred past mlp4. - // - // tile_expert_ids_dy is overwritten between MLP's release_dy and - // comm's next-lap producer_release on the dY pipe. - // Phase 2 mlp3 reads this (mlp3 reads dY) and is - // safe as long as release_dy is deferred past mlp3. - // - // Splitting the array is what lets release_src happen between mlp4 and - // mlp3 without races: comm refilling tile_expert_ids_x in that window - // can't poison mlp3, because mlp3 reads the dY-pipe array. - int* tile_expert_ids_x; - int* tile_expert_ids_dy; - int* tile_valid_rows_x; - - // Remote data pointers (symmetric memory) - void* remote_x; // get source: X on remote PE - void* remote_dy; // get source: dY on remote PE - void* remote_dx; // put destination: dX on remote PE - - // Reused from forward - int* all_expert_offsets; - int* all_expert_counts; - int all_expert_offsets_stride; - - int my_pe; - int num_pes; - nvshmem_team_t team; - - int hidden_dim; - int num_experts; - int experts_per_pe; - - // GEMM/comm-active CTA count = floor_NS = grid_x · NSplit (NOT the - // launched grid size). Under the flat 1-D launch gridDim.x = num_blocks - // (= num_sms) ≥ n_gemm, so MC = n_gemm / NC and the MLP column count - // grid_x = n_gemm / NSplit must come from here, not from gridDim. - int n_gemm; - - // Used by the dummy_consumer_bwd in mlp_comms_bwd.cu when it replicates - // the real BWD's batch loop (mlp_global_barrier + mlp_x_barrier). Not - // used by the comm warps themselves. - int* barrier_counter; - int* phase_counter; -}; - -// ═══════════════════════════════════════════════════════════════════ -// do_get_bwd_tma — TMA GET of one buffer (X or dY) for one tile -// ═══════════════════════════════════════════════════════════════════ -// -// Mirrors fwd do_get's TMA path (cuda::ptx intrinsics): peer → smem box → -// local staging, via TMA load + TMA store. Because the bounce box is small -// (kGetBoxRowsBwd rows), the single get warp loops over its TileM/NC-row -// band in kGetBoxRowsBwd-tall sub-bands × kGetKChunk-wide column boxes. -// chunk_idx == yc (the CTA's slot within its NC-cooperating group); with -// kNumGetWarpsBwd == 1 the per-warp row band is TileM/NC. -template -__device__ __forceinline__ void do_get_bwd_tma( - StagePipe& src_pipe, - const CommBuffersBwd& bufs, - const TileInfo& info, - Element* remote_base, - Element* staging_base, - int tile_elems, - int xc, - int MC, - int chunk_idx, // = yc (0..NC-1) - int& ib_seq, // advanced once per actual IB get call - int lane, - int mbar_index, - int* tile_expert_ids_dst, - int* tile_valid_rows_dst, - const GetTmaDescsBwd* descs, - const CUtensorMap* src_desc_per_pe, // descs->src_x_desc or src_dy_desc - const CUtensorMap* dst_desc, // descs->dst_x_staging or dst_dy_staging - char* bounce_warp) { - - namespace ptx = cuda::ptx; - using Bounce = GetBounceBwd; - constexpr int kWarpsPerTile = kNumGetWarpsBwd * NC; // = NC - constexpr int kRowBand = TileM / kWarpsPerTile; // rows this warp owns - constexpr int kBoxRows = Bounce::kBoxRows; - - src_pipe.producer_acquire(); - - int slot = xc + src_pipe.producer_stage() * MC; - Element* local_tile = staging_base + slot * tile_elems; - Element* remote_tile = remote_base + info.token_offset * bufs.hidden_dim; - - bool is_local = descs && descs->enabled && get_pe_is_local_bwd(descs, info.pe); - bool tma_ok = is_local && (bufs.hidden_dim % kGetKChunk == 0) && bounce_warp != nullptr; - - if (tma_ok) { - Element* sbuf = reinterpret_cast(bounce_warp); - uint64_t* mbar = reinterpret_cast( - bounce_warp + Bounce::kMbarOffset + - mbar_index * Bounce::kMbarStride); - - if (lane == 0) - ptx::mbarrier_init(mbar, 1u); - __syncwarp(); - ptx::fence_proxy_async(ptx::space_shared); - - const int n_k = bufs.hidden_dim / kGetKChunk; - const uint32_t txbytes = (uint32_t)Bounce::kBoxBytes; - int y_src_base = info.token_offset + chunk_idx * kRowBand; // peer row - int y_dst_base = slot * TileM + chunk_idx * kRowBand; // staging row - const int local_peer = info.pe % descs->gpus_per_node; - - uint32_t phase = 0; - for (int rb = 0; rb < kRowBand; rb += kBoxRows) { - for (int k = 0; k < n_k; ++k) { - int x = k * kGetKChunk; - if (lane == 0) { - (void)ptx::mbarrier_arrive_expect_tx( - ptx::sem_release, ptx::scope_cta, ptx::space_shared, - mbar, txbytes); - int32_t load_crd[2] = {x, y_src_base + rb}; - ptx::cp_async_bulk_tensor( - ptx::space_cluster, ptx::space_global, - sbuf, &src_desc_per_pe[local_peer], load_crd, mbar); - while (!ptx::mbarrier_try_wait_parity(mbar, phase)) {} - int32_t store_crd[2] = {x, y_dst_base + rb}; - ptx::cp_async_bulk_tensor( - ptx::space_global, ptx::space_shared, - dst_desc, store_crd, sbuf); - ptx::cp_async_bulk_commit_group(); - ptx::cp_async_bulk_wait_group(ptx::n32_t<0>{}); - } - phase ^= 1; - __syncwarp(); - } - } - } else { - constexpr int kWarpsPerTile = kNumGetWarpsBwd * NC; // one get warp per CTA - int chunk = tile_elems / kWarpsPerTile; - int offset = chunk_idx * chunk; - int global_pe = nvshmem_team_translate_pe(bufs.team, info.pe, NVSHMEM_TEAM_WORLD); - void* p2p_remote = nvshmem_ptr(remote_tile, global_pe); - if (p2p_remote != nullptr) { - nvshmemx_getmem_warp(local_tile + offset, remote_tile + offset, - (size_t)chunk * sizeof(Element), global_pe); - } else { - if (chunk_idx == ib_seq % kWarpsPerTile) { - int valid_elems = - max(0, min(TileM, info.valid_rows)) * bufs.hidden_dim; - for (int i = valid_elems + lane; i < tile_elems; i += 32) - local_tile[i] = Element{}; - __syncwarp(); - if (valid_elems > 0) - nvshmemx_getmem_warp( - local_tile, remote_tile, - (size_t)valid_elems * sizeof(Element), global_pe); - } - ++ib_seq; - } - } - - if (tile_expert_ids_dst != nullptr && chunk_idx == 0 && lane == 0) { - tile_expert_ids_dst[slot] = info.expert; - if (tile_valid_rows_dst != nullptr) - tile_valid_rows_dst[slot] = info.valid_rows; - } - - __syncwarp(); - src_pipe.producer_release(lane); -} - -// ═══════════════════════════════════════════════════════════════════ -// do_get_bwd — gets one buffer (X or dY) for one tile -// ═══════════════════════════════════════════════════════════════════ -// -// kWarpsPerTile = kNumGetWarpsPerCta * NC = 2·NC (warps 1+2 across -// NC cooperating CTAs). Slot in flat list = xc + producer_stage*MC. - -// `tile_expert_ids_dst` is the per-pipe expert-id array to write to (one of -// bufs.tile_expert_ids_x or bufs.tile_expert_ids_dy, depending on which pipe -// this call is for), or nullptr to skip the write. Exactly one comm warp per -// tile is expected to pass non-null — caller picks the writer to avoid double -// writes and false sharing. -// -// The write happens AFTER the getmem completes and BEFORE producer_release, -// so it is part of the slot's atomic publish: any consumer that's seen -// producer_release sees both the getmem'd staging bytes AND the expert id. -template -__device__ __forceinline__ void do_get_bwd( - StagePipe& src_pipe, - const CommBuffersBwd& bufs, - const TileInfo& info, - Element* remote_base, - Element* staging_base, // base of flat [L][TileM][hidden_dim] - int tile_elems, - int xc, // start slot in flat list - int MC, // stride between consecutive comm slots - int chunk_idx, // 0..2·NC-1 - int& ib_seq, // IB-tile counter (round-robin turn; advanced here) - int lane, - int* tile_expert_ids_dst) { - - constexpr int kWarpsPerTile = kNumGetWarpsPerCta * NC; - int chunk = tile_elems / kWarpsPerTile; - int offset = chunk_idx * chunk; - - src_pipe.producer_acquire(); - - // Slot in the flat list. Comm strides MC across pipe stages; - // MLP strides gridDim.x. Both resolve to the same physical - // slot for a given global tile since L = MC·K = gridDim.x·NumStages. - int slot = xc + src_pipe.producer_stage() * MC; - - int global_pe = nvshmem_team_translate_pe(bufs.team, info.pe, NVSHMEM_TEAM_WORLD); - - // Tile base (offset 0). P2P chunks add `offset`; the IB whole-tile - // transfer addresses the base directly. - Element* local_tile = staging_base + slot * tile_elems; - Element* remote_tile = remote_base + info.token_offset * bufs.hidden_dim; - - // Fast path: NVLink P2P → all 2·NC warps cooperate, each copying its - // chunk with direct int4 loads. IB path: one whole-tile getmem issued - // by a single warp, chosen round-robin across IB tiles only (ib_seq - // skips local/P2P tiles and is advanced in lockstep by every warp). - // See the matching fwd do_get for the full rationale. - void* p2p_remote = nvshmem_ptr(remote_tile, global_pe); - if (p2p_remote != nullptr) { - static_assert(sizeof(Element) == 2, - "P2P int4 copy assumes 2-byte elements (bf16/half)"); - constexpr int kElemsPerVec = sizeof(int4) / sizeof(Element); - int n_vec = chunk / kElemsPerVec; - const int4* src_v = reinterpret_cast( - static_cast(p2p_remote) + offset); - int4* dst_v = reinterpret_cast(local_tile + offset); - #pragma unroll 4 - for (int i = lane; i < n_vec; i += 32) { - dst_v[i] = src_v[i]; - } - __syncwarp(); - } else { - if (chunk_idx == ib_seq % kWarpsPerTile) { - nvshmemx_getmem_warp(local_tile, remote_tile, - (size_t)tile_elems * sizeof(Element), global_pe); - } - ++ib_seq; - } - - if (tile_expert_ids_dst != nullptr && chunk_idx == 0 && lane == 0) { - tile_expert_ids_dst[slot] = info.expert; - } - __threadfence(); - - src_pipe.producer_release(lane); -} - -// ═══════════════════════════════════════════════════════════════════ -// do_put_bwd — puts dX for one tile back to remote PE -// ═══════════════════════════════════════════════════════════════════ -// -// kWarpsPerTile = NC. NumProducers = NSplit (MLP-side), -// NumConsumers = NC (comm-side). - -// num_put_warps / put_local_idx keep the single-put-warp layout -// (kWarpsPerTile = NC, slice = yc): warp 2 owns dX puts while warp 3 is -// reserved for the SM100 MLP path. -template -__device__ __forceinline__ void do_put_bwd( - StagePipe& dst_pipe, - const CommBuffersBwd& bufs, - const TileInfo& info, - Element* dx_staging, // base of flat [L][TileM][hidden_dim] - int tile_elems, - int xc, - int MC, - int chunk_idx, // 0..NC-1 (= yc) - int& ib_seq, // IB-tile counter (round-robin turn; advanced here) - int lane, - int num_put_warps = 1, // 1 (fallback) or kNumPutWarpsBwd (TMA mode) - int put_local_idx = 0) { // 0..num_put_warps-1 - - int kWarpsPerTile = num_put_warps * NC; - int put_idx = chunk_idx * num_put_warps + put_local_idx; // 0..kWarpsPerTile-1 - int chunk = tile_elems / kWarpsPerTile; - int offset = put_idx * chunk; - int my_bytes = chunk * sizeof(Element); - - dst_pipe.consumer_acquire(); - - int slot = xc + dst_pipe.consumer_stage() * MC; - auto* remote_dx = static_cast(bufs.remote_dx); - Element* local_tile = dx_staging + slot * tile_elems; - Element* remote_tile = remote_dx + info.token_offset * bufs.hidden_dim; - - int global_pe = nvshmem_team_translate_pe(bufs.team, info.pe, NVSHMEM_TEAM_WORLD); - - // NVLink P2P peers keep the cooperative split; IB peers send the whole - // tile as one RDMA issued by a single put warp, round-robin over IB - // tiles only (see do_get_bwd / fwd do_put). - void* p2p_remote = nvshmem_ptr(remote_tile, global_pe); - if (p2p_remote != nullptr) { - nvshmemx_putmem_warp(remote_tile + offset, local_tile + offset, - my_bytes, global_pe); - } else { - if (put_idx == ib_seq % kWarpsPerTile) { - int valid_elems = - max(0, min(TileM, info.valid_rows)) * bufs.hidden_dim; - if (valid_elems > 0) - nvshmemx_putmem_warp( - remote_tile, local_tile, - (size_t)valid_elems * sizeof(Element), global_pe); - } - ++ib_seq; - } - - dst_pipe.consumer_release(lane); -} - -// ═══════════════════════════════════════════════════════════════════ -// nvshmem_comm_prologue_bwd — fetch remote offsets, compute tile counts -// ═══════════════════════════════════════════════════════════════════ -// -// Mirrors fwd nvshmem_comm_prologue. Computes both counts since uneven -// distribution can give comm and MLP different per-CTA tile counts: -// smem.total_tiles = comm count (start = xc, stride = MC) -// smem.per_cta_tiles = MLP count (start = blockIdx.x, stride = gridDim.x) -// -// Runs on warps 2-3 (kCommWarpStart..kCommWarpEnd). Grid is 2D -// (gridDim.x, NSplit). - -__device__ __forceinline__ int select_runtime_nsplit_bwd_from_tiles( - int total_tiles, - int num_n_tiles_1, int num_n_tiles_2t, int num_n_tiles_5, - int num_k_tiles_1, int num_k_tiles_2t, int num_k_tiles_5, - int experts_per_pe, - int num_blocks, int static_nsplit) { - if (total_tiles <= 0) return static_nsplit; - constexpr int candidates[5] = {2, 4, 6, 8, 16}; - int best_ns = static_nsplit; - long long best_cost = -1; - long long best_waste = -1; - int best_n_gemm = -1; - long long w1 = 2LL * num_k_tiles_1; // mlp1 recomputes gate + up - long long w2t = (long long)num_k_tiles_2t; - long long w5 = 2LL * num_k_tiles_5; // mlp5 computes dU@B + dV@C - #pragma unroll - for (int i = 0; i < 5; ++i) { - int ns = candidates[i]; - if (ns > num_blocks) continue; - int n_gemm = (num_blocks / ns) * ns; - int ms = n_gemm / ns; - if (ms <= 0) continue; - int m_waves = (total_tiles + ms - 1) / ms; - int n1_waves = (num_n_tiles_1 + ns - 1) / ns; - int n2t_waves = (num_n_tiles_2t + ns - 1) / ns; - int n5_waves = (num_n_tiles_5 + ns - 1) / ns; - long long phase_cost = w1 * n1_waves + w2t * n2t_waves + w5 * n5_waves; - long long cost = (long long)m_waves * phase_cost; - int min_n_tiles = min(num_n_tiles_1, min(num_n_tiles_2t, num_n_tiles_5)); - bool fragmented_experts = (experts_per_pe >= 8) && (min_n_tiles <= 8); - if (fragmented_experts && ns < 4) { - long long frag = ((long long)experts_per_pe * 1024LL) - / (long long)max(total_tiles, 1); - cost += (cost * frag) / 128LL; - } - long long waste1 = - (long long)m_waves * ms * n1_waves * ns - - (long long)total_tiles * num_n_tiles_1; - long long waste2t = - (long long)m_waves * ms * n2t_waves * ns - - (long long)total_tiles * num_n_tiles_2t; - long long waste5 = - (long long)m_waves * ms * n5_waves * ns - - (long long)total_tiles * num_n_tiles_5; - long long waste = w1 * waste1 + w2t * waste2t + w5 * waste5; - if (best_cost < 0 || - cost < best_cost || - (cost == best_cost && waste < best_waste) || - (cost == best_cost && waste == best_waste && n_gemm > best_n_gemm)) { - best_cost = cost; - best_waste = waste; - best_n_gemm = n_gemm; - best_ns = ns; - } - } - return best_ns; -} - -template -__device__ __forceinline__ void nvshmem_comm_prologue_bwd( - CommSmem& smem, - const CommBuffersBwd& bufs, - int static_nsplit) { - - int warp_id = threadIdx.x / 32; - if (warp_id < kCommWarpStart || warp_id > kCommWarpEnd) - return; - - int lane = threadIdx.x % 32; - int local_warp = warp_id - kCommWarpStart; - bool is_leader = (local_warp == 0 && lane == 0); - - int experts_per_pe = bufs.experts_per_pe; - int num_pes = bufs.num_pes; - int my_pe = bufs.my_pe; - - if (bufs.all_expert_offsets == nullptr) { - if (is_leader) { - smem.total_tiles = 0; - smem.per_cta_tiles = 0; - smem.global_total = 0; - smem.runtime_nsplit = static_nsplit; - smem.runtime_n_gemm = bufs.n_gemm; - smem.runtime_grid_x = bufs.n_gemm / static_nsplit; - } - __threadfence_block(); - __syncwarp(); - return; - } - // NOTE: num_pes == 1 is NO LONGER short-circuited (mirrors fwd - // nvshmem_comm_prologue). The shared TileIterator enumerates the LOCAL PE - // (p=0), so a single-GPU bwd run stages its local X/dY through the comm - // get path and puts dX back through the comm put path exactly like the - // multi-GPU remote path — the separate local MLP pass is gone. - - int local_e_start = my_pe * experts_per_pe; - int src_stride = bufs.num_experts + 1; - int dst_stride = experts_per_pe + 1; - { - int tid = threadIdx.x - kCommThreadStart; - for (int pe = 0; pe < num_pes; ++pe) { - for (int i = tid; i <= experts_per_pe; i += (kCommWarpsPerCta * 32)) - smem.remote_offsets[pe * dst_stride + i] = - bufs.all_expert_offsets[pe * src_stride + local_e_start + i]; - if (bufs.all_expert_counts != nullptr) { - for (int i = tid; i < experts_per_pe; i += (kCommWarpsPerCta * 32)) - smem.remote_counts[pe * experts_per_pe + i] = - bufs.all_expert_counts[ - pe * bufs.num_experts + local_e_start + i]; - } - } - } - __syncwarp(); - - // Flat 1-D launch: cta_id = flat_id = blockIdx.x (gridDim.y == 1). - // Comm uses all NC-complete launched CTAs, decoupled from the MLP runtime - // NS. MLP count below is only the static-seed default; moe_bwd_kernel - // overwrites it after selecting runtime NS from the actual tile count. - int cta_id = blockIdx.x * gridDim.y + blockIdx.y; - int n_comm = (gridDim.x / NC) * NC; - int xc = cta_id / NC; - int MC = n_comm / NC; - int grid_x = bufs.n_gemm / static_nsplit; - int mlp_col = cta_id / static_nsplit; - - TileIterator iter; - iter.remote_offsets = smem.remote_offsets; - iter.offsets_stride = dst_stride; - iter.remote_counts = - bufs.all_expert_counts != nullptr ? smem.remote_counts : nullptr; - iter.counts_stride = experts_per_pe; - iter.experts_per_pe = experts_per_pe; - iter.num_pes = num_pes; - iter.my_pe = my_pe; - - // Comm count: this CTA's per-(xc, yc) workload. - iter.init(xc, MC); - int comm_count = iter.total_tiles; - - // MLP count: this CTA's logical column (mlp_col) workload, stride grid_x. - iter.init(mlp_col, grid_x); - int mlp_count = iter.total_tiles; - - // Grid-wide remote tile count (start=0, stride=1 → counts every tile - // once). Same value on every CTA — used by moe_fused_bwd to gate the - // remote phase consistently across the grid. - iter.init(0, 1); - int grid_total = iter.total_tiles; - - if (is_leader) { - smem.total_tiles = comm_count; // read by comm_main_bwd - smem.per_cta_tiles = mlp_count; // read by MLP iterator - smem.global_total = grid_total; // read by moe_fused_bwd gate - smem.runtime_nsplit = static_nsplit; - smem.runtime_n_gemm = bufs.n_gemm; - smem.runtime_grid_x = grid_x; - } - // Block-scope fence so the leader's smem writes above are visible - // to all threads in this CTA (including MLP warps that early- - // returned at the top) AFTER the caller's __syncthreads(). Without - // this, the compiler/HW may reorder MLP warps' loads of - // smem.global_total to before the comm warp's store, giving stale - // (zero-initialized) values and causing CTA(0,0) to skip the remote - // phase while peer CTAs proceed — manifests as a multi-PE deadlock. - __threadfence_block(); - __syncwarp(); -} - -// ═══════════════════════════════════════════════════════════════════ -// nvshmem_comm_main_bwd — independent flat loops per warp role -// ═══════════════════════════════════════════════════════════════════ -// -// All three pipes are per-tile StagePipes (mirrors fwd). Comm at xc -// walks slots {xc, xc+MC, ..., xc+(K-1)·MC} of the [L]-sized signal -// arrays (stride MC). Get pipes use NumProducers = 2·NC, -// NumConsumers = NSplit; put pipe uses NumProducers = NSplit, -// NumConsumers = NC. - -template < - typename Element, - int TileM, - int NumStages, - int NC = 2, - int BoxRows = kGetBoxRowsBwdSm90> -__device__ __forceinline__ void nvshmem_comm_main_bwd( - CommSmem& smem, - const CommBuffersBwd& bufs, - const GetTmaDescsBwd* descs = nullptr, // nullptr/disabled → getmem layout - char* get_bounce = nullptr) { // bwd TMA bounce smem region - - // Comm-side pipe depth is NumStages — one comm CTA cycle covers - // the full MC × NumStages staging area, independent of NSplit and - // NC. The MLP side is ticket-based (RemoteMlpTileIteratorBwd) and - // shares no per-stage_id state, so no divisibility relationship is - // required between NumStages, NSplit and NC. - constexpr int K = NumStages; - - int warp_id = threadIdx.x / 32; - if (warp_id != kCommBwdGetWarp0 && warp_id != kCommBwdPutWarp) - return; - - int lane = threadIdx.x % 32; - bool is_leader = (lane == 0); - - // Flat 1-D launch: cta_id = flat_id = blockIdx.x (gridDim.y == 1). - int cta_id = blockIdx.x * gridDim.y + blockIdx.y; - int n_comm = (gridDim.x / NC) * NC; - if (cta_id >= n_comm) return; - int xc = cta_id / NC; - int yc = cta_id % NC; - int MC = n_comm / NC; - int runtime_nsplit = smem.runtime_nsplit; - - int total = smem.total_tiles; - if (total == 0) return; - - int tile_elems = TileM * bufs.hidden_dim; - - TileIterator iter; - iter.remote_offsets = smem.remote_offsets; - iter.offsets_stride = bufs.experts_per_pe + 1; - iter.remote_counts = - bufs.all_expert_counts != nullptr ? smem.remote_counts : nullptr; - iter.counts_stride = bufs.experts_per_pe; - iter.experts_per_pe = bufs.experts_per_pe; - iter.num_pes = bufs.num_pes; - iter.my_pe = bufs.my_pe; - iter.init(xc, MC); - - int my_total = iter.total_tiles; - if (my_total == 0) return; - - auto* x_staging = static_cast(bufs.x_staging); - auto* dy_staging = static_cast(bufs.dy_staging); - auto* dx_staging = static_cast(bufs.dx_staging); - auto* remote_x_base = static_cast(bufs.remote_x); - auto* remote_dy_base = static_cast(bufs.remote_dy); - - // Unified layout: warp 1 gets X then dY for each tile; warp 2 puts dX. - // do_get_bwd_tma chooses per tile: same-host TMA, P2P fallback, or IB - // whole-tile round-robin. This keeps X(t), dY(t), X(t+1), ... in one loop. - if (warp_id == kCommBwdGetWarp0) { - StagePipe x_pipe; - x_pipe.init(bufs.x_src_ready + xc, bufs.x_src_consumed + xc, - is_leader, MC, kNumGetWarpsBwd * NC, runtime_nsplit); - - StagePipe dy_pipe; - dy_pipe.init(bufs.dy_src_ready + xc, bufs.dy_src_consumed + xc, - is_leader, MC, kNumGetWarpsBwd * NC, runtime_nsplit); - - int chunk_idx = yc; // kNumGetWarpsBwd == 1 - // Exactly one writer per tile per pipe: warp 1 (kCommBwdGetWarp0), - // yc=0, lane=0. The X-pipe get writes tile_expert_ids_x; the dY-pipe - // get writes tile_expert_ids_dy. Both arrays end up holding the same - // info.expert for any given live slot — but each is gated by its - // own pipe's release window, so MLP can use the X-pipe array for - // Phase 2 mlp4 and the dY-pipe array for Phase 2 mlp3 with - // asymmetric release timing (release_src after mlp4, release_dy - // after mlp3) and no race in either read. - bool is_expert_id_writer = - (warp_id == kCommBwdGetWarp0) && (yc == 0); - int* x_expert_dst = is_expert_id_writer ? bufs.tile_expert_ids_x : nullptr; - int* dy_expert_dst = is_expert_id_writer ? bufs.tile_expert_ids_dy : nullptr; - int* x_valid_dst = is_expert_id_writer ? bufs.tile_valid_rows_x : nullptr; - - // One IB-tile counter shared by the X and dY gets, advanced per IB - // getmem so X(t), dY(t), X(t+1), … rotate across the NC get warps. - int ib_seq = 0; - for (int i = 0; i < my_total; ++i) { - TileInfo info = iter.next(); - - do_get_bwd_tma( - x_pipe, bufs, info, remote_x_base, x_staging, tile_elems, - xc, MC, chunk_idx, ib_seq, lane, - /*mbar_index=*/0, - x_expert_dst, x_valid_dst, descs, - descs ? descs->src_x_desc : nullptr, - descs ? &descs->dst_x_staging_desc : nullptr, - get_bounce); - - do_get_bwd_tma( - dy_pipe, bufs, info, remote_dy_base, dy_staging, tile_elems, - xc, MC, chunk_idx, ib_seq, lane, - /*mbar_index=*/1, - dy_expert_dst, nullptr, descs, - descs ? descs->src_dy_desc : nullptr, - descs ? &descs->dst_dy_staging_desc : nullptr, - get_bounce); - } - } else { // warp 2 → dX put warp - StagePipe dst_pipe; - dst_pipe.init(bufs.dst_ready + xc, bufs.dst_consumed + xc, - is_leader, MC, runtime_nsplit, kNumPutWarpsBwd * NC); - - int put_local_idx = 0; - int chunk_idx = yc; - int ib_seq = 0; // advances only on IB tiles (inside do_put_bwd) - for (int i = 0; i < my_total; ++i) { - TileInfo info = iter.next(); - do_put_bwd( - dst_pipe, bufs, info, - dx_staging, tile_elems, xc, MC, chunk_idx, ib_seq, lane, - kNumPutWarpsBwd, put_local_idx); - } - } -} - -} // namespace liger +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif + +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "mlp_comms_bwd_sm90.cuh" +#else +#include "mlp_comms_bwd_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm100.cuh new file mode 100644 index 00000000..291403af --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm100.cuh @@ -0,0 +1,772 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Backward MoE communication: gets X + dY, puts dX +// ═══════════════════════════════════════════════════════════════════ +// +// Warp roles: +// warp 1 → X/dY get warp +// warp 2 → dX put warp +// warp 3 → MLP-side SM100 UMMA producer (exits early on SM90) +// +// Grid is launched 2D as (gridDim.x, NSplit) where NSplit is the +// MLP-side cooperative count. Comms re-indexes the flat CTA id into +// a (MC, NC) logical grid: +// +// cta_id = blockIdx.x * gridDim.y + blockIdx.y +// xc = cta_id / NC in [0, MC) +// yc = cta_id % NC in [0, NC) +// MC = (gridDim.x * gridDim.y) / NC (multiply before divide so +// MC is exact when NC ∤ NSplit) +// +// NC CTAs (same xc, different yc) cooperate on one tile per get/put. +// MC > gridDim.x when NC < NSplit so excess comm CTAs run ahead and +// load the next tiles while the MLP is still on the current one. +// +// ── Tile-level locking (same as fwd) ─────────────────────────────── +// +// All three pipes (X, dY, dX) use the per-tile StagePipe over a flat +// [L = MC * NumStages]-sized signal array. Get pipes have +// NumProducers = kNumGetWarpsPerCta·NC, NumConsumers = NSplit. Put +// pipe has NumProducers = NSplit, NumConsumers = NC. Comm strides MC +// across stages (NumStages entries — full ring cycle, independent of +// NSplit and NC); MLP is ticket-based (RemoteMlpTileIteratorBwd) and +// keys on the absolute slot index = (m_base + idx·gridDim.x) mod L. +// Both sides cover the same flat L slots. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp_comms_sm100.cuh" + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// Bwd comm warp constants +// ═══════════════════════════════════════════════════════════════════ +// +// Main loop runs on warps 1-2 only. kCommBwdGetWarp1 remains for legacy +// prologue compatibility with shared helpers; it is not a main-loop role. + +static constexpr int kCommBwdGetWarp0 = 1; +static constexpr int kCommBwdGetWarp1 = 2; // prologue compatibility only +static constexpr int kCommBwdPutWarp = 2; +static constexpr int kCommBwdWarpStart = kCommBwdGetWarp0; +static constexpr int kCommBwdWarpEnd = kCommBwdPutWarp; + +// ── TMA-GET (NVLink) warp rebalance ───────────────────────────────── +// Mirrors the fwd #111 layout: one warp saturates the TMA GET path, so warp 1 +// gets X+dY and warp 2 puts dX. Warp 3 enters the MLP-side branch and exits +// there for Compute != 100. The MLP-side iterator picks producer/consumer +// counts from the same runtime tma_enabled flag. +static constexpr int kNumGetWarpsBwd = 1; // TMA mode: 1 get warp (X+dY) +static constexpr int kNumPutWarpsBwd = 1; // TMA mode: 1 put warp (dX) + +// ═══════════════════════════════════════════════════════════════════ +// GetTmaDescsBwd — raw CUtensorMap descriptors for the bwd TMA GET +// ═══════════════════════════════════════════════════════════════════ +// +// Per-same-host-peer source maps over nvshmem_ptr(x_sorted, peer) and +// nvshmem_ptr(dy_sorted, peer); one local dst map each over x_staging / +// dy_staging. Built host-side in moe_bwd.cu, passed __grid_constant__. +// Same as fwd GetTmaDescs: local peers index by pe % gpus_per_node; remote +// peers fall back to IB getmem in do_get_bwd_tma. +struct GetTmaDescsBwd { + CUtensorMap src_x_desc[kGetMaxPes]; // index by (TileInfo::pe % gpus_per_node) + CUtensorMap src_dy_desc[kGetMaxPes]; + CUtensorMap dst_x_staging_desc; // local x_staging + CUtensorMap dst_dy_staging_desc; // local dy_staging + int enabled; // 1 = local-peer TMA path available + int gpus_per_node; + int my_pe; +}; + +__host__ __device__ __forceinline__ bool get_pe_is_local_bwd( + const GetTmaDescsBwd* descs, int pe) { + if (descs == nullptr || descs->gpus_per_node <= 0) return true; + return (pe / descs->gpus_per_node) == (descs->my_pe / descs->gpus_per_node); +} + +// Box height (rows) of a single backward TMA transfer. SM100 mirrors +// forward's full [32×256] row-band box; SM90 uses a smaller 8-row box. +static constexpr int kGetBoxRowsBwdSm90 = 8; +static constexpr int kGetBoxRowsBwdSm100 = 32; + +// X and dY reuse one data box sequentially but keep distinct mbarriers. +// Reinitializing one mbarrier across the two TMA streams races the async proxy; +// the box itself is safe to reuse after cp.async.bulk.wait_group<0>. +template +struct GetBounceBwd { + static constexpr int kBoxRows = BoxRows; + static constexpr int kBoxElems = kBoxRows * kGetKChunk; + static constexpr int kBoxBytes = kBoxElems * (int)sizeof(Element); + static constexpr int kMbarOffset = kBoxBytes; + static constexpr int kMbarStride = sizeof(uint64_t); + static constexpr int kTotalBytes = + (kBoxBytes + 2 * kMbarStride + 127) & ~127; +}; + +// do_get_bwd_tma is defined below, after CommBuffersBwd (it needs the +// complete struct type for bufs.hidden_dim). + +// ═══════════════════════════════════════════════════════════════════ +// CommBuffersBwd — backward-specific buffer pointers +// ═══════════════════════════════════════════════════════════════════ + +struct CommBuffersBwd { + // Staging: flat slot list [L][TileM][hidden_dim] where + // L = gridDim.x * NumStages. Indexed by slot ∈ [0, L). + void* x_staging; // gets: X from remote PE + void* dy_staging; // gets: dY from remote PE + void* dx_staging; // puts: dX back to remote PE + + // Per-tile signals ([L] entries each), one counter pair per slot — + // same granularity as fwd src/dst pipes. + int* x_src_ready; + int* x_src_consumed; + int* dy_src_ready; + int* dy_src_consumed; + int* dst_ready; + int* dst_consumed; + + // Per-slot expert IDs ([L] each), one array per get pipe. + // + // They hold the SAME value for any given live slot — info.expert at the + // time the comm warp wrote that slot — but each is gated by its own + // pipe's StagePipe release window: + // + // tile_expert_ids_x is overwritten between MLP's release_src and + // comm's next-lap producer_release on the X pipe. + // Phase 2 mlp4 reads this (mlp4 reads X) and is + // safe as long as release_src is deferred past mlp4. + // + // tile_expert_ids_dy is overwritten between MLP's release_dy and + // comm's next-lap producer_release on the dY pipe. + // Phase 2 mlp3 reads this (mlp3 reads dY) and is + // safe as long as release_dy is deferred past mlp3. + // + // Splitting the array is what lets release_src happen between mlp4 and + // mlp3 without races: comm refilling tile_expert_ids_x in that window + // can't poison mlp3, because mlp3 reads the dY-pipe array. + int* tile_expert_ids_x; + int* tile_expert_ids_dy; + int* tile_valid_rows_x; + + // Remote data pointers (symmetric memory) + void* remote_x; // get source: X on remote PE + void* remote_dy; // get source: dY on remote PE + void* remote_dx; // put destination: dX on remote PE + + // Reused from forward + int* all_expert_offsets; + int* all_expert_counts; + int all_expert_offsets_stride; + + int my_pe; + int num_pes; + nvshmem_team_t team; + + int hidden_dim; + int num_experts; + int experts_per_pe; + + // GEMM/comm-active CTA count = floor_NS = grid_x · NSplit (NOT the + // launched grid size). Under the flat 1-D launch gridDim.x = num_blocks + // (= num_sms) ≥ n_gemm, so MC = n_gemm / NC and the MLP column count + // grid_x = n_gemm / NSplit must come from here, not from gridDim. + int n_gemm; + + // Used by the dummy_consumer_bwd in mlp_comms_bwd.cu when it replicates + // the real BWD's batch loop (mlp_global_barrier + mlp_x_barrier). Not + // used by the comm warps themselves. + int* barrier_counter; + int* phase_counter; +}; + +// ═══════════════════════════════════════════════════════════════════ +// do_get_bwd_tma — TMA GET of one buffer (X or dY) for one tile +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors fwd do_get's TMA path (cuda::ptx intrinsics): peer → smem box → +// local staging, via TMA load + TMA store. Because the bounce box is small +// (kGetBoxRowsBwd rows), the single get warp loops over its TileM/NC-row +// band in kGetBoxRowsBwd-tall sub-bands × kGetKChunk-wide column boxes. +// chunk_idx == yc (the CTA's slot within its NC-cooperating group); with +// kNumGetWarpsBwd == 1 the per-warp row band is TileM/NC. +template +__device__ __forceinline__ void do_get_bwd_tma( + StagePipe& src_pipe, + const CommBuffersBwd& bufs, + const TileInfo& info, + Element* remote_base, + Element* staging_base, + int tile_elems, + int xc, + int MC, + int chunk_idx, // = yc (0..NC-1) + int& ib_seq, // advanced once per actual IB get call + int lane, + int mbar_index, + int* tile_expert_ids_dst, + int* tile_valid_rows_dst, + const GetTmaDescsBwd* descs, + const CUtensorMap* src_desc_per_pe, // descs->src_x_desc or src_dy_desc + const CUtensorMap* dst_desc, // descs->dst_x_staging or dst_dy_staging + char* bounce_warp) { + + namespace ptx = cuda::ptx; + using Bounce = GetBounceBwd; + constexpr int kWarpsPerTile = kNumGetWarpsBwd * NC; // = NC + constexpr int kRowBand = TileM / kWarpsPerTile; // rows this warp owns + constexpr int kBoxRows = Bounce::kBoxRows; + + src_pipe.producer_acquire(); + + int slot = xc + src_pipe.producer_stage() * MC; + Element* local_tile = staging_base + slot * tile_elems; + Element* remote_tile = remote_base + info.token_offset * bufs.hidden_dim; + + bool is_local = descs && descs->enabled && get_pe_is_local_bwd(descs, info.pe); + bool tma_ok = is_local && (bufs.hidden_dim % kGetKChunk == 0) && bounce_warp != nullptr; + + if (tma_ok) { + Element* sbuf = reinterpret_cast(bounce_warp); + uint64_t* mbar = reinterpret_cast( + bounce_warp + Bounce::kMbarOffset + + mbar_index * Bounce::kMbarStride); + + if (lane == 0) + ptx::mbarrier_init(mbar, 1u); + __syncwarp(); + ptx::fence_proxy_async(ptx::space_shared); + + const int n_k = bufs.hidden_dim / kGetKChunk; + const uint32_t txbytes = (uint32_t)Bounce::kBoxBytes; + int y_src_base = info.token_offset + chunk_idx * kRowBand; // peer row + int y_dst_base = slot * TileM + chunk_idx * kRowBand; // staging row + const int local_peer = info.pe % descs->gpus_per_node; + + uint32_t phase = 0; + for (int rb = 0; rb < kRowBand; rb += kBoxRows) { + for (int k = 0; k < n_k; ++k) { + int x = k * kGetKChunk; + if (lane == 0) { + (void)ptx::mbarrier_arrive_expect_tx( + ptx::sem_release, ptx::scope_cta, ptx::space_shared, + mbar, txbytes); + int32_t load_crd[2] = {x, y_src_base + rb}; + ptx::cp_async_bulk_tensor( + ptx::space_cluster, ptx::space_global, + sbuf, &src_desc_per_pe[local_peer], load_crd, mbar); + while (!ptx::mbarrier_try_wait_parity(mbar, phase)) {} + int32_t store_crd[2] = {x, y_dst_base + rb}; + ptx::cp_async_bulk_tensor( + ptx::space_global, ptx::space_shared, + dst_desc, store_crd, sbuf); + ptx::cp_async_bulk_commit_group(); + ptx::cp_async_bulk_wait_group(ptx::n32_t<0>{}); + } + phase ^= 1; + __syncwarp(); + } + } + } else { + constexpr int kWarpsPerTile = kNumGetWarpsBwd * NC; // one get warp per CTA + int chunk = tile_elems / kWarpsPerTile; + int offset = chunk_idx * chunk; + int global_pe = nvshmem_team_translate_pe(bufs.team, info.pe, NVSHMEM_TEAM_WORLD); + void* p2p_remote = nvshmem_ptr(remote_tile, global_pe); + if (p2p_remote != nullptr) { + nvshmemx_getmem_warp(local_tile + offset, remote_tile + offset, + (size_t)chunk * sizeof(Element), global_pe); + } else { + if (chunk_idx == ib_seq % kWarpsPerTile) { + int valid_elems = + max(0, min(TileM, info.valid_rows)) * bufs.hidden_dim; + for (int i = valid_elems + lane; i < tile_elems; i += 32) + local_tile[i] = Element{}; + __syncwarp(); + if (valid_elems > 0) + nvshmemx_getmem_warp( + local_tile, remote_tile, + (size_t)valid_elems * sizeof(Element), global_pe); + } + ++ib_seq; + } + } + + if (tile_expert_ids_dst != nullptr && chunk_idx == 0 && lane == 0) { + tile_expert_ids_dst[slot] = info.expert; + if (tile_valid_rows_dst != nullptr) + tile_valid_rows_dst[slot] = info.valid_rows; + } + + __syncwarp(); + src_pipe.producer_release(lane); +} + +// ═══════════════════════════════════════════════════════════════════ +// do_get_bwd — gets one buffer (X or dY) for one tile +// ═══════════════════════════════════════════════════════════════════ +// +// kWarpsPerTile = kNumGetWarpsPerCta * NC = 2·NC (warps 1+2 across +// NC cooperating CTAs). Slot in flat list = xc + producer_stage*MC. + +// `tile_expert_ids_dst` is the per-pipe expert-id array to write to (one of +// bufs.tile_expert_ids_x or bufs.tile_expert_ids_dy, depending on which pipe +// this call is for), or nullptr to skip the write. Exactly one comm warp per +// tile is expected to pass non-null — caller picks the writer to avoid double +// writes and false sharing. +// +// The write happens AFTER the getmem completes and BEFORE producer_release, +// so it is part of the slot's atomic publish: any consumer that's seen +// producer_release sees both the getmem'd staging bytes AND the expert id. +template +__device__ __forceinline__ void do_get_bwd( + StagePipe& src_pipe, + const CommBuffersBwd& bufs, + const TileInfo& info, + Element* remote_base, + Element* staging_base, // base of flat [L][TileM][hidden_dim] + int tile_elems, + int xc, // start slot in flat list + int MC, // stride between consecutive comm slots + int chunk_idx, // 0..2·NC-1 + int& ib_seq, // IB-tile counter (round-robin turn; advanced here) + int lane, + int* tile_expert_ids_dst) { + + constexpr int kWarpsPerTile = kNumGetWarpsPerCta * NC; + int chunk = tile_elems / kWarpsPerTile; + int offset = chunk_idx * chunk; + + src_pipe.producer_acquire(); + + // Slot in the flat list. Comm strides MC across pipe stages; + // MLP strides gridDim.x. Both resolve to the same physical + // slot for a given global tile since L = MC·K = gridDim.x·NumStages. + int slot = xc + src_pipe.producer_stage() * MC; + + int global_pe = nvshmem_team_translate_pe(bufs.team, info.pe, NVSHMEM_TEAM_WORLD); + + // Tile base (offset 0). P2P chunks add `offset`; the IB whole-tile + // transfer addresses the base directly. + Element* local_tile = staging_base + slot * tile_elems; + Element* remote_tile = remote_base + info.token_offset * bufs.hidden_dim; + + // Fast path: NVLink P2P → all 2·NC warps cooperate, each copying its + // chunk with direct int4 loads. IB path: one whole-tile getmem issued + // by a single warp, chosen round-robin across IB tiles only (ib_seq + // skips local/P2P tiles and is advanced in lockstep by every warp). + // See the matching fwd do_get for the full rationale. + void* p2p_remote = nvshmem_ptr(remote_tile, global_pe); + if (p2p_remote != nullptr) { + static_assert(sizeof(Element) == 2, + "P2P int4 copy assumes 2-byte elements (bf16/half)"); + constexpr int kElemsPerVec = sizeof(int4) / sizeof(Element); + int n_vec = chunk / kElemsPerVec; + const int4* src_v = reinterpret_cast( + static_cast(p2p_remote) + offset); + int4* dst_v = reinterpret_cast(local_tile + offset); + #pragma unroll 4 + for (int i = lane; i < n_vec; i += 32) { + dst_v[i] = src_v[i]; + } + __syncwarp(); + } else { + if (chunk_idx == ib_seq % kWarpsPerTile) { + nvshmemx_getmem_warp(local_tile, remote_tile, + (size_t)tile_elems * sizeof(Element), global_pe); + } + ++ib_seq; + } + + if (tile_expert_ids_dst != nullptr && chunk_idx == 0 && lane == 0) { + tile_expert_ids_dst[slot] = info.expert; + } + __threadfence(); + + src_pipe.producer_release(lane); +} + +// ═══════════════════════════════════════════════════════════════════ +// do_put_bwd — puts dX for one tile back to remote PE +// ═══════════════════════════════════════════════════════════════════ +// +// kWarpsPerTile = NC. NumProducers = NSplit (MLP-side), +// NumConsumers = NC (comm-side). + +// num_put_warps / put_local_idx keep the single-put-warp layout +// (kWarpsPerTile = NC, slice = yc): warp 2 owns dX puts while warp 3 is +// reserved for the SM100 MLP path. +template +__device__ __forceinline__ void do_put_bwd( + StagePipe& dst_pipe, + const CommBuffersBwd& bufs, + const TileInfo& info, + Element* dx_staging, // base of flat [L][TileM][hidden_dim] + int tile_elems, + int xc, + int MC, + int chunk_idx, // 0..NC-1 (= yc) + int& ib_seq, // IB-tile counter (round-robin turn; advanced here) + int lane, + int num_put_warps = 1, // 1 (fallback) or kNumPutWarpsBwd (TMA mode) + int put_local_idx = 0) { // 0..num_put_warps-1 + + int kWarpsPerTile = num_put_warps * NC; + int put_idx = chunk_idx * num_put_warps + put_local_idx; // 0..kWarpsPerTile-1 + int chunk = tile_elems / kWarpsPerTile; + int offset = put_idx * chunk; + int my_bytes = chunk * sizeof(Element); + + dst_pipe.consumer_acquire(); + + int slot = xc + dst_pipe.consumer_stage() * MC; + auto* remote_dx = static_cast(bufs.remote_dx); + Element* local_tile = dx_staging + slot * tile_elems; + Element* remote_tile = remote_dx + info.token_offset * bufs.hidden_dim; + + int global_pe = nvshmem_team_translate_pe(bufs.team, info.pe, NVSHMEM_TEAM_WORLD); + + // NVLink P2P peers keep the cooperative split; IB peers send the whole + // tile as one RDMA issued by a single put warp, round-robin over IB + // tiles only (see do_get_bwd / fwd do_put). + void* p2p_remote = nvshmem_ptr(remote_tile, global_pe); + if (p2p_remote != nullptr) { + nvshmemx_putmem_warp(remote_tile + offset, local_tile + offset, + my_bytes, global_pe); + } else { + if (put_idx == ib_seq % kWarpsPerTile) { + int valid_elems = + max(0, min(TileM, info.valid_rows)) * bufs.hidden_dim; + if (valid_elems > 0) + nvshmemx_putmem_warp( + remote_tile, local_tile, + (size_t)valid_elems * sizeof(Element), global_pe); + } + ++ib_seq; + } + + dst_pipe.consumer_release(lane); +} + +// ═══════════════════════════════════════════════════════════════════ +// nvshmem_comm_prologue_bwd — fetch remote offsets, compute tile counts +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors fwd nvshmem_comm_prologue. Computes both counts since uneven +// distribution can give comm and MLP different per-CTA tile counts: +// smem.total_tiles = comm count (start = xc, stride = MC) +// smem.per_cta_tiles = MLP count (start = blockIdx.x, stride = gridDim.x) +// +// Runs on warps 2-3 (kCommWarpStart..kCommWarpEnd). Grid is 2D +// (gridDim.x, NSplit). + +__device__ __forceinline__ int select_runtime_nsplit_bwd_from_tiles( + int total_tiles, + int num_n_tiles_1, int num_n_tiles_2t, int num_n_tiles_5, + int num_k_tiles_1, int num_k_tiles_2t, int num_k_tiles_5, + int experts_per_pe, + int num_blocks, int static_nsplit) { + if (total_tiles <= 0) return static_nsplit; + constexpr int candidates[5] = {2, 4, 6, 8, 16}; + int best_ns = static_nsplit; + long long best_cost = -1; + long long best_waste = -1; + int best_n_gemm = -1; + long long w1 = 2LL * num_k_tiles_1; // mlp1 recomputes gate + up + long long w2t = (long long)num_k_tiles_2t; + long long w5 = 2LL * num_k_tiles_5; // mlp5 computes dU@B + dV@C + #pragma unroll + for (int i = 0; i < 5; ++i) { + int ns = candidates[i]; + if (ns > num_blocks) continue; + int n_gemm = (num_blocks / ns) * ns; + int ms = n_gemm / ns; + if (ms <= 0) continue; + int m_waves = (total_tiles + ms - 1) / ms; + int n1_waves = (num_n_tiles_1 + ns - 1) / ns; + int n2t_waves = (num_n_tiles_2t + ns - 1) / ns; + int n5_waves = (num_n_tiles_5 + ns - 1) / ns; + long long phase_cost = w1 * n1_waves + w2t * n2t_waves + w5 * n5_waves; + long long cost = (long long)m_waves * phase_cost; + int min_n_tiles = min(num_n_tiles_1, min(num_n_tiles_2t, num_n_tiles_5)); + bool fragmented_experts = (experts_per_pe >= 8) && (min_n_tiles <= 8); + if (fragmented_experts && ns < 4) { + long long frag = ((long long)experts_per_pe * 1024LL) + / (long long)max(total_tiles, 1); + cost += (cost * frag) / 128LL; + } + long long waste1 = + (long long)m_waves * ms * n1_waves * ns + - (long long)total_tiles * num_n_tiles_1; + long long waste2t = + (long long)m_waves * ms * n2t_waves * ns + - (long long)total_tiles * num_n_tiles_2t; + long long waste5 = + (long long)m_waves * ms * n5_waves * ns + - (long long)total_tiles * num_n_tiles_5; + long long waste = w1 * waste1 + w2t * waste2t + w5 * waste5; + if (best_cost < 0 || + cost < best_cost || + (cost == best_cost && waste < best_waste) || + (cost == best_cost && waste == best_waste && n_gemm > best_n_gemm)) { + best_cost = cost; + best_waste = waste; + best_n_gemm = n_gemm; + best_ns = ns; + } + } + return best_ns; +} + +template +__device__ __forceinline__ void nvshmem_comm_prologue_bwd( + CommSmem& smem, + const CommBuffersBwd& bufs, + int static_nsplit) { + + int warp_id = threadIdx.x / 32; + if (warp_id < kCommWarpStart || warp_id > kCommWarpEnd) + return; + + int lane = threadIdx.x % 32; + int local_warp = warp_id - kCommWarpStart; + bool is_leader = (local_warp == 0 && lane == 0); + + int experts_per_pe = bufs.experts_per_pe; + int num_pes = bufs.num_pes; + int my_pe = bufs.my_pe; + + if (bufs.all_expert_offsets == nullptr) { + if (is_leader) { + smem.total_tiles = 0; + smem.per_cta_tiles = 0; + smem.global_total = 0; + smem.runtime_nsplit = static_nsplit; + smem.runtime_n_gemm = bufs.n_gemm; + smem.runtime_grid_x = bufs.n_gemm / static_nsplit; + } + __threadfence_block(); + __syncwarp(); + return; + } + // NOTE: num_pes == 1 is NO LONGER short-circuited (mirrors fwd + // nvshmem_comm_prologue). The shared TileIterator enumerates the LOCAL PE + // (p=0), so a single-GPU bwd run stages its local X/dY through the comm + // get path and puts dX back through the comm put path exactly like the + // multi-GPU remote path — the separate local MLP pass is gone. + + int local_e_start = my_pe * experts_per_pe; + int src_stride = bufs.num_experts + 1; + int dst_stride = experts_per_pe + 1; + { + int tid = threadIdx.x - kCommThreadStart; + for (int pe = 0; pe < num_pes; ++pe) { + for (int i = tid; i <= experts_per_pe; i += (kCommWarpsPerCta * 32)) + smem.remote_offsets[pe * dst_stride + i] = + bufs.all_expert_offsets[pe * src_stride + local_e_start + i]; + if (bufs.all_expert_counts != nullptr) { + for (int i = tid; i < experts_per_pe; i += (kCommWarpsPerCta * 32)) + smem.remote_counts[pe * experts_per_pe + i] = + bufs.all_expert_counts[ + pe * bufs.num_experts + local_e_start + i]; + } + } + } + __syncwarp(); + + // Flat 1-D launch: cta_id = flat_id = blockIdx.x (gridDim.y == 1). + // Comm uses all NC-complete launched CTAs, decoupled from the MLP runtime + // NS. MLP count below is only the static-seed default; moe_bwd_kernel + // overwrites it after selecting runtime NS from the actual tile count. + int cta_id = blockIdx.x * gridDim.y + blockIdx.y; + int n_comm = (gridDim.x / NC) * NC; + int xc = cta_id / NC; + int MC = n_comm / NC; + int grid_x = bufs.n_gemm / static_nsplit; + int mlp_col = cta_id / static_nsplit; + + TileIterator iter; + iter.remote_offsets = smem.remote_offsets; + iter.offsets_stride = dst_stride; + iter.remote_counts = + bufs.all_expert_counts != nullptr ? smem.remote_counts : nullptr; + iter.counts_stride = experts_per_pe; + iter.experts_per_pe = experts_per_pe; + iter.num_pes = num_pes; + iter.my_pe = my_pe; + + // Comm count: this CTA's per-(xc, yc) workload. + iter.init(xc, MC); + int comm_count = iter.total_tiles; + + // MLP count: this CTA's logical column (mlp_col) workload, stride grid_x. + iter.init(mlp_col, grid_x); + int mlp_count = iter.total_tiles; + + // Grid-wide remote tile count (start=0, stride=1 → counts every tile + // once). Same value on every CTA — used by moe_fused_bwd to gate the + // remote phase consistently across the grid. + iter.init(0, 1); + int grid_total = iter.total_tiles; + + if (is_leader) { + smem.total_tiles = comm_count; // read by comm_main_bwd + smem.per_cta_tiles = mlp_count; // read by MLP iterator + smem.global_total = grid_total; // read by moe_fused_bwd gate + smem.runtime_nsplit = static_nsplit; + smem.runtime_n_gemm = bufs.n_gemm; + smem.runtime_grid_x = grid_x; + } + // Block-scope fence so the leader's smem writes above are visible + // to all threads in this CTA (including MLP warps that early- + // returned at the top) AFTER the caller's __syncthreads(). Without + // this, the compiler/HW may reorder MLP warps' loads of + // smem.global_total to before the comm warp's store, giving stale + // (zero-initialized) values and causing CTA(0,0) to skip the remote + // phase while peer CTAs proceed — manifests as a multi-PE deadlock. + __threadfence_block(); + __syncwarp(); +} + +// ═══════════════════════════════════════════════════════════════════ +// nvshmem_comm_main_bwd — independent flat loops per warp role +// ═══════════════════════════════════════════════════════════════════ +// +// All three pipes are per-tile StagePipes (mirrors fwd). Comm at xc +// walks slots {xc, xc+MC, ..., xc+(K-1)·MC} of the [L]-sized signal +// arrays (stride MC). Get pipes use NumProducers = 2·NC, +// NumConsumers = NSplit; put pipe uses NumProducers = NSplit, +// NumConsumers = NC. + +template < + typename Element, + int TileM, + int NumStages, + int NC = 2, + int BoxRows = kGetBoxRowsBwdSm90> +__device__ __forceinline__ void nvshmem_comm_main_bwd( + CommSmem& smem, + const CommBuffersBwd& bufs, + const GetTmaDescsBwd* descs = nullptr, // nullptr/disabled → getmem layout + char* get_bounce = nullptr) { // bwd TMA bounce smem region + + // Comm-side pipe depth is NumStages — one comm CTA cycle covers + // the full MC × NumStages staging area, independent of NSplit and + // NC. The MLP side is ticket-based (RemoteMlpTileIteratorBwd) and + // shares no per-stage_id state, so no divisibility relationship is + // required between NumStages, NSplit and NC. + constexpr int K = NumStages; + + int warp_id = threadIdx.x / 32; + if (warp_id != kCommBwdGetWarp0 && warp_id != kCommBwdPutWarp) + return; + + int lane = threadIdx.x % 32; + bool is_leader = (lane == 0); + + // Flat 1-D launch: cta_id = flat_id = blockIdx.x (gridDim.y == 1). + int cta_id = blockIdx.x * gridDim.y + blockIdx.y; + int n_comm = (gridDim.x / NC) * NC; + if (cta_id >= n_comm) return; + int xc = cta_id / NC; + int yc = cta_id % NC; + int MC = n_comm / NC; + int runtime_nsplit = smem.runtime_nsplit; + + int total = smem.total_tiles; + if (total == 0) return; + + int tile_elems = TileM * bufs.hidden_dim; + + TileIterator iter; + iter.remote_offsets = smem.remote_offsets; + iter.offsets_stride = bufs.experts_per_pe + 1; + iter.remote_counts = + bufs.all_expert_counts != nullptr ? smem.remote_counts : nullptr; + iter.counts_stride = bufs.experts_per_pe; + iter.experts_per_pe = bufs.experts_per_pe; + iter.num_pes = bufs.num_pes; + iter.my_pe = bufs.my_pe; + iter.init(xc, MC); + + int my_total = iter.total_tiles; + if (my_total == 0) return; + + auto* x_staging = static_cast(bufs.x_staging); + auto* dy_staging = static_cast(bufs.dy_staging); + auto* dx_staging = static_cast(bufs.dx_staging); + auto* remote_x_base = static_cast(bufs.remote_x); + auto* remote_dy_base = static_cast(bufs.remote_dy); + + // Unified layout: warp 1 gets X then dY for each tile; warp 2 puts dX. + // do_get_bwd_tma chooses per tile: same-host TMA, P2P fallback, or IB + // whole-tile round-robin. This keeps X(t), dY(t), X(t+1), ... in one loop. + if (warp_id == kCommBwdGetWarp0) { + StagePipe x_pipe; + x_pipe.init(bufs.x_src_ready + xc, bufs.x_src_consumed + xc, + is_leader, MC, kNumGetWarpsBwd * NC, runtime_nsplit); + + StagePipe dy_pipe; + dy_pipe.init(bufs.dy_src_ready + xc, bufs.dy_src_consumed + xc, + is_leader, MC, kNumGetWarpsBwd * NC, runtime_nsplit); + + int chunk_idx = yc; // kNumGetWarpsBwd == 1 + // Exactly one writer per tile per pipe: warp 1 (kCommBwdGetWarp0), + // yc=0, lane=0. The X-pipe get writes tile_expert_ids_x; the dY-pipe + // get writes tile_expert_ids_dy. Both arrays end up holding the same + // info.expert for any given live slot — but each is gated by its + // own pipe's release window, so MLP can use the X-pipe array for + // Phase 2 mlp4 and the dY-pipe array for Phase 2 mlp3 with + // asymmetric release timing (release_src after mlp4, release_dy + // after mlp3) and no race in either read. + bool is_expert_id_writer = + (warp_id == kCommBwdGetWarp0) && (yc == 0); + int* x_expert_dst = is_expert_id_writer ? bufs.tile_expert_ids_x : nullptr; + int* dy_expert_dst = is_expert_id_writer ? bufs.tile_expert_ids_dy : nullptr; + int* x_valid_dst = is_expert_id_writer ? bufs.tile_valid_rows_x : nullptr; + + // One IB-tile counter shared by the X and dY gets, advanced per IB + // getmem so X(t), dY(t), X(t+1), … rotate across the NC get warps. + int ib_seq = 0; + for (int i = 0; i < my_total; ++i) { + TileInfo info = iter.next(); + + do_get_bwd_tma( + x_pipe, bufs, info, remote_x_base, x_staging, tile_elems, + xc, MC, chunk_idx, ib_seq, lane, + /*mbar_index=*/0, + x_expert_dst, x_valid_dst, descs, + descs ? descs->src_x_desc : nullptr, + descs ? &descs->dst_x_staging_desc : nullptr, + get_bounce); + + do_get_bwd_tma( + dy_pipe, bufs, info, remote_dy_base, dy_staging, tile_elems, + xc, MC, chunk_idx, ib_seq, lane, + /*mbar_index=*/1, + dy_expert_dst, nullptr, descs, + descs ? descs->src_dy_desc : nullptr, + descs ? &descs->dst_dy_staging_desc : nullptr, + get_bounce); + } + } else { // warp 2 → dX put warp + StagePipe dst_pipe; + dst_pipe.init(bufs.dst_ready + xc, bufs.dst_consumed + xc, + is_leader, MC, runtime_nsplit, kNumPutWarpsBwd * NC); + + int put_local_idx = 0; + int chunk_idx = yc; + int ib_seq = 0; // advances only on IB tiles (inside do_put_bwd) + for (int i = 0; i < my_total; ++i) { + TileInfo info = iter.next(); + do_put_bwd( + dst_pipe, bufs, info, + dx_staging, tile_elems, xc, MC, chunk_idx, ib_seq, lane, + kNumPutWarpsBwd, put_local_idx); + } + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm90.cuh new file mode 100644 index 00000000..62f2881a --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm90.cuh @@ -0,0 +1,772 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Backward MoE communication: gets X + dY, puts dX +// ═══════════════════════════════════════════════════════════════════ +// +// Warp roles: +// warp 1 → X/dY get warp +// warp 2 → dX put warp +// warp 3 → MLP-side SM100 UMMA producer (exits early on SM90) +// +// Grid is launched 2D as (gridDim.x, NSplit) where NSplit is the +// MLP-side cooperative count. Comms re-indexes the flat CTA id into +// a (MC, NC) logical grid: +// +// cta_id = blockIdx.x * gridDim.y + blockIdx.y +// xc = cta_id / NC in [0, MC) +// yc = cta_id % NC in [0, NC) +// MC = (gridDim.x * gridDim.y) / NC (multiply before divide so +// MC is exact when NC ∤ NSplit) +// +// NC CTAs (same xc, different yc) cooperate on one tile per get/put. +// MC > gridDim.x when NC < NSplit so excess comm CTAs run ahead and +// load the next tiles while the MLP is still on the current one. +// +// ── Tile-level locking (same as fwd) ─────────────────────────────── +// +// All three pipes (X, dY, dX) use the per-tile StagePipe over a flat +// [L = MC * NumStages]-sized signal array. Get pipes have +// NumProducers = kNumGetWarpsPerCta·NC, NumConsumers = NSplit. Put +// pipe has NumProducers = NSplit, NumConsumers = NC. Comm strides MC +// across stages (NumStages entries — full ring cycle, independent of +// NSplit and NC); MLP is ticket-based (RemoteMlpTileIteratorBwd) and +// keys on the absolute slot index = (m_base + idx·gridDim.x) mod L. +// Both sides cover the same flat L slots. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp_comms_sm90.cuh" + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// Bwd comm warp constants +// ═══════════════════════════════════════════════════════════════════ +// +// Main loop runs on warps 1-2 only. kCommBwdGetWarp1 remains for legacy +// prologue compatibility with shared helpers; it is not a main-loop role. + +static constexpr int kCommBwdGetWarp0 = 1; +static constexpr int kCommBwdGetWarp1 = 2; // prologue compatibility only +static constexpr int kCommBwdPutWarp = 2; +static constexpr int kCommBwdWarpStart = kCommBwdGetWarp0; +static constexpr int kCommBwdWarpEnd = kCommBwdPutWarp; + +// ── TMA-GET (NVLink) warp rebalance ───────────────────────────────── +// Mirrors the fwd #111 layout: one warp saturates the TMA GET path, so warp 1 +// gets X+dY and warp 2 puts dX. Warp 3 enters the MLP-side branch and exits +// there for Compute != 100. The MLP-side iterator picks producer/consumer +// counts from the same runtime tma_enabled flag. +static constexpr int kNumGetWarpsBwd = 1; // TMA mode: 1 get warp (X+dY) +static constexpr int kNumPutWarpsBwd = 1; // TMA mode: 1 put warp (dX) + +// ═══════════════════════════════════════════════════════════════════ +// GetTmaDescsBwd — raw CUtensorMap descriptors for the bwd TMA GET +// ═══════════════════════════════════════════════════════════════════ +// +// Per-same-host-peer source maps over nvshmem_ptr(x_sorted, peer) and +// nvshmem_ptr(dy_sorted, peer); one local dst map each over x_staging / +// dy_staging. Built host-side in moe_bwd.cu, passed __grid_constant__. +// Same as fwd GetTmaDescs: local peers index by pe % gpus_per_node; remote +// peers fall back to IB getmem in do_get_bwd_tma. +struct GetTmaDescsBwd { + CUtensorMap src_x_desc[kGetMaxPes]; // index by (TileInfo::pe % gpus_per_node) + CUtensorMap src_dy_desc[kGetMaxPes]; + CUtensorMap dst_x_staging_desc; // local x_staging + CUtensorMap dst_dy_staging_desc; // local dy_staging + int enabled; // 1 = local-peer TMA path available + int gpus_per_node; + int my_pe; +}; + +__host__ __device__ __forceinline__ bool get_pe_is_local_bwd( + const GetTmaDescsBwd* descs, int pe) { + if (descs == nullptr || descs->gpus_per_node <= 0) return true; + return (pe / descs->gpus_per_node) == (descs->my_pe / descs->gpus_per_node); +} + +// Box height (rows) of a single backward TMA transfer. SM100 mirrors +// forward's full [32×256] row-band box; SM90 uses a smaller 8-row box. +static constexpr int kGetBoxRowsBwdSm90 = 8; +static constexpr int kGetBoxRowsBwdSm100 = 32; + +// X and dY reuse one data box sequentially but keep distinct mbarriers. +// Reinitializing one mbarrier across the two TMA streams races the async proxy; +// the box itself is safe to reuse after cp.async.bulk.wait_group<0>. +template +struct GetBounceBwd { + static constexpr int kBoxRows = BoxRows; + static constexpr int kBoxElems = kBoxRows * kGetKChunk; + static constexpr int kBoxBytes = kBoxElems * (int)sizeof(Element); + static constexpr int kMbarOffset = kBoxBytes; + static constexpr int kMbarStride = sizeof(uint64_t); + static constexpr int kTotalBytes = + (kBoxBytes + 2 * kMbarStride + 127) & ~127; +}; + +// do_get_bwd_tma is defined below, after CommBuffersBwd (it needs the +// complete struct type for bufs.hidden_dim). + +// ═══════════════════════════════════════════════════════════════════ +// CommBuffersBwd — backward-specific buffer pointers +// ═══════════════════════════════════════════════════════════════════ + +struct CommBuffersBwd { + // Staging: flat slot list [L][TileM][hidden_dim] where + // L = gridDim.x * NumStages. Indexed by slot ∈ [0, L). + void* x_staging; // gets: X from remote PE + void* dy_staging; // gets: dY from remote PE + void* dx_staging; // puts: dX back to remote PE + + // Per-tile signals ([L] entries each), one counter pair per slot — + // same granularity as fwd src/dst pipes. + int* x_src_ready; + int* x_src_consumed; + int* dy_src_ready; + int* dy_src_consumed; + int* dst_ready; + int* dst_consumed; + + // Per-slot expert IDs ([L] each), one array per get pipe. + // + // They hold the SAME value for any given live slot — info.expert at the + // time the comm warp wrote that slot — but each is gated by its own + // pipe's StagePipe release window: + // + // tile_expert_ids_x is overwritten between MLP's release_src and + // comm's next-lap producer_release on the X pipe. + // Phase 2 mlp4 reads this (mlp4 reads X) and is + // safe as long as release_src is deferred past mlp4. + // + // tile_expert_ids_dy is overwritten between MLP's release_dy and + // comm's next-lap producer_release on the dY pipe. + // Phase 2 mlp3 reads this (mlp3 reads dY) and is + // safe as long as release_dy is deferred past mlp3. + // + // Splitting the array is what lets release_src happen between mlp4 and + // mlp3 without races: comm refilling tile_expert_ids_x in that window + // can't poison mlp3, because mlp3 reads the dY-pipe array. + int* tile_expert_ids_x; + int* tile_expert_ids_dy; + int* tile_valid_rows_x; + + // Remote data pointers (symmetric memory) + void* remote_x; // get source: X on remote PE + void* remote_dy; // get source: dY on remote PE + void* remote_dx; // put destination: dX on remote PE + + // Reused from forward + int* all_expert_offsets; + int* all_expert_counts; + int all_expert_offsets_stride; + + int my_pe; + int num_pes; + nvshmem_team_t team; + + int hidden_dim; + int num_experts; + int experts_per_pe; + + // GEMM/comm-active CTA count = floor_NS = grid_x · NSplit (NOT the + // launched grid size). Under the flat 1-D launch gridDim.x = num_blocks + // (= num_sms) ≥ n_gemm, so MC = n_gemm / NC and the MLP column count + // grid_x = n_gemm / NSplit must come from here, not from gridDim. + int n_gemm; + + // Used by the dummy_consumer_bwd in mlp_comms_bwd.cu when it replicates + // the real BWD's batch loop (mlp_global_barrier + mlp_x_barrier). Not + // used by the comm warps themselves. + int* barrier_counter; + int* phase_counter; +}; + +// ═══════════════════════════════════════════════════════════════════ +// do_get_bwd_tma — TMA GET of one buffer (X or dY) for one tile +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors fwd do_get's TMA path (cuda::ptx intrinsics): peer → smem box → +// local staging, via TMA load + TMA store. Because the bounce box is small +// (kGetBoxRowsBwd rows), the single get warp loops over its TileM/NC-row +// band in kGetBoxRowsBwd-tall sub-bands × kGetKChunk-wide column boxes. +// chunk_idx == yc (the CTA's slot within its NC-cooperating group); with +// kNumGetWarpsBwd == 1 the per-warp row band is TileM/NC. +template +__device__ __forceinline__ void do_get_bwd_tma( + StagePipe& src_pipe, + const CommBuffersBwd& bufs, + const TileInfo& info, + Element* remote_base, + Element* staging_base, + int tile_elems, + int xc, + int MC, + int chunk_idx, // = yc (0..NC-1) + int& ib_seq, // advanced once per actual IB get call + int lane, + int mbar_index, + int* tile_expert_ids_dst, + int* tile_valid_rows_dst, + const GetTmaDescsBwd* descs, + const CUtensorMap* src_desc_per_pe, // descs->src_x_desc or src_dy_desc + const CUtensorMap* dst_desc, // descs->dst_x_staging or dst_dy_staging + char* bounce_warp) { + + namespace ptx = cuda::ptx; + using Bounce = GetBounceBwd; + constexpr int kWarpsPerTile = kNumGetWarpsBwd * NC; // = NC + constexpr int kRowBand = TileM / kWarpsPerTile; // rows this warp owns + constexpr int kBoxRows = Bounce::kBoxRows; + + src_pipe.producer_acquire(); + + int slot = xc + src_pipe.producer_stage() * MC; + Element* local_tile = staging_base + slot * tile_elems; + Element* remote_tile = remote_base + info.token_offset * bufs.hidden_dim; + + bool is_local = descs && descs->enabled && get_pe_is_local_bwd(descs, info.pe); + bool tma_ok = is_local && (bufs.hidden_dim % kGetKChunk == 0) && bounce_warp != nullptr; + + if (tma_ok) { + Element* sbuf = reinterpret_cast(bounce_warp); + uint64_t* mbar = reinterpret_cast( + bounce_warp + Bounce::kMbarOffset + + mbar_index * Bounce::kMbarStride); + + if (lane == 0) + ptx::mbarrier_init(mbar, 1u); + __syncwarp(); + ptx::fence_proxy_async(ptx::space_shared); + + const int n_k = bufs.hidden_dim / kGetKChunk; + const uint32_t txbytes = (uint32_t)Bounce::kBoxBytes; + int y_src_base = info.token_offset + chunk_idx * kRowBand; // peer row + int y_dst_base = slot * TileM + chunk_idx * kRowBand; // staging row + const int local_peer = info.pe % descs->gpus_per_node; + + uint32_t phase = 0; + for (int rb = 0; rb < kRowBand; rb += kBoxRows) { + for (int k = 0; k < n_k; ++k) { + int x = k * kGetKChunk; + if (lane == 0) { + (void)ptx::mbarrier_arrive_expect_tx( + ptx::sem_release, ptx::scope_cta, ptx::space_shared, + mbar, txbytes); + int32_t load_crd[2] = {x, y_src_base + rb}; + ptx::cp_async_bulk_tensor( + ptx::space_cluster, ptx::space_global, + sbuf, &src_desc_per_pe[local_peer], load_crd, mbar); + while (!ptx::mbarrier_try_wait_parity(mbar, phase)) {} + int32_t store_crd[2] = {x, y_dst_base + rb}; + ptx::cp_async_bulk_tensor( + ptx::space_global, ptx::space_shared, + dst_desc, store_crd, sbuf); + ptx::cp_async_bulk_commit_group(); + ptx::cp_async_bulk_wait_group(ptx::n32_t<0>{}); + } + phase ^= 1; + __syncwarp(); + } + } + } else { + constexpr int kWarpsPerTile = kNumGetWarpsBwd * NC; // one get warp per CTA + int chunk = tile_elems / kWarpsPerTile; + int offset = chunk_idx * chunk; + int global_pe = nvshmem_team_translate_pe(bufs.team, info.pe, NVSHMEM_TEAM_WORLD); + void* p2p_remote = nvshmem_ptr(remote_tile, global_pe); + if (p2p_remote != nullptr) { + nvshmemx_getmem_warp(local_tile + offset, remote_tile + offset, + (size_t)chunk * sizeof(Element), global_pe); + } else { + if (chunk_idx == ib_seq % kWarpsPerTile) { + int valid_elems = + max(0, min(TileM, info.valid_rows)) * bufs.hidden_dim; + for (int i = valid_elems + lane; i < tile_elems; i += 32) + local_tile[i] = Element{}; + __syncwarp(); + if (valid_elems > 0) + nvshmemx_getmem_warp( + local_tile, remote_tile, + (size_t)valid_elems * sizeof(Element), global_pe); + } + ++ib_seq; + } + } + + if (tile_expert_ids_dst != nullptr && chunk_idx == 0 && lane == 0) { + tile_expert_ids_dst[slot] = info.expert; + if (tile_valid_rows_dst != nullptr) + tile_valid_rows_dst[slot] = info.valid_rows; + } + + __syncwarp(); + src_pipe.producer_release(lane); +} + +// ═══════════════════════════════════════════════════════════════════ +// do_get_bwd — gets one buffer (X or dY) for one tile +// ═══════════════════════════════════════════════════════════════════ +// +// kWarpsPerTile = kNumGetWarpsPerCta * NC = 2·NC (warps 1+2 across +// NC cooperating CTAs). Slot in flat list = xc + producer_stage*MC. + +// `tile_expert_ids_dst` is the per-pipe expert-id array to write to (one of +// bufs.tile_expert_ids_x or bufs.tile_expert_ids_dy, depending on which pipe +// this call is for), or nullptr to skip the write. Exactly one comm warp per +// tile is expected to pass non-null — caller picks the writer to avoid double +// writes and false sharing. +// +// The write happens AFTER the getmem completes and BEFORE producer_release, +// so it is part of the slot's atomic publish: any consumer that's seen +// producer_release sees both the getmem'd staging bytes AND the expert id. +template +__device__ __forceinline__ void do_get_bwd( + StagePipe& src_pipe, + const CommBuffersBwd& bufs, + const TileInfo& info, + Element* remote_base, + Element* staging_base, // base of flat [L][TileM][hidden_dim] + int tile_elems, + int xc, // start slot in flat list + int MC, // stride between consecutive comm slots + int chunk_idx, // 0..2·NC-1 + int& ib_seq, // IB-tile counter (round-robin turn; advanced here) + int lane, + int* tile_expert_ids_dst) { + + constexpr int kWarpsPerTile = kNumGetWarpsPerCta * NC; + int chunk = tile_elems / kWarpsPerTile; + int offset = chunk_idx * chunk; + + src_pipe.producer_acquire(); + + // Slot in the flat list. Comm strides MC across pipe stages; + // MLP strides gridDim.x. Both resolve to the same physical + // slot for a given global tile since L = MC·K = gridDim.x·NumStages. + int slot = xc + src_pipe.producer_stage() * MC; + + int global_pe = nvshmem_team_translate_pe(bufs.team, info.pe, NVSHMEM_TEAM_WORLD); + + // Tile base (offset 0). P2P chunks add `offset`; the IB whole-tile + // transfer addresses the base directly. + Element* local_tile = staging_base + slot * tile_elems; + Element* remote_tile = remote_base + info.token_offset * bufs.hidden_dim; + + // Fast path: NVLink P2P → all 2·NC warps cooperate, each copying its + // chunk with direct int4 loads. IB path: one whole-tile getmem issued + // by a single warp, chosen round-robin across IB tiles only (ib_seq + // skips local/P2P tiles and is advanced in lockstep by every warp). + // See the matching fwd do_get for the full rationale. + void* p2p_remote = nvshmem_ptr(remote_tile, global_pe); + if (p2p_remote != nullptr) { + static_assert(sizeof(Element) == 2, + "P2P int4 copy assumes 2-byte elements (bf16/half)"); + constexpr int kElemsPerVec = sizeof(int4) / sizeof(Element); + int n_vec = chunk / kElemsPerVec; + const int4* src_v = reinterpret_cast( + static_cast(p2p_remote) + offset); + int4* dst_v = reinterpret_cast(local_tile + offset); + #pragma unroll 4 + for (int i = lane; i < n_vec; i += 32) { + dst_v[i] = src_v[i]; + } + __syncwarp(); + } else { + if (chunk_idx == ib_seq % kWarpsPerTile) { + nvshmemx_getmem_warp(local_tile, remote_tile, + (size_t)tile_elems * sizeof(Element), global_pe); + } + ++ib_seq; + } + + if (tile_expert_ids_dst != nullptr && chunk_idx == 0 && lane == 0) { + tile_expert_ids_dst[slot] = info.expert; + } + __threadfence(); + + src_pipe.producer_release(lane); +} + +// ═══════════════════════════════════════════════════════════════════ +// do_put_bwd — puts dX for one tile back to remote PE +// ═══════════════════════════════════════════════════════════════════ +// +// kWarpsPerTile = NC. NumProducers = NSplit (MLP-side), +// NumConsumers = NC (comm-side). + +// num_put_warps / put_local_idx keep the single-put-warp layout +// (kWarpsPerTile = NC, slice = yc): warp 2 owns dX puts while warp 3 is +// reserved for the SM100 MLP path. +template +__device__ __forceinline__ void do_put_bwd( + StagePipe& dst_pipe, + const CommBuffersBwd& bufs, + const TileInfo& info, + Element* dx_staging, // base of flat [L][TileM][hidden_dim] + int tile_elems, + int xc, + int MC, + int chunk_idx, // 0..NC-1 (= yc) + int& ib_seq, // IB-tile counter (round-robin turn; advanced here) + int lane, + int num_put_warps = 1, // 1 (fallback) or kNumPutWarpsBwd (TMA mode) + int put_local_idx = 0) { // 0..num_put_warps-1 + + int kWarpsPerTile = num_put_warps * NC; + int put_idx = chunk_idx * num_put_warps + put_local_idx; // 0..kWarpsPerTile-1 + int chunk = tile_elems / kWarpsPerTile; + int offset = put_idx * chunk; + int my_bytes = chunk * sizeof(Element); + + dst_pipe.consumer_acquire(); + + int slot = xc + dst_pipe.consumer_stage() * MC; + auto* remote_dx = static_cast(bufs.remote_dx); + Element* local_tile = dx_staging + slot * tile_elems; + Element* remote_tile = remote_dx + info.token_offset * bufs.hidden_dim; + + int global_pe = nvshmem_team_translate_pe(bufs.team, info.pe, NVSHMEM_TEAM_WORLD); + + // NVLink P2P peers keep the cooperative split; IB peers send the whole + // tile as one RDMA issued by a single put warp, round-robin over IB + // tiles only (see do_get_bwd / fwd do_put). + void* p2p_remote = nvshmem_ptr(remote_tile, global_pe); + if (p2p_remote != nullptr) { + nvshmemx_putmem_warp(remote_tile + offset, local_tile + offset, + my_bytes, global_pe); + } else { + if (put_idx == ib_seq % kWarpsPerTile) { + int valid_elems = + max(0, min(TileM, info.valid_rows)) * bufs.hidden_dim; + if (valid_elems > 0) + nvshmemx_putmem_warp( + remote_tile, local_tile, + (size_t)valid_elems * sizeof(Element), global_pe); + } + ++ib_seq; + } + + dst_pipe.consumer_release(lane); +} + +// ═══════════════════════════════════════════════════════════════════ +// nvshmem_comm_prologue_bwd — fetch remote offsets, compute tile counts +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors fwd nvshmem_comm_prologue. Computes both counts since uneven +// distribution can give comm and MLP different per-CTA tile counts: +// smem.total_tiles = comm count (start = xc, stride = MC) +// smem.per_cta_tiles = MLP count (start = blockIdx.x, stride = gridDim.x) +// +// Runs on warps 2-3 (kCommWarpStart..kCommWarpEnd). Grid is 2D +// (gridDim.x, NSplit). + +__device__ __forceinline__ int select_runtime_nsplit_bwd_from_tiles( + int total_tiles, + int num_n_tiles_1, int num_n_tiles_2t, int num_n_tiles_5, + int num_k_tiles_1, int num_k_tiles_2t, int num_k_tiles_5, + int experts_per_pe, + int num_blocks, int static_nsplit) { + if (total_tiles <= 0) return static_nsplit; + constexpr int candidates[5] = {2, 4, 6, 8, 16}; + int best_ns = static_nsplit; + long long best_cost = -1; + long long best_waste = -1; + int best_n_gemm = -1; + long long w1 = 2LL * num_k_tiles_1; // mlp1 recomputes gate + up + long long w2t = (long long)num_k_tiles_2t; + long long w5 = 2LL * num_k_tiles_5; // mlp5 computes dU@B + dV@C + #pragma unroll + for (int i = 0; i < 5; ++i) { + int ns = candidates[i]; + if (ns > num_blocks) continue; + int n_gemm = (num_blocks / ns) * ns; + int ms = n_gemm / ns; + if (ms <= 0) continue; + int m_waves = (total_tiles + ms - 1) / ms; + int n1_waves = (num_n_tiles_1 + ns - 1) / ns; + int n2t_waves = (num_n_tiles_2t + ns - 1) / ns; + int n5_waves = (num_n_tiles_5 + ns - 1) / ns; + long long phase_cost = w1 * n1_waves + w2t * n2t_waves + w5 * n5_waves; + long long cost = (long long)m_waves * phase_cost; + int min_n_tiles = min(num_n_tiles_1, min(num_n_tiles_2t, num_n_tiles_5)); + bool fragmented_experts = (experts_per_pe >= 8) && (min_n_tiles <= 8); + if (fragmented_experts && ns < 4) { + long long frag = ((long long)experts_per_pe * 1024LL) + / (long long)max(total_tiles, 1); + cost += (cost * frag) / 128LL; + } + long long waste1 = + (long long)m_waves * ms * n1_waves * ns + - (long long)total_tiles * num_n_tiles_1; + long long waste2t = + (long long)m_waves * ms * n2t_waves * ns + - (long long)total_tiles * num_n_tiles_2t; + long long waste5 = + (long long)m_waves * ms * n5_waves * ns + - (long long)total_tiles * num_n_tiles_5; + long long waste = w1 * waste1 + w2t * waste2t + w5 * waste5; + if (best_cost < 0 || + cost < best_cost || + (cost == best_cost && waste < best_waste) || + (cost == best_cost && waste == best_waste && n_gemm > best_n_gemm)) { + best_cost = cost; + best_waste = waste; + best_n_gemm = n_gemm; + best_ns = ns; + } + } + return best_ns; +} + +template +__device__ __forceinline__ void nvshmem_comm_prologue_bwd( + CommSmem& smem, + const CommBuffersBwd& bufs, + int static_nsplit) { + + int warp_id = threadIdx.x / 32; + if (warp_id < kCommWarpStart || warp_id > kCommWarpEnd) + return; + + int lane = threadIdx.x % 32; + int local_warp = warp_id - kCommWarpStart; + bool is_leader = (local_warp == 0 && lane == 0); + + int experts_per_pe = bufs.experts_per_pe; + int num_pes = bufs.num_pes; + int my_pe = bufs.my_pe; + + if (bufs.all_expert_offsets == nullptr) { + if (is_leader) { + smem.total_tiles = 0; + smem.per_cta_tiles = 0; + smem.global_total = 0; + smem.runtime_nsplit = static_nsplit; + smem.runtime_n_gemm = bufs.n_gemm; + smem.runtime_grid_x = bufs.n_gemm / static_nsplit; + } + __threadfence_block(); + __syncwarp(); + return; + } + // NOTE: num_pes == 1 is NO LONGER short-circuited (mirrors fwd + // nvshmem_comm_prologue). The shared TileIterator enumerates the LOCAL PE + // (p=0), so a single-GPU bwd run stages its local X/dY through the comm + // get path and puts dX back through the comm put path exactly like the + // multi-GPU remote path — the separate local MLP pass is gone. + + int local_e_start = my_pe * experts_per_pe; + int src_stride = bufs.num_experts + 1; + int dst_stride = experts_per_pe + 1; + { + int tid = threadIdx.x - kCommThreadStart; + for (int pe = 0; pe < num_pes; ++pe) { + for (int i = tid; i <= experts_per_pe; i += (kCommWarpsPerCta * 32)) + smem.remote_offsets[pe * dst_stride + i] = + bufs.all_expert_offsets[pe * src_stride + local_e_start + i]; + if (bufs.all_expert_counts != nullptr) { + for (int i = tid; i < experts_per_pe; i += (kCommWarpsPerCta * 32)) + smem.remote_counts[pe * experts_per_pe + i] = + bufs.all_expert_counts[ + pe * bufs.num_experts + local_e_start + i]; + } + } + } + __syncwarp(); + + // Flat 1-D launch: cta_id = flat_id = blockIdx.x (gridDim.y == 1). + // Comm uses all NC-complete launched CTAs, decoupled from the MLP runtime + // NS. MLP count below is only the static-seed default; moe_bwd_kernel + // overwrites it after selecting runtime NS from the actual tile count. + int cta_id = blockIdx.x * gridDim.y + blockIdx.y; + int n_comm = (gridDim.x / NC) * NC; + int xc = cta_id / NC; + int MC = n_comm / NC; + int grid_x = bufs.n_gemm / static_nsplit; + int mlp_col = cta_id / static_nsplit; + + TileIterator iter; + iter.remote_offsets = smem.remote_offsets; + iter.offsets_stride = dst_stride; + iter.remote_counts = + bufs.all_expert_counts != nullptr ? smem.remote_counts : nullptr; + iter.counts_stride = experts_per_pe; + iter.experts_per_pe = experts_per_pe; + iter.num_pes = num_pes; + iter.my_pe = my_pe; + + // Comm count: this CTA's per-(xc, yc) workload. + iter.init(xc, MC); + int comm_count = iter.total_tiles; + + // MLP count: this CTA's logical column (mlp_col) workload, stride grid_x. + iter.init(mlp_col, grid_x); + int mlp_count = iter.total_tiles; + + // Grid-wide remote tile count (start=0, stride=1 → counts every tile + // once). Same value on every CTA — used by moe_fused_bwd to gate the + // remote phase consistently across the grid. + iter.init(0, 1); + int grid_total = iter.total_tiles; + + if (is_leader) { + smem.total_tiles = comm_count; // read by comm_main_bwd + smem.per_cta_tiles = mlp_count; // read by MLP iterator + smem.global_total = grid_total; // read by moe_fused_bwd gate + smem.runtime_nsplit = static_nsplit; + smem.runtime_n_gemm = bufs.n_gemm; + smem.runtime_grid_x = grid_x; + } + // Block-scope fence so the leader's smem writes above are visible + // to all threads in this CTA (including MLP warps that early- + // returned at the top) AFTER the caller's __syncthreads(). Without + // this, the compiler/HW may reorder MLP warps' loads of + // smem.global_total to before the comm warp's store, giving stale + // (zero-initialized) values and causing CTA(0,0) to skip the remote + // phase while peer CTAs proceed — manifests as a multi-PE deadlock. + __threadfence_block(); + __syncwarp(); +} + +// ═══════════════════════════════════════════════════════════════════ +// nvshmem_comm_main_bwd — independent flat loops per warp role +// ═══════════════════════════════════════════════════════════════════ +// +// All three pipes are per-tile StagePipes (mirrors fwd). Comm at xc +// walks slots {xc, xc+MC, ..., xc+(K-1)·MC} of the [L]-sized signal +// arrays (stride MC). Get pipes use NumProducers = 2·NC, +// NumConsumers = NSplit; put pipe uses NumProducers = NSplit, +// NumConsumers = NC. + +template < + typename Element, + int TileM, + int NumStages, + int NC = 2, + int BoxRows = kGetBoxRowsBwdSm90> +__device__ __forceinline__ void nvshmem_comm_main_bwd( + CommSmem& smem, + const CommBuffersBwd& bufs, + const GetTmaDescsBwd* descs = nullptr, // nullptr/disabled → getmem layout + char* get_bounce = nullptr) { // bwd TMA bounce smem region + + // Comm-side pipe depth is NumStages — one comm CTA cycle covers + // the full MC × NumStages staging area, independent of NSplit and + // NC. The MLP side is ticket-based (RemoteMlpTileIteratorBwd) and + // shares no per-stage_id state, so no divisibility relationship is + // required between NumStages, NSplit and NC. + constexpr int K = NumStages; + + int warp_id = threadIdx.x / 32; + if (warp_id != kCommBwdGetWarp0 && warp_id != kCommBwdPutWarp) + return; + + int lane = threadIdx.x % 32; + bool is_leader = (lane == 0); + + // Flat 1-D launch: cta_id = flat_id = blockIdx.x (gridDim.y == 1). + int cta_id = blockIdx.x * gridDim.y + blockIdx.y; + int n_comm = (gridDim.x / NC) * NC; + if (cta_id >= n_comm) return; + int xc = cta_id / NC; + int yc = cta_id % NC; + int MC = n_comm / NC; + int runtime_nsplit = smem.runtime_nsplit; + + int total = smem.total_tiles; + if (total == 0) return; + + int tile_elems = TileM * bufs.hidden_dim; + + TileIterator iter; + iter.remote_offsets = smem.remote_offsets; + iter.offsets_stride = bufs.experts_per_pe + 1; + iter.remote_counts = + bufs.all_expert_counts != nullptr ? smem.remote_counts : nullptr; + iter.counts_stride = bufs.experts_per_pe; + iter.experts_per_pe = bufs.experts_per_pe; + iter.num_pes = bufs.num_pes; + iter.my_pe = bufs.my_pe; + iter.init(xc, MC); + + int my_total = iter.total_tiles; + if (my_total == 0) return; + + auto* x_staging = static_cast(bufs.x_staging); + auto* dy_staging = static_cast(bufs.dy_staging); + auto* dx_staging = static_cast(bufs.dx_staging); + auto* remote_x_base = static_cast(bufs.remote_x); + auto* remote_dy_base = static_cast(bufs.remote_dy); + + // Unified layout: warp 1 gets X then dY for each tile; warp 2 puts dX. + // do_get_bwd_tma chooses per tile: same-host TMA, P2P fallback, or IB + // whole-tile round-robin. This keeps X(t), dY(t), X(t+1), ... in one loop. + if (warp_id == kCommBwdGetWarp0) { + StagePipe x_pipe; + x_pipe.init(bufs.x_src_ready + xc, bufs.x_src_consumed + xc, + is_leader, MC, kNumGetWarpsBwd * NC, runtime_nsplit); + + StagePipe dy_pipe; + dy_pipe.init(bufs.dy_src_ready + xc, bufs.dy_src_consumed + xc, + is_leader, MC, kNumGetWarpsBwd * NC, runtime_nsplit); + + int chunk_idx = yc; // kNumGetWarpsBwd == 1 + // Exactly one writer per tile per pipe: warp 1 (kCommBwdGetWarp0), + // yc=0, lane=0. The X-pipe get writes tile_expert_ids_x; the dY-pipe + // get writes tile_expert_ids_dy. Both arrays end up holding the same + // info.expert for any given live slot — but each is gated by its + // own pipe's release window, so MLP can use the X-pipe array for + // Phase 2 mlp4 and the dY-pipe array for Phase 2 mlp3 with + // asymmetric release timing (release_src after mlp4, release_dy + // after mlp3) and no race in either read. + bool is_expert_id_writer = + (warp_id == kCommBwdGetWarp0) && (yc == 0); + int* x_expert_dst = is_expert_id_writer ? bufs.tile_expert_ids_x : nullptr; + int* dy_expert_dst = is_expert_id_writer ? bufs.tile_expert_ids_dy : nullptr; + int* x_valid_dst = is_expert_id_writer ? bufs.tile_valid_rows_x : nullptr; + + // One IB-tile counter shared by the X and dY gets, advanced per IB + // getmem so X(t), dY(t), X(t+1), … rotate across the NC get warps. + int ib_seq = 0; + for (int i = 0; i < my_total; ++i) { + TileInfo info = iter.next(); + + do_get_bwd_tma( + x_pipe, bufs, info, remote_x_base, x_staging, tile_elems, + xc, MC, chunk_idx, ib_seq, lane, + /*mbar_index=*/0, + x_expert_dst, x_valid_dst, descs, + descs ? descs->src_x_desc : nullptr, + descs ? &descs->dst_x_staging_desc : nullptr, + get_bounce); + + do_get_bwd_tma( + dy_pipe, bufs, info, remote_dy_base, dy_staging, tile_elems, + xc, MC, chunk_idx, ib_seq, lane, + /*mbar_index=*/1, + dy_expert_dst, nullptr, descs, + descs ? descs->src_dy_desc : nullptr, + descs ? &descs->dst_dy_staging_desc : nullptr, + get_bounce); + } + } else { // warp 2 → dX put warp + StagePipe dst_pipe; + dst_pipe.init(bufs.dst_ready + xc, bufs.dst_consumed + xc, + is_leader, MC, runtime_nsplit, kNumPutWarpsBwd * NC); + + int put_local_idx = 0; + int chunk_idx = yc; + int ib_seq = 0; // advances only on IB tiles (inside do_put_bwd) + for (int i = 0; i < my_total; ++i) { + TileInfo info = iter.next(); + do_put_bwd( + dst_pipe, bufs, info, + dx_staging, tile_elems, xc, MC, chunk_idx, ib_seq, lane, + kNumPutWarpsBwd, put_local_idx); + } + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm100.cuh new file mode 100644 index 00000000..907d68e3 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm100.cuh @@ -0,0 +1,1343 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MoE MLP with NVSHMEM communication +// ═══════════════════════════════════════════════════════════════════ +// +// Combines nvshmem get/put with MLP compute in a single persistent +// kernel. Warps 1-3 of WG0 handle nvshmem communication while the +// rest of the CTA runs the MLP1 pipeline. +// +// Warp roles (384 threads = 12 warps): +// +// WG0 (warps 0–3): +// warp 0 TMA producer (mlp1 pipelines) +// warp 1 NVSHMEM get warp 0 (fetches X into src_staging) +// warp 2 NVSHMEM get warp 1 (fetches X into src_staging) +// warp 3 NVSHMEM put warp (sends Y from dst_staging) +// +// WG1 (warps 4–7): Consumer A (mlp1 WGMMA) +// WG2 (warps 8–11): Consumer B (mlp1 WGMMA) +// +// Staging areas (per blockIdx.x, shared by both blockIdx.y): +// src_staging[NumStages][TileM][hidden_dim] +// dst_staging[NumStages][TileM][intermediate_dim] +// +// Schedule: the two get warps and the put warp each run an +// independent flat loop over the same tile sequence. The NumStages +// staging ring bounds how far the get warps can lead the put warp. +// +// NVSHMEM get throughput is roughly half of put; we counteract that +// by issuing 2× as many gets per tile. Each tile is chunked across +// 2·NC warps on the get side (2 warps per CTA × NC CTAs) but only NC +// warps on the put side. Get chunk = tile_bytes / (2·NC); put chunk +// = tile_bytes / NC. +// +// Transport-dependent chunking: +// * NVLink P2P (nvshmem_ptr != nullptr): all 2·NC get warps / NC put +// warps cooperate per tile, each copying its chunk with coalesced +// int4 loads/stores — many warps maximise NVLink throughput. +// * IB (nvshmem_ptr == nullptr): the per-warp chunking is counter- +// productive (small RDMA messages, per-tile rendezvous). Instead a +// SINGLE warp issues one whole-tile getmem/putmem, and the warp that +// issues rotates round-robin across the 2·NC (get) / NC (put) warps +// keyed by the tile sequence index, so consecutive tiles are issued +// by different warps. Their blocking transfers then overlap (bounded +// by the NumStages staging ring). The non-issuing warps do no I/O and +// just rubber-stamp the StagePipe so the producer/consumer counters +// stay intact. +// +// ═══════════════════════════════════════════════════════════════════ + +#include // CUtensorMap (raw TMA descriptor) +#include +#include +#include // typed TMA / mbarrier intrinsics (replaces inline PTX) +#include +#include +#include +#include "nvshmem_helpers.cuh" // g_dest_table + +#include "cta_barrier_sm100.cuh" // SyncThreadsCtaCounterBarrier + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// StagePipe — bidirectional acquire/release for a staging buffer +// ═══════════════════════════════════════════════════════════════════ +// +// Wraps a pair of atomic counters (ready[S], consumed[S]) shared +// between a producer and a consumer. Each side calls acquire/release: +// +// Producer: acquire() waits on consumed[] (slot free), then +// release() signals ready[] (data written). +// Consumer: acquire() waits on ready[] (data available), then +// release() signals consumed[] (slot freed). +// +// Internally tracks stage index and monotonic thresholds. +// is_leader is set at init — only the leader thread polls atomics. + +template +struct StagePipe { + int* ready; // producer signals, consumer waits + int* consumed; // consumer signals, producer waits + int stride; // step in slots between consecutive logical stages + bool is_leader; + int num_producers; + int num_consumers; + + int prod_count; + int cons_count; + int prod_expected[NumStages]; // threshold for consumed[] (producer waits on) + int cons_expected[NumStages]; // threshold for ready[] (consumer waits on) + + // Stride lets the pipe address into a flat slot list with non-unit + // stage spacing (e.g. comm pipes use stride = MC, MLP pipes use + // stride = gridDim.x). Default stride = 1 preserves the original + // contiguous-ring layout used by the bwd path. + __device__ void init(int* ready_, int* consumed_, bool is_leader_, + int stride_ = 1) { + ready = ready_; + consumed = consumed_; + stride = stride_; + is_leader = is_leader_; + num_producers = NumProducers; + num_consumers = NumConsumers; + prod_count = 0; + cons_count = 0; + for (int i = 0; i < NumStages; ++i) { + prod_expected[i] = 0; + cons_expected[i] = 0; + } + } + + __device__ void init(int* ready_, int* consumed_, bool is_leader_, + int stride_, int num_producers_, int num_consumers_) { + init(ready_, consumed_, is_leader_, stride_); + num_producers = num_producers_; + num_consumers = num_consumers_; + } + + // ── Producer side ─────────────────────────────────────── + + __device__ int producer_stage() const { + return prod_count % NumStages; + } + + // Wait until this slot is free (consumer has released it). + // + // Acquire-side memory ordering: only the leader does the atomic poll + // (cheaper than 32 lanes spinning), but ALL lanes need to see the + // counterparty's writes after this returns. __syncwarp alone does NOT + // establish global-memory ordering — non-leader lanes may still hit + // stale L1/L2. We add a __threadfence after the leader's poll AND after + // __syncwarp so every lane has a release/acquire pair. + __device__ void producer_acquire() { + int s = producer_stage(); + if (prod_count >= NumStages) { + if (is_leader) { + cuda::atomic_ref r(consumed[s * stride]); + while (r.load(cuda::memory_order_relaxed) < prod_expected[s]) {} + } + __syncwarp(); + __threadfence(); + } + prod_expected[s] += num_consumers; + } + + // Signal that data has been written to this slot. + __device__ void producer_release(int lane) { + int s = producer_stage(); + // Device-scope: this is the get-side publish (NIC-written src → MLP) or + // a slot-free signal — both GPU-internal w.r.t. the consumer. The NIC's + // inbound-write visibility is NOT fixable by a fence (needs GDRCopy), so + // system scope here only costs bandwidth without buying correctness. The + // one publish that genuinely crosses to the NIC (dst → put warp) keeps + // system scope in RemoteIter::release_dst. + __threadfence(); + if (lane == 0) { + atomicAdd(&ready[s * stride], 1); + } + prod_count++; + } + + // ── Consumer side ─────────────────────────────────────── + + __device__ int consumer_stage() const { + return cons_count % NumStages; + } + + // Wait until this slot has data (producer has released it). + // + // Acquire-side memory ordering: see comment on producer_acquire above. + // __syncwarp alone does NOT establish global-memory visibility for + // non-leader lanes — they may load stale data from L1/L2 even after + // the leader has confirmed the producer's release. The __threadfence + // after __syncwarp gives every lane in the warp a fresh view of the + // producer's writes (HBM coherence) before the warp-collective load + // in the subsequent put. + __device__ void consumer_acquire() { + int s = consumer_stage(); + if (is_leader) { + cuda::atomic_ref r(ready[s * stride]); + while (r.load(cuda::memory_order_relaxed) < cons_expected[s] + num_producers) {} + } + __syncwarp(); + // Device-scope: the put warp observes dst_staging (Y), already flushed + // to HBM by the MLP's system-scope release_dst. Observe is GPU-internal. + __threadfence(); + cons_expected[s] += num_producers; + } + + // Signal that this slot has been consumed and is free. + __device__ void consumer_release(int lane) { + int s = consumer_stage(); + __threadfence(); + if (lane == 0) { + atomicAdd(&consumed[s * stride], 1); + } + cons_count++; + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// CommSmem — per-CTA shared memory (smem-local, not shared across CTAs) +// ═══════════════════════════════════════════════════════════════════ + +struct CommSmem { + // Total number of remote tiles across all CTAs. Initialized to -1. + // Comm warps write the actual count; consumers wait for >= 0. + int total_tiles; + + // Per-CTA tile count (this blockIdx.x's share of total_tiles). + int per_cta_tiles; + + // Grid-wide remote tile count (sum across all CTAs/columns). Same + // value on every CTA in the grid; used by moe_fused_bwd to gate the + // remote phase consistently — the inner mlp_fused_bwd's cross-CTA + // mlp_global_barrier requires every CTA to make the same enter/skip + // decision, so per-CTA per_cta_tiles can't be the gate. + int global_total; + int runtime_nsplit; // N-split CTAs per GEMM M-row + int runtime_n_gemm; + int runtime_grid_x; + + // Pre-fetched remote expert offsets, assigned to dynamic smem at runtime. + // Layout: [num_pes][experts_per_pe + 1] + int* remote_offsets; + // Exact routed-row counts matching remote_offsets. + // Layout: [num_pes][experts_per_pe] + int* remote_counts; +}; + +__device__ __forceinline__ int select_runtime_nsplit_from_tiles( + int total_tiles, int num_n_tiles_1, int num_n_tiles_2, + int num_k_tiles_1, int num_k_tiles_2, + int num_blocks, int static_nsplit, int m_subtiles = 1) { + if (total_tiles <= 0) return 2 * m_subtiles; + constexpr int base_candidates[5] = {2, 4, 6, 8, 16}; + int best_base_ns = static_nsplit; + long long best_cost = -1; + long long best_waste = -1; + int best_n_gemm = -1; + long long phase1_weight = 2LL * num_k_tiles_1; // gate + up + long long phase2_weight = (long long)num_k_tiles_2; + #pragma unroll + for (int i = 0; i < 5; ++i) { + int base_ns = base_candidates[i]; + if (base_ns * m_subtiles > num_blocks) continue; + int n_gemm = (num_blocks / base_ns) * base_ns; + int ms = n_gemm / base_ns; + if (ms <= 0) continue; + int m_waves = (total_tiles + ms - 1) / ms; + int n1_waves = (num_n_tiles_1 + base_ns - 1) / base_ns; + int n2_waves = (num_n_tiles_2 + base_ns - 1) / base_ns; + + long long phase_cost = phase1_weight * n1_waves + + phase2_weight * n2_waves; + long long cost = (long long)m_waves * phase_cost; + long long phase1_waste = + (long long)m_waves * ms * n1_waves * base_ns + - (long long)total_tiles * num_n_tiles_1; + long long phase2_waste = + (long long)m_waves * ms * n2_waves * base_ns + - (long long)total_tiles * num_n_tiles_2; + long long waste = phase1_weight * phase1_waste + + phase2_weight * phase2_waste; + if (best_cost < 0 || + cost < best_cost || + (cost == best_cost && waste < best_waste) || + (cost == best_cost && waste == best_waste && n_gemm > best_n_gemm)) { + best_cost = cost; + best_waste = waste; + best_n_gemm = n_gemm; + best_base_ns = base_ns; + } + } + return best_base_ns * m_subtiles; +} + +// ═══════════════════════════════════════════════════════════════════ +// PeSync — device-level phase synchronization +// ═══════════════════════════════════════════════════════════════════ + +struct PeSync { + int* cta_counter; // [1] device memory — shared monotonic counter + int num_pes; + int my_pe; + nvshmem_team_t team; + + // ── barrier: cross-CTA global memory fence ────────── + // Defers to a persistent SyncThreadsCtaCounterBarrier owned by the + // caller — see cta_barrier.cuh for why the barrier object must be + // constructed ONCE per kernel and reused (target accumulates). + // PeSync cannot own the barrier itself because it lives inside a + // __grid_constant__ MoeParams; the caller builds the barrier on + // the kernel stack and threads it into every signal/wait call. + __device__ void barrier(SyncThreadsCtaCounterBarrier& b) const { + b.wait(); + } + + __device__ void barrier_all(SyncThreadsCtaCounterBarrier& b) const { + b.wait(); + int cta_id = blockIdx.x + blockIdx.y * gridDim.x; + if (cta_id == 0) { + if (num_pes > 1 && threadIdx.x == 0) + nvshmem_barrier(team); + __syncthreads(); + } + b.wait(); + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// CommBuffers — pointers passed to the kernel +// ═══════════════════════════════════════════════════════════════════ + +struct CommBuffers { + // Staging areas: per-blockIdx.x, shared by both blockIdx.y CTAs. + // Layout: [gridDim.x][NumStages][TileM][dim] + void* src_staging; // [gridDim.x * S * TileM * hidden_dim] + void* dst_staging; // [gridDim.x * S * TileM * intermediate_dim] + + // Stage pipe signals in device memory (HBM), per-blockIdx.x. + // Layout: [gridDim.x][NumStages] — shared by both blockIdx.y CTAs. + int* src_ready; // nvshmem → consumer: tile data written + int* src_consumed; // consumer → nvshmem: tile data read + int* dst_ready; // consumer → nvshmem: tile data written + int* dst_consumed; // nvshmem → consumer: put completed, slot free + + const int* expert_offsets; // [E + 1] local (from radix sort, TileM-aligned, symmetric) + const int* expert_counts; // [E] exact routed rows per expert + const int* sorted_token_ids; // [total_slots] local sorted token IDs + + void* local_tokens; // [T, hidden_dim] local input (symmetric) + void* local_output; // [total_slots, intermediate_dim] local output (symmetric) + + // Per-staging-slot expert IDs in HBM. Written by do_get, read by MLP. + // Layout: [gridDim.x][NumStages] + int* tile_expert_ids; + + // All PEs' expert offsets, gathered into local symmetric memory. + // Layout: [num_pes][num_experts + 1] — symmetric. + int* all_expert_offsets; + // All PEs' exact expert counts. Layout: [num_pes][num_experts]. + int* all_expert_counts; + + // Cross-PE synchronization for offset broadcast/consume. + PeSync pe_sync; + + int hidden_dim; + int num_experts; // total global experts + int experts_per_pe; // num_experts / num_pes (evenly split) + int all_expert_offsets_stride; // max_num_experts + 1 (fixed across configs) + // GEMM/comm-active CTA count = floor_NS (largest multiple of NSplit <= + // num_sms). The flat-grid launch may run more CTAs than this (grid-stride + // phases use them); GEMM and comm gate to the first n_gemm CTAs so the + // comm↔MLP staging ring geometry (grid_x = n_gemm/NSplit, MC = n_gemm/NC) + // stays consistent. Replaces the old gridDim.x·gridDim.y product. + int n_gemm; + + // Convenience accessors for fields stored in pe_sync. + __host__ __device__ __forceinline__ int my_pe() const { return pe_sync.my_pe; } + __host__ __device__ __forceinline__ int num_pes() const { return pe_sync.num_pes; } + __host__ __device__ __forceinline__ nvshmem_team_t team() const { return pe_sync.team; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// TileInfo — describes one tile to get or put +// ═══════════════════════════════════════════════════════════════════ + +struct TileInfo { + int pe; // remote PE + int token_offset; // offset into remote PE's token buffer (in tokens) + int expert; // local expert index (for weight selection) + int valid_rows; // valid prefix rows in this communication tile + int m_subtile; // GemmTileM row within this communication tile + int parent_tile; // flat communication-tile index +}; + +// ═══════════════════════════════════════════════════════════════════ +// TileIterator — enumerates tiles: expert-outer, PE-inner, strided +// ═══════════════════════════════════════════════════════════════════ +// +// Tiles within each (expert, pe) pair are strided by gridDim.x. +// Each blockIdx.x starts at its own offset and advances by gridDim.x. + +template +struct TileIterator { + static constexpr int TileM = TileM_; + + const int* remote_offsets; // smem: [num_pes][experts_per_pe + 1] + int offsets_stride; // experts_per_pe + 1 + const int* remote_counts; // smem: [num_pes][experts_per_pe] (nullable) + int counts_stride; // experts_per_pe + int experts_per_pe; + int num_pes; + int my_pe; + + int stride; // distance between consecutive tiles + + // Current position in the iteration. + int le; // local expert index + int pe_idx; // 1..num_pes-1 (skips my_pe) + int tile; // tile within current (expert, pe) pair + int remote_tiles; // total tiles for current (expert, pe) pair + int remote_start; // start offset for current (expert, pe) + int remote_count; // exact valid rows for current (expert, pe) + + int total_tiles; // this iterator's tile count + + __device__ void init(int start, int str) { + stride = str; + le = 0; + pe_idx = 0; + total_tiles = 0; + + // Count total tiles across all (expert, pe) pairs and + // compute this iterator's share. + // NOTE: p starts at 0 → the LOCAL PE (my_pe, via g_dest_table slot + // offset 0) is enumerated first for each expert, alongside the remote + // PEs (p=1..num_pes-1). Unified path: local experts are staged by the + // comm get warp from local symmetric memory just like remote tiles, so + // the separate local MLP pass is gone (see moe.cuh moe_fused_fwd). + int lane = threadIdx.x & 31; + int local_total = 0; + int n_experts = experts_per_pe * num_pes; + for (int idx = lane; idx < n_experts; idx += 32) { + int e = idx / num_pes; + int p = idx - e * num_pes; + int pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + p) % num_pes]; + int rs = remote_offsets[pe * offsets_stride + e]; + int re = remote_offsets[pe * offsets_stride + e + 1]; + local_total += (re - rs) / TileM; + } + #pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) { + local_total += __shfl_down_sync(0xffffffff, local_total, delta); + } + int global_total = __shfl_sync(0xffffffff, local_total, 0); + if (start < global_total) + total_tiles = (global_total - 1 - start) / stride + 1; + + // Find the starting position: walk the flat tile space + // to find where start lands. + int skip = start; + for (int e = 0; e < experts_per_pe; ++e) { + for (int p = 0; p < num_pes; ++p) { + int pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + p) % num_pes]; + int rs = remote_offsets[pe * offsets_stride + e]; + int re = remote_offsets[pe * offsets_stride + e + 1]; + int ntiles = (re - rs) / TileM; + if (skip < ntiles) { + le = e; + pe_idx = p; + tile = skip; + remote_start = rs; + remote_tiles = ntiles; + remote_count = remote_counts != nullptr + ? remote_counts[pe * counts_stride + e] + : ntiles * TileM; + return; + } + skip -= ntiles; + } + } + // No tiles for this iterator. + le = experts_per_pe; + } + + __device__ bool has_next() const { + return le < experts_per_pe; + } + + __device__ TileInfo next() { + int pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; + TileInfo info; + info.pe = pe; + info.token_offset = remote_start + tile * TileM; + info.expert = le; + info.valid_rows = max(0, min(TileM, remote_count - tile * TileM)); + info.m_subtile = 0; + info.parent_tile = -1; + + // Advance by stride in the flat tile space. + // Carry the remainder across (expert, pe) boundaries. + tile += stride; + while (tile >= remote_tiles && le < experts_per_pe) { + int remaining = tile - remote_tiles; + pe_idx++; + if (pe_idx >= num_pes) { + pe_idx = 0; + le++; + } + if (le >= experts_per_pe) break; + int next_pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; + remote_start = remote_offsets[next_pe * offsets_stride + le]; + int remote_end = remote_offsets[next_pe * offsets_stride + le + 1]; + remote_tiles = (remote_end - remote_start) / TileM; + remote_count = remote_counts != nullptr + ? remote_counts[next_pe * counts_stride + le] + : remote_tiles * TileM; + tile = remaining; + } + return info; + } +}; + +template +struct GemmTileIterator { + static constexpr int TileM = TileM_; + static constexpr int MSubTiles = MSubTiles_; + static constexpr int GemmTileM = TileM / MSubTiles; + + const int* remote_offsets; + int offsets_stride; + const int* remote_counts; + int counts_stride; + int experts_per_pe; + int num_pes; + int my_pe; + int stride; + + int le; + int pe_idx; + int tile; + int remote_tiles; + int remote_parent_tiles; + int remote_start; + int remote_count; + int parent_prefix; + int total_tiles; + + __device__ void init(int start, int str) { + stride = str; + le = 0; + pe_idx = 0; + total_tiles = 0; + + int lane = threadIdx.x & 31; + int local_total = 0; + int n_experts = experts_per_pe * num_pes; + for (int idx = lane; idx < n_experts; idx += 32) { + int e = idx / num_pes; + int p = idx - e * num_pes; + int pe = liger::g_dest_table[ + (liger::comm_slot_of(my_pe) + p) % num_pes]; + int count = remote_counts[pe * counts_stride + e]; + local_total += (count + GemmTileM - 1) / GemmTileM; + } + #pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) + local_total += __shfl_down_sync(0xffffffff, local_total, delta); + int global_total = __shfl_sync(0xffffffff, local_total, 0); + if (start < global_total) + total_tiles = (global_total - 1 - start) / stride + 1; + + int skip = start; + parent_prefix = 0; + for (int e = 0; e < experts_per_pe; ++e) { + for (int p = 0; p < num_pes; ++p) { + int pe = liger::g_dest_table[ + (liger::comm_slot_of(my_pe) + p) % num_pes]; + int rs = remote_offsets[pe * offsets_stride + e]; + int re = remote_offsets[pe * offsets_stride + e + 1]; + int count = remote_counts[pe * counts_stride + e]; + int ntiles = (count + GemmTileM - 1) / GemmTileM; + int nparents = (re - rs) / TileM; + if (skip < ntiles) { + le = e; + pe_idx = p; + tile = skip; + remote_tiles = ntiles; + remote_parent_tiles = nparents; + remote_start = rs; + remote_count = count; + return; + } + skip -= ntiles; + parent_prefix += nparents; + } + } + le = experts_per_pe; + } + + __device__ bool has_next() const { + return le < experts_per_pe; + } + + __device__ int current_parent_tile() const { + return parent_prefix + tile / MSubTiles; + } + + __device__ int current_m_subtile() const { + return tile % MSubTiles; + } + + __device__ TileInfo next() { + int pe = liger::g_dest_table[ + (liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; + int parent_in_expert = tile / MSubTiles; + TileInfo info; + info.pe = pe; + info.token_offset = remote_start + parent_in_expert * TileM; + info.expert = le; + info.valid_rows = max( + 0, min(TileM, remote_count - parent_in_expert * TileM)); + info.m_subtile = tile % MSubTiles; + info.parent_tile = parent_prefix + parent_in_expert; + + tile += stride; + while (tile >= remote_tiles && le < experts_per_pe) { + int remaining = tile - remote_tiles; + parent_prefix += remote_parent_tiles; + pe_idx++; + if (pe_idx >= num_pes) { + pe_idx = 0; + le++; + } + if (le >= experts_per_pe) + break; + int next_pe = liger::g_dest_table[ + (liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; + remote_start = + remote_offsets[next_pe * offsets_stride + le]; + int remote_end = + remote_offsets[next_pe * offsets_stride + le + 1]; + remote_count = + remote_counts[next_pe * counts_stride + le]; + remote_tiles = + (remote_count + GemmTileM - 1) / GemmTileM; + remote_parent_tiles = + (remote_end - remote_start) / TileM; + tile = remaining; + } + return info; + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// NVSHMEM Comm device functions — dedicated get/put warps +// ═══════════════════════════════════════════════════════════════════ +// +// Per CTA: warp 2 is the getter, warp 3 is the putter. +// +// The physical grid is launched as (gridDim.x, N_SPLIT) to match the +// GEMM's cooperative-WG layout. Comms re-indexes the flat CTA id into +// a (MC, NC) logical grid where NC must divide N_SPLIT: +// +// cta_id = blockIdx.x * gridDim.y + blockIdx.y +// xc = cta_id / NC in [0, MC) +// yc = cta_id % NC in [0, NC) — chunk index +// MC = gridDim.x * (gridDim.y / NC) +// +// NC CTAs (same xc, different yc) cooperate on one tile and chunk its +// bytes NC ways. MC groups stride through the remote-tile list, so +// more tiles are in flight per kernel launch when NC < N_SPLIT. +// +// A CTA's comm tile is chosen by xc; its GEMM tile is chosen by +// blockIdx.x. These can diverge when NC < N_SPLIT (the whole point of +// the re-indexing) — the staging ring is keyed by xc so MLP and comm +// rendezvous through the ring, not through shared CTA coordinates. +// +// The two roles run independent flat loops over all tiles. No 1F1B — +// the staging ring (NumStages) bounds how far the getter can lead +// the putter via the shared src/dst StagePipe signals. + +// Prologue + bwd use warps 2-3 only (kept for backward compatibility +// with the bwd path that has not been migrated to the 2-getter layout). +static constexpr int kCommWarpStart = 2; +static constexpr int kCommWarpEnd = 2; +static constexpr int kCommWarpsPerCta = kCommWarpEnd - kCommWarpStart + 1; // 2 +static constexpr int kCommThreadStart = kCommWarpStart * 32; + +// Fwd comm layout: 2 get warps + 1 put warp per CTA. +// warp 1 → get warp 0 +// warp 2 → get warp 1 (also prologue leader) +// warp 3 → put warp +static constexpr int kCommGetWarp = 1; +static constexpr int kCommPutWarp = 2; +static constexpr int kNumGetWarpsPerCta = 2; // shared w/ bwd — do NOT change + +// FWD-only active get-warp count. With the TMA GET (DMA engine, issue-bound +// not thread-bound), ONE warp saturates the per-SM TMA unit, so warp 2 is +// freed (idle in steady state; still helps the epilogue last-tile put drain). +// kNumGetWarpsPerCta stays 2 for the bwd path and for the 2-warp epilogue +// drain (kPutW); only the steady-state FWD producer count drops to this. +static constexpr int kNumGetWarpsFwd = 1; + +// FWD-only put-warp count. Keep one active put warp on Hopper; the freed warp +// enters the MLP-side branch and exits there for Compute != 100. +static constexpr int kNumPutWarpsFwd = 1; + +// Range of warps participating in fwd comm (used by callers to filter +// MLP vs comm warps). +static constexpr int kCommFwdWarpStart = kCommGetWarp; +static constexpr int kCommFwdWarpEnd = kCommPutWarp; + +// Get side has kNumGetWarpsPerCta × NC warps per tile (chunk = +// tile_bytes / (kNumGetWarpsPerCta × NC)); put side has NC warps per +// tile (chunk = tile_bytes / NC). + +// ═══════════════════════════════════════════════════════════════════ +// GET-via-TMA: fetch X over NVLink with the TMA engine, not warp LD/ST. +// ═══════════════════════════════════════════════════════════════════ +// +// The dedicated get warp drives a TMA load (peer X → smem) followed by a +// TMA store (smem → local src_staging HBM). TMA can't go global→global on +// Hopper, so the bounce through smem is mandatory; the store TARGET is +// LOCAL HBM, so completion (cp.async.bulk.wait_group) + a device-scope +// __threadfence in producer_release makes it visible to the MLP reader — +// the safe-completion path (unlike a store to a PEER, which failed before). +// +// Validated in microbench: ~370 GB/s, in- and out-of-kernel correct. +// +// kGetKChunk is the contiguous column box width (elements). Each of the +// kNumGetWarpsPerCta·NC warps owns a TileM/(kNumGetWarpsPerCta·NC) row +// band of the tile and loops over hidden_dim/kGetKChunk column boxes. +// Box = [kGetBoxRows, kGetKChunk]; descriptors built host-side in moe.cu. +static constexpr int kGetKChunk = 256; + +// Per-peer source descriptors + one local-staging dst descriptor, built on +// the host (nvshmem_ptr(x_sorted, peer)) and passed __grid_constant__. +// FWD-only: NOT referenced by bwd — kept out of CommBuffers so the shared +// struct's ABI is unchanged and moe_bwd.cu need not recompile. +// +// Multi-host: only SAME-HOST peers are P2P-reachable, so src_x_desc is indexed +// by intra-host local rank (pe % gpus_per_node), NOT by global team PE — the +// array therefore only ever needs gpus_per_node (≤ kGetMaxPes) entries +// regardless of total PE count. Remote (IB) peers fall back to getmem and never +// touch this array. Assumes ranks are host-contiguous and host-aligned to +// gpus_per_node (true for a WORLD-contiguous team); asserted host-side. +static constexpr int kGetMaxPes = 8; // upper bound on GPUs per host +struct GetTmaDescs { + CUtensorMap src_x_desc[kGetMaxPes]; // index by (TileInfo::pe % gpus_per_node) + CUtensorMap dst_staging_desc; // local src_staging + int enabled; // 1 = NVLink P2P TMA path available; 0 = getmem + int gpus_per_node; // local host PE count — host/local-rank divisor + int my_pe; // this PE in team space (for same-host test) +}; + +// Same-host predicate (pure arithmetic): true iff peer `pe` shares this PE's +// host. Used both to gate the TMA get path and to pick the release_dst fence +// scope (device for same-host NVLink, system for cross-host IB). +__host__ __device__ __forceinline__ bool get_pe_is_local( + const GetTmaDescs* descs, int pe) { + if (descs == nullptr || descs->gpus_per_node <= 0) return true; // single-host/test + return (pe / descs->gpus_per_node) == (descs->my_pe / descs->gpus_per_node); +} + +// Per-peer Y-store descriptors for the direct GEMM→peer TMA store (FWD-only). +// desc[r] is a CUTE TMA-store object built host-side over +// nvshmem_ptr(local_output, peer) for the same-host peer at intra-host local +// rank r (self included, rank = my_pe % gpus_per_node). When a tile is_local, +// MLP2 stores Y straight into the destination peer's symmetric local_output via +// desc[MlpTileInfo::peer_rank] instead of dst_staging — the put warp then only +// rubber-stamps. The direct store is gated per-tile on is_local (same-host → +// P2P guaranteed), so no runtime enable flag is needed. Templated on the CUTE +// store type so the same object the wrapper passes as tma_store_y populates it. +template +struct PeerYStoreDescs { + int enabled = 1; + TmaStoreY desc[kGetMaxPes]; // index by (TileInfo::pe % gpus_per_node) +}; + +// Smem bounce: per get-warp single box + its mbarrier, 128-aligned so the +// TMA dst smem address is aligned. Stride between the two get warps' regions. +template +struct GetBounce { + static constexpr int kWarpsPerTile = kNumGetWarpsFwd * NC; + static constexpr int kBoxRows = TileM / kWarpsPerTile; + static constexpr int kBoxElems = kBoxRows * kGetKChunk; + static constexpr int kBoxBytes = kBoxElems * (int)sizeof(Element); + // [box (128-aligned)] [mbar:8B], next warp's region 128-aligned. + static constexpr int kPerWarpStride = ((kBoxBytes + 8) + 127) & ~127; + static constexpr int kTotalBytes = kPerWarpStride * kNumGetWarpsFwd; +}; + +template +__device__ __forceinline__ void do_get( + StagePipe& src_pipe, + const CommBuffers& bufs, + const TileInfo& info, + Element* src_staging, // base of flat staging buffer + int src_tile_elems, + int xc, // start slot in flat list + int MC, // stride between consecutive comm stages + int chunk_idx, // 0..kNumGetWarpsPerCta * NC - 1 + int m_subtile, + int row_chunk_idx, + bool row_empty, + int& ib_seq, // IB-tile counter (round-robin turn; advanced here) + int lane, + const GetTmaDescs* descs, // per-peer src + local dst TMA descriptors (nullptr → getmem) + char* bounce_warp) { // this warp's GetBounce sub-region (128-aligned) + + using Bounce = GetBounce; + constexpr int kWarpsPerTile = Bounce::kWarpsPerTile; + constexpr int kBoxRows = Bounce::kBoxRows; + + src_pipe.producer_acquire(); + + // No per-tile signal wait — wait_dispatch() already confirmed + // all dispatch writes on all PEs are complete. + + int slot = xc + src_pipe.producer_stage() * MC; + if (row_empty) { + src_pipe.producer_release(lane); + return; + } + + // Locality of this tile's source PE (same-host NVLink vs cross-host IB). + // Pure arithmetic; gates the TMA get path below. The MLP's release_dst + // recomputes the SAME predicate independently (embedded TileIterator), so no + // per-slot flag is published here. + bool is_local = get_pe_is_local(descs, info.pe); + + // TMA path: peer X → smem → local src_staging, via the TMA engine. + // Gated on same-host NVLink P2P (is_local), TMA availability (descs->enabled), + // and hidden_dim divisibility; cross-host IB, odd shapes, or a null descs + // (e.g. the standalone test) fall back to getmem. + bool tma_ok = descs && descs->enabled && is_local + && (bufs.hidden_dim % kGetKChunk == 0); + if (tma_ok) { + // Typed TMA / mbarrier intrinsics from . These compile to the + // same SASS as the hand-written PTX they replace: the peer-X→smem load + // is the exact `cp.async.bulk.tensor.2d.shared::cluster.global.tile` + // instruction, and the smem→staging store the `.global.shared::cta.tile` + // one (`.tile` is the default tiled mode → identical encoding). + namespace ptx = cuda::ptx; + Element* sbuf = reinterpret_cast(bounce_warp); + uint64_t* mbar = reinterpret_cast(bounce_warp + Bounce::kBoxBytes); + + if (lane == 0) + ptx::mbarrier_init(mbar, 1u); + __syncwarp(); + // Make the barrier init visible to the async (TMA) proxy before use. + ptx::fence_proxy_async(ptx::space_shared); + + // Tile coords: each warp owns its kBoxRows-tall row band. + int y_src = info.token_offset + chunk_idx * kBoxRows; // peer X row + int y_dst = slot * TileM + chunk_idx * kBoxRows; // staging row + // Same-host (tma_ok ⇒ is_local) → index by intra-host local rank so the + // descriptor array stays gpus_per_node-sized regardless of total PEs. + const CUtensorMap* src_map = &descs->src_x_desc[info.pe % descs->gpus_per_node]; + const CUtensorMap* dst_map = &descs->dst_staging_desc; + + const int n_k = bufs.hidden_dim / kGetKChunk; + const uint32_t txbytes = (uint32_t)Bounce::kBoxBytes; + uint32_t phase = 0; + for (int k = 0; k < n_k; ++k) { + int x = k * kGetKChunk; + if (lane == 0) { + // Arm the barrier for the incoming box, then issue the + // peer-X → smem TMA load it tracks. + (void)ptx::mbarrier_arrive_expect_tx( + ptx::sem_release, ptx::scope_cta, ptx::space_shared, + mbar, txbytes); + int32_t load_crd[2] = {x, y_src}; + ptx::cp_async_bulk_tensor( + ptx::space_cluster, ptx::space_global, + sbuf, src_map, load_crd, mbar); + // Spin until the load completes (single-buffer bounce). + while (!ptx::mbarrier_try_wait_parity(mbar, phase)) {} + // smem now holds the loaded box; push it to local staging. + int32_t store_crd[2] = {x, y_dst}; + ptx::cp_async_bulk_tensor( + ptx::space_global, ptx::space_shared, + dst_map, store_crd, sbuf); + ptx::cp_async_bulk_commit_group(); + // Single-buffer: drain the store before the next load overwrites smem. + ptx::cp_async_bulk_wait_group(ptx::n32_t<0>{}); + } + phase ^= 1; + __syncwarp(); + } + + if (bufs.tile_expert_ids && chunk_idx == 0 && lane == 0) + bufs.tile_expert_ids[slot] = info.expert; + + // producer_release's __threadfence (device scope) orders the completed + // local TMA stores before the ready signal — MLP reader sees fresh X. + __syncwarp(); + src_pipe.producer_release(lane); + return; + } + + // ── Fallback: warp getmem (IB, or hidden_dim not divisible by kGetKChunk) ── + int chunk = src_tile_elems / kWarpsPerTile; + int offset = chunk_idx * chunk; + auto* local_tokens = static_cast(bufs.local_tokens); + Element* remote_base = local_tokens + info.token_offset * bufs.hidden_dim; + Element* local_base = src_staging + slot * src_tile_elems; + int global_pe = nvshmem_team_translate_pe(bufs.team(), info.pe, NVSHMEM_TEAM_WORLD); + bool is_ib = descs != nullptr && !is_local; + if (is_ib) { + constexpr int kRowNC = NC / MSubTiles; + constexpr int kRowsPerSubtile = TileM / MSubTiles; + int turn = ib_seq % kRowNC; + if (row_chunk_idx == turn) { + int row_start = m_subtile * kRowsPerSubtile; + int valid_rows = max( + 0, min(kRowsPerSubtile, info.valid_rows - row_start)); + int row_offset = row_start * bufs.hidden_dim; + nvshmemx_getmem_warp( + local_base + row_offset, remote_base + row_offset, + (size_t)valid_rows * bufs.hidden_dim * sizeof(Element), + global_pe); + } + ib_seq++; + } else { + nvshmemx_getmem_warp(local_base + offset, remote_base + offset, + (size_t)chunk * sizeof(Element), global_pe); + } + if (bufs.tile_expert_ids && chunk_idx == 0 && lane == 0) { + bufs.tile_expert_ids[slot] = info.expert; + } + src_pipe.producer_release(lane); +} + +template +__device__ __forceinline__ void do_put( + StagePipe& dst_pipe, + const CommBuffers& bufs, + const TileInfo& info, + Element* dst_staging, + int dst_tile_elems, + int xc, // start slot in flat list + int MC, // stride between consecutive comm stages + int yc, // 0..NC-1 — CTA's chunk index within the tile + int m_subtile, + int row_chunk_idx, + bool row_empty, + int& ib_seq, // IB-tile counter (round-robin turn; advanced here) + int lane, + int put_local_idx, // 0..kNumPutWarpsFwd-1 — put warp within the CTA + const GetTmaDescs* descs) { // null → standalone test (real put); else skip local + + // Two put warps per CTA: split the tile kNumPutWarpsFwd·NC ways, this + // warp owns slice (yc·kNumPutWarpsFwd + put_local_idx). + constexpr int kWarpsPerTile = kNumPutWarpsFwd * NC; + int chunk = dst_tile_elems / kWarpsPerTile; + int offset = (yc * kNumPutWarpsFwd + put_local_idx) * chunk; + int my_bytes = chunk * sizeof(Element); + + dst_pipe.consumer_acquire(); + if (row_empty) { + dst_pipe.consumer_release(lane); + return; + } + + // Direct-store: for same-host (NVLink) tiles the GEMM has already TMA-stored + // Y straight into the peer's symmetric local_output, so this put warp must + // NOT copy dst_staging over it (that would clobber the correct Y with a stale + // staging slot). It only RUBBER-STAMPS the dst pipe so the StagePipe counters + // stay consistent and the GEMM's acquire_dst never stalls. Cross-host (IB) + // tiles still need the real putmem. descs==nullptr is the standalone comm + // test (no direct store) → always do the real put. + bool skip_put = (descs != nullptr) && get_pe_is_local(descs, info.pe); + if (!skip_put) { + auto* local_output = static_cast(bufs.local_output); + int slot = xc + dst_pipe.consumer_stage() * MC; + Element* local_base = dst_staging + slot * dst_tile_elems; + Element* remote_base = local_output + info.token_offset * bufs.hidden_dim; + int global_pe = nvshmem_team_translate_pe(bufs.team(), info.pe, NVSHMEM_TEAM_WORLD); + bool is_ib = descs != nullptr; + if (is_ib) { + constexpr int kRowNC = NC / MSubTiles; + constexpr int kRowsPerSubtile = TileM / MSubTiles; + int turn = ib_seq % (kNumPutWarpsFwd * kRowNC); + int row_warp_idx = + row_chunk_idx * kNumPutWarpsFwd + put_local_idx; + if (row_warp_idx == turn) { + int row_start = m_subtile * kRowsPerSubtile; + int valid_rows = max( + 0, min(kRowsPerSubtile, info.valid_rows - row_start)); + int row_offset = row_start * bufs.hidden_dim; + nvshmemx_putmem_warp( + remote_base + row_offset, local_base + row_offset, + (size_t)valid_rows * bufs.hidden_dim * sizeof(Element), + global_pe); + } + ib_seq++; + } else { + nvshmemx_putmem_warp(remote_base + offset, local_base + offset, + my_bytes, global_pe); + } + } + + dst_pipe.consumer_release(lane); +} + +// Chunk-only put for the cooperative LAST-tile drain (epilogue get-warp put +// swap). The dedicated put warp owns the pipe (acquire/release); each of the +// PutW comm warps just writes its share — NO pipe ops here, so no extra +// dst_ready/dst_consumed signalling. P2P splits the tile PutW*NC ways; IB +// sends the whole tile from rank 0 only. dst_tile_elems must be divisible by +// PutW*NC (TileM*D divisible by 2*NC holds for the supported shapes). +template +__device__ __forceinline__ void do_put_chunk( + const CommBuffers& bufs, const TileInfo& info, + Element* dst_staging, int dst_tile_elems, + int slot, int yc, int put_rank, int lane) { + constexpr int kWarpsPerTile = PutW * NC; + int put_idx = yc * PutW + put_rank; + int chunk = dst_tile_elems / kWarpsPerTile; + int offset = put_idx * chunk; + Element* local_base = dst_staging + (size_t)slot * dst_tile_elems; + Element* remote_base = static_cast(bufs.local_output) + + (size_t)info.token_offset * bufs.hidden_dim; + int global_pe = nvshmem_team_translate_pe(bufs.team(), info.pe, NVSHMEM_TEAM_WORLD); + // nbi: the last tile's put has no downstream consumer waiting on it, so it + // need not block — completion is guaranteed by the nvshmem_quiet() each + // comm warp issues after comm_main returns (moe.cu). Single per-chunk put, + // no transport branch (NVLink-only build) → one call site. + nvshmemx_putmem_nbi_warp(remote_base + offset, local_base + offset, + (size_t)chunk * sizeof(Element), global_pe); +} + +// ═══════════════════════════════════════════════════════════════════ +// broadcast_expert_offsets — called after sort, before MLP phase. +// ═══════════════════════════════════════════════════════════════════ + +__device__ __forceinline__ void broadcast_expert_offsets( + const CommBuffers& bufs) { + + int cta_id = blockIdx.x + blockIdx.y * gridDim.x; + int num_experts = bufs.num_experts; + int num_pes = bufs.num_pes(); + int my_pe = bufs.my_pe(); + int offsets_stride = num_experts + 1; + const int* my_offsets = bufs.expert_offsets; + + // CTA 0: copy into own local all_expert_offsets[my_pe][:]. + if (cta_id == 0 && threadIdx.x == 0) { + int* local_dst = bufs.all_expert_offsets + my_pe * offsets_stride; + for (int i = 0; i <= num_experts; ++i) + local_dst[i] = my_offsets[i]; + int* local_counts_dst = + bufs.all_expert_counts + my_pe * num_experts; + for (int i = 0; i < num_experts; ++i) + local_counts_dst[i] = bufs.expert_counts[i]; + } + + // CTA i puts offsets to PE i (if valid remote PE). + // No signal here — signal_dispatch() after dispatch covers it. + // NVSHMEM ordering: when dispatch signal arrives, this put is visible. + int target_pe = cta_id; + if (target_pe >= num_pes || target_pe == my_pe) return; + + int global_target = nvshmem_team_translate_pe(bufs.team(), target_pe, NVSHMEM_TEAM_WORLD); + int* remote_dst = bufs.all_expert_offsets + my_pe * offsets_stride; + nvshmemx_int_put_nbi_block( + remote_dst, my_offsets, + offsets_stride, + global_target); + int* remote_counts_dst = + bufs.all_expert_counts + my_pe * num_experts; + nvshmemx_int_put_nbi_block( + remote_counts_dst, bufs.expert_counts, + num_experts, + global_target); +} + +// ═══════════════════════════════════════════════════════════════════ +// Prologue: waits for remote offsets, computes total_tiles. +// ═══════════════════════════════════════════════════════════════════ +// +// Called by all threads — comm warps do the work, others pass through. +// Must be followed by __syncthreads() before nvshmem_comm_main. + +template +__device__ __forceinline__ void nvshmem_comm_prologue( + CommSmem& smem, + const CommBuffers& bufs, + int static_nsplit) { + + int warp_id = threadIdx.x / 32; + if (warp_id < kCommWarpStart || warp_id > kCommWarpEnd) + return; + + int lane = threadIdx.x % 32; + int local_warp = warp_id - kCommWarpStart; + bool is_leader = (local_warp == 0 && lane == 0); + + int experts_per_pe = bufs.experts_per_pe; + int num_pes = bufs.num_pes(); + int my_pe = bufs.my_pe(); + + if (bufs.all_expert_offsets == nullptr) { + // Standalone test (no nvshmem offsets) — no tiles. + if (is_leader) { + smem.total_tiles = 0; + smem.per_cta_tiles = 0; + smem.global_total = 0; + smem.runtime_nsplit = static_nsplit; + smem.runtime_n_gemm = bufs.n_gemm; + smem.runtime_grid_x = bufs.n_gemm / static_nsplit; + } + __syncwarp(); + return; + } + // NOTE: num_pes == 1 is NO LONGER short-circuited. The unified TileIterator + // enumerates the LOCAL PE (p=0), so a single-GPU run stages its local tiles + // through the comm get/put path exactly like the multi-GPU remote path — + // this is the apples-to-apples isolation of the per-tile staging cost. + + // wait_dispatch() was called before entering the MLP phase, + // guaranteeing all remote PEs' offsets are visible. + + // Extract this PE's local expert offsets from each remote PE's full array. + // all_expert_offsets layout: [num_pes][num_experts + 1]. + // smem layout: [num_pes][experts_per_pe + 1] — only this PE's experts. + int local_e_start = my_pe * experts_per_pe; + int src_stride = bufs.num_experts + 1; + int dst_stride = experts_per_pe + 1; + { + int tid = threadIdx.x - kCommThreadStart; + for (int pe = 0; pe < num_pes; ++pe) { + for (int i = tid; i <= experts_per_pe; i += (kCommWarpsPerCta * 32)) + smem.remote_offsets[pe * dst_stride + i] = + bufs.all_expert_offsets[pe * src_stride + local_e_start + i]; + for (int i = tid; i < experts_per_pe; i += (kCommWarpsPerCta * 32)) + smem.remote_counts[pe * experts_per_pe + i] = + bufs.all_expert_counts[ + pe * bufs.num_experts + local_e_start + i]; + } + } + __syncwarp(); + + // Re-index the flat cta_id to the comm (xc, yc) grid. Comm walks + // tiles owned by (xc, yc) with stride MC. MLP walks tiles for this + // blockIdx.x column with stride gridDim.x — a different start AND + // a different stride, so the two counts must be computed separately + // (uneven distribution can make them differ by more than the ratio + // of the strides per CTA). + // Flat-grid launch: blockIdx.x is the flat cta_id; the logical comm/MLP + // grid comes from n_gemm + compile-time NC/N_SPLIT, not gridDim. + // comm: xc = cta_id / NC, MC = n_gemm / NC + // MLP: col = cta_id / N_SPLIT, grid_x = n_gemm / N_SPLIT + int cta_id = blockIdx.x; + // Comm grid is NC-aligned and DECOUPLED from the GEMM's NS-aligned n_gemm: + // recompute the comm count n_comm from the launched grid (NOT n_gemm). The + // comm side uses all NC-complete CTAs; the leftover remainder (xc >= MC) + // gets comm_count == 0 and self-gates. The MLP count below keeps n_gemm. + int n_comm = (gridDim.x / NC) * NC; + int xc = cta_id / NC; + int MC = n_comm / NC; + + TileIterator iter; + iter.remote_offsets = smem.remote_offsets; + iter.offsets_stride = dst_stride; + iter.remote_counts = smem.remote_counts; + iter.counts_stride = experts_per_pe; + iter.experts_per_pe = experts_per_pe; + iter.num_pes = num_pes; + iter.my_pe = my_pe; + + // Comm count: this CTA's per-(xc, yc) workload. + iter.init(xc, MC); + int comm_count = iter.total_tiles; + + // Grid-wide total tile count (start=0, stride=1 → all tiles once). + iter.init(0, 1); + int global_count = iter.total_tiles; + + // Default static-NS MLP count; moe_fused_kernel overwrites this after it + // selects a runtime NS using MLP N-tile divisibility. + int static_grid_x = bufs.n_gemm / static_nsplit; + iter.init(cta_id / static_nsplit, static_grid_x); + int mlp_count = iter.total_tiles; + + if (is_leader) { + smem.total_tiles = comm_count; // read by comm_main + smem.per_cta_tiles = mlp_count; // read by MLP iterator + smem.global_total = global_count; + smem.runtime_nsplit = static_nsplit; + smem.runtime_n_gemm = bufs.n_gemm; + smem.runtime_grid_x = static_grid_x; + } + __syncwarp(); +} + +// Main loop: warps 1, 2 (getters) and warp 3 (putter) each run an +// independent flat loop over the same tile sequence. Called after +// __syncthreads(). +// +// CTA re-indexing: +// cta_id = blockIdx.x * gridDim.y + blockIdx.y +// xc = cta_id / NC (tile-stream index in [0, MC)) +// yc = cta_id % NC (CTA's slot within its xc group) +// MC = (gridDim.x * gridDim.y) / NC (multiply before divide so MC is +// exact when NC ∤ gridDim.y) +// +// Each xc group strides through the remote-tile list by MC; all NC +// CTAs with the same xc cooperate on the same tile. Get side splits +// each tile kNumGetWarpsPerCta × NC ways, with +// chunk_idx = yc * kNumGetWarpsPerCta + local_get_idx. Put side +// splits NC ways with chunk_idx = yc. +// +// Signal arrays and staging are a flat list of L = MC * NumStages +// slots. Comm CTA at xc walks slots {xc, xc+MC, …, xc+(NumStages-1)·MC} +// (stride MC, NumStages entries — full column coverage of L). MLP at +// blockIdx.x walks slots {blockIdx.x, blockIdx.x+gridDim.x, …} (stride +// gridDim.x, L/gridDim.x = NumStages·N_SPLIT/NC entries — looping back +// across the full L). Both sides cover all L slots, and per-column +// subsets are disjoint across blockIdx.x (gridDim.x | L, since +// L = gridDim.x · NumStages · N_SPLIT/NC and NC | NumStages·N_SPLIT). +// +// The flat list bounds how far the getters can lead the putter: the +// getter blocks in producer_acquire once its NumStages-slot subset is +// pending MLP consumption, and the putter blocks in consumer_acquire +// until MLP releases a dst slot. +template +__device__ __forceinline__ void nvshmem_comm_main( + CommSmem& smem, + const CommBuffers& bufs, + const GetTmaDescs* get_descs, // per-peer X load + local staging store (nullptr → getmem) + char* get_bounce) { // base of this CTA's GetBounce smem region + + // Comm-side pipe depth is NumStages — one comm CTA cycle covers + // the full MC × NumStages staging area, independent of N_SPLIT and + // NC. The MLP side is ticket-based (RemoteMlpTileIterator) and + // shares no per-stage_id state, so no divisibility relationship is + // required between NumStages, N_SPLIT and NC. + constexpr int K = NumStages; + + int warp_id = threadIdx.x / 32; + if (warp_id != kCommGetWarp && warp_id != kCommPutWarp) + return; + + int lane = threadIdx.x % 32; + // Flat-grid launch: blockIdx.x is the flat cta_id. (xc, yc) comm grid and + // MC derive from n_gemm + compile-time NC, not gridDim. + int cta_id = blockIdx.x; + // n_comm: NC-aligned comm grid, decoupled from n_gemm, recomputed from the + // launched grid to match the prologue's comm_count. + int n_comm = (gridDim.x / NC) * NC; + int xc = cta_id / NC; + int yc = cta_id % NC; + int MC = n_comm / NC; + static_assert(NC % MSubTiles == 0, + "communication CTAs must divide evenly across GEMM M-subtiles"); + constexpr int kRowNC = NC / MSubTiles; + int m_subtile = yc / kRowNC; + int row_yc = yc % kRowNC; + if (cta_id >= n_comm) return; + + // Per-warp leader: each comm warp must independently spin in + // StagePipe acquire(). __syncwarp() is warp-local. + bool is_leader = (lane == 0); + + int total = smem.total_tiles; + if (total == 0) return; + + int tile_elems = TileM * bufs.hidden_dim; + + // ── Tile iterator: MC-strided through all remote tiles, starting at xc ── + // All three roles visit the same tile sequence in the same order. + TileIterator iter; + iter.remote_offsets = smem.remote_offsets; + iter.offsets_stride = bufs.experts_per_pe + 1; + iter.remote_counts = smem.remote_counts; + iter.counts_stride = bufs.experts_per_pe; + iter.experts_per_pe = bufs.experts_per_pe; + iter.num_pes = bufs.num_pes(); + iter.my_pe = bufs.my_pe(); + iter.init(xc, MC); + + int my_total = iter.total_tiles; + if (my_total == 0) return; + int runtime_nsplit = smem.runtime_nsplit; + + auto* src_staging_base = static_cast(bufs.src_staging); + auto* dst_staging_base = static_cast(bufs.dst_staging); + + // Warp roles (FWD, TMA GET): warp 1 = single TMA get warp; warp 3 = + // single put warp. Warp 2 is deliberately not a comm warp in this variant. + if (warp_id == kCommPutWarp) { + // One put warp per CTA. put_local_idx is 0; it puts its + // 1/(kNumPutWarpsFwd·NC) slice and independently drives the pipe. + // NumProducers = N_SPLIT (GEMM CTAs produce Y), NumConsumers = + // kNumPutWarpsFwd·NC. + const int put_local_idx = 0; + StagePipe dst_pipe; + dst_pipe.init( + bufs.dst_ready + xc * MSubTiles + m_subtile, + bufs.dst_consumed + xc * MSubTiles + m_subtile, + is_leader, MC * MSubTiles, + runtime_nsplit, kNumPutWarpsFwd * kRowNC); + + int ib_seq = 0; // advances only on IB tiles (inside do_put) + for (int i = 0; i < my_total; ++i) { + TileInfo info = iter.next(); + constexpr int kRowsPerSubtile = TileM / MSubTiles; + bool row_empty = + info.valid_rows <= m_subtile * kRowsPerSubtile; + if (row_empty && row_yc == 0 && lane == 0) { + int s = dst_pipe.consumer_stage(); + atomicAdd( + &dst_pipe.ready[s * dst_pipe.stride], + runtime_nsplit); + } + do_put( + dst_pipe, bufs, info, + dst_staging_base, tile_elems, xc, MC, yc, + m_subtile, row_yc, row_empty, ib_seq, lane, + put_local_idx, get_descs); + } + } else if (warp_id == kCommGetWarp) { // single TMA get warp + const int chunk_idx = yc; // kNumGetWarpsFwd == 1 → yc·1 + 0 + + // NumProducers = kNumGetWarpsFwd·NC (= NC: one get warp × NC CTAs), + // NumConsumers = N_SPLIT (all GEMM CTAs consume cooperatively). + StagePipe src_pipe; + src_pipe.init( + bufs.src_ready + xc * MSubTiles + m_subtile, + bufs.src_consumed + xc * MSubTiles + m_subtile, + is_leader, MC * MSubTiles, + kNumGetWarpsFwd * kRowNC, runtime_nsplit); + char* bounce_warp = get_bounce; // single warp → slice 0 + + int ib_seq = 0; // advances only on IB tiles (inside do_get) + for (int i = 0; i < my_total; ++i) { + TileInfo info = iter.next(); + constexpr int kRowsPerSubtile = TileM / MSubTiles; + bool row_empty = + info.valid_rows <= m_subtile * kRowsPerSubtile; + int s = src_pipe.producer_stage(); + do_get( + src_pipe, bufs, info, + src_staging_base, tile_elems, xc, MC, chunk_idx, + m_subtile, row_yc, row_empty, ib_seq, lane, + get_descs, bounce_warp); + if (row_empty && row_yc == 0 && lane == 0) { + atomicAdd( + &src_pipe.consumed[s * src_pipe.stride], + runtime_nsplit); + } + } + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm90.cuh new file mode 100644 index 00000000..a06f5bad --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm90.cuh @@ -0,0 +1,1343 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MoE MLP with NVSHMEM communication +// ═══════════════════════════════════════════════════════════════════ +// +// Combines nvshmem get/put with MLP compute in a single persistent +// kernel. Warps 1-3 of WG0 handle nvshmem communication while the +// rest of the CTA runs the MLP1 pipeline. +// +// Warp roles (384 threads = 12 warps): +// +// WG0 (warps 0–3): +// warp 0 TMA producer (mlp1 pipelines) +// warp 1 NVSHMEM get warp 0 (fetches X into src_staging) +// warp 2 NVSHMEM get warp 1 (fetches X into src_staging) +// warp 3 NVSHMEM put warp (sends Y from dst_staging) +// +// WG1 (warps 4–7): Consumer A (mlp1 WGMMA) +// WG2 (warps 8–11): Consumer B (mlp1 WGMMA) +// +// Staging areas (per blockIdx.x, shared by both blockIdx.y): +// src_staging[NumStages][TileM][hidden_dim] +// dst_staging[NumStages][TileM][intermediate_dim] +// +// Schedule: the two get warps and the put warp each run an +// independent flat loop over the same tile sequence. The NumStages +// staging ring bounds how far the get warps can lead the put warp. +// +// NVSHMEM get throughput is roughly half of put; we counteract that +// by issuing 2× as many gets per tile. Each tile is chunked across +// 2·NC warps on the get side (2 warps per CTA × NC CTAs) but only NC +// warps on the put side. Get chunk = tile_bytes / (2·NC); put chunk +// = tile_bytes / NC. +// +// Transport-dependent chunking: +// * NVLink P2P (nvshmem_ptr != nullptr): all 2·NC get warps / NC put +// warps cooperate per tile, each copying its chunk with coalesced +// int4 loads/stores — many warps maximise NVLink throughput. +// * IB (nvshmem_ptr == nullptr): the per-warp chunking is counter- +// productive (small RDMA messages, per-tile rendezvous). Instead a +// SINGLE warp issues one whole-tile getmem/putmem, and the warp that +// issues rotates round-robin across the 2·NC (get) / NC (put) warps +// keyed by the tile sequence index, so consecutive tiles are issued +// by different warps. Their blocking transfers then overlap (bounded +// by the NumStages staging ring). The non-issuing warps do no I/O and +// just rubber-stamp the StagePipe so the producer/consumer counters +// stay intact. +// +// ═══════════════════════════════════════════════════════════════════ + +#include // CUtensorMap (raw TMA descriptor) +#include +#include +#include // typed TMA / mbarrier intrinsics (replaces inline PTX) +#include +#include +#include +#include "nvshmem_helpers.cuh" // g_dest_table + +#include "cta_barrier_sm90.cuh" // SyncThreadsCtaCounterBarrier + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// StagePipe — bidirectional acquire/release for a staging buffer +// ═══════════════════════════════════════════════════════════════════ +// +// Wraps a pair of atomic counters (ready[S], consumed[S]) shared +// between a producer and a consumer. Each side calls acquire/release: +// +// Producer: acquire() waits on consumed[] (slot free), then +// release() signals ready[] (data written). +// Consumer: acquire() waits on ready[] (data available), then +// release() signals consumed[] (slot freed). +// +// Internally tracks stage index and monotonic thresholds. +// is_leader is set at init — only the leader thread polls atomics. + +template +struct StagePipe { + int* ready; // producer signals, consumer waits + int* consumed; // consumer signals, producer waits + int stride; // step in slots between consecutive logical stages + bool is_leader; + int num_producers; + int num_consumers; + + int prod_count; + int cons_count; + int prod_expected[NumStages]; // threshold for consumed[] (producer waits on) + int cons_expected[NumStages]; // threshold for ready[] (consumer waits on) + + // Stride lets the pipe address into a flat slot list with non-unit + // stage spacing (e.g. comm pipes use stride = MC, MLP pipes use + // stride = gridDim.x). Default stride = 1 preserves the original + // contiguous-ring layout used by the bwd path. + __device__ void init(int* ready_, int* consumed_, bool is_leader_, + int stride_ = 1) { + ready = ready_; + consumed = consumed_; + stride = stride_; + is_leader = is_leader_; + num_producers = NumProducers; + num_consumers = NumConsumers; + prod_count = 0; + cons_count = 0; + for (int i = 0; i < NumStages; ++i) { + prod_expected[i] = 0; + cons_expected[i] = 0; + } + } + + __device__ void init(int* ready_, int* consumed_, bool is_leader_, + int stride_, int num_producers_, int num_consumers_) { + init(ready_, consumed_, is_leader_, stride_); + num_producers = num_producers_; + num_consumers = num_consumers_; + } + + // ── Producer side ─────────────────────────────────────── + + __device__ int producer_stage() const { + return prod_count % NumStages; + } + + // Wait until this slot is free (consumer has released it). + // + // Acquire-side memory ordering: only the leader does the atomic poll + // (cheaper than 32 lanes spinning), but ALL lanes need to see the + // counterparty's writes after this returns. __syncwarp alone does NOT + // establish global-memory ordering — non-leader lanes may still hit + // stale L1/L2. We add a __threadfence after the leader's poll AND after + // __syncwarp so every lane has a release/acquire pair. + __device__ void producer_acquire() { + int s = producer_stage(); + if (prod_count >= NumStages) { + if (is_leader) { + cuda::atomic_ref r(consumed[s * stride]); + while (r.load(cuda::memory_order_relaxed) < prod_expected[s]) {} + } + __syncwarp(); + __threadfence(); + } + prod_expected[s] += num_consumers; + } + + // Signal that data has been written to this slot. + __device__ void producer_release(int lane) { + int s = producer_stage(); + // Device-scope: this is the get-side publish (NIC-written src → MLP) or + // a slot-free signal — both GPU-internal w.r.t. the consumer. The NIC's + // inbound-write visibility is NOT fixable by a fence (needs GDRCopy), so + // system scope here only costs bandwidth without buying correctness. The + // one publish that genuinely crosses to the NIC (dst → put warp) keeps + // system scope in RemoteIter::release_dst. + __threadfence(); + if (lane == 0) { + atomicAdd(&ready[s * stride], 1); + } + prod_count++; + } + + // ── Consumer side ─────────────────────────────────────── + + __device__ int consumer_stage() const { + return cons_count % NumStages; + } + + // Wait until this slot has data (producer has released it). + // + // Acquire-side memory ordering: see comment on producer_acquire above. + // __syncwarp alone does NOT establish global-memory visibility for + // non-leader lanes — they may load stale data from L1/L2 even after + // the leader has confirmed the producer's release. The __threadfence + // after __syncwarp gives every lane in the warp a fresh view of the + // producer's writes (HBM coherence) before the warp-collective load + // in the subsequent put. + __device__ void consumer_acquire() { + int s = consumer_stage(); + if (is_leader) { + cuda::atomic_ref r(ready[s * stride]); + while (r.load(cuda::memory_order_relaxed) < cons_expected[s] + num_producers) {} + } + __syncwarp(); + // Device-scope: the put warp observes dst_staging (Y), already flushed + // to HBM by the MLP's system-scope release_dst. Observe is GPU-internal. + __threadfence(); + cons_expected[s] += num_producers; + } + + // Signal that this slot has been consumed and is free. + __device__ void consumer_release(int lane) { + int s = consumer_stage(); + __threadfence(); + if (lane == 0) { + atomicAdd(&consumed[s * stride], 1); + } + cons_count++; + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// CommSmem — per-CTA shared memory (smem-local, not shared across CTAs) +// ═══════════════════════════════════════════════════════════════════ + +struct CommSmem { + // Total number of remote tiles across all CTAs. Initialized to -1. + // Comm warps write the actual count; consumers wait for >= 0. + int total_tiles; + + // Per-CTA tile count (this blockIdx.x's share of total_tiles). + int per_cta_tiles; + + // Grid-wide remote tile count (sum across all CTAs/columns). Same + // value on every CTA in the grid; used by moe_fused_bwd to gate the + // remote phase consistently — the inner mlp_fused_bwd's cross-CTA + // mlp_global_barrier requires every CTA to make the same enter/skip + // decision, so per-CTA per_cta_tiles can't be the gate. + int global_total; + int runtime_nsplit; // N-split CTAs per GEMM M-row + int runtime_n_gemm; + int runtime_grid_x; + + // Pre-fetched remote expert offsets, assigned to dynamic smem at runtime. + // Layout: [num_pes][experts_per_pe + 1] + int* remote_offsets; + // Exact routed-row counts matching remote_offsets. + // Layout: [num_pes][experts_per_pe] + int* remote_counts; +}; + +__device__ __forceinline__ int select_runtime_nsplit_from_tiles( + int total_tiles, int num_n_tiles_1, int num_n_tiles_2, + int num_k_tiles_1, int num_k_tiles_2, + int num_blocks, int static_nsplit, int m_subtiles = 1) { + if (total_tiles <= 0) return 2 * m_subtiles; + constexpr int base_candidates[5] = {2, 4, 6, 8, 16}; + int best_base_ns = static_nsplit; + long long best_cost = -1; + long long best_waste = -1; + int best_n_gemm = -1; + long long phase1_weight = 2LL * num_k_tiles_1; // gate + up + long long phase2_weight = (long long)num_k_tiles_2; + #pragma unroll + for (int i = 0; i < 5; ++i) { + int base_ns = base_candidates[i]; + if (base_ns * m_subtiles > num_blocks) continue; + int n_gemm = (num_blocks / base_ns) * base_ns; + int ms = n_gemm / base_ns; + if (ms <= 0) continue; + int m_waves = (total_tiles + ms - 1) / ms; + int n1_waves = (num_n_tiles_1 + base_ns - 1) / base_ns; + int n2_waves = (num_n_tiles_2 + base_ns - 1) / base_ns; + + long long phase_cost = phase1_weight * n1_waves + + phase2_weight * n2_waves; + long long cost = (long long)m_waves * phase_cost; + long long phase1_waste = + (long long)m_waves * ms * n1_waves * base_ns + - (long long)total_tiles * num_n_tiles_1; + long long phase2_waste = + (long long)m_waves * ms * n2_waves * base_ns + - (long long)total_tiles * num_n_tiles_2; + long long waste = phase1_weight * phase1_waste + + phase2_weight * phase2_waste; + if (best_cost < 0 || + cost < best_cost || + (cost == best_cost && waste < best_waste) || + (cost == best_cost && waste == best_waste && n_gemm > best_n_gemm)) { + best_cost = cost; + best_waste = waste; + best_n_gemm = n_gemm; + best_base_ns = base_ns; + } + } + return best_base_ns * m_subtiles; +} + +// ═══════════════════════════════════════════════════════════════════ +// PeSync — device-level phase synchronization +// ═══════════════════════════════════════════════════════════════════ + +struct PeSync { + int* cta_counter; // [1] device memory — shared monotonic counter + int num_pes; + int my_pe; + nvshmem_team_t team; + + // ── barrier: cross-CTA global memory fence ────────── + // Defers to a persistent SyncThreadsCtaCounterBarrier owned by the + // caller — see cta_barrier.cuh for why the barrier object must be + // constructed ONCE per kernel and reused (target accumulates). + // PeSync cannot own the barrier itself because it lives inside a + // __grid_constant__ MoeParams; the caller builds the barrier on + // the kernel stack and threads it into every signal/wait call. + __device__ void barrier(SyncThreadsCtaCounterBarrier& b) const { + b.wait(); + } + + __device__ void barrier_all(SyncThreadsCtaCounterBarrier& b) const { + b.wait(); + int cta_id = blockIdx.x + blockIdx.y * gridDim.x; + if (cta_id == 0) { + if (num_pes > 1 && threadIdx.x == 0) + nvshmem_barrier(team); + __syncthreads(); + } + b.wait(); + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// CommBuffers — pointers passed to the kernel +// ═══════════════════════════════════════════════════════════════════ + +struct CommBuffers { + // Staging areas: per-blockIdx.x, shared by both blockIdx.y CTAs. + // Layout: [gridDim.x][NumStages][TileM][dim] + void* src_staging; // [gridDim.x * S * TileM * hidden_dim] + void* dst_staging; // [gridDim.x * S * TileM * intermediate_dim] + + // Stage pipe signals in device memory (HBM), per-blockIdx.x. + // Layout: [gridDim.x][NumStages] — shared by both blockIdx.y CTAs. + int* src_ready; // nvshmem → consumer: tile data written + int* src_consumed; // consumer → nvshmem: tile data read + int* dst_ready; // consumer → nvshmem: tile data written + int* dst_consumed; // nvshmem → consumer: put completed, slot free + + const int* expert_offsets; // [E + 1] local (from radix sort, TileM-aligned, symmetric) + const int* expert_counts; // [E] exact routed rows per expert + const int* sorted_token_ids; // [total_slots] local sorted token IDs + + void* local_tokens; // [T, hidden_dim] local input (symmetric) + void* local_output; // [total_slots, intermediate_dim] local output (symmetric) + + // Per-staging-slot expert IDs in HBM. Written by do_get, read by MLP. + // Layout: [gridDim.x][NumStages] + int* tile_expert_ids; + + // All PEs' expert offsets, gathered into local symmetric memory. + // Layout: [num_pes][num_experts + 1] — symmetric. + int* all_expert_offsets; + // All PEs' exact expert counts. Layout: [num_pes][num_experts]. + int* all_expert_counts; + + // Cross-PE synchronization for offset broadcast/consume. + PeSync pe_sync; + + int hidden_dim; + int num_experts; // total global experts + int experts_per_pe; // num_experts / num_pes (evenly split) + int all_expert_offsets_stride; // max_num_experts + 1 (fixed across configs) + // GEMM/comm-active CTA count = floor_NS (largest multiple of NSplit <= + // num_sms). The flat-grid launch may run more CTAs than this (grid-stride + // phases use them); GEMM and comm gate to the first n_gemm CTAs so the + // comm↔MLP staging ring geometry (grid_x = n_gemm/NSplit, MC = n_gemm/NC) + // stays consistent. Replaces the old gridDim.x·gridDim.y product. + int n_gemm; + + // Convenience accessors for fields stored in pe_sync. + __host__ __device__ __forceinline__ int my_pe() const { return pe_sync.my_pe; } + __host__ __device__ __forceinline__ int num_pes() const { return pe_sync.num_pes; } + __host__ __device__ __forceinline__ nvshmem_team_t team() const { return pe_sync.team; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// TileInfo — describes one tile to get or put +// ═══════════════════════════════════════════════════════════════════ + +struct TileInfo { + int pe; // remote PE + int token_offset; // offset into remote PE's token buffer (in tokens) + int expert; // local expert index (for weight selection) + int valid_rows; // valid prefix rows in this communication tile + int m_subtile; // GemmTileM row within this communication tile + int parent_tile; // flat communication-tile index +}; + +// ═══════════════════════════════════════════════════════════════════ +// TileIterator — enumerates tiles: expert-outer, PE-inner, strided +// ═══════════════════════════════════════════════════════════════════ +// +// Tiles within each (expert, pe) pair are strided by gridDim.x. +// Each blockIdx.x starts at its own offset and advances by gridDim.x. + +template +struct TileIterator { + static constexpr int TileM = TileM_; + + const int* remote_offsets; // smem: [num_pes][experts_per_pe + 1] + int offsets_stride; // experts_per_pe + 1 + const int* remote_counts; // smem: [num_pes][experts_per_pe] (nullable) + int counts_stride; // experts_per_pe + int experts_per_pe; + int num_pes; + int my_pe; + + int stride; // distance between consecutive tiles + + // Current position in the iteration. + int le; // local expert index + int pe_idx; // 1..num_pes-1 (skips my_pe) + int tile; // tile within current (expert, pe) pair + int remote_tiles; // total tiles for current (expert, pe) pair + int remote_start; // start offset for current (expert, pe) + int remote_count; // exact valid rows for current (expert, pe) + + int total_tiles; // this iterator's tile count + + __device__ void init(int start, int str) { + stride = str; + le = 0; + pe_idx = 0; + total_tiles = 0; + + // Count total tiles across all (expert, pe) pairs and + // compute this iterator's share. + // NOTE: p starts at 0 → the LOCAL PE (my_pe, via g_dest_table slot + // offset 0) is enumerated first for each expert, alongside the remote + // PEs (p=1..num_pes-1). Unified path: local experts are staged by the + // comm get warp from local symmetric memory just like remote tiles, so + // the separate local MLP pass is gone (see moe.cuh moe_fused_fwd). + int lane = threadIdx.x & 31; + int local_total = 0; + int n_experts = experts_per_pe * num_pes; + for (int idx = lane; idx < n_experts; idx += 32) { + int e = idx / num_pes; + int p = idx - e * num_pes; + int pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + p) % num_pes]; + int rs = remote_offsets[pe * offsets_stride + e]; + int re = remote_offsets[pe * offsets_stride + e + 1]; + local_total += (re - rs) / TileM; + } + #pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) { + local_total += __shfl_down_sync(0xffffffff, local_total, delta); + } + int global_total = __shfl_sync(0xffffffff, local_total, 0); + if (start < global_total) + total_tiles = (global_total - 1 - start) / stride + 1; + + // Find the starting position: walk the flat tile space + // to find where start lands. + int skip = start; + for (int e = 0; e < experts_per_pe; ++e) { + for (int p = 0; p < num_pes; ++p) { + int pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + p) % num_pes]; + int rs = remote_offsets[pe * offsets_stride + e]; + int re = remote_offsets[pe * offsets_stride + e + 1]; + int ntiles = (re - rs) / TileM; + if (skip < ntiles) { + le = e; + pe_idx = p; + tile = skip; + remote_start = rs; + remote_tiles = ntiles; + remote_count = remote_counts != nullptr + ? remote_counts[pe * counts_stride + e] + : ntiles * TileM; + return; + } + skip -= ntiles; + } + } + // No tiles for this iterator. + le = experts_per_pe; + } + + __device__ bool has_next() const { + return le < experts_per_pe; + } + + __device__ TileInfo next() { + int pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; + TileInfo info; + info.pe = pe; + info.token_offset = remote_start + tile * TileM; + info.expert = le; + info.valid_rows = max(0, min(TileM, remote_count - tile * TileM)); + info.m_subtile = 0; + info.parent_tile = -1; + + // Advance by stride in the flat tile space. + // Carry the remainder across (expert, pe) boundaries. + tile += stride; + while (tile >= remote_tiles && le < experts_per_pe) { + int remaining = tile - remote_tiles; + pe_idx++; + if (pe_idx >= num_pes) { + pe_idx = 0; + le++; + } + if (le >= experts_per_pe) break; + int next_pe = liger::g_dest_table[(liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; + remote_start = remote_offsets[next_pe * offsets_stride + le]; + int remote_end = remote_offsets[next_pe * offsets_stride + le + 1]; + remote_tiles = (remote_end - remote_start) / TileM; + remote_count = remote_counts != nullptr + ? remote_counts[next_pe * counts_stride + le] + : remote_tiles * TileM; + tile = remaining; + } + return info; + } +}; + +template +struct GemmTileIterator { + static constexpr int TileM = TileM_; + static constexpr int MSubTiles = MSubTiles_; + static constexpr int GemmTileM = TileM / MSubTiles; + + const int* remote_offsets; + int offsets_stride; + const int* remote_counts; + int counts_stride; + int experts_per_pe; + int num_pes; + int my_pe; + int stride; + + int le; + int pe_idx; + int tile; + int remote_tiles; + int remote_parent_tiles; + int remote_start; + int remote_count; + int parent_prefix; + int total_tiles; + + __device__ void init(int start, int str) { + stride = str; + le = 0; + pe_idx = 0; + total_tiles = 0; + + int lane = threadIdx.x & 31; + int local_total = 0; + int n_experts = experts_per_pe * num_pes; + for (int idx = lane; idx < n_experts; idx += 32) { + int e = idx / num_pes; + int p = idx - e * num_pes; + int pe = liger::g_dest_table[ + (liger::comm_slot_of(my_pe) + p) % num_pes]; + int count = remote_counts[pe * counts_stride + e]; + local_total += (count + GemmTileM - 1) / GemmTileM; + } + #pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) + local_total += __shfl_down_sync(0xffffffff, local_total, delta); + int global_total = __shfl_sync(0xffffffff, local_total, 0); + if (start < global_total) + total_tiles = (global_total - 1 - start) / stride + 1; + + int skip = start; + parent_prefix = 0; + for (int e = 0; e < experts_per_pe; ++e) { + for (int p = 0; p < num_pes; ++p) { + int pe = liger::g_dest_table[ + (liger::comm_slot_of(my_pe) + p) % num_pes]; + int rs = remote_offsets[pe * offsets_stride + e]; + int re = remote_offsets[pe * offsets_stride + e + 1]; + int count = remote_counts[pe * counts_stride + e]; + int ntiles = (count + GemmTileM - 1) / GemmTileM; + int nparents = (re - rs) / TileM; + if (skip < ntiles) { + le = e; + pe_idx = p; + tile = skip; + remote_tiles = ntiles; + remote_parent_tiles = nparents; + remote_start = rs; + remote_count = count; + return; + } + skip -= ntiles; + parent_prefix += nparents; + } + } + le = experts_per_pe; + } + + __device__ bool has_next() const { + return le < experts_per_pe; + } + + __device__ int current_parent_tile() const { + return parent_prefix + tile / MSubTiles; + } + + __device__ int current_m_subtile() const { + return tile % MSubTiles; + } + + __device__ TileInfo next() { + int pe = liger::g_dest_table[ + (liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; + int parent_in_expert = tile / MSubTiles; + TileInfo info; + info.pe = pe; + info.token_offset = remote_start + parent_in_expert * TileM; + info.expert = le; + info.valid_rows = max( + 0, min(TileM, remote_count - parent_in_expert * TileM)); + info.m_subtile = tile % MSubTiles; + info.parent_tile = parent_prefix + parent_in_expert; + + tile += stride; + while (tile >= remote_tiles && le < experts_per_pe) { + int remaining = tile - remote_tiles; + parent_prefix += remote_parent_tiles; + pe_idx++; + if (pe_idx >= num_pes) { + pe_idx = 0; + le++; + } + if (le >= experts_per_pe) + break; + int next_pe = liger::g_dest_table[ + (liger::comm_slot_of(my_pe) + pe_idx) % num_pes]; + remote_start = + remote_offsets[next_pe * offsets_stride + le]; + int remote_end = + remote_offsets[next_pe * offsets_stride + le + 1]; + remote_count = + remote_counts[next_pe * counts_stride + le]; + remote_tiles = + (remote_count + GemmTileM - 1) / GemmTileM; + remote_parent_tiles = + (remote_end - remote_start) / TileM; + tile = remaining; + } + return info; + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// NVSHMEM Comm device functions — dedicated get/put warps +// ═══════════════════════════════════════════════════════════════════ +// +// Per CTA: warp 2 is the getter, warp 3 is the putter. +// +// The physical grid is launched as (gridDim.x, N_SPLIT) to match the +// GEMM's cooperative-WG layout. Comms re-indexes the flat CTA id into +// a (MC, NC) logical grid where NC must divide N_SPLIT: +// +// cta_id = blockIdx.x * gridDim.y + blockIdx.y +// xc = cta_id / NC in [0, MC) +// yc = cta_id % NC in [0, NC) — chunk index +// MC = gridDim.x * (gridDim.y / NC) +// +// NC CTAs (same xc, different yc) cooperate on one tile and chunk its +// bytes NC ways. MC groups stride through the remote-tile list, so +// more tiles are in flight per kernel launch when NC < N_SPLIT. +// +// A CTA's comm tile is chosen by xc; its GEMM tile is chosen by +// blockIdx.x. These can diverge when NC < N_SPLIT (the whole point of +// the re-indexing) — the staging ring is keyed by xc so MLP and comm +// rendezvous through the ring, not through shared CTA coordinates. +// +// The two roles run independent flat loops over all tiles. No 1F1B — +// the staging ring (NumStages) bounds how far the getter can lead +// the putter via the shared src/dst StagePipe signals. + +// Prologue + bwd use warps 2-3 only (kept for backward compatibility +// with the bwd path that has not been migrated to the 2-getter layout). +static constexpr int kCommWarpStart = 2; +static constexpr int kCommWarpEnd = 2; +static constexpr int kCommWarpsPerCta = kCommWarpEnd - kCommWarpStart + 1; // 2 +static constexpr int kCommThreadStart = kCommWarpStart * 32; + +// Fwd comm layout: 2 get warps + 1 put warp per CTA. +// warp 1 → get warp 0 +// warp 2 → get warp 1 (also prologue leader) +// warp 3 → put warp +static constexpr int kCommGetWarp = 1; +static constexpr int kCommPutWarp = 2; +static constexpr int kNumGetWarpsPerCta = 2; // shared w/ bwd — do NOT change + +// FWD-only active get-warp count. With the TMA GET (DMA engine, issue-bound +// not thread-bound), ONE warp saturates the per-SM TMA unit, so warp 2 is +// freed (idle in steady state; still helps the epilogue last-tile put drain). +// kNumGetWarpsPerCta stays 2 for the bwd path and for the 2-warp epilogue +// drain (kPutW); only the steady-state FWD producer count drops to this. +static constexpr int kNumGetWarpsFwd = 1; + +// FWD-only put-warp count. Keep one active put warp on Hopper; the freed warp +// enters the MLP-side branch and exits there for Compute != 100. +static constexpr int kNumPutWarpsFwd = 1; + +// Range of warps participating in fwd comm (used by callers to filter +// MLP vs comm warps). +static constexpr int kCommFwdWarpStart = kCommGetWarp; +static constexpr int kCommFwdWarpEnd = kCommPutWarp; + +// Get side has kNumGetWarpsPerCta × NC warps per tile (chunk = +// tile_bytes / (kNumGetWarpsPerCta × NC)); put side has NC warps per +// tile (chunk = tile_bytes / NC). + +// ═══════════════════════════════════════════════════════════════════ +// GET-via-TMA: fetch X over NVLink with the TMA engine, not warp LD/ST. +// ═══════════════════════════════════════════════════════════════════ +// +// The dedicated get warp drives a TMA load (peer X → smem) followed by a +// TMA store (smem → local src_staging HBM). TMA can't go global→global on +// Hopper, so the bounce through smem is mandatory; the store TARGET is +// LOCAL HBM, so completion (cp.async.bulk.wait_group) + a device-scope +// __threadfence in producer_release makes it visible to the MLP reader — +// the safe-completion path (unlike a store to a PEER, which failed before). +// +// Validated in microbench: ~370 GB/s, in- and out-of-kernel correct. +// +// kGetKChunk is the contiguous column box width (elements). Each of the +// kNumGetWarpsPerCta·NC warps owns a TileM/(kNumGetWarpsPerCta·NC) row +// band of the tile and loops over hidden_dim/kGetKChunk column boxes. +// Box = [kGetBoxRows, kGetKChunk]; descriptors built host-side in moe.cu. +static constexpr int kGetKChunk = 256; + +// Per-peer source descriptors + one local-staging dst descriptor, built on +// the host (nvshmem_ptr(x_sorted, peer)) and passed __grid_constant__. +// FWD-only: NOT referenced by bwd — kept out of CommBuffers so the shared +// struct's ABI is unchanged and moe_bwd.cu need not recompile. +// +// Multi-host: only SAME-HOST peers are P2P-reachable, so src_x_desc is indexed +// by intra-host local rank (pe % gpus_per_node), NOT by global team PE — the +// array therefore only ever needs gpus_per_node (≤ kGetMaxPes) entries +// regardless of total PE count. Remote (IB) peers fall back to getmem and never +// touch this array. Assumes ranks are host-contiguous and host-aligned to +// gpus_per_node (true for a WORLD-contiguous team); asserted host-side. +static constexpr int kGetMaxPes = 8; // upper bound on GPUs per host +struct GetTmaDescs { + CUtensorMap src_x_desc[kGetMaxPes]; // index by (TileInfo::pe % gpus_per_node) + CUtensorMap dst_staging_desc; // local src_staging + int enabled; // 1 = NVLink P2P TMA path available; 0 = getmem + int gpus_per_node; // local host PE count — host/local-rank divisor + int my_pe; // this PE in team space (for same-host test) +}; + +// Same-host predicate (pure arithmetic): true iff peer `pe` shares this PE's +// host. Used both to gate the TMA get path and to pick the release_dst fence +// scope (device for same-host NVLink, system for cross-host IB). +__host__ __device__ __forceinline__ bool get_pe_is_local( + const GetTmaDescs* descs, int pe) { + if (descs == nullptr || descs->gpus_per_node <= 0) return true; // single-host/test + return (pe / descs->gpus_per_node) == (descs->my_pe / descs->gpus_per_node); +} + +// Per-peer Y-store descriptors for the direct GEMM→peer TMA store (FWD-only). +// desc[r] is a CUTE TMA-store object built host-side over +// nvshmem_ptr(local_output, peer) for the same-host peer at intra-host local +// rank r (self included, rank = my_pe % gpus_per_node). When a tile is_local, +// MLP2 stores Y straight into the destination peer's symmetric local_output via +// desc[MlpTileInfo::peer_rank] instead of dst_staging — the put warp then only +// rubber-stamps. The direct store is gated per-tile on is_local (same-host → +// P2P guaranteed), so no runtime enable flag is needed. Templated on the CUTE +// store type so the same object the wrapper passes as tma_store_y populates it. +template +struct PeerYStoreDescs { + int enabled = 1; + TmaStoreY desc[kGetMaxPes]; // index by (TileInfo::pe % gpus_per_node) +}; + +// Smem bounce: per get-warp single box + its mbarrier, 128-aligned so the +// TMA dst smem address is aligned. Stride between the two get warps' regions. +template +struct GetBounce { + static constexpr int kWarpsPerTile = kNumGetWarpsFwd * NC; + static constexpr int kBoxRows = TileM / kWarpsPerTile; + static constexpr int kBoxElems = kBoxRows * kGetKChunk; + static constexpr int kBoxBytes = kBoxElems * (int)sizeof(Element); + // [box (128-aligned)] [mbar:8B], next warp's region 128-aligned. + static constexpr int kPerWarpStride = ((kBoxBytes + 8) + 127) & ~127; + static constexpr int kTotalBytes = kPerWarpStride * kNumGetWarpsFwd; +}; + +template +__device__ __forceinline__ void do_get( + StagePipe& src_pipe, + const CommBuffers& bufs, + const TileInfo& info, + Element* src_staging, // base of flat staging buffer + int src_tile_elems, + int xc, // start slot in flat list + int MC, // stride between consecutive comm stages + int chunk_idx, // 0..kNumGetWarpsPerCta * NC - 1 + int m_subtile, + int row_chunk_idx, + bool row_empty, + int& ib_seq, // IB-tile counter (round-robin turn; advanced here) + int lane, + const GetTmaDescs* descs, // per-peer src + local dst TMA descriptors (nullptr → getmem) + char* bounce_warp) { // this warp's GetBounce sub-region (128-aligned) + + using Bounce = GetBounce; + constexpr int kWarpsPerTile = Bounce::kWarpsPerTile; + constexpr int kBoxRows = Bounce::kBoxRows; + + src_pipe.producer_acquire(); + + // No per-tile signal wait — wait_dispatch() already confirmed + // all dispatch writes on all PEs are complete. + + int slot = xc + src_pipe.producer_stage() * MC; + if (row_empty) { + src_pipe.producer_release(lane); + return; + } + + // Locality of this tile's source PE (same-host NVLink vs cross-host IB). + // Pure arithmetic; gates the TMA get path below. The MLP's release_dst + // recomputes the SAME predicate independently (embedded TileIterator), so no + // per-slot flag is published here. + bool is_local = get_pe_is_local(descs, info.pe); + + // TMA path: peer X → smem → local src_staging, via the TMA engine. + // Gated on same-host NVLink P2P (is_local), TMA availability (descs->enabled), + // and hidden_dim divisibility; cross-host IB, odd shapes, or a null descs + // (e.g. the standalone test) fall back to getmem. + bool tma_ok = descs && descs->enabled && is_local + && (bufs.hidden_dim % kGetKChunk == 0); + if (tma_ok) { + // Typed TMA / mbarrier intrinsics from . These compile to the + // same SASS as the hand-written PTX they replace: the peer-X→smem load + // is the exact `cp.async.bulk.tensor.2d.shared::cluster.global.tile` + // instruction, and the smem→staging store the `.global.shared::cta.tile` + // one (`.tile` is the default tiled mode → identical encoding). + namespace ptx = cuda::ptx; + Element* sbuf = reinterpret_cast(bounce_warp); + uint64_t* mbar = reinterpret_cast(bounce_warp + Bounce::kBoxBytes); + + if (lane == 0) + ptx::mbarrier_init(mbar, 1u); + __syncwarp(); + // Make the barrier init visible to the async (TMA) proxy before use. + ptx::fence_proxy_async(ptx::space_shared); + + // Tile coords: each warp owns its kBoxRows-tall row band. + int y_src = info.token_offset + chunk_idx * kBoxRows; // peer X row + int y_dst = slot * TileM + chunk_idx * kBoxRows; // staging row + // Same-host (tma_ok ⇒ is_local) → index by intra-host local rank so the + // descriptor array stays gpus_per_node-sized regardless of total PEs. + const CUtensorMap* src_map = &descs->src_x_desc[info.pe % descs->gpus_per_node]; + const CUtensorMap* dst_map = &descs->dst_staging_desc; + + const int n_k = bufs.hidden_dim / kGetKChunk; + const uint32_t txbytes = (uint32_t)Bounce::kBoxBytes; + uint32_t phase = 0; + for (int k = 0; k < n_k; ++k) { + int x = k * kGetKChunk; + if (lane == 0) { + // Arm the barrier for the incoming box, then issue the + // peer-X → smem TMA load it tracks. + (void)ptx::mbarrier_arrive_expect_tx( + ptx::sem_release, ptx::scope_cta, ptx::space_shared, + mbar, txbytes); + int32_t load_crd[2] = {x, y_src}; + ptx::cp_async_bulk_tensor( + ptx::space_cluster, ptx::space_global, + sbuf, src_map, load_crd, mbar); + // Spin until the load completes (single-buffer bounce). + while (!ptx::mbarrier_try_wait_parity(mbar, phase)) {} + // smem now holds the loaded box; push it to local staging. + int32_t store_crd[2] = {x, y_dst}; + ptx::cp_async_bulk_tensor( + ptx::space_global, ptx::space_shared, + dst_map, store_crd, sbuf); + ptx::cp_async_bulk_commit_group(); + // Single-buffer: drain the store before the next load overwrites smem. + ptx::cp_async_bulk_wait_group(ptx::n32_t<0>{}); + } + phase ^= 1; + __syncwarp(); + } + + if (bufs.tile_expert_ids && chunk_idx == 0 && lane == 0) + bufs.tile_expert_ids[slot] = info.expert; + + // producer_release's __threadfence (device scope) orders the completed + // local TMA stores before the ready signal — MLP reader sees fresh X. + __syncwarp(); + src_pipe.producer_release(lane); + return; + } + + // ── Fallback: warp getmem (IB, or hidden_dim not divisible by kGetKChunk) ── + int chunk = src_tile_elems / kWarpsPerTile; + int offset = chunk_idx * chunk; + auto* local_tokens = static_cast(bufs.local_tokens); + Element* remote_base = local_tokens + info.token_offset * bufs.hidden_dim; + Element* local_base = src_staging + slot * src_tile_elems; + int global_pe = nvshmem_team_translate_pe(bufs.team(), info.pe, NVSHMEM_TEAM_WORLD); + bool is_ib = descs != nullptr && !is_local; + if (is_ib) { + constexpr int kRowNC = NC / MSubTiles; + constexpr int kRowsPerSubtile = TileM / MSubTiles; + int turn = ib_seq % kRowNC; + if (row_chunk_idx == turn) { + int row_start = m_subtile * kRowsPerSubtile; + int valid_rows = max( + 0, min(kRowsPerSubtile, info.valid_rows - row_start)); + int row_offset = row_start * bufs.hidden_dim; + nvshmemx_getmem_warp( + local_base + row_offset, remote_base + row_offset, + (size_t)valid_rows * bufs.hidden_dim * sizeof(Element), + global_pe); + } + ib_seq++; + } else { + nvshmemx_getmem_warp(local_base + offset, remote_base + offset, + (size_t)chunk * sizeof(Element), global_pe); + } + if (bufs.tile_expert_ids && chunk_idx == 0 && lane == 0) { + bufs.tile_expert_ids[slot] = info.expert; + } + src_pipe.producer_release(lane); +} + +template +__device__ __forceinline__ void do_put( + StagePipe& dst_pipe, + const CommBuffers& bufs, + const TileInfo& info, + Element* dst_staging, + int dst_tile_elems, + int xc, // start slot in flat list + int MC, // stride between consecutive comm stages + int yc, // 0..NC-1 — CTA's chunk index within the tile + int m_subtile, + int row_chunk_idx, + bool row_empty, + int& ib_seq, // IB-tile counter (round-robin turn; advanced here) + int lane, + int put_local_idx, // 0..kNumPutWarpsFwd-1 — put warp within the CTA + const GetTmaDescs* descs) { // null → standalone test (real put); else skip local + + // Two put warps per CTA: split the tile kNumPutWarpsFwd·NC ways, this + // warp owns slice (yc·kNumPutWarpsFwd + put_local_idx). + constexpr int kWarpsPerTile = kNumPutWarpsFwd * NC; + int chunk = dst_tile_elems / kWarpsPerTile; + int offset = (yc * kNumPutWarpsFwd + put_local_idx) * chunk; + int my_bytes = chunk * sizeof(Element); + + dst_pipe.consumer_acquire(); + if (row_empty) { + dst_pipe.consumer_release(lane); + return; + } + + // Direct-store: for same-host (NVLink) tiles the GEMM has already TMA-stored + // Y straight into the peer's symmetric local_output, so this put warp must + // NOT copy dst_staging over it (that would clobber the correct Y with a stale + // staging slot). It only RUBBER-STAMPS the dst pipe so the StagePipe counters + // stay consistent and the GEMM's acquire_dst never stalls. Cross-host (IB) + // tiles still need the real putmem. descs==nullptr is the standalone comm + // test (no direct store) → always do the real put. + bool skip_put = (descs != nullptr) && get_pe_is_local(descs, info.pe); + if (!skip_put) { + auto* local_output = static_cast(bufs.local_output); + int slot = xc + dst_pipe.consumer_stage() * MC; + Element* local_base = dst_staging + slot * dst_tile_elems; + Element* remote_base = local_output + info.token_offset * bufs.hidden_dim; + int global_pe = nvshmem_team_translate_pe(bufs.team(), info.pe, NVSHMEM_TEAM_WORLD); + bool is_ib = descs != nullptr; + if (is_ib) { + constexpr int kRowNC = NC / MSubTiles; + constexpr int kRowsPerSubtile = TileM / MSubTiles; + int turn = ib_seq % (kNumPutWarpsFwd * kRowNC); + int row_warp_idx = + row_chunk_idx * kNumPutWarpsFwd + put_local_idx; + if (row_warp_idx == turn) { + int row_start = m_subtile * kRowsPerSubtile; + int valid_rows = max( + 0, min(kRowsPerSubtile, info.valid_rows - row_start)); + int row_offset = row_start * bufs.hidden_dim; + nvshmemx_putmem_warp( + remote_base + row_offset, local_base + row_offset, + (size_t)valid_rows * bufs.hidden_dim * sizeof(Element), + global_pe); + } + ib_seq++; + } else { + nvshmemx_putmem_warp(remote_base + offset, local_base + offset, + my_bytes, global_pe); + } + } + + dst_pipe.consumer_release(lane); +} + +// Chunk-only put for the cooperative LAST-tile drain (epilogue get-warp put +// swap). The dedicated put warp owns the pipe (acquire/release); each of the +// PutW comm warps just writes its share — NO pipe ops here, so no extra +// dst_ready/dst_consumed signalling. P2P splits the tile PutW*NC ways; IB +// sends the whole tile from rank 0 only. dst_tile_elems must be divisible by +// PutW*NC (TileM*D divisible by 2*NC holds for the supported shapes). +template +__device__ __forceinline__ void do_put_chunk( + const CommBuffers& bufs, const TileInfo& info, + Element* dst_staging, int dst_tile_elems, + int slot, int yc, int put_rank, int lane) { + constexpr int kWarpsPerTile = PutW * NC; + int put_idx = yc * PutW + put_rank; + int chunk = dst_tile_elems / kWarpsPerTile; + int offset = put_idx * chunk; + Element* local_base = dst_staging + (size_t)slot * dst_tile_elems; + Element* remote_base = static_cast(bufs.local_output) + + (size_t)info.token_offset * bufs.hidden_dim; + int global_pe = nvshmem_team_translate_pe(bufs.team(), info.pe, NVSHMEM_TEAM_WORLD); + // nbi: the last tile's put has no downstream consumer waiting on it, so it + // need not block — completion is guaranteed by the nvshmem_quiet() each + // comm warp issues after comm_main returns (moe.cu). Single per-chunk put, + // no transport branch (NVLink-only build) → one call site. + nvshmemx_putmem_nbi_warp(remote_base + offset, local_base + offset, + (size_t)chunk * sizeof(Element), global_pe); +} + +// ═══════════════════════════════════════════════════════════════════ +// broadcast_expert_offsets — called after sort, before MLP phase. +// ═══════════════════════════════════════════════════════════════════ + +__device__ __forceinline__ void broadcast_expert_offsets( + const CommBuffers& bufs) { + + int cta_id = blockIdx.x + blockIdx.y * gridDim.x; + int num_experts = bufs.num_experts; + int num_pes = bufs.num_pes(); + int my_pe = bufs.my_pe(); + int offsets_stride = num_experts + 1; + const int* my_offsets = bufs.expert_offsets; + + // CTA 0: copy into own local all_expert_offsets[my_pe][:]. + if (cta_id == 0 && threadIdx.x == 0) { + int* local_dst = bufs.all_expert_offsets + my_pe * offsets_stride; + for (int i = 0; i <= num_experts; ++i) + local_dst[i] = my_offsets[i]; + int* local_counts_dst = + bufs.all_expert_counts + my_pe * num_experts; + for (int i = 0; i < num_experts; ++i) + local_counts_dst[i] = bufs.expert_counts[i]; + } + + // CTA i puts offsets to PE i (if valid remote PE). + // No signal here — signal_dispatch() after dispatch covers it. + // NVSHMEM ordering: when dispatch signal arrives, this put is visible. + int target_pe = cta_id; + if (target_pe >= num_pes || target_pe == my_pe) return; + + int global_target = nvshmem_team_translate_pe(bufs.team(), target_pe, NVSHMEM_TEAM_WORLD); + int* remote_dst = bufs.all_expert_offsets + my_pe * offsets_stride; + nvshmemx_int_put_nbi_block( + remote_dst, my_offsets, + offsets_stride, + global_target); + int* remote_counts_dst = + bufs.all_expert_counts + my_pe * num_experts; + nvshmemx_int_put_nbi_block( + remote_counts_dst, bufs.expert_counts, + num_experts, + global_target); +} + +// ═══════════════════════════════════════════════════════════════════ +// Prologue: waits for remote offsets, computes total_tiles. +// ═══════════════════════════════════════════════════════════════════ +// +// Called by all threads — comm warps do the work, others pass through. +// Must be followed by __syncthreads() before nvshmem_comm_main. + +template +__device__ __forceinline__ void nvshmem_comm_prologue( + CommSmem& smem, + const CommBuffers& bufs, + int static_nsplit) { + + int warp_id = threadIdx.x / 32; + if (warp_id < kCommWarpStart || warp_id > kCommWarpEnd) + return; + + int lane = threadIdx.x % 32; + int local_warp = warp_id - kCommWarpStart; + bool is_leader = (local_warp == 0 && lane == 0); + + int experts_per_pe = bufs.experts_per_pe; + int num_pes = bufs.num_pes(); + int my_pe = bufs.my_pe(); + + if (bufs.all_expert_offsets == nullptr) { + // Standalone test (no nvshmem offsets) — no tiles. + if (is_leader) { + smem.total_tiles = 0; + smem.per_cta_tiles = 0; + smem.global_total = 0; + smem.runtime_nsplit = static_nsplit; + smem.runtime_n_gemm = bufs.n_gemm; + smem.runtime_grid_x = bufs.n_gemm / static_nsplit; + } + __syncwarp(); + return; + } + // NOTE: num_pes == 1 is NO LONGER short-circuited. The unified TileIterator + // enumerates the LOCAL PE (p=0), so a single-GPU run stages its local tiles + // through the comm get/put path exactly like the multi-GPU remote path — + // this is the apples-to-apples isolation of the per-tile staging cost. + + // wait_dispatch() was called before entering the MLP phase, + // guaranteeing all remote PEs' offsets are visible. + + // Extract this PE's local expert offsets from each remote PE's full array. + // all_expert_offsets layout: [num_pes][num_experts + 1]. + // smem layout: [num_pes][experts_per_pe + 1] — only this PE's experts. + int local_e_start = my_pe * experts_per_pe; + int src_stride = bufs.num_experts + 1; + int dst_stride = experts_per_pe + 1; + { + int tid = threadIdx.x - kCommThreadStart; + for (int pe = 0; pe < num_pes; ++pe) { + for (int i = tid; i <= experts_per_pe; i += (kCommWarpsPerCta * 32)) + smem.remote_offsets[pe * dst_stride + i] = + bufs.all_expert_offsets[pe * src_stride + local_e_start + i]; + for (int i = tid; i < experts_per_pe; i += (kCommWarpsPerCta * 32)) + smem.remote_counts[pe * experts_per_pe + i] = + bufs.all_expert_counts[ + pe * bufs.num_experts + local_e_start + i]; + } + } + __syncwarp(); + + // Re-index the flat cta_id to the comm (xc, yc) grid. Comm walks + // tiles owned by (xc, yc) with stride MC. MLP walks tiles for this + // blockIdx.x column with stride gridDim.x — a different start AND + // a different stride, so the two counts must be computed separately + // (uneven distribution can make them differ by more than the ratio + // of the strides per CTA). + // Flat-grid launch: blockIdx.x is the flat cta_id; the logical comm/MLP + // grid comes from n_gemm + compile-time NC/N_SPLIT, not gridDim. + // comm: xc = cta_id / NC, MC = n_gemm / NC + // MLP: col = cta_id / N_SPLIT, grid_x = n_gemm / N_SPLIT + int cta_id = blockIdx.x; + // Comm grid is NC-aligned and DECOUPLED from the GEMM's NS-aligned n_gemm: + // recompute the comm count n_comm from the launched grid (NOT n_gemm). The + // comm side uses all NC-complete CTAs; the leftover remainder (xc >= MC) + // gets comm_count == 0 and self-gates. The MLP count below keeps n_gemm. + int n_comm = (gridDim.x / NC) * NC; + int xc = cta_id / NC; + int MC = n_comm / NC; + + TileIterator iter; + iter.remote_offsets = smem.remote_offsets; + iter.offsets_stride = dst_stride; + iter.remote_counts = smem.remote_counts; + iter.counts_stride = experts_per_pe; + iter.experts_per_pe = experts_per_pe; + iter.num_pes = num_pes; + iter.my_pe = my_pe; + + // Comm count: this CTA's per-(xc, yc) workload. + iter.init(xc, MC); + int comm_count = iter.total_tiles; + + // Grid-wide total tile count (start=0, stride=1 → all tiles once). + iter.init(0, 1); + int global_count = iter.total_tiles; + + // Default static-NS MLP count; moe_fused_kernel overwrites this after it + // selects a runtime NS using MLP N-tile divisibility. + int static_grid_x = bufs.n_gemm / static_nsplit; + iter.init(cta_id / static_nsplit, static_grid_x); + int mlp_count = iter.total_tiles; + + if (is_leader) { + smem.total_tiles = comm_count; // read by comm_main + smem.per_cta_tiles = mlp_count; // read by MLP iterator + smem.global_total = global_count; + smem.runtime_nsplit = static_nsplit; + smem.runtime_n_gemm = bufs.n_gemm; + smem.runtime_grid_x = static_grid_x; + } + __syncwarp(); +} + +// Main loop: warps 1, 2 (getters) and warp 3 (putter) each run an +// independent flat loop over the same tile sequence. Called after +// __syncthreads(). +// +// CTA re-indexing: +// cta_id = blockIdx.x * gridDim.y + blockIdx.y +// xc = cta_id / NC (tile-stream index in [0, MC)) +// yc = cta_id % NC (CTA's slot within its xc group) +// MC = (gridDim.x * gridDim.y) / NC (multiply before divide so MC is +// exact when NC ∤ gridDim.y) +// +// Each xc group strides through the remote-tile list by MC; all NC +// CTAs with the same xc cooperate on the same tile. Get side splits +// each tile kNumGetWarpsPerCta × NC ways, with +// chunk_idx = yc * kNumGetWarpsPerCta + local_get_idx. Put side +// splits NC ways with chunk_idx = yc. +// +// Signal arrays and staging are a flat list of L = MC * NumStages +// slots. Comm CTA at xc walks slots {xc, xc+MC, …, xc+(NumStages-1)·MC} +// (stride MC, NumStages entries — full column coverage of L). MLP at +// blockIdx.x walks slots {blockIdx.x, blockIdx.x+gridDim.x, …} (stride +// gridDim.x, L/gridDim.x = NumStages·N_SPLIT/NC entries — looping back +// across the full L). Both sides cover all L slots, and per-column +// subsets are disjoint across blockIdx.x (gridDim.x | L, since +// L = gridDim.x · NumStages · N_SPLIT/NC and NC | NumStages·N_SPLIT). +// +// The flat list bounds how far the getters can lead the putter: the +// getter blocks in producer_acquire once its NumStages-slot subset is +// pending MLP consumption, and the putter blocks in consumer_acquire +// until MLP releases a dst slot. +template +__device__ __forceinline__ void nvshmem_comm_main( + CommSmem& smem, + const CommBuffers& bufs, + const GetTmaDescs* get_descs, // per-peer X load + local staging store (nullptr → getmem) + char* get_bounce) { // base of this CTA's GetBounce smem region + + // Comm-side pipe depth is NumStages — one comm CTA cycle covers + // the full MC × NumStages staging area, independent of N_SPLIT and + // NC. The MLP side is ticket-based (RemoteMlpTileIterator) and + // shares no per-stage_id state, so no divisibility relationship is + // required between NumStages, N_SPLIT and NC. + constexpr int K = NumStages; + + int warp_id = threadIdx.x / 32; + if (warp_id != kCommGetWarp && warp_id != kCommPutWarp) + return; + + int lane = threadIdx.x % 32; + // Flat-grid launch: blockIdx.x is the flat cta_id. (xc, yc) comm grid and + // MC derive from n_gemm + compile-time NC, not gridDim. + int cta_id = blockIdx.x; + // n_comm: NC-aligned comm grid, decoupled from n_gemm, recomputed from the + // launched grid to match the prologue's comm_count. + int n_comm = (gridDim.x / NC) * NC; + int xc = cta_id / NC; + int yc = cta_id % NC; + int MC = n_comm / NC; + static_assert(NC % MSubTiles == 0, + "communication CTAs must divide evenly across GEMM M-subtiles"); + constexpr int kRowNC = NC / MSubTiles; + int m_subtile = yc / kRowNC; + int row_yc = yc % kRowNC; + if (cta_id >= n_comm) return; + + // Per-warp leader: each comm warp must independently spin in + // StagePipe acquire(). __syncwarp() is warp-local. + bool is_leader = (lane == 0); + + int total = smem.total_tiles; + if (total == 0) return; + + int tile_elems = TileM * bufs.hidden_dim; + + // ── Tile iterator: MC-strided through all remote tiles, starting at xc ── + // All three roles visit the same tile sequence in the same order. + TileIterator iter; + iter.remote_offsets = smem.remote_offsets; + iter.offsets_stride = bufs.experts_per_pe + 1; + iter.remote_counts = smem.remote_counts; + iter.counts_stride = bufs.experts_per_pe; + iter.experts_per_pe = bufs.experts_per_pe; + iter.num_pes = bufs.num_pes(); + iter.my_pe = bufs.my_pe(); + iter.init(xc, MC); + + int my_total = iter.total_tiles; + if (my_total == 0) return; + int runtime_nsplit = smem.runtime_nsplit; + + auto* src_staging_base = static_cast(bufs.src_staging); + auto* dst_staging_base = static_cast(bufs.dst_staging); + + // Warp roles (FWD, TMA GET): warp 1 = single TMA get warp; warp 3 = + // single put warp. Warp 2 is deliberately not a comm warp in this variant. + if (warp_id == kCommPutWarp) { + // One put warp per CTA. put_local_idx is 0; it puts its + // 1/(kNumPutWarpsFwd·NC) slice and independently drives the pipe. + // NumProducers = N_SPLIT (GEMM CTAs produce Y), NumConsumers = + // kNumPutWarpsFwd·NC. + const int put_local_idx = 0; + StagePipe dst_pipe; + dst_pipe.init( + bufs.dst_ready + xc * MSubTiles + m_subtile, + bufs.dst_consumed + xc * MSubTiles + m_subtile, + is_leader, MC * MSubTiles, + runtime_nsplit, kNumPutWarpsFwd * kRowNC); + + int ib_seq = 0; // advances only on IB tiles (inside do_put) + for (int i = 0; i < my_total; ++i) { + TileInfo info = iter.next(); + constexpr int kRowsPerSubtile = TileM / MSubTiles; + bool row_empty = + info.valid_rows <= m_subtile * kRowsPerSubtile; + if (row_empty && row_yc == 0 && lane == 0) { + int s = dst_pipe.consumer_stage(); + atomicAdd( + &dst_pipe.ready[s * dst_pipe.stride], + runtime_nsplit); + } + do_put( + dst_pipe, bufs, info, + dst_staging_base, tile_elems, xc, MC, yc, + m_subtile, row_yc, row_empty, ib_seq, lane, + put_local_idx, get_descs); + } + } else if (warp_id == kCommGetWarp) { // single TMA get warp + const int chunk_idx = yc; // kNumGetWarpsFwd == 1 → yc·1 + 0 + + // NumProducers = kNumGetWarpsFwd·NC (= NC: one get warp × NC CTAs), + // NumConsumers = N_SPLIT (all GEMM CTAs consume cooperatively). + StagePipe src_pipe; + src_pipe.init( + bufs.src_ready + xc * MSubTiles + m_subtile, + bufs.src_consumed + xc * MSubTiles + m_subtile, + is_leader, MC * MSubTiles, + kNumGetWarpsFwd * kRowNC, runtime_nsplit); + char* bounce_warp = get_bounce; // single warp → slice 0 + + int ib_seq = 0; // advances only on IB tiles (inside do_get) + for (int i = 0; i < my_total; ++i) { + TileInfo info = iter.next(); + constexpr int kRowsPerSubtile = TileM / MSubTiles; + bool row_empty = + info.valid_rows <= m_subtile * kRowsPerSubtile; + int s = src_pipe.producer_stage(); + do_get( + src_pipe, bufs, info, + src_staging_base, tile_elems, xc, MC, chunk_idx, + m_subtile, row_yc, row_empty, ib_seq, lane, + get_descs, bounce_warp); + if (row_empty && row_yc == 0 && lane == 0) { + atomicAdd( + &src_pipe.consumed[s * src_pipe.stride], + runtime_nsplit); + } + } + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_sm100.cuh new file mode 100644 index 00000000..8b90d834 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_sm100.cuh @@ -0,0 +1,417 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP: Y = (Z @ A^T) where Z = SiLU(B @ X) * (C @ X) +// ═══════════════════════════════════════════════════════════════════ +// +// Pipelines are constructed once in the prologue (requires all threads). +// Pipeline states are persistent across tiles — cooperative 2-WG + +// fused pipes reduce the state register footprint to two states per +// phase (producer + consumer), so hoisting them out of the tile loop +// still fits comfortably in registers. +// +// The tile loop is driven by a TileIterator (local or remote). +// The intermediate Z buffer is double-buffered: z_m alternates +// between blockIdx.x*ZBufferSlots + (iter % ZBufferSlots). +// +// WG0 warps 1-3 do NOT participate — reserved for NVSHMEM comm. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp1_fused_sm100.cuh" +#include "mlp2_fused_sm100.cuh" +#include "tile_iterator_sm100.cuh" +#include "cta_barrier_sm100.cuh" // target-based CtaCounterBarrier + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Fused shared memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct MlpFusedSmem { + union { + Mlp1FusedSmem mlp1; + Mlp2FusedSmem mlp2; + }; + + // Pipeline barrier storage — persistent across phases. + // Both phases use a single fused pipe (one mbarrier per stage). + typename Mlp1MainloopPipelineFor::SharedStorage p1_pipe; + typename Mlp2MainloopPipelineFor::SharedStorage p2_pipe; + alignas(16) uint32_t tmem_base; +}; + +static constexpr int kMlpBarrierId = 14; + +// Target-based monotonic barrier for the FWD MLP warps. Count is Compute-aware: +// Hopper 288 (w0 + w4-11); Blackwell 320 (w0 + w3-11, warp 3 = MLP1 UMMA). Must +// match the in-fn kMlpBarrierThreads — both fire kMlpBarrierId. See cta_barrier.cuh +// for why the modulo variant deadlocks under CTA skew. +template +using MlpFwdCtaBarrierT = CtaCounterBarrier<(Compute == 100 ? 10 : 9) * 32, kMlpBarrierId>; + +// ═══════════════════════════════════════════════════════════════════ +// MlpDims — invariant MLP dimensions, passed as __grid_constant__ +// ═══════════════════════════════════════════════════════════════════ + +struct MlpDims { + int hidden_dim; + int intermediate_dim; + int num_experts; // local experts represented by the weight tensors + int total_n_rows_1; // physical rows spanned by the strided B/C views + int total_n_rows_2; // experts_per_pe * hidden_dim + int expert_n_stride_1; // B/C expert stride in TileN1 units + int num_n_tiles_1; // intermediate_dim / TileN1 + int num_n_tiles_2; // ceildiv(hidden_dim/TileN2, NSub) — per-WG outer steps + int num_k_tiles_1; // hidden_dim / TileK1 + int num_k_tiles_2; // intermediate_dim / TileK2 + int y_buf_m_tiles; // max_total_slots / TileM — row-tile extent of the + // symmetric local_output (y_buf) descriptor, used as + // num_y_m_tiles for the direct-to-peer Y store. + int* phase_counter; +}; + +// ═══════════════════════════════════════════════════════════════════ +// mlp_fused_fwd — generic over TileIterator +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp_fused_fwd( + MlpFusedSmem& smem, + TileIter& iter, + // MLP1 TMA descriptors + TmaLoadX const& tma_load_x, + TmaLoadW1 const& tma_load_b, + TmaLoadW1 const& tma_load_c, + TmaStoreZ const& tma_store_z, + // MLP2 TMA descriptors + TmaLoadZ const& tma_load_z, + TmaLoadW2 const& tma_load_a, + TmaStoreY const& tma_store_y, + // Invariant MLP dims — const ref into __grid_constant__ .param space + const MlpDims& dims, + // Per-phase params (change between local and remote) + int num_tokens, + int num_m_tiles, + // Persistent per-CTA x_barrier — constructed once in moe_fused_fwd + // and threaded through here so target accumulates across all 3 + // mlp_fused_fwd calls (A1 / B / A2). Standalone callers (no + // remote phase) still pass their own barrier instance. + MlpFwdCtaBarrierT & x_barrier, + int local_expert_start = 0, + // Logical column / grid_x for the flat-grid fused launch. Defaults + // (-1) fall back to blockIdx.x / gridDim.x for standalone 2-D-grid + // callers (mlp.cu bench kernel). + int col = -1, + int grid_x = -1, + // N-split identity for the GEMM producers/consumers. Defaults (-1) + // fall back to blockIdx.y / gridDim.y (standalone 2-D-grid mlp.cu). + int split_idx = -1, + int num_splits = -1, + // Direct-to-peer Y store: array of per-peer TMA descriptors over + // nvshmem_ptr(local_output, peer) (indexed by MlpTileInfo::peer_rank). + // For is_local tiles MLP2 stores Y straight into the destination's + // symmetric local_output instead of dst_staging (tma_store_y). nullptr → + // staging path only (standalone mlp.cu callers have no peer descriptors). + const TmaStoreY* y_peer_desc = nullptr) { + + using Pipeline1 = Mlp1MainloopPipelineFor; + using PipeState1 = typename Traits1::PipelineState; + using Pipeline2 = Mlp2MainloopPipelineFor; + using PipeState2 = typename Traits2::PipelineState; + + int warp_id = threadIdx.x / Traits1::WarpSize; + + // MLP-warp NamedBarrier arrival count — must equal MlpFwdCtaBarrierT's + // template count (both fire kMlpBarrierId). Hopper: w0+w4-11 (288). Blackwell: + // w0+w3-11 (320, warp 3 = MLP1 UMMA). + constexpr int kMlpBarrierThreads = (Compute == 100 ? 10 : 9) * 32; + + cute::prefetch_tma_descriptor(tma_load_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_b.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_y.get_tma_descriptor()); + + // ── Prologue: construct all pipelines once ────────── + // Pipeline constructors set per-thread role (Producer / Consumer / + // NonParticipant) so warps 1-3 can safely skip the MLP body below. + auto p1_pipe = [&]() { + if constexpr (Compute == 100) + return mlp1_make_pipe_umma(smem.p1_pipe); + else + return mlp1_make_pipe(smem.p1_pipe); + }(); + auto p2_pipe = [&]() { + if constexpr (Compute == 100) + return mlp2_make_pipe_umma(smem.p2_pipe); + else + return mlp2_make_pipe(smem.p2_pipe); + }(); + + if constexpr (Compute == 90) { + if (warp_id == 3) { + return; + } + } + + // Sync MLP warps only (excludes comm warps 1-3 which may be in + // nvshmem_comm_main when called from moe_fused_fwd). + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + + cute::TMEM::Allocator1Sm tmem_alloc{}; + if constexpr (Compute == 100) { + constexpr int kMlp1Cols = Traits1::AccStages * (2 * Traits1::TileN); + constexpr int kMlp2Cols = Traits2::AccStages * Traits2::TileN; + constexpr int kTmemColumns = (kMlp1Cols > kMlp2Cols) + ? kMlp1Cols : kMlp2Cols; + if (warp_id == 3) { + tmem_alloc.allocate(kTmemColumns, &smem.tmem_base); + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + } + + // Z buffer: [grid_x, ZBufferSlots, SubTiles, GemmTileM, intermediate_dim]. + // Indexed in GemmTileM-row tiles: SubTiles Z regions per comm-tile slot so + // mlp1 can stage all sub-tiles before mlp2 drains them. z_rows is unchanged + // (CommTileM == SubTiles·GemmTileM); only the indexing splits. + int z_col = (col >= 0) ? col : (int)blockIdx.x; + int z_grid_x = (grid_x >= 0) ? grid_x : (int)gridDim.x; + int num_z_m_tiles = z_grid_x * ZBufferSlots * SubTiles; + int z_base = z_col * ZBufferSlots * SubTiles; + // N-split identity for the GEMM producers/consumers (flat-grid: NOT + // blockIdx.y / gridDim.y, which are 0 / 1). Each CTA computes only its + // stripe of N-tiles (n = n_split, n_split+n_count, ...). + int n_split = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_count = (num_splits >= 0) ? num_splits : (int)gridDim.y; + int tile_iter = 0; + + // ── Persistent pipeline states (hoisted out of the tile loop) ── + // With single fused pipes + cooperative consumer, we have only + // 2 states per phase. They persist across tiles since k-tile + // counts are fixed — the barrier phase bit cycles back to 0 + // after each tile's k-loop completes. + PipeState1 p1_state = (warp_id == 0) ? + cutlass::make_producer_start_state() : PipeState1{}; + PipeState1 p1_cons_state; + PipeState2 p2_state = (warp_id == 0) ? + cutlass::make_producer_start_state() : PipeState2{}; + PipeState2 p2_cons_state; + + // x_barrier is passed in by reference (constructed once per CTA + // in moe_fused_fwd) so target advances monotonically across + // the A1 / B / A2 phases that all share dims.phase_counter. + + // ── Tile loop ─────────────────────────────────────── + while (iter.has_next()) { + // Barrier: ensure previous phase 2 is done before reusing smem union. + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + iter.acquire_src(); + auto tile = iter.next(); + + int expert = tile.expert - local_expert_start; + int first_sub = 0; + int last_sub = SubTiles; + if constexpr (SubTiles > 1) { + first_sub = tile.m_subtile; + last_sub = first_sub + 1; + } + // Base Z region for this comm tile. SubTiles GemmTileM-row regions per + // slot, indexed z_m_base + sub. tile.y_m is the comm slot index in + // CommTileM(=GemmTileM·SubTiles) units; the GemmTileM-unit row coordinate + // of sub-tile `sub` in that slot is tile.y_m·SubTiles + sub. + int z_m_base = z_base + (tile_iter % ZBufferSlots) * SubTiles; + + // Phase 1: for each GemmTileM sub-tile, Z[z_m_base+sub] = SiLU(B@X)·(C@X). + // One comm semaphore (acquire_src above) covers all SubTiles; the GEMM + // just steps finer. SubTiles==1 collapses to the original single tile. + for (int sub = first_sub; sub < last_sub; ++sub) { + int x_mt = tile.y_m * SubTiles + sub; + int z_m = z_m_base + sub; + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits1::WarpSize == 0) + smem.mlp1.tmem_base = smem.tmem_base; + if (warp_id == 0) { + if constexpr (Compute == 100) { + mlp1_fused_producer( + p1_pipe, p1_state, + smem.mlp1, tma_load_x, tma_load_b, tma_load_c, + x_mt, expert * dims.expert_n_stride_1, + num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.num_experts, dims.total_n_rows_1, + dims.num_n_tiles_1, dims.num_k_tiles_1, n_split, n_count); + } else { + mlp1_fused_producer( + p1_pipe, p1_state, + smem.mlp1, tma_load_x, tma_load_b, tma_load_c, + x_mt, expert, + num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.num_experts, dims.total_n_rows_1, + dims.num_n_tiles_1, dims.num_k_tiles_1, n_split, n_count); + } + } + + // Fused MoE CTA warp map: warp 0 = TMA, warps 1..3 = NVSHMEM comm + // (warp 3 = put), warps 4..11 = MLP GEMM consumers. Warp 3 is a comm + // warp here, so the fused MLP1 keeps the ORIGINAL warp-4 UMMA design + // (single-warp UMMA + warps 4..11 epilogue). The warp-3 dual-WG + // rewrite is used ONLY by the isolated MLP1 kernel (no comm warps), + // so it neither frees warp 3 from comm nor changes the shared MLP + // barrier count — MLP2 and the rest of the fused pipeline are + // untouched. Compile-time flag (no register cost). Enabling warp-3 + // here would require the comm-warp + barrier migration (pipeline.md + // §5) — intentionally out of scope. + + if (warp_id >= 3 && warp_id <= 11) { + if constexpr (Compute == 100) { + mlp1_fused_consumer( + p1_pipe, p1_cons_state, + smem.mlp1, tma_store_z, z_m, dims.intermediate_dim, + num_z_m_tiles, dims.num_n_tiles_1, dims.num_k_tiles_1, + n_split, n_count); + } else { + mlp1_fused_consumer( + p1_pipe, p1_cons_state, + smem.mlp1, tma_store_z, z_m, dims.intermediate_dim, + num_z_m_tiles, dims.num_n_tiles_1, dims.num_k_tiles_1, + n_split, n_count); + } + } + + if constexpr (Compute == 90) { + if (sub + 1 < SubTiles) + cutlass::arch::NamedBarrier::sync( + kMlpBarrierThreads, kMlpBarrierId); + } + } + + // One phase barrier covers the full communication tile. The persistent + // TMA/UMMA pipeline states sequence its GEMM sub-tiles internally. + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + + // Release src staging (remote: free slot; local: no-op). + // Pass threadIdx.x directly so only thread 0 (TMA producer warp + // leader) has lane==0 and fires the atomicAdd exactly once per + // CTA — StagePipe expects NumConsumers = NC releases per slot. + iter.release_src(threadIdx.x); + + // Cross-CTA sync: wait for ALL y-CTAs sharing this blockIdx.x + // to finish phase 1 for this M-tile (all sub-tiles' Z regions). + x_barrier.wait(); + + // Acquire dst staging (remote: wait for put slot; local: no-op). + iter.acquire_dst(); + + // Y-store target select. Same-host (is_local) tiles store Y DIRECTLY + // into the destination peer's symmetric local_output via its per-peer + // descriptor at the token-tile row; all other tiles store to the + // staging ring (tma_store_y) at the slot row. mlp1's X-load above still + // uses tile.y_m (the staging slot) regardless — input path unchanged. + // Tile counts are in CommTileM units → scale by SubTiles for GemmTileM boxes. + bool y_direct = (y_peer_desc != nullptr && tile.is_local); + const TmaStoreY& y_desc = y_direct ? y_peer_desc[tile.peer_rank] : tma_store_y; + int y_base = y_direct ? tile.y_store_m : tile.y_m; + int y_ntiles = (y_direct ? dims.y_buf_m_tiles : num_m_tiles) * SubTiles; + // Phase 2: for each sub-tile, Y[y_base·SubTiles+sub] = Z[z_m_base+sub] @ A^T. + // Cooperative-M-split: both consumer WGs cooperate on the same (m,n) tile. + for (int sub = first_sub; sub < last_sub; ++sub) { + int z_m = z_m_base + sub; + int y_coord = y_base * SubTiles + sub; + constexpr int kM64PerGemmTile = Traits2::TileM / 64; + int valid_m64_subtiles = max( + 0, min(kM64PerGemmTile, + tile.valid_m64_subtiles - sub * kM64PerGemmTile)); + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits2::WarpSize == 0) + smem.mlp2.tmem_base = smem.tmem_base; + if (warp_id == 0) { + if constexpr (Compute == 100) { + mlp2_fused_producer( + p2_pipe, p2_state, + smem.mlp2, tma_load_z, tma_load_a, + z_m, expert * dims.num_n_tiles_2, + num_z_m_tiles * Traits1::TileM, dims.hidden_dim, + dims.intermediate_dim, dims.num_experts, dims.total_n_rows_2, + dims.num_n_tiles_2, dims.num_k_tiles_2, n_split, n_count); + } else { + mlp2_fused_producer( + p2_pipe, p2_state, + smem.mlp2, tma_load_z, tma_load_a, + z_m, expert, + num_z_m_tiles * Traits1::TileM, dims.hidden_dim, + dims.intermediate_dim, dims.num_experts, dims.total_n_rows_2, + dims.num_n_tiles_2, dims.num_k_tiles_2, n_split, n_count); + } + } + + if (warp_id >= 3 && warp_id <= 11) { + if constexpr (Compute == 100) { + mlp2_fused_consumer( + p2_pipe, p2_cons_state, + smem.mlp2, y_desc, + y_coord, dims.hidden_dim, + y_ntiles, dims.num_n_tiles_2, dims.num_k_tiles_2, + n_split, n_count, valid_m64_subtiles); + } else { + mlp2_fused_consumer( + p2_pipe, p2_cons_state, + smem.mlp2, y_desc, + y_coord, dims.hidden_dim, + y_ntiles, dims.num_n_tiles_2, dims.num_k_tiles_2, + n_split, n_count); + } + } + + if constexpr (Compute == 90) { + if (sub + 1 < SubTiles) + cutlass::arch::NamedBarrier::sync( + kMlpBarrierThreads, kMlpBarrierId); + } + } + + // Release dst staging (remote: cross-CTA barrier + signal; local: no-op). + // Pass threadIdx.x directly — exactly one atomicAdd per CTA + // (threadIdx.x == 0), matching NumProducers = NC on dst_pipe. + // Device-scope: the system-scope flush of Y to HBM for the NIC put is + // done by RemoteIter::release_dst below (all threads). This is just the + // intra-CTA ordering before the MLP-warp barrier. + __threadfence(); + // As in phase 1, synchronize once after all GEMM sub-tiles. + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + iter.release_dst(threadIdx.x); + + tile_iter++; + } + + if constexpr (Compute == 100) { + constexpr int kMlp1Cols = Traits1::AccStages * (2 * Traits1::TileN); + constexpr int kMlp2Cols = Traits2::AccStages * Traits2::TileN; + constexpr int kTmemColumns = (kMlp1Cols > kMlp2Cols) + ? kMlp1Cols : kMlp2Cols; + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + if (warp_id == 3) { + tmem_alloc.release_allocation_lock(); + tmem_alloc.free(smem.tmem_base, kTmemColumns); + } + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/mlp_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/mlp_sm90.cuh new file mode 100644 index 00000000..dfb7a424 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/mlp_sm90.cuh @@ -0,0 +1,417 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MLP: Y = (Z @ A^T) where Z = SiLU(B @ X) * (C @ X) +// ═══════════════════════════════════════════════════════════════════ +// +// Pipelines are constructed once in the prologue (requires all threads). +// Pipeline states are persistent across tiles — cooperative 2-WG + +// fused pipes reduce the state register footprint to two states per +// phase (producer + consumer), so hoisting them out of the tile loop +// still fits comfortably in registers. +// +// The tile loop is driven by a TileIterator (local or remote). +// The intermediate Z buffer is double-buffered: z_m alternates +// between blockIdx.x*ZBufferSlots + (iter % ZBufferSlots). +// +// WG0 warps 1-3 do NOT participate — reserved for NVSHMEM comm. +// +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp1_fused_sm90.cuh" +#include "mlp2_fused_sm90.cuh" +#include "tile_iterator_sm90.cuh" +#include "cta_barrier_sm90.cuh" // target-based CtaCounterBarrier + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// Fused shared memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct MlpFusedSmem { + union { + Mlp1FusedSmem mlp1; + Mlp2FusedSmem mlp2; + }; + + // Pipeline barrier storage — persistent across phases. + // Both phases use a single fused pipe (one mbarrier per stage). + typename Mlp1MainloopPipelineFor::SharedStorage p1_pipe; + typename Mlp2MainloopPipelineFor::SharedStorage p2_pipe; + alignas(16) uint32_t tmem_base; +}; + +static constexpr int kMlpBarrierId = 14; + +// Target-based monotonic barrier for the FWD MLP warps. Count is Compute-aware: +// Hopper 288 (w0 + w4-11); Blackwell 320 (w0 + w3-11, warp 3 = MLP1 UMMA). Must +// match the in-fn kMlpBarrierThreads — both fire kMlpBarrierId. See cta_barrier.cuh +// for why the modulo variant deadlocks under CTA skew. +template +using MlpFwdCtaBarrierT = CtaCounterBarrier<(Compute == 100 ? 10 : 9) * 32, kMlpBarrierId>; + +// ═══════════════════════════════════════════════════════════════════ +// MlpDims — invariant MLP dimensions, passed as __grid_constant__ +// ═══════════════════════════════════════════════════════════════════ + +struct MlpDims { + int hidden_dim; + int intermediate_dim; + int num_experts; // local experts represented by the weight tensors + int total_n_rows_1; // physical rows spanned by the strided B/C views + int total_n_rows_2; // experts_per_pe * hidden_dim + int expert_n_stride_1; // B/C expert stride in TileN1 units + int num_n_tiles_1; // intermediate_dim / TileN1 + int num_n_tiles_2; // ceildiv(hidden_dim/TileN2, NSub) — per-WG outer steps + int num_k_tiles_1; // hidden_dim / TileK1 + int num_k_tiles_2; // intermediate_dim / TileK2 + int y_buf_m_tiles; // max_total_slots / TileM — row-tile extent of the + // symmetric local_output (y_buf) descriptor, used as + // num_y_m_tiles for the direct-to-peer Y store. + int* phase_counter; +}; + +// ═══════════════════════════════════════════════════════════════════ +// mlp_fused_fwd — generic over TileIterator +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void mlp_fused_fwd( + MlpFusedSmem& smem, + TileIter& iter, + // MLP1 TMA descriptors + TmaLoadX const& tma_load_x, + TmaLoadW1 const& tma_load_b, + TmaLoadW1 const& tma_load_c, + TmaStoreZ const& tma_store_z, + // MLP2 TMA descriptors + TmaLoadZ const& tma_load_z, + TmaLoadW2 const& tma_load_a, + TmaStoreY const& tma_store_y, + // Invariant MLP dims — const ref into __grid_constant__ .param space + const MlpDims& dims, + // Per-phase params (change between local and remote) + int num_tokens, + int num_m_tiles, + // Persistent per-CTA x_barrier — constructed once in moe_fused_fwd + // and threaded through here so target accumulates across all 3 + // mlp_fused_fwd calls (A1 / B / A2). Standalone callers (no + // remote phase) still pass their own barrier instance. + MlpFwdCtaBarrierT & x_barrier, + int local_expert_start = 0, + // Logical column / grid_x for the flat-grid fused launch. Defaults + // (-1) fall back to blockIdx.x / gridDim.x for standalone 2-D-grid + // callers (mlp.cu bench kernel). + int col = -1, + int grid_x = -1, + // N-split identity for the GEMM producers/consumers. Defaults (-1) + // fall back to blockIdx.y / gridDim.y (standalone 2-D-grid mlp.cu). + int split_idx = -1, + int num_splits = -1, + // Direct-to-peer Y store: array of per-peer TMA descriptors over + // nvshmem_ptr(local_output, peer) (indexed by MlpTileInfo::peer_rank). + // For is_local tiles MLP2 stores Y straight into the destination's + // symmetric local_output instead of dst_staging (tma_store_y). nullptr → + // staging path only (standalone mlp.cu callers have no peer descriptors). + const TmaStoreY* y_peer_desc = nullptr) { + + using Pipeline1 = Mlp1MainloopPipelineFor; + using PipeState1 = typename Traits1::PipelineState; + using Pipeline2 = Mlp2MainloopPipelineFor; + using PipeState2 = typename Traits2::PipelineState; + + int warp_id = threadIdx.x / Traits1::WarpSize; + + // MLP-warp NamedBarrier arrival count — must equal MlpFwdCtaBarrierT's + // template count (both fire kMlpBarrierId). Hopper: w0+w4-11 (288). Blackwell: + // w0+w3-11 (320, warp 3 = MLP1 UMMA). + constexpr int kMlpBarrierThreads = (Compute == 100 ? 10 : 9) * 32; + + cute::prefetch_tma_descriptor(tma_load_x.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_b.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_c.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_z.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_load_a.get_tma_descriptor()); + cute::prefetch_tma_descriptor(tma_store_y.get_tma_descriptor()); + + // ── Prologue: construct all pipelines once ────────── + // Pipeline constructors set per-thread role (Producer / Consumer / + // NonParticipant) so warps 1-3 can safely skip the MLP body below. + auto p1_pipe = [&]() { + if constexpr (Compute == 100) + return mlp1_make_pipe_umma(smem.p1_pipe); + else + return mlp1_make_pipe(smem.p1_pipe); + }(); + auto p2_pipe = [&]() { + if constexpr (Compute == 100) + return mlp2_make_pipe_umma(smem.p2_pipe); + else + return mlp2_make_pipe(smem.p2_pipe); + }(); + + if constexpr (Compute == 90) { + if (warp_id == 3) { + return; + } + } + + // Sync MLP warps only (excludes comm warps 1-3 which may be in + // nvshmem_comm_main when called from moe_fused_fwd). + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + + cute::TMEM::Allocator1Sm tmem_alloc{}; + if constexpr (Compute == 100) { + constexpr int kMlp1Cols = Traits1::AccStages * (2 * Traits1::TileN); + constexpr int kMlp2Cols = Traits2::AccStages * Traits2::TileN; + constexpr int kTmemColumns = (kMlp1Cols > kMlp2Cols) + ? kMlp1Cols : kMlp2Cols; + if (warp_id == 3) { + tmem_alloc.allocate(kTmemColumns, &smem.tmem_base); + __syncwarp(); + } + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + } + + // Z buffer: [grid_x, ZBufferSlots, SubTiles, GemmTileM, intermediate_dim]. + // Indexed in GemmTileM-row tiles: SubTiles Z regions per comm-tile slot so + // mlp1 can stage all sub-tiles before mlp2 drains them. z_rows is unchanged + // (CommTileM == SubTiles·GemmTileM); only the indexing splits. + int z_col = (col >= 0) ? col : (int)blockIdx.x; + int z_grid_x = (grid_x >= 0) ? grid_x : (int)gridDim.x; + int num_z_m_tiles = z_grid_x * ZBufferSlots * SubTiles; + int z_base = z_col * ZBufferSlots * SubTiles; + // N-split identity for the GEMM producers/consumers (flat-grid: NOT + // blockIdx.y / gridDim.y, which are 0 / 1). Each CTA computes only its + // stripe of N-tiles (n = n_split, n_split+n_count, ...). + int n_split = (split_idx >= 0) ? split_idx : (int)blockIdx.y; + int n_count = (num_splits >= 0) ? num_splits : (int)gridDim.y; + int tile_iter = 0; + + // ── Persistent pipeline states (hoisted out of the tile loop) ── + // With single fused pipes + cooperative consumer, we have only + // 2 states per phase. They persist across tiles since k-tile + // counts are fixed — the barrier phase bit cycles back to 0 + // after each tile's k-loop completes. + PipeState1 p1_state = (warp_id == 0) ? + cutlass::make_producer_start_state() : PipeState1{}; + PipeState1 p1_cons_state; + PipeState2 p2_state = (warp_id == 0) ? + cutlass::make_producer_start_state() : PipeState2{}; + PipeState2 p2_cons_state; + + // x_barrier is passed in by reference (constructed once per CTA + // in moe_fused_fwd) so target advances monotonically across + // the A1 / B / A2 phases that all share dims.phase_counter. + + // ── Tile loop ─────────────────────────────────────── + while (iter.has_next()) { + // Barrier: ensure previous phase 2 is done before reusing smem union. + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + iter.acquire_src(); + auto tile = iter.next(); + + int expert = tile.expert - local_expert_start; + int first_sub = 0; + int last_sub = SubTiles; + if constexpr (SubTiles > 1) { + first_sub = tile.m_subtile; + last_sub = first_sub + 1; + } + // Base Z region for this comm tile. SubTiles GemmTileM-row regions per + // slot, indexed z_m_base + sub. tile.y_m is the comm slot index in + // CommTileM(=GemmTileM·SubTiles) units; the GemmTileM-unit row coordinate + // of sub-tile `sub` in that slot is tile.y_m·SubTiles + sub. + int z_m_base = z_base + (tile_iter % ZBufferSlots) * SubTiles; + + // Phase 1: for each GemmTileM sub-tile, Z[z_m_base+sub] = SiLU(B@X)·(C@X). + // One comm semaphore (acquire_src above) covers all SubTiles; the GEMM + // just steps finer. SubTiles==1 collapses to the original single tile. + for (int sub = first_sub; sub < last_sub; ++sub) { + int x_mt = tile.y_m * SubTiles + sub; + int z_m = z_m_base + sub; + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits1::WarpSize == 0) + smem.mlp1.tmem_base = smem.tmem_base; + if (warp_id == 0) { + if constexpr (Compute == 100) { + mlp1_fused_producer( + p1_pipe, p1_state, + smem.mlp1, tma_load_x, tma_load_b, tma_load_c, + x_mt, expert * dims.expert_n_stride_1, + num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.num_experts, dims.total_n_rows_1, + dims.num_n_tiles_1, dims.num_k_tiles_1, n_split, n_count); + } else { + mlp1_fused_producer( + p1_pipe, p1_state, + smem.mlp1, tma_load_x, tma_load_b, tma_load_c, + x_mt, expert, + num_tokens, dims.hidden_dim, dims.intermediate_dim, + dims.num_experts, dims.total_n_rows_1, + dims.num_n_tiles_1, dims.num_k_tiles_1, n_split, n_count); + } + } + + // Fused MoE CTA warp map: warp 0 = TMA, warps 1..3 = NVSHMEM comm + // (warp 3 = put), warps 4..11 = MLP GEMM consumers. Warp 3 is a comm + // warp here, so the fused MLP1 keeps the ORIGINAL warp-4 UMMA design + // (single-warp UMMA + warps 4..11 epilogue). The warp-3 dual-WG + // rewrite is used ONLY by the isolated MLP1 kernel (no comm warps), + // so it neither frees warp 3 from comm nor changes the shared MLP + // barrier count — MLP2 and the rest of the fused pipeline are + // untouched. Compile-time flag (no register cost). Enabling warp-3 + // here would require the comm-warp + barrier migration (pipeline.md + // §5) — intentionally out of scope. + + if (warp_id >= 3 && warp_id <= 11) { + if constexpr (Compute == 100) { + mlp1_fused_consumer( + p1_pipe, p1_cons_state, + smem.mlp1, tma_store_z, z_m, dims.intermediate_dim, + num_z_m_tiles, dims.num_n_tiles_1, dims.num_k_tiles_1, + n_split, n_count); + } else { + mlp1_fused_consumer( + p1_pipe, p1_cons_state, + smem.mlp1, tma_store_z, z_m, dims.intermediate_dim, + num_z_m_tiles, dims.num_n_tiles_1, dims.num_k_tiles_1, + n_split, n_count); + } + } + + if constexpr (Compute == 90) { + if (sub + 1 < SubTiles) + cutlass::arch::NamedBarrier::sync( + kMlpBarrierThreads, kMlpBarrierId); + } + } + + // One phase barrier covers the full communication tile. The persistent + // TMA/UMMA pipeline states sequence its GEMM sub-tiles internally. + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + + // Release src staging (remote: free slot; local: no-op). + // Pass threadIdx.x directly so only thread 0 (TMA producer warp + // leader) has lane==0 and fires the atomicAdd exactly once per + // CTA — StagePipe expects NumConsumers = NC releases per slot. + iter.release_src(threadIdx.x); + + // Cross-CTA sync: wait for ALL y-CTAs sharing this blockIdx.x + // to finish phase 1 for this M-tile (all sub-tiles' Z regions). + x_barrier.wait(); + + // Acquire dst staging (remote: wait for put slot; local: no-op). + iter.acquire_dst(); + + // Y-store target select. Same-host (is_local) tiles store Y DIRECTLY + // into the destination peer's symmetric local_output via its per-peer + // descriptor at the token-tile row; all other tiles store to the + // staging ring (tma_store_y) at the slot row. mlp1's X-load above still + // uses tile.y_m (the staging slot) regardless — input path unchanged. + // Tile counts are in CommTileM units → scale by SubTiles for GemmTileM boxes. + bool y_direct = (y_peer_desc != nullptr && tile.is_local); + const TmaStoreY& y_desc = y_direct ? y_peer_desc[tile.peer_rank] : tma_store_y; + int y_base = y_direct ? tile.y_store_m : tile.y_m; + int y_ntiles = (y_direct ? dims.y_buf_m_tiles : num_m_tiles) * SubTiles; + // Phase 2: for each sub-tile, Y[y_base·SubTiles+sub] = Z[z_m_base+sub] @ A^T. + // Cooperative-M-split: both consumer WGs cooperate on the same (m,n) tile. + for (int sub = first_sub; sub < last_sub; ++sub) { + int z_m = z_m_base + sub; + int y_coord = y_base * SubTiles + sub; + constexpr int kM64PerGemmTile = Traits2::TileM / 64; + int valid_m64_subtiles = max( + 0, min(kM64PerGemmTile, + tile.valid_m64_subtiles - sub * kM64PerGemmTile)); + if constexpr (Compute == 100) + if (warp_id == 3 && threadIdx.x % Traits2::WarpSize == 0) + smem.mlp2.tmem_base = smem.tmem_base; + if (warp_id == 0) { + if constexpr (Compute == 100) { + mlp2_fused_producer( + p2_pipe, p2_state, + smem.mlp2, tma_load_z, tma_load_a, + z_m, expert * dims.num_n_tiles_2, + num_z_m_tiles * Traits1::TileM, dims.hidden_dim, + dims.intermediate_dim, dims.num_experts, dims.total_n_rows_2, + dims.num_n_tiles_2, dims.num_k_tiles_2, n_split, n_count); + } else { + mlp2_fused_producer( + p2_pipe, p2_state, + smem.mlp2, tma_load_z, tma_load_a, + z_m, expert, + num_z_m_tiles * Traits1::TileM, dims.hidden_dim, + dims.intermediate_dim, dims.num_experts, dims.total_n_rows_2, + dims.num_n_tiles_2, dims.num_k_tiles_2, n_split, n_count); + } + } + + if (warp_id >= 3 && warp_id <= 11) { + if constexpr (Compute == 100) { + mlp2_fused_consumer( + p2_pipe, p2_cons_state, + smem.mlp2, y_desc, + y_coord, dims.hidden_dim, + y_ntiles, dims.num_n_tiles_2, dims.num_k_tiles_2, + n_split, n_count, valid_m64_subtiles); + } else { + mlp2_fused_consumer( + p2_pipe, p2_cons_state, + smem.mlp2, y_desc, + y_coord, dims.hidden_dim, + y_ntiles, dims.num_n_tiles_2, dims.num_k_tiles_2, + n_split, n_count); + } + } + + if constexpr (Compute == 90) { + if (sub + 1 < SubTiles) + cutlass::arch::NamedBarrier::sync( + kMlpBarrierThreads, kMlpBarrierId); + } + } + + // Release dst staging (remote: cross-CTA barrier + signal; local: no-op). + // Pass threadIdx.x directly — exactly one atomicAdd per CTA + // (threadIdx.x == 0), matching NumProducers = NC on dst_pipe. + // Device-scope: the system-scope flush of Y to HBM for the NIC put is + // done by RemoteIter::release_dst below (all threads). This is just the + // intra-CTA ordering before the MLP-warp barrier. + __threadfence(); + // As in phase 1, synchronize once after all GEMM sub-tiles. + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + iter.release_dst(threadIdx.x); + + tile_iter++; + } + + if constexpr (Compute == 100) { + constexpr int kMlp1Cols = Traits1::AccStages * (2 * Traits1::TileN); + constexpr int kMlp2Cols = Traits2::AccStages * Traits2::TileN; + constexpr int kTmemColumns = (kMlp1Cols > kMlp2Cols) + ? kMlp1Cols : kMlp2Cols; + cutlass::arch::NamedBarrier::sync(kMlpBarrierThreads, kMlpBarrierId); + if (warp_id == 3) { + tmem_alloc.release_allocation_lock(); + tmem_alloc.free(smem.tmem_base, kTmemColumns); + } + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/models.cuh b/liger_cute_kernels/csrc/core/src/moe/models.cuh index d53d4ac0..b51fa7cf 100644 --- a/liger_cute_kernels/csrc/core/src/moe/models.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/models.cuh @@ -1,189 +1,11 @@ #pragma once -// #ifndef NDEBUG -// #define NDEBUG -// #endif +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -#include -#include -#include -#include -#include -#include -#include - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// GMMA atom selection -// ═══════════════════════════════════════════════════════════════════ - -template -struct GmmaSelector; - -static constexpr auto GmmaK = GMMA::Major::K; -static constexpr auto GmmaMN = GMMA::Major::MN; - -template <> struct GmmaSelector { using Atom = SM90_64x8x16_F32BF16BF16_SS; }; -template <> struct GmmaSelector { using Atom = SM90_64x16x16_F32BF16BF16_SS; }; -template <> struct GmmaSelector { using Atom = SM90_64x32x16_F32BF16BF16_SS; }; -template <> struct GmmaSelector { using Atom = SM90_64x64x16_F32BF16BF16_SS; }; -template <> struct GmmaSelector { using Atom = SM90_64x128x16_F32BF16BF16_SS; }; -template <> struct GmmaSelector { using Atom = SM90_64x192x16_F32BF16BF16_SS; }; -template <> struct GmmaSelector { using Atom = SM90_64x256x16_F32BF16BF16_SS; }; - -template <> struct GmmaSelector { using Atom = SM90_64x8x16_F32F16F16_SS; }; -template <> struct GmmaSelector { using Atom = SM90_64x16x16_F32F16F16_SS; }; -template <> struct GmmaSelector { using Atom = SM90_64x32x16_F32F16F16_SS; }; -template <> struct GmmaSelector { using Atom = SM90_64x64x16_F32F16F16_SS; }; - -// ═══════════════════════════════════════════════════════════════════ -// GMMA atom selection — A=K-major, B=MN-major (SS) -// Used when B operand has MN-contiguous smem (e.g. mlp2_t_fused). -// ═══════════════════════════════════════════════════════════════════ - -template -struct GmmaSelectorKMN; - -template <> struct GmmaSelectorKMN { using Atom = SM90_64x8x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorKMN { using Atom = SM90_64x16x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorKMN { using Atom = SM90_64x32x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorKMN { using Atom = SM90_64x64x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorKMN { using Atom = SM90_64x128x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorKMN { using Atom = SM90_64x256x16_F32BF16BF16_SS; }; - -template <> struct GmmaSelectorKMN { using Atom = SM90_64x8x16_F32F16F16_SS; }; -template <> struct GmmaSelectorKMN { using Atom = SM90_64x16x16_F32F16F16_SS; }; -template <> struct GmmaSelectorKMN { using Atom = SM90_64x32x16_F32F16F16_SS; }; -template <> struct GmmaSelectorKMN { using Atom = SM90_64x64x16_F32F16F16_SS; }; - -// Backward compat alias -template -struct GmmaSelectorMN : GmmaSelectorKMN {}; - -// ═══════════════════════════════════════════════════════════════════ -// GMMA atom selection — both MN-major (SS) -// Used when both operands have MN-contiguous smem (e.g. mlp3). -// ═══════════════════════════════════════════════════════════════════ - -template -struct GmmaSelectorMNMN; - -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x8x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x16x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x32x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x64x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x128x16_F32BF16BF16_SS; }; -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x256x16_F32BF16BF16_SS; }; - -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x8x16_F32F16F16_SS; }; -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x16x16_F32F16F16_SS; }; -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x32x16_F32F16F16_SS; }; -template <> struct GmmaSelectorMNMN { using Atom = SM90_64x64x16_F32F16F16_SS; }; - -// ═══════════════════════════════════════════════════════════════════ -// FusedMoETraits -// ═══════════════════════════════════════════════════════════════════ - -template < - typename Element_, - int TileN_, - int TileK_ = 64, - int Stages_ = 2, - int MaxTopK_ = 8 -> -struct FusedMoETraits { - using Element = Element_; - using ElementAccum = float; - - static constexpr int TileM = 64; - static constexpr int TileN = TileN_; - static constexpr int TileK = TileK_; - static constexpr int Stages = Stages_; - static constexpr int MaxTopK = MaxTopK_; - - using GmmaAtom = typename GmmaSelector::Atom; - using TiledMma = TiledMMA, Layout>>; - - using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom; - - using SmemLayoutA_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{})); - using SmemLayoutB_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{})); - - using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{})); - using SmemLayoutB = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{})); - - using MainloopPipeline = cutlass::PipelineTmaAsync; - using PipelineState = cutlass::PipelineState; - - static constexpr int TmaTransactionBytesA = static_cast(size(SmemLayoutA_1{}) * sizeof(Element)); - static constexpr int TmaTransactionBytesB = static_cast(size(SmemLayoutB_1{}) * sizeof(Element)); - static constexpr int TmaTransactionBytes = TmaTransactionBytesA + TmaTransactionBytesB; - - static constexpr int WarpSize = 32; - static constexpr int WarpGroupSize = 128; - static constexpr int NumThreads = 384; - static constexpr int EpilogueThreads = 2 * WarpSize; // warps 1+2 (64 threads) - static constexpr int EpilogueBarrier = 3; // NamedBarrier ID for epilogue sync - static constexpr int SortWarpId = 3; // warp 3 (threads 96-127) -}; - -// ═══════════════════════════════════════════════════════════════════ -// RouterSmem — shared memory layout for the fused MoE kernel -// ═══════════════════════════════════════════════════════════════════ - -template -struct RouterSmem { - using Element = typename Traits::Element; - - static constexpr int smem_A_size = cosize_v; - static constexpr int smem_B_size = cosize_v; - - // Consumer A smem (even K-tiles) - alignas(128) Element smem_A_a[smem_A_size]; - alignas(128) Element smem_B_a[smem_B_size]; - // Consumer B smem (odd K-tiles) - alignas(128) Element smem_A_b[smem_A_size]; - alignas(128) Element smem_B_b[smem_B_size]; - - // Two TMA pipelines (one per consumer) - typename Traits::MainloopPipeline::SharedStorage pipe_a_storage; - typename Traits::MainloopPipeline::SharedStorage pipe_b_storage; - - alignas(16) float score_strip[2][Traits::TileM][Traits::TileN]; - - alignas(16) float running_max[Traits::TileM]; - alignas(16) float running_sum[Traits::TileM]; - alignas(16) float top_scores[Traits::TileM][Traits::MaxTopK]; - alignas(16) int top_indices[Traits::TileM][Traits::MaxTopK]; - - alignas(16) int strip_ready[2]; - alignas(16) int strip_consumed[2]; - - alignas(8) uint64_t mbar_debug; - - CUTE_DEVICE Element* A_a() { return &smem_A_a[0]; } - CUTE_DEVICE Element* B_a() { return &smem_B_a[0]; } - CUTE_DEVICE Element* A_b() { return &smem_A_b[0]; } - CUTE_DEVICE Element* B_b() { return &smem_B_b[0]; } -}; - -// ═══════════════════════════════════════════════════════════════════ -// SortBuffers — global memory pointers for the sort phase -// ═══════════════════════════════════════════════════════════════════ - -struct SortBuffers { - int* tile_expert_counts; // [num_m_tiles * E] — per-tile histograms, - // modified in-place for prefix sums - int* sorted_token_ids; // [total_slots] — final output: slot → token - int* token_expert_slots; // [T * K] — final output: (token, k) → slot - int* expert_offsets; // [E + 1] — aligned expert boundaries (symmetric) - int* expert_counts; // [E] — exact routed rows per expert (nullable) - int* tile_expert_ids; // [max_total_tiles] — expert ID per M-tile (nullable) - int* cta_done; // [num_blocks] — per-CTA completion flags (Blelloch tree) - int* cta_sums; // [num_blocks * E] — per-CTA expert totals for prefix sum -}; - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "models_sm90.cuh" +#else +#include "models_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/models_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/models_sm100.cuh new file mode 100644 index 00000000..d53d4ac0 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/models_sm100.cuh @@ -0,0 +1,189 @@ +#pragma once + +// #ifndef NDEBUG +// #define NDEBUG +// #endif + +#include +#include +#include +#include +#include +#include +#include + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// GMMA atom selection +// ═══════════════════════════════════════════════════════════════════ + +template +struct GmmaSelector; + +static constexpr auto GmmaK = GMMA::Major::K; +static constexpr auto GmmaMN = GMMA::Major::MN; + +template <> struct GmmaSelector { using Atom = SM90_64x8x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x16x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x32x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x64x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x128x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x192x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x256x16_F32BF16BF16_SS; }; + +template <> struct GmmaSelector { using Atom = SM90_64x8x16_F32F16F16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x16x16_F32F16F16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x32x16_F32F16F16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x64x16_F32F16F16_SS; }; + +// ═══════════════════════════════════════════════════════════════════ +// GMMA atom selection — A=K-major, B=MN-major (SS) +// Used when B operand has MN-contiguous smem (e.g. mlp2_t_fused). +// ═══════════════════════════════════════════════════════════════════ + +template +struct GmmaSelectorKMN; + +template <> struct GmmaSelectorKMN { using Atom = SM90_64x8x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x16x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x32x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x64x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x128x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x256x16_F32BF16BF16_SS; }; + +template <> struct GmmaSelectorKMN { using Atom = SM90_64x8x16_F32F16F16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x16x16_F32F16F16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x32x16_F32F16F16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x64x16_F32F16F16_SS; }; + +// Backward compat alias +template +struct GmmaSelectorMN : GmmaSelectorKMN {}; + +// ═══════════════════════════════════════════════════════════════════ +// GMMA atom selection — both MN-major (SS) +// Used when both operands have MN-contiguous smem (e.g. mlp3). +// ═══════════════════════════════════════════════════════════════════ + +template +struct GmmaSelectorMNMN; + +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x8x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x16x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x32x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x64x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x128x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x256x16_F32BF16BF16_SS; }; + +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x8x16_F32F16F16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x16x16_F32F16F16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x32x16_F32F16F16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x64x16_F32F16F16_SS; }; + +// ═══════════════════════════════════════════════════════════════════ +// FusedMoETraits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileN_, + int TileK_ = 64, + int Stages_ = 2, + int MaxTopK_ = 8 +> +struct FusedMoETraits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = 64; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int MaxTopK = MaxTopK_; + + using GmmaAtom = typename GmmaSelector::Atom; + using TiledMma = TiledMMA, Layout>>; + + using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom; + + using SmemLayoutA_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{})); + using SmemLayoutB_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{})); + + using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{})); + using SmemLayoutB = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{})); + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + static constexpr int TmaTransactionBytesA = static_cast(size(SmemLayoutA_1{}) * sizeof(Element)); + static constexpr int TmaTransactionBytesB = static_cast(size(SmemLayoutB_1{}) * sizeof(Element)); + static constexpr int TmaTransactionBytes = TmaTransactionBytesA + TmaTransactionBytesB; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumThreads = 384; + static constexpr int EpilogueThreads = 2 * WarpSize; // warps 1+2 (64 threads) + static constexpr int EpilogueBarrier = 3; // NamedBarrier ID for epilogue sync + static constexpr int SortWarpId = 3; // warp 3 (threads 96-127) +}; + +// ═══════════════════════════════════════════════════════════════════ +// RouterSmem — shared memory layout for the fused MoE kernel +// ═══════════════════════════════════════════════════════════════════ + +template +struct RouterSmem { + using Element = typename Traits::Element; + + static constexpr int smem_A_size = cosize_v; + static constexpr int smem_B_size = cosize_v; + + // Consumer A smem (even K-tiles) + alignas(128) Element smem_A_a[smem_A_size]; + alignas(128) Element smem_B_a[smem_B_size]; + // Consumer B smem (odd K-tiles) + alignas(128) Element smem_A_b[smem_A_size]; + alignas(128) Element smem_B_b[smem_B_size]; + + // Two TMA pipelines (one per consumer) + typename Traits::MainloopPipeline::SharedStorage pipe_a_storage; + typename Traits::MainloopPipeline::SharedStorage pipe_b_storage; + + alignas(16) float score_strip[2][Traits::TileM][Traits::TileN]; + + alignas(16) float running_max[Traits::TileM]; + alignas(16) float running_sum[Traits::TileM]; + alignas(16) float top_scores[Traits::TileM][Traits::MaxTopK]; + alignas(16) int top_indices[Traits::TileM][Traits::MaxTopK]; + + alignas(16) int strip_ready[2]; + alignas(16) int strip_consumed[2]; + + alignas(8) uint64_t mbar_debug; + + CUTE_DEVICE Element* A_a() { return &smem_A_a[0]; } + CUTE_DEVICE Element* B_a() { return &smem_B_a[0]; } + CUTE_DEVICE Element* A_b() { return &smem_A_b[0]; } + CUTE_DEVICE Element* B_b() { return &smem_B_b[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// SortBuffers — global memory pointers for the sort phase +// ═══════════════════════════════════════════════════════════════════ + +struct SortBuffers { + int* tile_expert_counts; // [num_m_tiles * E] — per-tile histograms, + // modified in-place for prefix sums + int* sorted_token_ids; // [total_slots] — final output: slot → token + int* token_expert_slots; // [T * K] — final output: (token, k) → slot + int* expert_offsets; // [E + 1] — aligned expert boundaries (symmetric) + int* expert_counts; // [E] — exact routed rows per expert (nullable) + int* tile_expert_ids; // [max_total_tiles] — expert ID per M-tile (nullable) + int* cta_done; // [num_blocks] — per-CTA completion flags (Blelloch tree) + int* cta_sums; // [num_blocks * E] — per-CTA expert totals for prefix sum +}; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/models_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/models_sm90.cuh new file mode 100644 index 00000000..d53d4ac0 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/models_sm90.cuh @@ -0,0 +1,189 @@ +#pragma once + +// #ifndef NDEBUG +// #define NDEBUG +// #endif + +#include +#include +#include +#include +#include +#include +#include + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// GMMA atom selection +// ═══════════════════════════════════════════════════════════════════ + +template +struct GmmaSelector; + +static constexpr auto GmmaK = GMMA::Major::K; +static constexpr auto GmmaMN = GMMA::Major::MN; + +template <> struct GmmaSelector { using Atom = SM90_64x8x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x16x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x32x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x64x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x128x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x192x16_F32BF16BF16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x256x16_F32BF16BF16_SS; }; + +template <> struct GmmaSelector { using Atom = SM90_64x8x16_F32F16F16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x16x16_F32F16F16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x32x16_F32F16F16_SS; }; +template <> struct GmmaSelector { using Atom = SM90_64x64x16_F32F16F16_SS; }; + +// ═══════════════════════════════════════════════════════════════════ +// GMMA atom selection — A=K-major, B=MN-major (SS) +// Used when B operand has MN-contiguous smem (e.g. mlp2_t_fused). +// ═══════════════════════════════════════════════════════════════════ + +template +struct GmmaSelectorKMN; + +template <> struct GmmaSelectorKMN { using Atom = SM90_64x8x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x16x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x32x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x64x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x128x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x256x16_F32BF16BF16_SS; }; + +template <> struct GmmaSelectorKMN { using Atom = SM90_64x8x16_F32F16F16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x16x16_F32F16F16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x32x16_F32F16F16_SS; }; +template <> struct GmmaSelectorKMN { using Atom = SM90_64x64x16_F32F16F16_SS; }; + +// Backward compat alias +template +struct GmmaSelectorMN : GmmaSelectorKMN {}; + +// ═══════════════════════════════════════════════════════════════════ +// GMMA atom selection — both MN-major (SS) +// Used when both operands have MN-contiguous smem (e.g. mlp3). +// ═══════════════════════════════════════════════════════════════════ + +template +struct GmmaSelectorMNMN; + +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x8x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x16x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x32x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x64x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x128x16_F32BF16BF16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x256x16_F32BF16BF16_SS; }; + +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x8x16_F32F16F16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x16x16_F32F16F16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x32x16_F32F16F16_SS; }; +template <> struct GmmaSelectorMNMN { using Atom = SM90_64x64x16_F32F16F16_SS; }; + +// ═══════════════════════════════════════════════════════════════════ +// FusedMoETraits +// ═══════════════════════════════════════════════════════════════════ + +template < + typename Element_, + int TileN_, + int TileK_ = 64, + int Stages_ = 2, + int MaxTopK_ = 8 +> +struct FusedMoETraits { + using Element = Element_; + using ElementAccum = float; + + static constexpr int TileM = 64; + static constexpr int TileN = TileN_; + static constexpr int TileK = TileK_; + static constexpr int Stages = Stages_; + static constexpr int MaxTopK = MaxTopK_; + + using GmmaAtom = typename GmmaSelector::Atom; + using TiledMma = TiledMMA, Layout>>; + + using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom; + + using SmemLayoutA_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{})); + using SmemLayoutB_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int>{})); + + using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{})); + using SmemLayoutB = decltype(tile_to_shape(SmemLayoutAtom{}, Shape, Int, Int>{})); + + using MainloopPipeline = cutlass::PipelineTmaAsync; + using PipelineState = cutlass::PipelineState; + + static constexpr int TmaTransactionBytesA = static_cast(size(SmemLayoutA_1{}) * sizeof(Element)); + static constexpr int TmaTransactionBytesB = static_cast(size(SmemLayoutB_1{}) * sizeof(Element)); + static constexpr int TmaTransactionBytes = TmaTransactionBytesA + TmaTransactionBytesB; + + static constexpr int WarpSize = 32; + static constexpr int WarpGroupSize = 128; + static constexpr int NumThreads = 384; + static constexpr int EpilogueThreads = 2 * WarpSize; // warps 1+2 (64 threads) + static constexpr int EpilogueBarrier = 3; // NamedBarrier ID for epilogue sync + static constexpr int SortWarpId = 3; // warp 3 (threads 96-127) +}; + +// ═══════════════════════════════════════════════════════════════════ +// RouterSmem — shared memory layout for the fused MoE kernel +// ═══════════════════════════════════════════════════════════════════ + +template +struct RouterSmem { + using Element = typename Traits::Element; + + static constexpr int smem_A_size = cosize_v; + static constexpr int smem_B_size = cosize_v; + + // Consumer A smem (even K-tiles) + alignas(128) Element smem_A_a[smem_A_size]; + alignas(128) Element smem_B_a[smem_B_size]; + // Consumer B smem (odd K-tiles) + alignas(128) Element smem_A_b[smem_A_size]; + alignas(128) Element smem_B_b[smem_B_size]; + + // Two TMA pipelines (one per consumer) + typename Traits::MainloopPipeline::SharedStorage pipe_a_storage; + typename Traits::MainloopPipeline::SharedStorage pipe_b_storage; + + alignas(16) float score_strip[2][Traits::TileM][Traits::TileN]; + + alignas(16) float running_max[Traits::TileM]; + alignas(16) float running_sum[Traits::TileM]; + alignas(16) float top_scores[Traits::TileM][Traits::MaxTopK]; + alignas(16) int top_indices[Traits::TileM][Traits::MaxTopK]; + + alignas(16) int strip_ready[2]; + alignas(16) int strip_consumed[2]; + + alignas(8) uint64_t mbar_debug; + + CUTE_DEVICE Element* A_a() { return &smem_A_a[0]; } + CUTE_DEVICE Element* B_a() { return &smem_B_a[0]; } + CUTE_DEVICE Element* A_b() { return &smem_A_b[0]; } + CUTE_DEVICE Element* B_b() { return &smem_B_b[0]; } +}; + +// ═══════════════════════════════════════════════════════════════════ +// SortBuffers — global memory pointers for the sort phase +// ═══════════════════════════════════════════════════════════════════ + +struct SortBuffers { + int* tile_expert_counts; // [num_m_tiles * E] — per-tile histograms, + // modified in-place for prefix sums + int* sorted_token_ids; // [total_slots] — final output: slot → token + int* token_expert_slots; // [T * K] — final output: (token, k) → slot + int* expert_offsets; // [E + 1] — aligned expert boundaries (symmetric) + int* expert_counts; // [E] — exact routed rows per expert (nullable) + int* tile_expert_ids; // [max_total_tiles] — expert ID per M-tile (nullable) + int* cta_done; // [num_blocks] — per-CTA completion flags (Blelloch tree) + int* cta_sums; // [num_blocks * E] — per-CTA expert totals for prefix sum +}; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe.cu b/liger_cute_kernels/csrc/core/src/moe/moe.cu index ddc23795..6626d880 100644 --- a/liger_cute_kernels/csrc/core/src/moe/moe.cu +++ b/liger_cute_kernels/csrc/core/src/moe/moe.cu @@ -1,1775 +1,9 @@ -// Fused MoE: sort+dispatch → MLP+combine — torch-free core (liger_cute). -// -// Ported from LigerCommKernels' moe.cu. The device kernels and the templated -// host launcher are upstream-verbatim (namespace liger, internal to the .so); -// only the boundary changed: the remaining public entry points here are flat -// extern "C" control/configuration APIs. Tensor-carrying MoE calls are exposed -// through TVM FFI and lower directly into the internal MoeFwdArgs/MoeBwdArgs -// launch bundles. torch::Tensor never appears here. -#define LIGER_CUTE_BUILDING 1 - -#include -#include // pulls in nvshmemi_device_state_d, required by the - // inline device helpers in . -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -#include "liger_cute/moe.h" -#include "liger_cute/check.h" -#include "liger_cute/detail/status.h" -#include "liger_cute/detail/comm_schedule.cuh" - -#include "moe.cuh" -#include "moe_symm_config.cuh" -#include "moe_launch.h" // MoeFwdArgs / MoeBwdArgs (shared with the tuner) -#include "moe_utils.cuh" -#include "moe_fwd_bwd_tune_configs.hpp" -#include "moe_dispatch_configs_sm90.cuh" -#include "moe_dispatch_configs_sm100.cuh" -#include "tma_copy_atom.cuh" -#include "radix_sort.cuh" -#include "dispatch.cuh" -#include "combine.cuh" -#include "buffer_pool.cuh" // shim: liger:: aliases of liger_cute::detail infra -#include "fused_memset.cuh" - -using namespace cute; - -namespace liger { - -// 0 = build every dispatch family (developer fallback). CMake sets this to 90 -// for Hopper builds and 100 for Blackwell builds so normal native builds do not -// instantiate kernels for hardware that cannot run in that binary. #ifndef LIGER_CUTE_DISPATCH_COMPUTE #define LIGER_CUTE_DISPATCH_COMPUTE 0 #endif -// init_comm_schedule lives in liger_cute::detail (comm_schedule.cuh); the host -// configure path below calls it through this alias. -using liger_cute::detail::init_comm_schedule; - -// ============================================================================ -// Global iteration counter for PeSync -// ============================================================================ -// -// PeSync's dispatch_done / mlp_done symmetric arrays must increase -// monotonically across calls (the wait-side uses NVSHMEM_CMP_GE). The -// iteration value is what each call SETs the signal to via -// nvshmemx_signal_op(SET). -// -// The counter lives in device memory (allocated by BufferPool as -// "moe_fwd_iter_counter") and is atomicAdded at the end of every -// kernel by CTA 0 thread 0 — see PeSync::advance_iteration_counter(). -// This is the form that survives CUDA graph capture: a host-side -// atomic would be read once on the host before the kernel launch and -// baked into the captured graph, so all replays would re-SET the same -// already-met value and the cross-PE wait would race; a GPU-resident -// counter incremented inside the captured kernel re-evaluates on -// every replay, giving each launch a fresh monotonic target. -// -// Buffer pool keys "moe_fwd_iter_counter" (fwd) and -// "moe_bwd_iter_counter" (bwd) name the two separate counters. - -// ============================================================================ -// MoeFusedConfig — bundles all tunable kernel parameters -// ============================================================================ -// -// TileM is fixed at 64 by the WGMMA hardware. -// NumThreads is fixed at 384 (10 MLP warps + 2 comm warps). -// All other parameters are exposed for autotuning. - -template < - typename Element_ = bfloat16_t, - int TileM_ = 128, - int TileN1_ = 128, - int TileK1_ = 64, - int Stages1_ = 4, - int EpiChunkN1_ = 64, - int TileN2_ = 256, - int TileK2_ = 64, - int Stages2_ = 4, - int EpiChunkN2_ = 64, - int ZBufferSlots_ = 4, - int CommNumStages_ = 2, - int NC_ = 4, - // GemmTileM_ decouples the GEMM compute tile (mlp1/mlp2 wgmma + TMA) from - // the COMM staging tile (TileM_). TileM_ is the comm tile — staging ring - // slots, per-tile semaphores, NC, sort/dispatch — and may be wider than the - // GEMM tile. GemmTileM_ ∈ {64,128}: when < TileM_, the GEMM waits on one comm - // semaphore once, then steps through kSubTiles = TileM_/GemmTileM_ sub-tiles - // of GemmTileM_ rows each (finer TMA + wgmma) within that one comm slot. - // Default GemmTileM_ = TileM_ → kSubTiles = 1 → bit-identical to the old - // single-TileM path. - int GemmTileM_ = TileM_, - int Compute_ = 90> -struct MoeFusedConfig { - using Element = Element_; - // GEMM tile (mlp1/mlp2): GemmTileM_ ∈ {64,128}, independent of the comm tile. - using Traits1 = Mlp1Traits< - Element, GemmTileM_, TileN1_, TileK1_, Stages1_, EpiChunkN1_>; - using Traits2 = Mlp2Traits; - - // No compile-time constraint between NC, NSplit and CommNumStages - // — the GEMM-side iterator is ticket-based (slot = T mod L, - // ticket = T / L) and synchronizes per-slot via raw atomics - // computed from the global tile index, so the only requirement - // (runtime) is that NC evenly divides gridDim.x · gridDim.y so MC - // is integer. CommNumStages is now a free buffering knob. - - // Compile-time smem budget. H100 dynamic-smem cap is 228 KiB; we leave - // ~4 KiB for the runtime-sized remote_offsets region appended by the - // launcher (see moe.cu: smem_mlp += num_pes * (experts_per_pe+1) * 4). - // Catches over-budget configs at link time instead of cudaErrorInvalidValue - // at launch. - static constexpr size_t kSmemBytes = sizeof(MoeSmem); - static constexpr size_t kSmemBudget = 224 * 1024; - static_assert(kSmemBytes <= kSmemBudget, - "MoE fused fwd kernel smem footprint exceeds 224 KiB budget — " - "reduce TileN/TileK/Stages for this config"); - - static constexpr int kNC = NC_; - static constexpr int kZBufferSlots = ZBufferSlots_; - static constexpr int kCommNumStages = CommNumStages_; - // kTileM is the COMM tile (staging ring slots, per-tile semaphores, NC, - // sort/dispatch granularity). - static constexpr int kTileM = TileM_; - static constexpr int kGemmTileM = GemmTileM_; // gemm tile ∈ {64,128} - static constexpr int kSubTiles = TileM_ / GemmTileM_; - static constexpr int kCompute = Compute_; - static constexpr int kNumThreads = 384; // fixed - - static_assert(TileM_ % GemmTileM_ == 0, - "CommTileM must be an integer multiple of GemmTileM"); - static_assert(TileM_ % NC_ == 0, - "CommTileM must be divisible by NC"); - static_assert(kSubTiles == 1 || kSubTiles == 2, - "CommTileM/GemmTileM must be 1 or 2"); - static_assert(GemmTileM_ == 64 || GemmTileM_ == 128, - "GemmTileM must be 64 or 128"); - - // Iterators operate on the COMM tile (kTileM), independent of the GEMM - // sub-tile size. - using LocalIter = LocalMlpTileIterator; - using RemoteIter = RemoteMlpTileIterator< - Element, kCommNumStages, kNC, kTileM, kSubTiles>; -}; - -// ============================================================================ -// MoeParams -- sort + dispatch + MLP + combine (single kernel) -// ============================================================================ - -template -struct MoeParams { - using Element = typename Config::Element; - - // Router outputs (read-only, from caller's tensors). - const int* expert_indices; // [T, K] - const Element* expert_weights; // [T, K] - - // Sort buffers. - SortBuffers sort; - - // Dispatch. - const Element* tokens; // [T, D] input (symmetric) - Element* dispatch_dst; // [total_slots, D] sorted output (symmetric) - - // Combine output. - Element* combine_output; // [T, D] final output - - // MLP + comm. - CommBuffers comm; - - // Dimensions. - int num_tokens; - int intermediate_dim; - int top_k; - int total_slots; - int num_m_tiles; // sort input tiles -}; - -// ============================================================================ -// MoeBuffers -- all device/symmetric allocations -// ============================================================================ - -template -struct MoeBuffers { - using Element = typename Config::Element; - - // MLP intermediate Z. - Element* z_buf; - int z_rows; - - // Sorted X (dispatch destination) -- symmetric, remote PEs get from here. - Element* x_sorted; - - // Output Y -- symmetric, local MLP writes here, remote PEs put here. - Element* y_buf; - - // Phase counters. - int* phase_counter; - - // Comm staging. - Element* src_staging; - Element* dst_staging; - int* src_ready; - int* src_consumed; - int* dst_ready; - int* dst_consumed; - int* tile_expert_ids; - - // All PEs' expert offsets -- symmetric. - int* all_expert_offsets; - int* all_expert_counts; - int* cta_counter; // [1] device memory — shared monotonic - - // Sort buffers. - int* tile_expert_counts; - int* sorted_token_ids; - int* token_expert_slots; - int* expert_offsets; // symmetric - int* expert_counts; // symmetric exact counts - int* cta_done; - int* cta_sums; - - // Per-tile expert IDs for MLP local iterator. - int* mlp_tile_expert_ids; -}; - -// MoeSymmConfig + get_symm_config() declaration live in the shared header so -// moe.cu and moe_bwd.cu can't disagree on the layout (see header). This TU -// owns the single instance. -MoeSymmConfig& get_symm_config() { - static MoeSymmConfig cfg; - return cfg; -} - -void moe_configure_symmetric( - int max_tokens, - int hidden_dim, - int max_num_experts, - int max_top_k, - int num_pes, - int num_hosts, - int gpus_per_host) { - static constexpr int kTileM = kMaxMoeCommTileM; - - // num_hosts × gpus_per_host must equal the team size so the comm - // schedule (g_dest_table / g_rank_table) addresses exactly the PEs - // in the team. Mismatch would silently map ranks to non-existent - // destinations (or skip some). - LIGER_CHECK(num_hosts * gpus_per_host == num_pes, - "moe_configure_symmetric: num_hosts (", num_hosts, ") * gpus_per_host (", - gpus_per_host, ") = ", num_hosts * gpus_per_host, - " must equal num_pes (", num_pes, ")"); - - // Populate g_dest_table / g_rank_table in device constant memory. - // Idempotent — safe to re-call with the same (N, M). - init_comm_schedule(num_hosts, gpus_per_host); - - auto& cfg = get_symm_config(); - cfg.max_total_slots = max_tokens * max_top_k + max_num_experts * kTileM; - cfg.max_num_experts = max_num_experts; - cfg.hidden_dim = hidden_dim; - cfg.num_pes = num_pes; - cfg.experts_per_pe = max_num_experts / num_pes; - cfg.max_top_k = max_top_k; - cfg.initialized = true; - - // Preset sizes for the symmetric stack used to hold X / Y / routing metadata - // across fwd → bwd. Element type is bf16 (= 2 bytes); kept in sync - // with allocate_moe_buffers below. - constexpr std::size_t kElemBytes = 2; // sizeof(__nv_bfloat16) - std::size_t symm_slot_bytes = - static_cast(cfg.max_total_slots) * cfg.hidden_dim * kElemBytes; - std::size_t all_off_bytes = - static_cast(cfg.num_pes) * (cfg.max_num_experts + 1) * sizeof(int); - std::size_t all_count_bytes = - static_cast(cfg.num_pes) * cfg.max_num_experts * sizeof(int); - - auto& stack = global_symmetric_stack(); - stack.set_size(kSymmKeyXSorted, symm_slot_bytes); - stack.set_size(kSymmKeyYBuf, symm_slot_bytes); - stack.set_size(kSymmKeyAllExpertOffsets, all_off_bytes); - stack.set_size(kSymmKeyAllExpertCounts, all_count_bytes); -} - -// Releases the forward metadata entries pushed onto the symmetric stack by the -// most recent fwd call. Must be called collectively -// on every PE in LIFO order vs. the fwd that produced them. After this, -// the buffers can be reused by a subsequent fwd. The torch::Tensor views -// returned by fwd remain valid pointers but will alias whatever lives in -// those slots after the next put. -void moe_pop_fwd() { - auto& stack = global_symmetric_stack(); - // LIFO order: bottom pushes first. Reverse the put order in - // allocate_moe_buffers (x_sorted, y_buf, all_expert_offsets, - // all_expert_counts). - stack.pop(kSymmKeyAllExpertCounts); - stack.pop(kSymmKeyAllExpertOffsets); - stack.pop(kSymmKeyYBuf); - stack.pop(kSymmKeyXSorted); -} - -template -static MoeBuffers allocate_moe_buffers( - BufferPool& pool, - cudaStream_t stream, - int grid_x, - int num_blocks, - int hidden_dim, - int intermediate_dim, - int num_tokens, - int top_k, - int num_experts, - int experts_per_pe, - int num_pes, - int total_slots, - int num_m_tiles, - // Caller-owned torch storage for the two sort outputs that bwd - // needs to consume. Passed in so we don't lease buffer-pool - // slots that the next fwd would overwrite. - int* token_expert_slots_ptr, - int* mlp_tile_expert_ids_ptr) { - - using Element = typename Config::Element; - static constexpr int kTileM = Config::kTileM; - static constexpr int kZBufferSlots = Config::kZBufferSlots; - - // Staging is a flat ring of L = MC · kCommNumStages slots, where - // MC = num_blocks / NC (num_blocks == launched gridDim.x). This MUST equal - // the device-side comm producer's stride (mlp_comms.cuh: MC = n_comm/NC, - // n_comm = (gridDim.x/NC)·NC) AND the MLP consumer's ring stride - // (RemoteMlpTileIterator::init: mc = gridDim.x/NC), so every tile's - // slot = g % L agrees across producer / consumer / allocation. Sizing this - // on grid_x·NSplit (= n_gemm) instead under-sizes the signal arrays whenever - // gridDim.x/NC > n_gemm/NC — i.e. num_sms % NSplit != 0, e.g. NSplit=8 on - // 132 SMs (33 vs 32) — and the gridDim.x-based ring then indexes past the - // end of src_ready / tile_expert_ids (OOB). num_blocks/NC == gridDim.x/NC - // == n_gemm/NC in the divisible (legacy) case, so this is a no-op there. - // Comm strides MC through NumStages per xc; MLP (ticket iterator) wraps mod L. - // Raw SM count — used as the device-independent upper bound on the comm - // staging row count (mc <= sm_count) so the symmetric staging size is a - // worst-case bound across configs (see staging sizing below). - int sm_count = 0; - { - int dev = 0; - cudaGetDevice(&dev); - cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev); - } - - MoeBuffers b; - - // Symmetric allocations use max sizes from MoeSymmConfig - // so all PEs allocate the same sizes regardless of per-call config. - auto& scfg = get_symm_config(); - LIGER_CHECK(scfg.initialized, "Call moe_configure_symmetric before moe_fused_fwd_bf16"); - - // Persistent symmetric tensors and routing metadata live in the symmetric - // stack so they can survive across fwd calls into bwd. - auto& sym_stack = global_symmetric_stack(); - - // Z buffer (device). - // size_t cast on b.z_rows: at NSplit=2 ZBuf=16 I=16384, the product - // z_rows · intermediate_dim · 2 = ~4.4 GB which overflows int32 if - // the multiply happens in int (b.z_rows and intermediate_dim are - // both int). Cast first so the entire product is in size_t. - int max_runtime_grid_x = (num_blocks + 1) / 2; // min candidate NS = 2 - int z_m_tiles = max_runtime_grid_x * kZBufferSlots; - b.z_rows = z_m_tiles * kTileM; - size_t z_bytes = static_cast(b.z_rows) * intermediate_dim * sizeof(Element); - b.z_buf = reinterpret_cast(pool.get_device("moe_z_buf", z_bytes)); - - // Sorted X -- symmetric, dispatch writes here, remote PEs get from here. - b.x_sorted = reinterpret_cast(sym_stack.put(kSymmKeyXSorted)); - - // Output Y -- symmetric, local MLP writes here, remote PEs put here. - b.y_buf = reinterpret_cast(sym_stack.put(kSymmKeyYBuf)); - - // Phase counters. - size_t counter_bytes = max_runtime_grid_x * sizeof(int); - b.phase_counter = reinterpret_cast( - pool.get_device("moe_phase_counter", counter_bytes)); - if (cudaError_t e = cudaMemsetAsync(b.phase_counter, 0, counter_bytes, stream); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(phase_counter) failed: ", cudaGetErrorString(e)); - - // Flat staging list — L = grid_x * kCommNumStages slots, each - // holding TileM rows × hidden_dim columns. - // Size by the CONFIGURED MAX on every axis (scfg.hidden_dim, max_comm_stages, - // max_tile_m), NOT the per-call/per-config values: the symmetric pool caches - // by name and aborts on grow, so a small first config must not lock in a - // staging buffer a later larger config can't grow into. This covers all three - // independently-varying axes that the tuner sweeps: - // - hidden_dim: D=512 then D=8192 - // - CommNumStages: CS=2 then CS=8 (NS=16 config) - // - TileM: 64 then 128 - // The kernel indexes slots by the per-call (comm_l, tile_elems <= max), so an - // over-sized buffer is always safe (symmetric memory may be larger than - // needed, never smaller). - size_t staging_rows_ub = - (size_t)sm_count * scfg.max_comm_stages * scfg.max_tile_m; - size_t src_bytes = staging_rows_ub * scfg.hidden_dim * sizeof(Element); - size_t dst_bytes = staging_rows_ub * scfg.hidden_dim * sizeof(Element); - b.src_staging = reinterpret_cast(pool.get_symmetric("moe_src_staging", src_bytes)); - b.dst_staging = reinterpret_cast(pool.get_symmetric("moe_dst_staging", dst_bytes)); - - // Device buffers are cached by key across tuning/config changes, so size and - // clear them at a configuration-independent upper bound. Since - // kSubTiles <= NC, (num_blocks / NC) * stages * kSubTiles is bounded by - // sm_count * max_comm_stages. - size_t signal_l_ub = - static_cast(sm_count) * scfg.max_comm_stages; - size_t sig_bytes = signal_l_ub * sizeof(int); - b.src_ready = static_cast(pool.get_device("moe_src_ready", sig_bytes)); - b.src_consumed = static_cast(pool.get_device("moe_src_consumed", sig_bytes)); - b.dst_ready = static_cast(pool.get_device("moe_dst_ready", sig_bytes)); - b.dst_consumed = static_cast(pool.get_device("moe_dst_consumed", sig_bytes)); - if (cudaError_t e = cudaMemsetAsync(b.src_ready, 0, sig_bytes, stream); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(src_ready) failed: ", cudaGetErrorString(e)); - if (cudaError_t e = cudaMemsetAsync(b.src_consumed, 0, sig_bytes, stream); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(src_consumed) failed: ", cudaGetErrorString(e)); - if (cudaError_t e = cudaMemsetAsync(b.dst_ready, 0, sig_bytes, stream); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(dst_ready) failed: ", cudaGetErrorString(e)); - if (cudaError_t e = cudaMemsetAsync(b.dst_consumed, 0, sig_bytes, stream); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(dst_consumed) failed: ", cudaGetErrorString(e)); - - // Per-parent-slot expert IDs use the same conservative upper bound. - size_t eid_bytes = signal_l_ub * sizeof(int); - b.tile_expert_ids = static_cast(pool.get_device("moe_tile_expert_ids", eid_bytes)); - - // All PEs' expert offsets (symmetric). Use max experts for stable sizing. - b.all_expert_offsets = static_cast(sym_stack.put(kSymmKeyAllExpertOffsets)); - b.all_expert_counts = static_cast(sym_stack.put(kSymmKeyAllExpertCounts)); - - // CTA counter for local all-CTA barriers around NVSHMEM collectives. - b.cta_counter = static_cast( - pool.get_device("moe_cta_counter", sizeof(int))); - - // Sort buffers -- use scfg max sizes to avoid reallocation across configs. - int max_sorted_slots = scfg.max_total_slots; - int max_m_tiles = scfg.max_total_slots / kTileM; - int max_tile_counts = max_m_tiles * scfg.max_num_experts; - int max_token_slots_ub = scfg.max_total_slots; - - (void)max_token_slots_ub; - b.tile_expert_counts = static_cast( - pool.get_device("moe_sort_tile_counts", max_tile_counts * sizeof(int))); - b.sorted_token_ids = static_cast( - pool.get_device("moe_sort_sorted_ids", max_sorted_slots * sizeof(int))); - // token_expert_slots: caller-owned torch tensor (passed in). Used by - // the bwd path after this call returns — must not alias a pool slot. - b.token_expert_slots = token_expert_slots_ptr; - b.expert_offsets = static_cast( - pool.get_symmetric("moe_sort_offsets", (scfg.max_num_experts + 1) * sizeof(int))); - b.expert_counts = static_cast( - pool.get_symmetric("moe_sort_counts", scfg.max_num_experts * sizeof(int))); - b.cta_done = static_cast( - pool.get_device("moe_sort_cta_done", num_blocks * sizeof(int))); - b.cta_sums = static_cast( - pool.get_device("moe_sort_cta_sums", num_blocks * scfg.max_num_experts * sizeof(int))); - - // Per-tile expert IDs for MLP — caller-owned torch tensor. - b.mlp_tile_expert_ids = mlp_tile_expert_ids_ptr; - - if (cudaError_t e = cudaMemsetAsync(b.cta_done, 0, num_blocks * sizeof(int), stream); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(cta_done) failed: ", cudaGetErrorString(e)); - if (cudaError_t e = cudaMemsetAsync(b.sorted_token_ids, 0xff, - max_sorted_slots * sizeof(int), stream); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(sorted_token_ids) failed: ", cudaGetErrorString(e)); - - return b; -} - -// ============================================================================ -// GET-via-TMA descriptor builder (raw CUtensorMap, FWD-only) -// ============================================================================ -// -// Box = [box_rows, kGetKChunk] over a row-major [n_rows, hidden_dim] bf16 -// tensor. SWIZZLE_NONE keeps the global bytes row-major, so the MLP's -// swizzled reader (tma_load_x_remote) sees correct data. Returns false on -// encode failure → caller disables the TMA-get path (falls back to getmem). -static inline bool make_get_tma_map(CUtensorMap& m, void* base, - int hidden_dim, int n_rows, int box_rows) { - uint64_t gdim[2] = { (uint64_t)hidden_dim, (uint64_t)n_rows }; - uint64_t gstride[1] = { (uint64_t)hidden_dim * sizeof(bfloat16_t) }; - uint32_t bdim[2] = { (uint32_t)liger::kGetKChunk, (uint32_t)box_rows }; - uint32_t estride[2] = { 1, 1 }; - CUresult r = cuTensorMapEncodeTiled( - &m, CU_TENSOR_MAP_DATA_TYPE_UINT16, 2, base, gdim, gstride, bdim, estride, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, - CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - return r == CUDA_SUCCESS; -} - -// ============================================================================ -// Fused kernel: sort → dispatch → MLP + combine -// ============================================================================ - -template -__global__ void __launch_bounds__(Config::kNumThreads) -moe_fused_kernel( - __grid_constant__ TmaLoadX const tma_load_x, - __grid_constant__ TmaLoadW1 const tma_load_b, - __grid_constant__ TmaLoadW1 const tma_load_c, - __grid_constant__ TmaStoreZ const tma_store_z, - __grid_constant__ TmaLoadZ const tma_load_z, - __grid_constant__ TmaLoadW2 const tma_load_a, - __grid_constant__ TmaStoreY const tma_store_y, - __grid_constant__ TmaLoadX const tma_load_x_remote, - __grid_constant__ TmaStoreY const tma_store_y_remote, - __grid_constant__ TmaStoreYCombine const tma_store_y_combine, - __grid_constant__ MlpDims const mlp_dims, - __grid_constant__ MoeParams const p, - __grid_constant__ int const static_nsplit, - __grid_constant__ GetTmaDescs const get_descs, - __grid_constant__ PeerYStoreDescs const peer_y) { - - using Element = typename Config::Element; - using Traits1 = typename Config::Traits1; - using Traits2 = typename Config::Traits2; - using LocalIter = typename Config::LocalIter; - using RemoteIter = typename Config::RemoteIter; - constexpr int kTileM = Config::kTileM; - constexpr int kNumThreads = Config::kNumThreads; - constexpr int kZBufferSlots = Config::kZBufferSlots; - constexpr int kCommNumStages = Config::kCommNumStages; - constexpr int kNC = Config::kNC; - - // -- Sort ---------------------------------------------------------------- - { - extern __shared__ int sort_smem[]; - sort_tokens( - sort_smem, p.expert_indices, p.sort, - p.num_tokens, p.comm.num_experts, p.top_k, p.num_m_tiles); - } - __threadfence_system(); // device-scope (NVLink-only build; IB needs system scope) - __syncthreads(); - - // Target-based cross-CTA barrier — single instance per kernel call, - // reused by every PeSync::barrier / signal_dispatch / signal_mlp - // invocation below. Constructing a fresh barrier per call resets - // the per-CTA target to 0 while the shared device counter keeps - // growing → fast CTAs spin past the next-call target instantly. - SyncThreadsCtaCounterBarrier pe_barrier(p.comm.pe_sync.cta_counter, - (int)gridDim.x * (int)gridDim.y); - - // Cross-CTA barrier: sort writes sorted_token_ids across CTAs; - // dispatch and combine read them from other CTAs. - p.comm.pe_sync.barrier(pe_barrier); - - // -- Broadcast expert offsets to all PEs ---------------------------------- - broadcast_expert_offsets(p.comm); - __syncthreads(); - - // -- Dispatch ------------------------------------------------------------- - dispatch_tokens( - p.tokens, p.dispatch_dst, - p.sort.sorted_token_ids, - p.comm.hidden_dim, p.total_slots); - __syncthreads(); - - // H2 (IB coherence): dispatch_tokens wrote dispatch_dst (= x_sorted) with - // plain GPU stores; remote PEs RDMA-read it in do_get after wait_dispatch. - // Flush those stores to HBM so the remote NIC reads fresh data — the quiet - // below only orders this PE's OUTGOING nvshmem ops, not these local - // stores, and __syncthreads is device-scope. The offset-broadcast source - // is already covered by radix_sort's __threadfence_system. NVLink P2P - // reads are L2-coherent, so single-host never exercised this. - __threadfence_system(); // device-scope (NVLink-only build; IB needs system scope) - - p.comm.pe_sync.barrier_all(pe_barrier); - - // -- MLP (local + remote) ------------------------------------------------ - { - extern __shared__ char raw_smem[]; - auto& smem = *reinterpret_cast*>(raw_smem); - - size_t offsets_region = sizeof(MoeSmem); - offsets_region = (offsets_region + sizeof(int) - 1) & ~(sizeof(int) - 1); - size_t offsets_count = - (size_t)p.comm.num_pes() * (p.comm.experts_per_pe + 1); - if (threadIdx.x == 0) { - smem.comm.remote_offsets = - reinterpret_cast(raw_smem + offsets_region); - smem.comm.remote_counts = - smem.comm.remote_offsets + offsets_count; - } - __syncthreads(); - - // GET-via-TMA bounce region: placed after remote_offsets, 128-aligned - // (TMA dst smem must be 128B-aligned). MUST match the host smem_size - // computation (same sizeof(MoeSmem), num_pes, experts_per_pe). Only - // touched by the comm warps' do_get when get_descs.enabled. - char* get_bounce = nullptr; - if (get_descs.enabled) { - size_t boff = offsets_region - + (offsets_count - + (size_t)p.comm.num_pes() * p.comm.experts_per_pe) - * sizeof(int); - boff = (boff + 127) & ~((size_t)127); - get_bounce = raw_smem + boff; - } - - //if (threadIdx.x == 0) - // smem.comm.total_tiles = -1; - - // Flat-grid logical coordinates for this CTA. GEMM and comm gate to - // the first n_gemm = floor_NS CTAs (the comm↔MLP staging ring needs - // grid_x = n_gemm/static_nsplit and MC = n_gemm/NC consistent); the gap CTAs - // [n_gemm, num_blocks) fall through to the shared pe_barrier only. - int flat_id = blockIdx.x; - int n_gemm = p.comm.n_gemm; - int grid_x = n_gemm / static_nsplit; - int col = flat_id / static_nsplit; - int runtime_nsplit = static_nsplit; - - // gemm_active removed: comm warps self-gate on their precomputed tile - // count (smem.comm.total_tiles) so a CTA can do comm-only work without - // participating in GEMM (the point of the comm/GEMM grid relaxation). - // Run the prologue unconditionally so every launched CTA has total_tiles - // initialized; the MLP/GEMM path still gates on GEMM-grid membership - // (flat_id < n_gemm) below. - nvshmem_comm_prologue(smem.comm, p.comm, static_nsplit); - __syncthreads(); - - if (threadIdx.x == 0) { - int total_m_tiles = smem.comm.global_total; - int mlp_tiles = total_m_tiles; - if constexpr (Config::kSubTiles > 1) { - mlp_tiles = 0; - for (int pe = 0; pe < p.comm.num_pes(); ++pe) { - for (int e = 0; e < p.comm.experts_per_pe; ++e) { - int count = smem.comm.remote_counts[ - pe * p.comm.experts_per_pe + e]; - mlp_tiles += - (count + Config::kGemmTileM - 1) / - Config::kGemmTileM; - } - } - } - runtime_nsplit = select_runtime_nsplit_from_tiles( - mlp_tiles, - mlp_dims.num_n_tiles_1, - mlp_dims.num_n_tiles_2, - mlp_dims.num_k_tiles_1, - mlp_dims.num_k_tiles_2, - (int)gridDim.x, - static_nsplit); - n_gemm = ((int)gridDim.x / runtime_nsplit) * runtime_nsplit; - grid_x = n_gemm / runtime_nsplit; - col = flat_id / runtime_nsplit; - - int mlp_count = 0; - if (flat_id < n_gemm && col < mlp_tiles) - mlp_count = (mlp_tiles - 1 - col) / grid_x + 1; - smem.comm.per_cta_tiles = mlp_count; - smem.comm.runtime_nsplit = runtime_nsplit; - smem.comm.runtime_n_gemm = n_gemm; - smem.comm.runtime_grid_x = grid_x; - } - __syncthreads(); - runtime_nsplit = smem.comm.runtime_nsplit; - n_gemm = smem.comm.runtime_n_gemm; - grid_x = smem.comm.runtime_grid_x; - col = flat_id / runtime_nsplit; - - // Diverge: comm warps handle NVSHMEM, MLP warps handle compute. - // Separate scopes prevent CommBuffers from inflating MLP register budget. - // Fwd uses warps 1-3 (2 getters + 1 putter). - int warp_id = threadIdx.x / 32; - bool is_comm_warp = (warp_id == kCommGetWarp || warp_id == kCommPutWarp); - - if (is_comm_warp) { - // Comm warps: dedicated get/put — full CommBuffers access. Self-gate - // on the precomputed comm tile count — no GEMM coupling, so a - // comm-only CTA (no local GEMM work) still drives its comm tiles. - if (smem.comm.total_tiles > 0) { - nvshmem_comm_main< - Element, kTileM, kCommNumStages, kNC, Config::kSubTiles>( - smem.comm, p.comm, &get_descs, get_bounce); - } - } else if (flat_id < n_gemm) { - // MLP warps: no CommBuffers in scope. Gap CTAs (flat_id >= n_gemm) - // skip this entirely and fall through to the pe_barrier below. - int local_e_start = p.comm.my_pe() * p.comm.experts_per_pe; - int local_e_end = local_e_start + p.comm.experts_per_pe; - int local_offset = p.sort.expert_offsets[local_e_start]; - int local_end = p.sort.expert_offsets[local_e_end]; - int local_m_tiles = (local_end - local_offset) / kTileM; - int local_tile_off = local_offset / kTileM; - - LocalIter local_iter; - local_iter.init( - p.dispatch_dst + local_offset * p.comm.hidden_dim, - p.sort.tile_expert_ids + local_tile_off, - local_m_tiles, p.comm.hidden_dim, - /*col_stride=*/grid_x, /*start_col=*/col, local_tile_off); - - moe_fused_fwd( - smem, local_iter, - p.comm.src_ready, p.comm.src_consumed, - p.comm.dst_ready, p.comm.dst_consumed, - p.comm.experts_per_pe, p.comm.num_pes(), p.comm.my_pe(), - get_descs.gpus_per_node, - tma_load_x, tma_load_b, tma_load_c, tma_store_z, - tma_load_z, tma_load_a, tma_store_y, - tma_load_x_remote, tma_store_y_remote, - mlp_dims, - local_m_tiles * kTileM, local_m_tiles, - col, grid_x, /*split=*/flat_id % runtime_nsplit, runtime_nsplit, - local_e_start, - /*y_peer_desc=*/peer_y.enabled ? peer_y.desc : nullptr); - } - - __threadfence(); - __syncthreads(); - p.comm.pe_sync.barrier_all(pe_barrier); - } - - // -- Combine -------------------------------------------------------------- - // Reuse the kernel's dynamic SMEM region as CombineSmem. wait_mlp above - // is a cross-CTA + cross-PE barrier that ends in its own __syncthreads, so - // the CTA is already aligned and the MLP/Comm SMEM contents are dead by - // this point — no extra __syncthreads needed before the reinterpret. - { - constexpr int kCombineNumWGs = - kNumThreads / (kCombineWarpsPerWG * kCombineWarpSize); - extern __shared__ char raw_smem[]; - auto& combine_smem = - *reinterpret_cast*>(raw_smem); - combine_tokens_tma( - reinterpret_cast(p.comm.local_output), - p.expert_weights, - p.sort.token_expert_slots, - tma_store_y_combine, - combine_smem, - p.comm.hidden_dim, - p.num_tokens, - p.top_k); - } - -} - -// ============================================================================ -// Host wrapper -// ============================================================================ - -// MoeFwdArgs (the POD launch bundle) is defined in moe_launch.h, shared with the -// tuner. The runtime dispatch table holds a single function-pointer type -// (MoeFwdFn) while every tuned config is a distinct template instantiation. - -template < - typename Element_ = bfloat16_t, - int TileN1 = 128, - int TileK1 = 64, - int Stages1 = 4, - int TileN2 = 256, - int TileK2 = 64, - int Stages2 = 4, - int ZBufferSlots = 4, - int CommNumStages = 1, - int EpiChunkN1 = 64, // mlp1 epilogue chunk; must divide TileN1 - int EpiChunkN2 = 64, // mlp2 epilogue chunk; must divide TileN2 - int TileM = 128, // COMM tile (staging/semaphores/NC); 128 - int GemmTileM = TileM, // GEMM tile ∈ {64,128}; default == comm (no sub-tiling) - int Compute = 90> -void -moe_fused_fwd_bf16(const MoeFwdArgs& a, int static_nsplit) { - - // Unpack the POD argument bundle into the names the upstream body uses. - const int num_experts = a.num_experts; - const int top_k = a.top_k; - const nvshmem_team_t team = a.team; - cudaStream_t stream = a.stream; - - // NC = number of cooperating comm CTAs per (xc, yc) tile. With the - // ticket-based GEMM-side iterator, runtime NS and NC no longer need to - // share a divisibility relationship — the only runtime requirement - // is NC | (gridDim.x · gridDim.y) so MC is integer. - // - // Wrapper param order (10 "core" args + trailing EpiChunkN + TileM): - // Element, TN1, TK1, S1, TN2, TK2, S2, ZBuf, CS, - // EpiChunkN1, EpiChunkN2, TileM - // MoeFusedConfig interleaves TileM/EpiChunkN inside each phase: - // Element, TileM, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, - // ZBuf, CS, NC - // The remapping happens here. TileM is shared by MLP1 and MLP2. - // - // NC choice (verified via bench_mlp_comms BW sweep): - // - TileM=128: NC=4 gives ~2× the comm bandwidth of NC=2 because - // 4 cooperating comm CTAs split each tile, deepening the comm - // pipeline (more in-flight slots = grid.x/NC). - // - TileM=64: NC=2 already matches NC=4 throughput because tiles - // are half-size, so twice as many flow through the same pipeline. - // Both architectures use CommTileM=128 with NC=4. Keeping CommTileM uniform per - // architecture preserves - // all-to-all dispatch padding (x_sorted / expert_offsets, padded to the - // comm tile) across PEs. The 256/8 ratio keeps each SM100 comm CTA's row - // band and GET-bounce footprint unchanged. - using CommGeometry = MoeCommGeometry; - using Config = MoeFusedConfig< - Element_, /*CommTileM=*/CommGeometry::TileM, - TileN1, TileK1, Stages1, EpiChunkN1, - TileN2, TileK2, Stages2, EpiChunkN2, - ZBufferSlots, - CommNumStages * CommGeometry::RingStageScale - + CommGeometry::RingStageExtra, - /*NC=*/CommGeometry::NC, - /*GemmTileM=*/GemmTileM, /*Compute_=*/Compute>; - - using Element = typename Config::Element; - using Traits1 = typename Config::Traits1; - using Traits2 = typename Config::Traits2; - using TmaLoadAtom = TmaLoadAtomForCompute; - using TmaStoreAtom = TmaStoreAtomForCompute; - constexpr int kTileM = Config::kTileM; - constexpr int kNumThreads = Config::kNumThreads; - - // Iteration counter lives on the device and is incremented at the - // end of the kernel — see PeSync::advance_iteration_counter() and - // the comment block at the top of this file. - int num_tokens = a.num_tokens; - int hidden_dim = a.hidden_dim; - int experts_per_pe = a.experts_per_pe; // local experts (from weight matrices) - int intermediate_dim = a.intermediate_dim; - - // Upper bound for total_slots: T * K + E * TileM (TileM padding per expert). - // Actual value is computed by sort (expert_offsets[E]) inside the kernel. - int max_total_slots = num_tokens * top_k + num_experts * kTileM; - int num_m_tiles_in = (num_tokens + kTileM - 1) / kTileM; // sort input tiles - - if constexpr (Compute == 100) { - LIGER_CHECK(hidden_dim % Traits1::TileK == 0, - "hidden_dim (", hidden_dim, ") must be divisible by MLP1 TileK (", Traits1::TileK, ")"); - LIGER_CHECK(intermediate_dim % Traits1::TileN == 0, - "intermediate_dim (", intermediate_dim, ") must be divisible by MLP1 TileN (", Traits1::TileN, ")"); - LIGER_CHECK(intermediate_dim % Traits2::TileK == 0, - "intermediate_dim (", intermediate_dim, ") must be divisible by MLP2 TileK (", Traits2::TileK, ")"); - LIGER_CHECK(hidden_dim % Traits2::TileN == 0, - "hidden_dim (", hidden_dim, ") must be divisible by MLP2 TileN (", Traits2::TileN, ")"); - } - LIGER_CHECK(hidden_dim % 8 == 0, - "hidden_dim (", hidden_dim, ") must be divisible by 8 (int4 vectorization in dispatch/combine)"); - LIGER_CHECK(intermediate_dim % 8 == 0, - "intermediate_dim (", intermediate_dim, ") must be divisible by 8 (int4 vectorization)"); - LIGER_CHECK(top_k <= kCombineMaxTopK, - "top_k (", top_k, ") exceeds kCombineMaxTopK (", kCombineMaxTopK, - "). Bump the constant in combine.cuh."); - - int num_n_tiles_1 = (Compute == 90) - ? (intermediate_dim + Traits1::TileN - 1) / Traits1::TileN - : intermediate_dim / Traits1::TileN; - int num_n_tiles_2 = (Compute == 90) - ? (hidden_dim + Traits2::TileN - 1) / Traits2::TileN - : hidden_dim / Traits2::TileN; - - // Relaxed: runtime NS > num_n_tiles is allowed; surplus split CTAs run - // empty inner loops but still drive their TMA-load/store mbarriers so the - // pipeline stays balanced across the column. - - LIGER_CHECK(a.weight_expert_stride >= - static_cast(intermediate_dim) * hidden_dim, - "weight expert stride (", a.weight_expert_stride, - ") must cover at least one [I,D] expert matrix"); - LIGER_CHECK(a.weight_expert_stride % hidden_dim == 0, - "weight expert stride (", a.weight_expert_stride, - ") must contain complete hidden-dimension rows (D=", hidden_dim, ")"); - int weight_expert_stride_rows = - static_cast(a.weight_expert_stride / hidden_dim); - LIGER_CHECK(weight_expert_stride_rows % Traits1::TileN == 0, - "weight expert stride in rows (", weight_expert_stride_rows, - ") must be divisible by MLP1 TileN (", Traits1::TileN, ")"); - int total_n_rows_1 = - (experts_per_pe - 1) * weight_expert_stride_rows + intermediate_dim; - int total_n_rows_2 = experts_per_pe * hidden_dim; - - int num_pes = nvshmem_team_n_pes(team); - int my_pe = nvshmem_team_my_pe(team); - - // Grid — flat 1-D launch. static_nsplit only seeds n_gemm/grid_x before - // the kernel selects the runtime split from the actual routed tile count. - // - // GEMM and comm gate to the first n_gemm = floor_NS CTAs: the comm↔MLP - // staging ring couples producer and consumer (it needs grid_x = n_gemm/static_nsplit - // and MC = n_gemm/NC to be consistent), so they share one count. The - // sort/dispatch/combine grid-stride phases and the cross-CTA pe_barrier use - // ALL num_blocks launched CTAs. The runtime selector may choose a different - // NS inside the kernel and rewrites n_gemm/grid_x consistently for MLP. - int num_sms; - cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, a.device); - int floor_ns = num_sms - num_sms % static_nsplit; - int floor_nc = num_sms - num_sms % Config::kNC; - int n_gemm = floor_ns; // GEMM/comm-active CTAs - int grid_x = n_gemm / static_nsplit; // initial logical GEMM columns - int num_blocks = (floor_ns > floor_nc) ? floor_ns : floor_nc; // launched CTAs - dim3 grid(num_blocks); - - // stream comes from a.stream (the binding passes the current torch stream). - auto& pool = global_buffer_pool(); - - // The two sort outputs the bwd path consumes (token_expert_slots - // [max_total_slots], tile_expert_ids [max_total_slots/kTileM]) are - // caller-owned: the binding pre-allocated them at exactly these sizes and - // passed their device pointers in a. They must NOT alias pool slots that the - // next forward would overwrite. - int* token_expert_slots_ptr = a.token_expert_slots; - int* tile_expert_ids_ptr = a.tile_expert_ids; - - // Allocate all buffers. - auto buf = allocate_moe_buffers(pool, stream, grid_x, num_blocks, - hidden_dim, intermediate_dim, num_tokens, top_k, num_experts, - experts_per_pe, num_pes, - max_total_slots, num_m_tiles_in, - token_expert_slots_ptr, - tile_expert_ids_ptr); - if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: CUDA error before TMA descriptor setup: ", cudaGetErrorString(e)); - - // Final combine output [T, D] — caller-owned (binding-allocated) Y buffer. - Element* output_ptr = reinterpret_cast(a.Y); - Element* y_buf_ptr = buf.y_buf; - - auto& scfg = get_symm_config(); - - // ---- Build MoeParams ------------------------------------------------------- - MoeParams p; - p.expert_indices = a.expert_indices; - p.expert_weights = reinterpret_cast(a.expert_weights); - - p.sort.tile_expert_counts = buf.tile_expert_counts; - p.sort.sorted_token_ids = buf.sorted_token_ids; - p.sort.token_expert_slots = buf.token_expert_slots; - p.sort.expert_offsets = buf.expert_offsets; - p.sort.expert_counts = buf.expert_counts; - p.sort.tile_expert_ids = buf.mlp_tile_expert_ids; - p.sort.cta_done = buf.cta_done; - p.sort.cta_sums = buf.cta_sums; - - p.tokens = reinterpret_cast(a.X); - p.dispatch_dst = buf.x_sorted; - p.combine_output = output_ptr; - - p.comm.src_staging = buf.src_staging; - p.comm.dst_staging = buf.dst_staging; - p.comm.src_ready = buf.src_ready; - p.comm.src_consumed = buf.src_consumed; - p.comm.dst_ready = buf.dst_ready; - p.comm.dst_consumed = buf.dst_consumed; - p.comm.expert_offsets = buf.expert_offsets; - p.comm.expert_counts = buf.expert_counts; - p.comm.sorted_token_ids = buf.sorted_token_ids; - p.comm.local_tokens = buf.x_sorted; - p.comm.local_output = y_buf_ptr; - p.comm.tile_expert_ids = buf.tile_expert_ids; - p.comm.all_expert_offsets = buf.all_expert_offsets; - p.comm.all_expert_counts = buf.all_expert_counts; - p.comm.pe_sync.cta_counter = buf.cta_counter; - p.comm.pe_sync.num_pes = num_pes; - p.comm.pe_sync.my_pe = my_pe; - p.comm.pe_sync.team = team; - p.comm.hidden_dim = hidden_dim; - p.comm.num_experts = num_experts; - p.comm.experts_per_pe = experts_per_pe; - p.comm.n_gemm = n_gemm; // GEMM/comm-active CTA count (floor_NS) - - p.num_tokens = num_tokens; - p.intermediate_dim = intermediate_dim; - p.top_k = top_k; - p.total_slots = max_total_slots; - p.num_m_tiles = num_m_tiles_in; - - // ---- MLP invariant dims (passed as __grid_constant__) ---------------------- - MlpDims mlp_dims; - mlp_dims.hidden_dim = hidden_dim; - mlp_dims.intermediate_dim = intermediate_dim; - mlp_dims.num_experts = experts_per_pe; - mlp_dims.total_n_rows_1 = total_n_rows_1; - mlp_dims.total_n_rows_2 = total_n_rows_2; - mlp_dims.expert_n_stride_1 = weight_expert_stride_rows / Traits1::TileN; - mlp_dims.num_n_tiles_1 = num_n_tiles_1; - mlp_dims.num_n_tiles_2 = num_n_tiles_2; - mlp_dims.num_k_tiles_1 = (Compute == 90) - ? (hidden_dim + Traits1::TileK - 1) / Traits1::TileK - : hidden_dim / Traits1::TileK; - mlp_dims.num_k_tiles_2 = (Compute == 90) - ? (intermediate_dim + Traits2::TileK - 1) / Traits2::TileK - : intermediate_dim / Traits2::TileK; - // Row-tile extent of the symmetric local_output (y_buf) for the - // direct-to-peer Y store (num_y_m_tiles when storing at the token-tile row). - mlp_dims.y_buf_m_tiles = max_total_slots / kTileM; - mlp_dims.phase_counter = buf.phase_counter; - - // ---- MLP TMA descriptors ----------------------------------------------- - // X for MLP is x_sorted (dispatch destination, symmetric). - // All shape/stride values are cast to int64_t — see the long comment - // below the tma_load_c block for the rationale. B and C in particular - // are large: their row count = experts_per_pe · intermediate_dim, - // times hidden_dim → exceeds INT_MAX at production shapes. - auto tma_load_x = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(buf.x_sorted)), - make_shape(static_cast(max_total_slots), - static_cast(hidden_dim)), - make_stride(static_cast(hidden_dim), Int<1>{})), - typename Traits1::SmemLayoutX_1{}); - - auto* ptr_B = reinterpret_cast(a.all_B); - auto* ptr_C = reinterpret_cast(a.all_C); - // Hopper keeps the expert as an explicit tensor mode. A tail TMA box can - // therefore go OOB only within that expert (loads zero-fill) instead of - // spilling into the next expert's flat rows. - auto tma_load_b = [&]() { - if constexpr (Compute == 90) { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(ptr_B), - make_shape(static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(experts_per_pe)), - make_stride( - static_cast(hidden_dim), - Int<1>{}, - a.weight_expert_stride)), - typename Traits1::SmemLayoutW_1{}); - } else { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(ptr_B), - make_shape(static_cast(total_n_rows_1), - static_cast(hidden_dim)), - make_stride(static_cast(hidden_dim), Int<1>{})), - typename Traits1::SmemLayoutW_1{}); - } - }(); - auto tma_load_c = [&]() { - if constexpr (Compute == 90) { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(ptr_C), - make_shape(static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(experts_per_pe)), - make_stride( - static_cast(hidden_dim), - Int<1>{}, - a.weight_expert_stride)), - typename Traits1::SmemLayoutW_1{}); - } else { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(ptr_C), - make_shape(static_cast(total_n_rows_1), - static_cast(hidden_dim)), - make_stride(static_cast(hidden_dim), Int<1>{})), - typename Traits1::SmemLayoutW_1{}); - } - }(); - - // New design's TMA store box = per-WG epi slot (AtomTileM, EpiChunkN). - // - // Shape / stride values are cast to int64_t so CUTE's layout math - // (cosize, byte-offset, total-elements) uses 64-bit arithmetic. The - // Z tensor in particular can exceed INT_MAX elements when - // grid_x · ZBuf · TileM · intermediate_dim > 2^31 — e.g. NSplit=2 - // ZBuf=16 at intermediate_dim=16384 wraps to a negative stride and - // IMAs the kernel. Casting one operand promotes the multiply to - // int64_t; we cast both to be safe across CUTE versions. - auto tma_store_z = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(buf.z_buf), - make_shape(static_cast(buf.z_rows), - static_cast(intermediate_dim)), - make_stride(static_cast(intermediate_dim), Int<1>{})), - typename Traits1::SmemLayoutStoreSlot{}); - - auto tma_load_z = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(buf.z_buf)), - make_shape(static_cast(buf.z_rows), - static_cast(intermediate_dim)), - make_stride(static_cast(intermediate_dim), Int<1>{})), - typename Traits2::SmemLayoutZ_1{}); - - auto* ptr_A = reinterpret_cast(a.all_A); - // Same expert isolation for the down projection, including K-tail loads. - auto tma_load_a = [&]() { - if constexpr (Compute == 90) { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(ptr_A), - make_shape(static_cast(hidden_dim), - static_cast(intermediate_dim), - static_cast(experts_per_pe)), - make_stride( - static_cast(intermediate_dim), - Int<1>{}, - static_cast(hidden_dim) * intermediate_dim)), - typename Traits2::SmemLayoutW_1{}); - } else { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(ptr_A), - make_shape(static_cast(total_n_rows_2), - static_cast(intermediate_dim)), - make_stride(static_cast(intermediate_dim), Int<1>{})), - typename Traits2::SmemLayoutW_1{}); - } - }(); - - auto tma_store_y = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(y_buf_ptr), - make_shape(static_cast(max_total_slots), - static_cast(hidden_dim)), - make_stride(static_cast(hidden_dim), Int<1>{})), - typename Traits2::SmemLayoutStore_1{}); - - // Remote staging is a flat list of L = MC · kCommNumStages slots, - // each TileM rows tall. MC = num_blocks / NC (== gridDim.x / NC) — MUST - // match the ring size everywhere else (allocate_moe_buffers, the per-launch - // memset, the consumer's ring_len). Sizing this TMA descriptor on n_gemm - // would leave the wrap-around slots outside its row extent (TMA OOB) when - // num_blocks/NC > n_gemm/NC (num_sms % NSplit != 0). - int mc_remote = num_blocks / Config::kNC; - int remote_total_rows = mc_remote * Config::kCommNumStages * kTileM; - - auto tma_load_x_remote = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(buf.src_staging)), - make_shape(static_cast(remote_total_rows), - static_cast(hidden_dim)), - make_stride(static_cast(hidden_dim), Int<1>{})), - typename Traits1::SmemLayoutX_1{}); - - auto tma_store_y_remote = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(buf.dst_staging), - make_shape(static_cast(remote_total_rows), - static_cast(hidden_dim)), - make_stride(static_cast(hidden_dim), Int<1>{})), - typename Traits2::SmemLayoutStore_1{}); - - // ---- GET-via-TMA descriptors (FWD-only, raw CUtensorMap) ----------------- - // SAME-HOST source maps over nvshmem_ptr(x_sorted, peer) + one dst map over - // local src_staging. TMA get is enabled only for same-host NVLink P2P peers - // (cross-host peers fall back to warp getmem); also needs divisible - // hidden_dim + smem budget (checked below). - // - // src_x_desc is indexed by INTRA-HOST local rank (pe % gpus_per_node), so the - // array only ever needs gpus_per_node (≤ kGetMaxPes) entries regardless of - // total PE count. Assumes host-contiguous, host-aligned ranks (true for a - // WORLD-contiguous team); the do_get / RemoteIter is_local predicate uses the - // same arithmetic. - GetTmaDescs get_descs{}; - get_descs.enabled = 0; - int gpus_per_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); - LIGER_CHECK(gpus_per_node >= 1 && gpus_per_node <= liger::kGetMaxPes, - "gpus_per_node (", gpus_per_node, ") must be in [1, kGetMaxPes=", - liger::kGetMaxPes, "]"); - get_descs.gpus_per_node = gpus_per_node; - get_descs.my_pe = my_pe; - constexpr int kGetBoxRows = - Config::kTileM / (liger::kNumGetWarpsFwd * Config::kNC); - size_t get_bounce_bytes = (size_t) - liger::GetBounce::kTotalBytes; - { - bool ok = (hidden_dim % liger::kGetKChunk == 0); - for (int tp = 0; ok && tp < num_pes; ++tp) { - // Only same-host peers are P2P-reachable → only they get a TMA - // descriptor (indexed by local rank). Cross-host peers skip it and - // use getmem at runtime. - bool same_host = (tp / gpus_per_node) == (my_pe / gpus_per_node); - if (!same_host) continue; - int gpe = nvshmem_team_translate_pe(team, tp, NVSHMEM_TEAM_WORLD); - void* pp = nvshmem_ptr(buf.x_sorted, gpe); // self → own ptr (non-null) - if (pp == nullptr) { ok = false; break; } // same-host but no P2P → disable TMA - ok = make_get_tma_map(get_descs.src_x_desc[tp % gpus_per_node], pp, - hidden_dim, max_total_slots, kGetBoxRows); - } - if (ok) - ok = make_get_tma_map(get_descs.dst_staging_desc, buf.src_staging, - hidden_dim, remote_total_rows, kGetBoxRows); - get_descs.enabled = ok ? 1 : 0; - if (const char* e = std::getenv("LIGER_GET_TMA")) - if (e[0] == '0') get_descs.enabled = 0; - } - - // ---- Direct-to-peer Y store descriptors (FWD-only) ----------------------- - // Per-peer TMA-store maps over nvshmem_ptr(y_buf, peer) for SAME-HOST peers - // (self included), indexed by intra-host local rank. When a tile is_local, - // MLP2 stores Y straight into the destination's symmetric local_output via - // its peer descriptor (token-tile row), bypassing dst_staging + the put warp - // (which then only rubber-stamps). Same shape/layout as tma_store_y; only the - // base pointer differs (the peer's y_buf). For single-PE the only peer is - // self → desc[my_pe%gpus_per_node] == tma_store_y. - PeerYStoreDescs peer_y{}; - peer_y.enabled = 1; - for (int tp = 0; tp < num_pes; ++tp) { - bool same_host = (tp / gpus_per_node) == (my_pe / gpus_per_node); - if (!same_host) continue; // cross-host peers keep the dst_staging+putmem path - int gpe = nvshmem_team_translate_pe(team, tp, NVSHMEM_TEAM_WORLD); - void* py = nvshmem_ptr(buf.y_buf, gpe); // same-host → P2P ptr (non-null on NVLink) - LIGER_CHECK(py != nullptr, - "direct-to-peer Y store: nvshmem_ptr(y_buf, same-host pe ", tp, - ") is null (no P2P) — unsupported for this build"); - peer_y.desc[tp % gpus_per_node] = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(py)), - make_shape(static_cast(max_total_slots), - static_cast(hidden_dim)), - make_stride(static_cast(hidden_dim), Int<1>{})), - typename Traits2::SmemLayoutStore_1{}); - } - - // TMA descriptor for the final combine output [num_tokens, hidden_dim]. - // Box shape matches combine_tokens_tma's WG tile: 32 tokens × 256 hidden. - auto* ptr_combine_out = output_ptr; - auto tma_store_y_combine = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(ptr_combine_out), - make_shape(static_cast(num_tokens), - static_cast(hidden_dim)), - make_stride(static_cast(hidden_dim), Int<1>{})), - CombineSmemLayoutSlot{}); - - // ---- Dynamic smem -------------------------------------------------------- - int sort_num_warps = kNumThreads / 32; - size_t smem_sort = sort_num_warps * num_experts * sizeof(int); - - size_t smem_mlp = sizeof(MoeSmem); - smem_mlp = (smem_mlp + sizeof(int) - 1) & ~(sizeof(int) - 1); - smem_mlp += num_pes * (experts_per_pe + 1) * sizeof(int); // remote_offsets - smem_mlp += num_pes * experts_per_pe * sizeof(int); // remote_counts - - // GET-via-TMA bounce: append a 128-aligned region (must match the kernel's - // get_bounce offset). Disable the TMA-get path if it would exceed the - // opt-in dynamic-smem cap (227 KiB) — that config just keeps warp getmem. - if (get_descs.enabled) { - size_t mlp_with_bounce = ((smem_mlp + 127) & ~((size_t)127)) + get_bounce_bytes; - if (mlp_with_bounce <= (size_t)227 * 1024) smem_mlp = mlp_with_bounce; - else get_descs.enabled = 0; - } - - // Combine phase reinterprets the same SMEM region as CombineSmem. - constexpr int kCombineNumWGs = - Config::kNumThreads / (kCombineWarpsPerWG * kCombineWarpSize); - size_t smem_combine = sizeof(CombineSmem); - - size_t smem_size = std::max({smem_sort, smem_mlp, smem_combine}); - - // ---- Launch fused kernel ----------------------------------------------- - auto kernel_fn = moe_fused_kernel; - - if (cudaError_t e = cudaFuncSetAttribute(kernel_fn, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: cudaFuncSetAttribute(MaxDynamicSharedMemorySize=", smem_size, - ") failed: ", cudaGetErrorString(e)); - - // Reset per-call buffers in a single fused-memset kernel — one CTA per - // region. Replaces 11 cudaMemsetAsync calls that would otherwise - // serialize on the stream. - // - // Notes: - // - dispatch_done / mlp_done are already reset upstream in - // allocate_moe_buffers (the original "monotonic, never reset" design - // was retired — wait_until could exit on stale residual signal - // values in multi-PE shape-transition workloads). - // - all_expert_offsets stays unreset (overwritten by - // broadcast_expert_offsets on every call). - // - y_buf stays unreset (slots referenced by token_expert_slots[t,k] - // are always written by local MLP or the source PE's NVSHMEM - // put-back; padding slots are never read by combine_tokens_tma). - // - src_staging / dst_staging / z_buf intentionally NOT zero'd (see - // comment block below). - // - // Per-region reasons: - // - cta_counter: PeSync.barrier() does `% num_ctas`. Residual values - // from a prior call with different num_ctas misalign to the new - // modulus, hanging barrier() forever. - // - tile_expert_ids (0xff): sentinel-fill so unfetched slots produce - // expert_id = -1; matches moe_bwd.cu's 0xff fill of - // tile_expert_ids_{x,dy}. Without this, stale slot values from prior - // calls can match a valid expert id and confuse the MLP-side remote - // iterator at higher PE counts. - // - src/dst_ready/consumed: atomically accumulated; stale ready/ - // consumed values make the consumer-acquire wait fall through. - // - tile_expert_counts / cta_sums: atomicAdd'd during sort; must - // start at 0. - { - auto& scfg = get_symm_config(); - // L = MC · kCommNumStages, MC = num_blocks / NC (== gridDim.x / NC). - // MUST match the ring size in allocate_moe_buffers and the device-side - // comm/consumer stride — otherwise this memset under-covers the signal / - // tile_expert_ids arrays (slots [n_gemm/NC·CS, gridDim.x/NC·CS) left - // uninitialized) and the consumer reads stale src_ready / garbage expert - // ids for those slots (see the OOB note in allocate_moe_buffers). - int mc_local = num_blocks / Config::kNC; - int comm_l = mc_local * Config::kCommNumStages; - size_t sig_bytes = - static_cast(comm_l) * Config::kSubTiles * sizeof(int); - size_t eid_bytes = static_cast(comm_l) * sizeof(int); - int max_m_tiles = scfg.max_total_slots / kTileM; - int max_tile_counts = max_m_tiles * scfg.max_num_experts; - - FusedMemsetArgs ma{}; - fused_memset_add(ma, buf.phase_counter, - ((num_blocks + 1) / 2) * sizeof(int), 0x00); - fused_memset_add(ma, buf.cta_counter, - sizeof(int), 0x00); - fused_memset_add(ma, buf.cta_done, - num_blocks * sizeof(int), 0x00); - fused_memset_add(ma, buf.sorted_token_ids, - static_cast(max_total_slots) * sizeof(int), 0xff); - fused_memset_add(ma, buf.src_ready, sig_bytes, 0x00); - fused_memset_add(ma, buf.src_consumed, sig_bytes, 0x00); - fused_memset_add(ma, buf.dst_ready, sig_bytes, 0x00); - fused_memset_add(ma, buf.dst_consumed, sig_bytes, 0x00); - fused_memset_add(ma, buf.tile_expert_ids, eid_bytes, 0xff); - fused_memset_add(ma, buf.tile_expert_counts, - static_cast(max_tile_counts) * sizeof(int), 0x00); - fused_memset_add(ma, buf.cta_sums, - static_cast(num_blocks) * scfg.max_num_experts * sizeof(int), 0x00); - fused_memset_launch(ma, stream); - } - - // src_staging / dst_staging / z_buf intentionally NOT zero'd here: - // - // src_staging: comm warps fill the slots THIS CTA fetches; unfetched - // slots are gated by tile_expert_ids = -1 (set via the 0xff memset - // above) and RemoteIter::has_next() walks only per_cta_tiles, so the - // remote MLP never reads unfetched slots. Stale prev-call contents - // are inert under both bounds checks. - // - // dst_staging: write-then-read for routed slots only. Remote MLP - // writes Y for tiles RemoteIter returns; comm warps put back only - // those tiles. Padding/unused slots are never read at either end. - // - // z_buf: per-tile write→read within a single MLP launch. LocalIter is - // range-bounded to [local_e_start, local_e_end) experts via - // set_range; the remote pass is bounded by per_cta_tiles. Tiles - // outside the current call's active range are never written and - // never read — buffer-pool size growth from previous (larger) calls - // is invisible to the kernel. - // - // The tile_expert_ids 0xff fill above is what makes this safe — keep - // that. - - kernel_fn<<>>( - tma_load_x, tma_load_b, tma_load_c, tma_store_z, - tma_load_z, tma_load_a, tma_store_y, - tma_load_x_remote, tma_store_y_remote, - tma_store_y_combine, - mlp_dims, p, static_nsplit, get_descs, peer_y); - if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_fwd: kernel launch failed for grid=", grid.x, - " threads=", kNumThreads, " smem=", smem_size, ": ", - cudaGetErrorString(e)); - - // Post-launch global barrier: ensures all NVSHMEM operations (tile puts, - // signal ops) across all PEs are complete before any PE starts the next - // call. Without this, a fast PE could overwrite symmetric buffers - // (x_sorted, y_buf, tile signals) while a slow PE is still reading them. - // - // Stream-ordered variant — capturable into a CUDA graph. The host-side - // nvshmem_team_sync(team) would do an internal cudaStreamSynchronize and - // therefore fail with "operation not permitted when stream is capturing". - //if (num_pes > 1) { - // nvshmemx_barrier_on_stream(team, stream); - //} - - // Hand the persistent symmetric buffers back to the ABI wrapper so it can - // alias them as the *_out_symm views (shapes use the static MoeSymmConfig - // max; the views do NOT own the storage — SymmetricMemoryStack does). Y and - // the two sort outputs were written in place into the caller-owned buffers. - *a.x_sorted_out = buf.x_sorted; - *a.y_buf_out = y_buf_ptr; - *a.all_expert_offsets_out = buf.all_expert_offsets; - if (a.all_expert_counts_out != nullptr) - *a.all_expert_counts_out = buf.all_expert_counts; -} - -// ── Runtime dispatch table for tuned-config lookup ──────────────────── - -using MoeFwdFn = void (*)(const MoeFwdArgs&, int static_nsplit); -static constexpr int kFwdRuntimeNsplitSeed = 8; - -struct DispatchEntry { - int Compute; - int TileN1, TileK1, Stages1, EpiChunkN1; - int TileN2, TileK2, Stages2, EpiChunkN2; - int ZBufferSlots, CommNumStages; - int TileM; // tuned GEMM-tile field (legacy dispatch-table name) - int GemmTileM; // GEMM tile ∈ {64,128} - MoeFwdFn fn; -}; - -// 13-arg core: explicit GemmTileM. 12-arg shim sets GemmTileM = comm TileM. -#define LIGER_MOE_DISPATCH_ENTRY_G_C(Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) \ - { Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM, \ - &moe_fused_fwd_bf16 }, -#define LIGER_MOE_DISPATCH_ENTRY_C(Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) \ - LIGER_MOE_DISPATCH_ENTRY_G_C(Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, TileM) -#define LIGER_MOE_DISPATCH_ENTRY_G_SM90(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) \ - LIGER_MOE_DISPATCH_ENTRY_G_C(90, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) -#define LIGER_MOE_DISPATCH_ENTRY_SM90(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) \ - LIGER_MOE_DISPATCH_ENTRY_C(90, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) -#define LIGER_MOE_DISPATCH_ENTRY_G_SM100(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) \ - LIGER_MOE_DISPATCH_ENTRY_G_C(100, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) -#define LIGER_MOE_DISPATCH_ENTRY_SM100(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) \ - LIGER_MOE_DISPATCH_ENTRY_C(100, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) - -static const DispatchEntry kDispatchTable[] = { -#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 90 - LIGER_MOE_FWD_DISPATCH_CONFIGS_SM90(LIGER_MOE_DISPATCH_ENTRY_SM90, LIGER_MOE_DISPATCH_ENTRY_G_SM90) +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "moe_sm90.inc" +#else +#include "moe_sm100.inc" #endif -#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 100 - LIGER_MOE_FWD_DISPATCH_CONFIGS_SM100(LIGER_MOE_DISPATCH_ENTRY_SM100, LIGER_MOE_DISPATCH_ENTRY_G_SM100) -#endif -}; -#undef LIGER_MOE_DISPATCH_ENTRY_SM100 -#undef LIGER_MOE_DISPATCH_ENTRY_G_SM100 -#undef LIGER_MOE_DISPATCH_ENTRY_SM90 -#undef LIGER_MOE_DISPATCH_ENTRY_G_SM90 -#undef LIGER_MOE_DISPATCH_ENTRY_C -#undef LIGER_MOE_DISPATCH_ENTRY_G_C -static constexpr int kNumDispatchEntries = - sizeof(kDispatchTable) / sizeof(kDispatchTable[0]); - -} // namespace liger - -// Joint fwd+bwd tuning table (auto-generated by benchmarks/tune_moe_fwd_bwd). -// We only consume the Fwd_* fields here; moe_bwd.cu consumes the same table's -// Bwd_* fields. Picking fwd and bwd from a single table guarantees the two -// directions are tuned as a pair (so their TileM agrees and their per-shape -// templates were measured together, not independently). -#include "moe_fwd_bwd_tuning_configs.cuh" - -namespace liger { - -// TunedConfig is the fwd-only projection of TunedConfigFwdBwd. The rest of -// the dispatch code (validity check, dispatch-entry match) was written -// against this struct and is reused unchanged. -struct TunedConfig { - int Compute; - int TK, TKE, D, I; - int TileN1, TileK1, Stages1, EpiChunkN1; - int TileN2, TileK2, Stages2, EpiChunkN2; - int ZBufferSlots, CommNumStages; - int TileM; // comm tile - int GemmTileM; // gemm tile (auto path: == TileM; FORCE_CONFIG may set 64) -}; - -// Project a TunedConfigFwdBwd row down to its fwd template params. The tuned -// table has no GemmTileM column, so the auto path keeps gemm == comm (TileM). -static TunedConfig project_fwd(const TunedConfigFwdBwd& c, int compute) { - return TunedConfig{ - compute, - c.TK, c.TKE, c.D, c.I, - c.Fwd_TileN1, c.Fwd_TileK1, c.Fwd_Stages1, c.Fwd_EpiChunkN1, - c.Fwd_TileN2, c.Fwd_TileK2, c.Fwd_Stages2, c.Fwd_EpiChunkN2, - c.Fwd_ZBufferSlots, c.Fwd_CommNumStages, - c.Fwd_TileM, c.Fwd_TileM, - }; -} - -// Find the dispatch table entry matching a TunedConfig's fields. -static const DispatchEntry* find_dispatch_entry(const TunedConfig& c) { - for (int i = 0; i < kNumDispatchEntries; ++i) { - const auto& e = kDispatchTable[i]; - if (e.Compute == c.Compute && - e.TileN1 == c.TileN1 && e.TileK1 == c.TileK1 && e.Stages1 == c.Stages1 && - e.EpiChunkN1 == c.EpiChunkN1 && - e.TileN2 == c.TileN2 && e.TileK2 == c.TileK2 && e.Stages2 == c.Stages2 && - e.EpiChunkN2 == c.EpiChunkN2 && - e.ZBufferSlots == c.ZBufferSlots && e.CommNumStages == c.CommNumStages && - e.TileM == c.TileM && e.GemmTileM == c.GemmTileM) { - return &e; - } - } - return nullptr; -} - -// Mirrors shape_valid() in tune_moe.cu: structural constraints the kernel -// requires on (D, I) for a given config. Configs that violate these for -// the runtime shape are skipped during nearest-neighbor lookup so we -// don't dispatch to a template instantiation the kernel can't run. -// -// Cooperative 2-WG consumer layout switches on TileM: -// - TileM=128 (M-split, Layout<_2,_1,_1>): each WG owns AtomTileM=64 -// rows × TileN cols. EpiChunkN must divide TileN. -// - TileM=64 (N-split, Layout<_1,_2,_1>): each WG owns TileM=64 rows -// × TileN/2 cols. EpiChunkN must divide TileN/2; TileN must be even. -// - Hopper uses ceil-div tile counts; SM100 retains exact divisibility. -// - store_buf is (AtomTileM=64, EpiChunkN) per WG in both modes. -static bool tuned_config_valid(int D, int I, const TunedConfig& c) { - if (c.TileM != 64 && c.TileM != 128) return false; - // The current Blackwell UMMA MLP1 consumer only supports the 128-row comm - // tile. Keep mirrored SM90 TileM=64 tuning rows out of SM100 auto-dispatch - // until that TMEM epilogue layout is implemented. - if (c.Compute == 100 && c.TileM != 128) return false; - if (c.Compute == 100) { - if (D % c.TileK1 != 0) return false; - if (I % c.TileN1 != 0) return false; - if (I % c.TileK2 != 0) return false; - if (D % c.TileN2 != 0) return false; - } - if (D % 8 != 0 || I % 8 != 0) return false; - // NC selection mirrors moe_fused_fwd_bf16: TileM=128 → NC=4, - // TileM=64 → NC=2. With the ticket-based GEMM iterator the only - // runtime requirement is NC | (gridDim.x · gridDim.y); runtime NS / - // CommNumStages no longer need to share any divisibility with NC. - const int nc = (c.TileM == 128) ? 4 : 2; - (void)nc; - // EpiChunkN must divide the per-WG N extent (WgTileN). - const int wg_tile_n1 = (c.TileM == 128) ? c.TileN1 : c.TileN1 / 2; - const int wg_tile_n2 = (c.TileM == 128) ? c.TileN2 : c.TileN2 / 2; - if (wg_tile_n1 <= 0 || wg_tile_n1 % c.EpiChunkN1 != 0) return false; - if (wg_tile_n2 <= 0 || wg_tile_n2 % c.EpiChunkN2 != 0) return false; - // N-split requires TileN to be evenly halvable (TileN/2 must be ≥ the - // WGMMA atom-N min, which is 8 for bf16 but 64 for our smem layout). - if (c.TileM == 64 && (c.TileN1 % 2 != 0 || c.TileN2 % 2 != 0)) return false; - - // SMEM cap (H100 = 228 KiB; budget 224 KiB inside MoeFusedConfig). - // Per-CTA, bf16 (2 B/elem). TileM scales the X/Z staging: - // MLP1 = 2 * (TileM·TK1·S1 + 2·TN1·TK1·S1 + 2·64·EC1) - // MLP2 = 2 * (TileM·TK2·S2 + TN2·TK2·S2 + 2·64·EC2) - // smem ≈ max(MLP1, MLP2) - int mlp1_bytes = 2 * (c.TileM * c.TileK1 * c.Stages1 - + 2 * c.TileN1 * c.TileK1 * c.Stages1 - + 2 * 64 * c.EpiChunkN1); - int mlp2_bytes = 2 * (c.TileM * c.TileK2 * c.Stages2 - + c.TileN2 * c.TileK2 * c.Stages2 - + 2 * 64 * c.EpiChunkN2); - int smem = mlp1_bytes > mlp2_bytes ? mlp1_bytes : mlp2_bytes; - constexpr int kSmemBudget = 224 * 1024; - if (smem > kSmemBudget) return false; - return true; -} - -static bool find_tuned_table(int compute, int n_pes, - const TunedConfigFwdBwd** table, - int* table_n) { - const TunedConfigFwdBwdTable* tables = - (n_pes <= 1) ? kTunedConfigTablesSingle : kTunedConfigTablesMulti; - const int tables_n = - (n_pes <= 1) ? kNumTunedConfigTablesSingle : kNumTunedConfigTablesMulti; - for (int i = 0; i < tables_n; ++i) { - if (tables[i].Compute == compute) { - *table = tables[i].configs; - *table_n = tables[i].count; - return *table_n > 0; - } - } - return false; -} - -// Nearest-neighbor lookup in log-space (TK, TKE, D, I), where TK = T*top_k -// is the total routed-token count (the actual M-axis the kernel sees) and -// TKE = TK / E_local (E_local = E / n_pes) is the per-LOCAL-expert avg -// K-range. Adding TKE to the distance lets the dispatcher distinguish -// shapes whose mlp3/mlp4 work distribution differs even when (TK, D, I) -// is the same — e.g. same total tokens, different expert count or PE -// count. -// Configs that violate structural constraints for the runtime (D, I) are -// skipped — the nearest tuned point is not always valid for the actual -// shape (e.g. tuned at I=4096 with TileN1=256, runtime I=768). -// Out-param API instead of pointer-into-static-array because each candidate -// TunedConfig is projected on the fly from kTunedConfigsFwdBwd (no stable -// address). Returns true when a valid config was found. -static bool find_nearest_tuned(int compute, int TK, int TKE, int D, int I, int n_pes, - TunedConfig* out) { - // World-size class: single-GPU (n_pes <= 1) and multi-GPU tables are tuned - // and stored separately, since the comm path (and thus the optimal NSplit/ - // CommNumStages/GEMM-tile) differs sharply between one PE and many. - const TunedConfigFwdBwd* table = nullptr; - int table_n = 0; - if (!find_tuned_table(compute, n_pes, &table, &table_n)) return false; - bool found = false; - double best_dist = 1e30; - double lTK = std::log((double)TK); - double lTKE = std::log((double)TKE); - double lD = std::log((double)D); - double lI = std::log((double)I); - for (int i = 0; i < table_n; ++i) { - TunedConfig c = project_fwd(table[i], compute); - if (!tuned_config_valid(D, I, c)) continue; - // Guard against TKE=0 in legacy rows (would yield -inf in log). - double dTK = std::log((double)c.TK) - lTK; - double dTKE = c.TKE > 0 - ? std::log((double)c.TKE) - lTKE - : 0.0; // legacy / placeholder rows contribute 0 to the TKE axis - double dD = std::log((double)c.D) - lD; - double dI = std::log((double)c.I) - lI; - double d = dTK*dTK + dTKE*dTKE + dD*dD + dI*dI; - if (d < best_dist) { best_dist = d; *out = c; found = true; } - } - return found; -} - -// ============================================================================ -// Auto-dispatch wrapper — pure tuned-table lookup -// ============================================================================ -// -// Picks the nearest tuned config (log-space distance over (T, D, I)) from -// moe_fwd_bwd_tuning_configs.cuh and dispatches to the matching template -// instantiation in kDispatchTable. Both tables are populated by the -// tune_moe_fwd_bwd sweep — -// shapes outside the swept range error out instead of falling back to a -// heuristic. Add the missing shape to the sweep (or to LIGER_MOE_TUNE_CONFIGS -// directly) to extend coverage. - -// Internal auto-dispatch: picks the nearest tuned config for the runtime shape -// (or honours LIGER_MOE_FORCE_CONFIG) and invokes the matching template -// instantiation via the dispatch table. Writes the chosen comm TileM to -// *chosen_tile_m. All shape info comes from the MoeFwdArgs bundle; the launcher -// writes the outputs in place. -void moe_fused_fwd_dispatch(const MoeFwdArgs& a, int* chosen_tile_m) { - const int T = a.num_tokens; - const int D = a.hidden_dim; - const int I = a.intermediate_dim; - const int top_k = a.top_k; - const int num_experts = a.num_experts; - const int TK = T * top_k; // total routed tokens — the M-axis the kernel sees - // TKE = TK / E_local where E_local = num_experts / n_pes (per-LOCAL-expert - // avg K-range). Matches the tuner's definition so the nearest-neighbor - // lookup picks rows recorded at comparable per-PE load. - const int n_pes = nvshmem_team_n_pes(a.team); - const int E_local = (n_pes > 0) ? std::max(1, num_experts / n_pes) : num_experts; - // The lookup operates in log space, so clamp the per-expert average to - // the smallest representable workload for tiny but valid token batches. - const int TKE = std::max(1, TK / E_local); - const int compute = moe_detect_compute_dispatch_key(a.device, "moe_fused_fwd_bf16_auto"); - - // ── Override path: LIGER_MOE_FORCE_CONFIG=TN1,TK1,S1,EC1,TN2,TK2,S2,EC2,ZB,CS,TM[,GemmTM] - // Lets a test suite sweep every compiled config without rebuilding. The - // optional final int is GemmTileM (comm/gemm tile decouple); omitted → == TM. - if (const char* s = std::getenv("LIGER_MOE_FORCE_CONFIG")) { - TunedConfig tc{}; - tc.Compute = compute; - tc.TileM = 128; // default when omitted (backward compat) - tc.GemmTileM = -1; // sentinel → default to TileM after parse - int n = std::sscanf(s, "%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d", - &tc.TileN1, &tc.TileK1, &tc.Stages1, &tc.EpiChunkN1, - &tc.TileN2, &tc.TileK2, &tc.Stages2, &tc.EpiChunkN2, - &tc.ZBufferSlots, &tc.CommNumStages, - &tc.TileM, &tc.GemmTileM); - LIGER_CHECK(n == 10 || n == 11 || n == 12, - "LIGER_MOE_FORCE_CONFIG must have 10, 11 or 12 comma-separated ints " - "(TN1,TK1,S1,EC1,TN2,TK2,S2,EC2,ZB,CS[,TM[,GemmTM]]); got '", s, "'"); - if (tc.GemmTileM <= 0) tc.GemmTileM = tc.TileM; // omitted → gemm == comm - const DispatchEntry* de = find_dispatch_entry(tc); - LIGER_CHECK(de != nullptr, - "LIGER_MOE_FORCE_CONFIG=", s, " is not in the compiled dispatch " - "table. Add it to LIGER_MOE_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); - de->fn(a, kFwdRuntimeNsplitSeed); - *chosen_tile_m = de->Compute == 100 - ? MoeCommGeometry<100>::TileM - : MoeCommGeometry<90>::TileM; - return; - } - - TunedConfig tc{}; - LIGER_CHECK(find_nearest_tuned(compute, TK, TKE, D, I, n_pes, &tc), - "moe_fused_fwd_bf16_auto: no valid tuned config for compute=", compute, - " shape (T=", T, - " K=", top_k, " TK=", TK, " TKE=", TKE, " D=", D, " I=", I, - "). Either moe_fwd_bwd_tuning_configs.cuh is empty (run tune_moe_fwd_bwd), " - "or every tuned entry violates structural constraints for this (D, I) " - "(tile/layout compatibility and shared-memory limits). " - "Add a compatible config to LIGER_MOE_TUNE_CONFIGS and re-tune."); - - const DispatchEntry* de = find_dispatch_entry(tc); - LIGER_CHECK(de != nullptr, - "moe_fused_fwd_bf16_auto: tuned config (Compute=", tc.Compute, - " TN1=", tc.TileN1, "/", tc.TileK1, "/", tc.Stages1, "/EC", tc.EpiChunkN1, - " TN2=", tc.TileN2, "/", tc.TileK2, "/", tc.Stages2, "/EC", tc.EpiChunkN2, - " ZB=", tc.ZBufferSlots, " CS=", tc.CommNumStages, - ") for shape (T=", T, " K=", top_k, " TK=", TK, " D=", D, " I=", I, - ") is not in the compiled dispatch table. " - "Add it to LIGER_MOE_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); - - de->fn(a, kFwdRuntimeNsplitSeed); - *chosen_tile_m = de->Compute == 100 - ? MoeCommGeometry<100>::TileM - : MoeCommGeometry<90>::TileM; -} - -// ───────────────────────────────────────────────────────────────────── -// Default (untuned) wrapper — single hardcoded config per TileM. Use -// this for the regular bench_moe / test_moe paths so they don't depend -// on the autotuner dispatch table. Returns the same 6-tuple as the auto -// wrapper. -// -// Each config below must have a matching explicit instantiation. The -// TileM=64 row is in LIGER_MOE_TUNE_CONFIGS (so LIGER_MOE_INSTANTIATE -// emits it automatically); the TileM=128 row is emitted as a direct -// `template` instantiation right below this function. -// ───────────────────────────────────────────────────────────────────── - -// Returns the compiled MoE config table as a list of 12-int rows -// (Compute, TileN1, TileK1, Stages1, EpiChunkN1, TileN2, TileK2, Stages2, -// EpiChunkN2, ZBufferSlots, CommNumStages, TileM). Kept for the tuner / a future -// sweep binding so they stay in sync with the dispatch table without manual -// duplication. Internal (std::vector) — not part of the flat ABI. -std::vector> -moe_list_compiled_configs() { - std::vector> out; - out.reserve(kNumDispatchEntries); - for (int i = 0; i < kNumDispatchEntries; ++i) { - const auto& e = kDispatchTable[i]; - out.emplace_back(e.Compute, e.TileN1, e.TileK1, e.Stages1, e.EpiChunkN1, - e.TileN2, e.TileK2, e.Stages2, e.EpiChunkN2, - e.ZBufferSlots, e.CommNumStages, e.TileM); - } - return out; -} - -} // namespace liger - -// ============================================================================ -// Flat extern "C" control ABI (declared in liger_cute/moe.h) -// ============================================================================ - -namespace { - -using liger_cute::detail::guarded; - -} // namespace - -extern "C" { - -liger_cute_status_t liger_cute_moe_get_symm_config(liger_cute_moe_symm_config_t* out) { - return guarded([&]() -> liger_cute_status_t { - LIGER_CHECK(out != nullptr, "moe_get_symm_config: out is null"); - const liger::MoeSymmConfig& cfg = liger::get_symm_config(); - out->max_total_slots = cfg.max_total_slots; - out->max_num_experts = cfg.max_num_experts; - out->hidden_dim = cfg.hidden_dim; - out->num_pes = cfg.num_pes; - out->experts_per_pe = cfg.experts_per_pe; - out->max_top_k = cfg.max_top_k; - out->initialized = cfg.initialized ? 1 : 0; - return LIGER_CUTE_OK; - }); -} - -liger_cute_status_t liger_cute_moe_configure_symmetric( - const int max_tokens, const int hidden_dim, const int max_num_experts, - const int max_top_k, const int num_pes, const int num_hosts, - const int gpus_per_host) { - return guarded([&]() -> liger_cute_status_t { - liger::moe_configure_symmetric(max_tokens, hidden_dim, max_num_experts, - max_top_k, num_pes, num_hosts, gpus_per_host); - return LIGER_CUTE_OK; - }); -} - -liger_cute_status_t liger_cute_moe_pop_fwd(void) { - return guarded([&]() -> liger_cute_status_t { - liger::moe_pop_fwd(); - return LIGER_CUTE_OK; - }); -} - -} // extern "C" diff --git a/liger_cute_kernels/csrc/core/src/moe/moe.cuh b/liger_cute_kernels/csrc/core/src/moe/moe.cuh index ca547ba9..28101f62 100644 --- a/liger_cute_kernels/csrc/core/src/moe/moe.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/moe.cuh @@ -1,200 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// Fused MoE: MLP compute for local + remote tiles -// ═══════════════════════════════════════════════════════════════════ -// -// MLP-only device function — called by MLP warps (0,1,4-11). -// Comm warps (2-3) are handled separately at the kernel level to -// keep CommBuffers out of the MLP register scope. -// -// Execution flow (when remote tiles exist): -// a. mlp_fused_fwd with LocalMlpTileIterator (first half of column) -// b. NamedBarrier (MLP warps only) -// c. mlp_fused_fwd with RemoteMlpTileIterator -// d. NamedBarrier (MLP warps only) -// e. mlp_fused_fwd with LocalMlpTileIterator (second half of column) -// -// Splitting the local pass around the remote pass lets the comm warps -// overlap remote gets with the first local half and remote puts with -// the second local half. When there are no remote tiles, the kernel -// runs a single full local pass and returns. -// -// For remote tiles: -// - X input comes from src_staging (filled by comm get) -// - Y output goes to dst_staging (drained by comm put) -// - The RemoteMlpTileIterator acquire/release hooks synchronize -// with comm warps via StagePipe -// ═══════════════════════════════════════════════════════════════════ +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -#include "mlp.cuh" - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// SymmetricMemoryStack keys for the persistent fwd→bwd buffers. -// Centralized so set_size / put / pop / top sites all agree. -// ═══════════════════════════════════════════════════════════════════ -inline constexpr const char* kSymmKeyXSorted = "x_sorted"; -inline constexpr const char* kSymmKeyYBuf = "y_buf"; -inline constexpr const char* kSymmKeyAllExpertOffsets = "all_expert_offsets"; -inline constexpr const char* kSymmKeyAllExpertCounts = "all_expert_counts"; - -// ═══════════════════════════════════════════════════════════════════ -// MoE shared memory — union of MLP smem + comm smem -// ═══════════════════════════════════════════════════════════════════ - -template -struct MoeSmem { - // MLP shared memory (fused phase 1 + 2 with pipeline barriers). - MlpFusedSmem mlp; - - // Comm shared memory (total_tiles count + remote_offsets pointer). - CommSmem comm; -}; - -// ═══════════════════════════════════════════════════════════════════ -// moe_fused_fwd — MLP warps only (no CommBuffers in scope) -// ═══════════════════════════════════════════════════════════════════ - -template 1 means the GEMM steps - // through SubTiles sub-tiles of Traits1::TileM (=GemmTileM) rows per - // communication staging slot. 1 → unchanged single-tile path. - int SubTiles, - int Compute = 90, - typename LocalIter, typename RemoteIter, - typename TmaLoadX, typename TmaLoadW1, typename TmaStoreZ, - typename TmaLoadZ, typename TmaLoadW2, typename TmaStoreY> -__device__ __forceinline__ void moe_fused_fwd( - MoeSmem& smem, - LocalIter& local_iter, - // Remote-phase comm pointers (no CommBuffers — keeps MLP register scope clean) - int* src_ready, - int* src_consumed, - int* dst_ready, - int* dst_consumed, - // RemoteIter re-derives (expert, pe) and the per-tile is_local fence-scope - // predicate from remote_offsets (SMEM) using these — no per-slot HBM. - int experts_per_pe, - int num_pes, - int my_pe, - int gpus_per_node, - // MLP1 TMA descriptors - TmaLoadX const& tma_load_x, - TmaLoadW1 const& tma_load_b, - TmaLoadW1 const& tma_load_c, - TmaStoreZ const& tma_store_z, - // MLP2 TMA descriptors - TmaLoadZ const& tma_load_z, - TmaLoadW2 const& tma_load_a, - TmaStoreY const& tma_store_y, - // Remote-phase TMA descriptors (X loaded from staging, Y stored to staging) - TmaLoadX const& tma_load_x_remote, - TmaStoreY const& tma_store_y_remote, - // Invariant MLP dims — const ref into __grid_constant__ .param space - const MlpDims& dims, - // Per-phase params (change between local and remote) - int num_tokens, - int num_m_tiles, - // Flat-grid logical coordinates: col = flat_id / NSplit (this CTA's - // GEMM column), grid_x = n_gemm / NSplit (number of columns). Under - // the flat launch gridDim.x = num_blocks and gridDim.y = 1, so these - // replace every blockIdx.x / gridDim.x / gridDim.y read below. - int col, - int grid_x, - int split, // this CTA's N-split index = flat_id % runtime_nsplit - int runtime_nsplit, - int local_expert_start = 0, - // Direct-to-peer Y store: per-peer TMA descriptors over - // nvshmem_ptr(local_output, peer). Gated per-tile on is_local downstream. - const TmaStoreY* y_peer_desc = nullptr) { - - int warp_id = threadIdx.x / Traits1::WarpSize; - - // Construct the per-CTA x_barrier ONCE here so its target counter - // accumulates monotonically across all 3 mlp_fused_fwd calls - // (A1, B, A2). Each phase invokes mlp_fused_fwd which calls - // x_barrier.wait() inside its tile loop — they must share the - // same barrier object so target stays aligned with the shared - // device counter. Constructing the barrier inside mlp_fused_fwd - // would reset target=0 per call while the counter keeps growing, - // and the second mlp_fused_fwd would see counter > target on - // its very first poll and exit without actually waiting. - MlpFwdCtaBarrierT x_barrier(&dims.phase_counter[col], - runtime_nsplit); - - // UNIFIED PATH: the local MLP pass has been removed. The comm-side - // TileIterator now enumerates LOCAL experts (p=0) alongside remote ones, - // and the get warp stages local X from local symmetric memory (do_get - // is_local TMA path) exactly like remote tiles. So EVERY tile this column - // owns — local or remote — flows through the RemoteMlpTileIterator / - // src_staging→MLP→dst_staging pipeline below. local_iter is unused now. - // smem.comm.per_cta_tiles is this column's total tile count (stride - // grid_x); a column with no tiles returns immediately. - (void)local_iter; (void)tma_load_x; (void)tma_store_y; (void)num_tokens; - (void)num_m_tiles; (void)local_expert_start; - int total_remote = smem.comm.per_cta_tiles; - if (total_remote == 0) return; - - bool is_leader = (warp_id == 0 && (threadIdx.x % Traits1::WarpSize) == 0); - - // Build RemoteMlpTileIterator. The staging is a flat ring of - // L = MC · CommNumStages slots, where MC = (gridDim.x · gridDim.y) / NC, - // indexed by - // slot = (blockIdx.x + j · gridDim.x) mod L (MLP, ticket-based) - // slot = xc + (j mod CommNumStages) · MC - // (comm, NumStages cycle). - // Both sides resolve to the same slot for a given global tile T and - // cover the full L. No divisibility relationship is required between - // NC, NSplit and CommNumStages (runtime: NC must divide gridDim.x · - // gridDim.y so MC is integer). - // n_gemm = grid_x · runtime_nsplit — the GEMM/comm-active CTA count. - int n_gemm = grid_x * runtime_nsplit; - RemoteIter remote_iter; - remote_iter.init( - // remote_offsets in SMEM drives the embedded walker (null → no derivation). - smem.comm.remote_offsets, smem.comm.remote_counts, - experts_per_pe, num_pes, my_pe, gpus_per_node, - smem.comm.per_cta_tiles, - col, - n_gemm, - src_ready, src_consumed, - dst_ready, dst_consumed, - is_leader, - runtime_nsplit); - - // Remote MLP dimensions: TMA descriptors cover the full flat staging - // list = MC · CommNumStages · TileM rows. MC = gridDim.x / NC — MUST match - // the ring size in allocate_moe_buffers and the consumer's ring_len - // (RemoteMlpTileIterator mc = gridDim.x/NC); sizing on n_gemm here would - // leave the wrap-around slots [n_gemm/NC·CS, gridDim.x/NC·CS) outside the - // staging coverage (OOB) when gridDim.x/NC > n_gemm/NC. - int mc = gridDim.x / NC; - int remote_num_m_tiles = mc * CommNumStages; - - // NOTE: phase_counter is monotonic across phase A and phase B — - // the modulo-gridDim.y check in mlp_fused_fwd is invariant - // under any shared starting value. The host zeroes phase_counter - // before each kernel launch (see moe.cu); resetting it here would - // race across the gridDim.y CTAs in the column. - - // num_tokens = total staging rows = remote_num_m_tiles COMM tiles × CommTileM. - // CommTileM = Traits1::TileM (=GemmTileM) × SubTiles. num_m_tiles passed is the - // COMM tile count; mlp_fused_fwd scales it by SubTiles for the GEMM sub-tiles. - mlp_fused_fwd( - smem.mlp, remote_iter, - tma_load_x_remote, tma_load_b, tma_load_c, tma_store_z, - tma_load_z, tma_load_a, tma_store_y_remote, - dims, - remote_num_m_tiles * Traits1::TileM * SubTiles, remote_num_m_tiles, x_barrier, - /*local_expert_start=*/0, col, grid_x, split, runtime_nsplit, - /*y_peer_desc=*/y_peer_desc); -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "moe_sm90.cuh" +#else +#include "moe_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_bwd.cu b/liger_cute_kernels/csrc/core/src/moe/moe_bwd.cu index af92ccc5..69ff5227 100644 --- a/liger_cute_kernels/csrc/core/src/moe/moe_bwd.cu +++ b/liger_cute_kernels/csrc/core/src/moe/moe_bwd.cu @@ -1,2092 +1,9 @@ -// Fused MoE Backward: sort → fwd-dispatch (for X) → reverse-combine (for dY) -// → MLP_bwd + comm overlap → reverse-dispatch (for dX). -// Compiled with separable compilation. Do NOT link liger_comm_kernels library. -// -// Torch-free core (liger_cute): ported from LigerCommKernels' moe_bwd.cu. The -// device kernel + templated launcher are upstream-verbatim (namespace liger). -// Tensor-carrying backward calls are exported through TVM FFI and lower directly -// into MoeBwdArgs. -#define LIGER_CUTE_BUILDING 1 - -#include -#include // pulls in nvshmemi_device_state_d, required by the - // inline device helpers in . -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -// SM100 2SM (paired-CTA cluster) TMA descriptor factories (make_tma_copy_A_sm100 -// / make_tma_copy_B_sm100) and the SM100_TMA_2SM_LOAD copy atom used below to -// build pair-aware Phase-2 mlp3/mlp4 descriptors. Already pulled in -// transitively via moe_bwd.cuh -> mlp_bwd.cuh -> mlp3.cuh/mlp4.cuh; included -// directly here too since this TU calls these factories itself. -#include - -#include "liger_cute/moe.h" -#include "liger_cute/check.h" -#include "liger_cute/detail/status.h" - -#include "moe.cuh" -#include "moe_bwd.cuh" -#include "moe_symm_config.cuh" -#include "moe_launch.h" // MoeFwdArgs / MoeBwdArgs (shared with the tuner) -#include "moe_utils.cuh" -#include "moe_fwd_bwd_tune_configs.hpp" -#include "moe_dispatch_configs_sm90.cuh" -#include "moe_dispatch_configs_sm100.cuh" -#include "mlp_comms_bwd.cuh" -#include "tile_iterator_bwd.cuh" -#include "dispatch.cuh" -#include "dispatch_bwd.cuh" -#include "combine_bwd.cuh" -#include "tma_copy_atom.cuh" -#include "radix_sort.cuh" -#include "buffer_pool.cuh" // shim: liger:: aliases of liger_cute::detail infra - -using namespace cute; - -namespace liger { - -// 0 = build every dispatch family (developer fallback). CMake sets this to 90 -// for Hopper builds and 100 for Blackwell builds so normal native builds do not -// instantiate kernels for hardware that cannot run in that binary. #ifndef LIGER_CUTE_DISPATCH_COMPUTE #define LIGER_CUTE_DISPATCH_COMPUTE 0 #endif -// External symbol: moe_configure_symmetric (defined in moe.cu) is reused. -// MoeSymmConfig + get_symm_config() now come from the shared header (included -// below), so this TU can't drift from moe.cu's layout. - -// Global iteration counter — same role as the fwd counter in moe.cu. -// Lives in device memory (allocated from BufferPool under the key -// "moe_bwd_iter_counter"), atomicAdded at the end of each kernel by -// CTA 0 thread 0. See PeSync::advance_iteration_counter() and the -// detailed comment in mlp_comms.cuh for why the counter has to be -// device-resident: a host-side atomic would be baked into a captured -// CUDA graph as a constant, breaking the cross-PE wait-until-greater -// protocol on every replay. - -// ============================================================================ -// MoeBwdConfig — bundles bwd-side traits + comm/grid constants -// ============================================================================ -// -// All five MLP_bwd phase traits + grid params. Mirrors MoeFusedConfig. -// Trait choices match the values that mlp_bwd.cu landed on after the -// 2026-04-26 perf tuning round (see project_fused_bwd_tile_bump.md). - -// Trait defaults match mlp_bwd.cu's validated set (lines 28–33). MoE -// backward shares mlp_fused_bwd; any drift from those values has hit -// illegal-instruction / smem-overrun bugs in the past. -template < - int NSplit2_ = 4, // Phase 2 (mlp3/mlp4) virtual n-split (autotuning surface) - // SM100 uses Stage-3 Phase-1 traits to leave room for independent full-size - // X/dY TMA bounce buffers. SM90 dispatch rows provide their tuned values. - int TileM1_ = 128, - int TileN1_ = 128, - int TileK1_ = 64, - int Stages1_ = 3, - int EpiChunkN1_ = 32, // mlp1_act - int TileN2_ = 256, - int TileK2_ = 64, - int Stages2_ = 3, - int TileM2T_ = 128, // Phase 1b' (mlp2_t): TileM ∈ {64, 128} - // mlp3/mlp4 cooperative layout (#102) supports only (256,128) or (128,256) - // and needs Stages=2 to fit the 228 KiB smem cap — Stages=4 at TileN=256 - // alone is ~384 KiB. (256,128) M-split is the smallest-smem legal choice - // (~192 KiB; the binding floor is mlp1_act ≈ 216 KiB). - int TileM3_ = 256, // mlp3 cooperative: (256,128) M-split - int TileN3_ = 128, - int TileK3_ = 64, - int Stages3_ = 2, - int TileM4_ = 128, // mlp4 cooperative: swap of mlp3 → (128,256) - int TileN4_ = 256, - int TileK4_ = 64, - int Stages4_ = 2, - int TileM5_ = 128, - int TileN5_ = 256, - int TileK5_ = 64, - int Stages5_ = 3, - int EpiChunkN25_ = 64, // shared by mlp2_t and mlp5 (same smem shape) - int EpiChunkN34_ = 64, // shared by mlp3 and mlp4 (same smem shape) - typename Element_ = bfloat16_t, - int CommNumStages_ = 2, - int NC_ = 4, - // Phase-1 sub-batch grouping count: run SubBatch Phase-1 sub-batches - // (mlp1/2/5) before one Phase-2 (mlp3/4) over their combined K-window. - // 1 = original per-batch loop. The only structural requirement is a ring - // deeper than the group span (CommNumStages·mc_remote > SubBatch·grid_x); - // the grouped window wraps the ring exactly as the un-grouped loop does. - int SubBatch_ = 2, - // CommTileM_ decouples the COMM staging tile (ring slots, semaphores, - // tile_expert_ids granularity) from the Phase-1 GEMM tile (Traits1::TileM = - // mlp1a/mlp2t/mlp5 = the GemmTileM hyperparameter ∈ {64,128}). When - // Traits1::TileM < CommTileM_, Phase-1 steps through - // kSubTiles = CommTileM_/Traits1::TileM sub-tiles of GemmTileM rows per - // communication slot (mirrors forward). Default == TileM1_ → kSubTiles=1 - // → bit-identical to the old coupled path. - int CommTileM_ = 128, - int Compute_ = 90> -struct MoeBwdConfig { - using Element = Element_; - static constexpr int kActStoreTileM = 64; - using Traits1 = Mlp1Traits < - Element, TileM1_, TileN1_, TileK1_, Stages1_, EpiChunkN1_, - kActStoreTileM>; // Phase 1a - using Traits2T = Mlp2TTraits; // Phase 1b' - // Blackwell always uses the paired-CTA MLP3/MLP4 path. Hopper keeps the - // existing single-CTA traits and never instantiates SM100-only machinery. - // The existing TileK3_/Stages3_/EpiChunkN34_ fields remain the shared - // tuning surface for both weight-gradient kernels. - static constexpr bool kUsesTwoSm = Compute_ == 100; - static constexpr int kClusterM = kUsesTwoSm ? 2 : 1; - using Traits3 = cute::conditional_t< - kUsesTwoSm, - Mlp3Traits2Sm, - Mlp3Traits >; // Phase 2 mlp3 - using Traits4 = cute::conditional_t< - kUsesTwoSm, - Mlp4Traits2Sm, - Mlp4Traits >; // Phase 2 mlp4 - using Traits5 = Mlp5Traits ; // Phase 1d - - // No divisibility constraints between NSplit/NC/CommNumStages. - // MLP side is ticket-based (RemoteMlpTileIteratorBwd); comm side - // uses K = CommNumStages. Only runtime requirement is that NC - // evenly divides gridDim.x · gridDim.y so MC is integer. - - static constexpr int kNSplit2 = NSplit2_; - static constexpr int kNC = NC_; - static constexpr int kCommNumStages = CommNumStages_; - static constexpr int kSubBatch = SubBatch_; - // kTileM is the COMM staging tile (ring slots / semaphores / tile_expert_ids - // granularity), decoupled from the Phase-1 GEMM tile. kGemmTileM = - // Traits1::TileM (mlp1/2/5 GEMM tile); - // the mlp3/4 ratios (efkb_stride / bk↔Kblock / ring_kb) are computed from it. - static constexpr int kTileM = CommTileM_; - static constexpr int kGemmTileM = Traits1::TileM; // Phase-1 gemm tile (64/128) - static constexpr int kSubTiles = CommTileM_ / Traits1::TileM; - static constexpr int kCompute = Compute_; - static constexpr int kNumThreads = 384; - - static_assert(CommTileM_ % Traits1::TileM == 0, - "CommTileM must be an integer multiple of the Phase-1 GEMM TileM (Traits1::TileM)"); - static_assert(kSubTiles == 1 || kSubTiles == 2, - "CommTileM/GemmTileM must be 1 or 2"); - // mlp1a/mlp2t/mlp5 must share the same Phase-1 GEMM tile so their sub-tiling - // and the Z/dU/dV/dZ buffer rows stay aligned. - static_assert(Traits1::TileM == Traits2T::TileM && Traits1::TileM == Traits5::TileM, - "Phase-1 GEMM TileM must match across mlp1a / mlp2t / mlp5"); - - // Compile-time smem budget. H100 dynamic-smem cap is 228 KiB. The bwd - // union sits very close to that ceiling (Mlp3/Mlp4 with TileK=128 are - // 224 KiB on their own, and pipe storage pushes the total a bit - // further), so we use the full hardware cap here. The runtime-sized - // remote_offsets region appended by the launcher is ≤ a few hundred - // bytes for autotuner-scale E and num_pes; relying on the kernel's - // cudaFuncSetAttribute call to surface any overflow. - static constexpr size_t kSmemBytes = - sizeof(MoeBwdSmem); - static constexpr size_t kSmemBudget = 228 * 1024; - static_assert(kSmemBytes <= kSmemBudget, - "MoE fused bwd kernel smem footprint exceeds H100 228 KiB cap — " - "reduce TileN/TileK/Stages of one of the bwd phase traits"); - - using LocalIter = LocalMlpTileIteratorBwd; - using RemoteIter = RemoteMlpTileIteratorBwd< - Element, kCommNumStages, kNC, kTileM>; - using FusedIter = FusedMlpTileIteratorBwd< - Element, kCommNumStages, kNC, kTileM>; -}; - -// ============================================================================ -// MoeBwdParams — input/output pointers + sort buffers + comm + dims -// ============================================================================ - -template -struct MoeBwdParams { - using Element = typename Config::Element; - - // Router outputs (caller-provided). - const int* expert_indices; // [T, K] - const Element* expert_weights; // [T, K] - - // Sort buffers. - SortBuffers sort; - - // Inputs. - // Note: there is no `tokens` (X) field — BWD reads token data only - // via `x_sorted` (FWD's expert-sorted view from its symmetric stack - // entry). Removed 2026-05-18 after confirming p.tokens was set in the - // host wrapper and never read by the kernel or any helper. - const Element* dY; // [T, D] (regular device) - // Forward MoE output Y_fwd. Caller is responsible for providing this - // — either retained from the most recent forward pass or recomputed - // under gradient checkpointing. Read by combine_tokens_bwd to compute - // dW[t,k] = Σ_d dO[t,d] * Y_fwd[slot(t,k), d]. Layout: [total_slots, D] - // (one row per routed (token, expert-slot) pair, in the same - // expert-sorted order produced by the forward dispatch). - const Element* output_buffer; - - // Outputs (caller-allocated). - Element* dX_out; // [T, D] (regular device) - Element* dB_out; // [epp, I, D] - Element* dC_out; // [epp, I, D] - Element* dA_out; // [epp, D, I] - Element* dW_out; // [T, K] - - // Symmetric staging buffers (allocated by allocate_moe_bwd_buffers). - Element* x_sorted; // [total_slots, D] — fwd dispatch dst - Element* dy_sorted; // [total_slots, D] — combine_bwd dst - Element* dx_sorted; // [total_slots, D] — MLP_bwd dX dst - - // Comm (forward, for sort/dispatch/pe_sync) and bwd (for X+dY gets, dX puts). - CommBuffers comm; - CommBuffersBwd comm_bwd; - - // Dimensions. - int num_tokens; - int hidden_dim; - int intermediate_dim; - int top_k; - int total_slots; - int num_m_tiles; // sort input tiles = T / TileM - - // Output stride for dB/dC/dA per-expert. - int experts_per_pe; -}; - -// ============================================================================ -// MoeBwdBuffers — owned device + symmetric allocations -// ============================================================================ - -template -struct MoeBwdBuffers { - using Element = typename Config::Element; - - // Symmetric (NVSHMEM). - Element* x_sorted; - Element* dy_sorted; - Element* dx_sorted; - - // Symmetric offsets. - int* expert_offsets_sym; // for sort - - // Device-only. - int* cta_counter; - int* phase_counter; // mlp_bwd uses this - int* barrier_counter; // mlp_global_barrier - Element* z_buf; // [scratch_rows, I] - Element* du_buf; // [scratch_rows, I] -> dU after silu_bwd - Element* dv_buf; // [scratch_rows, I] - Element* dz_buf; // [scratch_rows, I] - int scratch_rows; // max(local padded rows, remote ring rows) - int* expert_for_k_block; // num_k_blocks; populated post-sort - - // Phase 2 per-expert K-range scratch. Filled in bwd preamble (once - // per pass) by mlp_bwd_phase_2_build_ranges; shared between mlp3 and - // mlp4 (they share TileK). Sized experts_per_pe. - int* expert_k_starts; - int* expert_k_ends; - - // Comm staging (X + dY gets, dX puts). - Element* x_staging; - Element* dy_staging; - Element* dx_staging; - - int* x_src_ready; - int* x_src_consumed; - int* dy_src_ready; - int* dy_src_consumed; - int* dst_ready; - int* dst_consumed; - // Two per-staging-slot expert-id arrays, one per get pipe. Comm warps - // write the X-side array during the X getmem and the dY-side array - // during the dY getmem. Phase 2 mlp4 reads tile_expert_ids_x; Phase 2 - // mlp3 reads tile_expert_ids_dy. See CommBuffersBwd docstring for why. - int* tile_expert_ids_x; - int* tile_expert_ids_dy; - int* tile_valid_rows_x; - - // Sort-internal scratch. - int* tile_expert_counts; - int* sorted_token_ids; - int* token_expert_slots; - int* cta_done; - int* cta_sums; -}; - -template -static MoeBwdBuffers allocate_moe_bwd_buffers( - BufferPool& pool, cudaStream_t stream, - int grid_x, int num_blocks, - int hidden_dim, int intermediate_dim, int num_tokens, int top_k, - int experts_per_pe, int total_slots, int num_m_tiles) { - - using Element = typename Config::Element; - constexpr int kTileM = Config::kTileM; - // kCommNumStages / kNSplit / kNC are not needed here: the comm staging is - // sized at the config-independent UB (kCommStagesUB · kTileMUB) below. - - MoeBwdBuffers b; - - auto& scfg = get_symm_config(); - LIGER_CHECK(scfg.initialized, - "Call moe_configure_symmetric before moe_bwd_fwd_bf16"); - - size_t symm_slot_bytes = - static_cast(scfg.max_total_slots) * scfg.hidden_dim * sizeof(Element); - - // ── Symmetric ──────────────────────────────────────── - // x_sorted comes from the FWD pass — the caller passes its symm - // pointer in the `x_sorted` torch tensor (FWD's pool entry under - // key "moe_x_sorted"). BWD reads it directly and never allocates - // or writes to its own copy. - b.dy_sorted = static_cast(pool.get_symmetric("moe_bwd_dy_sorted", symm_slot_bytes)); - b.dx_sorted = static_cast(pool.get_symmetric("moe_bwd_dx_sorted", symm_slot_bytes)); - - // dy_sorted MUST be zeroed: combine_tokens_bwd only writes routed slots - // (top_k per token). Padding slots stay 0, contributing nothing to MLP_bwd. - cudaMemsetAsync(b.dy_sorted, 0, symm_slot_bytes, stream); - - // all_expert_offsets is no longer allocated here — the caller passes - // the broadcasted offsets in via the BWD's `expert_offsets` tensor - // (populated by FWD's broadcast_expert_offsets and returned as part - // of its intermediates). - b.expert_offsets_sym = static_cast( - pool.get_symmetric("moe_sort_offsets", - (scfg.max_num_experts + 1) * sizeof(int))); - - // ── Device-only ────────────────────────────────────── - b.cta_counter = static_cast(pool.get_device("moe_bwd_cta_counter", sizeof(int))); - b.phase_counter = static_cast(pool.get_device("moe_bwd_phase_counter", grid_x * sizeof(int))); - b.barrier_counter = static_cast(pool.get_device("moe_bwd_barrier", sizeof(int))); - cudaMemsetAsync(b.cta_counter, 0, sizeof(int), stream); - cudaMemsetAsync(b.phase_counter, 0, grid_x * sizeof(int), stream); - cudaMemsetAsync(b.barrier_counter, 0, sizeof(int), stream); - - // Local Phase 1 indexes scratch by padded token row, while the remote pass - // indexes the same buffers by staging-ring slot. A deep comm ring can exceed - // the local padded extent (for example, 33 MC columns * CS8 * TileM128 = - // 33792 rows versus 17408 local rows for Mixtral-8x7B T=8K). Size for both - // coordinate spaces or a valid remote slot writes past z/du/dv/dz. - const int remote_m_tiles = - (num_blocks / Config::kNC) * Config::kCommNumStages; - const size_t remote_rows = - static_cast(remote_m_tiles) * kTileM; - const size_t scratch_rows = - std::max(static_cast(total_slots), remote_rows); - b.scratch_rows = static_cast(scratch_rows); - size_t intermediate_bytes = - scratch_rows * intermediate_dim * sizeof(Element); - b.z_buf = static_cast(pool.get_device("moe_bwd_z_buf", intermediate_bytes)); - b.du_buf = static_cast(pool.get_device("moe_bwd_du_buf", intermediate_bytes)); - b.dv_buf = static_cast(pool.get_device("moe_bwd_dv_buf", intermediate_bytes)); - b.dz_buf = static_cast(pool.get_device("moe_bwd_dz_buf", intermediate_bytes)); - - // expert_for_k_block: one int per K-block. The K-block / m-tile size - // ratio depends on Traits3::TileK vs Traits1::TileM; see - // build_expert_for_k_block which uses both dims directly. - // Size by scfg.max_total_slots, not per-call total_slots — the buffer - // pool caches by name and won't resize on later calls. Otherwise the - // first (smallest) config locks in a tiny allocation and a later - // larger config (e.g. E=64, K=8) reads OOB past the cached size. - int max_num_k_blocks = scfg.max_total_slots / Config::Traits3::TileK; - b.expert_for_k_block = static_cast( - pool.get_device("moe_bwd_expert_for_k_block", max_num_k_blocks * sizeof(int))); - - // Phase 2 per-expert K-range scratch — sized by scfg.max_num_experts - // (the pool caches by name, so we have to size to the worst case - // across configs; same rationale as expert_for_k_block above). - // Shared between mlp3 and mlp4 (they share TileK). - size_t expert_range_bytes = scfg.max_num_experts * sizeof(int); - b.expert_k_starts = static_cast( - pool.get_device("moe_bwd_expert_k_starts", expert_range_bytes)); - b.expert_k_ends = static_cast( - pool.get_device("moe_bwd_expert_k_ends", expert_range_bytes)); - - // ── Comm staging ───────────────────────────────────── - // Flat staging list — L = MC · kCommNumStages slots, where - // MC = (grid_x · NSplit) / NC. Comm strides MC through NumStages - // stages; MLP is ticket-based and keys on the absolute slot index - // modulo L. With NC = NSplit (legacy config), MC = grid_x and this - // collapses to grid_x · kCommNumStages — unchanged behavior. The live - // length MC·kCommNumStages (MC = (grid_x·NSplit)/NC) is bounded by - // comm_l_ub below; the buffers are sized at that UB, not the live value. - // Staging must be SYMMETRIC: the comm path issues nvshmemx_getmem/putmem - // whose LOCAL side (x/dy/dx staging) is the NIC-registered buffer. - // Device-only memory faults with IBV_WC_LOC_PROT_ERR over IB (works on - // NVLink either way, which masked it single-host). Size at an - // NC-independent upper bound (mc <= sm_count) so the symmetric - // allocation never has to grow across NC variants — it aborts on grow; - // larger-than-needed is fine. Mirrors the fwd moe.cu comm_l_ub sizing. - int sm_count = 0; - { - int dev = 0; - cudaGetDevice(&dev); - cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev); - } - // Worst-case sizing across the whole tuner sweep: the symmetric/device - // pools cache by key and ABORT on grow, so the FIRST config to allocate a - // key locks its size. A small-TileM (64) or small-CommStages config must - // not lock in a buffer that a later TileM=256 / deeper-pipe config can't - // grow into. Size every comm buffer at the centralized TileM/stage maxima; - // the kernel indexes by the per-call - // (mc·kCommNumStages, kTileM) ≤ these UBs, so over-allocation is safe. - // Mirrors moe.cu's MoeSymmConfig.max_tile_m / max_comm_stages fix. - constexpr int kCommStagesUB = kMaxMoeCommStages; - constexpr int kTileMUB = kMaxMoeCommTileM; - int comm_l_ub = sm_count * kCommStagesUB; - size_t stage_slots_bytes = - (size_t)comm_l_ub * kTileMUB * scfg.hidden_dim * sizeof(Element); - b.x_staging = static_cast(pool.get_symmetric("moe_bwd_x_staging", stage_slots_bytes)); - b.dy_staging = static_cast(pool.get_symmetric("moe_bwd_dy_staging", stage_slots_bytes)); - b.dx_staging = static_cast(pool.get_symmetric("moe_bwd_dx_staging", stage_slots_bytes)); - - // X / dY pipe signals: per-slot, [L = MC * NumStages] ints. Comm - // producer indexes by `+ xc` (stride MC); MLP consumer is - // ticket-based and walks slot = (m_base + idx·gridDim.x) mod L. - // Allocate at the UB (comm_l_ub ≥ MC·kCommNumStages); the kernel only - // touches [0, MC·kCommNumStages). Per-call resets (below) zero the live prefix. - size_t tile_sig_bytes = (size_t)comm_l_ub * sizeof(int); - b.x_src_ready = static_cast(pool.get_device("moe_bwd_x_src_ready", tile_sig_bytes)); - b.x_src_consumed = static_cast(pool.get_device("moe_bwd_x_src_consumed", tile_sig_bytes)); - b.dy_src_ready = static_cast(pool.get_device("moe_bwd_dy_src_ready", tile_sig_bytes)); - b.dy_src_consumed= static_cast(pool.get_device("moe_bwd_dy_src_consumed",tile_sig_bytes)); - cudaMemsetAsync(b.x_src_ready, 0, tile_sig_bytes, stream); - cudaMemsetAsync(b.x_src_consumed, 0, tile_sig_bytes, stream); - cudaMemsetAsync(b.dy_src_ready, 0, tile_sig_bytes, stream); - cudaMemsetAsync(b.dy_src_consumed,0, tile_sig_bytes, stream); - - // dX pipe signals: per-slot, [L = MC * NumStages] ints. UB-sized (see above). - size_t tile_pipe_bytes = (size_t)comm_l_ub * sizeof(int); - b.dst_ready = static_cast(pool.get_device("moe_bwd_dst_ready", tile_pipe_bytes)); - b.dst_consumed = static_cast(pool.get_device("moe_bwd_dst_consumed", tile_pipe_bytes)); - cudaMemsetAsync(b.dst_ready, 0, tile_pipe_bytes, stream); - cudaMemsetAsync(b.dst_consumed, 0, tile_pipe_bytes, stream); - - size_t eid_bytes = (size_t)comm_l_ub * sizeof(int); // UB-sized (see above) - // One array per get pipe (see MoeBwdBuffers / CommBuffersBwd docstrings). - // Same shape as the staging signal arrays. - b.tile_expert_ids_x = static_cast( - pool.get_device("moe_bwd_tile_expert_ids_x", eid_bytes)); - b.tile_expert_ids_dy = static_cast( - pool.get_device("moe_bwd_tile_expert_ids_dy", eid_bytes)); - b.tile_valid_rows_x = static_cast( - pool.get_device("moe_bwd_tile_valid_rows_x", eid_bytes)); - - // ── Sort scratch (sized to scfg max for stability across calls) ── - int max_sorted_slots = scfg.max_total_slots; - int max_m_tiles = scfg.max_total_slots / kTileM; - int max_tile_counts = max_m_tiles * scfg.max_num_experts; - - b.tile_expert_counts = static_cast( - pool.get_device("moe_bwd_sort_tile_counts", max_tile_counts * sizeof(int))); - b.sorted_token_ids = static_cast( - pool.get_device("moe_bwd_sort_sorted_ids", max_sorted_slots * sizeof(int))); - b.token_expert_slots = static_cast( - pool.get_device("moe_bwd_sort_token_slots", max_sorted_slots * sizeof(int))); - b.cta_done = static_cast( - pool.get_device("moe_bwd_sort_cta_done", num_blocks * sizeof(int))); - b.cta_sums = static_cast( - pool.get_device("moe_bwd_sort_cta_sums", - num_blocks * scfg.max_num_experts * sizeof(int))); - - cudaMemsetAsync(b.cta_done, 0, num_blocks * sizeof(int), stream); - cudaMemsetAsync(b.sorted_token_ids, 0xff, max_sorted_slots * sizeof(int), stream); - - return b; -} - -// ============================================================================ -// build_expert_for_k_block — small helper kernel section -// ============================================================================ -// -// Post-sort: derive expert_for_k_block[k] = tile_expert_ids[(k * TileK) / TileM]. -// Works for both directions: -// TileK > TileM (K-block spans multiple m-tiles): uses the first m-tile's expert. -// TileK < TileM (m-tile spans multiple K-blocks): all K-blocks in a m-tile -// resolve to the same tile_expert_ids slot via integer division. -// Assumes per-expert ranges are TileK-aligned, OR that any K-block straddling -// an expert boundary has only padding tokens in its second m-tile (ensured -// by zeroed dy_sorted + zero-filled X-padding from forward dispatch). - -__device__ __forceinline__ void build_expert_for_k_block( - int* expert_for_k_block, - const int* tile_expert_ids, - int num_k_blocks, - int num_actual_k_blocks, - int tile_k, - int tile_m, - int num_threads) { - // For k < num_actual_k_blocks, copy the expert from the sort output. - // For padding entries (k beyond expert_offsets[num_experts] / TileK), - // write a sentinel that the Phase 3/4 bounds check will reject. Without - // this, stale memory in tile_expert_ids[k] could pass the bounds check - // and cause Phase 4 to load garbage X/dU into the dB/dC accumulators. - for (int k = threadIdx.x; k < num_k_blocks; k += num_threads) { - expert_for_k_block[k] = (k < num_actual_k_blocks) - ? tile_expert_ids[(k * tile_k) / tile_m] - : -1; - } -} - -// ============================================================================ -// Fused kernel: sort → fwd-dispatch → combine_bwd → MLP_bwd+comm → dispatch_bwd -// ============================================================================ - -template -__global__ void __launch_bounds__(Config::kNumThreads) -moe_bwd_kernel( - // Local TMA - __grid_constant__ TmaLoadX const tma_load_x, - __grid_constant__ TmaLoadW1 const tma_load_b_fwd, - __grid_constant__ TmaLoadW1 const tma_load_c_fwd, - __grid_constant__ TmaStoreZ const tma_store_z, - __grid_constant__ TmaStoreU const tma_store_u, - __grid_constant__ TmaStoreV const tma_store_v, - __grid_constant__ TmaLoadDY const tma_load_dy, - __grid_constant__ TmaLoadW2T const tma_load_a_col, - __grid_constant__ TmaStoreDZ const tma_store_dz, - __grid_constant__ TmaLoadDU const tma_load_du, - __grid_constant__ TmaLoadDV const tma_load_dv, - __grid_constant__ TmaLoadB const tma_load_b_col, - __grid_constant__ TmaLoadC const tma_load_c_col, - __grid_constant__ TmaStoreDX const tma_store_dx, - __grid_constant__ TmaLoadDYT3 const tma_load_dyt3, - __grid_constant__ TmaLoadZT3 const tma_load_zt3, - __grid_constant__ TmaReduceDA const tma_reduce_da, - __grid_constant__ TmaLoadXT4 const tma_load_xt4, - __grid_constant__ TmaLoaddUT4 const tma_load_dut4, - __grid_constant__ TmaLoaddVT4 const tma_load_dvt4, - __grid_constant__ TmaReduceDB const tma_reduce_db, - __grid_constant__ TmaReduceDC const tma_reduce_dc, - // Remote TMA - __grid_constant__ TmaLoadXR const tma_load_x_remote, - __grid_constant__ TmaLoadDYR const tma_load_dy_remote, - __grid_constant__ TmaStoreDXR const tma_store_dx_remote, - // Remote Phase 2 TMA (staging-backed dY/X for dA/dB/dC) - __grid_constant__ TmaLoadDYT3R const tma_load_dyt3_remote, - __grid_constant__ TmaLoadXT4R const tma_load_xt4_remote, - // Dims + bufs - __grid_constant__ MlpBwdDims const mlp_dims, - __grid_constant__ MlpBwdBufs const mlp_bufs, - __grid_constant__ MoeBwdParams const p, - __grid_constant__ int const static_nsplit, - __grid_constant__ GetTmaDescsBwd const get_descs_bwd) { - - using Element = typename Config::Element; - using Traits1 = typename Config::Traits1; - using Traits2T = typename Config::Traits2T; - using Traits3 = typename Config::Traits3; - using Traits4 = typename Config::Traits4; - using Traits5 = typename Config::Traits5; - using FusedIter = typename Config::FusedIter; - - constexpr int kTileM = Config::kTileM; - constexpr int kNumThreads = Config::kNumThreads; - constexpr int kCommNumStages = Config::kCommNumStages; - constexpr int kNC = Config::kNC; - constexpr int kGetBoxRowsBwd = Config::kCompute == 100 - ? liger::kGetBoxRowsBwdSm100 - : liger::kGetBoxRowsBwdSm90; - - // Target-based cross-CTA barrier — single instance per kernel call. - SyncThreadsCtaCounterBarrier pe_barrier(p.comm.pe_sync.cta_counter, - (int)gridDim.x * (int)gridDim.y); - - // ── Phase 0: Sort — SKIPPED. - // Caller provides token_expert_slots, tile_expert_ids, and - // expert_offsets from the most recent forward call (these come out of - // moe_fused_fwd_bf16, populated by the forward - // kernel's sort_tokens). We trust those instead of re-sorting. - - // ── Phase 1: Broadcast expert offsets — SKIPPED. - // Forward returned the broadcasted offsets in the `expert_offsets` - // tensor we forward verbatim to p.comm.all_expert_offsets, and - // forward ended with nvshmem_team_sync(team), so every PE's writes - // are visible here. No bwd-side re-broadcast needed. - - // ── Phase 2a: Build expert_for_k_block (post-sort, before MLP) ─── - // Operates on caller-provided tile_expert_ids (forward's sort output, - // threaded through p.sort.tile_expert_ids). BWD_TileM is constrained - // to match FWD_TileM (enforced in moe_bwd_fwd_bf16_auto), so the - // FWD-populated array indexes correctly at BWD's tile granularity. - { - int num_k_blocks = p.total_slots / Traits3::TileK; - // Real data ends at expert_offsets[num_experts]; everything beyond - // is buffer-pool padding from the symmetric x_sorted allocation. - int total_real_rows = p.sort.expert_offsets[p.comm.num_experts]; - int num_actual_k_blocks = total_real_rows / Traits3::TileK; - // Both _mlp4 and _mlp3 pointers alias buf.expert_for_k_block in the - // local pass (set above), so writing through either populates the - // shared post-sort array. - build_expert_for_k_block( - const_cast(mlp_bufs.expert_for_k_block_mlp4), - p.sort.tile_expert_ids, - num_k_blocks, num_actual_k_blocks, - Traits3::TileK, kTileM, kNumThreads); - } - __syncthreads(); - - // ── Phase 2b: Forward dispatch — SKIPPED. - // p.x_sorted points to the forward's x_sorted (already dispatched). - // MLP TMAs read from there directly; nothing to do here. - - // ── Phase 3: Reverse-combine — dY → dy_sorted (+ dW side-output) ─ - // Caller-provided output_buffer is Y_fwd in expert-sorted layout - // [total_slots, D] (must match the expert-sort produced by the - // forward dispatch). combine_tokens_bwd computes: - // dW[t,k] = Σ_d dY[t,d] * Y_fwd[slot(t,k), d] → p.dW_out - // dy_sorted[slot,d]= dY[t,d] * w[t,k] → fed to mlp_fused_bwd - combine_tokens_bwd( - p.dY, - /*expert_out=*/p.output_buffer, - p.expert_weights, - p.sort.token_expert_slots, - p.dW_out, - p.dy_sorted, - p.hidden_dim, - p.num_tokens, - p.top_k); - __threadfence_system(); - __syncthreads(); - - p.comm.pe_sync.barrier_all(pe_barrier); - - // ── Phase 4: MLP_bwd with comm overlap ─────────────────────────── - { - extern __shared__ char raw_smem[]; - auto& smem = *reinterpret_cast*>(raw_smem); - - size_t offsets_region = sizeof(MoeBwdSmem); - offsets_region = (offsets_region + sizeof(int) - 1) & ~(sizeof(int) - 1); - size_t offsets_count = - (size_t)p.comm_bwd.num_pes * (p.comm_bwd.experts_per_pe + 1); - if (threadIdx.x == 0) { - smem.comm.remote_offsets = - reinterpret_cast(raw_smem + offsets_region); - smem.comm.remote_counts = - smem.comm.remote_offsets + offsets_count; - } - __syncthreads(); - - // Bwd TMA-GET bounce: 128-aligned region after offsets and counts. - // Mirrors the host smem sizing (align128(smem_mlp) + bounce). Only - // dereferenced when get_descs_bwd.enabled; harmless otherwise. - size_t bounce_off = offsets_region - + (offsets_count - + (size_t)p.comm_bwd.num_pes * p.comm_bwd.experts_per_pe) - * sizeof(int); - bounce_off = (bounce_off + 127) & ~((size_t)127); - char* get_bounce = raw_smem + bounce_off; - - // No sentinel write here: nvshmem_comm_prologue_bwd's warp-2 leader - // is guaranteed to write smem.comm.total_tiles (either to 0 in the - // early-exit case or to comm_count in the normal case) before the - // __syncthreads below. A `smem.comm.total_tiles = -1` from - // threadIdx.x == 0 racing with the prologue's warp-2 write is what - // caused the BWD-TM=64 + 2 PE deadlock: when thread 0's sentinel - // write landed AFTER warp 2's write, total_tiles ended up as -1, - // comm skipped nvshmem_comm_main_bwd, the comm pipes (x_src_ready - // etc.) never got a producer_release, and the MLP remote pass' - // iter.acquire_src spun forever. - nvshmem_comm_prologue_bwd(smem.comm, p.comm_bwd, static_nsplit); - __syncthreads(); - - int flat_id = (int)blockIdx.x; - int runtime_nsplit = static_nsplit; - int n_gemm = p.comm_bwd.n_gemm; - int grid_x = n_gemm / runtime_nsplit; - int col = flat_id / runtime_nsplit; - if (threadIdx.x == 0) { - int total_m_tiles = smem.comm.global_total; - runtime_nsplit = select_runtime_nsplit_bwd_from_tiles( - total_m_tiles, - mlp_dims.num_n_tiles_1, - mlp_dims.num_n_tiles_2t, - mlp_dims.num_n_tiles_5, - mlp_dims.num_k_tiles_1, - mlp_dims.num_k_tiles_2t, - mlp_dims.num_k_tiles_5, - p.comm_bwd.experts_per_pe, - (int)gridDim.x, - static_nsplit); - n_gemm = ((int)gridDim.x / runtime_nsplit) * runtime_nsplit; - grid_x = n_gemm / runtime_nsplit; - col = flat_id / runtime_nsplit; - int mlp_count = 0; - if (flat_id < n_gemm && col < total_m_tiles) - mlp_count = (total_m_tiles - 1 - col) / grid_x + 1; - smem.comm.per_cta_tiles = mlp_count; - smem.comm.runtime_nsplit = runtime_nsplit; - smem.comm.runtime_n_gemm = n_gemm; - smem.comm.runtime_grid_x = grid_x; - } - __syncthreads(); - runtime_nsplit = smem.comm.runtime_nsplit; - n_gemm = smem.comm.runtime_n_gemm; - grid_x = smem.comm.runtime_grid_x; - col = flat_id / runtime_nsplit; - int split = flat_id % runtime_nsplit; - bool gemm_active = flat_id < n_gemm; - - int warp_id = threadIdx.x / 32; - bool is_comm_warp = (warp_id == kCommBwdGetWarp0 || warp_id == kCommBwdPutWarp); - - if (is_comm_warp) { - if (smem.comm.total_tiles > 0) { - nvshmem_comm_main_bwd< - Element, kTileM, kCommNumStages, kNC, - kGetBoxRowsBwd>( - smem.comm, p.comm_bwd, &get_descs_bwd, get_bounce); - nvshmem_quiet(); - } - } else { - // MLP warps. The fused iter is initialized for the staging ring - // inside moe_fused_bwd (init_remote); every tile — local or remote — - // flows through that staging path, so there is no local sub-iter to - // seed here. Gap CTAs skip Phase 1; moe_fused_bwd gates it on - // gemm_active. - FusedIter fused_iter; - MlpBwdDims runtime_mlp_dims = mlp_dims; - runtime_mlp_dims.grid_x = grid_x; - runtime_mlp_dims.n_gemm = n_gemm; - - moe_fused_bwd( - smem, fused_iter, - tma_load_x, tma_load_b_fwd, tma_load_c_fwd, - tma_store_z, tma_store_u, tma_store_v, - tma_load_dy, tma_load_a_col, tma_store_dz, - tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, - tma_load_dyt3, tma_load_zt3, tma_reduce_da, - tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, - tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, - tma_load_dyt3_remote, tma_load_xt4_remote, - p.comm_bwd.x_src_ready, p.comm_bwd.x_src_consumed, - p.comm_bwd.dy_src_ready, p.comm_bwd.dy_src_consumed, - p.comm_bwd.dst_ready, p.comm_bwd.dst_consumed, - p.comm_bwd.tile_expert_ids_x, - p.comm_bwd.tile_expert_ids_dy, - p.comm_bwd.tile_valid_rows_x, - runtime_mlp_dims, mlp_bufs, - col, grid_x, split, runtime_nsplit, gemm_active); - } - - __threadfence(); - __syncthreads(); - p.comm.pe_sync.barrier_all(pe_barrier); - } - - // ── Phase 5: Reverse-dispatch — dx_sorted → dX_out (gather-sum) ── - dispatch_tokens_bwd( - p.dx_sorted, - p.sort.token_expert_slots, - p.dX_out, - p.hidden_dim, - p.num_tokens, - p.top_k); - -} - -// ============================================================================ -// GET-via-TMA descriptor builder for the bwd comm path (raw CUtensorMap). -// ============================================================================ -// -// Box = [kGetBoxRowsBwd, kGetKChunk] over a row-major [n_rows, hidden_dim] -// bf16 tensor (SWIZZLE_NONE keeps global bytes row-major so the MLP's -// swizzled reader sees correct data). Mirrors moe.cu's make_get_tma_map but -// with the bwd's small box height. Returns false on encode failure → caller -// disables the bwd TMA-get path (falls back to getmem/int4). -static inline bool make_get_tma_map_bwd(CUtensorMap& m, void* base, - int hidden_dim, int n_rows, int box_rows) { - uint64_t gdim[2] = { (uint64_t)hidden_dim, (uint64_t)n_rows }; - uint64_t gstride[1] = { (uint64_t)hidden_dim * sizeof(bfloat16_t) }; - uint32_t bdim[2] = { (uint32_t)liger::kGetKChunk, (uint32_t)box_rows }; - uint32_t estride[2] = { 1, 1 }; - CUresult r = cuTensorMapEncodeTiled( - &m, CU_TENSOR_MAP_DATA_TYPE_UINT16, 2, base, gdim, gstride, bdim, estride, - CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, - CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - return r == CUDA_SUCCESS; -} - -// ============================================================================ -// Host wrapper -// ============================================================================ - -// MoeBwdArgs (the POD launch bundle) is defined in moe_launch.h, shared with the -// tuner (mirrors MoeFwdArgs). dB/dC/dA/dW are gradient accumulation targets — -// the launcher zeroes them before the kernel. - -template > -void -moe_bwd_fwd_bf16(const MoeBwdArgs& a, int static_nsplit) { - - const int num_experts = a.num_experts; - const int top_k = a.top_k; - const nvshmem_team_t team = a.team; - cudaStream_t stream = a.stream; - cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; - if (cudaError_t e = cudaStreamIsCapturing(stream, &capture_status); - e != cudaSuccess) - LIGER_FAIL_CUDA("moe_bwd: cudaStreamIsCapturing failed: ", - cudaGetErrorString(e)); - if (capture_status == cudaStreamCaptureStatusNone) { - if (cudaError_t e = cudaStreamSynchronize(stream); e != cudaSuccess) - LIGER_FAIL_CUDA( - "moe_bwd: synchronizing the forward-to-backward NVSHMEM " - "handoff failed: ", cudaGetErrorString(e)); - } - - using Element = typename Config::Element; - using Traits1 = typename Config::Traits1; - using Traits2T = typename Config::Traits2T; - using Traits3 = typename Config::Traits3; - using Traits4 = typename Config::Traits4; - using Traits5 = typename Config::Traits5; - using TmaLoadAtom = TmaLoadAtomForCompute; - using TmaStoreAtom = TmaStoreAtomForCompute; - using TmaReduceAddAtom = TmaReduceAddAtomForCompute; - constexpr int kTileM = Config::kTileM; - constexpr int kNumThreads = Config::kNumThreads; - - // Iteration counter lives on the device; CTA 0 thread 0 atomicAdds it at - // the end of the kernel via PeSync::advance_iteration_counter(). - int num_tokens = a.num_tokens; - int hidden_dim = a.hidden_dim; - int experts_per_pe = a.experts_per_pe; - int intermediate_dim = a.intermediate_dim; - - // max_total_slots uses kTileM because forward dispatch always pads each - // expert range to its own TileM. moe_bwd_fwd_bf16_auto enforces - // BWD_TileM == FWD_TileM, so kTileM is the right alignment here. - int max_total_slots = num_tokens * top_k + num_experts * kTileM; - int num_m_tiles_in = (num_tokens + kTileM - 1) / kTileM; - - int num_pes = nvshmem_team_n_pes(team); - int my_pe = nvshmem_team_my_pe(team); - int local_gpus = 0; - cudaGetDeviceCount(&local_gpus); - int gpus_per_node = - (local_gpus > 0 && num_pes % local_gpus == 0) ? local_gpus : num_pes; - - int num_sms; - cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, a.device); - // Flat 1-D launch: use the largest NC-complete SM subset. Phase 1 + comm partition over the - // GEMM-active subset n_gemm = floor_NS (largest multiple of static_nsplit ≤ - // num_blocks); Phase 2 (mlp3/mlp4) grid-strides over all num_blocks CTAs - // (no divisibility needed — cell stride = num_blocks); the gap CTAs - // [n_gemm, num_blocks) fall through Phase 1/comm but still run Phase 2 - // and every barrier. NC-complete launch geometry is required because comm - // ring lengths, signal arrays, and TMA descriptors all use gridDim.x / NC. - // ClusterM divides NC for every supported configuration, so this also - // guarantees complete SM100 CTA pairs. - static_assert(Config::kNC % Config::kClusterM == 0); - int num_blocks = num_sms - num_sms % Config::kNC; - int n_gemm = num_blocks - num_blocks % static_nsplit; - int grid_x = n_gemm / static_nsplit; - int max_runtime_grid_x = (num_blocks + 1) / 2; // min runtime NS candidate = 2 - dim3 grid(num_blocks); - LIGER_CHECK(top_k <= kCombineMaxTopK, - "top_k (", top_k, ") exceeds kCombineMaxTopK (", kCombineMaxTopK, - "). Bump the constant in combine.cuh."); - - // stream comes from a.stream (the binding passes the current torch stream). - auto& pool = global_buffer_pool(); - - auto buf = allocate_moe_bwd_buffers(pool, stream, max_runtime_grid_x, num_blocks, - hidden_dim, intermediate_dim, num_tokens, top_k, - experts_per_pe, max_total_slots, num_m_tiles_in); - - // Caller-owned gradient outputs (binding-allocated). dB/dC/dA/dW are - // accumulation targets — zero them on the stream (upstream used torch::zeros); - // dX is fully overwritten so it needs no pre-zero. - constexpr std::size_t kElemBytes = 2; // sizeof(bf16) - cudaMemsetAsync(a.dB, 0, - static_cast(experts_per_pe) * intermediate_dim * hidden_dim * kElemBytes, stream); - cudaMemsetAsync(a.dC, 0, - static_cast(experts_per_pe) * intermediate_dim * hidden_dim * kElemBytes, stream); - cudaMemsetAsync(a.dA, 0, - static_cast(experts_per_pe) * hidden_dim * intermediate_dim * kElemBytes, stream); - cudaMemsetAsync(a.dW, 0, - static_cast(num_tokens) * top_k * kElemBytes, stream); - - // ── Build MoeBwdParams ────────────────────────────────────────── - MoeBwdParams p; - p.expert_indices = a.expert_indices; - p.expert_weights = reinterpret_cast(a.expert_weights); - // The BWD kernel never reads the original X token layout — token data - // is read only from x_sorted (FWD's expert-sorted view of X). Shape - // and device for dX_out come from dY. - p.dY = reinterpret_cast(a.dY); - p.output_buffer = reinterpret_cast(a.Y_fwd); - - p.dX_out = reinterpret_cast(a.dX); - p.dB_out = reinterpret_cast(a.dB); - p.dC_out = reinterpret_cast(a.dC); - p.dA_out = reinterpret_cast(a.dA); - p.dW_out = reinterpret_cast(a.dW); - - // Caller-provided forward intermediates (x_sorted, sort outputs). - // These come from moe_fused_fwd_bf16 and replace - // what the bwd allocator + internal sort/dispatch used to populate. - p.x_sorted = reinterpret_cast(a.x_sorted); - p.dy_sorted = buf.dy_sorted; - p.dx_sorted = buf.dx_sorted; - - // Sort outputs come from forward; we don't run sort_tokens in bwd. - // The fields we leave from buf.* (tile_expert_counts, sorted_token_ids, - // cta_done, cta_sums) are only consumed by sort_tokens — now skipped. - // expert_offsets is now the broadcasted [num_pes, num_experts + 1] - // view from FWD. The local-PE single-row view is offset by - // my_pe * (num_experts + 1) — needed by sort lookup - // (total_real_rows = local_offsets[num_experts]) and by comm code - // that walks this PE's expert range. - int* local_expert_offsets = - a.expert_offsets + my_pe * (num_experts + 1); - - p.sort.tile_expert_counts = buf.tile_expert_counts; - p.sort.sorted_token_ids = buf.sorted_token_ids; - p.sort.token_expert_slots = a.token_expert_slots; - p.sort.expert_offsets = local_expert_offsets; - p.sort.expert_counts = nullptr; - p.sort.tile_expert_ids = a.tile_expert_ids; - p.sort.cta_done = buf.cta_done; - p.sort.cta_sums = buf.cta_sums; - - // Forward CommBuffers (used for pe_sync only — bwd no longer - // re-broadcasts expert offsets; it reads what fwd wrote). - p.comm.src_staging = nullptr; - p.comm.dst_staging = nullptr; - p.comm.src_ready = nullptr; - p.comm.src_consumed = nullptr; - p.comm.dst_ready = nullptr; - p.comm.dst_consumed = nullptr; - p.comm.expert_offsets = local_expert_offsets; - p.comm.expert_counts = nullptr; - p.comm.sorted_token_ids = buf.sorted_token_ids; - p.comm.local_tokens = reinterpret_cast(a.x_sorted); - p.comm.local_output = buf.dx_sorted; - // Forward CommBuffers field — preserved from before the bwd-side split. - // The X-pipe array is the canonical "tile_expert_ids" (the dY-pipe array - // is a bwd-only mirror gated by release_dy). - p.comm.tile_expert_ids = buf.tile_expert_ids_x; - p.comm.all_expert_offsets = a.expert_offsets; - p.comm.all_expert_counts = a.expert_counts; - p.comm.pe_sync.cta_counter = buf.cta_counter; - p.comm.pe_sync.num_pes = num_pes; - p.comm.pe_sync.my_pe = my_pe; - p.comm.pe_sync.team = team; - p.comm.hidden_dim = hidden_dim; - p.comm.num_experts = num_experts; - p.comm.experts_per_pe = experts_per_pe; - p.comm.all_expert_offsets_stride = num_experts + 1; - - // Backward CommBuffers (used for X+dY gets and dX puts). - p.comm_bwd.x_staging = buf.x_staging; - p.comm_bwd.dy_staging = buf.dy_staging; - p.comm_bwd.dx_staging = buf.dx_staging; - p.comm_bwd.x_src_ready = buf.x_src_ready; - p.comm_bwd.x_src_consumed = buf.x_src_consumed; - p.comm_bwd.dy_src_ready = buf.dy_src_ready; - p.comm_bwd.dy_src_consumed = buf.dy_src_consumed; - p.comm_bwd.dst_ready = buf.dst_ready; - p.comm_bwd.dst_consumed = buf.dst_consumed; - p.comm_bwd.tile_expert_ids_x = buf.tile_expert_ids_x; - p.comm_bwd.tile_expert_ids_dy = buf.tile_expert_ids_dy; - p.comm_bwd.tile_valid_rows_x = buf.tile_valid_rows_x; - // Cross-PE pull sources. Both x_sorted and dy_sorted are in the - // same global-expert-sorted layout on every PE: - // • x_sorted: forward's sort_tokens → dispatch wrote it. - // • dy_sorted: combine_tokens_bwd writes [slot, d] using the same - // token_expert_slots → matches forward's per-PE expert sort. - // all_expert_offsets is the broadcasted offset tensor produced by - // fwd (broadcast_expert_offsets) and passed through here verbatim — - // it tells remote PEs the [start, end) slice in our buffers that - // holds data for their experts. - p.comm_bwd.remote_x = p.x_sorted; - p.comm_bwd.remote_dy = buf.dy_sorted; - p.comm_bwd.remote_dx = buf.dx_sorted; - p.comm_bwd.all_expert_offsets = a.expert_offsets; - p.comm_bwd.all_expert_counts = a.expert_counts; - p.comm_bwd.all_expert_offsets_stride = num_experts + 1; - p.comm_bwd.my_pe = my_pe; - p.comm_bwd.num_pes = num_pes; - p.comm_bwd.team = team; - p.comm_bwd.hidden_dim = hidden_dim; - p.comm_bwd.num_experts = num_experts; - p.comm_bwd.experts_per_pe = experts_per_pe; - // GEMM/comm-active count: comm derives MC = n_gemm/NC and the MLP column - // count grid_x = n_gemm/NSplit from this (the flat launch over-covers it). - p.comm_bwd.n_gemm = n_gemm; - - p.num_tokens = num_tokens; - p.hidden_dim = hidden_dim; - p.intermediate_dim = intermediate_dim; - p.top_k = top_k; - p.total_slots = max_total_slots; - p.num_m_tiles = num_m_tiles_in; - p.experts_per_pe = experts_per_pe; - - // ── MlpBwdDims (passed to inner mlp_fused_bwd) ────────────────── - int total_n_rows_1 = experts_per_pe * intermediate_dim; - int total_k_cols_2 = experts_per_pe * hidden_dim; - - MlpBwdDims mlp_dims; - mlp_dims.hidden_dim = hidden_dim; - mlp_dims.intermediate_dim = intermediate_dim; - mlp_dims.total_n_rows_1 = total_n_rows_1; - mlp_dims.num_n_tiles_1 = (Config::kCompute == 90) - ? (intermediate_dim + Traits1::TileN - 1) / Traits1::TileN - : intermediate_dim / Traits1::TileN; - mlp_dims.num_k_tiles_1 = (Config::kCompute == 90) - ? (hidden_dim + Traits1::TileK - 1) / Traits1::TileK - : hidden_dim / Traits1::TileK; - mlp_dims.total_k_cols_2t = total_k_cols_2; - mlp_dims.num_n_tiles_2t = (Config::kCompute == 90) - ? (intermediate_dim + Traits2T::TileN - 1) / Traits2T::TileN - : intermediate_dim / Traits2T::TileN; - mlp_dims.num_k_tiles_2t = (Config::kCompute == 90) - ? (hidden_dim + Traits2T::TileK - 1) / Traits2T::TileK - : hidden_dim / Traits2T::TileK; - mlp_dims.total_k_cols_5 = total_n_rows_1; - mlp_dims.num_n_tiles_5 = (Config::kCompute == 90) - ? (hidden_dim + Traits5::TileN - 1) / Traits5::TileN - : hidden_dim / Traits5::TileN; - mlp_dims.num_k_tiles_5 = (Config::kCompute == 90) - ? (intermediate_dim + Traits5::TileK - 1) / Traits5::TileK - : intermediate_dim / Traits5::TileK; - mlp_dims.num_m_tiles_3 = (Config::kCompute == 90) - ? (hidden_dim + Traits3::TileM - 1) / Traits3::TileM - : hidden_dim / Traits3::TileM; - mlp_dims.num_n_tiles_3 = (Config::kCompute == 90) - ? (intermediate_dim + Traits3::TileN - 1) / Traits3::TileN - : intermediate_dim / Traits3::TileN; - mlp_dims.total_n_rows_3 = total_k_cols_2; - mlp_dims.num_m_tiles_4 = (Config::kCompute == 90) - ? (intermediate_dim + Traits4::TileM - 1) / Traits4::TileM - : intermediate_dim / Traits4::TileM; - mlp_dims.num_n_tiles_4 = (Config::kCompute == 90) - ? (hidden_dim + Traits4::TileN - 1) / Traits4::TileN - : hidden_dim / Traits4::TileN; - mlp_dims.total_m_rows_4 = total_n_rows_1; - mlp_dims.phase_counter = buf.phase_counter; - // Flat-grid geometry: Phase 1 columns = grid_x, GEMM-active count = n_gemm. - mlp_dims.grid_x = grid_x; - mlp_dims.n_gemm = n_gemm; - mlp_dims.num_tokens = max_total_slots; - mlp_dims.num_m_tiles = max_total_slots / kTileM; - // Local pass: ring/buffer size and the K-block iteration bound coincide, - // so total_m_tiles_in_pass equals num_m_tiles. The remote pass overrides - // this in moe_bwd.cuh to comm.global_total. - mlp_dims.total_m_tiles_in_pass = mlp_dims.num_m_tiles; - // Phase 1 processes grid_x m-tiles per batch (one per column). - mlp_dims.num_batches = (mlp_dims.num_m_tiles + grid_x - 1) / grid_x; - // Sort emits GLOBAL expert IDs; mlp_fused_bwd computes - // `local_expert = tile.expert - local_expert_start` to index this PE's - // local all_B/all_C/all_A. Forward kernel does the same (moe.cu:469). - mlp_dims.local_expert_start = my_pe * experts_per_pe; - mlp_dims.num_pes = num_pes; - mlp_dims.experts_per_pe = experts_per_pe; - // Local pass: expert_for_k_block is built per-K-block by - // build_expert_for_k_block (one entry per Traits{3,4}::TileK), so the - // K-block index can be used directly as the array index. The remote pass - // overrides these via EfkbStride{3,4} template params in moe_fused_bwd - // because its array is per-comm-tile. - - // Host-computable constants (grid-constant — saves device-side arithmetic). - // phase2_num_blocks is no longer used for the Phase 2 cell stride (that now - // reads gridDim.x·gridDim.y = num_blocks directly), kept for compatibility. - mlp_dims.phase2_num_blocks = num_blocks; - // Remote pass: TMA descriptors cover the full flat staging list of - // L = MC · kCommNumStages slots, each TileM rows tall. - // MC = launched NC-complete comm columns (matches nvshmem_comm_main_bwd). - int mc_remote = num_blocks / Config::kNC; - mlp_dims.remote_num_m_tiles = mc_remote * Config::kCommNumStages; - mlp_dims.remote_num_tokens = mlp_dims.remote_num_m_tiles * kTileM; - - // Phase-1 sub-batch grouping (mlp_bwd.cuh, SubBatch): the remote pass groups - // up to SubBatch Phase-1 sub-batches before one Phase-2 over their combined - // K-window. The kernel ADAPTS the group size to the ring depth (a group must - // fit in the ring so its grid_x·gcount live slots map to distinct ring slots), - // degrading to fewer sub-batches — down to 1 = the original per-batch loop — - // on shallow rings. So no host-side ring constraint is needed here. - - MlpBwdBufs mlp_bufs; - mlp_bufs.du_buf = buf.du_buf; - mlp_bufs.dv_buf = buf.dv_buf; - mlp_bufs.dz_buf = buf.dz_buf; - // Local pass: no comm staging, no race — both Phase 2 sub-passes alias - // the same post-sort expert_for_k_block. The remote sub-call in - // moe_fused_bwd overrides these to the per-pipe staging arrays. - mlp_bufs.expert_for_k_block_mlp4 = buf.expert_for_k_block; - mlp_bufs.expert_for_k_block_mlp3 = buf.expert_for_k_block; - mlp_bufs.valid_rows_mlp4 = nullptr; - mlp_bufs.expert_k_starts = buf.expert_k_starts; - mlp_bufs.expert_k_ends = buf.expert_k_ends; - mlp_bufs.gmem_da = reinterpret_cast(a.dA); - mlp_bufs.gmem_db = reinterpret_cast(a.dB); - mlp_bufs.gmem_dc = reinterpret_cast(a.dC); - mlp_bufs.barrier_counter = buf.barrier_counter; - - // ── TMA descriptors ───────────────────────────────────────────── - // X comes from caller's x_sorted (forward's pool, post-dispatch). - // dy_sorted is bwd-allocated (combine_tokens_bwd will populate it). - auto* pX = reinterpret_cast(a.x_sorted); // sorted X - auto* pDY = reinterpret_cast(buf.dy_sorted); // expanded dY - auto* pB = reinterpret_cast(a.all_B); - auto* pC = reinterpret_cast(a.all_C); - auto* pA = reinterpret_cast(a.all_A); - auto* pDX = reinterpret_cast(buf.dx_sorted); - - // All Hopper expert weights and weight-gradient targets use rank-3 tensor - // maps. TMA handles partial edge boxes as zero-filled loads / dropped - // stores, while the explicit expert coordinate prevents cross-expert tails. - // (See mlp_bwd.cu for the full pattern; here we mirror it but the - // "X" tensor is the symmetric x_sorted and "dY" is the symmetric - // dy_sorted, both of shape [max_total_slots, hidden_dim].) - auto tma_load_x = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pX), make_shape(max_total_slots, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits1::SmemLayoutX_1{}); - - auto tma_load_b_fwd = [&]() { - if constexpr (Config::kCompute == 90) { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pB), - make_shape(static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(experts_per_pe)), - make_stride( - static_cast(hidden_dim), - Int<1>{}, - static_cast(intermediate_dim) * hidden_dim)), - typename Traits1::SmemLayoutW_1{}); - } else { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pB), - make_shape(total_n_rows_1, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits1::SmemLayoutW_1{}); - } - }(); - auto tma_load_c_fwd = [&]() { - if constexpr (Config::kCompute == 90) { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pC), - make_shape(static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(experts_per_pe)), - make_stride( - static_cast(hidden_dim), - Int<1>{}, - static_cast(intermediate_dim) * hidden_dim)), - typename Traits1::SmemLayoutW_1{}); - } else { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pC), - make_shape(total_n_rows_1, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits1::SmemLayoutW_1{}); - } - }(); - - auto tma_store_z = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(buf.z_buf), make_shape(buf.scratch_rows, intermediate_dim), - make_stride(intermediate_dim, Int<1>{})), - typename Traits1::SmemLayoutActStoreSlot{}); - auto tma_store_u = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(buf.du_buf), make_shape(buf.scratch_rows, intermediate_dim), - make_stride(intermediate_dim, Int<1>{})), - typename Traits1::SmemLayoutActStoreSlot{}); - auto tma_store_v = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(buf.dv_buf), make_shape(buf.scratch_rows, intermediate_dim), - make_stride(intermediate_dim, Int<1>{})), - typename Traits1::SmemLayoutActStoreSlot{}); - - auto tma_load_dy = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pDY), make_shape(max_total_slots, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits2T::SmemLayoutZ_1{}); - auto tma_load_a_col = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pA), - make_shape(static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(experts_per_pe)), - make_stride( - Int<1>{}, - static_cast(intermediate_dim), - static_cast(hidden_dim) * intermediate_dim)), - typename Traits2T::SmemLayoutW_1{}); - // One TMA box per sub-tile; consumer issues NSub stores per outer N-tile. - auto tma_store_dz = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(buf.dz_buf), - make_shape(buf.scratch_rows, intermediate_dim), - make_stride(intermediate_dim, Int<1>{})), - typename Traits2T::SmemLayoutStore_1{}); - - auto tma_load_du = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(buf.du_buf)), - make_shape(buf.scratch_rows, intermediate_dim), - make_stride(intermediate_dim, Int<1>{})), - typename Traits5::SmemLayoutZ_1{}); - auto tma_load_dv = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(buf.dv_buf)), - make_shape(buf.scratch_rows, intermediate_dim), - make_stride(intermediate_dim, Int<1>{})), - typename Traits5::SmemLayoutZ_1{}); - auto tma_load_b_col = [&]() { - if constexpr (Config::kCompute == 90) { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pB), - make_shape(static_cast(hidden_dim), - static_cast(intermediate_dim), - static_cast(experts_per_pe)), - make_stride( - Int<1>{}, - static_cast(hidden_dim), - static_cast(intermediate_dim) * hidden_dim)), - typename Traits5::SmemLayoutW_1{}); - } else { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pB), - make_shape(hidden_dim, total_n_rows_1), - make_stride(Int<1>{}, hidden_dim)), - typename Traits5::SmemLayoutW_1{}); - } - }(); - auto tma_load_c_col = [&]() { - if constexpr (Config::kCompute == 90) { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pC), - make_shape(static_cast(hidden_dim), - static_cast(intermediate_dim), - static_cast(experts_per_pe)), - make_stride( - Int<1>{}, - static_cast(hidden_dim), - static_cast(intermediate_dim) * hidden_dim)), - typename Traits5::SmemLayoutW_1{}); - } else { - return make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(pC), - make_shape(hidden_dim, total_n_rows_1), - make_stride(Int<1>{}, hidden_dim)), - typename Traits5::SmemLayoutW_1{}); - } - }(); - // dX TMA: split-N store (TileN_half boxes per Mlp5::SmemLayoutStore_1). - auto tma_store_dx = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(pDX), - make_shape(max_total_slots, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits5::SmemLayoutStore_1{}); - - // Phase 2 TMAs. mlp3's dY^T / Z operands and mlp4's X / dU^T / dV^T - // operands need pair-aware SM100_TMA_2SM_LOAD descriptors (built via - // make_tma_copy_A_sm100 / make_tma_copy_B_sm100, keyed off the joined - // TileShape + TiledMma2Sm) when Config::kUsesTwoSm (all SM100 configs); - // the dA/dB/dC TMA_REDUCE_ADD outputs are unchanged either way (each - // peer CTA still reduce-adds its own CtaTileM rows independently). - auto tma_load_dyt3 = [&] { - auto tensor = make_tensor(make_gmem_ptr(pDY), - make_shape(hidden_dim, max_total_slots), - make_stride(Int<1>{}, hidden_dim)); - if constexpr (Config::kUsesTwoSm) { - return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, - typename Traits3::SmemLayoutDYT_1{}, - typename Traits3::TileShape{}, - typename Traits3::TiledMma2Sm{}); - } else { - return make_tma_copy(TmaLoadAtom{}, tensor, - typename Traits3::SmemLayoutDYT_1{}); - } - }(); - auto tma_load_zt3 = [&] { - auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.z_buf)), - make_shape(intermediate_dim, buf.scratch_rows), - make_stride(Int<1>{}, intermediate_dim)); - if constexpr (Config::kUsesTwoSm) { - return make_tma_copy_B_sm100(SM100_TMA_2SM_LOAD{}, tensor, - typename Traits3::SmemLayoutZ_1{}, - typename Traits3::TileShape{}, - typename Traits3::TiledMma2Sm{}); - } else { - return make_tma_copy(TmaLoadAtom{}, tensor, - typename Traits3::SmemLayoutZ_1{}); - } - }(); - auto tma_reduce_da = [&]() { - if constexpr (Config::kCompute == 90) { - return make_tma_copy(TmaReduceAddAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(a.dA)), - make_shape(static_cast(hidden_dim), - static_cast(intermediate_dim), - static_cast(experts_per_pe)), - make_stride( - static_cast(intermediate_dim), - Int<1>{}, - static_cast(hidden_dim) * intermediate_dim)), - typename Traits3::SmemLayoutStore{}); - } else { - return make_tma_copy(TmaReduceAddAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(a.dA)), - make_shape(total_k_cols_2, intermediate_dim), - make_stride(intermediate_dim, Int<1>{})), - typename Traits3::SmemLayoutStore{}); - } - }(); - - auto tma_load_xt4 = [&] { - auto tensor = make_tensor(make_gmem_ptr(pX), - make_shape(hidden_dim, max_total_slots), - make_stride(Int<1>{}, hidden_dim)); - if constexpr (Config::kUsesTwoSm) { - return make_tma_copy_B_sm100(SM100_TMA_2SM_LOAD{}, tensor, - typename Traits4::SmemLayoutX_1{}, - typename Traits4::TileShape{}, - typename Traits4::TiledMma2Sm{}); - } else { - return make_tma_copy(TmaLoadAtom{}, tensor, - typename Traits4::SmemLayoutX_1{}); - } - }(); - auto tma_load_dut4 = [&] { - auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.du_buf)), - make_shape(intermediate_dim, buf.scratch_rows), - make_stride(Int<1>{}, intermediate_dim)); - if constexpr (Config::kUsesTwoSm) { - return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, - typename Traits4::SmemLayoutA_1{}, - typename Traits4::TileShape{}, - typename Traits4::TiledMma2Sm{}); - } else { - return make_tma_copy(TmaLoadAtom{}, tensor, - typename Traits4::SmemLayoutA_1{}); - } - }(); - auto tma_load_dvt4 = [&] { - auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.dv_buf)), - make_shape(intermediate_dim, buf.scratch_rows), - make_stride(Int<1>{}, intermediate_dim)); - if constexpr (Config::kUsesTwoSm) { - return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, - typename Traits4::SmemLayoutA_1{}, - typename Traits4::TileShape{}, - typename Traits4::TiledMma2Sm{}); - } else { - return make_tma_copy(TmaLoadAtom{}, tensor, - typename Traits4::SmemLayoutA_1{}); - } - }(); - auto tma_reduce_db = [&]() { - if constexpr (Config::kCompute == 90) { - return make_tma_copy(TmaReduceAddAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(a.dB)), - make_shape(static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(experts_per_pe)), - make_stride( - static_cast(hidden_dim), - Int<1>{}, - static_cast(intermediate_dim) * hidden_dim)), - typename Traits4::SmemLayoutStore{}); - } else { - return make_tma_copy(TmaReduceAddAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(a.dB)), - make_shape(total_n_rows_1, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits4::SmemLayoutStore{}); - } - }(); - auto tma_reduce_dc = [&]() { - if constexpr (Config::kCompute == 90) { - return make_tma_copy(TmaReduceAddAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(a.dC)), - make_shape(static_cast(intermediate_dim), - static_cast(hidden_dim), - static_cast(experts_per_pe)), - make_stride( - static_cast(hidden_dim), - Int<1>{}, - static_cast(intermediate_dim) * hidden_dim)), - typename Traits4::SmemLayoutStore{}); - } else { - return make_tma_copy(TmaReduceAddAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(a.dC)), - make_shape(total_n_rows_1, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits4::SmemLayoutStore{}); - } - }(); - - // Remote (staging) TMAs — buffers cover L = MC · CommNumStages slots, - // each TileM rows tall. MC = (grid.x · NSplit) / NC. - int remote_total_rows = mc_remote * Config::kCommNumStages * kTileM; - auto tma_load_x_remote = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(buf.x_staging)), - make_shape(remote_total_rows, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits1::SmemLayoutX_1{}); - auto tma_load_dy_remote = make_tma_copy(TmaLoadAtom{}, - make_tensor(make_gmem_ptr(reinterpret_cast(buf.dy_staging)), - make_shape(remote_total_rows, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits2T::SmemLayoutZ_1{}); - // Remote dX TMA: split-N store (matches local tma_store_dx). - auto tma_store_dx_remote = make_tma_copy(TmaStoreAtom{}, - make_tensor(make_gmem_ptr(buf.dx_staging), - make_shape(remote_total_rows, hidden_dim), - make_stride(hidden_dim, Int<1>{})), - typename Traits5::SmemLayoutStore_1{}); - - // Remote Phase 3/4 TMAs: transposed staging sources. The remote phase - // of moe_fused_bwd previously reused tma_load_dyt3/tma_load_xt4 (which - // point at LOCAL dy_sorted/x_sorted), so dA/dB/dC for cross-PE-routed - // tokens used the wrong rows. These descriptors pull from the same - // staging buffers that Phase 1/2 of the remote pass already use. - auto tma_load_dyt3_remote = [&] { - auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.dy_staging)), - make_shape(hidden_dim, remote_total_rows), - make_stride(Int<1>{}, hidden_dim)); - if constexpr (Config::kUsesTwoSm) { - return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, - typename Traits3::SmemLayoutDYT_1{}, - typename Traits3::TileShape{}, - typename Traits3::TiledMma2Sm{}); - } else { - return make_tma_copy(TmaLoadAtom{}, tensor, - typename Traits3::SmemLayoutDYT_1{}); - } - }(); - auto tma_load_xt4_remote = [&] { - auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.x_staging)), - make_shape(hidden_dim, remote_total_rows), - make_stride(Int<1>{}, hidden_dim)); - if constexpr (Config::kUsesTwoSm) { - return make_tma_copy_B_sm100(SM100_TMA_2SM_LOAD{}, tensor, - typename Traits4::SmemLayoutX_1{}, - typename Traits4::TileShape{}, - typename Traits4::TiledMma2Sm{}); - } else { - return make_tma_copy(TmaLoadAtom{}, tensor, - typename Traits4::SmemLayoutX_1{}); - } - }(); - - // ── GET-via-TMA descriptors for the bwd comm path (#111 port) ──── - // Per-peer source maps over nvshmem_ptr(x_sorted/dy_sorted, peer) + one - // local dst map each over x_staging / dy_staging. Enabled only on NVLink - // P2P + divisible hidden_dim + bounce smem fit (checked after smem sizing). - GetTmaDescsBwd get_descs_bwd{}; - get_descs_bwd.enabled = 0; - get_descs_bwd.gpus_per_node = gpus_per_node; - get_descs_bwd.my_pe = my_pe; - constexpr int kGetBoxRowsBwd = Config::kCompute == 100 - ? liger::kGetBoxRowsBwdSm100 - : liger::kGetBoxRowsBwdSm90; - size_t get_bounce_bytes_bwd = - (size_t)liger::GetBounceBwd::kTotalBytes; - { - // The single get warp's row band is TileM/NC; the small TMA box height - // must divide it (else the last band would over-read). Holds for NC≤8 - // at TileM=128; larger NC (band < 16) disables the path. - constexpr int kRowBandBwd = Config::kTileM / Config::kNC; - bool ok = (hidden_dim % liger::kGetKChunk == 0) - && (gpus_per_node <= liger::kGetMaxPes) - && (kRowBandBwd % kGetBoxRowsBwd == 0); - int host_start = (my_pe / gpus_per_node) * gpus_per_node; - for (int local_peer = 0; ok && local_peer < gpus_per_node; ++local_peer) { - int tp = host_start + local_peer; - if (tp >= num_pes) { ok = false; break; } - int gpe = nvshmem_team_translate_pe(team, tp, NVSHMEM_TEAM_WORLD); - void* px_peer = nvshmem_ptr(p.comm_bwd.remote_x, gpe); // peer x_sorted - void* pdy_peer = nvshmem_ptr(p.comm_bwd.remote_dy, gpe); // peer dy_sorted - if (px_peer == nullptr || pdy_peer == nullptr) { ok = false; break; } - ok = make_get_tma_map_bwd(get_descs_bwd.src_x_desc[local_peer], px_peer, - hidden_dim, max_total_slots, kGetBoxRowsBwd); - if (ok) ok = make_get_tma_map_bwd(get_descs_bwd.src_dy_desc[local_peer], pdy_peer, - hidden_dim, max_total_slots, kGetBoxRowsBwd); - } - if (ok) ok = make_get_tma_map_bwd(get_descs_bwd.dst_x_staging_desc, - buf.x_staging, hidden_dim, remote_total_rows, kGetBoxRowsBwd); - if (ok) ok = make_get_tma_map_bwd(get_descs_bwd.dst_dy_staging_desc, - buf.dy_staging, hidden_dim, remote_total_rows, kGetBoxRowsBwd); - get_descs_bwd.enabled = ok ? 1 : 0; - if (const char* e = std::getenv("LIGER_GET_TMA")) - if (e[0] == '0') get_descs_bwd.enabled = 0; - } - - // ── Smem sizing ───────────────────────────────────────────────── - int sort_num_warps = kNumThreads / 32; - size_t smem_sort = sort_num_warps * num_experts * sizeof(int); - - size_t smem_mlp = sizeof(MoeBwdSmem); - smem_mlp = (smem_mlp + sizeof(int) - 1) & ~(sizeof(int) - 1); - smem_mlp += num_pes * (experts_per_pe + 1) * sizeof(int); // remote_offsets - smem_mlp += num_pes * experts_per_pe * sizeof(int); // remote_counts - - // Query the device's actual opt-in dynamic-smem cap rather than trust - // Config::kSmemBudget's compile-time 228 KiB constant: the 2SM paired-CTA - // path (Mlp3Traits2Sm / Mlp4Traits2Sm) pushes some candidates close to - // that ceiling, and the real per-SKU limit is what - // cudaFuncAttributeMaxDynamicSharedMemorySize is actually bounded by. - int device = 0; - int dynamic_smem_limit = 0; - if (cudaError_t e = cudaGetDevice(&device); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_bwd: cudaGetDevice failed: ", cudaGetErrorString(e)); - if (cudaError_t e = cudaDeviceGetAttribute(&dynamic_smem_limit, - cudaDevAttrMaxSharedMemoryPerBlockOptin, device); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_bwd: querying the opt-in shared-memory limit failed: ", - cudaGetErrorString(e)); - - // GET-via-TMA bounce: append a 128-aligned region after remote_offsets - // (the kernel computes the same offset). The bwd union sits near the - // device's dynamic-smem cap, so disable the TMA-get path if the bounce - // won't fit against the QUERIED limit (not the static 228 KiB constant) - // — that config just keeps the original 2-get/1-put getmem layout (no - // regression). - if (get_descs_bwd.enabled) { - size_t mlp_with_bounce = - ((smem_mlp + 127) & ~((size_t)127)) + get_bounce_bytes_bwd; - if (mlp_with_bounce <= static_cast(dynamic_smem_limit)) smem_mlp = mlp_with_bounce; - else get_descs_bwd.enabled = 0; - } - // Unified BWD comm layout always uses one get warp and one put warp. - // get_descs_bwd.enabled only controls whether same-host tiles use TMA; - // remote/unsupported tiles fall back inside do_get_bwd_tma. - mlp_dims.tma_get_enabled = 1; - - size_t smem_size = std::max(smem_sort, smem_mlp); - - // ── Launch ────────────────────────────────────────────────────── - auto kernel_fn = moe_bwd_kernel; - - if (cudaError_t e = cudaFuncSetAttribute(kernel_fn, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_bwd: cudaFuncSetAttribute(MaxDynamicSharedMemorySize=", - smem_size, ", device opt-in limit=", dynamic_smem_limit, ") failed: ", - cudaGetErrorString(e)); - // The SM100 2SM paired-CTA cluster path launches with an explicit - // clusterDim; opt in to non-portable cluster sizes (the CUDA default - // portable max is 8, so ClusterM=2 is already portable in practice, - // but every 2SM launch enables this defensively, mirroring CUTLASS's - // own cluster-launch helpers). - if constexpr (Config::kUsesTwoSm) { - if (cudaError_t e = cudaFuncSetAttribute(kernel_fn, - cudaFuncAttributeNonPortableClusterSizeAllowed, 1); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_bwd: enabling non-portable cluster size (ClusterM=", - Config::kClusterM, ") failed: ", cudaGetErrorString(e)); - } - - // Fresh per-call counter resets. - cudaMemsetAsync(buf.cta_counter, 0, sizeof(int), stream); - cudaMemsetAsync(buf.barrier_counter, 0, sizeof(int), stream); - cudaMemsetAsync(buf.phase_counter, 0, max_runtime_grid_x * sizeof(int), stream); - // L = MC · CommNumStages, MC = launched NC-complete comm columns. - int mc_local = ((int)grid.x / Config::kNC); - int comm_l_local = mc_local * Config::kCommNumStages; - size_t per_call_tile_sig_bytes = - static_cast(comm_l_local) * sizeof(int); - cudaMemsetAsync(buf.x_src_ready, 0, per_call_tile_sig_bytes, stream); - cudaMemsetAsync(buf.x_src_consumed, 0, per_call_tile_sig_bytes, stream); - cudaMemsetAsync(buf.dy_src_ready, 0, per_call_tile_sig_bytes, stream); - cudaMemsetAsync(buf.dy_src_consumed,0, per_call_tile_sig_bytes, stream); - cudaMemsetAsync(buf.dst_ready, 0, per_call_tile_sig_bytes, stream); - cudaMemsetAsync(buf.dst_consumed, 0, per_call_tile_sig_bytes, stream); - cudaMemsetAsync(buf.cta_done, 0, num_blocks * sizeof(int), stream); - cudaMemsetAsync(buf.sorted_token_ids, 0xff, max_total_slots * sizeof(int), stream); - // dy_sorted MUST be zeroed: combine_tokens_bwd writes routed slots - // only (no padding fill — unlike forward's dispatch_tokens). Partial- - // expert tile rows of dy_sorted are read by Phase 1b's TMA load; if - // those rows aren't zero, the resulting dZ/dU/dV poison Phase 2 dA/dB/dC - // for K-blocks that straddle an expert boundary (see the - // build_expert_for_k_block comment for the alignment assumption). - cudaMemsetAsync(buf.dy_sorted, 0, static_cast(max_total_slots) * hidden_dim * sizeof(Element), stream); - size_t intermediate_bytes = - static_cast(buf.scratch_rows) * intermediate_dim * sizeof(Element); - cudaMemsetAsync(buf.z_buf, 0, intermediate_bytes, stream); - cudaMemsetAsync(buf.du_buf, 0, intermediate_bytes, stream); - cudaMemsetAsync(buf.dv_buf, 0, intermediate_bytes, stream); - cudaMemsetAsync(buf.dz_buf, 0, intermediate_bytes, stream); - - // Sentinel-fill BOTH comm-side tile_expert_ids arrays so unfetched - // slots produce expert_id = -1, which the Phase 3/4 bounds check - // rejects. Without this, stale memory could land in [0, experts_per_pe) - // and trigger spurious dB/dC/dA writes for empty staging slots. - size_t comm_eid_bytes = - static_cast(comm_l_local) * sizeof(int); - if (cudaError_t e = cudaMemsetAsync(buf.tile_expert_ids_x, 0xff, comm_eid_bytes, stream); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_bwd: cudaMemsetAsync(tile_expert_ids_x) failed: ", cudaGetErrorString(e)); - if (cudaError_t e = cudaMemsetAsync(buf.tile_expert_ids_dy, 0xff, comm_eid_bytes, stream); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_bwd: cudaMemsetAsync(tile_expert_ids_dy) failed: ", cudaGetErrorString(e)); - cudaMemsetAsync(buf.tile_valid_rows_x, 0, comm_eid_bytes, stream); - if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_bwd: CUDA error before kernel launch: ", cudaGetErrorString(e)); - - // Compute=90 keeps the ordinary <<<>>> launch. Compute=100 always uses - // the explicit paired-CTA cluster launch below. - if constexpr (Config::kUsesTwoSm) { - cudaLaunchConfig_t launch_config = {}; - launch_config.gridDim = grid; - launch_config.blockDim = kNumThreads; - launch_config.dynamicSmemBytes = smem_size; - launch_config.stream = stream; - cudaLaunchAttribute cluster_attr = {}; - cluster_attr.id = cudaLaunchAttributeClusterDimension; - cluster_attr.val.clusterDim.x = Config::kClusterM; - cluster_attr.val.clusterDim.y = 1; - cluster_attr.val.clusterDim.z = 1; - launch_config.attrs = &cluster_attr; - launch_config.numAttrs = 1; - cudaLaunchKernelEx(&launch_config, kernel_fn, - tma_load_x, tma_load_b_fwd, tma_load_c_fwd, - tma_store_z, tma_store_u, tma_store_v, - tma_load_dy, tma_load_a_col, tma_store_dz, - tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, - tma_load_dyt3, tma_load_zt3, tma_reduce_da, - tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, - tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, - tma_load_dyt3_remote, tma_load_xt4_remote, - mlp_dims, mlp_bufs, p, static_nsplit, get_descs_bwd); - } else { - kernel_fn<<>>( - tma_load_x, tma_load_b_fwd, tma_load_c_fwd, - tma_store_z, tma_store_u, tma_store_v, - tma_load_dy, tma_load_a_col, tma_store_dz, - tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, - tma_load_dyt3, tma_load_zt3, tma_reduce_da, - tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, - tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, - tma_load_dyt3_remote, tma_load_xt4_remote, - mlp_dims, mlp_bufs, p, static_nsplit, get_descs_bwd); - } - if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) - LIGER_FAIL_CUDA("moe_bwd: kernel launch failed for grid=", grid.x, - " threads=", kNumThreads, " smem=", smem_size, ": ", - cudaGetErrorString(e)); - - // Post-kernel stream-ordered barrier: mirrors the fwd wrapper. The - // kernel's PeSync wait_mlp guarantees this PE has received peer - // signals, but does not guarantee this PE's outbound signals have - // landed on peers. Under back-to-back captured replays (fwd → bwd → - // fwd), bwd's tail-end signals to peers can race against the next - // fwd's wait_dispatch: a peer's wait_until(>= target) may briefly - // see a stale value if bwd's older signal arrives after a newer - // one. Forcing host-side completion here serialises remote - // signal delivery against the next captured kernel. - if (num_pes > 1) { - nvshmemx_barrier_on_stream(team, stream); - } -} - -// ============================================================================ -// Tuned-config wrapper: 13 free knobs → full MoeBwdConfig -// ============================================================================ -// -// Function template that takes the X-macro's free knobs and derives the -// dependent traits per the symmetric-traits constraints documented in -// moe_fwd_bwd_tune_configs.hpp: -// Phase 1a: TileN1, TileK1, Stages1, EpiChunkN1 -// Phase 1b': TileN = 256 (fixed), TileK = TileK1, Stages = Stages1, EpiChunkN25 -// Phase 1d (mlp5): TileM=128, TileN=256, TileK = TileK1, Stages = Stages1, EpiChunkN25 -// Phase 2 mlp3: TileM3, TileN3, TileK3, Stages3, EpiChunkN34 -// Phase 2 mlp4: TileM = TileN3, TileN = TileM3, TileK3, Stages3, EpiChunkN34 -// (mlp3/mlp4 swap) -// NC = NSplit -// -// The tuner registry and dispatch table are built by emitting one explicit -// instantiation of this template per X-macro entry. - -template < - int NSplit2, - int TileN1, int TileK1, int Stages1, - int TileM3, int TileN3, int TileK3, int Stages3, - int EpiChunkN1, int EpiChunkN25, int EpiChunkN34, - int CommNumStages, - int TileM, // COMM tile (staging/semaphores/NC); must match FWD's TileM. 64 or 128. - int GemmTileM = TileM, // Phase-1 GEMM tile (mlp1a/mlp2t/mlp5) ∈ {64,128}; default == comm. - int Compute = 90> -void -moe_bwd_fwd_bf16_tuned(const MoeBwdArgs& a, int static_nsplit) { - // NC paired with TileM (mirrors moe.cu's moe_fused_fwd_bf16): - // TileM=128 → NC=4 (~2× the comm bandwidth of NC=2; deeper comm pipeline). - // TileM=64 → NC=2 (tiles are half-size, NC=2 matches NC=4 throughput). - // Ticket-based iterator + K=NumStages comm means NSplit and NC have no - // divisibility constraint — only requirement is that NC evenly divides - // gridDim.x · gridDim.y so MC is integer (e.g. NS=6, NC=4: MC = 22·6/4 - // = 33). - // Both architectures use CommTileM=128/NC4, matching forward metadata. - using CommGeometry = MoeCommGeometry; - static constexpr int kPhase1TileM = Compute == 100 ? 128 : GemmTileM; - static constexpr int kSubBatch = 2; - using Config = MoeBwdConfig< - NSplit2, - kPhase1TileM, TileN1, TileK1, Stages1, EpiChunkN1, // Phase 1a (mlp1_act) — GEMM tile - /*TileN2=*/256, TileK1, Stages1, // Phase 1b' (Mlp2TTraits caps TileN ≤ 256) - /*TileM2T=*/kPhase1TileM, // Phase 1b' GEMM tile - TileM3, TileN3, TileK3, Stages3, // Phase 2 mlp3 - TileN3, TileM3, TileK3, Stages3, // Phase 2 mlp4 (mlp3/mlp4 swap; cooperative) - kPhase1TileM, 256, TileK1, Stages1, // Phase 1d mlp5 — GEMM tile - EpiChunkN25, // shared mlp2_t + mlp5 - EpiChunkN34, // shared mlp3 + mlp4 - bfloat16_t, - CommNumStages * CommGeometry::RingStageScale - + CommGeometry::RingStageExtra, - /*NC_=*/CommGeometry::NC, - /*SubBatch_=*/kSubBatch, - /*CommTileM_=*/CommGeometry::TileM, - /*Compute_=*/Compute>; - moe_bwd_fwd_bf16(a, static_nsplit); -} - -// ── Runtime dispatch table for tuned-config lookup ──────────────────── - -using MoeBwdFwdFn = void (*)(const MoeBwdArgs&, int static_nsplit); - -struct DispatchEntryBwd { - int Compute; - int NSplit2; - int TileN1, TileK1, Stages1; - int TileM3, TileN3, TileK3, Stages3; - int EpiChunkN1, EpiChunkN25, EpiChunkN34; - int CommNumStages; - int TileM; // COMM tile (staging/semaphores/NC). - int GemmTileM; // Phase-1 GEMM tile (mlp1a/mlp2t/mlp5); == TileM for coupled configs. - MoeBwdFwdFn fn; -}; - -#define LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) \ - { Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM, \ - &moe_bwd_fwd_bf16_tuned }, -#define LIGER_MOE_BWD_DISPATCH_ENTRY_C(Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) \ - LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, TM) -#define LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM90(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) \ - LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(90, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) -#define LIGER_MOE_BWD_DISPATCH_ENTRY_SM90(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) \ - LIGER_MOE_BWD_DISPATCH_ENTRY_C(90, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) -#define LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM100(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) \ - LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(100, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) -#define LIGER_MOE_BWD_DISPATCH_ENTRY_SM100(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) \ - LIGER_MOE_BWD_DISPATCH_ENTRY_C(100, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) - -static const DispatchEntryBwd kDispatchTableBwd[] = { -#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 90 - LIGER_MOE_BWD_DISPATCH_CONFIGS_SM90(LIGER_MOE_BWD_DISPATCH_ENTRY_SM90, LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM90) +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "moe_bwd_sm90.inc" +#else +#include "moe_bwd_sm100.inc" #endif -#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 100 - LIGER_MOE_BWD_DISPATCH_CONFIGS_SM100(LIGER_MOE_BWD_DISPATCH_ENTRY_SM100, LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM100) -#endif -}; -#undef LIGER_MOE_BWD_DISPATCH_ENTRY_SM100 -#undef LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM100 -#undef LIGER_MOE_BWD_DISPATCH_ENTRY_SM90 -#undef LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM90 -#undef LIGER_MOE_BWD_DISPATCH_ENTRY_C -#undef LIGER_MOE_BWD_DISPATCH_ENTRY_G_C -static constexpr int kNumDispatchEntriesBwd = - sizeof(kDispatchTableBwd) / sizeof(kDispatchTableBwd[0]); - -} // namespace liger - -// Joint fwd+bwd tuning table (auto-generated by benchmarks/tune_moe_fwd_bwd). -// We only consume the Bwd_* fields here; moe.cu consumes the same table's -// Fwd_* fields. Picking fwd and bwd from a single table guarantees the two -// directions are tuned as a pair — in particular their TileM agrees. -#include "moe_fwd_bwd_tuning_configs.cuh" - -namespace liger { - -static constexpr int kBwdRuntimeNsplitSeed = 8; - -// TunedConfigBwd is the bwd-only projection of TunedConfigFwdBwd. The rest -// of the dispatch code (validity check, dispatch-entry match) was written -// against this struct and is reused unchanged. -struct TunedConfigBwd { - int Compute; - int TK, TKE, D, I; - int NSplit2; - int TileN1, TileK1, Stages1; - int TileM3, TileN3, TileK3, Stages3; - int EpiChunkN1, EpiChunkN25, EpiChunkN34; - int CommNumStages; - int TileM; // comm tile - int GemmTileM; // gemm tile (auto path: == TileM; FORCE_CONFIG may set 64) -}; - -// Project a TunedConfigFwdBwd row down to its bwd template params. The tuned -// table has no GemmTileM column, so the auto path keeps gemm == comm (TileM). -static TunedConfigBwd project_bwd(const TunedConfigFwdBwd& c, int compute) { - return TunedConfigBwd{ - compute, - c.TK, c.TKE, c.D, c.I, - c.Bwd_NSplit2, - c.Bwd_TileN1, c.Bwd_TileK1, c.Bwd_Stages1, - c.Bwd_TileM3, c.Bwd_TileN3, c.Bwd_TileK3, c.Bwd_Stages3, - c.Bwd_EpiChunkN1, c.Bwd_EpiChunkN25, c.Bwd_EpiChunkN34, - c.Bwd_CommNumStages, - c.Bwd_TileM, c.Bwd_TileM, - }; -} - -// Find the dispatch table entry matching a TunedConfigBwd's knobs. -static const DispatchEntryBwd* find_dispatch_entry_bwd(const TunedConfigBwd& c) { - for (int i = 0; i < kNumDispatchEntriesBwd; ++i) { - const auto& e = kDispatchTableBwd[i]; - if (e.Compute == c.Compute && - e.NSplit2 == c.NSplit2 && - e.TileN1 == c.TileN1 && e.TileK1 == c.TileK1 && e.Stages1 == c.Stages1 && - e.TileM3 == c.TileM3 && e.TileN3 == c.TileN3 && e.TileK3 == c.TileK3 && - e.Stages3 == c.Stages3 && - e.EpiChunkN1 == c.EpiChunkN1 && - e.EpiChunkN25 == c.EpiChunkN25 && e.EpiChunkN34 == c.EpiChunkN34 && - e.CommNumStages == c.CommNumStages && - e.TileM == c.TileM && e.GemmTileM == c.GemmTileM) { - return &e; - } - } - return nullptr; -} - -// Structural divisibility checks on (D, I) for a given tuned config. -// Mirrors the constraint list in moe_fwd_bwd_tune_configs.hpp. -static bool tuned_config_valid_bwd(int D, int I, const TunedConfigBwd& c) { - // The current Blackwell UMMA MLP1/MLP1-act consumers only support the - // 128-row comm tile. Keep mirrored SM90 TileM=64 tuning rows out of SM100 - // auto-dispatch until that TMEM epilogue layout is implemented. - if (c.Compute == 100 && c.TileM != 128) return false; - if (c.Compute == 100) { - if (D % c.TileK1 != 0) return false; // mlp1/mlp2_t/mlp5 K - if (D % (2 * c.TileN1) != 0) return false; // mlp5 N (= TileN2) - if (I % (2 * c.TileN1) != 0) return false; // mlp2_t N - if (I % c.TileK1 != 0) return false; // mlp5 K - if (I % c.TileK3 != 0) return false; // mlp3/mlp4 K - } - if (D % 8 != 0 || I % 8 != 0) return false; // vectorization - if (c.Compute == 100) { - if (D % 256 != 0 || I % 256 != 0) return false; - } else { - if (!((c.TileM3 == 256 && c.TileN3 == 128) || - (c.TileM3 == 128 && c.TileN3 == 256))) return false; - if (D % c.TileM3 != 0 || I % c.TileN3 != 0) return false; - } - return true; -} - -static bool find_tuned_table_bwd(int compute, int n_pes, - const TunedConfigFwdBwd** table, - int* table_n) { - const TunedConfigFwdBwdTable* tables = - (n_pes <= 1) ? kTunedConfigTablesSingle : kTunedConfigTablesMulti; - const int tables_n = - (n_pes <= 1) ? kNumTunedConfigTablesSingle : kNumTunedConfigTablesMulti; - for (int i = 0; i < tables_n; ++i) { - if (tables[i].Compute == compute) { - *table = tables[i].configs; - *table_n = tables[i].count; - return *table_n > 0; - } - } - return false; -} - -// Nearest-neighbor lookup in log-space (TK, TKE, D, I), where TK = T*top_k -// is the total routed-token count (the bwd's M-axis post-dispatch) and -// TKE = TK / E_local (E_local = E / n_pes) is the per-LOCAL-expert avg -// K-range. mlp3/mlp4 walk each PE's local experts, so adding TKE to the -// distance lets the dispatcher distinguish configs whose walk-scheduler -// load differs even when (TK, D, I) is the same. -// Configs that violate structural constraints for the runtime (D, I) are -// skipped — the nearest tuned point is not always valid for the actual shape. -// If fwd_tile_m > 0, configs with TileM != fwd_tile_m are filtered out: -// BWD's sort-output reads (tile_expert_ids granularity, expert_offsets -// alignment, x_sorted slot layout) must match the forward's TileM. -// Out-param API instead of pointer-into-static-array because each candidate -// TunedConfigBwd is projected on the fly from kTunedConfigsFwdBwd (no stable -// address). Returns true when a valid config was found. -static bool find_nearest_tuned_bwd(int compute, int TK, int TKE, int D, int I, - int fwd_tile_m, int n_pes, - TunedConfigBwd* out) { - // World-size class: single-GPU (n_pes <= 1) and multi-GPU tables are tuned - // and stored separately (mirrors find_nearest_tuned on the fwd side). - const TunedConfigFwdBwd* table = nullptr; - int table_n = 0; - if (!find_tuned_table_bwd(compute, n_pes, &table, &table_n)) return false; - bool found = false; - double best_dist = 1e30; - double lTK = std::log((double)TK); - double lTKE = std::log((double)TKE); - double lD = std::log((double)D); - double lI = std::log((double)I); - for (int i = 0; i < table_n; ++i) { - TunedConfigBwd c = project_bwd(table[i], compute); - // Forward and backward select the same architecture-specific comm tile - // (M128 on both architectures), independent of the tuned GEMM tile. - (void)fwd_tile_m; - if (!tuned_config_valid_bwd(D, I, c)) continue; - // Guard against TKE=0 in legacy rows (would yield -inf in log). - double dTK = std::log((double)c.TK) - lTK; - double dTKE = c.TKE > 0 - ? std::log((double)c.TKE) - lTKE - : 0.0; // legacy / placeholder rows contribute 0 to the TKE axis - double dD = std::log((double)c.D) - lD; - double dI = std::log((double)c.I) - lI; - double d = dTK*dTK + dTKE*dTKE + dD*dD + dI*dI; - if (d < best_dist) { best_dist = d; *out = c; found = true; } - } - return found; -} - -// ============================================================================ -// Auto-dispatch wrapper — pure tuned-table lookup -// ============================================================================ -// -// Picks the nearest tuned config (log-space distance over (TK, D, I)) from -// moe_fwd_bwd_tuning_configs.cuh and dispatches to the matching template -// instantiation in kDispatchTableBwd. Both tables are populated by the -// tune_moe_bwd sweep — shapes outside the swept range error out instead of -// falling back to a heuristic. -// -// fwd_tile_m: BWD's TileM must match the forward's TileM. Forward sort -// pads expert ranges to FWD's TileM, populates tile_expert_ids at that -// granularity, and writes x_sorted in FWD-TileM-aligned slots. BWD -// reads all three at the same granularity, so a mismatch corrupts both -// the LocalIter expert lookup and the symmetric-memory bounds. -// Callers must pass the FWD's TileM (= 128 for the default fwd path). -// -// Override path: LIGER_MOE_BWD_FORCE_CONFIG=NS2,TN1,TK1,S1,TM3,TN3,TK3,S3,EN1,EN25,EN34,CS,TM -// Lets a test suite sweep every compiled config without rebuilding. - -// Internal auto-dispatch: nearest tuned bwd config for the runtime shape (or -// LIGER_MOE_BWD_FORCE_CONFIG), then invoke the matching template instantiation. -// All shape info comes from the MoeBwdArgs bundle; the launcher writes the grads. -void moe_bwd_dispatch(const MoeBwdArgs& a, int fwd_tile_m) { - const int T = a.num_tokens; - const int D = a.hidden_dim; - const int I = a.intermediate_dim; - const int top_k = a.top_k; - const int num_experts = a.num_experts; - const int TK = T * top_k; - const int n_pes = nvshmem_team_n_pes(a.team); - const int E_local = (n_pes > 0) ? std::max(1, num_experts / n_pes) : num_experts; - // Keep backward lookup consistent with forward for tiny token batches. - const int TKE = std::max(1, TK / E_local); - const int compute = moe_detect_compute_dispatch_key(a.device, "moe_fused_bwd_bf16_auto"); - if (const char* s = std::getenv("LIGER_MOE_BWD_FORCE_CONFIG")) { - TunedConfigBwd tc{}; - tc.Compute = compute; - tc.GemmTileM = -1; // sentinel → default to comm TileM after parse - int vals[14] = {}; - int n = std::sscanf(s, "%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d", - &vals[0], &vals[1], &vals[2], &vals[3], &vals[4], &vals[5], &vals[6], - &vals[7], &vals[8], &vals[9], &vals[10], &vals[11], &vals[12], &vals[13]); - LIGER_CHECK(n == 13 || n == 14, - "LIGER_MOE_BWD_FORCE_CONFIG must have 13 or 14 comma-separated ints " - "(NS2,TN1,TK1,S1,TM3,TN3,TK3,S3,EN1,EN25,EN34,CS,TM[,GemmTM]); got '", s, "'"); - tc.NSplit2 = vals[0]; - tc.TileN1 = vals[1]; - tc.TileK1 = vals[2]; - tc.Stages1 = vals[3]; - tc.TileM3 = vals[4]; - tc.TileN3 = vals[5]; - tc.TileK3 = vals[6]; - tc.Stages3 = vals[7]; - tc.EpiChunkN1 = vals[8]; - tc.EpiChunkN25 = vals[9]; - tc.EpiChunkN34 = vals[10]; - tc.CommNumStages = vals[11]; - tc.TileM = vals[12]; - tc.GemmTileM = (n == 14) ? vals[13] : tc.TileM; - if (tc.GemmTileM <= 0) tc.GemmTileM = tc.TileM; // omitted → gemm == comm - // The wrapper selects the architecture-specific comm tile consistently - // with forward; the parsed TileM remains the GEMM tuning field. - (void)fwd_tile_m; - const DispatchEntryBwd* de = find_dispatch_entry_bwd(tc); - LIGER_CHECK(de != nullptr, - "LIGER_MOE_BWD_FORCE_CONFIG=", s, " is not in the compiled dispatch " - "table. Add it to LIGER_MOE_BWD_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); - de->fn(a, kBwdRuntimeNsplitSeed); - return; - } - - TunedConfigBwd tc{}; - LIGER_CHECK(find_nearest_tuned_bwd(compute, TK, TKE, D, I, fwd_tile_m, n_pes, &tc), - "moe_bwd_fwd_bf16_auto: no valid tuned config for compute=", compute, - " shape (T=", T, - " K=", top_k, " TK=", TK, " TKE=", TKE, " D=", D, " I=", I, - " fwd_tile_m=", fwd_tile_m, - "). Either moe_fwd_bwd_tuning_configs.cuh has no entry with TileM=", - fwd_tile_m, ", or every matching entry violates structural " - "constraints for this (D, I). Add a compatible config to " - "LIGER_MOE_BWD_TUNE_CONFIGS and re-tune."); - - const DispatchEntryBwd* de = find_dispatch_entry_bwd(tc); - LIGER_CHECK(de != nullptr, - "moe_bwd_fwd_bf16_auto: tuned config (Compute=", tc.Compute, - " NS2=", tc.NSplit2, - " TN1=", tc.TileN1, "/", tc.TileK1, "/", tc.Stages1, - " TM3=", tc.TileM3, " TN3=", tc.TileN3, "/", tc.TileK3, - "/", tc.Stages3, - " EN1=", tc.EpiChunkN1, - " EN25=", tc.EpiChunkN25, " EN34=", tc.EpiChunkN34, - " CS=", tc.CommNumStages, " TM=", tc.TileM, - ") for shape (T=", T, " K=", top_k, " TK=", TK, " D=", D, " I=", I, - ") is not in the compiled dispatch table. " - "Add it to LIGER_MOE_BWD_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); - - de->fn(a, kBwdRuntimeNsplitSeed); -} - -// Returns the compiled bwd config table as a list of 13-tuples -// (Compute, NSplit2, TileN1, TileK1, Stages1, TileM3, TileN3, TileK3, -// Stages3, EpiChunkN1, EpiChunkN25, EpiChunkN34, CommNumStages). -std::vector> -moe_bwd_list_compiled_configs() { - std::vector> out; - out.reserve(kNumDispatchEntriesBwd); - for (int i = 0; i < kNumDispatchEntriesBwd; ++i) { - const auto& e = kDispatchTableBwd[i]; - out.emplace_back(e.Compute, e.NSplit2, - e.TileN1, e.TileK1, e.Stages1, - e.TileM3, e.TileN3, e.TileK3, e.Stages3, - e.EpiChunkN1, e.EpiChunkN25, e.EpiChunkN34, - e.CommNumStages); - } - return out; -} - -} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_bwd.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_bwd.cuh index 726c9d4c..8670c4ed 100644 --- a/liger_cute_kernels/csrc/core/src/moe/moe_bwd.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/moe_bwd.cuh @@ -1,219 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// Fused MoE Backward: MLP compute for local + remote tiles -// ═══════════════════════════════════════════════════════════════════ -// -// MLP-only device function — called by non-comm warps. In BWD comm uses -// warps 1-2; warp 3 enters the MLP path on SM100 and exits immediately on SM90. -// -// Execution flow (UNIFIED — local tile loop removed, mirrors moe_fused_fwd): -// EVERY tile (local or remote) flows through the comm get→staging→MLP→put -// path, driven by mlp_fused_bwd_dual over the RemoteMlpTileIteratorBwd / -// staging pipeline (there is no separate materialized local pass). -// -// mlp_fused_bwd_dual runs the grouped batch loop: -// SubBatch × (Phase 1 tile) → Phase 2 weight grads over the grouped window -// -// ═══════════════════════════════════════════════════════════════════ +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -#include "mlp_bwd.cuh" - -namespace liger { - -using namespace cute; - -// ═══════════════════════════════════════════════════════════════════ -// MoE backward shared memory -// ═══════════════════════════════════════════════════════════════════ - -template -struct MoeBwdSmem { - MlpFusedBwdSmem mlp; - CommSmem comm; -}; - -// ═══════════════════════════════════════════════════════════════════ -// moe_fused_bwd — local phase → barrier → remote phase -// ═══════════════════════════════════════════════════════════════════ - -template 1 → Phase-1 (mlp1a/mlp2t/ - // silu/mlp5) steps through SubTiles GemmTileM-row sub-tiles per - // comm staging slot; the mlp3/4 ratios use Traits1::TileM (= GemmTileM). - int SubTiles, - int Compute = 90, - typename FusedIter, - // Phase 1 TMA - typename TmaLoadX, typename TmaLoadW1, - typename TmaStoreZ, typename TmaStoreU, typename TmaStoreV, - typename TmaLoadDY, typename TmaLoadW2T, typename TmaStoreDZ, - typename TmaLoadDU, typename TmaLoadDV, - typename TmaLoadB, typename TmaLoadC, typename TmaStoreDX, - // Phase 2 TMA - typename TmaLoadDYT3, typename TmaLoadZT3, typename TmaStoredA, - typename TmaLoadXT4, typename TmaLoaddUT4, typename TmaLoaddVT4, - typename TmaStoredB, typename TmaStoredC, - // Remote TMA: SAME C++ types as their local counterparts - // (the cute Layouts for both x_sorted and the staging ring use - // dynamic shape + _1 col stride + identical smem layout), so we - // can reuse the TmaLoadX/DY/StoreDX/etc. template params. Only - // the runtime gmem pointer differs between local and remote. - typename TmaLoadXR = TmaLoadX, - typename TmaLoadDYR = TmaLoadDY, - typename TmaStoreDXR = TmaStoreDX, - typename TmaLoadDYT3R = TmaLoadDYT3, - typename TmaLoadXT4R = TmaLoadXT4> -__device__ __forceinline__ void moe_fused_bwd( - MoeBwdSmem& smem, - FusedIter& iter, - // Phase 1 TMA (local) - TmaLoadX const& tma_load_x, TmaLoadW1 const& tma_load_b_fwd, - TmaLoadW1 const& tma_load_c_fwd, TmaStoreZ const& tma_store_z, - TmaStoreU const& tma_store_u, TmaStoreV const& tma_store_v, - TmaLoadDY const& tma_load_dy, TmaLoadW2T const& tma_load_a_col, - TmaStoreDZ const& tma_store_dz, - TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, - TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, - TmaStoreDX const& tma_store_dx, - // Phase 2 TMA - TmaLoadDYT3 const& tma_load_dyt3, TmaLoadZT3 const& tma_load_zt3, - TmaStoredA const& tma_store_da, - TmaLoadXT4 const& tma_load_xt4, - TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, - TmaStoredB const& tma_store_db, TmaStoredC const& tma_store_dc, - // Remote TMA - TmaLoadXR const& tma_load_x_remote, - TmaLoadDYR const& tma_load_dy_remote, - TmaStoreDXR const& tma_store_dx_remote, - // Remote Phase 2 TMA: transposed staging sources for dA (Phase 3) - // and dB/dC (Phase 4). Without these, the remote phase reuses the - // local-only x_sorted/dy_sorted descriptors and produces wrong - // gradients for tokens routed across PEs. - TmaLoadDYT3R const& tma_load_dyt3_remote, - TmaLoadXT4R const& tma_load_xt4_remote, - // Comm pipe pointers (for building remote iter) - int* x_src_ready, int* x_src_consumed, - int* dy_src_ready, int* dy_src_consumed, - int* dst_ready, int* dst_consumed, - // Two per-pipe expert-id arrays (see CommBuffersBwd docstring). - // remote_tile_expert_ids_x is the X-pipe array used for the iter's - // tile.expert lookup (and Phase 2 mlp4 in this remote run). - // remote_tile_expert_ids_dy is the dY-pipe array used for Phase 2 - // mlp3 — gated by release_dy, so it stays stable across mlp3 even - // when release_src has already let comm refill the X-pipe array. - const int* remote_tile_expert_ids_x, - const int* remote_tile_expert_ids_dy, - const int* remote_tile_valid_rows_x, - // Dims + buffers - const MlpBwdDims& dims, - const MlpBwdBufs& bufs, - int col, - int grid_x, - int split, - int runtime_nsplit, - bool gemm_active) { - - // Target-based barriers — constructed ONCE so the per-CTA target - // counters accumulate monotonically across: - // local pass → global_barrier.wait() → remote pass - // Constructing a fresh barrier per wait() resets target=0 and is racy. - // Flat 1-D launch coordinates. blockIdx.x = flat_id ∈ [0, num_blocks). - // global_barrier counts ALL launched CTAs (num_blocks = gridDim.x). - // x_barrier is per logical Phase-1 column (col), stride runtime_nsplit. - int flat_id = (int)blockIdx.x; - - MlpBwdCtaBarrierT global_barrier(bufs.barrier_counter, - (int)gridDim.x * (int)gridDim.y); - MlpBwdCtaBarrierT x_barrier(&dims.phase_counter[col], - runtime_nsplit); - - // Initialize the fused iter's remote sub-iter from comm pipe pointers. - // (Local sub-iter was init'd by the caller in moe_bwd_kernel.) Comm warps - // filled staging during this CTA's local pass. Ticket-based iterator - // keys on absolute slot index, so pointer offsets are NOT pre-applied - // (mirrors RemoteMlpTileIterator's init in moe.cuh). - iter.init_remote( - remote_tile_expert_ids_x, - remote_tile_valid_rows_x, - smem.comm.per_cta_tiles, - col, - dims.n_gemm, - x_src_ready, x_src_consumed, - dy_src_ready, dy_src_consumed, - dst_ready, dst_consumed, - /*is_leader=*/(threadIdx.x == 0), - runtime_nsplit, - /*tma_enabled=*/(dims.tma_get_enabled != 0)); - - // Remote dims override: num_tokens / num_m_tiles use the staging- - // ring shape; total_m_tiles_in_pass uses the actual remote tile - // count (allows batches ≥ NumStages to run with bk_start wrapped - // into the live ring window). local_expert_start = 0 because - // RemoteIter reads tile_expert_ids written by nvshmem_comm_main_bwd - // (= TileIterator's LOCAL expert index). - MlpBwdDims remote_dims = dims; - remote_dims.num_tokens = dims.remote_num_tokens; - remote_dims.num_m_tiles = dims.remote_num_m_tiles; - remote_dims.total_m_tiles_in_pass = smem.comm.global_total; - remote_dims.num_batches = - (smem.comm.global_total + grid_x - 1) / grid_x; - remote_dims.local_expert_start = 0; - - // Alias per-pipe expert-id arrays into MlpBwdBufs. Phase 2 mlp4 - // reads X-pipe gated array (stable across mlp4 because release_src - // is deferred past mlp4). Phase 2 mlp3 reads dY-pipe gated array. - MlpBwdBufs remote_bufs = bufs; - remote_bufs.expert_for_k_block_mlp4 = remote_tile_expert_ids_x; - remote_bufs.expert_for_k_block_mlp3 = remote_tile_expert_ids_dy; - remote_bufs.valid_rows_mlp4 = remote_tile_valid_rows_x; - - // Grid-uniform activity predicate. The unified path stages local tiles even - // at world size one, so every valid non-empty launch is active. - bool remote_active = - (dims.num_pes > 0) && (dims.num_tokens > 0); - - // The remote staging ring (L = MC·CommNumStages tiles, MC = n_gemm/NC) is a - // STREAMING buffer: when global_total > L it wraps and physical slots are - // reused across laps. This is correct — mlp_fused_bwd_dual rescans the live - // per-slot expert ids PER BATCH (see its loop), so reused slots are always - // attributed to the expert currently resident, and each slot's reduction is - // committed before comm refills it. No L ≥ global_total requirement. - - // UNIFIED PATH: there is no separate materialized local MLP pass. The shared - // TileIterator enumerates LOCAL experts (p=0) alongside remote ones; the comm - // get warps stage local X/dY from local symmetric memory and the put warps - // drain local dX back to local dx_sorted, so EVERY tile (local or remote) - // flows through the RemoteMlpTileIteratorBwd / staging pipeline that - // mlp_fused_bwd_dual drives. Only the staging-ring (remote) TMAs and dims/bufs - // are passed; the weight/intermediate descriptors are shared across all tiles. - // efkb_stride maps a (GemmTileM-granular) Phase-2 K-block to its parent - // COMM-tile slot in tile_expert_ids (which the comm warps emit once per slot). - // = CommTileM/TileK = Traits1::TileM·SubTiles/Traits4::TileK. The integer - // division floors each GEMM sub-tile to its parent communication slot's expert (the - // interpolation). SubTiles=1 → Traits1::TileM/TileK, unchanged. - constexpr int kEfkbStride4Remote = Traits1::TileM * SubTiles / Traits4::TileK; - mlp_fused_bwd_dual< - Traits1, Traits2T, Traits3, Traits4, Traits5, NSplit2, SubBatch, SubTiles, Compute>( - smem.mlp, iter, remote_active, - kEfkbStride4Remote, - // Phase 1 shared weights/intermediates - tma_load_b_fwd, tma_load_c_fwd, - tma_store_z, tma_store_u, tma_store_v, - tma_load_a_col, tma_store_dz, - tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, - // Phase 2 shared - tma_load_zt3, tma_store_da, - tma_load_dut4, tma_load_dvt4, tma_store_db, tma_store_dc, - // Staging-ring data TMAs (X/dY/dX, dYᵀ/Xᵀ) - tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, - tma_load_dyt3_remote, tma_load_xt4_remote, - remote_dims, remote_bufs, - global_barrier, x_barrier, - flat_id, col, grid_x, split, /*num_splits=*/runtime_nsplit, gemm_active); -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "moe_bwd_sm90.cuh" +#else +#include "moe_bwd_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.cuh new file mode 100644 index 00000000..e81fa5d9 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.cuh @@ -0,0 +1,219 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MoE Backward: MLP compute for local + remote tiles +// ═══════════════════════════════════════════════════════════════════ +// +// MLP-only device function — called by non-comm warps. In BWD comm uses +// warps 1-2; warp 3 enters the MLP path on SM100 and exits immediately on SM90. +// +// Execution flow (UNIFIED — local tile loop removed, mirrors moe_fused_fwd): +// EVERY tile (local or remote) flows through the comm get→staging→MLP→put +// path, driven by mlp_fused_bwd_dual over the RemoteMlpTileIteratorBwd / +// staging pipeline (there is no separate materialized local pass). +// +// mlp_fused_bwd_dual runs the grouped batch loop: +// SubBatch × (Phase 1 tile) → Phase 2 weight grads over the grouped window +// +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp_bwd_sm100.cuh" + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// MoE backward shared memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct MoeBwdSmem { + MlpFusedBwdSmem mlp; + CommSmem comm; +}; + +// ═══════════════════════════════════════════════════════════════════ +// moe_fused_bwd — local phase → barrier → remote phase +// ═══════════════════════════════════════════════════════════════════ + +template 1 → Phase-1 (mlp1a/mlp2t/ + // silu/mlp5) steps through SubTiles GemmTileM-row sub-tiles per + // comm staging slot; the mlp3/4 ratios use Traits1::TileM (= GemmTileM). + int SubTiles, + int Compute = 90, + typename FusedIter, + // Phase 1 TMA + typename TmaLoadX, typename TmaLoadW1, + typename TmaStoreZ, typename TmaStoreU, typename TmaStoreV, + typename TmaLoadDY, typename TmaLoadW2T, typename TmaStoreDZ, + typename TmaLoadDU, typename TmaLoadDV, + typename TmaLoadB, typename TmaLoadC, typename TmaStoreDX, + // Phase 2 TMA + typename TmaLoadDYT3, typename TmaLoadZT3, typename TmaStoredA, + typename TmaLoadXT4, typename TmaLoaddUT4, typename TmaLoaddVT4, + typename TmaStoredB, typename TmaStoredC, + // Remote TMA: SAME C++ types as their local counterparts + // (the cute Layouts for both x_sorted and the staging ring use + // dynamic shape + _1 col stride + identical smem layout), so we + // can reuse the TmaLoadX/DY/StoreDX/etc. template params. Only + // the runtime gmem pointer differs between local and remote. + typename TmaLoadXR = TmaLoadX, + typename TmaLoadDYR = TmaLoadDY, + typename TmaStoreDXR = TmaStoreDX, + typename TmaLoadDYT3R = TmaLoadDYT3, + typename TmaLoadXT4R = TmaLoadXT4> +__device__ __forceinline__ void moe_fused_bwd( + MoeBwdSmem& smem, + FusedIter& iter, + // Phase 1 TMA (local) + TmaLoadX const& tma_load_x, TmaLoadW1 const& tma_load_b_fwd, + TmaLoadW1 const& tma_load_c_fwd, TmaStoreZ const& tma_store_z, + TmaStoreU const& tma_store_u, TmaStoreV const& tma_store_v, + TmaLoadDY const& tma_load_dy, TmaLoadW2T const& tma_load_a_col, + TmaStoreDZ const& tma_store_dz, + TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, + TmaStoreDX const& tma_store_dx, + // Phase 2 TMA + TmaLoadDYT3 const& tma_load_dyt3, TmaLoadZT3 const& tma_load_zt3, + TmaStoredA const& tma_store_da, + TmaLoadXT4 const& tma_load_xt4, + TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, + TmaStoredB const& tma_store_db, TmaStoredC const& tma_store_dc, + // Remote TMA + TmaLoadXR const& tma_load_x_remote, + TmaLoadDYR const& tma_load_dy_remote, + TmaStoreDXR const& tma_store_dx_remote, + // Remote Phase 2 TMA: transposed staging sources for dA (Phase 3) + // and dB/dC (Phase 4). Without these, the remote phase reuses the + // local-only x_sorted/dy_sorted descriptors and produces wrong + // gradients for tokens routed across PEs. + TmaLoadDYT3R const& tma_load_dyt3_remote, + TmaLoadXT4R const& tma_load_xt4_remote, + // Comm pipe pointers (for building remote iter) + int* x_src_ready, int* x_src_consumed, + int* dy_src_ready, int* dy_src_consumed, + int* dst_ready, int* dst_consumed, + // Two per-pipe expert-id arrays (see CommBuffersBwd docstring). + // remote_tile_expert_ids_x is the X-pipe array used for the iter's + // tile.expert lookup (and Phase 2 mlp4 in this remote run). + // remote_tile_expert_ids_dy is the dY-pipe array used for Phase 2 + // mlp3 — gated by release_dy, so it stays stable across mlp3 even + // when release_src has already let comm refill the X-pipe array. + const int* remote_tile_expert_ids_x, + const int* remote_tile_expert_ids_dy, + const int* remote_tile_valid_rows_x, + // Dims + buffers + const MlpBwdDims& dims, + const MlpBwdBufs& bufs, + int col, + int grid_x, + int split, + int runtime_nsplit, + bool gemm_active) { + + // Target-based barriers — constructed ONCE so the per-CTA target + // counters accumulate monotonically across: + // local pass → global_barrier.wait() → remote pass + // Constructing a fresh barrier per wait() resets target=0 and is racy. + // Flat 1-D launch coordinates. blockIdx.x = flat_id ∈ [0, num_blocks). + // global_barrier counts ALL launched CTAs (num_blocks = gridDim.x). + // x_barrier is per logical Phase-1 column (col), stride runtime_nsplit. + int flat_id = (int)blockIdx.x; + + MlpBwdCtaBarrierT global_barrier(bufs.barrier_counter, + (int)gridDim.x * (int)gridDim.y); + MlpBwdCtaBarrierT x_barrier(&dims.phase_counter[col], + runtime_nsplit); + + // Initialize the fused iter's remote sub-iter from comm pipe pointers. + // (Local sub-iter was init'd by the caller in moe_bwd_kernel.) Comm warps + // filled staging during this CTA's local pass. Ticket-based iterator + // keys on absolute slot index, so pointer offsets are NOT pre-applied + // (mirrors RemoteMlpTileIterator's init in moe.cuh). + iter.init_remote( + remote_tile_expert_ids_x, + remote_tile_valid_rows_x, + smem.comm.per_cta_tiles, + col, + dims.n_gemm, + x_src_ready, x_src_consumed, + dy_src_ready, dy_src_consumed, + dst_ready, dst_consumed, + /*is_leader=*/(threadIdx.x == 0), + runtime_nsplit, + /*tma_enabled=*/(dims.tma_get_enabled != 0)); + + // Remote dims override: num_tokens / num_m_tiles use the staging- + // ring shape; total_m_tiles_in_pass uses the actual remote tile + // count (allows batches ≥ NumStages to run with bk_start wrapped + // into the live ring window). local_expert_start = 0 because + // RemoteIter reads tile_expert_ids written by nvshmem_comm_main_bwd + // (= TileIterator's LOCAL expert index). + MlpBwdDims remote_dims = dims; + remote_dims.num_tokens = dims.remote_num_tokens; + remote_dims.num_m_tiles = dims.remote_num_m_tiles; + remote_dims.total_m_tiles_in_pass = smem.comm.global_total; + remote_dims.num_batches = + (smem.comm.global_total + grid_x - 1) / grid_x; + remote_dims.local_expert_start = 0; + + // Alias per-pipe expert-id arrays into MlpBwdBufs. Phase 2 mlp4 + // reads X-pipe gated array (stable across mlp4 because release_src + // is deferred past mlp4). Phase 2 mlp3 reads dY-pipe gated array. + MlpBwdBufs remote_bufs = bufs; + remote_bufs.expert_for_k_block_mlp4 = remote_tile_expert_ids_x; + remote_bufs.expert_for_k_block_mlp3 = remote_tile_expert_ids_dy; + remote_bufs.valid_rows_mlp4 = remote_tile_valid_rows_x; + + // Grid-uniform activity predicate. The unified path stages local tiles even + // at world size one, so every valid non-empty launch is active. + bool remote_active = + (dims.num_pes > 0) && (dims.num_tokens > 0); + + // The remote staging ring (L = MC·CommNumStages tiles, MC = n_gemm/NC) is a + // STREAMING buffer: when global_total > L it wraps and physical slots are + // reused across laps. This is correct — mlp_fused_bwd_dual rescans the live + // per-slot expert ids PER BATCH (see its loop), so reused slots are always + // attributed to the expert currently resident, and each slot's reduction is + // committed before comm refills it. No L ≥ global_total requirement. + + // UNIFIED PATH: there is no separate materialized local MLP pass. The shared + // TileIterator enumerates LOCAL experts (p=0) alongside remote ones; the comm + // get warps stage local X/dY from local symmetric memory and the put warps + // drain local dX back to local dx_sorted, so EVERY tile (local or remote) + // flows through the RemoteMlpTileIteratorBwd / staging pipeline that + // mlp_fused_bwd_dual drives. Only the staging-ring (remote) TMAs and dims/bufs + // are passed; the weight/intermediate descriptors are shared across all tiles. + // efkb_stride maps a (GemmTileM-granular) Phase-2 K-block to its parent + // COMM-tile slot in tile_expert_ids (which the comm warps emit once per slot). + // = CommTileM/TileK = Traits1::TileM·SubTiles/Traits4::TileK. The integer + // division floors each GEMM sub-tile to its parent communication slot's expert (the + // interpolation). SubTiles=1 → Traits1::TileM/TileK, unchanged. + constexpr int kEfkbStride4Remote = Traits1::TileM * SubTiles / Traits4::TileK; + mlp_fused_bwd_dual< + Traits1, Traits2T, Traits3, Traits4, Traits5, NSplit2, SubBatch, SubTiles, Compute>( + smem.mlp, iter, remote_active, + kEfkbStride4Remote, + // Phase 1 shared weights/intermediates + tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_u, tma_store_v, + tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, + // Phase 2 shared + tma_load_zt3, tma_store_da, + tma_load_dut4, tma_load_dvt4, tma_store_db, tma_store_dc, + // Staging-ring data TMAs (X/dY/dX, dYᵀ/Xᵀ) + tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, + tma_load_dyt3_remote, tma_load_xt4_remote, + remote_dims, remote_bufs, + global_barrier, x_barrier, + flat_id, col, grid_x, split, /*num_splits=*/runtime_nsplit, gemm_active); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.inc b/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.inc new file mode 100644 index 00000000..98cbb69f --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.inc @@ -0,0 +1,2091 @@ +// Fused MoE Backward: sort → fwd-dispatch (for X) → reverse-combine (for dY) +// → MLP_bwd + comm overlap → reverse-dispatch (for dX). +// Compiled with separable compilation. Do NOT link liger_comm_kernels library. +// +// Torch-free core (liger_cute): ported from LigerCommKernels' moe_bwd.cu. The +// device kernel + templated launcher are upstream-verbatim (namespace liger). +// Tensor-carrying backward calls are exported through TVM FFI and lower directly +// into MoeBwdArgs. +#define LIGER_CUTE_BUILDING 1 + +#include +#include // pulls in nvshmemi_device_state_d, required by the + // inline device helpers in . +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +// SM100 2SM (paired-CTA cluster) TMA descriptor factories (make_tma_copy_A_sm100 +// / make_tma_copy_B_sm100) and the SM100_TMA_2SM_LOAD copy atom used below to +// build pair-aware Phase-2 mlp3/mlp4 descriptors. Already pulled in +// transitively via moe_bwd.cuh -> mlp_bwd.cuh -> mlp3.cuh/mlp4.cuh; included +// directly here too since this TU calls these factories itself. +#include + +#include "liger_cute/moe.h" +#include "liger_cute/check.h" +#include "liger_cute/detail/status.h" + +#include "moe_sm100.cuh" +#include "moe_bwd_sm100.cuh" +#include "moe_symm_config_sm100.cuh" +#include "moe_launch.h" // MoeFwdArgs / MoeBwdArgs (shared with the tuner) +#include "moe_utils.cuh" +#include "moe_fwd_bwd_tune_configs_sm100.hpp" +#include "moe_dispatch_configs_sm100.cuh" +#include "mlp_comms_bwd_sm100.cuh" +#include "tile_iterator_bwd_sm100.cuh" +#include "dispatch.cuh" +#include "dispatch_bwd.cuh" +#include "combine_bwd.cuh" +#include "tma_copy_atom_sm100.cuh" +#include "radix_sort.cuh" +#include "buffer_pool.cuh" // shim: liger:: aliases of liger_cute::detail infra + +using namespace cute; + +namespace liger { + +// 0 = build every dispatch family (developer fallback). CMake sets this to 90 +// for Hopper builds and 100 for Blackwell builds so normal native builds do not +// instantiate kernels for hardware that cannot run in that binary. +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif + +// External symbol: moe_configure_symmetric (defined in moe.cu) is reused. +// MoeSymmConfig + get_symm_config() now come from the shared header (included +// below), so this TU can't drift from moe.cu's layout. + +// Global iteration counter — same role as the fwd counter in moe.cu. +// Lives in device memory (allocated from BufferPool under the key +// "moe_bwd_iter_counter"), atomicAdded at the end of each kernel by +// CTA 0 thread 0. See PeSync::advance_iteration_counter() and the +// detailed comment in mlp_comms.cuh for why the counter has to be +// device-resident: a host-side atomic would be baked into a captured +// CUDA graph as a constant, breaking the cross-PE wait-until-greater +// protocol on every replay. + +// ============================================================================ +// MoeBwdConfig — bundles bwd-side traits + comm/grid constants +// ============================================================================ +// +// All five MLP_bwd phase traits + grid params. Mirrors MoeFusedConfig. +// Trait choices match the values that mlp_bwd.cu landed on after the +// 2026-04-26 perf tuning round (see project_fused_bwd_tile_bump.md). + +// Trait defaults match mlp_bwd.cu's validated set (lines 28–33). MoE +// backward shares mlp_fused_bwd; any drift from those values has hit +// illegal-instruction / smem-overrun bugs in the past. +template < + int NSplit2_ = 4, // Phase 2 (mlp3/mlp4) virtual n-split (autotuning surface) + // SM100 uses Stage-3 Phase-1 traits to leave room for independent full-size + // X/dY TMA bounce buffers. SM90 dispatch rows provide their tuned values. + int TileM1_ = 128, + int TileN1_ = 128, + int TileK1_ = 64, + int Stages1_ = 3, + int EpiChunkN1_ = 32, // mlp1_act + int TileN2_ = 256, + int TileK2_ = 64, + int Stages2_ = 3, + int TileM2T_ = 128, // Phase 1b' (mlp2_t): TileM ∈ {64, 128} + // mlp3/mlp4 cooperative layout (#102) supports only (256,128) or (128,256) + // and needs Stages=2 to fit the 228 KiB smem cap — Stages=4 at TileN=256 + // alone is ~384 KiB. (256,128) M-split is the smallest-smem legal choice + // (~192 KiB; the binding floor is mlp1_act ≈ 216 KiB). + int TileM3_ = 256, // mlp3 cooperative: (256,128) M-split + int TileN3_ = 128, + int TileK3_ = 64, + int Stages3_ = 2, + int TileM4_ = 128, // mlp4 cooperative: swap of mlp3 → (128,256) + int TileN4_ = 256, + int TileK4_ = 64, + int Stages4_ = 2, + int TileM5_ = 128, + int TileN5_ = 256, + int TileK5_ = 64, + int Stages5_ = 3, + int EpiChunkN25_ = 64, // shared by mlp2_t and mlp5 (same smem shape) + int EpiChunkN34_ = 64, // shared by mlp3 and mlp4 (same smem shape) + typename Element_ = bfloat16_t, + int CommNumStages_ = 2, + int NC_ = 4, + // Phase-1 sub-batch grouping count: run SubBatch Phase-1 sub-batches + // (mlp1/2/5) before one Phase-2 (mlp3/4) over their combined K-window. + // 1 = original per-batch loop. The only structural requirement is a ring + // deeper than the group span (CommNumStages·mc_remote > SubBatch·grid_x); + // the grouped window wraps the ring exactly as the un-grouped loop does. + int SubBatch_ = 2, + // CommTileM_ decouples the COMM staging tile (ring slots, semaphores, + // tile_expert_ids granularity) from the Phase-1 GEMM tile (Traits1::TileM = + // mlp1a/mlp2t/mlp5 = the GemmTileM hyperparameter ∈ {64,128}). When + // Traits1::TileM < CommTileM_, Phase-1 steps through + // kSubTiles = CommTileM_/Traits1::TileM sub-tiles of GemmTileM rows per + // communication slot (mirrors forward). Default == TileM1_ → kSubTiles=1 + // → bit-identical to the old coupled path. + int CommTileM_ = 128, + int Compute_ = 90> +struct MoeBwdConfig { + using Element = Element_; + static constexpr int kActStoreTileM = 64; + using Traits1 = Mlp1Traits < + Element, TileM1_, TileN1_, TileK1_, Stages1_, EpiChunkN1_, + kActStoreTileM>; // Phase 1a + using Traits2T = Mlp2TTraits; // Phase 1b' + // Blackwell always uses the paired-CTA MLP3/MLP4 path. Hopper keeps the + // existing single-CTA traits and never instantiates SM100-only machinery. + // The existing TileK3_/Stages3_/EpiChunkN34_ fields remain the shared + // tuning surface for both weight-gradient kernels. + static constexpr bool kUsesTwoSm = Compute_ == 100; + static constexpr int kClusterM = kUsesTwoSm ? 2 : 1; + using Traits3 = cute::conditional_t< + kUsesTwoSm, + Mlp3Traits2Sm, + Mlp3Traits >; // Phase 2 mlp3 + using Traits4 = cute::conditional_t< + kUsesTwoSm, + Mlp4Traits2Sm, + Mlp4Traits >; // Phase 2 mlp4 + using Traits5 = Mlp5Traits ; // Phase 1d + + // No divisibility constraints between NSplit/NC/CommNumStages. + // MLP side is ticket-based (RemoteMlpTileIteratorBwd); comm side + // uses K = CommNumStages. Only runtime requirement is that NC + // evenly divides gridDim.x · gridDim.y so MC is integer. + + static constexpr int kNSplit2 = NSplit2_; + static constexpr int kNC = NC_; + static constexpr int kCommNumStages = CommNumStages_; + static constexpr int kSubBatch = SubBatch_; + // kTileM is the COMM staging tile (ring slots / semaphores / tile_expert_ids + // granularity), decoupled from the Phase-1 GEMM tile. kGemmTileM = + // Traits1::TileM (mlp1/2/5 GEMM tile); + // the mlp3/4 ratios (efkb_stride / bk↔Kblock / ring_kb) are computed from it. + static constexpr int kTileM = CommTileM_; + static constexpr int kGemmTileM = Traits1::TileM; // Phase-1 gemm tile (64/128) + static constexpr int kSubTiles = CommTileM_ / Traits1::TileM; + static constexpr int kCompute = Compute_; + static constexpr int kNumThreads = 384; + + static_assert(CommTileM_ % Traits1::TileM == 0, + "CommTileM must be an integer multiple of the Phase-1 GEMM TileM (Traits1::TileM)"); + static_assert(kSubTiles == 1 || kSubTiles == 2, + "CommTileM/GemmTileM must be 1 or 2"); + // mlp1a/mlp2t/mlp5 must share the same Phase-1 GEMM tile so their sub-tiling + // and the Z/dU/dV/dZ buffer rows stay aligned. + static_assert(Traits1::TileM == Traits2T::TileM && Traits1::TileM == Traits5::TileM, + "Phase-1 GEMM TileM must match across mlp1a / mlp2t / mlp5"); + + // Compile-time smem budget. H100 dynamic-smem cap is 228 KiB. The bwd + // union sits very close to that ceiling (Mlp3/Mlp4 with TileK=128 are + // 224 KiB on their own, and pipe storage pushes the total a bit + // further), so we use the full hardware cap here. The runtime-sized + // remote_offsets region appended by the launcher is ≤ a few hundred + // bytes for autotuner-scale E and num_pes; relying on the kernel's + // cudaFuncSetAttribute call to surface any overflow. + static constexpr size_t kSmemBytes = + sizeof(MoeBwdSmem); + static constexpr size_t kSmemBudget = 228 * 1024; + static_assert(kSmemBytes <= kSmemBudget, + "MoE fused bwd kernel smem footprint exceeds H100 228 KiB cap — " + "reduce TileN/TileK/Stages of one of the bwd phase traits"); + + using LocalIter = LocalMlpTileIteratorBwd; + using RemoteIter = RemoteMlpTileIteratorBwd< + Element, kCommNumStages, kNC, kTileM>; + using FusedIter = FusedMlpTileIteratorBwd< + Element, kCommNumStages, kNC, kTileM>; +}; + +// ============================================================================ +// MoeBwdParams — input/output pointers + sort buffers + comm + dims +// ============================================================================ + +template +struct MoeBwdParams { + using Element = typename Config::Element; + + // Router outputs (caller-provided). + const int* expert_indices; // [T, K] + const Element* expert_weights; // [T, K] + + // Sort buffers. + SortBuffers sort; + + // Inputs. + // Note: there is no `tokens` (X) field — BWD reads token data only + // via `x_sorted` (FWD's expert-sorted view from its symmetric stack + // entry). Removed 2026-05-18 after confirming p.tokens was set in the + // host wrapper and never read by the kernel or any helper. + const Element* dY; // [T, D] (regular device) + // Forward MoE output Y_fwd. Caller is responsible for providing this + // — either retained from the most recent forward pass or recomputed + // under gradient checkpointing. Read by combine_tokens_bwd to compute + // dW[t,k] = Σ_d dO[t,d] * Y_fwd[slot(t,k), d]. Layout: [total_slots, D] + // (one row per routed (token, expert-slot) pair, in the same + // expert-sorted order produced by the forward dispatch). + const Element* output_buffer; + + // Outputs (caller-allocated). + Element* dX_out; // [T, D] (regular device) + Element* dB_out; // [epp, I, D] + Element* dC_out; // [epp, I, D] + Element* dA_out; // [epp, D, I] + Element* dW_out; // [T, K] + + // Symmetric staging buffers (allocated by allocate_moe_bwd_buffers). + Element* x_sorted; // [total_slots, D] — fwd dispatch dst + Element* dy_sorted; // [total_slots, D] — combine_bwd dst + Element* dx_sorted; // [total_slots, D] — MLP_bwd dX dst + + // Comm (forward, for sort/dispatch/pe_sync) and bwd (for X+dY gets, dX puts). + CommBuffers comm; + CommBuffersBwd comm_bwd; + + // Dimensions. + int num_tokens; + int hidden_dim; + int intermediate_dim; + int top_k; + int total_slots; + int num_m_tiles; // sort input tiles = T / TileM + + // Output stride for dB/dC/dA per-expert. + int experts_per_pe; +}; + +// ============================================================================ +// MoeBwdBuffers — owned device + symmetric allocations +// ============================================================================ + +template +struct MoeBwdBuffers { + using Element = typename Config::Element; + + // Symmetric (NVSHMEM). + Element* x_sorted; + Element* dy_sorted; + Element* dx_sorted; + + // Symmetric offsets. + int* expert_offsets_sym; // for sort + + // Device-only. + int* cta_counter; + int* phase_counter; // mlp_bwd uses this + int* barrier_counter; // mlp_global_barrier + Element* z_buf; // [scratch_rows, I] + Element* du_buf; // [scratch_rows, I] -> dU after silu_bwd + Element* dv_buf; // [scratch_rows, I] + Element* dz_buf; // [scratch_rows, I] + int scratch_rows; // max(local padded rows, remote ring rows) + int* expert_for_k_block; // num_k_blocks; populated post-sort + + // Phase 2 per-expert K-range scratch. Filled in bwd preamble (once + // per pass) by mlp_bwd_phase_2_build_ranges; shared between mlp3 and + // mlp4 (they share TileK). Sized experts_per_pe. + int* expert_k_starts; + int* expert_k_ends; + + // Comm staging (X + dY gets, dX puts). + Element* x_staging; + Element* dy_staging; + Element* dx_staging; + + int* x_src_ready; + int* x_src_consumed; + int* dy_src_ready; + int* dy_src_consumed; + int* dst_ready; + int* dst_consumed; + // Two per-staging-slot expert-id arrays, one per get pipe. Comm warps + // write the X-side array during the X getmem and the dY-side array + // during the dY getmem. Phase 2 mlp4 reads tile_expert_ids_x; Phase 2 + // mlp3 reads tile_expert_ids_dy. See CommBuffersBwd docstring for why. + int* tile_expert_ids_x; + int* tile_expert_ids_dy; + int* tile_valid_rows_x; + + // Sort-internal scratch. + int* tile_expert_counts; + int* sorted_token_ids; + int* token_expert_slots; + int* cta_done; + int* cta_sums; +}; + +template +static MoeBwdBuffers allocate_moe_bwd_buffers( + BufferPool& pool, cudaStream_t stream, + int grid_x, int num_blocks, + int hidden_dim, int intermediate_dim, int num_tokens, int top_k, + int experts_per_pe, int total_slots, int num_m_tiles) { + + using Element = typename Config::Element; + constexpr int kTileM = Config::kTileM; + // kCommNumStages / kNSplit / kNC are not needed here: the comm staging is + // sized at the config-independent UB (kCommStagesUB · kTileMUB) below. + + MoeBwdBuffers b; + + auto& scfg = get_symm_config(); + LIGER_CHECK(scfg.initialized, + "Call moe_configure_symmetric before moe_bwd_fwd_bf16"); + + size_t symm_slot_bytes = + static_cast(scfg.max_total_slots) * scfg.hidden_dim * sizeof(Element); + + // ── Symmetric ──────────────────────────────────────── + // x_sorted comes from the FWD pass — the caller passes its symm + // pointer in the `x_sorted` torch tensor (FWD's pool entry under + // key "moe_x_sorted"). BWD reads it directly and never allocates + // or writes to its own copy. + b.dy_sorted = static_cast(pool.get_symmetric("moe_bwd_dy_sorted", symm_slot_bytes)); + b.dx_sorted = static_cast(pool.get_symmetric("moe_bwd_dx_sorted", symm_slot_bytes)); + + // dy_sorted MUST be zeroed: combine_tokens_bwd only writes routed slots + // (top_k per token). Padding slots stay 0, contributing nothing to MLP_bwd. + cudaMemsetAsync(b.dy_sorted, 0, symm_slot_bytes, stream); + + // all_expert_offsets is no longer allocated here — the caller passes + // the broadcasted offsets in via the BWD's `expert_offsets` tensor + // (populated by FWD's broadcast_expert_offsets and returned as part + // of its intermediates). + b.expert_offsets_sym = static_cast( + pool.get_symmetric("moe_sort_offsets", + (scfg.max_num_experts + 1) * sizeof(int))); + + // ── Device-only ────────────────────────────────────── + b.cta_counter = static_cast(pool.get_device("moe_bwd_cta_counter", sizeof(int))); + b.phase_counter = static_cast(pool.get_device("moe_bwd_phase_counter", grid_x * sizeof(int))); + b.barrier_counter = static_cast(pool.get_device("moe_bwd_barrier", sizeof(int))); + cudaMemsetAsync(b.cta_counter, 0, sizeof(int), stream); + cudaMemsetAsync(b.phase_counter, 0, grid_x * sizeof(int), stream); + cudaMemsetAsync(b.barrier_counter, 0, sizeof(int), stream); + + // Local Phase 1 indexes scratch by padded token row, while the remote pass + // indexes the same buffers by staging-ring slot. A deep comm ring can exceed + // the local padded extent (for example, 33 MC columns * CS8 * TileM128 = + // 33792 rows versus 17408 local rows for Mixtral-8x7B T=8K). Size for both + // coordinate spaces or a valid remote slot writes past z/du/dv/dz. + const int remote_m_tiles = + (num_blocks / Config::kNC) * Config::kCommNumStages; + const size_t remote_rows = + static_cast(remote_m_tiles) * kTileM; + const size_t scratch_rows = + std::max(static_cast(total_slots), remote_rows); + b.scratch_rows = static_cast(scratch_rows); + size_t intermediate_bytes = + scratch_rows * intermediate_dim * sizeof(Element); + b.z_buf = static_cast(pool.get_device("moe_bwd_z_buf", intermediate_bytes)); + b.du_buf = static_cast(pool.get_device("moe_bwd_du_buf", intermediate_bytes)); + b.dv_buf = static_cast(pool.get_device("moe_bwd_dv_buf", intermediate_bytes)); + b.dz_buf = static_cast(pool.get_device("moe_bwd_dz_buf", intermediate_bytes)); + + // expert_for_k_block: one int per K-block. The K-block / m-tile size + // ratio depends on Traits3::TileK vs Traits1::TileM; see + // build_expert_for_k_block which uses both dims directly. + // Size by scfg.max_total_slots, not per-call total_slots — the buffer + // pool caches by name and won't resize on later calls. Otherwise the + // first (smallest) config locks in a tiny allocation and a later + // larger config (e.g. E=64, K=8) reads OOB past the cached size. + int max_num_k_blocks = scfg.max_total_slots / Config::Traits3::TileK; + b.expert_for_k_block = static_cast( + pool.get_device("moe_bwd_expert_for_k_block", max_num_k_blocks * sizeof(int))); + + // Phase 2 per-expert K-range scratch — sized by scfg.max_num_experts + // (the pool caches by name, so we have to size to the worst case + // across configs; same rationale as expert_for_k_block above). + // Shared between mlp3 and mlp4 (they share TileK). + size_t expert_range_bytes = scfg.max_num_experts * sizeof(int); + b.expert_k_starts = static_cast( + pool.get_device("moe_bwd_expert_k_starts", expert_range_bytes)); + b.expert_k_ends = static_cast( + pool.get_device("moe_bwd_expert_k_ends", expert_range_bytes)); + + // ── Comm staging ───────────────────────────────────── + // Flat staging list — L = MC · kCommNumStages slots, where + // MC = (grid_x · NSplit) / NC. Comm strides MC through NumStages + // stages; MLP is ticket-based and keys on the absolute slot index + // modulo L. With NC = NSplit (legacy config), MC = grid_x and this + // collapses to grid_x · kCommNumStages — unchanged behavior. The live + // length MC·kCommNumStages (MC = (grid_x·NSplit)/NC) is bounded by + // comm_l_ub below; the buffers are sized at that UB, not the live value. + // Staging must be SYMMETRIC: the comm path issues nvshmemx_getmem/putmem + // whose LOCAL side (x/dy/dx staging) is the NIC-registered buffer. + // Device-only memory faults with IBV_WC_LOC_PROT_ERR over IB (works on + // NVLink either way, which masked it single-host). Size at an + // NC-independent upper bound (mc <= sm_count) so the symmetric + // allocation never has to grow across NC variants — it aborts on grow; + // larger-than-needed is fine. Mirrors the fwd moe.cu comm_l_ub sizing. + int sm_count = 0; + { + int dev = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev); + } + // Worst-case sizing across the whole tuner sweep: the symmetric/device + // pools cache by key and ABORT on grow, so the FIRST config to allocate a + // key locks its size. A small-TileM (64) or small-CommStages config must + // not lock in a buffer that a later TileM=256 / deeper-pipe config can't + // grow into. Size every comm buffer at the centralized TileM/stage maxima; + // the kernel indexes by the per-call + // (mc·kCommNumStages, kTileM) ≤ these UBs, so over-allocation is safe. + // Mirrors moe.cu's MoeSymmConfig.max_tile_m / max_comm_stages fix. + constexpr int kCommStagesUB = kMaxMoeCommStages; + constexpr int kTileMUB = kMaxMoeCommTileM; + int comm_l_ub = sm_count * kCommStagesUB; + size_t stage_slots_bytes = + (size_t)comm_l_ub * kTileMUB * scfg.hidden_dim * sizeof(Element); + b.x_staging = static_cast(pool.get_symmetric("moe_bwd_x_staging", stage_slots_bytes)); + b.dy_staging = static_cast(pool.get_symmetric("moe_bwd_dy_staging", stage_slots_bytes)); + b.dx_staging = static_cast(pool.get_symmetric("moe_bwd_dx_staging", stage_slots_bytes)); + + // X / dY pipe signals: per-slot, [L = MC * NumStages] ints. Comm + // producer indexes by `+ xc` (stride MC); MLP consumer is + // ticket-based and walks slot = (m_base + idx·gridDim.x) mod L. + // Allocate at the UB (comm_l_ub ≥ MC·kCommNumStages); the kernel only + // touches [0, MC·kCommNumStages). Per-call resets (below) zero the live prefix. + size_t tile_sig_bytes = (size_t)comm_l_ub * sizeof(int); + b.x_src_ready = static_cast(pool.get_device("moe_bwd_x_src_ready", tile_sig_bytes)); + b.x_src_consumed = static_cast(pool.get_device("moe_bwd_x_src_consumed", tile_sig_bytes)); + b.dy_src_ready = static_cast(pool.get_device("moe_bwd_dy_src_ready", tile_sig_bytes)); + b.dy_src_consumed= static_cast(pool.get_device("moe_bwd_dy_src_consumed",tile_sig_bytes)); + cudaMemsetAsync(b.x_src_ready, 0, tile_sig_bytes, stream); + cudaMemsetAsync(b.x_src_consumed, 0, tile_sig_bytes, stream); + cudaMemsetAsync(b.dy_src_ready, 0, tile_sig_bytes, stream); + cudaMemsetAsync(b.dy_src_consumed,0, tile_sig_bytes, stream); + + // dX pipe signals: per-slot, [L = MC * NumStages] ints. UB-sized (see above). + size_t tile_pipe_bytes = (size_t)comm_l_ub * sizeof(int); + b.dst_ready = static_cast(pool.get_device("moe_bwd_dst_ready", tile_pipe_bytes)); + b.dst_consumed = static_cast(pool.get_device("moe_bwd_dst_consumed", tile_pipe_bytes)); + cudaMemsetAsync(b.dst_ready, 0, tile_pipe_bytes, stream); + cudaMemsetAsync(b.dst_consumed, 0, tile_pipe_bytes, stream); + + size_t eid_bytes = (size_t)comm_l_ub * sizeof(int); // UB-sized (see above) + // One array per get pipe (see MoeBwdBuffers / CommBuffersBwd docstrings). + // Same shape as the staging signal arrays. + b.tile_expert_ids_x = static_cast( + pool.get_device("moe_bwd_tile_expert_ids_x", eid_bytes)); + b.tile_expert_ids_dy = static_cast( + pool.get_device("moe_bwd_tile_expert_ids_dy", eid_bytes)); + b.tile_valid_rows_x = static_cast( + pool.get_device("moe_bwd_tile_valid_rows_x", eid_bytes)); + + // ── Sort scratch (sized to scfg max for stability across calls) ── + int max_sorted_slots = scfg.max_total_slots; + int max_m_tiles = scfg.max_total_slots / kTileM; + int max_tile_counts = max_m_tiles * scfg.max_num_experts; + + b.tile_expert_counts = static_cast( + pool.get_device("moe_bwd_sort_tile_counts", max_tile_counts * sizeof(int))); + b.sorted_token_ids = static_cast( + pool.get_device("moe_bwd_sort_sorted_ids", max_sorted_slots * sizeof(int))); + b.token_expert_slots = static_cast( + pool.get_device("moe_bwd_sort_token_slots", max_sorted_slots * sizeof(int))); + b.cta_done = static_cast( + pool.get_device("moe_bwd_sort_cta_done", num_blocks * sizeof(int))); + b.cta_sums = static_cast( + pool.get_device("moe_bwd_sort_cta_sums", + num_blocks * scfg.max_num_experts * sizeof(int))); + + cudaMemsetAsync(b.cta_done, 0, num_blocks * sizeof(int), stream); + cudaMemsetAsync(b.sorted_token_ids, 0xff, max_sorted_slots * sizeof(int), stream); + + return b; +} + +// ============================================================================ +// build_expert_for_k_block — small helper kernel section +// ============================================================================ +// +// Post-sort: derive expert_for_k_block[k] = tile_expert_ids[(k * TileK) / TileM]. +// Works for both directions: +// TileK > TileM (K-block spans multiple m-tiles): uses the first m-tile's expert. +// TileK < TileM (m-tile spans multiple K-blocks): all K-blocks in a m-tile +// resolve to the same tile_expert_ids slot via integer division. +// Assumes per-expert ranges are TileK-aligned, OR that any K-block straddling +// an expert boundary has only padding tokens in its second m-tile (ensured +// by zeroed dy_sorted + zero-filled X-padding from forward dispatch). + +__device__ __forceinline__ void build_expert_for_k_block( + int* expert_for_k_block, + const int* tile_expert_ids, + int num_k_blocks, + int num_actual_k_blocks, + int tile_k, + int tile_m, + int num_threads) { + // For k < num_actual_k_blocks, copy the expert from the sort output. + // For padding entries (k beyond expert_offsets[num_experts] / TileK), + // write a sentinel that the Phase 3/4 bounds check will reject. Without + // this, stale memory in tile_expert_ids[k] could pass the bounds check + // and cause Phase 4 to load garbage X/dU into the dB/dC accumulators. + for (int k = threadIdx.x; k < num_k_blocks; k += num_threads) { + expert_for_k_block[k] = (k < num_actual_k_blocks) + ? tile_expert_ids[(k * tile_k) / tile_m] + : -1; + } +} + +// ============================================================================ +// Fused kernel: sort → fwd-dispatch → combine_bwd → MLP_bwd+comm → dispatch_bwd +// ============================================================================ + +template +__global__ void __launch_bounds__(Config::kNumThreads) +moe_bwd_kernel( + // Local TMA + __grid_constant__ TmaLoadX const tma_load_x, + __grid_constant__ TmaLoadW1 const tma_load_b_fwd, + __grid_constant__ TmaLoadW1 const tma_load_c_fwd, + __grid_constant__ TmaStoreZ const tma_store_z, + __grid_constant__ TmaStoreU const tma_store_u, + __grid_constant__ TmaStoreV const tma_store_v, + __grid_constant__ TmaLoadDY const tma_load_dy, + __grid_constant__ TmaLoadW2T const tma_load_a_col, + __grid_constant__ TmaStoreDZ const tma_store_dz, + __grid_constant__ TmaLoadDU const tma_load_du, + __grid_constant__ TmaLoadDV const tma_load_dv, + __grid_constant__ TmaLoadB const tma_load_b_col, + __grid_constant__ TmaLoadC const tma_load_c_col, + __grid_constant__ TmaStoreDX const tma_store_dx, + __grid_constant__ TmaLoadDYT3 const tma_load_dyt3, + __grid_constant__ TmaLoadZT3 const tma_load_zt3, + __grid_constant__ TmaReduceDA const tma_reduce_da, + __grid_constant__ TmaLoadXT4 const tma_load_xt4, + __grid_constant__ TmaLoaddUT4 const tma_load_dut4, + __grid_constant__ TmaLoaddVT4 const tma_load_dvt4, + __grid_constant__ TmaReduceDB const tma_reduce_db, + __grid_constant__ TmaReduceDC const tma_reduce_dc, + // Remote TMA + __grid_constant__ TmaLoadXR const tma_load_x_remote, + __grid_constant__ TmaLoadDYR const tma_load_dy_remote, + __grid_constant__ TmaStoreDXR const tma_store_dx_remote, + // Remote Phase 2 TMA (staging-backed dY/X for dA/dB/dC) + __grid_constant__ TmaLoadDYT3R const tma_load_dyt3_remote, + __grid_constant__ TmaLoadXT4R const tma_load_xt4_remote, + // Dims + bufs + __grid_constant__ MlpBwdDims const mlp_dims, + __grid_constant__ MlpBwdBufs const mlp_bufs, + __grid_constant__ MoeBwdParams const p, + __grid_constant__ int const static_nsplit, + __grid_constant__ GetTmaDescsBwd const get_descs_bwd) { + + using Element = typename Config::Element; + using Traits1 = typename Config::Traits1; + using Traits2T = typename Config::Traits2T; + using Traits3 = typename Config::Traits3; + using Traits4 = typename Config::Traits4; + using Traits5 = typename Config::Traits5; + using FusedIter = typename Config::FusedIter; + + constexpr int kTileM = Config::kTileM; + constexpr int kNumThreads = Config::kNumThreads; + constexpr int kCommNumStages = Config::kCommNumStages; + constexpr int kNC = Config::kNC; + constexpr int kGetBoxRowsBwd = Config::kCompute == 100 + ? liger::kGetBoxRowsBwdSm100 + : liger::kGetBoxRowsBwdSm90; + + // Target-based cross-CTA barrier — single instance per kernel call. + SyncThreadsCtaCounterBarrier pe_barrier(p.comm.pe_sync.cta_counter, + (int)gridDim.x * (int)gridDim.y); + + // ── Phase 0: Sort — SKIPPED. + // Caller provides token_expert_slots, tile_expert_ids, and + // expert_offsets from the most recent forward call (these come out of + // moe_fused_fwd_bf16, populated by the forward + // kernel's sort_tokens). We trust those instead of re-sorting. + + // ── Phase 1: Broadcast expert offsets — SKIPPED. + // Forward returned the broadcasted offsets in the `expert_offsets` + // tensor we forward verbatim to p.comm.all_expert_offsets, and + // forward ended with nvshmem_team_sync(team), so every PE's writes + // are visible here. No bwd-side re-broadcast needed. + + // ── Phase 2a: Build expert_for_k_block (post-sort, before MLP) ─── + // Operates on caller-provided tile_expert_ids (forward's sort output, + // threaded through p.sort.tile_expert_ids). BWD_TileM is constrained + // to match FWD_TileM (enforced in moe_bwd_fwd_bf16_auto), so the + // FWD-populated array indexes correctly at BWD's tile granularity. + { + int num_k_blocks = p.total_slots / Traits3::TileK; + // Real data ends at expert_offsets[num_experts]; everything beyond + // is buffer-pool padding from the symmetric x_sorted allocation. + int total_real_rows = p.sort.expert_offsets[p.comm.num_experts]; + int num_actual_k_blocks = total_real_rows / Traits3::TileK; + // Both _mlp4 and _mlp3 pointers alias buf.expert_for_k_block in the + // local pass (set above), so writing through either populates the + // shared post-sort array. + build_expert_for_k_block( + const_cast(mlp_bufs.expert_for_k_block_mlp4), + p.sort.tile_expert_ids, + num_k_blocks, num_actual_k_blocks, + Traits3::TileK, kTileM, kNumThreads); + } + __syncthreads(); + + // ── Phase 2b: Forward dispatch — SKIPPED. + // p.x_sorted points to the forward's x_sorted (already dispatched). + // MLP TMAs read from there directly; nothing to do here. + + // ── Phase 3: Reverse-combine — dY → dy_sorted (+ dW side-output) ─ + // Caller-provided output_buffer is Y_fwd in expert-sorted layout + // [total_slots, D] (must match the expert-sort produced by the + // forward dispatch). combine_tokens_bwd computes: + // dW[t,k] = Σ_d dY[t,d] * Y_fwd[slot(t,k), d] → p.dW_out + // dy_sorted[slot,d]= dY[t,d] * w[t,k] → fed to mlp_fused_bwd + combine_tokens_bwd( + p.dY, + /*expert_out=*/p.output_buffer, + p.expert_weights, + p.sort.token_expert_slots, + p.dW_out, + p.dy_sorted, + p.hidden_dim, + p.num_tokens, + p.top_k); + __threadfence_system(); + __syncthreads(); + + p.comm.pe_sync.barrier_all(pe_barrier); + + // ── Phase 4: MLP_bwd with comm overlap ─────────────────────────── + { + extern __shared__ char raw_smem[]; + auto& smem = *reinterpret_cast*>(raw_smem); + + size_t offsets_region = sizeof(MoeBwdSmem); + offsets_region = (offsets_region + sizeof(int) - 1) & ~(sizeof(int) - 1); + size_t offsets_count = + (size_t)p.comm_bwd.num_pes * (p.comm_bwd.experts_per_pe + 1); + if (threadIdx.x == 0) { + smem.comm.remote_offsets = + reinterpret_cast(raw_smem + offsets_region); + smem.comm.remote_counts = + smem.comm.remote_offsets + offsets_count; + } + __syncthreads(); + + // Bwd TMA-GET bounce: 128-aligned region after offsets and counts. + // Mirrors the host smem sizing (align128(smem_mlp) + bounce). Only + // dereferenced when get_descs_bwd.enabled; harmless otherwise. + size_t bounce_off = offsets_region + + (offsets_count + + (size_t)p.comm_bwd.num_pes * p.comm_bwd.experts_per_pe) + * sizeof(int); + bounce_off = (bounce_off + 127) & ~((size_t)127); + char* get_bounce = raw_smem + bounce_off; + + // No sentinel write here: nvshmem_comm_prologue_bwd's warp-2 leader + // is guaranteed to write smem.comm.total_tiles (either to 0 in the + // early-exit case or to comm_count in the normal case) before the + // __syncthreads below. A `smem.comm.total_tiles = -1` from + // threadIdx.x == 0 racing with the prologue's warp-2 write is what + // caused the BWD-TM=64 + 2 PE deadlock: when thread 0's sentinel + // write landed AFTER warp 2's write, total_tiles ended up as -1, + // comm skipped nvshmem_comm_main_bwd, the comm pipes (x_src_ready + // etc.) never got a producer_release, and the MLP remote pass' + // iter.acquire_src spun forever. + nvshmem_comm_prologue_bwd(smem.comm, p.comm_bwd, static_nsplit); + __syncthreads(); + + int flat_id = (int)blockIdx.x; + int runtime_nsplit = static_nsplit; + int n_gemm = p.comm_bwd.n_gemm; + int grid_x = n_gemm / runtime_nsplit; + int col = flat_id / runtime_nsplit; + if (threadIdx.x == 0) { + int total_m_tiles = smem.comm.global_total; + runtime_nsplit = select_runtime_nsplit_bwd_from_tiles( + total_m_tiles, + mlp_dims.num_n_tiles_1, + mlp_dims.num_n_tiles_2t, + mlp_dims.num_n_tiles_5, + mlp_dims.num_k_tiles_1, + mlp_dims.num_k_tiles_2t, + mlp_dims.num_k_tiles_5, + p.comm_bwd.experts_per_pe, + (int)gridDim.x, + static_nsplit); + n_gemm = ((int)gridDim.x / runtime_nsplit) * runtime_nsplit; + grid_x = n_gemm / runtime_nsplit; + col = flat_id / runtime_nsplit; + int mlp_count = 0; + if (flat_id < n_gemm && col < total_m_tiles) + mlp_count = (total_m_tiles - 1 - col) / grid_x + 1; + smem.comm.per_cta_tiles = mlp_count; + smem.comm.runtime_nsplit = runtime_nsplit; + smem.comm.runtime_n_gemm = n_gemm; + smem.comm.runtime_grid_x = grid_x; + } + __syncthreads(); + runtime_nsplit = smem.comm.runtime_nsplit; + n_gemm = smem.comm.runtime_n_gemm; + grid_x = smem.comm.runtime_grid_x; + col = flat_id / runtime_nsplit; + int split = flat_id % runtime_nsplit; + bool gemm_active = flat_id < n_gemm; + + int warp_id = threadIdx.x / 32; + bool is_comm_warp = (warp_id == kCommBwdGetWarp0 || warp_id == kCommBwdPutWarp); + + if (is_comm_warp) { + if (smem.comm.total_tiles > 0) { + nvshmem_comm_main_bwd< + Element, kTileM, kCommNumStages, kNC, + kGetBoxRowsBwd>( + smem.comm, p.comm_bwd, &get_descs_bwd, get_bounce); + nvshmem_quiet(); + } + } else { + // MLP warps. The fused iter is initialized for the staging ring + // inside moe_fused_bwd (init_remote); every tile — local or remote — + // flows through that staging path, so there is no local sub-iter to + // seed here. Gap CTAs skip Phase 1; moe_fused_bwd gates it on + // gemm_active. + FusedIter fused_iter; + MlpBwdDims runtime_mlp_dims = mlp_dims; + runtime_mlp_dims.grid_x = grid_x; + runtime_mlp_dims.n_gemm = n_gemm; + + moe_fused_bwd( + smem, fused_iter, + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_u, tma_store_v, + tma_load_dy, tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, + tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, + tma_load_dyt3_remote, tma_load_xt4_remote, + p.comm_bwd.x_src_ready, p.comm_bwd.x_src_consumed, + p.comm_bwd.dy_src_ready, p.comm_bwd.dy_src_consumed, + p.comm_bwd.dst_ready, p.comm_bwd.dst_consumed, + p.comm_bwd.tile_expert_ids_x, + p.comm_bwd.tile_expert_ids_dy, + p.comm_bwd.tile_valid_rows_x, + runtime_mlp_dims, mlp_bufs, + col, grid_x, split, runtime_nsplit, gemm_active); + } + + __threadfence(); + __syncthreads(); + p.comm.pe_sync.barrier_all(pe_barrier); + } + + // ── Phase 5: Reverse-dispatch — dx_sorted → dX_out (gather-sum) ── + dispatch_tokens_bwd( + p.dx_sorted, + p.sort.token_expert_slots, + p.dX_out, + p.hidden_dim, + p.num_tokens, + p.top_k); + +} + +// ============================================================================ +// GET-via-TMA descriptor builder for the bwd comm path (raw CUtensorMap). +// ============================================================================ +// +// Box = [kGetBoxRowsBwd, kGetKChunk] over a row-major [n_rows, hidden_dim] +// bf16 tensor (SWIZZLE_NONE keeps global bytes row-major so the MLP's +// swizzled reader sees correct data). Mirrors moe.cu's make_get_tma_map but +// with the bwd's small box height. Returns false on encode failure → caller +// disables the bwd TMA-get path (falls back to getmem/int4). +static inline bool make_get_tma_map_bwd(CUtensorMap& m, void* base, + int hidden_dim, int n_rows, int box_rows) { + uint64_t gdim[2] = { (uint64_t)hidden_dim, (uint64_t)n_rows }; + uint64_t gstride[1] = { (uint64_t)hidden_dim * sizeof(bfloat16_t) }; + uint32_t bdim[2] = { (uint32_t)liger::kGetKChunk, (uint32_t)box_rows }; + uint32_t estride[2] = { 1, 1 }; + CUresult r = cuTensorMapEncodeTiled( + &m, CU_TENSOR_MAP_DATA_TYPE_UINT16, 2, base, gdim, gstride, bdim, estride, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + return r == CUDA_SUCCESS; +} + +// ============================================================================ +// Host wrapper +// ============================================================================ + +// MoeBwdArgs (the POD launch bundle) is defined in moe_launch.h, shared with the +// tuner (mirrors MoeFwdArgs). dB/dC/dA/dW are gradient accumulation targets — +// the launcher zeroes them before the kernel. + +template > +void +moe_bwd_fwd_bf16(const MoeBwdArgs& a, int static_nsplit) { + + const int num_experts = a.num_experts; + const int top_k = a.top_k; + const nvshmem_team_t team = a.team; + cudaStream_t stream = a.stream; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + if (cudaError_t e = cudaStreamIsCapturing(stream, &capture_status); + e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaStreamIsCapturing failed: ", + cudaGetErrorString(e)); + if (capture_status == cudaStreamCaptureStatusNone) { + if (cudaError_t e = cudaStreamSynchronize(stream); e != cudaSuccess) + LIGER_FAIL_CUDA( + "moe_bwd: synchronizing the forward-to-backward NVSHMEM " + "handoff failed: ", cudaGetErrorString(e)); + } + + using Element = typename Config::Element; + using Traits1 = typename Config::Traits1; + using Traits2T = typename Config::Traits2T; + using Traits3 = typename Config::Traits3; + using Traits4 = typename Config::Traits4; + using Traits5 = typename Config::Traits5; + using TmaLoadAtom = TmaLoadAtomForCompute; + using TmaStoreAtom = TmaStoreAtomForCompute; + using TmaReduceAddAtom = TmaReduceAddAtomForCompute; + constexpr int kTileM = Config::kTileM; + constexpr int kNumThreads = Config::kNumThreads; + + // Iteration counter lives on the device; CTA 0 thread 0 atomicAdds it at + // the end of the kernel via PeSync::advance_iteration_counter(). + int num_tokens = a.num_tokens; + int hidden_dim = a.hidden_dim; + int experts_per_pe = a.experts_per_pe; + int intermediate_dim = a.intermediate_dim; + + // max_total_slots uses kTileM because forward dispatch always pads each + // expert range to its own TileM. moe_bwd_fwd_bf16_auto enforces + // BWD_TileM == FWD_TileM, so kTileM is the right alignment here. + int max_total_slots = num_tokens * top_k + num_experts * kTileM; + int num_m_tiles_in = (num_tokens + kTileM - 1) / kTileM; + + int num_pes = nvshmem_team_n_pes(team); + int my_pe = nvshmem_team_my_pe(team); + int local_gpus = 0; + cudaGetDeviceCount(&local_gpus); + int gpus_per_node = + (local_gpus > 0 && num_pes % local_gpus == 0) ? local_gpus : num_pes; + + int num_sms; + cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, a.device); + // Flat 1-D launch: use the largest NC-complete SM subset. Phase 1 + comm partition over the + // GEMM-active subset n_gemm = floor_NS (largest multiple of static_nsplit ≤ + // num_blocks); Phase 2 (mlp3/mlp4) grid-strides over all num_blocks CTAs + // (no divisibility needed — cell stride = num_blocks); the gap CTAs + // [n_gemm, num_blocks) fall through Phase 1/comm but still run Phase 2 + // and every barrier. NC-complete launch geometry is required because comm + // ring lengths, signal arrays, and TMA descriptors all use gridDim.x / NC. + // ClusterM divides NC for every supported configuration, so this also + // guarantees complete SM100 CTA pairs. + static_assert(Config::kNC % Config::kClusterM == 0); + int num_blocks = num_sms - num_sms % Config::kNC; + int n_gemm = num_blocks - num_blocks % static_nsplit; + int grid_x = n_gemm / static_nsplit; + int max_runtime_grid_x = (num_blocks + 1) / 2; // min runtime NS candidate = 2 + dim3 grid(num_blocks); + LIGER_CHECK(top_k <= kCombineMaxTopK, + "top_k (", top_k, ") exceeds kCombineMaxTopK (", kCombineMaxTopK, + "). Bump the constant in combine.cuh."); + + // stream comes from a.stream (the binding passes the current torch stream). + auto& pool = global_buffer_pool(); + + auto buf = allocate_moe_bwd_buffers(pool, stream, max_runtime_grid_x, num_blocks, + hidden_dim, intermediate_dim, num_tokens, top_k, + experts_per_pe, max_total_slots, num_m_tiles_in); + + // Caller-owned gradient outputs (binding-allocated). dB/dC/dA/dW are + // accumulation targets — zero them on the stream (upstream used torch::zeros); + // dX is fully overwritten so it needs no pre-zero. + constexpr std::size_t kElemBytes = 2; // sizeof(bf16) + cudaMemsetAsync(a.dB, 0, + static_cast(experts_per_pe) * intermediate_dim * hidden_dim * kElemBytes, stream); + cudaMemsetAsync(a.dC, 0, + static_cast(experts_per_pe) * intermediate_dim * hidden_dim * kElemBytes, stream); + cudaMemsetAsync(a.dA, 0, + static_cast(experts_per_pe) * hidden_dim * intermediate_dim * kElemBytes, stream); + cudaMemsetAsync(a.dW, 0, + static_cast(num_tokens) * top_k * kElemBytes, stream); + + // ── Build MoeBwdParams ────────────────────────────────────────── + MoeBwdParams p; + p.expert_indices = a.expert_indices; + p.expert_weights = reinterpret_cast(a.expert_weights); + // The BWD kernel never reads the original X token layout — token data + // is read only from x_sorted (FWD's expert-sorted view of X). Shape + // and device for dX_out come from dY. + p.dY = reinterpret_cast(a.dY); + p.output_buffer = reinterpret_cast(a.Y_fwd); + + p.dX_out = reinterpret_cast(a.dX); + p.dB_out = reinterpret_cast(a.dB); + p.dC_out = reinterpret_cast(a.dC); + p.dA_out = reinterpret_cast(a.dA); + p.dW_out = reinterpret_cast(a.dW); + + // Caller-provided forward intermediates (x_sorted, sort outputs). + // These come from moe_fused_fwd_bf16 and replace + // what the bwd allocator + internal sort/dispatch used to populate. + p.x_sorted = reinterpret_cast(a.x_sorted); + p.dy_sorted = buf.dy_sorted; + p.dx_sorted = buf.dx_sorted; + + // Sort outputs come from forward; we don't run sort_tokens in bwd. + // The fields we leave from buf.* (tile_expert_counts, sorted_token_ids, + // cta_done, cta_sums) are only consumed by sort_tokens — now skipped. + // expert_offsets is now the broadcasted [num_pes, num_experts + 1] + // view from FWD. The local-PE single-row view is offset by + // my_pe * (num_experts + 1) — needed by sort lookup + // (total_real_rows = local_offsets[num_experts]) and by comm code + // that walks this PE's expert range. + int* local_expert_offsets = + a.expert_offsets + my_pe * (num_experts + 1); + + p.sort.tile_expert_counts = buf.tile_expert_counts; + p.sort.sorted_token_ids = buf.sorted_token_ids; + p.sort.token_expert_slots = a.token_expert_slots; + p.sort.expert_offsets = local_expert_offsets; + p.sort.expert_counts = nullptr; + p.sort.tile_expert_ids = a.tile_expert_ids; + p.sort.cta_done = buf.cta_done; + p.sort.cta_sums = buf.cta_sums; + + // Forward CommBuffers (used for pe_sync only — bwd no longer + // re-broadcasts expert offsets; it reads what fwd wrote). + p.comm.src_staging = nullptr; + p.comm.dst_staging = nullptr; + p.comm.src_ready = nullptr; + p.comm.src_consumed = nullptr; + p.comm.dst_ready = nullptr; + p.comm.dst_consumed = nullptr; + p.comm.expert_offsets = local_expert_offsets; + p.comm.expert_counts = nullptr; + p.comm.sorted_token_ids = buf.sorted_token_ids; + p.comm.local_tokens = reinterpret_cast(a.x_sorted); + p.comm.local_output = buf.dx_sorted; + // Forward CommBuffers field — preserved from before the bwd-side split. + // The X-pipe array is the canonical "tile_expert_ids" (the dY-pipe array + // is a bwd-only mirror gated by release_dy). + p.comm.tile_expert_ids = buf.tile_expert_ids_x; + p.comm.all_expert_offsets = a.expert_offsets; + p.comm.all_expert_counts = a.expert_counts; + p.comm.pe_sync.cta_counter = buf.cta_counter; + p.comm.pe_sync.num_pes = num_pes; + p.comm.pe_sync.my_pe = my_pe; + p.comm.pe_sync.team = team; + p.comm.hidden_dim = hidden_dim; + p.comm.num_experts = num_experts; + p.comm.experts_per_pe = experts_per_pe; + p.comm.all_expert_offsets_stride = num_experts + 1; + + // Backward CommBuffers (used for X+dY gets and dX puts). + p.comm_bwd.x_staging = buf.x_staging; + p.comm_bwd.dy_staging = buf.dy_staging; + p.comm_bwd.dx_staging = buf.dx_staging; + p.comm_bwd.x_src_ready = buf.x_src_ready; + p.comm_bwd.x_src_consumed = buf.x_src_consumed; + p.comm_bwd.dy_src_ready = buf.dy_src_ready; + p.comm_bwd.dy_src_consumed = buf.dy_src_consumed; + p.comm_bwd.dst_ready = buf.dst_ready; + p.comm_bwd.dst_consumed = buf.dst_consumed; + p.comm_bwd.tile_expert_ids_x = buf.tile_expert_ids_x; + p.comm_bwd.tile_expert_ids_dy = buf.tile_expert_ids_dy; + p.comm_bwd.tile_valid_rows_x = buf.tile_valid_rows_x; + // Cross-PE pull sources. Both x_sorted and dy_sorted are in the + // same global-expert-sorted layout on every PE: + // • x_sorted: forward's sort_tokens → dispatch wrote it. + // • dy_sorted: combine_tokens_bwd writes [slot, d] using the same + // token_expert_slots → matches forward's per-PE expert sort. + // all_expert_offsets is the broadcasted offset tensor produced by + // fwd (broadcast_expert_offsets) and passed through here verbatim — + // it tells remote PEs the [start, end) slice in our buffers that + // holds data for their experts. + p.comm_bwd.remote_x = p.x_sorted; + p.comm_bwd.remote_dy = buf.dy_sorted; + p.comm_bwd.remote_dx = buf.dx_sorted; + p.comm_bwd.all_expert_offsets = a.expert_offsets; + p.comm_bwd.all_expert_counts = a.expert_counts; + p.comm_bwd.all_expert_offsets_stride = num_experts + 1; + p.comm_bwd.my_pe = my_pe; + p.comm_bwd.num_pes = num_pes; + p.comm_bwd.team = team; + p.comm_bwd.hidden_dim = hidden_dim; + p.comm_bwd.num_experts = num_experts; + p.comm_bwd.experts_per_pe = experts_per_pe; + // GEMM/comm-active count: comm derives MC = n_gemm/NC and the MLP column + // count grid_x = n_gemm/NSplit from this (the flat launch over-covers it). + p.comm_bwd.n_gemm = n_gemm; + + p.num_tokens = num_tokens; + p.hidden_dim = hidden_dim; + p.intermediate_dim = intermediate_dim; + p.top_k = top_k; + p.total_slots = max_total_slots; + p.num_m_tiles = num_m_tiles_in; + p.experts_per_pe = experts_per_pe; + + // ── MlpBwdDims (passed to inner mlp_fused_bwd) ────────────────── + int total_n_rows_1 = experts_per_pe * intermediate_dim; + int total_k_cols_2 = experts_per_pe * hidden_dim; + + MlpBwdDims mlp_dims; + mlp_dims.hidden_dim = hidden_dim; + mlp_dims.intermediate_dim = intermediate_dim; + mlp_dims.total_n_rows_1 = total_n_rows_1; + mlp_dims.num_n_tiles_1 = (Config::kCompute == 90) + ? (intermediate_dim + Traits1::TileN - 1) / Traits1::TileN + : intermediate_dim / Traits1::TileN; + mlp_dims.num_k_tiles_1 = (Config::kCompute == 90) + ? (hidden_dim + Traits1::TileK - 1) / Traits1::TileK + : hidden_dim / Traits1::TileK; + mlp_dims.total_k_cols_2t = total_k_cols_2; + mlp_dims.num_n_tiles_2t = (Config::kCompute == 90) + ? (intermediate_dim + Traits2T::TileN - 1) / Traits2T::TileN + : intermediate_dim / Traits2T::TileN; + mlp_dims.num_k_tiles_2t = (Config::kCompute == 90) + ? (hidden_dim + Traits2T::TileK - 1) / Traits2T::TileK + : hidden_dim / Traits2T::TileK; + mlp_dims.total_k_cols_5 = total_n_rows_1; + mlp_dims.num_n_tiles_5 = (Config::kCompute == 90) + ? (hidden_dim + Traits5::TileN - 1) / Traits5::TileN + : hidden_dim / Traits5::TileN; + mlp_dims.num_k_tiles_5 = (Config::kCompute == 90) + ? (intermediate_dim + Traits5::TileK - 1) / Traits5::TileK + : intermediate_dim / Traits5::TileK; + mlp_dims.num_m_tiles_3 = (Config::kCompute == 90) + ? (hidden_dim + Traits3::TileM - 1) / Traits3::TileM + : hidden_dim / Traits3::TileM; + mlp_dims.num_n_tiles_3 = (Config::kCompute == 90) + ? (intermediate_dim + Traits3::TileN - 1) / Traits3::TileN + : intermediate_dim / Traits3::TileN; + mlp_dims.total_n_rows_3 = total_k_cols_2; + mlp_dims.num_m_tiles_4 = (Config::kCompute == 90) + ? (intermediate_dim + Traits4::TileM - 1) / Traits4::TileM + : intermediate_dim / Traits4::TileM; + mlp_dims.num_n_tiles_4 = (Config::kCompute == 90) + ? (hidden_dim + Traits4::TileN - 1) / Traits4::TileN + : hidden_dim / Traits4::TileN; + mlp_dims.total_m_rows_4 = total_n_rows_1; + mlp_dims.phase_counter = buf.phase_counter; + // Flat-grid geometry: Phase 1 columns = grid_x, GEMM-active count = n_gemm. + mlp_dims.grid_x = grid_x; + mlp_dims.n_gemm = n_gemm; + mlp_dims.num_tokens = max_total_slots; + mlp_dims.num_m_tiles = max_total_slots / kTileM; + // Local pass: ring/buffer size and the K-block iteration bound coincide, + // so total_m_tiles_in_pass equals num_m_tiles. The remote pass overrides + // this in moe_bwd.cuh to comm.global_total. + mlp_dims.total_m_tiles_in_pass = mlp_dims.num_m_tiles; + // Phase 1 processes grid_x m-tiles per batch (one per column). + mlp_dims.num_batches = (mlp_dims.num_m_tiles + grid_x - 1) / grid_x; + // Sort emits GLOBAL expert IDs; mlp_fused_bwd computes + // `local_expert = tile.expert - local_expert_start` to index this PE's + // local all_B/all_C/all_A. Forward kernel does the same (moe.cu:469). + mlp_dims.local_expert_start = my_pe * experts_per_pe; + mlp_dims.num_pes = num_pes; + mlp_dims.experts_per_pe = experts_per_pe; + // Local pass: expert_for_k_block is built per-K-block by + // build_expert_for_k_block (one entry per Traits{3,4}::TileK), so the + // K-block index can be used directly as the array index. The remote pass + // overrides these via EfkbStride{3,4} template params in moe_fused_bwd + // because its array is per-comm-tile. + + // Host-computable constants (grid-constant — saves device-side arithmetic). + // phase2_num_blocks is no longer used for the Phase 2 cell stride (that now + // reads gridDim.x·gridDim.y = num_blocks directly), kept for compatibility. + mlp_dims.phase2_num_blocks = num_blocks; + // Remote pass: TMA descriptors cover the full flat staging list of + // L = MC · kCommNumStages slots, each TileM rows tall. + // MC = launched NC-complete comm columns (matches nvshmem_comm_main_bwd). + int mc_remote = num_blocks / Config::kNC; + mlp_dims.remote_num_m_tiles = mc_remote * Config::kCommNumStages; + mlp_dims.remote_num_tokens = mlp_dims.remote_num_m_tiles * kTileM; + + // Phase-1 sub-batch grouping (mlp_bwd.cuh, SubBatch): the remote pass groups + // up to SubBatch Phase-1 sub-batches before one Phase-2 over their combined + // K-window. The kernel ADAPTS the group size to the ring depth (a group must + // fit in the ring so its grid_x·gcount live slots map to distinct ring slots), + // degrading to fewer sub-batches — down to 1 = the original per-batch loop — + // on shallow rings. So no host-side ring constraint is needed here. + + MlpBwdBufs mlp_bufs; + mlp_bufs.du_buf = buf.du_buf; + mlp_bufs.dv_buf = buf.dv_buf; + mlp_bufs.dz_buf = buf.dz_buf; + // Local pass: no comm staging, no race — both Phase 2 sub-passes alias + // the same post-sort expert_for_k_block. The remote sub-call in + // moe_fused_bwd overrides these to the per-pipe staging arrays. + mlp_bufs.expert_for_k_block_mlp4 = buf.expert_for_k_block; + mlp_bufs.expert_for_k_block_mlp3 = buf.expert_for_k_block; + mlp_bufs.valid_rows_mlp4 = nullptr; + mlp_bufs.expert_k_starts = buf.expert_k_starts; + mlp_bufs.expert_k_ends = buf.expert_k_ends; + mlp_bufs.gmem_da = reinterpret_cast(a.dA); + mlp_bufs.gmem_db = reinterpret_cast(a.dB); + mlp_bufs.gmem_dc = reinterpret_cast(a.dC); + mlp_bufs.barrier_counter = buf.barrier_counter; + + // ── TMA descriptors ───────────────────────────────────────────── + // X comes from caller's x_sorted (forward's pool, post-dispatch). + // dy_sorted is bwd-allocated (combine_tokens_bwd will populate it). + auto* pX = reinterpret_cast(a.x_sorted); // sorted X + auto* pDY = reinterpret_cast(buf.dy_sorted); // expanded dY + auto* pB = reinterpret_cast(a.all_B); + auto* pC = reinterpret_cast(a.all_C); + auto* pA = reinterpret_cast(a.all_A); + auto* pDX = reinterpret_cast(buf.dx_sorted); + + // All Hopper expert weights and weight-gradient targets use rank-3 tensor + // maps. TMA handles partial edge boxes as zero-filled loads / dropped + // stores, while the explicit expert coordinate prevents cross-expert tails. + // (See mlp_bwd.cu for the full pattern; here we mirror it but the + // "X" tensor is the symmetric x_sorted and "dY" is the symmetric + // dy_sorted, both of shape [max_total_slots, hidden_dim].) + auto tma_load_x = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pX), make_shape(max_total_slots, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits1::SmemLayoutX_1{}); + + auto tma_load_b_fwd = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pB), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + static_cast(intermediate_dim) * hidden_dim)), + typename Traits1::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pB), + make_shape(total_n_rows_1, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits1::SmemLayoutW_1{}); + } + }(); + auto tma_load_c_fwd = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pC), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + static_cast(intermediate_dim) * hidden_dim)), + typename Traits1::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pC), + make_shape(total_n_rows_1, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits1::SmemLayoutW_1{}); + } + }(); + + auto tma_store_z = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.z_buf), make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits1::SmemLayoutActStoreSlot{}); + auto tma_store_u = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.du_buf), make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits1::SmemLayoutActStoreSlot{}); + auto tma_store_v = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.dv_buf), make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits1::SmemLayoutActStoreSlot{}); + + auto tma_load_dy = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pDY), make_shape(max_total_slots, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits2T::SmemLayoutZ_1{}); + auto tma_load_a_col = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pA), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + Int<1>{}, + static_cast(intermediate_dim), + static_cast(hidden_dim) * intermediate_dim)), + typename Traits2T::SmemLayoutW_1{}); + // One TMA box per sub-tile; consumer issues NSub stores per outer N-tile. + auto tma_store_dz = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.dz_buf), + make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits2T::SmemLayoutStore_1{}); + + auto tma_load_du = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.du_buf)), + make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits5::SmemLayoutZ_1{}); + auto tma_load_dv = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.dv_buf)), + make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits5::SmemLayoutZ_1{}); + auto tma_load_b_col = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pB), + make_shape(static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(experts_per_pe)), + make_stride( + Int<1>{}, + static_cast(hidden_dim), + static_cast(intermediate_dim) * hidden_dim)), + typename Traits5::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pB), + make_shape(hidden_dim, total_n_rows_1), + make_stride(Int<1>{}, hidden_dim)), + typename Traits5::SmemLayoutW_1{}); + } + }(); + auto tma_load_c_col = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pC), + make_shape(static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(experts_per_pe)), + make_stride( + Int<1>{}, + static_cast(hidden_dim), + static_cast(intermediate_dim) * hidden_dim)), + typename Traits5::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pC), + make_shape(hidden_dim, total_n_rows_1), + make_stride(Int<1>{}, hidden_dim)), + typename Traits5::SmemLayoutW_1{}); + } + }(); + // dX TMA: split-N store (TileN_half boxes per Mlp5::SmemLayoutStore_1). + auto tma_store_dx = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(pDX), + make_shape(max_total_slots, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits5::SmemLayoutStore_1{}); + + // Phase 2 TMAs. mlp3's dY^T / Z operands and mlp4's X / dU^T / dV^T + // operands need pair-aware SM100_TMA_2SM_LOAD descriptors (built via + // make_tma_copy_A_sm100 / make_tma_copy_B_sm100, keyed off the joined + // TileShape + TiledMma2Sm) when Config::kUsesTwoSm (all SM100 configs); + // the dA/dB/dC TMA_REDUCE_ADD outputs are unchanged either way (each + // peer CTA still reduce-adds its own CtaTileM rows independently). + auto tma_load_dyt3 = [&] { + auto tensor = make_tensor(make_gmem_ptr(pDY), + make_shape(hidden_dim, max_total_slots), + make_stride(Int<1>{}, hidden_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits3::SmemLayoutDYT_1{}, + typename Traits3::TileShape{}, + typename Traits3::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits3::SmemLayoutDYT_1{}); + } + }(); + auto tma_load_zt3 = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.z_buf)), + make_shape(intermediate_dim, buf.scratch_rows), + make_stride(Int<1>{}, intermediate_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_B_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits3::SmemLayoutZ_1{}, + typename Traits3::TileShape{}, + typename Traits3::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits3::SmemLayoutZ_1{}); + } + }(); + auto tma_reduce_da = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dA)), + make_shape(static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(intermediate_dim), + Int<1>{}, + static_cast(hidden_dim) * intermediate_dim)), + typename Traits3::SmemLayoutStore{}); + } else { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dA)), + make_shape(total_k_cols_2, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits3::SmemLayoutStore{}); + } + }(); + + auto tma_load_xt4 = [&] { + auto tensor = make_tensor(make_gmem_ptr(pX), + make_shape(hidden_dim, max_total_slots), + make_stride(Int<1>{}, hidden_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_B_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits4::SmemLayoutX_1{}, + typename Traits4::TileShape{}, + typename Traits4::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits4::SmemLayoutX_1{}); + } + }(); + auto tma_load_dut4 = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.du_buf)), + make_shape(intermediate_dim, buf.scratch_rows), + make_stride(Int<1>{}, intermediate_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits4::SmemLayoutA_1{}, + typename Traits4::TileShape{}, + typename Traits4::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits4::SmemLayoutA_1{}); + } + }(); + auto tma_load_dvt4 = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.dv_buf)), + make_shape(intermediate_dim, buf.scratch_rows), + make_stride(Int<1>{}, intermediate_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits4::SmemLayoutA_1{}, + typename Traits4::TileShape{}, + typename Traits4::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits4::SmemLayoutA_1{}); + } + }(); + auto tma_reduce_db = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dB)), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + static_cast(intermediate_dim) * hidden_dim)), + typename Traits4::SmemLayoutStore{}); + } else { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dB)), + make_shape(total_n_rows_1, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits4::SmemLayoutStore{}); + } + }(); + auto tma_reduce_dc = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dC)), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + static_cast(intermediate_dim) * hidden_dim)), + typename Traits4::SmemLayoutStore{}); + } else { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dC)), + make_shape(total_n_rows_1, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits4::SmemLayoutStore{}); + } + }(); + + // Remote (staging) TMAs — buffers cover L = MC · CommNumStages slots, + // each TileM rows tall. MC = (grid.x · NSplit) / NC. + int remote_total_rows = mc_remote * Config::kCommNumStages * kTileM; + auto tma_load_x_remote = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.x_staging)), + make_shape(remote_total_rows, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits1::SmemLayoutX_1{}); + auto tma_load_dy_remote = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.dy_staging)), + make_shape(remote_total_rows, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits2T::SmemLayoutZ_1{}); + // Remote dX TMA: split-N store (matches local tma_store_dx). + auto tma_store_dx_remote = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.dx_staging), + make_shape(remote_total_rows, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits5::SmemLayoutStore_1{}); + + // Remote Phase 3/4 TMAs: transposed staging sources. The remote phase + // of moe_fused_bwd previously reused tma_load_dyt3/tma_load_xt4 (which + // point at LOCAL dy_sorted/x_sorted), so dA/dB/dC for cross-PE-routed + // tokens used the wrong rows. These descriptors pull from the same + // staging buffers that Phase 1/2 of the remote pass already use. + auto tma_load_dyt3_remote = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.dy_staging)), + make_shape(hidden_dim, remote_total_rows), + make_stride(Int<1>{}, hidden_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits3::SmemLayoutDYT_1{}, + typename Traits3::TileShape{}, + typename Traits3::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits3::SmemLayoutDYT_1{}); + } + }(); + auto tma_load_xt4_remote = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.x_staging)), + make_shape(hidden_dim, remote_total_rows), + make_stride(Int<1>{}, hidden_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_B_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits4::SmemLayoutX_1{}, + typename Traits4::TileShape{}, + typename Traits4::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits4::SmemLayoutX_1{}); + } + }(); + + // ── GET-via-TMA descriptors for the bwd comm path (#111 port) ──── + // Per-peer source maps over nvshmem_ptr(x_sorted/dy_sorted, peer) + one + // local dst map each over x_staging / dy_staging. Enabled only on NVLink + // P2P + divisible hidden_dim + bounce smem fit (checked after smem sizing). + GetTmaDescsBwd get_descs_bwd{}; + get_descs_bwd.enabled = 0; + get_descs_bwd.gpus_per_node = gpus_per_node; + get_descs_bwd.my_pe = my_pe; + constexpr int kGetBoxRowsBwd = Config::kCompute == 100 + ? liger::kGetBoxRowsBwdSm100 + : liger::kGetBoxRowsBwdSm90; + size_t get_bounce_bytes_bwd = + (size_t)liger::GetBounceBwd::kTotalBytes; + { + // The single get warp's row band is TileM/NC; the small TMA box height + // must divide it (else the last band would over-read). Holds for NC≤8 + // at TileM=128; larger NC (band < 16) disables the path. + constexpr int kRowBandBwd = Config::kTileM / Config::kNC; + bool ok = (hidden_dim % liger::kGetKChunk == 0) + && (gpus_per_node <= liger::kGetMaxPes) + && (kRowBandBwd % kGetBoxRowsBwd == 0); + int host_start = (my_pe / gpus_per_node) * gpus_per_node; + for (int local_peer = 0; ok && local_peer < gpus_per_node; ++local_peer) { + int tp = host_start + local_peer; + if (tp >= num_pes) { ok = false; break; } + int gpe = nvshmem_team_translate_pe(team, tp, NVSHMEM_TEAM_WORLD); + void* px_peer = nvshmem_ptr(p.comm_bwd.remote_x, gpe); // peer x_sorted + void* pdy_peer = nvshmem_ptr(p.comm_bwd.remote_dy, gpe); // peer dy_sorted + if (px_peer == nullptr || pdy_peer == nullptr) { ok = false; break; } + ok = make_get_tma_map_bwd(get_descs_bwd.src_x_desc[local_peer], px_peer, + hidden_dim, max_total_slots, kGetBoxRowsBwd); + if (ok) ok = make_get_tma_map_bwd(get_descs_bwd.src_dy_desc[local_peer], pdy_peer, + hidden_dim, max_total_slots, kGetBoxRowsBwd); + } + if (ok) ok = make_get_tma_map_bwd(get_descs_bwd.dst_x_staging_desc, + buf.x_staging, hidden_dim, remote_total_rows, kGetBoxRowsBwd); + if (ok) ok = make_get_tma_map_bwd(get_descs_bwd.dst_dy_staging_desc, + buf.dy_staging, hidden_dim, remote_total_rows, kGetBoxRowsBwd); + get_descs_bwd.enabled = ok ? 1 : 0; + if (const char* e = std::getenv("LIGER_GET_TMA")) + if (e[0] == '0') get_descs_bwd.enabled = 0; + } + + // ── Smem sizing ───────────────────────────────────────────────── + int sort_num_warps = kNumThreads / 32; + size_t smem_sort = sort_num_warps * num_experts * sizeof(int); + + size_t smem_mlp = sizeof(MoeBwdSmem); + smem_mlp = (smem_mlp + sizeof(int) - 1) & ~(sizeof(int) - 1); + smem_mlp += num_pes * (experts_per_pe + 1) * sizeof(int); // remote_offsets + smem_mlp += num_pes * experts_per_pe * sizeof(int); // remote_counts + + // Query the device's actual opt-in dynamic-smem cap rather than trust + // Config::kSmemBudget's compile-time 228 KiB constant: the 2SM paired-CTA + // path (Mlp3Traits2Sm / Mlp4Traits2Sm) pushes some candidates close to + // that ceiling, and the real per-SKU limit is what + // cudaFuncAttributeMaxDynamicSharedMemorySize is actually bounded by. + int device = 0; + int dynamic_smem_limit = 0; + if (cudaError_t e = cudaGetDevice(&device); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaGetDevice failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaDeviceGetAttribute(&dynamic_smem_limit, + cudaDevAttrMaxSharedMemoryPerBlockOptin, device); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: querying the opt-in shared-memory limit failed: ", + cudaGetErrorString(e)); + + // GET-via-TMA bounce: append a 128-aligned region after remote_offsets + // (the kernel computes the same offset). The bwd union sits near the + // device's dynamic-smem cap, so disable the TMA-get path if the bounce + // won't fit against the QUERIED limit (not the static 228 KiB constant) + // — that config just keeps the original 2-get/1-put getmem layout (no + // regression). + if (get_descs_bwd.enabled) { + size_t mlp_with_bounce = + ((smem_mlp + 127) & ~((size_t)127)) + get_bounce_bytes_bwd; + if (mlp_with_bounce <= static_cast(dynamic_smem_limit)) smem_mlp = mlp_with_bounce; + else get_descs_bwd.enabled = 0; + } + // Unified BWD comm layout always uses one get warp and one put warp. + // get_descs_bwd.enabled only controls whether same-host tiles use TMA; + // remote/unsupported tiles fall back inside do_get_bwd_tma. + mlp_dims.tma_get_enabled = 1; + + size_t smem_size = std::max(smem_sort, smem_mlp); + + // ── Launch ────────────────────────────────────────────────────── + auto kernel_fn = moe_bwd_kernel; + + if (cudaError_t e = cudaFuncSetAttribute(kernel_fn, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaFuncSetAttribute(MaxDynamicSharedMemorySize=", + smem_size, ", device opt-in limit=", dynamic_smem_limit, ") failed: ", + cudaGetErrorString(e)); + // The SM100 2SM paired-CTA cluster path launches with an explicit + // clusterDim; opt in to non-portable cluster sizes (the CUDA default + // portable max is 8, so ClusterM=2 is already portable in practice, + // but every 2SM launch enables this defensively, mirroring CUTLASS's + // own cluster-launch helpers). + if constexpr (Config::kUsesTwoSm) { + if (cudaError_t e = cudaFuncSetAttribute(kernel_fn, + cudaFuncAttributeNonPortableClusterSizeAllowed, 1); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: enabling non-portable cluster size (ClusterM=", + Config::kClusterM, ") failed: ", cudaGetErrorString(e)); + } + + // Fresh per-call counter resets. + cudaMemsetAsync(buf.cta_counter, 0, sizeof(int), stream); + cudaMemsetAsync(buf.barrier_counter, 0, sizeof(int), stream); + cudaMemsetAsync(buf.phase_counter, 0, max_runtime_grid_x * sizeof(int), stream); + // L = MC · CommNumStages, MC = launched NC-complete comm columns. + int mc_local = ((int)grid.x / Config::kNC); + int comm_l_local = mc_local * Config::kCommNumStages; + size_t per_call_tile_sig_bytes = + static_cast(comm_l_local) * sizeof(int); + cudaMemsetAsync(buf.x_src_ready, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.x_src_consumed, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.dy_src_ready, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.dy_src_consumed,0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.dst_ready, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.dst_consumed, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.cta_done, 0, num_blocks * sizeof(int), stream); + cudaMemsetAsync(buf.sorted_token_ids, 0xff, max_total_slots * sizeof(int), stream); + // dy_sorted MUST be zeroed: combine_tokens_bwd writes routed slots + // only (no padding fill — unlike forward's dispatch_tokens). Partial- + // expert tile rows of dy_sorted are read by Phase 1b's TMA load; if + // those rows aren't zero, the resulting dZ/dU/dV poison Phase 2 dA/dB/dC + // for K-blocks that straddle an expert boundary (see the + // build_expert_for_k_block comment for the alignment assumption). + cudaMemsetAsync(buf.dy_sorted, 0, static_cast(max_total_slots) * hidden_dim * sizeof(Element), stream); + size_t intermediate_bytes = + static_cast(buf.scratch_rows) * intermediate_dim * sizeof(Element); + cudaMemsetAsync(buf.z_buf, 0, intermediate_bytes, stream); + cudaMemsetAsync(buf.du_buf, 0, intermediate_bytes, stream); + cudaMemsetAsync(buf.dv_buf, 0, intermediate_bytes, stream); + cudaMemsetAsync(buf.dz_buf, 0, intermediate_bytes, stream); + + // Sentinel-fill BOTH comm-side tile_expert_ids arrays so unfetched + // slots produce expert_id = -1, which the Phase 3/4 bounds check + // rejects. Without this, stale memory could land in [0, experts_per_pe) + // and trigger spurious dB/dC/dA writes for empty staging slots. + size_t comm_eid_bytes = + static_cast(comm_l_local) * sizeof(int); + if (cudaError_t e = cudaMemsetAsync(buf.tile_expert_ids_x, 0xff, comm_eid_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaMemsetAsync(tile_expert_ids_x) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(buf.tile_expert_ids_dy, 0xff, comm_eid_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaMemsetAsync(tile_expert_ids_dy) failed: ", cudaGetErrorString(e)); + cudaMemsetAsync(buf.tile_valid_rows_x, 0, comm_eid_bytes, stream); + if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: CUDA error before kernel launch: ", cudaGetErrorString(e)); + + // Compute=90 keeps the ordinary <<<>>> launch. Compute=100 always uses + // the explicit paired-CTA cluster launch below. + if constexpr (Config::kUsesTwoSm) { + cudaLaunchConfig_t launch_config = {}; + launch_config.gridDim = grid; + launch_config.blockDim = kNumThreads; + launch_config.dynamicSmemBytes = smem_size; + launch_config.stream = stream; + cudaLaunchAttribute cluster_attr = {}; + cluster_attr.id = cudaLaunchAttributeClusterDimension; + cluster_attr.val.clusterDim.x = Config::kClusterM; + cluster_attr.val.clusterDim.y = 1; + cluster_attr.val.clusterDim.z = 1; + launch_config.attrs = &cluster_attr; + launch_config.numAttrs = 1; + cudaLaunchKernelEx(&launch_config, kernel_fn, + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_u, tma_store_v, + tma_load_dy, tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, + tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, + tma_load_dyt3_remote, tma_load_xt4_remote, + mlp_dims, mlp_bufs, p, static_nsplit, get_descs_bwd); + } else { + kernel_fn<<>>( + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_u, tma_store_v, + tma_load_dy, tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, + tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, + tma_load_dyt3_remote, tma_load_xt4_remote, + mlp_dims, mlp_bufs, p, static_nsplit, get_descs_bwd); + } + if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: kernel launch failed for grid=", grid.x, + " threads=", kNumThreads, " smem=", smem_size, ": ", + cudaGetErrorString(e)); + + // Post-kernel stream-ordered barrier: mirrors the fwd wrapper. The + // kernel's PeSync wait_mlp guarantees this PE has received peer + // signals, but does not guarantee this PE's outbound signals have + // landed on peers. Under back-to-back captured replays (fwd → bwd → + // fwd), bwd's tail-end signals to peers can race against the next + // fwd's wait_dispatch: a peer's wait_until(>= target) may briefly + // see a stale value if bwd's older signal arrives after a newer + // one. Forcing host-side completion here serialises remote + // signal delivery against the next captured kernel. + if (num_pes > 1) { + nvshmemx_barrier_on_stream(team, stream); + } +} + +// ============================================================================ +// Tuned-config wrapper: 13 free knobs → full MoeBwdConfig +// ============================================================================ +// +// Function template that takes the X-macro's free knobs and derives the +// dependent traits per the symmetric-traits constraints documented in +// moe_fwd_bwd_tune_configs.hpp: +// Phase 1a: TileN1, TileK1, Stages1, EpiChunkN1 +// Phase 1b': TileN = 256 (fixed), TileK = TileK1, Stages = Stages1, EpiChunkN25 +// Phase 1d (mlp5): TileM=128, TileN=256, TileK = TileK1, Stages = Stages1, EpiChunkN25 +// Phase 2 mlp3: TileM3, TileN3, TileK3, Stages3, EpiChunkN34 +// Phase 2 mlp4: TileM = TileN3, TileN = TileM3, TileK3, Stages3, EpiChunkN34 +// (mlp3/mlp4 swap) +// NC = NSplit +// +// The tuner registry and dispatch table are built by emitting one explicit +// instantiation of this template per X-macro entry. + +template < + int NSplit2, + int TileN1, int TileK1, int Stages1, + int TileM3, int TileN3, int TileK3, int Stages3, + int EpiChunkN1, int EpiChunkN25, int EpiChunkN34, + int CommNumStages, + int TileM, // COMM tile (staging/semaphores/NC); must match FWD's TileM. 64 or 128. + int GemmTileM = TileM, // Phase-1 GEMM tile (mlp1a/mlp2t/mlp5) ∈ {64,128}; default == comm. + int Compute = 90> +void +moe_bwd_fwd_bf16_tuned(const MoeBwdArgs& a, int static_nsplit) { + // NC paired with TileM (mirrors moe.cu's moe_fused_fwd_bf16): + // TileM=128 → NC=4 (~2× the comm bandwidth of NC=2; deeper comm pipeline). + // TileM=64 → NC=2 (tiles are half-size, NC=2 matches NC=4 throughput). + // Ticket-based iterator + K=NumStages comm means NSplit and NC have no + // divisibility constraint — only requirement is that NC evenly divides + // gridDim.x · gridDim.y so MC is integer (e.g. NS=6, NC=4: MC = 22·6/4 + // = 33). + // Both architectures use CommTileM=128/NC4, matching forward metadata. + using CommGeometry = MoeCommGeometry; + static constexpr int kPhase1TileM = Compute == 100 ? 128 : GemmTileM; + static constexpr int kSubBatch = 2; + using Config = MoeBwdConfig< + NSplit2, + kPhase1TileM, TileN1, TileK1, Stages1, EpiChunkN1, // Phase 1a (mlp1_act) — GEMM tile + /*TileN2=*/256, TileK1, Stages1, // Phase 1b' (Mlp2TTraits caps TileN ≤ 256) + /*TileM2T=*/kPhase1TileM, // Phase 1b' GEMM tile + TileM3, TileN3, TileK3, Stages3, // Phase 2 mlp3 + TileN3, TileM3, TileK3, Stages3, // Phase 2 mlp4 (mlp3/mlp4 swap; cooperative) + kPhase1TileM, 256, TileK1, Stages1, // Phase 1d mlp5 — GEMM tile + EpiChunkN25, // shared mlp2_t + mlp5 + EpiChunkN34, // shared mlp3 + mlp4 + bfloat16_t, + CommNumStages * CommGeometry::RingStageScale + + CommGeometry::RingStageExtra, + /*NC_=*/CommGeometry::NC, + /*SubBatch_=*/kSubBatch, + /*CommTileM_=*/CommGeometry::TileM, + /*Compute_=*/Compute>; + moe_bwd_fwd_bf16(a, static_nsplit); +} + +// ── Runtime dispatch table for tuned-config lookup ──────────────────── + +using MoeBwdFwdFn = void (*)(const MoeBwdArgs&, int static_nsplit); + +struct DispatchEntryBwd { + int Compute; + int NSplit2; + int TileN1, TileK1, Stages1; + int TileM3, TileN3, TileK3, Stages3; + int EpiChunkN1, EpiChunkN25, EpiChunkN34; + int CommNumStages; + int TileM; // COMM tile (staging/semaphores/NC). + int GemmTileM; // Phase-1 GEMM tile (mlp1a/mlp2t/mlp5); == TileM for coupled configs. + MoeBwdFwdFn fn; +}; + +#define LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) \ + { Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM, \ + &moe_bwd_fwd_bf16_tuned }, +#define LIGER_MOE_BWD_DISPATCH_ENTRY_C(Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, TM) +#define LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM90(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(90, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) +#define LIGER_MOE_BWD_DISPATCH_ENTRY_SM90(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_C(90, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) +#define LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM100(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(100, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) +#define LIGER_MOE_BWD_DISPATCH_ENTRY_SM100(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_C(100, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) + +static const DispatchEntryBwd kDispatchTableBwd[] = { +#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 90 + LIGER_MOE_BWD_DISPATCH_CONFIGS_SM90(LIGER_MOE_BWD_DISPATCH_ENTRY_SM90, LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM90) +#endif +#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 100 + LIGER_MOE_BWD_DISPATCH_CONFIGS_SM100(LIGER_MOE_BWD_DISPATCH_ENTRY_SM100, LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM100) +#endif +}; +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_SM100 +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM100 +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_SM90 +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM90 +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_C +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_G_C +static constexpr int kNumDispatchEntriesBwd = + sizeof(kDispatchTableBwd) / sizeof(kDispatchTableBwd[0]); + +} // namespace liger + +// Joint fwd+bwd tuning table (auto-generated by benchmarks/tune_moe_fwd_bwd). +// We only consume the Bwd_* fields here; moe.cu consumes the same table's +// Fwd_* fields. Picking fwd and bwd from a single table guarantees the two +// directions are tuned as a pair — in particular their TileM agrees. +#include "moe_fwd_bwd_tuning_configs_sm100.cuh" + +namespace liger { + +static constexpr int kBwdRuntimeNsplitSeed = 8; + +// TunedConfigBwd is the bwd-only projection of TunedConfigFwdBwd. The rest +// of the dispatch code (validity check, dispatch-entry match) was written +// against this struct and is reused unchanged. +struct TunedConfigBwd { + int Compute; + int TK, TKE, D, I; + int NSplit2; + int TileN1, TileK1, Stages1; + int TileM3, TileN3, TileK3, Stages3; + int EpiChunkN1, EpiChunkN25, EpiChunkN34; + int CommNumStages; + int TileM; // comm tile + int GemmTileM; // gemm tile (auto path: == TileM; FORCE_CONFIG may set 64) +}; + +// Project a TunedConfigFwdBwd row down to its bwd template params. The tuned +// table has no GemmTileM column, so the auto path keeps gemm == comm (TileM). +static TunedConfigBwd project_bwd(const TunedConfigFwdBwd& c, int compute) { + return TunedConfigBwd{ + compute, + c.TK, c.TKE, c.D, c.I, + c.Bwd_NSplit2, + c.Bwd_TileN1, c.Bwd_TileK1, c.Bwd_Stages1, + c.Bwd_TileM3, c.Bwd_TileN3, c.Bwd_TileK3, c.Bwd_Stages3, + c.Bwd_EpiChunkN1, c.Bwd_EpiChunkN25, c.Bwd_EpiChunkN34, + c.Bwd_CommNumStages, + c.Bwd_TileM, c.Bwd_TileM, + }; +} + +// Find the dispatch table entry matching a TunedConfigBwd's knobs. +static const DispatchEntryBwd* find_dispatch_entry_bwd(const TunedConfigBwd& c) { + for (int i = 0; i < kNumDispatchEntriesBwd; ++i) { + const auto& e = kDispatchTableBwd[i]; + if (e.Compute == c.Compute && + e.NSplit2 == c.NSplit2 && + e.TileN1 == c.TileN1 && e.TileK1 == c.TileK1 && e.Stages1 == c.Stages1 && + e.TileM3 == c.TileM3 && e.TileN3 == c.TileN3 && e.TileK3 == c.TileK3 && + e.Stages3 == c.Stages3 && + e.EpiChunkN1 == c.EpiChunkN1 && + e.EpiChunkN25 == c.EpiChunkN25 && e.EpiChunkN34 == c.EpiChunkN34 && + e.CommNumStages == c.CommNumStages && + e.TileM == c.TileM && e.GemmTileM == c.GemmTileM) { + return &e; + } + } + return nullptr; +} + +// Structural divisibility checks on (D, I) for a given tuned config. +// Mirrors the constraint list in moe_fwd_bwd_tune_configs.hpp. +static bool tuned_config_valid_bwd(int D, int I, const TunedConfigBwd& c) { + // The current Blackwell UMMA MLP1/MLP1-act consumers only support the + // 128-row comm tile. Keep mirrored SM90 TileM=64 tuning rows out of SM100 + // auto-dispatch until that TMEM epilogue layout is implemented. + if (c.Compute == 100 && c.TileM != 128) return false; + if (c.Compute == 100) { + if (D % c.TileK1 != 0) return false; // mlp1/mlp2_t/mlp5 K + if (D % (2 * c.TileN1) != 0) return false; // mlp5 N (= TileN2) + if (I % (2 * c.TileN1) != 0) return false; // mlp2_t N + if (I % c.TileK1 != 0) return false; // mlp5 K + if (I % c.TileK3 != 0) return false; // mlp3/mlp4 K + } + if (D % 8 != 0 || I % 8 != 0) return false; // vectorization + if (c.Compute == 100) { + if (D % 256 != 0 || I % 256 != 0) return false; + } else { + if (!((c.TileM3 == 256 && c.TileN3 == 128) || + (c.TileM3 == 128 && c.TileN3 == 256))) return false; + if (D % c.TileM3 != 0 || I % c.TileN3 != 0) return false; + } + return true; +} + +static bool find_tuned_table_bwd(int compute, int n_pes, + const TunedConfigFwdBwd** table, + int* table_n) { + const TunedConfigFwdBwdTable* tables = + (n_pes <= 1) ? kTunedConfigTablesSingle : kTunedConfigTablesMulti; + const int tables_n = + (n_pes <= 1) ? kNumTunedConfigTablesSingle : kNumTunedConfigTablesMulti; + for (int i = 0; i < tables_n; ++i) { + if (tables[i].Compute == compute) { + *table = tables[i].configs; + *table_n = tables[i].count; + return *table_n > 0; + } + } + return false; +} + +// Nearest-neighbor lookup in log-space (TK, TKE, D, I), where TK = T*top_k +// is the total routed-token count (the bwd's M-axis post-dispatch) and +// TKE = TK / E_local (E_local = E / n_pes) is the per-LOCAL-expert avg +// K-range. mlp3/mlp4 walk each PE's local experts, so adding TKE to the +// distance lets the dispatcher distinguish configs whose walk-scheduler +// load differs even when (TK, D, I) is the same. +// Configs that violate structural constraints for the runtime (D, I) are +// skipped — the nearest tuned point is not always valid for the actual shape. +// If fwd_tile_m > 0, configs with TileM != fwd_tile_m are filtered out: +// BWD's sort-output reads (tile_expert_ids granularity, expert_offsets +// alignment, x_sorted slot layout) must match the forward's TileM. +// Out-param API instead of pointer-into-static-array because each candidate +// TunedConfigBwd is projected on the fly from kTunedConfigsFwdBwd (no stable +// address). Returns true when a valid config was found. +static bool find_nearest_tuned_bwd(int compute, int TK, int TKE, int D, int I, + int fwd_tile_m, int n_pes, + TunedConfigBwd* out) { + // World-size class: single-GPU (n_pes <= 1) and multi-GPU tables are tuned + // and stored separately (mirrors find_nearest_tuned on the fwd side). + const TunedConfigFwdBwd* table = nullptr; + int table_n = 0; + if (!find_tuned_table_bwd(compute, n_pes, &table, &table_n)) return false; + bool found = false; + double best_dist = 1e30; + double lTK = std::log((double)TK); + double lTKE = std::log((double)TKE); + double lD = std::log((double)D); + double lI = std::log((double)I); + for (int i = 0; i < table_n; ++i) { + TunedConfigBwd c = project_bwd(table[i], compute); + // Forward and backward select the same architecture-specific comm tile + // (M128 on both architectures), independent of the tuned GEMM tile. + (void)fwd_tile_m; + if (!tuned_config_valid_bwd(D, I, c)) continue; + // Guard against TKE=0 in legacy rows (would yield -inf in log). + double dTK = std::log((double)c.TK) - lTK; + double dTKE = c.TKE > 0 + ? std::log((double)c.TKE) - lTKE + : 0.0; // legacy / placeholder rows contribute 0 to the TKE axis + double dD = std::log((double)c.D) - lD; + double dI = std::log((double)c.I) - lI; + double d = dTK*dTK + dTKE*dTKE + dD*dD + dI*dI; + if (d < best_dist) { best_dist = d; *out = c; found = true; } + } + return found; +} + +// ============================================================================ +// Auto-dispatch wrapper — pure tuned-table lookup +// ============================================================================ +// +// Picks the nearest tuned config (log-space distance over (TK, D, I)) from +// moe_fwd_bwd_tuning_configs.cuh and dispatches to the matching template +// instantiation in kDispatchTableBwd. Both tables are populated by the +// tune_moe_bwd sweep — shapes outside the swept range error out instead of +// falling back to a heuristic. +// +// fwd_tile_m: BWD's TileM must match the forward's TileM. Forward sort +// pads expert ranges to FWD's TileM, populates tile_expert_ids at that +// granularity, and writes x_sorted in FWD-TileM-aligned slots. BWD +// reads all three at the same granularity, so a mismatch corrupts both +// the LocalIter expert lookup and the symmetric-memory bounds. +// Callers must pass the FWD's TileM (= 128 for the default fwd path). +// +// Override path: LIGER_MOE_BWD_FORCE_CONFIG=NS2,TN1,TK1,S1,TM3,TN3,TK3,S3,EN1,EN25,EN34,CS,TM +// Lets a test suite sweep every compiled config without rebuilding. + +// Internal auto-dispatch: nearest tuned bwd config for the runtime shape (or +// LIGER_MOE_BWD_FORCE_CONFIG), then invoke the matching template instantiation. +// All shape info comes from the MoeBwdArgs bundle; the launcher writes the grads. +void moe_bwd_dispatch(const MoeBwdArgs& a, int fwd_tile_m) { + const int T = a.num_tokens; + const int D = a.hidden_dim; + const int I = a.intermediate_dim; + const int top_k = a.top_k; + const int num_experts = a.num_experts; + const int TK = T * top_k; + const int n_pes = nvshmem_team_n_pes(a.team); + const int E_local = (n_pes > 0) ? std::max(1, num_experts / n_pes) : num_experts; + // Keep backward lookup consistent with forward for tiny token batches. + const int TKE = std::max(1, TK / E_local); + const int compute = moe_detect_compute_dispatch_key(a.device, "moe_fused_bwd_bf16_auto"); + if (const char* s = std::getenv("LIGER_MOE_BWD_FORCE_CONFIG")) { + TunedConfigBwd tc{}; + tc.Compute = compute; + tc.GemmTileM = -1; // sentinel → default to comm TileM after parse + int vals[14] = {}; + int n = std::sscanf(s, "%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d", + &vals[0], &vals[1], &vals[2], &vals[3], &vals[4], &vals[5], &vals[6], + &vals[7], &vals[8], &vals[9], &vals[10], &vals[11], &vals[12], &vals[13]); + LIGER_CHECK(n == 13 || n == 14, + "LIGER_MOE_BWD_FORCE_CONFIG must have 13 or 14 comma-separated ints " + "(NS2,TN1,TK1,S1,TM3,TN3,TK3,S3,EN1,EN25,EN34,CS,TM[,GemmTM]); got '", s, "'"); + tc.NSplit2 = vals[0]; + tc.TileN1 = vals[1]; + tc.TileK1 = vals[2]; + tc.Stages1 = vals[3]; + tc.TileM3 = vals[4]; + tc.TileN3 = vals[5]; + tc.TileK3 = vals[6]; + tc.Stages3 = vals[7]; + tc.EpiChunkN1 = vals[8]; + tc.EpiChunkN25 = vals[9]; + tc.EpiChunkN34 = vals[10]; + tc.CommNumStages = vals[11]; + tc.TileM = vals[12]; + tc.GemmTileM = (n == 14) ? vals[13] : tc.TileM; + if (tc.GemmTileM <= 0) tc.GemmTileM = tc.TileM; // omitted → gemm == comm + // The wrapper selects the architecture-specific comm tile consistently + // with forward; the parsed TileM remains the GEMM tuning field. + (void)fwd_tile_m; + const DispatchEntryBwd* de = find_dispatch_entry_bwd(tc); + LIGER_CHECK(de != nullptr, + "LIGER_MOE_BWD_FORCE_CONFIG=", s, " is not in the compiled dispatch " + "table. Add it to LIGER_MOE_BWD_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); + de->fn(a, kBwdRuntimeNsplitSeed); + return; + } + + TunedConfigBwd tc{}; + LIGER_CHECK(find_nearest_tuned_bwd(compute, TK, TKE, D, I, fwd_tile_m, n_pes, &tc), + "moe_bwd_fwd_bf16_auto: no valid tuned config for compute=", compute, + " shape (T=", T, + " K=", top_k, " TK=", TK, " TKE=", TKE, " D=", D, " I=", I, + " fwd_tile_m=", fwd_tile_m, + "). Either moe_fwd_bwd_tuning_configs.cuh has no entry with TileM=", + fwd_tile_m, ", or every matching entry violates structural " + "constraints for this (D, I). Add a compatible config to " + "LIGER_MOE_BWD_TUNE_CONFIGS and re-tune."); + + const DispatchEntryBwd* de = find_dispatch_entry_bwd(tc); + LIGER_CHECK(de != nullptr, + "moe_bwd_fwd_bf16_auto: tuned config (Compute=", tc.Compute, + " NS2=", tc.NSplit2, + " TN1=", tc.TileN1, "/", tc.TileK1, "/", tc.Stages1, + " TM3=", tc.TileM3, " TN3=", tc.TileN3, "/", tc.TileK3, + "/", tc.Stages3, + " EN1=", tc.EpiChunkN1, + " EN25=", tc.EpiChunkN25, " EN34=", tc.EpiChunkN34, + " CS=", tc.CommNumStages, " TM=", tc.TileM, + ") for shape (T=", T, " K=", top_k, " TK=", TK, " D=", D, " I=", I, + ") is not in the compiled dispatch table. " + "Add it to LIGER_MOE_BWD_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); + + de->fn(a, kBwdRuntimeNsplitSeed); +} + +// Returns the compiled bwd config table as a list of 13-tuples +// (Compute, NSplit2, TileN1, TileK1, Stages1, TileM3, TileN3, TileK3, +// Stages3, EpiChunkN1, EpiChunkN25, EpiChunkN34, CommNumStages). +std::vector> +moe_bwd_list_compiled_configs() { + std::vector> out; + out.reserve(kNumDispatchEntriesBwd); + for (int i = 0; i < kNumDispatchEntriesBwd; ++i) { + const auto& e = kDispatchTableBwd[i]; + out.emplace_back(e.Compute, e.NSplit2, + e.TileN1, e.TileK1, e.Stages1, + e.TileM3, e.TileN3, e.TileK3, e.Stages3, + e.EpiChunkN1, e.EpiChunkN25, e.EpiChunkN34, + e.CommNumStages); + } + return out; +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.cuh new file mode 100644 index 00000000..aede8fe5 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.cuh @@ -0,0 +1,219 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MoE Backward: MLP compute for local + remote tiles +// ═══════════════════════════════════════════════════════════════════ +// +// MLP-only device function — called by non-comm warps. In BWD comm uses +// warps 1-2; warp 3 enters the MLP path on SM100 and exits immediately on SM90. +// +// Execution flow (UNIFIED — local tile loop removed, mirrors moe_fused_fwd): +// EVERY tile (local or remote) flows through the comm get→staging→MLP→put +// path, driven by mlp_fused_bwd_dual over the RemoteMlpTileIteratorBwd / +// staging pipeline (there is no separate materialized local pass). +// +// mlp_fused_bwd_dual runs the grouped batch loop: +// SubBatch × (Phase 1 tile) → Phase 2 weight grads over the grouped window +// +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp_bwd_sm90.cuh" + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// MoE backward shared memory +// ═══════════════════════════════════════════════════════════════════ + +template +struct MoeBwdSmem { + MlpFusedBwdSmem mlp; + CommSmem comm; +}; + +// ═══════════════════════════════════════════════════════════════════ +// moe_fused_bwd — local phase → barrier → remote phase +// ═══════════════════════════════════════════════════════════════════ + +template 1 → Phase-1 (mlp1a/mlp2t/ + // silu/mlp5) steps through SubTiles GemmTileM-row sub-tiles per + // comm staging slot; the mlp3/4 ratios use Traits1::TileM (= GemmTileM). + int SubTiles, + int Compute = 90, + typename FusedIter, + // Phase 1 TMA + typename TmaLoadX, typename TmaLoadW1, + typename TmaStoreZ, typename TmaStoreU, typename TmaStoreV, + typename TmaLoadDY, typename TmaLoadW2T, typename TmaStoreDZ, + typename TmaLoadDU, typename TmaLoadDV, + typename TmaLoadB, typename TmaLoadC, typename TmaStoreDX, + // Phase 2 TMA + typename TmaLoadDYT3, typename TmaLoadZT3, typename TmaStoredA, + typename TmaLoadXT4, typename TmaLoaddUT4, typename TmaLoaddVT4, + typename TmaStoredB, typename TmaStoredC, + // Remote TMA: SAME C++ types as their local counterparts + // (the cute Layouts for both x_sorted and the staging ring use + // dynamic shape + _1 col stride + identical smem layout), so we + // can reuse the TmaLoadX/DY/StoreDX/etc. template params. Only + // the runtime gmem pointer differs between local and remote. + typename TmaLoadXR = TmaLoadX, + typename TmaLoadDYR = TmaLoadDY, + typename TmaStoreDXR = TmaStoreDX, + typename TmaLoadDYT3R = TmaLoadDYT3, + typename TmaLoadXT4R = TmaLoadXT4> +__device__ __forceinline__ void moe_fused_bwd( + MoeBwdSmem& smem, + FusedIter& iter, + // Phase 1 TMA (local) + TmaLoadX const& tma_load_x, TmaLoadW1 const& tma_load_b_fwd, + TmaLoadW1 const& tma_load_c_fwd, TmaStoreZ const& tma_store_z, + TmaStoreU const& tma_store_u, TmaStoreV const& tma_store_v, + TmaLoadDY const& tma_load_dy, TmaLoadW2T const& tma_load_a_col, + TmaStoreDZ const& tma_store_dz, + TmaLoadDU const& tma_load_du, TmaLoadDV const& tma_load_dv, + TmaLoadB const& tma_load_b_col, TmaLoadC const& tma_load_c_col, + TmaStoreDX const& tma_store_dx, + // Phase 2 TMA + TmaLoadDYT3 const& tma_load_dyt3, TmaLoadZT3 const& tma_load_zt3, + TmaStoredA const& tma_store_da, + TmaLoadXT4 const& tma_load_xt4, + TmaLoaddUT4 const& tma_load_dut4, TmaLoaddVT4 const& tma_load_dvt4, + TmaStoredB const& tma_store_db, TmaStoredC const& tma_store_dc, + // Remote TMA + TmaLoadXR const& tma_load_x_remote, + TmaLoadDYR const& tma_load_dy_remote, + TmaStoreDXR const& tma_store_dx_remote, + // Remote Phase 2 TMA: transposed staging sources for dA (Phase 3) + // and dB/dC (Phase 4). Without these, the remote phase reuses the + // local-only x_sorted/dy_sorted descriptors and produces wrong + // gradients for tokens routed across PEs. + TmaLoadDYT3R const& tma_load_dyt3_remote, + TmaLoadXT4R const& tma_load_xt4_remote, + // Comm pipe pointers (for building remote iter) + int* x_src_ready, int* x_src_consumed, + int* dy_src_ready, int* dy_src_consumed, + int* dst_ready, int* dst_consumed, + // Two per-pipe expert-id arrays (see CommBuffersBwd docstring). + // remote_tile_expert_ids_x is the X-pipe array used for the iter's + // tile.expert lookup (and Phase 2 mlp4 in this remote run). + // remote_tile_expert_ids_dy is the dY-pipe array used for Phase 2 + // mlp3 — gated by release_dy, so it stays stable across mlp3 even + // when release_src has already let comm refill the X-pipe array. + const int* remote_tile_expert_ids_x, + const int* remote_tile_expert_ids_dy, + const int* remote_tile_valid_rows_x, + // Dims + buffers + const MlpBwdDims& dims, + const MlpBwdBufs& bufs, + int col, + int grid_x, + int split, + int runtime_nsplit, + bool gemm_active) { + + // Target-based barriers — constructed ONCE so the per-CTA target + // counters accumulate monotonically across: + // local pass → global_barrier.wait() → remote pass + // Constructing a fresh barrier per wait() resets target=0 and is racy. + // Flat 1-D launch coordinates. blockIdx.x = flat_id ∈ [0, num_blocks). + // global_barrier counts ALL launched CTAs (num_blocks = gridDim.x). + // x_barrier is per logical Phase-1 column (col), stride runtime_nsplit. + int flat_id = (int)blockIdx.x; + + MlpBwdCtaBarrierT global_barrier(bufs.barrier_counter, + (int)gridDim.x * (int)gridDim.y); + MlpBwdCtaBarrierT x_barrier(&dims.phase_counter[col], + runtime_nsplit); + + // Initialize the fused iter's remote sub-iter from comm pipe pointers. + // (Local sub-iter was init'd by the caller in moe_bwd_kernel.) Comm warps + // filled staging during this CTA's local pass. Ticket-based iterator + // keys on absolute slot index, so pointer offsets are NOT pre-applied + // (mirrors RemoteMlpTileIterator's init in moe.cuh). + iter.init_remote( + remote_tile_expert_ids_x, + remote_tile_valid_rows_x, + smem.comm.per_cta_tiles, + col, + dims.n_gemm, + x_src_ready, x_src_consumed, + dy_src_ready, dy_src_consumed, + dst_ready, dst_consumed, + /*is_leader=*/(threadIdx.x == 0), + runtime_nsplit, + /*tma_enabled=*/(dims.tma_get_enabled != 0)); + + // Remote dims override: num_tokens / num_m_tiles use the staging- + // ring shape; total_m_tiles_in_pass uses the actual remote tile + // count (allows batches ≥ NumStages to run with bk_start wrapped + // into the live ring window). local_expert_start = 0 because + // RemoteIter reads tile_expert_ids written by nvshmem_comm_main_bwd + // (= TileIterator's LOCAL expert index). + MlpBwdDims remote_dims = dims; + remote_dims.num_tokens = dims.remote_num_tokens; + remote_dims.num_m_tiles = dims.remote_num_m_tiles; + remote_dims.total_m_tiles_in_pass = smem.comm.global_total; + remote_dims.num_batches = + (smem.comm.global_total + grid_x - 1) / grid_x; + remote_dims.local_expert_start = 0; + + // Alias per-pipe expert-id arrays into MlpBwdBufs. Phase 2 mlp4 + // reads X-pipe gated array (stable across mlp4 because release_src + // is deferred past mlp4). Phase 2 mlp3 reads dY-pipe gated array. + MlpBwdBufs remote_bufs = bufs; + remote_bufs.expert_for_k_block_mlp4 = remote_tile_expert_ids_x; + remote_bufs.expert_for_k_block_mlp3 = remote_tile_expert_ids_dy; + remote_bufs.valid_rows_mlp4 = remote_tile_valid_rows_x; + + // Grid-uniform activity predicate. The unified path stages local tiles even + // at world size one, so every valid non-empty launch is active. + bool remote_active = + (dims.num_pes > 0) && (dims.num_tokens > 0); + + // The remote staging ring (L = MC·CommNumStages tiles, MC = n_gemm/NC) is a + // STREAMING buffer: when global_total > L it wraps and physical slots are + // reused across laps. This is correct — mlp_fused_bwd_dual rescans the live + // per-slot expert ids PER BATCH (see its loop), so reused slots are always + // attributed to the expert currently resident, and each slot's reduction is + // committed before comm refills it. No L ≥ global_total requirement. + + // UNIFIED PATH: there is no separate materialized local MLP pass. The shared + // TileIterator enumerates LOCAL experts (p=0) alongside remote ones; the comm + // get warps stage local X/dY from local symmetric memory and the put warps + // drain local dX back to local dx_sorted, so EVERY tile (local or remote) + // flows through the RemoteMlpTileIteratorBwd / staging pipeline that + // mlp_fused_bwd_dual drives. Only the staging-ring (remote) TMAs and dims/bufs + // are passed; the weight/intermediate descriptors are shared across all tiles. + // efkb_stride maps a (GemmTileM-granular) Phase-2 K-block to its parent + // COMM-tile slot in tile_expert_ids (which the comm warps emit once per slot). + // = CommTileM/TileK = Traits1::TileM·SubTiles/Traits4::TileK. The integer + // division floors each GEMM sub-tile to its parent communication slot's expert (the + // interpolation). SubTiles=1 → Traits1::TileM/TileK, unchanged. + constexpr int kEfkbStride4Remote = Traits1::TileM * SubTiles / Traits4::TileK; + mlp_fused_bwd_dual< + Traits1, Traits2T, Traits3, Traits4, Traits5, NSplit2, SubBatch, SubTiles, Compute>( + smem.mlp, iter, remote_active, + kEfkbStride4Remote, + // Phase 1 shared weights/intermediates + tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_u, tma_store_v, + tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, + // Phase 2 shared + tma_load_zt3, tma_store_da, + tma_load_dut4, tma_load_dvt4, tma_store_db, tma_store_dc, + // Staging-ring data TMAs (X/dY/dX, dYᵀ/Xᵀ) + tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, + tma_load_dyt3_remote, tma_load_xt4_remote, + remote_dims, remote_bufs, + global_barrier, x_barrier, + flat_id, col, grid_x, split, /*num_splits=*/runtime_nsplit, gemm_active); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.inc b/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.inc new file mode 100644 index 00000000..7ea7040e --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.inc @@ -0,0 +1,2091 @@ +// Fused MoE Backward: sort → fwd-dispatch (for X) → reverse-combine (for dY) +// → MLP_bwd + comm overlap → reverse-dispatch (for dX). +// Compiled with separable compilation. Do NOT link liger_comm_kernels library. +// +// Torch-free core (liger_cute): ported from LigerCommKernels' moe_bwd.cu. The +// device kernel + templated launcher are upstream-verbatim (namespace liger). +// Tensor-carrying backward calls are exported through TVM FFI and lower directly +// into MoeBwdArgs. +#define LIGER_CUTE_BUILDING 1 + +#include +#include // pulls in nvshmemi_device_state_d, required by the + // inline device helpers in . +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +// SM100 2SM (paired-CTA cluster) TMA descriptor factories (make_tma_copy_A_sm100 +// / make_tma_copy_B_sm100) and the SM100_TMA_2SM_LOAD copy atom used below to +// build pair-aware Phase-2 mlp3/mlp4 descriptors. Already pulled in +// transitively via moe_bwd.cuh -> mlp_bwd.cuh -> mlp3.cuh/mlp4.cuh; included +// directly here too since this TU calls these factories itself. +#include + +#include "liger_cute/moe.h" +#include "liger_cute/check.h" +#include "liger_cute/detail/status.h" + +#include "moe_sm90.cuh" +#include "moe_bwd_sm90.cuh" +#include "moe_symm_config_sm90.cuh" +#include "moe_launch.h" // MoeFwdArgs / MoeBwdArgs (shared with the tuner) +#include "moe_utils.cuh" +#include "moe_fwd_bwd_tune_configs_sm90.hpp" +#include "moe_dispatch_configs_sm90.cuh" +#include "mlp_comms_bwd_sm90.cuh" +#include "tile_iterator_bwd_sm90.cuh" +#include "dispatch.cuh" +#include "dispatch_bwd.cuh" +#include "combine_bwd.cuh" +#include "tma_copy_atom_sm90.cuh" +#include "radix_sort.cuh" +#include "buffer_pool.cuh" // shim: liger:: aliases of liger_cute::detail infra + +using namespace cute; + +namespace liger { + +// 0 = build every dispatch family (developer fallback). CMake sets this to 90 +// for Hopper builds and 100 for Blackwell builds so normal native builds do not +// instantiate kernels for hardware that cannot run in that binary. +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif + +// External symbol: moe_configure_symmetric (defined in moe.cu) is reused. +// MoeSymmConfig + get_symm_config() now come from the shared header (included +// below), so this TU can't drift from moe.cu's layout. + +// Global iteration counter — same role as the fwd counter in moe.cu. +// Lives in device memory (allocated from BufferPool under the key +// "moe_bwd_iter_counter"), atomicAdded at the end of each kernel by +// CTA 0 thread 0. See PeSync::advance_iteration_counter() and the +// detailed comment in mlp_comms.cuh for why the counter has to be +// device-resident: a host-side atomic would be baked into a captured +// CUDA graph as a constant, breaking the cross-PE wait-until-greater +// protocol on every replay. + +// ============================================================================ +// MoeBwdConfig — bundles bwd-side traits + comm/grid constants +// ============================================================================ +// +// All five MLP_bwd phase traits + grid params. Mirrors MoeFusedConfig. +// Trait choices match the values that mlp_bwd.cu landed on after the +// 2026-04-26 perf tuning round (see project_fused_bwd_tile_bump.md). + +// Trait defaults match mlp_bwd.cu's validated set (lines 28–33). MoE +// backward shares mlp_fused_bwd; any drift from those values has hit +// illegal-instruction / smem-overrun bugs in the past. +template < + int NSplit2_ = 4, // Phase 2 (mlp3/mlp4) virtual n-split (autotuning surface) + // SM100 uses Stage-3 Phase-1 traits to leave room for independent full-size + // X/dY TMA bounce buffers. SM90 dispatch rows provide their tuned values. + int TileM1_ = 128, + int TileN1_ = 128, + int TileK1_ = 64, + int Stages1_ = 3, + int EpiChunkN1_ = 32, // mlp1_act + int TileN2_ = 256, + int TileK2_ = 64, + int Stages2_ = 3, + int TileM2T_ = 128, // Phase 1b' (mlp2_t): TileM ∈ {64, 128} + // mlp3/mlp4 cooperative layout (#102) supports only (256,128) or (128,256) + // and needs Stages=2 to fit the 228 KiB smem cap — Stages=4 at TileN=256 + // alone is ~384 KiB. (256,128) M-split is the smallest-smem legal choice + // (~192 KiB; the binding floor is mlp1_act ≈ 216 KiB). + int TileM3_ = 256, // mlp3 cooperative: (256,128) M-split + int TileN3_ = 128, + int TileK3_ = 64, + int Stages3_ = 2, + int TileM4_ = 128, // mlp4 cooperative: swap of mlp3 → (128,256) + int TileN4_ = 256, + int TileK4_ = 64, + int Stages4_ = 2, + int TileM5_ = 128, + int TileN5_ = 256, + int TileK5_ = 64, + int Stages5_ = 3, + int EpiChunkN25_ = 64, // shared by mlp2_t and mlp5 (same smem shape) + int EpiChunkN34_ = 64, // shared by mlp3 and mlp4 (same smem shape) + typename Element_ = bfloat16_t, + int CommNumStages_ = 2, + int NC_ = 4, + // Phase-1 sub-batch grouping count: run SubBatch Phase-1 sub-batches + // (mlp1/2/5) before one Phase-2 (mlp3/4) over their combined K-window. + // 1 = original per-batch loop. The only structural requirement is a ring + // deeper than the group span (CommNumStages·mc_remote > SubBatch·grid_x); + // the grouped window wraps the ring exactly as the un-grouped loop does. + int SubBatch_ = 2, + // CommTileM_ decouples the COMM staging tile (ring slots, semaphores, + // tile_expert_ids granularity) from the Phase-1 GEMM tile (Traits1::TileM = + // mlp1a/mlp2t/mlp5 = the GemmTileM hyperparameter ∈ {64,128}). When + // Traits1::TileM < CommTileM_, Phase-1 steps through + // kSubTiles = CommTileM_/Traits1::TileM sub-tiles of GemmTileM rows per + // communication slot (mirrors forward). Default == TileM1_ → kSubTiles=1 + // → bit-identical to the old coupled path. + int CommTileM_ = 128, + int Compute_ = 90> +struct MoeBwdConfig { + using Element = Element_; + static constexpr int kActStoreTileM = 64; + using Traits1 = Mlp1Traits < + Element, TileM1_, TileN1_, TileK1_, Stages1_, EpiChunkN1_, + kActStoreTileM>; // Phase 1a + using Traits2T = Mlp2TTraits; // Phase 1b' + // Blackwell always uses the paired-CTA MLP3/MLP4 path. Hopper keeps the + // existing single-CTA traits and never instantiates SM100-only machinery. + // The existing TileK3_/Stages3_/EpiChunkN34_ fields remain the shared + // tuning surface for both weight-gradient kernels. + static constexpr bool kUsesTwoSm = Compute_ == 100; + static constexpr int kClusterM = kUsesTwoSm ? 2 : 1; + using Traits3 = cute::conditional_t< + kUsesTwoSm, + Mlp3Traits2Sm, + Mlp3Traits >; // Phase 2 mlp3 + using Traits4 = cute::conditional_t< + kUsesTwoSm, + Mlp4Traits2Sm, + Mlp4Traits >; // Phase 2 mlp4 + using Traits5 = Mlp5Traits ; // Phase 1d + + // No divisibility constraints between NSplit/NC/CommNumStages. + // MLP side is ticket-based (RemoteMlpTileIteratorBwd); comm side + // uses K = CommNumStages. Only runtime requirement is that NC + // evenly divides gridDim.x · gridDim.y so MC is integer. + + static constexpr int kNSplit2 = NSplit2_; + static constexpr int kNC = NC_; + static constexpr int kCommNumStages = CommNumStages_; + static constexpr int kSubBatch = SubBatch_; + // kTileM is the COMM staging tile (ring slots / semaphores / tile_expert_ids + // granularity), decoupled from the Phase-1 GEMM tile. kGemmTileM = + // Traits1::TileM (mlp1/2/5 GEMM tile); + // the mlp3/4 ratios (efkb_stride / bk↔Kblock / ring_kb) are computed from it. + static constexpr int kTileM = CommTileM_; + static constexpr int kGemmTileM = Traits1::TileM; // Phase-1 gemm tile (64/128) + static constexpr int kSubTiles = CommTileM_ / Traits1::TileM; + static constexpr int kCompute = Compute_; + static constexpr int kNumThreads = 384; + + static_assert(CommTileM_ % Traits1::TileM == 0, + "CommTileM must be an integer multiple of the Phase-1 GEMM TileM (Traits1::TileM)"); + static_assert(kSubTiles == 1 || kSubTiles == 2, + "CommTileM/GemmTileM must be 1 or 2"); + // mlp1a/mlp2t/mlp5 must share the same Phase-1 GEMM tile so their sub-tiling + // and the Z/dU/dV/dZ buffer rows stay aligned. + static_assert(Traits1::TileM == Traits2T::TileM && Traits1::TileM == Traits5::TileM, + "Phase-1 GEMM TileM must match across mlp1a / mlp2t / mlp5"); + + // Compile-time smem budget. H100 dynamic-smem cap is 228 KiB. The bwd + // union sits very close to that ceiling (Mlp3/Mlp4 with TileK=128 are + // 224 KiB on their own, and pipe storage pushes the total a bit + // further), so we use the full hardware cap here. The runtime-sized + // remote_offsets region appended by the launcher is ≤ a few hundred + // bytes for autotuner-scale E and num_pes; relying on the kernel's + // cudaFuncSetAttribute call to surface any overflow. + static constexpr size_t kSmemBytes = + sizeof(MoeBwdSmem); + static constexpr size_t kSmemBudget = 228 * 1024; + static_assert(kSmemBytes <= kSmemBudget, + "MoE fused bwd kernel smem footprint exceeds H100 228 KiB cap — " + "reduce TileN/TileK/Stages of one of the bwd phase traits"); + + using LocalIter = LocalMlpTileIteratorBwd; + using RemoteIter = RemoteMlpTileIteratorBwd< + Element, kCommNumStages, kNC, kTileM>; + using FusedIter = FusedMlpTileIteratorBwd< + Element, kCommNumStages, kNC, kTileM>; +}; + +// ============================================================================ +// MoeBwdParams — input/output pointers + sort buffers + comm + dims +// ============================================================================ + +template +struct MoeBwdParams { + using Element = typename Config::Element; + + // Router outputs (caller-provided). + const int* expert_indices; // [T, K] + const Element* expert_weights; // [T, K] + + // Sort buffers. + SortBuffers sort; + + // Inputs. + // Note: there is no `tokens` (X) field — BWD reads token data only + // via `x_sorted` (FWD's expert-sorted view from its symmetric stack + // entry). Removed 2026-05-18 after confirming p.tokens was set in the + // host wrapper and never read by the kernel or any helper. + const Element* dY; // [T, D] (regular device) + // Forward MoE output Y_fwd. Caller is responsible for providing this + // — either retained from the most recent forward pass or recomputed + // under gradient checkpointing. Read by combine_tokens_bwd to compute + // dW[t,k] = Σ_d dO[t,d] * Y_fwd[slot(t,k), d]. Layout: [total_slots, D] + // (one row per routed (token, expert-slot) pair, in the same + // expert-sorted order produced by the forward dispatch). + const Element* output_buffer; + + // Outputs (caller-allocated). + Element* dX_out; // [T, D] (regular device) + Element* dB_out; // [epp, I, D] + Element* dC_out; // [epp, I, D] + Element* dA_out; // [epp, D, I] + Element* dW_out; // [T, K] + + // Symmetric staging buffers (allocated by allocate_moe_bwd_buffers). + Element* x_sorted; // [total_slots, D] — fwd dispatch dst + Element* dy_sorted; // [total_slots, D] — combine_bwd dst + Element* dx_sorted; // [total_slots, D] — MLP_bwd dX dst + + // Comm (forward, for sort/dispatch/pe_sync) and bwd (for X+dY gets, dX puts). + CommBuffers comm; + CommBuffersBwd comm_bwd; + + // Dimensions. + int num_tokens; + int hidden_dim; + int intermediate_dim; + int top_k; + int total_slots; + int num_m_tiles; // sort input tiles = T / TileM + + // Output stride for dB/dC/dA per-expert. + int experts_per_pe; +}; + +// ============================================================================ +// MoeBwdBuffers — owned device + symmetric allocations +// ============================================================================ + +template +struct MoeBwdBuffers { + using Element = typename Config::Element; + + // Symmetric (NVSHMEM). + Element* x_sorted; + Element* dy_sorted; + Element* dx_sorted; + + // Symmetric offsets. + int* expert_offsets_sym; // for sort + + // Device-only. + int* cta_counter; + int* phase_counter; // mlp_bwd uses this + int* barrier_counter; // mlp_global_barrier + Element* z_buf; // [scratch_rows, I] + Element* du_buf; // [scratch_rows, I] -> dU after silu_bwd + Element* dv_buf; // [scratch_rows, I] + Element* dz_buf; // [scratch_rows, I] + int scratch_rows; // max(local padded rows, remote ring rows) + int* expert_for_k_block; // num_k_blocks; populated post-sort + + // Phase 2 per-expert K-range scratch. Filled in bwd preamble (once + // per pass) by mlp_bwd_phase_2_build_ranges; shared between mlp3 and + // mlp4 (they share TileK). Sized experts_per_pe. + int* expert_k_starts; + int* expert_k_ends; + + // Comm staging (X + dY gets, dX puts). + Element* x_staging; + Element* dy_staging; + Element* dx_staging; + + int* x_src_ready; + int* x_src_consumed; + int* dy_src_ready; + int* dy_src_consumed; + int* dst_ready; + int* dst_consumed; + // Two per-staging-slot expert-id arrays, one per get pipe. Comm warps + // write the X-side array during the X getmem and the dY-side array + // during the dY getmem. Phase 2 mlp4 reads tile_expert_ids_x; Phase 2 + // mlp3 reads tile_expert_ids_dy. See CommBuffersBwd docstring for why. + int* tile_expert_ids_x; + int* tile_expert_ids_dy; + int* tile_valid_rows_x; + + // Sort-internal scratch. + int* tile_expert_counts; + int* sorted_token_ids; + int* token_expert_slots; + int* cta_done; + int* cta_sums; +}; + +template +static MoeBwdBuffers allocate_moe_bwd_buffers( + BufferPool& pool, cudaStream_t stream, + int grid_x, int num_blocks, + int hidden_dim, int intermediate_dim, int num_tokens, int top_k, + int experts_per_pe, int total_slots, int num_m_tiles) { + + using Element = typename Config::Element; + constexpr int kTileM = Config::kTileM; + // kCommNumStages / kNSplit / kNC are not needed here: the comm staging is + // sized at the config-independent UB (kCommStagesUB · kTileMUB) below. + + MoeBwdBuffers b; + + auto& scfg = get_symm_config(); + LIGER_CHECK(scfg.initialized, + "Call moe_configure_symmetric before moe_bwd_fwd_bf16"); + + size_t symm_slot_bytes = + static_cast(scfg.max_total_slots) * scfg.hidden_dim * sizeof(Element); + + // ── Symmetric ──────────────────────────────────────── + // x_sorted comes from the FWD pass — the caller passes its symm + // pointer in the `x_sorted` torch tensor (FWD's pool entry under + // key "moe_x_sorted"). BWD reads it directly and never allocates + // or writes to its own copy. + b.dy_sorted = static_cast(pool.get_symmetric("moe_bwd_dy_sorted", symm_slot_bytes)); + b.dx_sorted = static_cast(pool.get_symmetric("moe_bwd_dx_sorted", symm_slot_bytes)); + + // dy_sorted MUST be zeroed: combine_tokens_bwd only writes routed slots + // (top_k per token). Padding slots stay 0, contributing nothing to MLP_bwd. + cudaMemsetAsync(b.dy_sorted, 0, symm_slot_bytes, stream); + + // all_expert_offsets is no longer allocated here — the caller passes + // the broadcasted offsets in via the BWD's `expert_offsets` tensor + // (populated by FWD's broadcast_expert_offsets and returned as part + // of its intermediates). + b.expert_offsets_sym = static_cast( + pool.get_symmetric("moe_sort_offsets", + (scfg.max_num_experts + 1) * sizeof(int))); + + // ── Device-only ────────────────────────────────────── + b.cta_counter = static_cast(pool.get_device("moe_bwd_cta_counter", sizeof(int))); + b.phase_counter = static_cast(pool.get_device("moe_bwd_phase_counter", grid_x * sizeof(int))); + b.barrier_counter = static_cast(pool.get_device("moe_bwd_barrier", sizeof(int))); + cudaMemsetAsync(b.cta_counter, 0, sizeof(int), stream); + cudaMemsetAsync(b.phase_counter, 0, grid_x * sizeof(int), stream); + cudaMemsetAsync(b.barrier_counter, 0, sizeof(int), stream); + + // Local Phase 1 indexes scratch by padded token row, while the remote pass + // indexes the same buffers by staging-ring slot. A deep comm ring can exceed + // the local padded extent (for example, 33 MC columns * CS8 * TileM128 = + // 33792 rows versus 17408 local rows for Mixtral-8x7B T=8K). Size for both + // coordinate spaces or a valid remote slot writes past z/du/dv/dz. + const int remote_m_tiles = + (num_blocks / Config::kNC) * Config::kCommNumStages; + const size_t remote_rows = + static_cast(remote_m_tiles) * kTileM; + const size_t scratch_rows = + std::max(static_cast(total_slots), remote_rows); + b.scratch_rows = static_cast(scratch_rows); + size_t intermediate_bytes = + scratch_rows * intermediate_dim * sizeof(Element); + b.z_buf = static_cast(pool.get_device("moe_bwd_z_buf", intermediate_bytes)); + b.du_buf = static_cast(pool.get_device("moe_bwd_du_buf", intermediate_bytes)); + b.dv_buf = static_cast(pool.get_device("moe_bwd_dv_buf", intermediate_bytes)); + b.dz_buf = static_cast(pool.get_device("moe_bwd_dz_buf", intermediate_bytes)); + + // expert_for_k_block: one int per K-block. The K-block / m-tile size + // ratio depends on Traits3::TileK vs Traits1::TileM; see + // build_expert_for_k_block which uses both dims directly. + // Size by scfg.max_total_slots, not per-call total_slots — the buffer + // pool caches by name and won't resize on later calls. Otherwise the + // first (smallest) config locks in a tiny allocation and a later + // larger config (e.g. E=64, K=8) reads OOB past the cached size. + int max_num_k_blocks = scfg.max_total_slots / Config::Traits3::TileK; + b.expert_for_k_block = static_cast( + pool.get_device("moe_bwd_expert_for_k_block", max_num_k_blocks * sizeof(int))); + + // Phase 2 per-expert K-range scratch — sized by scfg.max_num_experts + // (the pool caches by name, so we have to size to the worst case + // across configs; same rationale as expert_for_k_block above). + // Shared between mlp3 and mlp4 (they share TileK). + size_t expert_range_bytes = scfg.max_num_experts * sizeof(int); + b.expert_k_starts = static_cast( + pool.get_device("moe_bwd_expert_k_starts", expert_range_bytes)); + b.expert_k_ends = static_cast( + pool.get_device("moe_bwd_expert_k_ends", expert_range_bytes)); + + // ── Comm staging ───────────────────────────────────── + // Flat staging list — L = MC · kCommNumStages slots, where + // MC = (grid_x · NSplit) / NC. Comm strides MC through NumStages + // stages; MLP is ticket-based and keys on the absolute slot index + // modulo L. With NC = NSplit (legacy config), MC = grid_x and this + // collapses to grid_x · kCommNumStages — unchanged behavior. The live + // length MC·kCommNumStages (MC = (grid_x·NSplit)/NC) is bounded by + // comm_l_ub below; the buffers are sized at that UB, not the live value. + // Staging must be SYMMETRIC: the comm path issues nvshmemx_getmem/putmem + // whose LOCAL side (x/dy/dx staging) is the NIC-registered buffer. + // Device-only memory faults with IBV_WC_LOC_PROT_ERR over IB (works on + // NVLink either way, which masked it single-host). Size at an + // NC-independent upper bound (mc <= sm_count) so the symmetric + // allocation never has to grow across NC variants — it aborts on grow; + // larger-than-needed is fine. Mirrors the fwd moe.cu comm_l_ub sizing. + int sm_count = 0; + { + int dev = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev); + } + // Worst-case sizing across the whole tuner sweep: the symmetric/device + // pools cache by key and ABORT on grow, so the FIRST config to allocate a + // key locks its size. A small-TileM (64) or small-CommStages config must + // not lock in a buffer that a later TileM=256 / deeper-pipe config can't + // grow into. Size every comm buffer at the centralized TileM/stage maxima; + // the kernel indexes by the per-call + // (mc·kCommNumStages, kTileM) ≤ these UBs, so over-allocation is safe. + // Mirrors moe.cu's MoeSymmConfig.max_tile_m / max_comm_stages fix. + constexpr int kCommStagesUB = kMaxMoeCommStages; + constexpr int kTileMUB = kMaxMoeCommTileM; + int comm_l_ub = sm_count * kCommStagesUB; + size_t stage_slots_bytes = + (size_t)comm_l_ub * kTileMUB * scfg.hidden_dim * sizeof(Element); + b.x_staging = static_cast(pool.get_symmetric("moe_bwd_x_staging", stage_slots_bytes)); + b.dy_staging = static_cast(pool.get_symmetric("moe_bwd_dy_staging", stage_slots_bytes)); + b.dx_staging = static_cast(pool.get_symmetric("moe_bwd_dx_staging", stage_slots_bytes)); + + // X / dY pipe signals: per-slot, [L = MC * NumStages] ints. Comm + // producer indexes by `+ xc` (stride MC); MLP consumer is + // ticket-based and walks slot = (m_base + idx·gridDim.x) mod L. + // Allocate at the UB (comm_l_ub ≥ MC·kCommNumStages); the kernel only + // touches [0, MC·kCommNumStages). Per-call resets (below) zero the live prefix. + size_t tile_sig_bytes = (size_t)comm_l_ub * sizeof(int); + b.x_src_ready = static_cast(pool.get_device("moe_bwd_x_src_ready", tile_sig_bytes)); + b.x_src_consumed = static_cast(pool.get_device("moe_bwd_x_src_consumed", tile_sig_bytes)); + b.dy_src_ready = static_cast(pool.get_device("moe_bwd_dy_src_ready", tile_sig_bytes)); + b.dy_src_consumed= static_cast(pool.get_device("moe_bwd_dy_src_consumed",tile_sig_bytes)); + cudaMemsetAsync(b.x_src_ready, 0, tile_sig_bytes, stream); + cudaMemsetAsync(b.x_src_consumed, 0, tile_sig_bytes, stream); + cudaMemsetAsync(b.dy_src_ready, 0, tile_sig_bytes, stream); + cudaMemsetAsync(b.dy_src_consumed,0, tile_sig_bytes, stream); + + // dX pipe signals: per-slot, [L = MC * NumStages] ints. UB-sized (see above). + size_t tile_pipe_bytes = (size_t)comm_l_ub * sizeof(int); + b.dst_ready = static_cast(pool.get_device("moe_bwd_dst_ready", tile_pipe_bytes)); + b.dst_consumed = static_cast(pool.get_device("moe_bwd_dst_consumed", tile_pipe_bytes)); + cudaMemsetAsync(b.dst_ready, 0, tile_pipe_bytes, stream); + cudaMemsetAsync(b.dst_consumed, 0, tile_pipe_bytes, stream); + + size_t eid_bytes = (size_t)comm_l_ub * sizeof(int); // UB-sized (see above) + // One array per get pipe (see MoeBwdBuffers / CommBuffersBwd docstrings). + // Same shape as the staging signal arrays. + b.tile_expert_ids_x = static_cast( + pool.get_device("moe_bwd_tile_expert_ids_x", eid_bytes)); + b.tile_expert_ids_dy = static_cast( + pool.get_device("moe_bwd_tile_expert_ids_dy", eid_bytes)); + b.tile_valid_rows_x = static_cast( + pool.get_device("moe_bwd_tile_valid_rows_x", eid_bytes)); + + // ── Sort scratch (sized to scfg max for stability across calls) ── + int max_sorted_slots = scfg.max_total_slots; + int max_m_tiles = scfg.max_total_slots / kTileM; + int max_tile_counts = max_m_tiles * scfg.max_num_experts; + + b.tile_expert_counts = static_cast( + pool.get_device("moe_bwd_sort_tile_counts", max_tile_counts * sizeof(int))); + b.sorted_token_ids = static_cast( + pool.get_device("moe_bwd_sort_sorted_ids", max_sorted_slots * sizeof(int))); + b.token_expert_slots = static_cast( + pool.get_device("moe_bwd_sort_token_slots", max_sorted_slots * sizeof(int))); + b.cta_done = static_cast( + pool.get_device("moe_bwd_sort_cta_done", num_blocks * sizeof(int))); + b.cta_sums = static_cast( + pool.get_device("moe_bwd_sort_cta_sums", + num_blocks * scfg.max_num_experts * sizeof(int))); + + cudaMemsetAsync(b.cta_done, 0, num_blocks * sizeof(int), stream); + cudaMemsetAsync(b.sorted_token_ids, 0xff, max_sorted_slots * sizeof(int), stream); + + return b; +} + +// ============================================================================ +// build_expert_for_k_block — small helper kernel section +// ============================================================================ +// +// Post-sort: derive expert_for_k_block[k] = tile_expert_ids[(k * TileK) / TileM]. +// Works for both directions: +// TileK > TileM (K-block spans multiple m-tiles): uses the first m-tile's expert. +// TileK < TileM (m-tile spans multiple K-blocks): all K-blocks in a m-tile +// resolve to the same tile_expert_ids slot via integer division. +// Assumes per-expert ranges are TileK-aligned, OR that any K-block straddling +// an expert boundary has only padding tokens in its second m-tile (ensured +// by zeroed dy_sorted + zero-filled X-padding from forward dispatch). + +__device__ __forceinline__ void build_expert_for_k_block( + int* expert_for_k_block, + const int* tile_expert_ids, + int num_k_blocks, + int num_actual_k_blocks, + int tile_k, + int tile_m, + int num_threads) { + // For k < num_actual_k_blocks, copy the expert from the sort output. + // For padding entries (k beyond expert_offsets[num_experts] / TileK), + // write a sentinel that the Phase 3/4 bounds check will reject. Without + // this, stale memory in tile_expert_ids[k] could pass the bounds check + // and cause Phase 4 to load garbage X/dU into the dB/dC accumulators. + for (int k = threadIdx.x; k < num_k_blocks; k += num_threads) { + expert_for_k_block[k] = (k < num_actual_k_blocks) + ? tile_expert_ids[(k * tile_k) / tile_m] + : -1; + } +} + +// ============================================================================ +// Fused kernel: sort → fwd-dispatch → combine_bwd → MLP_bwd+comm → dispatch_bwd +// ============================================================================ + +template +__global__ void __launch_bounds__(Config::kNumThreads) +moe_bwd_kernel( + // Local TMA + __grid_constant__ TmaLoadX const tma_load_x, + __grid_constant__ TmaLoadW1 const tma_load_b_fwd, + __grid_constant__ TmaLoadW1 const tma_load_c_fwd, + __grid_constant__ TmaStoreZ const tma_store_z, + __grid_constant__ TmaStoreU const tma_store_u, + __grid_constant__ TmaStoreV const tma_store_v, + __grid_constant__ TmaLoadDY const tma_load_dy, + __grid_constant__ TmaLoadW2T const tma_load_a_col, + __grid_constant__ TmaStoreDZ const tma_store_dz, + __grid_constant__ TmaLoadDU const tma_load_du, + __grid_constant__ TmaLoadDV const tma_load_dv, + __grid_constant__ TmaLoadB const tma_load_b_col, + __grid_constant__ TmaLoadC const tma_load_c_col, + __grid_constant__ TmaStoreDX const tma_store_dx, + __grid_constant__ TmaLoadDYT3 const tma_load_dyt3, + __grid_constant__ TmaLoadZT3 const tma_load_zt3, + __grid_constant__ TmaReduceDA const tma_reduce_da, + __grid_constant__ TmaLoadXT4 const tma_load_xt4, + __grid_constant__ TmaLoaddUT4 const tma_load_dut4, + __grid_constant__ TmaLoaddVT4 const tma_load_dvt4, + __grid_constant__ TmaReduceDB const tma_reduce_db, + __grid_constant__ TmaReduceDC const tma_reduce_dc, + // Remote TMA + __grid_constant__ TmaLoadXR const tma_load_x_remote, + __grid_constant__ TmaLoadDYR const tma_load_dy_remote, + __grid_constant__ TmaStoreDXR const tma_store_dx_remote, + // Remote Phase 2 TMA (staging-backed dY/X for dA/dB/dC) + __grid_constant__ TmaLoadDYT3R const tma_load_dyt3_remote, + __grid_constant__ TmaLoadXT4R const tma_load_xt4_remote, + // Dims + bufs + __grid_constant__ MlpBwdDims const mlp_dims, + __grid_constant__ MlpBwdBufs const mlp_bufs, + __grid_constant__ MoeBwdParams const p, + __grid_constant__ int const static_nsplit, + __grid_constant__ GetTmaDescsBwd const get_descs_bwd) { + + using Element = typename Config::Element; + using Traits1 = typename Config::Traits1; + using Traits2T = typename Config::Traits2T; + using Traits3 = typename Config::Traits3; + using Traits4 = typename Config::Traits4; + using Traits5 = typename Config::Traits5; + using FusedIter = typename Config::FusedIter; + + constexpr int kTileM = Config::kTileM; + constexpr int kNumThreads = Config::kNumThreads; + constexpr int kCommNumStages = Config::kCommNumStages; + constexpr int kNC = Config::kNC; + constexpr int kGetBoxRowsBwd = Config::kCompute == 100 + ? liger::kGetBoxRowsBwdSm100 + : liger::kGetBoxRowsBwdSm90; + + // Target-based cross-CTA barrier — single instance per kernel call. + SyncThreadsCtaCounterBarrier pe_barrier(p.comm.pe_sync.cta_counter, + (int)gridDim.x * (int)gridDim.y); + + // ── Phase 0: Sort — SKIPPED. + // Caller provides token_expert_slots, tile_expert_ids, and + // expert_offsets from the most recent forward call (these come out of + // moe_fused_fwd_bf16, populated by the forward + // kernel's sort_tokens). We trust those instead of re-sorting. + + // ── Phase 1: Broadcast expert offsets — SKIPPED. + // Forward returned the broadcasted offsets in the `expert_offsets` + // tensor we forward verbatim to p.comm.all_expert_offsets, and + // forward ended with nvshmem_team_sync(team), so every PE's writes + // are visible here. No bwd-side re-broadcast needed. + + // ── Phase 2a: Build expert_for_k_block (post-sort, before MLP) ─── + // Operates on caller-provided tile_expert_ids (forward's sort output, + // threaded through p.sort.tile_expert_ids). BWD_TileM is constrained + // to match FWD_TileM (enforced in moe_bwd_fwd_bf16_auto), so the + // FWD-populated array indexes correctly at BWD's tile granularity. + { + int num_k_blocks = p.total_slots / Traits3::TileK; + // Real data ends at expert_offsets[num_experts]; everything beyond + // is buffer-pool padding from the symmetric x_sorted allocation. + int total_real_rows = p.sort.expert_offsets[p.comm.num_experts]; + int num_actual_k_blocks = total_real_rows / Traits3::TileK; + // Both _mlp4 and _mlp3 pointers alias buf.expert_for_k_block in the + // local pass (set above), so writing through either populates the + // shared post-sort array. + build_expert_for_k_block( + const_cast(mlp_bufs.expert_for_k_block_mlp4), + p.sort.tile_expert_ids, + num_k_blocks, num_actual_k_blocks, + Traits3::TileK, kTileM, kNumThreads); + } + __syncthreads(); + + // ── Phase 2b: Forward dispatch — SKIPPED. + // p.x_sorted points to the forward's x_sorted (already dispatched). + // MLP TMAs read from there directly; nothing to do here. + + // ── Phase 3: Reverse-combine — dY → dy_sorted (+ dW side-output) ─ + // Caller-provided output_buffer is Y_fwd in expert-sorted layout + // [total_slots, D] (must match the expert-sort produced by the + // forward dispatch). combine_tokens_bwd computes: + // dW[t,k] = Σ_d dY[t,d] * Y_fwd[slot(t,k), d] → p.dW_out + // dy_sorted[slot,d]= dY[t,d] * w[t,k] → fed to mlp_fused_bwd + combine_tokens_bwd( + p.dY, + /*expert_out=*/p.output_buffer, + p.expert_weights, + p.sort.token_expert_slots, + p.dW_out, + p.dy_sorted, + p.hidden_dim, + p.num_tokens, + p.top_k); + __threadfence_system(); + __syncthreads(); + + p.comm.pe_sync.barrier_all(pe_barrier); + + // ── Phase 4: MLP_bwd with comm overlap ─────────────────────────── + { + extern __shared__ char raw_smem[]; + auto& smem = *reinterpret_cast*>(raw_smem); + + size_t offsets_region = sizeof(MoeBwdSmem); + offsets_region = (offsets_region + sizeof(int) - 1) & ~(sizeof(int) - 1); + size_t offsets_count = + (size_t)p.comm_bwd.num_pes * (p.comm_bwd.experts_per_pe + 1); + if (threadIdx.x == 0) { + smem.comm.remote_offsets = + reinterpret_cast(raw_smem + offsets_region); + smem.comm.remote_counts = + smem.comm.remote_offsets + offsets_count; + } + __syncthreads(); + + // Bwd TMA-GET bounce: 128-aligned region after offsets and counts. + // Mirrors the host smem sizing (align128(smem_mlp) + bounce). Only + // dereferenced when get_descs_bwd.enabled; harmless otherwise. + size_t bounce_off = offsets_region + + (offsets_count + + (size_t)p.comm_bwd.num_pes * p.comm_bwd.experts_per_pe) + * sizeof(int); + bounce_off = (bounce_off + 127) & ~((size_t)127); + char* get_bounce = raw_smem + bounce_off; + + // No sentinel write here: nvshmem_comm_prologue_bwd's warp-2 leader + // is guaranteed to write smem.comm.total_tiles (either to 0 in the + // early-exit case or to comm_count in the normal case) before the + // __syncthreads below. A `smem.comm.total_tiles = -1` from + // threadIdx.x == 0 racing with the prologue's warp-2 write is what + // caused the BWD-TM=64 + 2 PE deadlock: when thread 0's sentinel + // write landed AFTER warp 2's write, total_tiles ended up as -1, + // comm skipped nvshmem_comm_main_bwd, the comm pipes (x_src_ready + // etc.) never got a producer_release, and the MLP remote pass' + // iter.acquire_src spun forever. + nvshmem_comm_prologue_bwd(smem.comm, p.comm_bwd, static_nsplit); + __syncthreads(); + + int flat_id = (int)blockIdx.x; + int runtime_nsplit = static_nsplit; + int n_gemm = p.comm_bwd.n_gemm; + int grid_x = n_gemm / runtime_nsplit; + int col = flat_id / runtime_nsplit; + if (threadIdx.x == 0) { + int total_m_tiles = smem.comm.global_total; + runtime_nsplit = select_runtime_nsplit_bwd_from_tiles( + total_m_tiles, + mlp_dims.num_n_tiles_1, + mlp_dims.num_n_tiles_2t, + mlp_dims.num_n_tiles_5, + mlp_dims.num_k_tiles_1, + mlp_dims.num_k_tiles_2t, + mlp_dims.num_k_tiles_5, + p.comm_bwd.experts_per_pe, + (int)gridDim.x, + static_nsplit); + n_gemm = ((int)gridDim.x / runtime_nsplit) * runtime_nsplit; + grid_x = n_gemm / runtime_nsplit; + col = flat_id / runtime_nsplit; + int mlp_count = 0; + if (flat_id < n_gemm && col < total_m_tiles) + mlp_count = (total_m_tiles - 1 - col) / grid_x + 1; + smem.comm.per_cta_tiles = mlp_count; + smem.comm.runtime_nsplit = runtime_nsplit; + smem.comm.runtime_n_gemm = n_gemm; + smem.comm.runtime_grid_x = grid_x; + } + __syncthreads(); + runtime_nsplit = smem.comm.runtime_nsplit; + n_gemm = smem.comm.runtime_n_gemm; + grid_x = smem.comm.runtime_grid_x; + col = flat_id / runtime_nsplit; + int split = flat_id % runtime_nsplit; + bool gemm_active = flat_id < n_gemm; + + int warp_id = threadIdx.x / 32; + bool is_comm_warp = (warp_id == kCommBwdGetWarp0 || warp_id == kCommBwdPutWarp); + + if (is_comm_warp) { + if (smem.comm.total_tiles > 0) { + nvshmem_comm_main_bwd< + Element, kTileM, kCommNumStages, kNC, + kGetBoxRowsBwd>( + smem.comm, p.comm_bwd, &get_descs_bwd, get_bounce); + nvshmem_quiet(); + } + } else { + // MLP warps. The fused iter is initialized for the staging ring + // inside moe_fused_bwd (init_remote); every tile — local or remote — + // flows through that staging path, so there is no local sub-iter to + // seed here. Gap CTAs skip Phase 1; moe_fused_bwd gates it on + // gemm_active. + FusedIter fused_iter; + MlpBwdDims runtime_mlp_dims = mlp_dims; + runtime_mlp_dims.grid_x = grid_x; + runtime_mlp_dims.n_gemm = n_gemm; + + moe_fused_bwd( + smem, fused_iter, + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_u, tma_store_v, + tma_load_dy, tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, + tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, + tma_load_dyt3_remote, tma_load_xt4_remote, + p.comm_bwd.x_src_ready, p.comm_bwd.x_src_consumed, + p.comm_bwd.dy_src_ready, p.comm_bwd.dy_src_consumed, + p.comm_bwd.dst_ready, p.comm_bwd.dst_consumed, + p.comm_bwd.tile_expert_ids_x, + p.comm_bwd.tile_expert_ids_dy, + p.comm_bwd.tile_valid_rows_x, + runtime_mlp_dims, mlp_bufs, + col, grid_x, split, runtime_nsplit, gemm_active); + } + + __threadfence(); + __syncthreads(); + p.comm.pe_sync.barrier_all(pe_barrier); + } + + // ── Phase 5: Reverse-dispatch — dx_sorted → dX_out (gather-sum) ── + dispatch_tokens_bwd( + p.dx_sorted, + p.sort.token_expert_slots, + p.dX_out, + p.hidden_dim, + p.num_tokens, + p.top_k); + +} + +// ============================================================================ +// GET-via-TMA descriptor builder for the bwd comm path (raw CUtensorMap). +// ============================================================================ +// +// Box = [kGetBoxRowsBwd, kGetKChunk] over a row-major [n_rows, hidden_dim] +// bf16 tensor (SWIZZLE_NONE keeps global bytes row-major so the MLP's +// swizzled reader sees correct data). Mirrors moe.cu's make_get_tma_map but +// with the bwd's small box height. Returns false on encode failure → caller +// disables the bwd TMA-get path (falls back to getmem/int4). +static inline bool make_get_tma_map_bwd(CUtensorMap& m, void* base, + int hidden_dim, int n_rows, int box_rows) { + uint64_t gdim[2] = { (uint64_t)hidden_dim, (uint64_t)n_rows }; + uint64_t gstride[1] = { (uint64_t)hidden_dim * sizeof(bfloat16_t) }; + uint32_t bdim[2] = { (uint32_t)liger::kGetKChunk, (uint32_t)box_rows }; + uint32_t estride[2] = { 1, 1 }; + CUresult r = cuTensorMapEncodeTiled( + &m, CU_TENSOR_MAP_DATA_TYPE_UINT16, 2, base, gdim, gstride, bdim, estride, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + return r == CUDA_SUCCESS; +} + +// ============================================================================ +// Host wrapper +// ============================================================================ + +// MoeBwdArgs (the POD launch bundle) is defined in moe_launch.h, shared with the +// tuner (mirrors MoeFwdArgs). dB/dC/dA/dW are gradient accumulation targets — +// the launcher zeroes them before the kernel. + +template > +void +moe_bwd_fwd_bf16(const MoeBwdArgs& a, int static_nsplit) { + + const int num_experts = a.num_experts; + const int top_k = a.top_k; + const nvshmem_team_t team = a.team; + cudaStream_t stream = a.stream; + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + if (cudaError_t e = cudaStreamIsCapturing(stream, &capture_status); + e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaStreamIsCapturing failed: ", + cudaGetErrorString(e)); + if (capture_status == cudaStreamCaptureStatusNone) { + if (cudaError_t e = cudaStreamSynchronize(stream); e != cudaSuccess) + LIGER_FAIL_CUDA( + "moe_bwd: synchronizing the forward-to-backward NVSHMEM " + "handoff failed: ", cudaGetErrorString(e)); + } + + using Element = typename Config::Element; + using Traits1 = typename Config::Traits1; + using Traits2T = typename Config::Traits2T; + using Traits3 = typename Config::Traits3; + using Traits4 = typename Config::Traits4; + using Traits5 = typename Config::Traits5; + using TmaLoadAtom = TmaLoadAtomForCompute; + using TmaStoreAtom = TmaStoreAtomForCompute; + using TmaReduceAddAtom = TmaReduceAddAtomForCompute; + constexpr int kTileM = Config::kTileM; + constexpr int kNumThreads = Config::kNumThreads; + + // Iteration counter lives on the device; CTA 0 thread 0 atomicAdds it at + // the end of the kernel via PeSync::advance_iteration_counter(). + int num_tokens = a.num_tokens; + int hidden_dim = a.hidden_dim; + int experts_per_pe = a.experts_per_pe; + int intermediate_dim = a.intermediate_dim; + + // max_total_slots uses kTileM because forward dispatch always pads each + // expert range to its own TileM. moe_bwd_fwd_bf16_auto enforces + // BWD_TileM == FWD_TileM, so kTileM is the right alignment here. + int max_total_slots = num_tokens * top_k + num_experts * kTileM; + int num_m_tiles_in = (num_tokens + kTileM - 1) / kTileM; + + int num_pes = nvshmem_team_n_pes(team); + int my_pe = nvshmem_team_my_pe(team); + int local_gpus = 0; + cudaGetDeviceCount(&local_gpus); + int gpus_per_node = + (local_gpus > 0 && num_pes % local_gpus == 0) ? local_gpus : num_pes; + + int num_sms; + cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, a.device); + // Flat 1-D launch: use the largest NC-complete SM subset. Phase 1 + comm partition over the + // GEMM-active subset n_gemm = floor_NS (largest multiple of static_nsplit ≤ + // num_blocks); Phase 2 (mlp3/mlp4) grid-strides over all num_blocks CTAs + // (no divisibility needed — cell stride = num_blocks); the gap CTAs + // [n_gemm, num_blocks) fall through Phase 1/comm but still run Phase 2 + // and every barrier. NC-complete launch geometry is required because comm + // ring lengths, signal arrays, and TMA descriptors all use gridDim.x / NC. + // ClusterM divides NC for every supported configuration, so this also + // guarantees complete SM100 CTA pairs. + static_assert(Config::kNC % Config::kClusterM == 0); + int num_blocks = num_sms - num_sms % Config::kNC; + int n_gemm = num_blocks - num_blocks % static_nsplit; + int grid_x = n_gemm / static_nsplit; + int max_runtime_grid_x = (num_blocks + 1) / 2; // min runtime NS candidate = 2 + dim3 grid(num_blocks); + LIGER_CHECK(top_k <= kCombineMaxTopK, + "top_k (", top_k, ") exceeds kCombineMaxTopK (", kCombineMaxTopK, + "). Bump the constant in combine.cuh."); + + // stream comes from a.stream (the binding passes the current torch stream). + auto& pool = global_buffer_pool(); + + auto buf = allocate_moe_bwd_buffers(pool, stream, max_runtime_grid_x, num_blocks, + hidden_dim, intermediate_dim, num_tokens, top_k, + experts_per_pe, max_total_slots, num_m_tiles_in); + + // Caller-owned gradient outputs (binding-allocated). dB/dC/dA/dW are + // accumulation targets — zero them on the stream (upstream used torch::zeros); + // dX is fully overwritten so it needs no pre-zero. + constexpr std::size_t kElemBytes = 2; // sizeof(bf16) + cudaMemsetAsync(a.dB, 0, + static_cast(experts_per_pe) * intermediate_dim * hidden_dim * kElemBytes, stream); + cudaMemsetAsync(a.dC, 0, + static_cast(experts_per_pe) * intermediate_dim * hidden_dim * kElemBytes, stream); + cudaMemsetAsync(a.dA, 0, + static_cast(experts_per_pe) * hidden_dim * intermediate_dim * kElemBytes, stream); + cudaMemsetAsync(a.dW, 0, + static_cast(num_tokens) * top_k * kElemBytes, stream); + + // ── Build MoeBwdParams ────────────────────────────────────────── + MoeBwdParams p; + p.expert_indices = a.expert_indices; + p.expert_weights = reinterpret_cast(a.expert_weights); + // The BWD kernel never reads the original X token layout — token data + // is read only from x_sorted (FWD's expert-sorted view of X). Shape + // and device for dX_out come from dY. + p.dY = reinterpret_cast(a.dY); + p.output_buffer = reinterpret_cast(a.Y_fwd); + + p.dX_out = reinterpret_cast(a.dX); + p.dB_out = reinterpret_cast(a.dB); + p.dC_out = reinterpret_cast(a.dC); + p.dA_out = reinterpret_cast(a.dA); + p.dW_out = reinterpret_cast(a.dW); + + // Caller-provided forward intermediates (x_sorted, sort outputs). + // These come from moe_fused_fwd_bf16 and replace + // what the bwd allocator + internal sort/dispatch used to populate. + p.x_sorted = reinterpret_cast(a.x_sorted); + p.dy_sorted = buf.dy_sorted; + p.dx_sorted = buf.dx_sorted; + + // Sort outputs come from forward; we don't run sort_tokens in bwd. + // The fields we leave from buf.* (tile_expert_counts, sorted_token_ids, + // cta_done, cta_sums) are only consumed by sort_tokens — now skipped. + // expert_offsets is now the broadcasted [num_pes, num_experts + 1] + // view from FWD. The local-PE single-row view is offset by + // my_pe * (num_experts + 1) — needed by sort lookup + // (total_real_rows = local_offsets[num_experts]) and by comm code + // that walks this PE's expert range. + int* local_expert_offsets = + a.expert_offsets + my_pe * (num_experts + 1); + + p.sort.tile_expert_counts = buf.tile_expert_counts; + p.sort.sorted_token_ids = buf.sorted_token_ids; + p.sort.token_expert_slots = a.token_expert_slots; + p.sort.expert_offsets = local_expert_offsets; + p.sort.expert_counts = nullptr; + p.sort.tile_expert_ids = a.tile_expert_ids; + p.sort.cta_done = buf.cta_done; + p.sort.cta_sums = buf.cta_sums; + + // Forward CommBuffers (used for pe_sync only — bwd no longer + // re-broadcasts expert offsets; it reads what fwd wrote). + p.comm.src_staging = nullptr; + p.comm.dst_staging = nullptr; + p.comm.src_ready = nullptr; + p.comm.src_consumed = nullptr; + p.comm.dst_ready = nullptr; + p.comm.dst_consumed = nullptr; + p.comm.expert_offsets = local_expert_offsets; + p.comm.expert_counts = nullptr; + p.comm.sorted_token_ids = buf.sorted_token_ids; + p.comm.local_tokens = reinterpret_cast(a.x_sorted); + p.comm.local_output = buf.dx_sorted; + // Forward CommBuffers field — preserved from before the bwd-side split. + // The X-pipe array is the canonical "tile_expert_ids" (the dY-pipe array + // is a bwd-only mirror gated by release_dy). + p.comm.tile_expert_ids = buf.tile_expert_ids_x; + p.comm.all_expert_offsets = a.expert_offsets; + p.comm.all_expert_counts = a.expert_counts; + p.comm.pe_sync.cta_counter = buf.cta_counter; + p.comm.pe_sync.num_pes = num_pes; + p.comm.pe_sync.my_pe = my_pe; + p.comm.pe_sync.team = team; + p.comm.hidden_dim = hidden_dim; + p.comm.num_experts = num_experts; + p.comm.experts_per_pe = experts_per_pe; + p.comm.all_expert_offsets_stride = num_experts + 1; + + // Backward CommBuffers (used for X+dY gets and dX puts). + p.comm_bwd.x_staging = buf.x_staging; + p.comm_bwd.dy_staging = buf.dy_staging; + p.comm_bwd.dx_staging = buf.dx_staging; + p.comm_bwd.x_src_ready = buf.x_src_ready; + p.comm_bwd.x_src_consumed = buf.x_src_consumed; + p.comm_bwd.dy_src_ready = buf.dy_src_ready; + p.comm_bwd.dy_src_consumed = buf.dy_src_consumed; + p.comm_bwd.dst_ready = buf.dst_ready; + p.comm_bwd.dst_consumed = buf.dst_consumed; + p.comm_bwd.tile_expert_ids_x = buf.tile_expert_ids_x; + p.comm_bwd.tile_expert_ids_dy = buf.tile_expert_ids_dy; + p.comm_bwd.tile_valid_rows_x = buf.tile_valid_rows_x; + // Cross-PE pull sources. Both x_sorted and dy_sorted are in the + // same global-expert-sorted layout on every PE: + // • x_sorted: forward's sort_tokens → dispatch wrote it. + // • dy_sorted: combine_tokens_bwd writes [slot, d] using the same + // token_expert_slots → matches forward's per-PE expert sort. + // all_expert_offsets is the broadcasted offset tensor produced by + // fwd (broadcast_expert_offsets) and passed through here verbatim — + // it tells remote PEs the [start, end) slice in our buffers that + // holds data for their experts. + p.comm_bwd.remote_x = p.x_sorted; + p.comm_bwd.remote_dy = buf.dy_sorted; + p.comm_bwd.remote_dx = buf.dx_sorted; + p.comm_bwd.all_expert_offsets = a.expert_offsets; + p.comm_bwd.all_expert_counts = a.expert_counts; + p.comm_bwd.all_expert_offsets_stride = num_experts + 1; + p.comm_bwd.my_pe = my_pe; + p.comm_bwd.num_pes = num_pes; + p.comm_bwd.team = team; + p.comm_bwd.hidden_dim = hidden_dim; + p.comm_bwd.num_experts = num_experts; + p.comm_bwd.experts_per_pe = experts_per_pe; + // GEMM/comm-active count: comm derives MC = n_gemm/NC and the MLP column + // count grid_x = n_gemm/NSplit from this (the flat launch over-covers it). + p.comm_bwd.n_gemm = n_gemm; + + p.num_tokens = num_tokens; + p.hidden_dim = hidden_dim; + p.intermediate_dim = intermediate_dim; + p.top_k = top_k; + p.total_slots = max_total_slots; + p.num_m_tiles = num_m_tiles_in; + p.experts_per_pe = experts_per_pe; + + // ── MlpBwdDims (passed to inner mlp_fused_bwd) ────────────────── + int total_n_rows_1 = experts_per_pe * intermediate_dim; + int total_k_cols_2 = experts_per_pe * hidden_dim; + + MlpBwdDims mlp_dims; + mlp_dims.hidden_dim = hidden_dim; + mlp_dims.intermediate_dim = intermediate_dim; + mlp_dims.total_n_rows_1 = total_n_rows_1; + mlp_dims.num_n_tiles_1 = (Config::kCompute == 90) + ? (intermediate_dim + Traits1::TileN - 1) / Traits1::TileN + : intermediate_dim / Traits1::TileN; + mlp_dims.num_k_tiles_1 = (Config::kCompute == 90) + ? (hidden_dim + Traits1::TileK - 1) / Traits1::TileK + : hidden_dim / Traits1::TileK; + mlp_dims.total_k_cols_2t = total_k_cols_2; + mlp_dims.num_n_tiles_2t = (Config::kCompute == 90) + ? (intermediate_dim + Traits2T::TileN - 1) / Traits2T::TileN + : intermediate_dim / Traits2T::TileN; + mlp_dims.num_k_tiles_2t = (Config::kCompute == 90) + ? (hidden_dim + Traits2T::TileK - 1) / Traits2T::TileK + : hidden_dim / Traits2T::TileK; + mlp_dims.total_k_cols_5 = total_n_rows_1; + mlp_dims.num_n_tiles_5 = (Config::kCompute == 90) + ? (hidden_dim + Traits5::TileN - 1) / Traits5::TileN + : hidden_dim / Traits5::TileN; + mlp_dims.num_k_tiles_5 = (Config::kCompute == 90) + ? (intermediate_dim + Traits5::TileK - 1) / Traits5::TileK + : intermediate_dim / Traits5::TileK; + mlp_dims.num_m_tiles_3 = (Config::kCompute == 90) + ? (hidden_dim + Traits3::TileM - 1) / Traits3::TileM + : hidden_dim / Traits3::TileM; + mlp_dims.num_n_tiles_3 = (Config::kCompute == 90) + ? (intermediate_dim + Traits3::TileN - 1) / Traits3::TileN + : intermediate_dim / Traits3::TileN; + mlp_dims.total_n_rows_3 = total_k_cols_2; + mlp_dims.num_m_tiles_4 = (Config::kCompute == 90) + ? (intermediate_dim + Traits4::TileM - 1) / Traits4::TileM + : intermediate_dim / Traits4::TileM; + mlp_dims.num_n_tiles_4 = (Config::kCompute == 90) + ? (hidden_dim + Traits4::TileN - 1) / Traits4::TileN + : hidden_dim / Traits4::TileN; + mlp_dims.total_m_rows_4 = total_n_rows_1; + mlp_dims.phase_counter = buf.phase_counter; + // Flat-grid geometry: Phase 1 columns = grid_x, GEMM-active count = n_gemm. + mlp_dims.grid_x = grid_x; + mlp_dims.n_gemm = n_gemm; + mlp_dims.num_tokens = max_total_slots; + mlp_dims.num_m_tiles = max_total_slots / kTileM; + // Local pass: ring/buffer size and the K-block iteration bound coincide, + // so total_m_tiles_in_pass equals num_m_tiles. The remote pass overrides + // this in moe_bwd.cuh to comm.global_total. + mlp_dims.total_m_tiles_in_pass = mlp_dims.num_m_tiles; + // Phase 1 processes grid_x m-tiles per batch (one per column). + mlp_dims.num_batches = (mlp_dims.num_m_tiles + grid_x - 1) / grid_x; + // Sort emits GLOBAL expert IDs; mlp_fused_bwd computes + // `local_expert = tile.expert - local_expert_start` to index this PE's + // local all_B/all_C/all_A. Forward kernel does the same (moe.cu:469). + mlp_dims.local_expert_start = my_pe * experts_per_pe; + mlp_dims.num_pes = num_pes; + mlp_dims.experts_per_pe = experts_per_pe; + // Local pass: expert_for_k_block is built per-K-block by + // build_expert_for_k_block (one entry per Traits{3,4}::TileK), so the + // K-block index can be used directly as the array index. The remote pass + // overrides these via EfkbStride{3,4} template params in moe_fused_bwd + // because its array is per-comm-tile. + + // Host-computable constants (grid-constant — saves device-side arithmetic). + // phase2_num_blocks is no longer used for the Phase 2 cell stride (that now + // reads gridDim.x·gridDim.y = num_blocks directly), kept for compatibility. + mlp_dims.phase2_num_blocks = num_blocks; + // Remote pass: TMA descriptors cover the full flat staging list of + // L = MC · kCommNumStages slots, each TileM rows tall. + // MC = launched NC-complete comm columns (matches nvshmem_comm_main_bwd). + int mc_remote = num_blocks / Config::kNC; + mlp_dims.remote_num_m_tiles = mc_remote * Config::kCommNumStages; + mlp_dims.remote_num_tokens = mlp_dims.remote_num_m_tiles * kTileM; + + // Phase-1 sub-batch grouping (mlp_bwd.cuh, SubBatch): the remote pass groups + // up to SubBatch Phase-1 sub-batches before one Phase-2 over their combined + // K-window. The kernel ADAPTS the group size to the ring depth (a group must + // fit in the ring so its grid_x·gcount live slots map to distinct ring slots), + // degrading to fewer sub-batches — down to 1 = the original per-batch loop — + // on shallow rings. So no host-side ring constraint is needed here. + + MlpBwdBufs mlp_bufs; + mlp_bufs.du_buf = buf.du_buf; + mlp_bufs.dv_buf = buf.dv_buf; + mlp_bufs.dz_buf = buf.dz_buf; + // Local pass: no comm staging, no race — both Phase 2 sub-passes alias + // the same post-sort expert_for_k_block. The remote sub-call in + // moe_fused_bwd overrides these to the per-pipe staging arrays. + mlp_bufs.expert_for_k_block_mlp4 = buf.expert_for_k_block; + mlp_bufs.expert_for_k_block_mlp3 = buf.expert_for_k_block; + mlp_bufs.valid_rows_mlp4 = nullptr; + mlp_bufs.expert_k_starts = buf.expert_k_starts; + mlp_bufs.expert_k_ends = buf.expert_k_ends; + mlp_bufs.gmem_da = reinterpret_cast(a.dA); + mlp_bufs.gmem_db = reinterpret_cast(a.dB); + mlp_bufs.gmem_dc = reinterpret_cast(a.dC); + mlp_bufs.barrier_counter = buf.barrier_counter; + + // ── TMA descriptors ───────────────────────────────────────────── + // X comes from caller's x_sorted (forward's pool, post-dispatch). + // dy_sorted is bwd-allocated (combine_tokens_bwd will populate it). + auto* pX = reinterpret_cast(a.x_sorted); // sorted X + auto* pDY = reinterpret_cast(buf.dy_sorted); // expanded dY + auto* pB = reinterpret_cast(a.all_B); + auto* pC = reinterpret_cast(a.all_C); + auto* pA = reinterpret_cast(a.all_A); + auto* pDX = reinterpret_cast(buf.dx_sorted); + + // All Hopper expert weights and weight-gradient targets use rank-3 tensor + // maps. TMA handles partial edge boxes as zero-filled loads / dropped + // stores, while the explicit expert coordinate prevents cross-expert tails. + // (See mlp_bwd.cu for the full pattern; here we mirror it but the + // "X" tensor is the symmetric x_sorted and "dY" is the symmetric + // dy_sorted, both of shape [max_total_slots, hidden_dim].) + auto tma_load_x = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pX), make_shape(max_total_slots, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits1::SmemLayoutX_1{}); + + auto tma_load_b_fwd = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pB), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + static_cast(intermediate_dim) * hidden_dim)), + typename Traits1::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pB), + make_shape(total_n_rows_1, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits1::SmemLayoutW_1{}); + } + }(); + auto tma_load_c_fwd = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pC), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + static_cast(intermediate_dim) * hidden_dim)), + typename Traits1::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pC), + make_shape(total_n_rows_1, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits1::SmemLayoutW_1{}); + } + }(); + + auto tma_store_z = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.z_buf), make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits1::SmemLayoutActStoreSlot{}); + auto tma_store_u = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.du_buf), make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits1::SmemLayoutActStoreSlot{}); + auto tma_store_v = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.dv_buf), make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits1::SmemLayoutActStoreSlot{}); + + auto tma_load_dy = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pDY), make_shape(max_total_slots, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits2T::SmemLayoutZ_1{}); + auto tma_load_a_col = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pA), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + Int<1>{}, + static_cast(intermediate_dim), + static_cast(hidden_dim) * intermediate_dim)), + typename Traits2T::SmemLayoutW_1{}); + // One TMA box per sub-tile; consumer issues NSub stores per outer N-tile. + auto tma_store_dz = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.dz_buf), + make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits2T::SmemLayoutStore_1{}); + + auto tma_load_du = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.du_buf)), + make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits5::SmemLayoutZ_1{}); + auto tma_load_dv = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.dv_buf)), + make_shape(buf.scratch_rows, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits5::SmemLayoutZ_1{}); + auto tma_load_b_col = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pB), + make_shape(static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(experts_per_pe)), + make_stride( + Int<1>{}, + static_cast(hidden_dim), + static_cast(intermediate_dim) * hidden_dim)), + typename Traits5::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pB), + make_shape(hidden_dim, total_n_rows_1), + make_stride(Int<1>{}, hidden_dim)), + typename Traits5::SmemLayoutW_1{}); + } + }(); + auto tma_load_c_col = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pC), + make_shape(static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(experts_per_pe)), + make_stride( + Int<1>{}, + static_cast(hidden_dim), + static_cast(intermediate_dim) * hidden_dim)), + typename Traits5::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(pC), + make_shape(hidden_dim, total_n_rows_1), + make_stride(Int<1>{}, hidden_dim)), + typename Traits5::SmemLayoutW_1{}); + } + }(); + // dX TMA: split-N store (TileN_half boxes per Mlp5::SmemLayoutStore_1). + auto tma_store_dx = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(pDX), + make_shape(max_total_slots, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits5::SmemLayoutStore_1{}); + + // Phase 2 TMAs. mlp3's dY^T / Z operands and mlp4's X / dU^T / dV^T + // operands need pair-aware SM100_TMA_2SM_LOAD descriptors (built via + // make_tma_copy_A_sm100 / make_tma_copy_B_sm100, keyed off the joined + // TileShape + TiledMma2Sm) when Config::kUsesTwoSm (all SM100 configs); + // the dA/dB/dC TMA_REDUCE_ADD outputs are unchanged either way (each + // peer CTA still reduce-adds its own CtaTileM rows independently). + auto tma_load_dyt3 = [&] { + auto tensor = make_tensor(make_gmem_ptr(pDY), + make_shape(hidden_dim, max_total_slots), + make_stride(Int<1>{}, hidden_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits3::SmemLayoutDYT_1{}, + typename Traits3::TileShape{}, + typename Traits3::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits3::SmemLayoutDYT_1{}); + } + }(); + auto tma_load_zt3 = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.z_buf)), + make_shape(intermediate_dim, buf.scratch_rows), + make_stride(Int<1>{}, intermediate_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_B_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits3::SmemLayoutZ_1{}, + typename Traits3::TileShape{}, + typename Traits3::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits3::SmemLayoutZ_1{}); + } + }(); + auto tma_reduce_da = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dA)), + make_shape(static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(intermediate_dim), + Int<1>{}, + static_cast(hidden_dim) * intermediate_dim)), + typename Traits3::SmemLayoutStore{}); + } else { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dA)), + make_shape(total_k_cols_2, intermediate_dim), + make_stride(intermediate_dim, Int<1>{})), + typename Traits3::SmemLayoutStore{}); + } + }(); + + auto tma_load_xt4 = [&] { + auto tensor = make_tensor(make_gmem_ptr(pX), + make_shape(hidden_dim, max_total_slots), + make_stride(Int<1>{}, hidden_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_B_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits4::SmemLayoutX_1{}, + typename Traits4::TileShape{}, + typename Traits4::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits4::SmemLayoutX_1{}); + } + }(); + auto tma_load_dut4 = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.du_buf)), + make_shape(intermediate_dim, buf.scratch_rows), + make_stride(Int<1>{}, intermediate_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits4::SmemLayoutA_1{}, + typename Traits4::TileShape{}, + typename Traits4::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits4::SmemLayoutA_1{}); + } + }(); + auto tma_load_dvt4 = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.dv_buf)), + make_shape(intermediate_dim, buf.scratch_rows), + make_stride(Int<1>{}, intermediate_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits4::SmemLayoutA_1{}, + typename Traits4::TileShape{}, + typename Traits4::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits4::SmemLayoutA_1{}); + } + }(); + auto tma_reduce_db = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dB)), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + static_cast(intermediate_dim) * hidden_dim)), + typename Traits4::SmemLayoutStore{}); + } else { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dB)), + make_shape(total_n_rows_1, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits4::SmemLayoutStore{}); + } + }(); + auto tma_reduce_dc = [&]() { + if constexpr (Config::kCompute == 90) { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dC)), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + static_cast(intermediate_dim) * hidden_dim)), + typename Traits4::SmemLayoutStore{}); + } else { + return make_tma_copy(TmaReduceAddAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(a.dC)), + make_shape(total_n_rows_1, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits4::SmemLayoutStore{}); + } + }(); + + // Remote (staging) TMAs — buffers cover L = MC · CommNumStages slots, + // each TileM rows tall. MC = (grid.x · NSplit) / NC. + int remote_total_rows = mc_remote * Config::kCommNumStages * kTileM; + auto tma_load_x_remote = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.x_staging)), + make_shape(remote_total_rows, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits1::SmemLayoutX_1{}); + auto tma_load_dy_remote = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.dy_staging)), + make_shape(remote_total_rows, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits2T::SmemLayoutZ_1{}); + // Remote dX TMA: split-N store (matches local tma_store_dx). + auto tma_store_dx_remote = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.dx_staging), + make_shape(remote_total_rows, hidden_dim), + make_stride(hidden_dim, Int<1>{})), + typename Traits5::SmemLayoutStore_1{}); + + // Remote Phase 3/4 TMAs: transposed staging sources. The remote phase + // of moe_fused_bwd previously reused tma_load_dyt3/tma_load_xt4 (which + // point at LOCAL dy_sorted/x_sorted), so dA/dB/dC for cross-PE-routed + // tokens used the wrong rows. These descriptors pull from the same + // staging buffers that Phase 1/2 of the remote pass already use. + auto tma_load_dyt3_remote = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.dy_staging)), + make_shape(hidden_dim, remote_total_rows), + make_stride(Int<1>{}, hidden_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_A_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits3::SmemLayoutDYT_1{}, + typename Traits3::TileShape{}, + typename Traits3::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits3::SmemLayoutDYT_1{}); + } + }(); + auto tma_load_xt4_remote = [&] { + auto tensor = make_tensor(make_gmem_ptr(reinterpret_cast(buf.x_staging)), + make_shape(hidden_dim, remote_total_rows), + make_stride(Int<1>{}, hidden_dim)); + if constexpr (Config::kUsesTwoSm) { + return make_tma_copy_B_sm100(SM100_TMA_2SM_LOAD{}, tensor, + typename Traits4::SmemLayoutX_1{}, + typename Traits4::TileShape{}, + typename Traits4::TiledMma2Sm{}); + } else { + return make_tma_copy(TmaLoadAtom{}, tensor, + typename Traits4::SmemLayoutX_1{}); + } + }(); + + // ── GET-via-TMA descriptors for the bwd comm path (#111 port) ──── + // Per-peer source maps over nvshmem_ptr(x_sorted/dy_sorted, peer) + one + // local dst map each over x_staging / dy_staging. Enabled only on NVLink + // P2P + divisible hidden_dim + bounce smem fit (checked after smem sizing). + GetTmaDescsBwd get_descs_bwd{}; + get_descs_bwd.enabled = 0; + get_descs_bwd.gpus_per_node = gpus_per_node; + get_descs_bwd.my_pe = my_pe; + constexpr int kGetBoxRowsBwd = Config::kCompute == 100 + ? liger::kGetBoxRowsBwdSm100 + : liger::kGetBoxRowsBwdSm90; + size_t get_bounce_bytes_bwd = + (size_t)liger::GetBounceBwd::kTotalBytes; + { + // The single get warp's row band is TileM/NC; the small TMA box height + // must divide it (else the last band would over-read). Holds for NC≤8 + // at TileM=128; larger NC (band < 16) disables the path. + constexpr int kRowBandBwd = Config::kTileM / Config::kNC; + bool ok = (hidden_dim % liger::kGetKChunk == 0) + && (gpus_per_node <= liger::kGetMaxPes) + && (kRowBandBwd % kGetBoxRowsBwd == 0); + int host_start = (my_pe / gpus_per_node) * gpus_per_node; + for (int local_peer = 0; ok && local_peer < gpus_per_node; ++local_peer) { + int tp = host_start + local_peer; + if (tp >= num_pes) { ok = false; break; } + int gpe = nvshmem_team_translate_pe(team, tp, NVSHMEM_TEAM_WORLD); + void* px_peer = nvshmem_ptr(p.comm_bwd.remote_x, gpe); // peer x_sorted + void* pdy_peer = nvshmem_ptr(p.comm_bwd.remote_dy, gpe); // peer dy_sorted + if (px_peer == nullptr || pdy_peer == nullptr) { ok = false; break; } + ok = make_get_tma_map_bwd(get_descs_bwd.src_x_desc[local_peer], px_peer, + hidden_dim, max_total_slots, kGetBoxRowsBwd); + if (ok) ok = make_get_tma_map_bwd(get_descs_bwd.src_dy_desc[local_peer], pdy_peer, + hidden_dim, max_total_slots, kGetBoxRowsBwd); + } + if (ok) ok = make_get_tma_map_bwd(get_descs_bwd.dst_x_staging_desc, + buf.x_staging, hidden_dim, remote_total_rows, kGetBoxRowsBwd); + if (ok) ok = make_get_tma_map_bwd(get_descs_bwd.dst_dy_staging_desc, + buf.dy_staging, hidden_dim, remote_total_rows, kGetBoxRowsBwd); + get_descs_bwd.enabled = ok ? 1 : 0; + if (const char* e = std::getenv("LIGER_GET_TMA")) + if (e[0] == '0') get_descs_bwd.enabled = 0; + } + + // ── Smem sizing ───────────────────────────────────────────────── + int sort_num_warps = kNumThreads / 32; + size_t smem_sort = sort_num_warps * num_experts * sizeof(int); + + size_t smem_mlp = sizeof(MoeBwdSmem); + smem_mlp = (smem_mlp + sizeof(int) - 1) & ~(sizeof(int) - 1); + smem_mlp += num_pes * (experts_per_pe + 1) * sizeof(int); // remote_offsets + smem_mlp += num_pes * experts_per_pe * sizeof(int); // remote_counts + + // Query the device's actual opt-in dynamic-smem cap rather than trust + // Config::kSmemBudget's compile-time 228 KiB constant: the 2SM paired-CTA + // path (Mlp3Traits2Sm / Mlp4Traits2Sm) pushes some candidates close to + // that ceiling, and the real per-SKU limit is what + // cudaFuncAttributeMaxDynamicSharedMemorySize is actually bounded by. + int device = 0; + int dynamic_smem_limit = 0; + if (cudaError_t e = cudaGetDevice(&device); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaGetDevice failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaDeviceGetAttribute(&dynamic_smem_limit, + cudaDevAttrMaxSharedMemoryPerBlockOptin, device); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: querying the opt-in shared-memory limit failed: ", + cudaGetErrorString(e)); + + // GET-via-TMA bounce: append a 128-aligned region after remote_offsets + // (the kernel computes the same offset). The bwd union sits near the + // device's dynamic-smem cap, so disable the TMA-get path if the bounce + // won't fit against the QUERIED limit (not the static 228 KiB constant) + // — that config just keeps the original 2-get/1-put getmem layout (no + // regression). + if (get_descs_bwd.enabled) { + size_t mlp_with_bounce = + ((smem_mlp + 127) & ~((size_t)127)) + get_bounce_bytes_bwd; + if (mlp_with_bounce <= static_cast(dynamic_smem_limit)) smem_mlp = mlp_with_bounce; + else get_descs_bwd.enabled = 0; + } + // Unified BWD comm layout always uses one get warp and one put warp. + // get_descs_bwd.enabled only controls whether same-host tiles use TMA; + // remote/unsupported tiles fall back inside do_get_bwd_tma. + mlp_dims.tma_get_enabled = 1; + + size_t smem_size = std::max(smem_sort, smem_mlp); + + // ── Launch ────────────────────────────────────────────────────── + auto kernel_fn = moe_bwd_kernel; + + if (cudaError_t e = cudaFuncSetAttribute(kernel_fn, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaFuncSetAttribute(MaxDynamicSharedMemorySize=", + smem_size, ", device opt-in limit=", dynamic_smem_limit, ") failed: ", + cudaGetErrorString(e)); + // The SM100 2SM paired-CTA cluster path launches with an explicit + // clusterDim; opt in to non-portable cluster sizes (the CUDA default + // portable max is 8, so ClusterM=2 is already portable in practice, + // but every 2SM launch enables this defensively, mirroring CUTLASS's + // own cluster-launch helpers). + if constexpr (Config::kUsesTwoSm) { + if (cudaError_t e = cudaFuncSetAttribute(kernel_fn, + cudaFuncAttributeNonPortableClusterSizeAllowed, 1); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: enabling non-portable cluster size (ClusterM=", + Config::kClusterM, ") failed: ", cudaGetErrorString(e)); + } + + // Fresh per-call counter resets. + cudaMemsetAsync(buf.cta_counter, 0, sizeof(int), stream); + cudaMemsetAsync(buf.barrier_counter, 0, sizeof(int), stream); + cudaMemsetAsync(buf.phase_counter, 0, max_runtime_grid_x * sizeof(int), stream); + // L = MC · CommNumStages, MC = launched NC-complete comm columns. + int mc_local = ((int)grid.x / Config::kNC); + int comm_l_local = mc_local * Config::kCommNumStages; + size_t per_call_tile_sig_bytes = + static_cast(comm_l_local) * sizeof(int); + cudaMemsetAsync(buf.x_src_ready, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.x_src_consumed, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.dy_src_ready, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.dy_src_consumed,0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.dst_ready, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.dst_consumed, 0, per_call_tile_sig_bytes, stream); + cudaMemsetAsync(buf.cta_done, 0, num_blocks * sizeof(int), stream); + cudaMemsetAsync(buf.sorted_token_ids, 0xff, max_total_slots * sizeof(int), stream); + // dy_sorted MUST be zeroed: combine_tokens_bwd writes routed slots + // only (no padding fill — unlike forward's dispatch_tokens). Partial- + // expert tile rows of dy_sorted are read by Phase 1b's TMA load; if + // those rows aren't zero, the resulting dZ/dU/dV poison Phase 2 dA/dB/dC + // for K-blocks that straddle an expert boundary (see the + // build_expert_for_k_block comment for the alignment assumption). + cudaMemsetAsync(buf.dy_sorted, 0, static_cast(max_total_slots) * hidden_dim * sizeof(Element), stream); + size_t intermediate_bytes = + static_cast(buf.scratch_rows) * intermediate_dim * sizeof(Element); + cudaMemsetAsync(buf.z_buf, 0, intermediate_bytes, stream); + cudaMemsetAsync(buf.du_buf, 0, intermediate_bytes, stream); + cudaMemsetAsync(buf.dv_buf, 0, intermediate_bytes, stream); + cudaMemsetAsync(buf.dz_buf, 0, intermediate_bytes, stream); + + // Sentinel-fill BOTH comm-side tile_expert_ids arrays so unfetched + // slots produce expert_id = -1, which the Phase 3/4 bounds check + // rejects. Without this, stale memory could land in [0, experts_per_pe) + // and trigger spurious dB/dC/dA writes for empty staging slots. + size_t comm_eid_bytes = + static_cast(comm_l_local) * sizeof(int); + if (cudaError_t e = cudaMemsetAsync(buf.tile_expert_ids_x, 0xff, comm_eid_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaMemsetAsync(tile_expert_ids_x) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(buf.tile_expert_ids_dy, 0xff, comm_eid_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: cudaMemsetAsync(tile_expert_ids_dy) failed: ", cudaGetErrorString(e)); + cudaMemsetAsync(buf.tile_valid_rows_x, 0, comm_eid_bytes, stream); + if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: CUDA error before kernel launch: ", cudaGetErrorString(e)); + + // Compute=90 keeps the ordinary <<<>>> launch. Compute=100 always uses + // the explicit paired-CTA cluster launch below. + if constexpr (Config::kUsesTwoSm) { + cudaLaunchConfig_t launch_config = {}; + launch_config.gridDim = grid; + launch_config.blockDim = kNumThreads; + launch_config.dynamicSmemBytes = smem_size; + launch_config.stream = stream; + cudaLaunchAttribute cluster_attr = {}; + cluster_attr.id = cudaLaunchAttributeClusterDimension; + cluster_attr.val.clusterDim.x = Config::kClusterM; + cluster_attr.val.clusterDim.y = 1; + cluster_attr.val.clusterDim.z = 1; + launch_config.attrs = &cluster_attr; + launch_config.numAttrs = 1; + cudaLaunchKernelEx(&launch_config, kernel_fn, + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_u, tma_store_v, + tma_load_dy, tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, + tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, + tma_load_dyt3_remote, tma_load_xt4_remote, + mlp_dims, mlp_bufs, p, static_nsplit, get_descs_bwd); + } else { + kernel_fn<<>>( + tma_load_x, tma_load_b_fwd, tma_load_c_fwd, + tma_store_z, tma_store_u, tma_store_v, + tma_load_dy, tma_load_a_col, tma_store_dz, + tma_load_du, tma_load_dv, tma_load_b_col, tma_load_c_col, tma_store_dx, + tma_load_dyt3, tma_load_zt3, tma_reduce_da, + tma_load_xt4, tma_load_dut4, tma_load_dvt4, tma_reduce_db, tma_reduce_dc, + tma_load_x_remote, tma_load_dy_remote, tma_store_dx_remote, + tma_load_dyt3_remote, tma_load_xt4_remote, + mlp_dims, mlp_bufs, p, static_nsplit, get_descs_bwd); + } + if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_bwd: kernel launch failed for grid=", grid.x, + " threads=", kNumThreads, " smem=", smem_size, ": ", + cudaGetErrorString(e)); + + // Post-kernel stream-ordered barrier: mirrors the fwd wrapper. The + // kernel's PeSync wait_mlp guarantees this PE has received peer + // signals, but does not guarantee this PE's outbound signals have + // landed on peers. Under back-to-back captured replays (fwd → bwd → + // fwd), bwd's tail-end signals to peers can race against the next + // fwd's wait_dispatch: a peer's wait_until(>= target) may briefly + // see a stale value if bwd's older signal arrives after a newer + // one. Forcing host-side completion here serialises remote + // signal delivery against the next captured kernel. + if (num_pes > 1) { + nvshmemx_barrier_on_stream(team, stream); + } +} + +// ============================================================================ +// Tuned-config wrapper: 13 free knobs → full MoeBwdConfig +// ============================================================================ +// +// Function template that takes the X-macro's free knobs and derives the +// dependent traits per the symmetric-traits constraints documented in +// moe_fwd_bwd_tune_configs.hpp: +// Phase 1a: TileN1, TileK1, Stages1, EpiChunkN1 +// Phase 1b': TileN = 256 (fixed), TileK = TileK1, Stages = Stages1, EpiChunkN25 +// Phase 1d (mlp5): TileM=128, TileN=256, TileK = TileK1, Stages = Stages1, EpiChunkN25 +// Phase 2 mlp3: TileM3, TileN3, TileK3, Stages3, EpiChunkN34 +// Phase 2 mlp4: TileM = TileN3, TileN = TileM3, TileK3, Stages3, EpiChunkN34 +// (mlp3/mlp4 swap) +// NC = NSplit +// +// The tuner registry and dispatch table are built by emitting one explicit +// instantiation of this template per X-macro entry. + +template < + int NSplit2, + int TileN1, int TileK1, int Stages1, + int TileM3, int TileN3, int TileK3, int Stages3, + int EpiChunkN1, int EpiChunkN25, int EpiChunkN34, + int CommNumStages, + int TileM, // COMM tile (staging/semaphores/NC); must match FWD's TileM. 64 or 128. + int GemmTileM = TileM, // Phase-1 GEMM tile (mlp1a/mlp2t/mlp5) ∈ {64,128}; default == comm. + int Compute = 90> +void +moe_bwd_fwd_bf16_tuned(const MoeBwdArgs& a, int static_nsplit) { + // NC paired with TileM (mirrors moe.cu's moe_fused_fwd_bf16): + // TileM=128 → NC=4 (~2× the comm bandwidth of NC=2; deeper comm pipeline). + // TileM=64 → NC=2 (tiles are half-size, NC=2 matches NC=4 throughput). + // Ticket-based iterator + K=NumStages comm means NSplit and NC have no + // divisibility constraint — only requirement is that NC evenly divides + // gridDim.x · gridDim.y so MC is integer (e.g. NS=6, NC=4: MC = 22·6/4 + // = 33). + // Both architectures use CommTileM=128/NC4, matching forward metadata. + using CommGeometry = MoeCommGeometry; + static constexpr int kPhase1TileM = Compute == 100 ? 128 : GemmTileM; + static constexpr int kSubBatch = 2; + using Config = MoeBwdConfig< + NSplit2, + kPhase1TileM, TileN1, TileK1, Stages1, EpiChunkN1, // Phase 1a (mlp1_act) — GEMM tile + /*TileN2=*/256, TileK1, Stages1, // Phase 1b' (Mlp2TTraits caps TileN ≤ 256) + /*TileM2T=*/kPhase1TileM, // Phase 1b' GEMM tile + TileM3, TileN3, TileK3, Stages3, // Phase 2 mlp3 + TileN3, TileM3, TileK3, Stages3, // Phase 2 mlp4 (mlp3/mlp4 swap; cooperative) + kPhase1TileM, 256, TileK1, Stages1, // Phase 1d mlp5 — GEMM tile + EpiChunkN25, // shared mlp2_t + mlp5 + EpiChunkN34, // shared mlp3 + mlp4 + bfloat16_t, + CommNumStages * CommGeometry::RingStageScale + + CommGeometry::RingStageExtra, + /*NC_=*/CommGeometry::NC, + /*SubBatch_=*/kSubBatch, + /*CommTileM_=*/CommGeometry::TileM, + /*Compute_=*/Compute>; + moe_bwd_fwd_bf16(a, static_nsplit); +} + +// ── Runtime dispatch table for tuned-config lookup ──────────────────── + +using MoeBwdFwdFn = void (*)(const MoeBwdArgs&, int static_nsplit); + +struct DispatchEntryBwd { + int Compute; + int NSplit2; + int TileN1, TileK1, Stages1; + int TileM3, TileN3, TileK3, Stages3; + int EpiChunkN1, EpiChunkN25, EpiChunkN34; + int CommNumStages; + int TileM; // COMM tile (staging/semaphores/NC). + int GemmTileM; // Phase-1 GEMM tile (mlp1a/mlp2t/mlp5); == TileM for coupled configs. + MoeBwdFwdFn fn; +}; + +#define LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) \ + { Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM, \ + &moe_bwd_fwd_bf16_tuned }, +#define LIGER_MOE_BWD_DISPATCH_ENTRY_C(Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(Compute, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, TM) +#define LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM90(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(90, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) +#define LIGER_MOE_BWD_DISPATCH_ENTRY_SM90(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_C(90, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) +#define LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM100(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_G_C(100, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM, GTM) +#define LIGER_MOE_BWD_DISPATCH_ENTRY_SM100(NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) \ + LIGER_MOE_BWD_DISPATCH_ENTRY_C(100, NS2, TN1, TK1, S1, TM3, TN3, TK3, S3, EN1, EN25, EN34, CS, TM) + +static const DispatchEntryBwd kDispatchTableBwd[] = { +#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 90 + LIGER_MOE_BWD_DISPATCH_CONFIGS_SM90(LIGER_MOE_BWD_DISPATCH_ENTRY_SM90, LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM90) +#endif +#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 100 + LIGER_MOE_BWD_DISPATCH_CONFIGS_SM100(LIGER_MOE_BWD_DISPATCH_ENTRY_SM100, LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM100) +#endif +}; +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_SM100 +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM100 +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_SM90 +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_G_SM90 +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_C +#undef LIGER_MOE_BWD_DISPATCH_ENTRY_G_C +static constexpr int kNumDispatchEntriesBwd = + sizeof(kDispatchTableBwd) / sizeof(kDispatchTableBwd[0]); + +} // namespace liger + +// Joint fwd+bwd tuning table (auto-generated by benchmarks/tune_moe_fwd_bwd). +// We only consume the Bwd_* fields here; moe.cu consumes the same table's +// Fwd_* fields. Picking fwd and bwd from a single table guarantees the two +// directions are tuned as a pair — in particular their TileM agrees. +#include "moe_fwd_bwd_tuning_configs_sm90.cuh" + +namespace liger { + +static constexpr int kBwdRuntimeNsplitSeed = 8; + +// TunedConfigBwd is the bwd-only projection of TunedConfigFwdBwd. The rest +// of the dispatch code (validity check, dispatch-entry match) was written +// against this struct and is reused unchanged. +struct TunedConfigBwd { + int Compute; + int TK, TKE, D, I; + int NSplit2; + int TileN1, TileK1, Stages1; + int TileM3, TileN3, TileK3, Stages3; + int EpiChunkN1, EpiChunkN25, EpiChunkN34; + int CommNumStages; + int TileM; // comm tile + int GemmTileM; // gemm tile (auto path: == TileM; FORCE_CONFIG may set 64) +}; + +// Project a TunedConfigFwdBwd row down to its bwd template params. The tuned +// table has no GemmTileM column, so the auto path keeps gemm == comm (TileM). +static TunedConfigBwd project_bwd(const TunedConfigFwdBwd& c, int compute) { + return TunedConfigBwd{ + compute, + c.TK, c.TKE, c.D, c.I, + c.Bwd_NSplit2, + c.Bwd_TileN1, c.Bwd_TileK1, c.Bwd_Stages1, + c.Bwd_TileM3, c.Bwd_TileN3, c.Bwd_TileK3, c.Bwd_Stages3, + c.Bwd_EpiChunkN1, c.Bwd_EpiChunkN25, c.Bwd_EpiChunkN34, + c.Bwd_CommNumStages, + c.Bwd_TileM, c.Bwd_TileM, + }; +} + +// Find the dispatch table entry matching a TunedConfigBwd's knobs. +static const DispatchEntryBwd* find_dispatch_entry_bwd(const TunedConfigBwd& c) { + for (int i = 0; i < kNumDispatchEntriesBwd; ++i) { + const auto& e = kDispatchTableBwd[i]; + if (e.Compute == c.Compute && + e.NSplit2 == c.NSplit2 && + e.TileN1 == c.TileN1 && e.TileK1 == c.TileK1 && e.Stages1 == c.Stages1 && + e.TileM3 == c.TileM3 && e.TileN3 == c.TileN3 && e.TileK3 == c.TileK3 && + e.Stages3 == c.Stages3 && + e.EpiChunkN1 == c.EpiChunkN1 && + e.EpiChunkN25 == c.EpiChunkN25 && e.EpiChunkN34 == c.EpiChunkN34 && + e.CommNumStages == c.CommNumStages && + e.TileM == c.TileM && e.GemmTileM == c.GemmTileM) { + return &e; + } + } + return nullptr; +} + +// Structural divisibility checks on (D, I) for a given tuned config. +// Mirrors the constraint list in moe_fwd_bwd_tune_configs.hpp. +static bool tuned_config_valid_bwd(int D, int I, const TunedConfigBwd& c) { + // The current Blackwell UMMA MLP1/MLP1-act consumers only support the + // 128-row comm tile. Keep mirrored SM90 TileM=64 tuning rows out of SM100 + // auto-dispatch until that TMEM epilogue layout is implemented. + if (c.Compute == 100 && c.TileM != 128) return false; + if (c.Compute == 100) { + if (D % c.TileK1 != 0) return false; // mlp1/mlp2_t/mlp5 K + if (D % (2 * c.TileN1) != 0) return false; // mlp5 N (= TileN2) + if (I % (2 * c.TileN1) != 0) return false; // mlp2_t N + if (I % c.TileK1 != 0) return false; // mlp5 K + if (I % c.TileK3 != 0) return false; // mlp3/mlp4 K + } + if (D % 8 != 0 || I % 8 != 0) return false; // vectorization + if (c.Compute == 100) { + if (D % 256 != 0 || I % 256 != 0) return false; + } else { + if (!((c.TileM3 == 256 && c.TileN3 == 128) || + (c.TileM3 == 128 && c.TileN3 == 256))) return false; + if (D % c.TileM3 != 0 || I % c.TileN3 != 0) return false; + } + return true; +} + +static bool find_tuned_table_bwd(int compute, int n_pes, + const TunedConfigFwdBwd** table, + int* table_n) { + const TunedConfigFwdBwdTable* tables = + (n_pes <= 1) ? kTunedConfigTablesSingle : kTunedConfigTablesMulti; + const int tables_n = + (n_pes <= 1) ? kNumTunedConfigTablesSingle : kNumTunedConfigTablesMulti; + for (int i = 0; i < tables_n; ++i) { + if (tables[i].Compute == compute) { + *table = tables[i].configs; + *table_n = tables[i].count; + return *table_n > 0; + } + } + return false; +} + +// Nearest-neighbor lookup in log-space (TK, TKE, D, I), where TK = T*top_k +// is the total routed-token count (the bwd's M-axis post-dispatch) and +// TKE = TK / E_local (E_local = E / n_pes) is the per-LOCAL-expert avg +// K-range. mlp3/mlp4 walk each PE's local experts, so adding TKE to the +// distance lets the dispatcher distinguish configs whose walk-scheduler +// load differs even when (TK, D, I) is the same. +// Configs that violate structural constraints for the runtime (D, I) are +// skipped — the nearest tuned point is not always valid for the actual shape. +// If fwd_tile_m > 0, configs with TileM != fwd_tile_m are filtered out: +// BWD's sort-output reads (tile_expert_ids granularity, expert_offsets +// alignment, x_sorted slot layout) must match the forward's TileM. +// Out-param API instead of pointer-into-static-array because each candidate +// TunedConfigBwd is projected on the fly from kTunedConfigsFwdBwd (no stable +// address). Returns true when a valid config was found. +static bool find_nearest_tuned_bwd(int compute, int TK, int TKE, int D, int I, + int fwd_tile_m, int n_pes, + TunedConfigBwd* out) { + // World-size class: single-GPU (n_pes <= 1) and multi-GPU tables are tuned + // and stored separately (mirrors find_nearest_tuned on the fwd side). + const TunedConfigFwdBwd* table = nullptr; + int table_n = 0; + if (!find_tuned_table_bwd(compute, n_pes, &table, &table_n)) return false; + bool found = false; + double best_dist = 1e30; + double lTK = std::log((double)TK); + double lTKE = std::log((double)TKE); + double lD = std::log((double)D); + double lI = std::log((double)I); + for (int i = 0; i < table_n; ++i) { + TunedConfigBwd c = project_bwd(table[i], compute); + // Forward and backward select the same architecture-specific comm tile + // (M128 on both architectures), independent of the tuned GEMM tile. + (void)fwd_tile_m; + if (!tuned_config_valid_bwd(D, I, c)) continue; + // Guard against TKE=0 in legacy rows (would yield -inf in log). + double dTK = std::log((double)c.TK) - lTK; + double dTKE = c.TKE > 0 + ? std::log((double)c.TKE) - lTKE + : 0.0; // legacy / placeholder rows contribute 0 to the TKE axis + double dD = std::log((double)c.D) - lD; + double dI = std::log((double)c.I) - lI; + double d = dTK*dTK + dTKE*dTKE + dD*dD + dI*dI; + if (d < best_dist) { best_dist = d; *out = c; found = true; } + } + return found; +} + +// ============================================================================ +// Auto-dispatch wrapper — pure tuned-table lookup +// ============================================================================ +// +// Picks the nearest tuned config (log-space distance over (TK, D, I)) from +// moe_fwd_bwd_tuning_configs.cuh and dispatches to the matching template +// instantiation in kDispatchTableBwd. Both tables are populated by the +// tune_moe_bwd sweep — shapes outside the swept range error out instead of +// falling back to a heuristic. +// +// fwd_tile_m: BWD's TileM must match the forward's TileM. Forward sort +// pads expert ranges to FWD's TileM, populates tile_expert_ids at that +// granularity, and writes x_sorted in FWD-TileM-aligned slots. BWD +// reads all three at the same granularity, so a mismatch corrupts both +// the LocalIter expert lookup and the symmetric-memory bounds. +// Callers must pass the FWD's TileM (= 128 for the default fwd path). +// +// Override path: LIGER_MOE_BWD_FORCE_CONFIG=NS2,TN1,TK1,S1,TM3,TN3,TK3,S3,EN1,EN25,EN34,CS,TM +// Lets a test suite sweep every compiled config without rebuilding. + +// Internal auto-dispatch: nearest tuned bwd config for the runtime shape (or +// LIGER_MOE_BWD_FORCE_CONFIG), then invoke the matching template instantiation. +// All shape info comes from the MoeBwdArgs bundle; the launcher writes the grads. +void moe_bwd_dispatch(const MoeBwdArgs& a, int fwd_tile_m) { + const int T = a.num_tokens; + const int D = a.hidden_dim; + const int I = a.intermediate_dim; + const int top_k = a.top_k; + const int num_experts = a.num_experts; + const int TK = T * top_k; + const int n_pes = nvshmem_team_n_pes(a.team); + const int E_local = (n_pes > 0) ? std::max(1, num_experts / n_pes) : num_experts; + // Keep backward lookup consistent with forward for tiny token batches. + const int TKE = std::max(1, TK / E_local); + const int compute = moe_detect_compute_dispatch_key(a.device, "moe_fused_bwd_bf16_auto"); + if (const char* s = std::getenv("LIGER_MOE_BWD_FORCE_CONFIG")) { + TunedConfigBwd tc{}; + tc.Compute = compute; + tc.GemmTileM = -1; // sentinel → default to comm TileM after parse + int vals[14] = {}; + int n = std::sscanf(s, "%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d", + &vals[0], &vals[1], &vals[2], &vals[3], &vals[4], &vals[5], &vals[6], + &vals[7], &vals[8], &vals[9], &vals[10], &vals[11], &vals[12], &vals[13]); + LIGER_CHECK(n == 13 || n == 14, + "LIGER_MOE_BWD_FORCE_CONFIG must have 13 or 14 comma-separated ints " + "(NS2,TN1,TK1,S1,TM3,TN3,TK3,S3,EN1,EN25,EN34,CS,TM[,GemmTM]); got '", s, "'"); + tc.NSplit2 = vals[0]; + tc.TileN1 = vals[1]; + tc.TileK1 = vals[2]; + tc.Stages1 = vals[3]; + tc.TileM3 = vals[4]; + tc.TileN3 = vals[5]; + tc.TileK3 = vals[6]; + tc.Stages3 = vals[7]; + tc.EpiChunkN1 = vals[8]; + tc.EpiChunkN25 = vals[9]; + tc.EpiChunkN34 = vals[10]; + tc.CommNumStages = vals[11]; + tc.TileM = vals[12]; + tc.GemmTileM = (n == 14) ? vals[13] : tc.TileM; + if (tc.GemmTileM <= 0) tc.GemmTileM = tc.TileM; // omitted → gemm == comm + // The wrapper selects the architecture-specific comm tile consistently + // with forward; the parsed TileM remains the GEMM tuning field. + (void)fwd_tile_m; + const DispatchEntryBwd* de = find_dispatch_entry_bwd(tc); + LIGER_CHECK(de != nullptr, + "LIGER_MOE_BWD_FORCE_CONFIG=", s, " is not in the compiled dispatch " + "table. Add it to LIGER_MOE_BWD_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); + de->fn(a, kBwdRuntimeNsplitSeed); + return; + } + + TunedConfigBwd tc{}; + LIGER_CHECK(find_nearest_tuned_bwd(compute, TK, TKE, D, I, fwd_tile_m, n_pes, &tc), + "moe_bwd_fwd_bf16_auto: no valid tuned config for compute=", compute, + " shape (T=", T, + " K=", top_k, " TK=", TK, " TKE=", TKE, " D=", D, " I=", I, + " fwd_tile_m=", fwd_tile_m, + "). Either moe_fwd_bwd_tuning_configs.cuh has no entry with TileM=", + fwd_tile_m, ", or every matching entry violates structural " + "constraints for this (D, I). Add a compatible config to " + "LIGER_MOE_BWD_TUNE_CONFIGS and re-tune."); + + const DispatchEntryBwd* de = find_dispatch_entry_bwd(tc); + LIGER_CHECK(de != nullptr, + "moe_bwd_fwd_bf16_auto: tuned config (Compute=", tc.Compute, + " NS2=", tc.NSplit2, + " TN1=", tc.TileN1, "/", tc.TileK1, "/", tc.Stages1, + " TM3=", tc.TileM3, " TN3=", tc.TileN3, "/", tc.TileK3, + "/", tc.Stages3, + " EN1=", tc.EpiChunkN1, + " EN25=", tc.EpiChunkN25, " EN34=", tc.EpiChunkN34, + " CS=", tc.CommNumStages, " TM=", tc.TileM, + ") for shape (T=", T, " K=", top_k, " TK=", TK, " D=", D, " I=", I, + ") is not in the compiled dispatch table. " + "Add it to LIGER_MOE_BWD_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); + + de->fn(a, kBwdRuntimeNsplitSeed); +} + +// Returns the compiled bwd config table as a list of 13-tuples +// (Compute, NSplit2, TileN1, TileK1, Stages1, TileM3, TileN3, TileK3, +// Stages3, EpiChunkN1, EpiChunkN25, EpiChunkN34, CommNumStages). +std::vector> +moe_bwd_list_compiled_configs() { + std::vector> out; + out.reserve(kNumDispatchEntriesBwd); + for (int i = 0; i < kNumDispatchEntriesBwd; ++i) { + const auto& e = kDispatchTableBwd[i]; + out.emplace_back(e.Compute, e.NSplit2, + e.TileN1, e.TileK1, e.Stages1, + e.TileM3, e.TileN3, e.TileK3, e.Stages3, + e.EpiChunkN1, e.EpiChunkN25, e.EpiChunkN34, + e.CommNumStages); + } + return out; +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_comm_config.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_comm_config.cuh index fea0e146..ebc532c3 100644 --- a/liger_cute_kernels/csrc/core/src/moe/moe_comm_config.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/moe_comm_config.cuh @@ -1,45 +1,11 @@ #pragma once -namespace liger { +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -template -struct MoeCommGeometry; - -template <> -struct MoeCommGeometry<90> { - static constexpr int TileM = 128; - static constexpr int NC = 4; - static constexpr int RingStageScale = 1; - static constexpr int RingStageExtra = 0; -}; - -template <> -struct MoeCommGeometry<100> { - static constexpr int TileM = 128; - static constexpr int NC = 4; - static constexpr int RingStageScale = 1; - static constexpr int RingStageExtra = 0; -}; - -inline constexpr int kMaxMoeCommTileM = - MoeCommGeometry<100>::TileM > MoeCommGeometry<90>::TileM - ? MoeCommGeometry<100>::TileM - : MoeCommGeometry<90>::TileM; -inline constexpr int kMaxMoeTunedCommStages = 8; -inline constexpr int kMaxMoeCommStagesSm90 = - kMaxMoeTunedCommStages * MoeCommGeometry<90>::RingStageScale - + MoeCommGeometry<90>::RingStageExtra; -inline constexpr int kMaxMoeCommStagesSm100 = - kMaxMoeTunedCommStages * MoeCommGeometry<100>::RingStageScale - + MoeCommGeometry<100>::RingStageExtra; -inline constexpr int kMaxMoeCommStages = - kMaxMoeCommStagesSm100 > kMaxMoeCommStagesSm90 - ? kMaxMoeCommStagesSm100 - : kMaxMoeCommStagesSm90; - -static_assert( - MoeCommGeometry<90>::TileM % MoeCommGeometry<90>::NC == 0); -static_assert( - MoeCommGeometry<100>::TileM % MoeCommGeometry<100>::NC == 0); - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "moe_comm_config_sm90.cuh" +#else +#include "moe_comm_config_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm100.cuh new file mode 100644 index 00000000..fea0e146 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm100.cuh @@ -0,0 +1,45 @@ +#pragma once + +namespace liger { + +template +struct MoeCommGeometry; + +template <> +struct MoeCommGeometry<90> { + static constexpr int TileM = 128; + static constexpr int NC = 4; + static constexpr int RingStageScale = 1; + static constexpr int RingStageExtra = 0; +}; + +template <> +struct MoeCommGeometry<100> { + static constexpr int TileM = 128; + static constexpr int NC = 4; + static constexpr int RingStageScale = 1; + static constexpr int RingStageExtra = 0; +}; + +inline constexpr int kMaxMoeCommTileM = + MoeCommGeometry<100>::TileM > MoeCommGeometry<90>::TileM + ? MoeCommGeometry<100>::TileM + : MoeCommGeometry<90>::TileM; +inline constexpr int kMaxMoeTunedCommStages = 8; +inline constexpr int kMaxMoeCommStagesSm90 = + kMaxMoeTunedCommStages * MoeCommGeometry<90>::RingStageScale + + MoeCommGeometry<90>::RingStageExtra; +inline constexpr int kMaxMoeCommStagesSm100 = + kMaxMoeTunedCommStages * MoeCommGeometry<100>::RingStageScale + + MoeCommGeometry<100>::RingStageExtra; +inline constexpr int kMaxMoeCommStages = + kMaxMoeCommStagesSm100 > kMaxMoeCommStagesSm90 + ? kMaxMoeCommStagesSm100 + : kMaxMoeCommStagesSm90; + +static_assert( + MoeCommGeometry<90>::TileM % MoeCommGeometry<90>::NC == 0); +static_assert( + MoeCommGeometry<100>::TileM % MoeCommGeometry<100>::NC == 0); + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm90.cuh new file mode 100644 index 00000000..fea0e146 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm90.cuh @@ -0,0 +1,45 @@ +#pragma once + +namespace liger { + +template +struct MoeCommGeometry; + +template <> +struct MoeCommGeometry<90> { + static constexpr int TileM = 128; + static constexpr int NC = 4; + static constexpr int RingStageScale = 1; + static constexpr int RingStageExtra = 0; +}; + +template <> +struct MoeCommGeometry<100> { + static constexpr int TileM = 128; + static constexpr int NC = 4; + static constexpr int RingStageScale = 1; + static constexpr int RingStageExtra = 0; +}; + +inline constexpr int kMaxMoeCommTileM = + MoeCommGeometry<100>::TileM > MoeCommGeometry<90>::TileM + ? MoeCommGeometry<100>::TileM + : MoeCommGeometry<90>::TileM; +inline constexpr int kMaxMoeTunedCommStages = 8; +inline constexpr int kMaxMoeCommStagesSm90 = + kMaxMoeTunedCommStages * MoeCommGeometry<90>::RingStageScale + + MoeCommGeometry<90>::RingStageExtra; +inline constexpr int kMaxMoeCommStagesSm100 = + kMaxMoeTunedCommStages * MoeCommGeometry<100>::RingStageScale + + MoeCommGeometry<100>::RingStageExtra; +inline constexpr int kMaxMoeCommStages = + kMaxMoeCommStagesSm100 > kMaxMoeCommStagesSm90 + ? kMaxMoeCommStagesSm100 + : kMaxMoeCommStagesSm90; + +static_assert( + MoeCommGeometry<90>::TileM % MoeCommGeometry<90>::NC == 0); +static_assert( + MoeCommGeometry<100>::TileM % MoeCommGeometry<100>::NC == 0); + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm100.cuh index f57f9f8d..ed12cb7d 100644 --- a/liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm100.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm100.cuh @@ -1,6 +1,6 @@ #pragma once -#include "moe_fwd_bwd_tune_configs.hpp" +#include "moe_fwd_bwd_tune_configs_sm100.hpp" // SM100 / Blackwell dispatch config set. Keep this separate from SM90 so // Blackwell-specific tuning can diverge without touching Hopper dispatch rows. diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm90.cuh index cd54162f..bd0cea71 100644 --- a/liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm90.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm90.cuh @@ -1,6 +1,6 @@ #pragma once -#include "moe_fwd_bwd_tune_configs.hpp" +#include "moe_fwd_bwd_tune_configs_sm90.hpp" // SM90 / Hopper dispatch config set. The action macros receive the same row // shapes as the base tuning menus: diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm100.hpp b/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm100.hpp new file mode 100644 index 00000000..093fbd68 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm100.hpp @@ -0,0 +1,182 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MoE fused forward + backward tuner config menu (X-macros) +// ═══════════════════════════════════════════════════════════════════ +// +// Single merged menu — replaces the old split moe_tune_configs.hpp + +// moe_bwd_tune_configs.hpp. Consumed by: +// moe.cu — LIGER_MOE_TUNE_CONFIGS → fwd template +// instantiations + runtime dispatch table. +// moe_bwd.cu — LIGER_MOE_BWD_TUNE_CONFIGS → bwd template +// instantiations + runtime dispatch table. +// tune_moe_fwd_bwd.cu — both, to build the fwd + bwd registries it +// times. The tuner filters each registry by +// TileM at runtime, so a single flat union per +// direction is all the .cu files need. +// test_moe_single.cu — LIGER_MOE_TUNE_CONFIGS for its config menu. +// +// ── Why FOUR groups ───────────────────────────────────────────────── +// The menu is split per direction and per GEMM TileM bucket: +// LIGER_MOE_FWD_TUNE_CONFIGS_TM64 / _TM128 +// LIGER_MOE_BWD_TUNE_CONFIGS_TM64 / _TM128 +// SM90-only forward candidates live in +// LIGER_MOE_FWD_TUNE_CONFIGS_SM90_ONLY +// so they do not instantiate unsupported Blackwell layouts. +// +// tune_moe_fwd_bwd tunes fwd and bwd independently and may select different +// GEMM TileM values. Their communication tile and sort-buffer granularity are +// fixed at 128, so GEMM TileM does not couple the directions. The .cu files +// don't care about the bucket split; the convenience unions below feed them +// the whole menu per direction: +// LIGER_MOE_TUNE_CONFIGS(X) = FWD_TM64(X) FWD_TM128(X) +// LIGER_MOE_BWD_TUNE_CONFIGS(X) = BWD_TM64(X) BWD_TM128(X) +// +// ── FWD row field legend (12 args) ────────────────────────────────── +// X(NSplit, +// TileN1, TileK1, Stages1, EpiChunkN1, +// TileN2, TileK2, Stages2, EpiChunkN2, +// ZBufferSlots, CommNumStages, TileM) +// +// - TileM ∈ {64, 128}; shared by MLP1 and MLP2. +// TileM=128: M-split (Layout<_2,_1,_1>), each WG owns 64 rows × TileN. +// TileM=64 : N-split (Layout<_1,_2,_1>), each WG owns 64 rows × TileN/2. +// - MLP1 TileN1 ∈ {64,128,256}. +// - MLP2 TileN2 ∈ {64,128,256}; SM90 additionally tunes TileN2=192. +// - Hopper uses ceil-div tile counts and rank-3 expert descriptors; TMA +// zero-fills load tails and drops store tails without crossing experts. +// - EpiChunkN must divide WgTileN (= TileN at TM128, TileN/2 at TM64). +// - TM64 requires TileN even (TileN/2 must be an integer N-half). +// - NC derived from TileM (NOT a knob): TM128→NC=4, TM64→NC=2. +// The NSplit % NC divisibility constraint was relaxed (ca1b943, +// ticket-based comm), so both buckets sweep NSplit ∈ {2,4,6,8,16}. +// +// ── BWD row field legend (14 args) ────────────────────────────────── +// X(NSplit, NSplit2, +// TileN1, TileK1, Stages1, +// TileM3, TileN3, TileK3, Stages3, +// EpiChunkN1, EpiChunkN25, EpiChunkN34, +// CommNumStages, TileM) +// +// EpiChunkN values grouped per smem-shape twin: +// EpiChunkN1 — mlp1_act ; EpiChunkN25 — mlp2_t + mlp5 ; EpiChunkN34 — mlp3 + mlp4. +// Derived phases: mlp2_t/mlp5 TileN=256, TileK=TileK1, Stages=Stages1; +// mlp4 swaps mlp3 (TileM=TileN3, TileN=TileM3). +// Constraints (mirrored in tuned_config_valid_bwd in moe_bwd.cu): +// D%TileK1==0, D%TileM3==0, D%256==0, I%256==0, I%TileN3==0, +// SM100 retains exact divisibility; Hopper permits TMA-padded tails; +// smem union ≤ 228 KiB. +// +// ── Per-kernel smem (Element = 2 bytes; H100 cap 228 KiB, budget 224) ─ +// FWD: MLP1 = 2·(TileM·TK1·S1 + 2·TN1·TK1·S1 + 2·64·EC1) +// MLP2 = 2·(TileM·TK2·S2 + TN2·TK2·S2 + 2·64·EC2); smem = max. +// BWD: mlp1_act = 2·(128 + 2·TN1)·TK1·S1 + 6·128·EC1 +// mlp2_t/mlp5 = 2·(128 + 256)·TK1·S1 + 2·128·EC25 +// mlp3 = 2·(TM3 + 2·TN3)·TK3·S3 + 4·TM3·EC34 +// mlp4 = 2·(2·TN3 + TM3)·TK3·S3 + 4·TN3·EC34; union = max. +// +// ── COOPERATIVE mlp3 / mlp4 CONSTRAINT ────────────────────────────── +// mlp3/mlp4 cooperative supports ONLY (TM3,TN3) = (256,128)/(128,256). +// Stages3 is PINNED at 2: at Stages3=4 the mlp3/mlp4 smem union is +// ~320-384 KiB, over the 228 KiB cap, so it fails the kernel's +// static_assert (won't compile) and bwd_shape_valid (would reject). PR +// #104 used Stages3=4 on its pre-K-split kernel; the K-loop split +// (4e39e50) redistributes work across CTAs but does NOT shrink per-CTA +// smem, so Stages3=2 is the deepest Phase-2 pipe that fits here — for +// the classic 1SM (single-CTA) kernel. SM100 always selects the paired-CTA +// 2SM path for both mlp3 and mlp4, using the canonical joined 256x256 shape; +// S3/EN34 drive the shared 2SM Stages/EpiChunkN tuning surface. +// +// ── X-macro structure caveat ──────────────────────────────────────── +// Inside a `#define`, the `\` line-continuation joins lines BEFORE +// comments are processed, so a `// X(...)` row in the middle silently +// swallows the rest of the joined line. Keep every group body's rows +// active and uncommented; park inactive rows in the candidate bank at +// the bottom of this file (outside any macro). + +// Pruned tuning menu: only rows referenced by the generated static dispatch +// tables remain here. This keeps auto-dispatch coverage identical while avoiding +// template instantiations for configs that no tuned row can select. + +// ── FWD · TileM=64 (N-split consumer, NC=2) ───────────────────────── +#define LIGER_MOE_FWD_TUNE_CONFIGS_TM64(X) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 4, 64) + +// ── FWD · TileM=128 (M-split consumer, NC=4 — divisibility relaxed) ── +#define LIGER_MOE_FWD_TUNE_CONFIGS_TM128(X) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 2, 128) \ + X(128, 64, 4, 64, 256, 64, 4, 64, 4, 2, 128) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 3, 128) \ + X(128, 64, 4, 64, 256, 64, 4, 64, 4, 3, 128) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 4, 128) \ + X(128, 64, 4, 64, 256, 64, 4, 64, 4, 4, 128) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 8, 128) \ + X(128, 64, 4, 64, 256, 64, 4, 64, 4, 8, 128) + +// Mixtral-8x22B T=4096 winner. Keep one CS variant: all four measured within +// 0.2%, and CS2 won at both E_local=1 and E_local=2. +#define LIGER_MOE_FWD_TUNE_CONFIGS_SM90_ONLY(X) \ + X(128, 64, 4, 64, 192, 64, 4, 64, 4, 2, 128) + +// ── BWD · TileM=64 (Phase-1 N-split, NC=2) ────────────────────────── +#define LIGER_MOE_BWD_TUNE_CONFIGS_TM64(X) \ + X(6, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) \ + X(2, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 64) \ + X(16, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 64) \ + X(8, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) \ + X(4, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) \ + X(8, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 64) \ + X(2, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) \ + X(16, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) + +// ── BWD · TileM=128 (Phase-1 M-split, NC=4 — divisibility relaxed) ── +#define LIGER_MOE_BWD_TUNE_CONFIGS_TM128(X) \ + X(2, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 128) \ + X(4, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(8, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(2, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(6, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(16, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 128) \ + X(16, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(8, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 128) \ + X(4, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 128) \ + X(2, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 3, 128) \ + X(4, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 3, 128) \ + X(2, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 128) \ + X(6, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 128) \ + X(16, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 128) \ + X(16, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 3, 128) + + +// ── Direction unions (whole menu per direction; fed to the .cu files) ─ +#define LIGER_MOE_TUNE_CONFIGS(X) \ + LIGER_MOE_FWD_TUNE_CONFIGS_TM64(X) \ + LIGER_MOE_FWD_TUNE_CONFIGS_TM128(X) + +#define LIGER_MOE_BWD_TUNE_CONFIGS(X) \ + LIGER_MOE_BWD_TUNE_CONFIGS_TM64(X) \ + LIGER_MOE_BWD_TUNE_CONFIGS_TM128(X) + +// ═══════════════════════════════════════════════════════════════════ +// Candidate bank — rows not winning any shape in the last sweep. +// ═══════════════════════════════════════════════════════════════════ +// To re-enable a row: copy it into the matching group macro above (as a +// new line ending in `\`, and add a `\` to the row that was previously +// last). The full pruned menus (~130 fwd / ~30 bwd variants) live in the +// git history of the deleted moe_tune_configs.hpp / moe_bwd_tune_configs.hpp; +// the highest-signal alternates are reproduced here. +// +// ── FWD alternates ────────────────────────────────────────────────── +// X(/*NS=*/ 4, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/32, /*TN2=*/128, /*TK2=*/64, /*S2=*/4, /*EC2=*/32, /*ZB=*/ 4, /*CS=*/2, /*TM=*/128) +// X(/*NS=*/ 4, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/128, /*TK2=*/64, /*S2=*/4, /*EC2=*/64, /*ZB=*/ 8, /*CS=*/2, /*TM=*/128) +// X(/*NS=*/ 4, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/256, /*TK2=*/64, /*S2=*/3, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/2, /*TM=*/128) +// X(/*NS=*/ 6, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/128, /*TK2=*/64, /*S2=*/4, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/2, /*TM=*/128) +// X(/*NS=*/ 8, /*TN1=*/128, /*TK1=*/64, /*S1=*/3, /*EC1=*/64, /*TN2=*/128, /*TK2=*/64, /*S2=*/3, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/4, /*TM=*/128) +// X(/*NS=*/ 8, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/256, /*TK2=*/64, /*S2=*/4, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/4, /*TM=*/ 64) +// X(/*NS=*/ 2, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/128, /*TK2=*/64, /*S2=*/4, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/1, /*TM=*/ 64) +// +// ── BWD alternates ────────────────────────────────────────────────── +// X(4, 8, 128, 64, 4, 128, 128, 64, 4, 32, 128, 64, 2, 128) // full-shape Phase-2 (TM3=TN3=128, S3=4) +// X(8, 4, 128, 64, 4, 128, 128, 64, 4, 32, 128, 64, 2, 128) +// X(6, 4, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 128) // NS=6 TM128 (Qwen3 I=768/TN1=128=6) +// X(2, 4, 128, 64, 4, 256, 128, 64, 2, 32, 128, 64, 2, 64) // TM64 EN25=128 diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm90.hpp b/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm90.hpp new file mode 100644 index 00000000..093fbd68 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm90.hpp @@ -0,0 +1,182 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MoE fused forward + backward tuner config menu (X-macros) +// ═══════════════════════════════════════════════════════════════════ +// +// Single merged menu — replaces the old split moe_tune_configs.hpp + +// moe_bwd_tune_configs.hpp. Consumed by: +// moe.cu — LIGER_MOE_TUNE_CONFIGS → fwd template +// instantiations + runtime dispatch table. +// moe_bwd.cu — LIGER_MOE_BWD_TUNE_CONFIGS → bwd template +// instantiations + runtime dispatch table. +// tune_moe_fwd_bwd.cu — both, to build the fwd + bwd registries it +// times. The tuner filters each registry by +// TileM at runtime, so a single flat union per +// direction is all the .cu files need. +// test_moe_single.cu — LIGER_MOE_TUNE_CONFIGS for its config menu. +// +// ── Why FOUR groups ───────────────────────────────────────────────── +// The menu is split per direction and per GEMM TileM bucket: +// LIGER_MOE_FWD_TUNE_CONFIGS_TM64 / _TM128 +// LIGER_MOE_BWD_TUNE_CONFIGS_TM64 / _TM128 +// SM90-only forward candidates live in +// LIGER_MOE_FWD_TUNE_CONFIGS_SM90_ONLY +// so they do not instantiate unsupported Blackwell layouts. +// +// tune_moe_fwd_bwd tunes fwd and bwd independently and may select different +// GEMM TileM values. Their communication tile and sort-buffer granularity are +// fixed at 128, so GEMM TileM does not couple the directions. The .cu files +// don't care about the bucket split; the convenience unions below feed them +// the whole menu per direction: +// LIGER_MOE_TUNE_CONFIGS(X) = FWD_TM64(X) FWD_TM128(X) +// LIGER_MOE_BWD_TUNE_CONFIGS(X) = BWD_TM64(X) BWD_TM128(X) +// +// ── FWD row field legend (12 args) ────────────────────────────────── +// X(NSplit, +// TileN1, TileK1, Stages1, EpiChunkN1, +// TileN2, TileK2, Stages2, EpiChunkN2, +// ZBufferSlots, CommNumStages, TileM) +// +// - TileM ∈ {64, 128}; shared by MLP1 and MLP2. +// TileM=128: M-split (Layout<_2,_1,_1>), each WG owns 64 rows × TileN. +// TileM=64 : N-split (Layout<_1,_2,_1>), each WG owns 64 rows × TileN/2. +// - MLP1 TileN1 ∈ {64,128,256}. +// - MLP2 TileN2 ∈ {64,128,256}; SM90 additionally tunes TileN2=192. +// - Hopper uses ceil-div tile counts and rank-3 expert descriptors; TMA +// zero-fills load tails and drops store tails without crossing experts. +// - EpiChunkN must divide WgTileN (= TileN at TM128, TileN/2 at TM64). +// - TM64 requires TileN even (TileN/2 must be an integer N-half). +// - NC derived from TileM (NOT a knob): TM128→NC=4, TM64→NC=2. +// The NSplit % NC divisibility constraint was relaxed (ca1b943, +// ticket-based comm), so both buckets sweep NSplit ∈ {2,4,6,8,16}. +// +// ── BWD row field legend (14 args) ────────────────────────────────── +// X(NSplit, NSplit2, +// TileN1, TileK1, Stages1, +// TileM3, TileN3, TileK3, Stages3, +// EpiChunkN1, EpiChunkN25, EpiChunkN34, +// CommNumStages, TileM) +// +// EpiChunkN values grouped per smem-shape twin: +// EpiChunkN1 — mlp1_act ; EpiChunkN25 — mlp2_t + mlp5 ; EpiChunkN34 — mlp3 + mlp4. +// Derived phases: mlp2_t/mlp5 TileN=256, TileK=TileK1, Stages=Stages1; +// mlp4 swaps mlp3 (TileM=TileN3, TileN=TileM3). +// Constraints (mirrored in tuned_config_valid_bwd in moe_bwd.cu): +// D%TileK1==0, D%TileM3==0, D%256==0, I%256==0, I%TileN3==0, +// SM100 retains exact divisibility; Hopper permits TMA-padded tails; +// smem union ≤ 228 KiB. +// +// ── Per-kernel smem (Element = 2 bytes; H100 cap 228 KiB, budget 224) ─ +// FWD: MLP1 = 2·(TileM·TK1·S1 + 2·TN1·TK1·S1 + 2·64·EC1) +// MLP2 = 2·(TileM·TK2·S2 + TN2·TK2·S2 + 2·64·EC2); smem = max. +// BWD: mlp1_act = 2·(128 + 2·TN1)·TK1·S1 + 6·128·EC1 +// mlp2_t/mlp5 = 2·(128 + 256)·TK1·S1 + 2·128·EC25 +// mlp3 = 2·(TM3 + 2·TN3)·TK3·S3 + 4·TM3·EC34 +// mlp4 = 2·(2·TN3 + TM3)·TK3·S3 + 4·TN3·EC34; union = max. +// +// ── COOPERATIVE mlp3 / mlp4 CONSTRAINT ────────────────────────────── +// mlp3/mlp4 cooperative supports ONLY (TM3,TN3) = (256,128)/(128,256). +// Stages3 is PINNED at 2: at Stages3=4 the mlp3/mlp4 smem union is +// ~320-384 KiB, over the 228 KiB cap, so it fails the kernel's +// static_assert (won't compile) and bwd_shape_valid (would reject). PR +// #104 used Stages3=4 on its pre-K-split kernel; the K-loop split +// (4e39e50) redistributes work across CTAs but does NOT shrink per-CTA +// smem, so Stages3=2 is the deepest Phase-2 pipe that fits here — for +// the classic 1SM (single-CTA) kernel. SM100 always selects the paired-CTA +// 2SM path for both mlp3 and mlp4, using the canonical joined 256x256 shape; +// S3/EN34 drive the shared 2SM Stages/EpiChunkN tuning surface. +// +// ── X-macro structure caveat ──────────────────────────────────────── +// Inside a `#define`, the `\` line-continuation joins lines BEFORE +// comments are processed, so a `// X(...)` row in the middle silently +// swallows the rest of the joined line. Keep every group body's rows +// active and uncommented; park inactive rows in the candidate bank at +// the bottom of this file (outside any macro). + +// Pruned tuning menu: only rows referenced by the generated static dispatch +// tables remain here. This keeps auto-dispatch coverage identical while avoiding +// template instantiations for configs that no tuned row can select. + +// ── FWD · TileM=64 (N-split consumer, NC=2) ───────────────────────── +#define LIGER_MOE_FWD_TUNE_CONFIGS_TM64(X) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 4, 64) + +// ── FWD · TileM=128 (M-split consumer, NC=4 — divisibility relaxed) ── +#define LIGER_MOE_FWD_TUNE_CONFIGS_TM128(X) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 2, 128) \ + X(128, 64, 4, 64, 256, 64, 4, 64, 4, 2, 128) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 3, 128) \ + X(128, 64, 4, 64, 256, 64, 4, 64, 4, 3, 128) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 4, 128) \ + X(128, 64, 4, 64, 256, 64, 4, 64, 4, 4, 128) \ + X(128, 64, 4, 64, 128, 64, 4, 64, 4, 8, 128) \ + X(128, 64, 4, 64, 256, 64, 4, 64, 4, 8, 128) + +// Mixtral-8x22B T=4096 winner. Keep one CS variant: all four measured within +// 0.2%, and CS2 won at both E_local=1 and E_local=2. +#define LIGER_MOE_FWD_TUNE_CONFIGS_SM90_ONLY(X) \ + X(128, 64, 4, 64, 192, 64, 4, 64, 4, 2, 128) + +// ── BWD · TileM=64 (Phase-1 N-split, NC=2) ────────────────────────── +#define LIGER_MOE_BWD_TUNE_CONFIGS_TM64(X) \ + X(6, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) \ + X(2, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 64) \ + X(16, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 64) \ + X(8, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) \ + X(4, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) \ + X(8, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 64) \ + X(2, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) \ + X(16, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 64) + +// ── BWD · TileM=128 (Phase-1 M-split, NC=4 — divisibility relaxed) ── +#define LIGER_MOE_BWD_TUNE_CONFIGS_TM128(X) \ + X(2, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 128) \ + X(4, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(8, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(2, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(6, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(16, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 128) \ + X(16, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 3, 128) \ + X(8, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 128) \ + X(4, 128, 64, 4, 128, 256, 64, 2, 32, 64, 64, 2, 128) \ + X(2, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 3, 128) \ + X(4, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 3, 128) \ + X(2, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 128) \ + X(6, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 128) \ + X(16, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 128) \ + X(16, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 3, 128) + + +// ── Direction unions (whole menu per direction; fed to the .cu files) ─ +#define LIGER_MOE_TUNE_CONFIGS(X) \ + LIGER_MOE_FWD_TUNE_CONFIGS_TM64(X) \ + LIGER_MOE_FWD_TUNE_CONFIGS_TM128(X) + +#define LIGER_MOE_BWD_TUNE_CONFIGS(X) \ + LIGER_MOE_BWD_TUNE_CONFIGS_TM64(X) \ + LIGER_MOE_BWD_TUNE_CONFIGS_TM128(X) + +// ═══════════════════════════════════════════════════════════════════ +// Candidate bank — rows not winning any shape in the last sweep. +// ═══════════════════════════════════════════════════════════════════ +// To re-enable a row: copy it into the matching group macro above (as a +// new line ending in `\`, and add a `\` to the row that was previously +// last). The full pruned menus (~130 fwd / ~30 bwd variants) live in the +// git history of the deleted moe_tune_configs.hpp / moe_bwd_tune_configs.hpp; +// the highest-signal alternates are reproduced here. +// +// ── FWD alternates ────────────────────────────────────────────────── +// X(/*NS=*/ 4, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/32, /*TN2=*/128, /*TK2=*/64, /*S2=*/4, /*EC2=*/32, /*ZB=*/ 4, /*CS=*/2, /*TM=*/128) +// X(/*NS=*/ 4, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/128, /*TK2=*/64, /*S2=*/4, /*EC2=*/64, /*ZB=*/ 8, /*CS=*/2, /*TM=*/128) +// X(/*NS=*/ 4, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/256, /*TK2=*/64, /*S2=*/3, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/2, /*TM=*/128) +// X(/*NS=*/ 6, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/128, /*TK2=*/64, /*S2=*/4, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/2, /*TM=*/128) +// X(/*NS=*/ 8, /*TN1=*/128, /*TK1=*/64, /*S1=*/3, /*EC1=*/64, /*TN2=*/128, /*TK2=*/64, /*S2=*/3, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/4, /*TM=*/128) +// X(/*NS=*/ 8, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/256, /*TK2=*/64, /*S2=*/4, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/4, /*TM=*/ 64) +// X(/*NS=*/ 2, /*TN1=*/128, /*TK1=*/64, /*S1=*/4, /*EC1=*/64, /*TN2=*/128, /*TK2=*/64, /*S2=*/4, /*EC2=*/64, /*ZB=*/ 4, /*CS=*/1, /*TM=*/ 64) +// +// ── BWD alternates ────────────────────────────────────────────────── +// X(4, 8, 128, 64, 4, 128, 128, 64, 4, 32, 128, 64, 2, 128) // full-shape Phase-2 (TM3=TN3=128, S3=4) +// X(8, 4, 128, 64, 4, 128, 128, 64, 4, 32, 128, 64, 2, 128) +// X(6, 4, 128, 64, 4, 256, 128, 64, 2, 32, 64, 64, 2, 128) // NS=6 TM128 (Qwen3 I=768/TN1=128=6) +// X(2, 4, 128, 64, 4, 256, 128, 64, 2, 32, 128, 64, 2, 64) // TM64 EN25=128 diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm100.cuh new file mode 100644 index 00000000..a9ff19bb --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm100.cuh @@ -0,0 +1,29 @@ +#pragma once + +#include "moe_fwd_bwd_tuning_config_types.cuh" +#include "moe_fwd_bwd_tuning_configs_single_sm100.cuh" +#include "moe_fwd_bwd_tuning_configs_multi_sm100.cuh" + +namespace liger { + +static constexpr TunedConfigFwdBwdTable kTunedConfigTablesSingle[] = { + {100, kTunedConfigsSingleSm100, kNumTunedConfigsSingleSm100}, +}; + +static constexpr int kNumTunedConfigTablesSingle = + sizeof(kTunedConfigTablesSingle) / sizeof(kTunedConfigTablesSingle[0]); + +static constexpr int kNumTunedConfigsSingle = + kNumTunedConfigsSingleSm100; + +static constexpr TunedConfigFwdBwdTable kTunedConfigTablesMulti[] = { + {100, kTunedConfigsMultiSm100, kNumTunedConfigsMultiSm100}, +}; + +static constexpr int kNumTunedConfigTablesMulti = + sizeof(kTunedConfigTablesMulti) / sizeof(kTunedConfigTablesMulti[0]); + +static constexpr int kNumTunedConfigsMulti = + kNumTunedConfigsMultiSm100; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm90.cuh new file mode 100644 index 00000000..32367059 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm90.cuh @@ -0,0 +1,29 @@ +#pragma once + +#include "moe_fwd_bwd_tuning_config_types.cuh" +#include "moe_fwd_bwd_tuning_configs_single_sm90.cuh" +#include "moe_fwd_bwd_tuning_configs_multi_sm90.cuh" + +namespace liger { + +static constexpr TunedConfigFwdBwdTable kTunedConfigTablesSingle[] = { + {90, kTunedConfigsSingleSm90, kNumTunedConfigsSingleSm90}, +}; + +static constexpr int kNumTunedConfigTablesSingle = + sizeof(kTunedConfigTablesSingle) / sizeof(kTunedConfigTablesSingle[0]); + +static constexpr int kNumTunedConfigsSingle = + kNumTunedConfigsSingleSm90; + +static constexpr TunedConfigFwdBwdTable kTunedConfigTablesMulti[] = { + {90, kTunedConfigsMultiSm90, kNumTunedConfigsMultiSm90}, +}; + +static constexpr int kNumTunedConfigTablesMulti = + sizeof(kTunedConfigTablesMulti) / sizeof(kTunedConfigTablesMulti[0]); + +static constexpr int kNumTunedConfigsMulti = + kNumTunedConfigsMultiSm90; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_sm100.cuh new file mode 100644 index 00000000..9add3862 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_sm100.cuh @@ -0,0 +1,200 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MoE: MLP compute for local + remote tiles +// ═══════════════════════════════════════════════════════════════════ +// +// MLP-only device function — called by MLP warps (0,1,4-11). +// Comm warps (2-3) are handled separately at the kernel level to +// keep CommBuffers out of the MLP register scope. +// +// Execution flow (when remote tiles exist): +// a. mlp_fused_fwd with LocalMlpTileIterator (first half of column) +// b. NamedBarrier (MLP warps only) +// c. mlp_fused_fwd with RemoteMlpTileIterator +// d. NamedBarrier (MLP warps only) +// e. mlp_fused_fwd with LocalMlpTileIterator (second half of column) +// +// Splitting the local pass around the remote pass lets the comm warps +// overlap remote gets with the first local half and remote puts with +// the second local half. When there are no remote tiles, the kernel +// runs a single full local pass and returns. +// +// For remote tiles: +// - X input comes from src_staging (filled by comm get) +// - Y output goes to dst_staging (drained by comm put) +// - The RemoteMlpTileIterator acquire/release hooks synchronize +// with comm warps via StagePipe +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp_sm100.cuh" + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// SymmetricMemoryStack keys for the persistent fwd→bwd buffers. +// Centralized so set_size / put / pop / top sites all agree. +// ═══════════════════════════════════════════════════════════════════ +inline constexpr const char* kSymmKeyXSorted = "x_sorted"; +inline constexpr const char* kSymmKeyYBuf = "y_buf"; +inline constexpr const char* kSymmKeyAllExpertOffsets = "all_expert_offsets"; +inline constexpr const char* kSymmKeyAllExpertCounts = "all_expert_counts"; + +// ═══════════════════════════════════════════════════════════════════ +// MoE shared memory — union of MLP smem + comm smem +// ═══════════════════════════════════════════════════════════════════ + +template +struct MoeSmem { + // MLP shared memory (fused phase 1 + 2 with pipeline barriers). + MlpFusedSmem mlp; + + // Comm shared memory (total_tiles count + remote_offsets pointer). + CommSmem comm; +}; + +// ═══════════════════════════════════════════════════════════════════ +// moe_fused_fwd — MLP warps only (no CommBuffers in scope) +// ═══════════════════════════════════════════════════════════════════ + +template 1 means the GEMM steps + // through SubTiles sub-tiles of Traits1::TileM (=GemmTileM) rows per + // communication staging slot. 1 → unchanged single-tile path. + int SubTiles, + int Compute = 90, + typename LocalIter, typename RemoteIter, + typename TmaLoadX, typename TmaLoadW1, typename TmaStoreZ, + typename TmaLoadZ, typename TmaLoadW2, typename TmaStoreY> +__device__ __forceinline__ void moe_fused_fwd( + MoeSmem& smem, + LocalIter& local_iter, + // Remote-phase comm pointers (no CommBuffers — keeps MLP register scope clean) + int* src_ready, + int* src_consumed, + int* dst_ready, + int* dst_consumed, + // RemoteIter re-derives (expert, pe) and the per-tile is_local fence-scope + // predicate from remote_offsets (SMEM) using these — no per-slot HBM. + int experts_per_pe, + int num_pes, + int my_pe, + int gpus_per_node, + // MLP1 TMA descriptors + TmaLoadX const& tma_load_x, + TmaLoadW1 const& tma_load_b, + TmaLoadW1 const& tma_load_c, + TmaStoreZ const& tma_store_z, + // MLP2 TMA descriptors + TmaLoadZ const& tma_load_z, + TmaLoadW2 const& tma_load_a, + TmaStoreY const& tma_store_y, + // Remote-phase TMA descriptors (X loaded from staging, Y stored to staging) + TmaLoadX const& tma_load_x_remote, + TmaStoreY const& tma_store_y_remote, + // Invariant MLP dims — const ref into __grid_constant__ .param space + const MlpDims& dims, + // Per-phase params (change between local and remote) + int num_tokens, + int num_m_tiles, + // Flat-grid logical coordinates: col = flat_id / NSplit (this CTA's + // GEMM column), grid_x = n_gemm / NSplit (number of columns). Under + // the flat launch gridDim.x = num_blocks and gridDim.y = 1, so these + // replace every blockIdx.x / gridDim.x / gridDim.y read below. + int col, + int grid_x, + int split, // this CTA's N-split index = flat_id % runtime_nsplit + int runtime_nsplit, + int local_expert_start = 0, + // Direct-to-peer Y store: per-peer TMA descriptors over + // nvshmem_ptr(local_output, peer). Gated per-tile on is_local downstream. + const TmaStoreY* y_peer_desc = nullptr) { + + int warp_id = threadIdx.x / Traits1::WarpSize; + + // Construct the per-CTA x_barrier ONCE here so its target counter + // accumulates monotonically across all 3 mlp_fused_fwd calls + // (A1, B, A2). Each phase invokes mlp_fused_fwd which calls + // x_barrier.wait() inside its tile loop — they must share the + // same barrier object so target stays aligned with the shared + // device counter. Constructing the barrier inside mlp_fused_fwd + // would reset target=0 per call while the counter keeps growing, + // and the second mlp_fused_fwd would see counter > target on + // its very first poll and exit without actually waiting. + MlpFwdCtaBarrierT x_barrier(&dims.phase_counter[col], + runtime_nsplit); + + // UNIFIED PATH: the local MLP pass has been removed. The comm-side + // TileIterator now enumerates LOCAL experts (p=0) alongside remote ones, + // and the get warp stages local X from local symmetric memory (do_get + // is_local TMA path) exactly like remote tiles. So EVERY tile this column + // owns — local or remote — flows through the RemoteMlpTileIterator / + // src_staging→MLP→dst_staging pipeline below. local_iter is unused now. + // smem.comm.per_cta_tiles is this column's total tile count (stride + // grid_x); a column with no tiles returns immediately. + (void)local_iter; (void)tma_load_x; (void)tma_store_y; (void)num_tokens; + (void)num_m_tiles; (void)local_expert_start; + int total_remote = smem.comm.per_cta_tiles; + if (total_remote == 0) return; + + bool is_leader = (warp_id == 0 && (threadIdx.x % Traits1::WarpSize) == 0); + + // Build RemoteMlpTileIterator. The staging is a flat ring of + // L = MC · CommNumStages slots, where MC = (gridDim.x · gridDim.y) / NC, + // indexed by + // slot = (blockIdx.x + j · gridDim.x) mod L (MLP, ticket-based) + // slot = xc + (j mod CommNumStages) · MC + // (comm, NumStages cycle). + // Both sides resolve to the same slot for a given global tile T and + // cover the full L. No divisibility relationship is required between + // NC, NSplit and CommNumStages (runtime: NC must divide gridDim.x · + // gridDim.y so MC is integer). + // n_gemm = grid_x · runtime_nsplit — the GEMM/comm-active CTA count. + int n_gemm = grid_x * runtime_nsplit; + RemoteIter remote_iter; + remote_iter.init( + // remote_offsets in SMEM drives the embedded walker (null → no derivation). + smem.comm.remote_offsets, smem.comm.remote_counts, + experts_per_pe, num_pes, my_pe, gpus_per_node, + smem.comm.per_cta_tiles, + col, + n_gemm, + src_ready, src_consumed, + dst_ready, dst_consumed, + is_leader, + runtime_nsplit); + + // Remote MLP dimensions: TMA descriptors cover the full flat staging + // list = MC · CommNumStages · TileM rows. MC = gridDim.x / NC — MUST match + // the ring size in allocate_moe_buffers and the consumer's ring_len + // (RemoteMlpTileIterator mc = gridDim.x/NC); sizing on n_gemm here would + // leave the wrap-around slots [n_gemm/NC·CS, gridDim.x/NC·CS) outside the + // staging coverage (OOB) when gridDim.x/NC > n_gemm/NC. + int mc = gridDim.x / NC; + int remote_num_m_tiles = mc * CommNumStages; + + // NOTE: phase_counter is monotonic across phase A and phase B — + // the modulo-gridDim.y check in mlp_fused_fwd is invariant + // under any shared starting value. The host zeroes phase_counter + // before each kernel launch (see moe.cu); resetting it here would + // race across the gridDim.y CTAs in the column. + + // num_tokens = total staging rows = remote_num_m_tiles COMM tiles × CommTileM. + // CommTileM = Traits1::TileM (=GemmTileM) × SubTiles. num_m_tiles passed is the + // COMM tile count; mlp_fused_fwd scales it by SubTiles for the GEMM sub-tiles. + mlp_fused_fwd( + smem.mlp, remote_iter, + tma_load_x_remote, tma_load_b, tma_load_c, tma_store_z, + tma_load_z, tma_load_a, tma_store_y_remote, + dims, + remote_num_m_tiles * Traits1::TileM * SubTiles, remote_num_m_tiles, x_barrier, + /*local_expert_start=*/0, col, grid_x, split, runtime_nsplit, + /*y_peer_desc=*/y_peer_desc); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_sm100.inc b/liger_cute_kernels/csrc/core/src/moe/moe_sm100.inc new file mode 100644 index 00000000..cd76316e --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_sm100.inc @@ -0,0 +1,1774 @@ +// Fused MoE: sort+dispatch → MLP+combine — torch-free core (liger_cute). +// +// Ported from LigerCommKernels' moe.cu. The device kernels and the templated +// host launcher are upstream-verbatim (namespace liger, internal to the .so); +// only the boundary changed: the remaining public entry points here are flat +// extern "C" control/configuration APIs. Tensor-carrying MoE calls are exposed +// through TVM FFI and lower directly into the internal MoeFwdArgs/MoeBwdArgs +// launch bundles. torch::Tensor never appears here. +#define LIGER_CUTE_BUILDING 1 + +#include +#include // pulls in nvshmemi_device_state_d, required by the + // inline device helpers in . +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "liger_cute/moe.h" +#include "liger_cute/check.h" +#include "liger_cute/detail/status.h" +#include "liger_cute/detail/comm_schedule.cuh" + +#include "moe_sm100.cuh" +#include "moe_symm_config_sm100.cuh" +#include "moe_launch.h" // MoeFwdArgs / MoeBwdArgs (shared with the tuner) +#include "moe_utils.cuh" +#include "moe_fwd_bwd_tune_configs_sm100.hpp" +#include "moe_dispatch_configs_sm100.cuh" +#include "tma_copy_atom_sm100.cuh" +#include "radix_sort.cuh" +#include "dispatch.cuh" +#include "combine.cuh" +#include "buffer_pool.cuh" // shim: liger:: aliases of liger_cute::detail infra +#include "fused_memset.cuh" + +using namespace cute; + +namespace liger { + +// 0 = build every dispatch family (developer fallback). CMake sets this to 90 +// for Hopper builds and 100 for Blackwell builds so normal native builds do not +// instantiate kernels for hardware that cannot run in that binary. +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif + +// init_comm_schedule lives in liger_cute::detail (comm_schedule.cuh); the host +// configure path below calls it through this alias. +using liger_cute::detail::init_comm_schedule; + +// ============================================================================ +// Global iteration counter for PeSync +// ============================================================================ +// +// PeSync's dispatch_done / mlp_done symmetric arrays must increase +// monotonically across calls (the wait-side uses NVSHMEM_CMP_GE). The +// iteration value is what each call SETs the signal to via +// nvshmemx_signal_op(SET). +// +// The counter lives in device memory (allocated by BufferPool as +// "moe_fwd_iter_counter") and is atomicAdded at the end of every +// kernel by CTA 0 thread 0 — see PeSync::advance_iteration_counter(). +// This is the form that survives CUDA graph capture: a host-side +// atomic would be read once on the host before the kernel launch and +// baked into the captured graph, so all replays would re-SET the same +// already-met value and the cross-PE wait would race; a GPU-resident +// counter incremented inside the captured kernel re-evaluates on +// every replay, giving each launch a fresh monotonic target. +// +// Buffer pool keys "moe_fwd_iter_counter" (fwd) and +// "moe_bwd_iter_counter" (bwd) name the two separate counters. + +// ============================================================================ +// MoeFusedConfig — bundles all tunable kernel parameters +// ============================================================================ +// +// TileM is fixed at 64 by the WGMMA hardware. +// NumThreads is fixed at 384 (10 MLP warps + 2 comm warps). +// All other parameters are exposed for autotuning. + +template < + typename Element_ = bfloat16_t, + int TileM_ = 128, + int TileN1_ = 128, + int TileK1_ = 64, + int Stages1_ = 4, + int EpiChunkN1_ = 64, + int TileN2_ = 256, + int TileK2_ = 64, + int Stages2_ = 4, + int EpiChunkN2_ = 64, + int ZBufferSlots_ = 4, + int CommNumStages_ = 2, + int NC_ = 4, + // GemmTileM_ decouples the GEMM compute tile (mlp1/mlp2 wgmma + TMA) from + // the COMM staging tile (TileM_). TileM_ is the comm tile — staging ring + // slots, per-tile semaphores, NC, sort/dispatch — and may be wider than the + // GEMM tile. GemmTileM_ ∈ {64,128}: when < TileM_, the GEMM waits on one comm + // semaphore once, then steps through kSubTiles = TileM_/GemmTileM_ sub-tiles + // of GemmTileM_ rows each (finer TMA + wgmma) within that one comm slot. + // Default GemmTileM_ = TileM_ → kSubTiles = 1 → bit-identical to the old + // single-TileM path. + int GemmTileM_ = TileM_, + int Compute_ = 90> +struct MoeFusedConfig { + using Element = Element_; + // GEMM tile (mlp1/mlp2): GemmTileM_ ∈ {64,128}, independent of the comm tile. + using Traits1 = Mlp1Traits< + Element, GemmTileM_, TileN1_, TileK1_, Stages1_, EpiChunkN1_>; + using Traits2 = Mlp2Traits; + + // No compile-time constraint between NC, NSplit and CommNumStages + // — the GEMM-side iterator is ticket-based (slot = T mod L, + // ticket = T / L) and synchronizes per-slot via raw atomics + // computed from the global tile index, so the only requirement + // (runtime) is that NC evenly divides gridDim.x · gridDim.y so MC + // is integer. CommNumStages is now a free buffering knob. + + // Compile-time smem budget. H100 dynamic-smem cap is 228 KiB; we leave + // ~4 KiB for the runtime-sized remote_offsets region appended by the + // launcher (see moe.cu: smem_mlp += num_pes * (experts_per_pe+1) * 4). + // Catches over-budget configs at link time instead of cudaErrorInvalidValue + // at launch. + static constexpr size_t kSmemBytes = sizeof(MoeSmem); + static constexpr size_t kSmemBudget = 224 * 1024; + static_assert(kSmemBytes <= kSmemBudget, + "MoE fused fwd kernel smem footprint exceeds 224 KiB budget — " + "reduce TileN/TileK/Stages for this config"); + + static constexpr int kNC = NC_; + static constexpr int kZBufferSlots = ZBufferSlots_; + static constexpr int kCommNumStages = CommNumStages_; + // kTileM is the COMM tile (staging ring slots, per-tile semaphores, NC, + // sort/dispatch granularity). + static constexpr int kTileM = TileM_; + static constexpr int kGemmTileM = GemmTileM_; // gemm tile ∈ {64,128} + static constexpr int kSubTiles = TileM_ / GemmTileM_; + static constexpr int kCompute = Compute_; + static constexpr int kNumThreads = 384; // fixed + + static_assert(TileM_ % GemmTileM_ == 0, + "CommTileM must be an integer multiple of GemmTileM"); + static_assert(TileM_ % NC_ == 0, + "CommTileM must be divisible by NC"); + static_assert(kSubTiles == 1 || kSubTiles == 2, + "CommTileM/GemmTileM must be 1 or 2"); + static_assert(GemmTileM_ == 64 || GemmTileM_ == 128, + "GemmTileM must be 64 or 128"); + + // Iterators operate on the COMM tile (kTileM), independent of the GEMM + // sub-tile size. + using LocalIter = LocalMlpTileIterator; + using RemoteIter = RemoteMlpTileIterator< + Element, kCommNumStages, kNC, kTileM, kSubTiles>; +}; + +// ============================================================================ +// MoeParams -- sort + dispatch + MLP + combine (single kernel) +// ============================================================================ + +template +struct MoeParams { + using Element = typename Config::Element; + + // Router outputs (read-only, from caller's tensors). + const int* expert_indices; // [T, K] + const Element* expert_weights; // [T, K] + + // Sort buffers. + SortBuffers sort; + + // Dispatch. + const Element* tokens; // [T, D] input (symmetric) + Element* dispatch_dst; // [total_slots, D] sorted output (symmetric) + + // Combine output. + Element* combine_output; // [T, D] final output + + // MLP + comm. + CommBuffers comm; + + // Dimensions. + int num_tokens; + int intermediate_dim; + int top_k; + int total_slots; + int num_m_tiles; // sort input tiles +}; + +// ============================================================================ +// MoeBuffers -- all device/symmetric allocations +// ============================================================================ + +template +struct MoeBuffers { + using Element = typename Config::Element; + + // MLP intermediate Z. + Element* z_buf; + int z_rows; + + // Sorted X (dispatch destination) -- symmetric, remote PEs get from here. + Element* x_sorted; + + // Output Y -- symmetric, local MLP writes here, remote PEs put here. + Element* y_buf; + + // Phase counters. + int* phase_counter; + + // Comm staging. + Element* src_staging; + Element* dst_staging; + int* src_ready; + int* src_consumed; + int* dst_ready; + int* dst_consumed; + int* tile_expert_ids; + + // All PEs' expert offsets -- symmetric. + int* all_expert_offsets; + int* all_expert_counts; + int* cta_counter; // [1] device memory — shared monotonic + + // Sort buffers. + int* tile_expert_counts; + int* sorted_token_ids; + int* token_expert_slots; + int* expert_offsets; // symmetric + int* expert_counts; // symmetric exact counts + int* cta_done; + int* cta_sums; + + // Per-tile expert IDs for MLP local iterator. + int* mlp_tile_expert_ids; +}; + +// MoeSymmConfig + get_symm_config() declaration live in the shared header so +// moe.cu and moe_bwd.cu can't disagree on the layout (see header). This TU +// owns the single instance. +MoeSymmConfig& get_symm_config() { + static MoeSymmConfig cfg; + return cfg; +} + +void moe_configure_symmetric( + int max_tokens, + int hidden_dim, + int max_num_experts, + int max_top_k, + int num_pes, + int num_hosts, + int gpus_per_host) { + static constexpr int kTileM = kMaxMoeCommTileM; + + // num_hosts × gpus_per_host must equal the team size so the comm + // schedule (g_dest_table / g_rank_table) addresses exactly the PEs + // in the team. Mismatch would silently map ranks to non-existent + // destinations (or skip some). + LIGER_CHECK(num_hosts * gpus_per_host == num_pes, + "moe_configure_symmetric: num_hosts (", num_hosts, ") * gpus_per_host (", + gpus_per_host, ") = ", num_hosts * gpus_per_host, + " must equal num_pes (", num_pes, ")"); + + // Populate g_dest_table / g_rank_table in device constant memory. + // Idempotent — safe to re-call with the same (N, M). + init_comm_schedule(num_hosts, gpus_per_host); + + auto& cfg = get_symm_config(); + cfg.max_total_slots = max_tokens * max_top_k + max_num_experts * kTileM; + cfg.max_num_experts = max_num_experts; + cfg.hidden_dim = hidden_dim; + cfg.num_pes = num_pes; + cfg.experts_per_pe = max_num_experts / num_pes; + cfg.max_top_k = max_top_k; + cfg.initialized = true; + + // Preset sizes for the symmetric stack used to hold X / Y / routing metadata + // across fwd → bwd. Element type is bf16 (= 2 bytes); kept in sync + // with allocate_moe_buffers below. + constexpr std::size_t kElemBytes = 2; // sizeof(__nv_bfloat16) + std::size_t symm_slot_bytes = + static_cast(cfg.max_total_slots) * cfg.hidden_dim * kElemBytes; + std::size_t all_off_bytes = + static_cast(cfg.num_pes) * (cfg.max_num_experts + 1) * sizeof(int); + std::size_t all_count_bytes = + static_cast(cfg.num_pes) * cfg.max_num_experts * sizeof(int); + + auto& stack = global_symmetric_stack(); + stack.set_size(kSymmKeyXSorted, symm_slot_bytes); + stack.set_size(kSymmKeyYBuf, symm_slot_bytes); + stack.set_size(kSymmKeyAllExpertOffsets, all_off_bytes); + stack.set_size(kSymmKeyAllExpertCounts, all_count_bytes); +} + +// Releases the forward metadata entries pushed onto the symmetric stack by the +// most recent fwd call. Must be called collectively +// on every PE in LIFO order vs. the fwd that produced them. After this, +// the buffers can be reused by a subsequent fwd. The torch::Tensor views +// returned by fwd remain valid pointers but will alias whatever lives in +// those slots after the next put. +void moe_pop_fwd() { + auto& stack = global_symmetric_stack(); + // LIFO order: bottom pushes first. Reverse the put order in + // allocate_moe_buffers (x_sorted, y_buf, all_expert_offsets, + // all_expert_counts). + stack.pop(kSymmKeyAllExpertCounts); + stack.pop(kSymmKeyAllExpertOffsets); + stack.pop(kSymmKeyYBuf); + stack.pop(kSymmKeyXSorted); +} + +template +static MoeBuffers allocate_moe_buffers( + BufferPool& pool, + cudaStream_t stream, + int grid_x, + int num_blocks, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int top_k, + int num_experts, + int experts_per_pe, + int num_pes, + int total_slots, + int num_m_tiles, + // Caller-owned torch storage for the two sort outputs that bwd + // needs to consume. Passed in so we don't lease buffer-pool + // slots that the next fwd would overwrite. + int* token_expert_slots_ptr, + int* mlp_tile_expert_ids_ptr) { + + using Element = typename Config::Element; + static constexpr int kTileM = Config::kTileM; + static constexpr int kZBufferSlots = Config::kZBufferSlots; + + // Staging is a flat ring of L = MC · kCommNumStages slots, where + // MC = num_blocks / NC (num_blocks == launched gridDim.x). This MUST equal + // the device-side comm producer's stride (mlp_comms.cuh: MC = n_comm/NC, + // n_comm = (gridDim.x/NC)·NC) AND the MLP consumer's ring stride + // (RemoteMlpTileIterator::init: mc = gridDim.x/NC), so every tile's + // slot = g % L agrees across producer / consumer / allocation. Sizing this + // on grid_x·NSplit (= n_gemm) instead under-sizes the signal arrays whenever + // gridDim.x/NC > n_gemm/NC — i.e. num_sms % NSplit != 0, e.g. NSplit=8 on + // 132 SMs (33 vs 32) — and the gridDim.x-based ring then indexes past the + // end of src_ready / tile_expert_ids (OOB). num_blocks/NC == gridDim.x/NC + // == n_gemm/NC in the divisible (legacy) case, so this is a no-op there. + // Comm strides MC through NumStages per xc; MLP (ticket iterator) wraps mod L. + // Raw SM count — used as the device-independent upper bound on the comm + // staging row count (mc <= sm_count) so the symmetric staging size is a + // worst-case bound across configs (see staging sizing below). + int sm_count = 0; + { + int dev = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev); + } + + MoeBuffers b; + + // Symmetric allocations use max sizes from MoeSymmConfig + // so all PEs allocate the same sizes regardless of per-call config. + auto& scfg = get_symm_config(); + LIGER_CHECK(scfg.initialized, "Call moe_configure_symmetric before moe_fused_fwd_bf16"); + + // Persistent symmetric tensors and routing metadata live in the symmetric + // stack so they can survive across fwd calls into bwd. + auto& sym_stack = global_symmetric_stack(); + + // Z buffer (device). + // size_t cast on b.z_rows: at NSplit=2 ZBuf=16 I=16384, the product + // z_rows · intermediate_dim · 2 = ~4.4 GB which overflows int32 if + // the multiply happens in int (b.z_rows and intermediate_dim are + // both int). Cast first so the entire product is in size_t. + int max_runtime_grid_x = (num_blocks + 1) / 2; // min candidate NS = 2 + int z_m_tiles = max_runtime_grid_x * kZBufferSlots; + b.z_rows = z_m_tiles * kTileM; + size_t z_bytes = static_cast(b.z_rows) * intermediate_dim * sizeof(Element); + b.z_buf = reinterpret_cast(pool.get_device("moe_z_buf", z_bytes)); + + // Sorted X -- symmetric, dispatch writes here, remote PEs get from here. + b.x_sorted = reinterpret_cast(sym_stack.put(kSymmKeyXSorted)); + + // Output Y -- symmetric, local MLP writes here, remote PEs put here. + b.y_buf = reinterpret_cast(sym_stack.put(kSymmKeyYBuf)); + + // Phase counters. + size_t counter_bytes = max_runtime_grid_x * sizeof(int); + b.phase_counter = reinterpret_cast( + pool.get_device("moe_phase_counter", counter_bytes)); + if (cudaError_t e = cudaMemsetAsync(b.phase_counter, 0, counter_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(phase_counter) failed: ", cudaGetErrorString(e)); + + // Flat staging list — L = grid_x * kCommNumStages slots, each + // holding TileM rows × hidden_dim columns. + // Size by the CONFIGURED MAX on every axis (scfg.hidden_dim, max_comm_stages, + // max_tile_m), NOT the per-call/per-config values: the symmetric pool caches + // by name and aborts on grow, so a small first config must not lock in a + // staging buffer a later larger config can't grow into. This covers all three + // independently-varying axes that the tuner sweeps: + // - hidden_dim: D=512 then D=8192 + // - CommNumStages: CS=2 then CS=8 (NS=16 config) + // - TileM: 64 then 128 + // The kernel indexes slots by the per-call (comm_l, tile_elems <= max), so an + // over-sized buffer is always safe (symmetric memory may be larger than + // needed, never smaller). + size_t staging_rows_ub = + (size_t)sm_count * scfg.max_comm_stages * scfg.max_tile_m; + size_t src_bytes = staging_rows_ub * scfg.hidden_dim * sizeof(Element); + size_t dst_bytes = staging_rows_ub * scfg.hidden_dim * sizeof(Element); + b.src_staging = reinterpret_cast(pool.get_symmetric("moe_src_staging", src_bytes)); + b.dst_staging = reinterpret_cast(pool.get_symmetric("moe_dst_staging", dst_bytes)); + + // Device buffers are cached by key across tuning/config changes, so size and + // clear them at a configuration-independent upper bound. Since + // kSubTiles <= NC, (num_blocks / NC) * stages * kSubTiles is bounded by + // sm_count * max_comm_stages. + size_t signal_l_ub = + static_cast(sm_count) * scfg.max_comm_stages; + size_t sig_bytes = signal_l_ub * sizeof(int); + b.src_ready = static_cast(pool.get_device("moe_src_ready", sig_bytes)); + b.src_consumed = static_cast(pool.get_device("moe_src_consumed", sig_bytes)); + b.dst_ready = static_cast(pool.get_device("moe_dst_ready", sig_bytes)); + b.dst_consumed = static_cast(pool.get_device("moe_dst_consumed", sig_bytes)); + if (cudaError_t e = cudaMemsetAsync(b.src_ready, 0, sig_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(src_ready) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(b.src_consumed, 0, sig_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(src_consumed) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(b.dst_ready, 0, sig_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(dst_ready) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(b.dst_consumed, 0, sig_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(dst_consumed) failed: ", cudaGetErrorString(e)); + + // Per-parent-slot expert IDs use the same conservative upper bound. + size_t eid_bytes = signal_l_ub * sizeof(int); + b.tile_expert_ids = static_cast(pool.get_device("moe_tile_expert_ids", eid_bytes)); + + // All PEs' expert offsets (symmetric). Use max experts for stable sizing. + b.all_expert_offsets = static_cast(sym_stack.put(kSymmKeyAllExpertOffsets)); + b.all_expert_counts = static_cast(sym_stack.put(kSymmKeyAllExpertCounts)); + + // CTA counter for local all-CTA barriers around NVSHMEM collectives. + b.cta_counter = static_cast( + pool.get_device("moe_cta_counter", sizeof(int))); + + // Sort buffers -- use scfg max sizes to avoid reallocation across configs. + int max_sorted_slots = scfg.max_total_slots; + int max_m_tiles = scfg.max_total_slots / kTileM; + int max_tile_counts = max_m_tiles * scfg.max_num_experts; + int max_token_slots_ub = scfg.max_total_slots; + + (void)max_token_slots_ub; + b.tile_expert_counts = static_cast( + pool.get_device("moe_sort_tile_counts", max_tile_counts * sizeof(int))); + b.sorted_token_ids = static_cast( + pool.get_device("moe_sort_sorted_ids", max_sorted_slots * sizeof(int))); + // token_expert_slots: caller-owned torch tensor (passed in). Used by + // the bwd path after this call returns — must not alias a pool slot. + b.token_expert_slots = token_expert_slots_ptr; + b.expert_offsets = static_cast( + pool.get_symmetric("moe_sort_offsets", (scfg.max_num_experts + 1) * sizeof(int))); + b.expert_counts = static_cast( + pool.get_symmetric("moe_sort_counts", scfg.max_num_experts * sizeof(int))); + b.cta_done = static_cast( + pool.get_device("moe_sort_cta_done", num_blocks * sizeof(int))); + b.cta_sums = static_cast( + pool.get_device("moe_sort_cta_sums", num_blocks * scfg.max_num_experts * sizeof(int))); + + // Per-tile expert IDs for MLP — caller-owned torch tensor. + b.mlp_tile_expert_ids = mlp_tile_expert_ids_ptr; + + if (cudaError_t e = cudaMemsetAsync(b.cta_done, 0, num_blocks * sizeof(int), stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(cta_done) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(b.sorted_token_ids, 0xff, + max_sorted_slots * sizeof(int), stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(sorted_token_ids) failed: ", cudaGetErrorString(e)); + + return b; +} + +// ============================================================================ +// GET-via-TMA descriptor builder (raw CUtensorMap, FWD-only) +// ============================================================================ +// +// Box = [box_rows, kGetKChunk] over a row-major [n_rows, hidden_dim] bf16 +// tensor. SWIZZLE_NONE keeps the global bytes row-major, so the MLP's +// swizzled reader (tma_load_x_remote) sees correct data. Returns false on +// encode failure → caller disables the TMA-get path (falls back to getmem). +static inline bool make_get_tma_map(CUtensorMap& m, void* base, + int hidden_dim, int n_rows, int box_rows) { + uint64_t gdim[2] = { (uint64_t)hidden_dim, (uint64_t)n_rows }; + uint64_t gstride[1] = { (uint64_t)hidden_dim * sizeof(bfloat16_t) }; + uint32_t bdim[2] = { (uint32_t)liger::kGetKChunk, (uint32_t)box_rows }; + uint32_t estride[2] = { 1, 1 }; + CUresult r = cuTensorMapEncodeTiled( + &m, CU_TENSOR_MAP_DATA_TYPE_UINT16, 2, base, gdim, gstride, bdim, estride, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + return r == CUDA_SUCCESS; +} + +// ============================================================================ +// Fused kernel: sort → dispatch → MLP + combine +// ============================================================================ + +template +__global__ void __launch_bounds__(Config::kNumThreads) +moe_fused_kernel( + __grid_constant__ TmaLoadX const tma_load_x, + __grid_constant__ TmaLoadW1 const tma_load_b, + __grid_constant__ TmaLoadW1 const tma_load_c, + __grid_constant__ TmaStoreZ const tma_store_z, + __grid_constant__ TmaLoadZ const tma_load_z, + __grid_constant__ TmaLoadW2 const tma_load_a, + __grid_constant__ TmaStoreY const tma_store_y, + __grid_constant__ TmaLoadX const tma_load_x_remote, + __grid_constant__ TmaStoreY const tma_store_y_remote, + __grid_constant__ TmaStoreYCombine const tma_store_y_combine, + __grid_constant__ MlpDims const mlp_dims, + __grid_constant__ MoeParams const p, + __grid_constant__ int const static_nsplit, + __grid_constant__ GetTmaDescs const get_descs, + __grid_constant__ PeerYStoreDescs const peer_y) { + + using Element = typename Config::Element; + using Traits1 = typename Config::Traits1; + using Traits2 = typename Config::Traits2; + using LocalIter = typename Config::LocalIter; + using RemoteIter = typename Config::RemoteIter; + constexpr int kTileM = Config::kTileM; + constexpr int kNumThreads = Config::kNumThreads; + constexpr int kZBufferSlots = Config::kZBufferSlots; + constexpr int kCommNumStages = Config::kCommNumStages; + constexpr int kNC = Config::kNC; + + // -- Sort ---------------------------------------------------------------- + { + extern __shared__ int sort_smem[]; + sort_tokens( + sort_smem, p.expert_indices, p.sort, + p.num_tokens, p.comm.num_experts, p.top_k, p.num_m_tiles); + } + __threadfence_system(); // device-scope (NVLink-only build; IB needs system scope) + __syncthreads(); + + // Target-based cross-CTA barrier — single instance per kernel call, + // reused by every PeSync::barrier / signal_dispatch / signal_mlp + // invocation below. Constructing a fresh barrier per call resets + // the per-CTA target to 0 while the shared device counter keeps + // growing → fast CTAs spin past the next-call target instantly. + SyncThreadsCtaCounterBarrier pe_barrier(p.comm.pe_sync.cta_counter, + (int)gridDim.x * (int)gridDim.y); + + // Cross-CTA barrier: sort writes sorted_token_ids across CTAs; + // dispatch and combine read them from other CTAs. + p.comm.pe_sync.barrier(pe_barrier); + + // -- Broadcast expert offsets to all PEs ---------------------------------- + broadcast_expert_offsets(p.comm); + __syncthreads(); + + // -- Dispatch ------------------------------------------------------------- + dispatch_tokens( + p.tokens, p.dispatch_dst, + p.sort.sorted_token_ids, + p.comm.hidden_dim, p.total_slots); + __syncthreads(); + + // H2 (IB coherence): dispatch_tokens wrote dispatch_dst (= x_sorted) with + // plain GPU stores; remote PEs RDMA-read it in do_get after wait_dispatch. + // Flush those stores to HBM so the remote NIC reads fresh data — the quiet + // below only orders this PE's OUTGOING nvshmem ops, not these local + // stores, and __syncthreads is device-scope. The offset-broadcast source + // is already covered by radix_sort's __threadfence_system. NVLink P2P + // reads are L2-coherent, so single-host never exercised this. + __threadfence_system(); // device-scope (NVLink-only build; IB needs system scope) + + p.comm.pe_sync.barrier_all(pe_barrier); + + // -- MLP (local + remote) ------------------------------------------------ + { + extern __shared__ char raw_smem[]; + auto& smem = *reinterpret_cast*>(raw_smem); + + size_t offsets_region = sizeof(MoeSmem); + offsets_region = (offsets_region + sizeof(int) - 1) & ~(sizeof(int) - 1); + size_t offsets_count = + (size_t)p.comm.num_pes() * (p.comm.experts_per_pe + 1); + if (threadIdx.x == 0) { + smem.comm.remote_offsets = + reinterpret_cast(raw_smem + offsets_region); + smem.comm.remote_counts = + smem.comm.remote_offsets + offsets_count; + } + __syncthreads(); + + // GET-via-TMA bounce region: placed after remote_offsets, 128-aligned + // (TMA dst smem must be 128B-aligned). MUST match the host smem_size + // computation (same sizeof(MoeSmem), num_pes, experts_per_pe). Only + // touched by the comm warps' do_get when get_descs.enabled. + char* get_bounce = nullptr; + if (get_descs.enabled) { + size_t boff = offsets_region + + (offsets_count + + (size_t)p.comm.num_pes() * p.comm.experts_per_pe) + * sizeof(int); + boff = (boff + 127) & ~((size_t)127); + get_bounce = raw_smem + boff; + } + + //if (threadIdx.x == 0) + // smem.comm.total_tiles = -1; + + // Flat-grid logical coordinates for this CTA. GEMM and comm gate to + // the first n_gemm = floor_NS CTAs (the comm↔MLP staging ring needs + // grid_x = n_gemm/static_nsplit and MC = n_gemm/NC consistent); the gap CTAs + // [n_gemm, num_blocks) fall through to the shared pe_barrier only. + int flat_id = blockIdx.x; + int n_gemm = p.comm.n_gemm; + int grid_x = n_gemm / static_nsplit; + int col = flat_id / static_nsplit; + int runtime_nsplit = static_nsplit; + + // gemm_active removed: comm warps self-gate on their precomputed tile + // count (smem.comm.total_tiles) so a CTA can do comm-only work without + // participating in GEMM (the point of the comm/GEMM grid relaxation). + // Run the prologue unconditionally so every launched CTA has total_tiles + // initialized; the MLP/GEMM path still gates on GEMM-grid membership + // (flat_id < n_gemm) below. + nvshmem_comm_prologue(smem.comm, p.comm, static_nsplit); + __syncthreads(); + + if (threadIdx.x == 0) { + int total_m_tiles = smem.comm.global_total; + int mlp_tiles = total_m_tiles; + if constexpr (Config::kSubTiles > 1) { + mlp_tiles = 0; + for (int pe = 0; pe < p.comm.num_pes(); ++pe) { + for (int e = 0; e < p.comm.experts_per_pe; ++e) { + int count = smem.comm.remote_counts[ + pe * p.comm.experts_per_pe + e]; + mlp_tiles += + (count + Config::kGemmTileM - 1) / + Config::kGemmTileM; + } + } + } + runtime_nsplit = select_runtime_nsplit_from_tiles( + mlp_tiles, + mlp_dims.num_n_tiles_1, + mlp_dims.num_n_tiles_2, + mlp_dims.num_k_tiles_1, + mlp_dims.num_k_tiles_2, + (int)gridDim.x, + static_nsplit); + n_gemm = ((int)gridDim.x / runtime_nsplit) * runtime_nsplit; + grid_x = n_gemm / runtime_nsplit; + col = flat_id / runtime_nsplit; + + int mlp_count = 0; + if (flat_id < n_gemm && col < mlp_tiles) + mlp_count = (mlp_tiles - 1 - col) / grid_x + 1; + smem.comm.per_cta_tiles = mlp_count; + smem.comm.runtime_nsplit = runtime_nsplit; + smem.comm.runtime_n_gemm = n_gemm; + smem.comm.runtime_grid_x = grid_x; + } + __syncthreads(); + runtime_nsplit = smem.comm.runtime_nsplit; + n_gemm = smem.comm.runtime_n_gemm; + grid_x = smem.comm.runtime_grid_x; + col = flat_id / runtime_nsplit; + + // Diverge: comm warps handle NVSHMEM, MLP warps handle compute. + // Separate scopes prevent CommBuffers from inflating MLP register budget. + // Fwd uses warps 1-3 (2 getters + 1 putter). + int warp_id = threadIdx.x / 32; + bool is_comm_warp = (warp_id == kCommGetWarp || warp_id == kCommPutWarp); + + if (is_comm_warp) { + // Comm warps: dedicated get/put — full CommBuffers access. Self-gate + // on the precomputed comm tile count — no GEMM coupling, so a + // comm-only CTA (no local GEMM work) still drives its comm tiles. + if (smem.comm.total_tiles > 0) { + nvshmem_comm_main< + Element, kTileM, kCommNumStages, kNC, Config::kSubTiles>( + smem.comm, p.comm, &get_descs, get_bounce); + } + } else if (flat_id < n_gemm) { + // MLP warps: no CommBuffers in scope. Gap CTAs (flat_id >= n_gemm) + // skip this entirely and fall through to the pe_barrier below. + int local_e_start = p.comm.my_pe() * p.comm.experts_per_pe; + int local_e_end = local_e_start + p.comm.experts_per_pe; + int local_offset = p.sort.expert_offsets[local_e_start]; + int local_end = p.sort.expert_offsets[local_e_end]; + int local_m_tiles = (local_end - local_offset) / kTileM; + int local_tile_off = local_offset / kTileM; + + LocalIter local_iter; + local_iter.init( + p.dispatch_dst + local_offset * p.comm.hidden_dim, + p.sort.tile_expert_ids + local_tile_off, + local_m_tiles, p.comm.hidden_dim, + /*col_stride=*/grid_x, /*start_col=*/col, local_tile_off); + + moe_fused_fwd( + smem, local_iter, + p.comm.src_ready, p.comm.src_consumed, + p.comm.dst_ready, p.comm.dst_consumed, + p.comm.experts_per_pe, p.comm.num_pes(), p.comm.my_pe(), + get_descs.gpus_per_node, + tma_load_x, tma_load_b, tma_load_c, tma_store_z, + tma_load_z, tma_load_a, tma_store_y, + tma_load_x_remote, tma_store_y_remote, + mlp_dims, + local_m_tiles * kTileM, local_m_tiles, + col, grid_x, /*split=*/flat_id % runtime_nsplit, runtime_nsplit, + local_e_start, + /*y_peer_desc=*/peer_y.enabled ? peer_y.desc : nullptr); + } + + __threadfence(); + __syncthreads(); + p.comm.pe_sync.barrier_all(pe_barrier); + } + + // -- Combine -------------------------------------------------------------- + // Reuse the kernel's dynamic SMEM region as CombineSmem. wait_mlp above + // is a cross-CTA + cross-PE barrier that ends in its own __syncthreads, so + // the CTA is already aligned and the MLP/Comm SMEM contents are dead by + // this point — no extra __syncthreads needed before the reinterpret. + { + constexpr int kCombineNumWGs = + kNumThreads / (kCombineWarpsPerWG * kCombineWarpSize); + extern __shared__ char raw_smem[]; + auto& combine_smem = + *reinterpret_cast*>(raw_smem); + combine_tokens_tma( + reinterpret_cast(p.comm.local_output), + p.expert_weights, + p.sort.token_expert_slots, + tma_store_y_combine, + combine_smem, + p.comm.hidden_dim, + p.num_tokens, + p.top_k); + } + +} + +// ============================================================================ +// Host wrapper +// ============================================================================ + +// MoeFwdArgs (the POD launch bundle) is defined in moe_launch.h, shared with the +// tuner. The runtime dispatch table holds a single function-pointer type +// (MoeFwdFn) while every tuned config is a distinct template instantiation. + +template < + typename Element_ = bfloat16_t, + int TileN1 = 128, + int TileK1 = 64, + int Stages1 = 4, + int TileN2 = 256, + int TileK2 = 64, + int Stages2 = 4, + int ZBufferSlots = 4, + int CommNumStages = 1, + int EpiChunkN1 = 64, // mlp1 epilogue chunk; must divide TileN1 + int EpiChunkN2 = 64, // mlp2 epilogue chunk; must divide TileN2 + int TileM = 128, // COMM tile (staging/semaphores/NC); 128 + int GemmTileM = TileM, // GEMM tile ∈ {64,128}; default == comm (no sub-tiling) + int Compute = 90> +void +moe_fused_fwd_bf16(const MoeFwdArgs& a, int static_nsplit) { + + // Unpack the POD argument bundle into the names the upstream body uses. + const int num_experts = a.num_experts; + const int top_k = a.top_k; + const nvshmem_team_t team = a.team; + cudaStream_t stream = a.stream; + + // NC = number of cooperating comm CTAs per (xc, yc) tile. With the + // ticket-based GEMM-side iterator, runtime NS and NC no longer need to + // share a divisibility relationship — the only runtime requirement + // is NC | (gridDim.x · gridDim.y) so MC is integer. + // + // Wrapper param order (10 "core" args + trailing EpiChunkN + TileM): + // Element, TN1, TK1, S1, TN2, TK2, S2, ZBuf, CS, + // EpiChunkN1, EpiChunkN2, TileM + // MoeFusedConfig interleaves TileM/EpiChunkN inside each phase: + // Element, TileM, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, + // ZBuf, CS, NC + // The remapping happens here. TileM is shared by MLP1 and MLP2. + // + // NC choice (verified via bench_mlp_comms BW sweep): + // - TileM=128: NC=4 gives ~2× the comm bandwidth of NC=2 because + // 4 cooperating comm CTAs split each tile, deepening the comm + // pipeline (more in-flight slots = grid.x/NC). + // - TileM=64: NC=2 already matches NC=4 throughput because tiles + // are half-size, so twice as many flow through the same pipeline. + // Both architectures use CommTileM=128 with NC=4. Keeping CommTileM uniform per + // architecture preserves + // all-to-all dispatch padding (x_sorted / expert_offsets, padded to the + // comm tile) across PEs. The 256/8 ratio keeps each SM100 comm CTA's row + // band and GET-bounce footprint unchanged. + using CommGeometry = MoeCommGeometry; + using Config = MoeFusedConfig< + Element_, /*CommTileM=*/CommGeometry::TileM, + TileN1, TileK1, Stages1, EpiChunkN1, + TileN2, TileK2, Stages2, EpiChunkN2, + ZBufferSlots, + CommNumStages * CommGeometry::RingStageScale + + CommGeometry::RingStageExtra, + /*NC=*/CommGeometry::NC, + /*GemmTileM=*/GemmTileM, /*Compute_=*/Compute>; + + using Element = typename Config::Element; + using Traits1 = typename Config::Traits1; + using Traits2 = typename Config::Traits2; + using TmaLoadAtom = TmaLoadAtomForCompute; + using TmaStoreAtom = TmaStoreAtomForCompute; + constexpr int kTileM = Config::kTileM; + constexpr int kNumThreads = Config::kNumThreads; + + // Iteration counter lives on the device and is incremented at the + // end of the kernel — see PeSync::advance_iteration_counter() and + // the comment block at the top of this file. + int num_tokens = a.num_tokens; + int hidden_dim = a.hidden_dim; + int experts_per_pe = a.experts_per_pe; // local experts (from weight matrices) + int intermediate_dim = a.intermediate_dim; + + // Upper bound for total_slots: T * K + E * TileM (TileM padding per expert). + // Actual value is computed by sort (expert_offsets[E]) inside the kernel. + int max_total_slots = num_tokens * top_k + num_experts * kTileM; + int num_m_tiles_in = (num_tokens + kTileM - 1) / kTileM; // sort input tiles + + if constexpr (Compute == 100) { + LIGER_CHECK(hidden_dim % Traits1::TileK == 0, + "hidden_dim (", hidden_dim, ") must be divisible by MLP1 TileK (", Traits1::TileK, ")"); + LIGER_CHECK(intermediate_dim % Traits1::TileN == 0, + "intermediate_dim (", intermediate_dim, ") must be divisible by MLP1 TileN (", Traits1::TileN, ")"); + LIGER_CHECK(intermediate_dim % Traits2::TileK == 0, + "intermediate_dim (", intermediate_dim, ") must be divisible by MLP2 TileK (", Traits2::TileK, ")"); + LIGER_CHECK(hidden_dim % Traits2::TileN == 0, + "hidden_dim (", hidden_dim, ") must be divisible by MLP2 TileN (", Traits2::TileN, ")"); + } + LIGER_CHECK(hidden_dim % 8 == 0, + "hidden_dim (", hidden_dim, ") must be divisible by 8 (int4 vectorization in dispatch/combine)"); + LIGER_CHECK(intermediate_dim % 8 == 0, + "intermediate_dim (", intermediate_dim, ") must be divisible by 8 (int4 vectorization)"); + LIGER_CHECK(top_k <= kCombineMaxTopK, + "top_k (", top_k, ") exceeds kCombineMaxTopK (", kCombineMaxTopK, + "). Bump the constant in combine.cuh."); + + int num_n_tiles_1 = (Compute == 90) + ? (intermediate_dim + Traits1::TileN - 1) / Traits1::TileN + : intermediate_dim / Traits1::TileN; + int num_n_tiles_2 = (Compute == 90) + ? (hidden_dim + Traits2::TileN - 1) / Traits2::TileN + : hidden_dim / Traits2::TileN; + + // Relaxed: runtime NS > num_n_tiles is allowed; surplus split CTAs run + // empty inner loops but still drive their TMA-load/store mbarriers so the + // pipeline stays balanced across the column. + + LIGER_CHECK(a.weight_expert_stride >= + static_cast(intermediate_dim) * hidden_dim, + "weight expert stride (", a.weight_expert_stride, + ") must cover at least one [I,D] expert matrix"); + LIGER_CHECK(a.weight_expert_stride % hidden_dim == 0, + "weight expert stride (", a.weight_expert_stride, + ") must contain complete hidden-dimension rows (D=", hidden_dim, ")"); + int weight_expert_stride_rows = + static_cast(a.weight_expert_stride / hidden_dim); + LIGER_CHECK(weight_expert_stride_rows % Traits1::TileN == 0, + "weight expert stride in rows (", weight_expert_stride_rows, + ") must be divisible by MLP1 TileN (", Traits1::TileN, ")"); + int total_n_rows_1 = + (experts_per_pe - 1) * weight_expert_stride_rows + intermediate_dim; + int total_n_rows_2 = experts_per_pe * hidden_dim; + + int num_pes = nvshmem_team_n_pes(team); + int my_pe = nvshmem_team_my_pe(team); + + // Grid — flat 1-D launch. static_nsplit only seeds n_gemm/grid_x before + // the kernel selects the runtime split from the actual routed tile count. + // + // GEMM and comm gate to the first n_gemm = floor_NS CTAs: the comm↔MLP + // staging ring couples producer and consumer (it needs grid_x = n_gemm/static_nsplit + // and MC = n_gemm/NC to be consistent), so they share one count. The + // sort/dispatch/combine grid-stride phases and the cross-CTA pe_barrier use + // ALL num_blocks launched CTAs. The runtime selector may choose a different + // NS inside the kernel and rewrites n_gemm/grid_x consistently for MLP. + int num_sms; + cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, a.device); + int floor_ns = num_sms - num_sms % static_nsplit; + int floor_nc = num_sms - num_sms % Config::kNC; + int n_gemm = floor_ns; // GEMM/comm-active CTAs + int grid_x = n_gemm / static_nsplit; // initial logical GEMM columns + int num_blocks = (floor_ns > floor_nc) ? floor_ns : floor_nc; // launched CTAs + dim3 grid(num_blocks); + + // stream comes from a.stream (the binding passes the current torch stream). + auto& pool = global_buffer_pool(); + + // The two sort outputs the bwd path consumes (token_expert_slots + // [max_total_slots], tile_expert_ids [max_total_slots/kTileM]) are + // caller-owned: the binding pre-allocated them at exactly these sizes and + // passed their device pointers in a. They must NOT alias pool slots that the + // next forward would overwrite. + int* token_expert_slots_ptr = a.token_expert_slots; + int* tile_expert_ids_ptr = a.tile_expert_ids; + + // Allocate all buffers. + auto buf = allocate_moe_buffers(pool, stream, grid_x, num_blocks, + hidden_dim, intermediate_dim, num_tokens, top_k, num_experts, + experts_per_pe, num_pes, + max_total_slots, num_m_tiles_in, + token_expert_slots_ptr, + tile_expert_ids_ptr); + if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: CUDA error before TMA descriptor setup: ", cudaGetErrorString(e)); + + // Final combine output [T, D] — caller-owned (binding-allocated) Y buffer. + Element* output_ptr = reinterpret_cast(a.Y); + Element* y_buf_ptr = buf.y_buf; + + auto& scfg = get_symm_config(); + + // ---- Build MoeParams ------------------------------------------------------- + MoeParams p; + p.expert_indices = a.expert_indices; + p.expert_weights = reinterpret_cast(a.expert_weights); + + p.sort.tile_expert_counts = buf.tile_expert_counts; + p.sort.sorted_token_ids = buf.sorted_token_ids; + p.sort.token_expert_slots = buf.token_expert_slots; + p.sort.expert_offsets = buf.expert_offsets; + p.sort.expert_counts = buf.expert_counts; + p.sort.tile_expert_ids = buf.mlp_tile_expert_ids; + p.sort.cta_done = buf.cta_done; + p.sort.cta_sums = buf.cta_sums; + + p.tokens = reinterpret_cast(a.X); + p.dispatch_dst = buf.x_sorted; + p.combine_output = output_ptr; + + p.comm.src_staging = buf.src_staging; + p.comm.dst_staging = buf.dst_staging; + p.comm.src_ready = buf.src_ready; + p.comm.src_consumed = buf.src_consumed; + p.comm.dst_ready = buf.dst_ready; + p.comm.dst_consumed = buf.dst_consumed; + p.comm.expert_offsets = buf.expert_offsets; + p.comm.expert_counts = buf.expert_counts; + p.comm.sorted_token_ids = buf.sorted_token_ids; + p.comm.local_tokens = buf.x_sorted; + p.comm.local_output = y_buf_ptr; + p.comm.tile_expert_ids = buf.tile_expert_ids; + p.comm.all_expert_offsets = buf.all_expert_offsets; + p.comm.all_expert_counts = buf.all_expert_counts; + p.comm.pe_sync.cta_counter = buf.cta_counter; + p.comm.pe_sync.num_pes = num_pes; + p.comm.pe_sync.my_pe = my_pe; + p.comm.pe_sync.team = team; + p.comm.hidden_dim = hidden_dim; + p.comm.num_experts = num_experts; + p.comm.experts_per_pe = experts_per_pe; + p.comm.n_gemm = n_gemm; // GEMM/comm-active CTA count (floor_NS) + + p.num_tokens = num_tokens; + p.intermediate_dim = intermediate_dim; + p.top_k = top_k; + p.total_slots = max_total_slots; + p.num_m_tiles = num_m_tiles_in; + + // ---- MLP invariant dims (passed as __grid_constant__) ---------------------- + MlpDims mlp_dims; + mlp_dims.hidden_dim = hidden_dim; + mlp_dims.intermediate_dim = intermediate_dim; + mlp_dims.num_experts = experts_per_pe; + mlp_dims.total_n_rows_1 = total_n_rows_1; + mlp_dims.total_n_rows_2 = total_n_rows_2; + mlp_dims.expert_n_stride_1 = weight_expert_stride_rows / Traits1::TileN; + mlp_dims.num_n_tiles_1 = num_n_tiles_1; + mlp_dims.num_n_tiles_2 = num_n_tiles_2; + mlp_dims.num_k_tiles_1 = (Compute == 90) + ? (hidden_dim + Traits1::TileK - 1) / Traits1::TileK + : hidden_dim / Traits1::TileK; + mlp_dims.num_k_tiles_2 = (Compute == 90) + ? (intermediate_dim + Traits2::TileK - 1) / Traits2::TileK + : intermediate_dim / Traits2::TileK; + // Row-tile extent of the symmetric local_output (y_buf) for the + // direct-to-peer Y store (num_y_m_tiles when storing at the token-tile row). + mlp_dims.y_buf_m_tiles = max_total_slots / kTileM; + mlp_dims.phase_counter = buf.phase_counter; + + // ---- MLP TMA descriptors ----------------------------------------------- + // X for MLP is x_sorted (dispatch destination, symmetric). + // All shape/stride values are cast to int64_t — see the long comment + // below the tma_load_c block for the rationale. B and C in particular + // are large: their row count = experts_per_pe · intermediate_dim, + // times hidden_dim → exceeds INT_MAX at production shapes. + auto tma_load_x = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.x_sorted)), + make_shape(static_cast(max_total_slots), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits1::SmemLayoutX_1{}); + + auto* ptr_B = reinterpret_cast(a.all_B); + auto* ptr_C = reinterpret_cast(a.all_C); + // Hopper keeps the expert as an explicit tensor mode. A tail TMA box can + // therefore go OOB only within that expert (loads zero-fill) instead of + // spilling into the next expert's flat rows. + auto tma_load_b = [&]() { + if constexpr (Compute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_B), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + a.weight_expert_stride)), + typename Traits1::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_B), + make_shape(static_cast(total_n_rows_1), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits1::SmemLayoutW_1{}); + } + }(); + auto tma_load_c = [&]() { + if constexpr (Compute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_C), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + a.weight_expert_stride)), + typename Traits1::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_C), + make_shape(static_cast(total_n_rows_1), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits1::SmemLayoutW_1{}); + } + }(); + + // New design's TMA store box = per-WG epi slot (AtomTileM, EpiChunkN). + // + // Shape / stride values are cast to int64_t so CUTE's layout math + // (cosize, byte-offset, total-elements) uses 64-bit arithmetic. The + // Z tensor in particular can exceed INT_MAX elements when + // grid_x · ZBuf · TileM · intermediate_dim > 2^31 — e.g. NSplit=2 + // ZBuf=16 at intermediate_dim=16384 wraps to a negative stride and + // IMAs the kernel. Casting one operand promotes the multiply to + // int64_t; we cast both to be safe across CUTE versions. + auto tma_store_z = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.z_buf), + make_shape(static_cast(buf.z_rows), + static_cast(intermediate_dim)), + make_stride(static_cast(intermediate_dim), Int<1>{})), + typename Traits1::SmemLayoutStoreSlot{}); + + auto tma_load_z = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.z_buf)), + make_shape(static_cast(buf.z_rows), + static_cast(intermediate_dim)), + make_stride(static_cast(intermediate_dim), Int<1>{})), + typename Traits2::SmemLayoutZ_1{}); + + auto* ptr_A = reinterpret_cast(a.all_A); + // Same expert isolation for the down projection, including K-tail loads. + auto tma_load_a = [&]() { + if constexpr (Compute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_A), + make_shape(static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(intermediate_dim), + Int<1>{}, + static_cast(hidden_dim) * intermediate_dim)), + typename Traits2::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_A), + make_shape(static_cast(total_n_rows_2), + static_cast(intermediate_dim)), + make_stride(static_cast(intermediate_dim), Int<1>{})), + typename Traits2::SmemLayoutW_1{}); + } + }(); + + auto tma_store_y = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(y_buf_ptr), + make_shape(static_cast(max_total_slots), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits2::SmemLayoutStore_1{}); + + // Remote staging is a flat list of L = MC · kCommNumStages slots, + // each TileM rows tall. MC = num_blocks / NC (== gridDim.x / NC) — MUST + // match the ring size everywhere else (allocate_moe_buffers, the per-launch + // memset, the consumer's ring_len). Sizing this TMA descriptor on n_gemm + // would leave the wrap-around slots outside its row extent (TMA OOB) when + // num_blocks/NC > n_gemm/NC (num_sms % NSplit != 0). + int mc_remote = num_blocks / Config::kNC; + int remote_total_rows = mc_remote * Config::kCommNumStages * kTileM; + + auto tma_load_x_remote = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.src_staging)), + make_shape(static_cast(remote_total_rows), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits1::SmemLayoutX_1{}); + + auto tma_store_y_remote = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.dst_staging), + make_shape(static_cast(remote_total_rows), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits2::SmemLayoutStore_1{}); + + // ---- GET-via-TMA descriptors (FWD-only, raw CUtensorMap) ----------------- + // SAME-HOST source maps over nvshmem_ptr(x_sorted, peer) + one dst map over + // local src_staging. TMA get is enabled only for same-host NVLink P2P peers + // (cross-host peers fall back to warp getmem); also needs divisible + // hidden_dim + smem budget (checked below). + // + // src_x_desc is indexed by INTRA-HOST local rank (pe % gpus_per_node), so the + // array only ever needs gpus_per_node (≤ kGetMaxPes) entries regardless of + // total PE count. Assumes host-contiguous, host-aligned ranks (true for a + // WORLD-contiguous team); the do_get / RemoteIter is_local predicate uses the + // same arithmetic. + GetTmaDescs get_descs{}; + get_descs.enabled = 0; + int gpus_per_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + LIGER_CHECK(gpus_per_node >= 1 && gpus_per_node <= liger::kGetMaxPes, + "gpus_per_node (", gpus_per_node, ") must be in [1, kGetMaxPes=", + liger::kGetMaxPes, "]"); + get_descs.gpus_per_node = gpus_per_node; + get_descs.my_pe = my_pe; + constexpr int kGetBoxRows = + Config::kTileM / (liger::kNumGetWarpsFwd * Config::kNC); + size_t get_bounce_bytes = (size_t) + liger::GetBounce::kTotalBytes; + { + bool ok = (hidden_dim % liger::kGetKChunk == 0); + for (int tp = 0; ok && tp < num_pes; ++tp) { + // Only same-host peers are P2P-reachable → only they get a TMA + // descriptor (indexed by local rank). Cross-host peers skip it and + // use getmem at runtime. + bool same_host = (tp / gpus_per_node) == (my_pe / gpus_per_node); + if (!same_host) continue; + int gpe = nvshmem_team_translate_pe(team, tp, NVSHMEM_TEAM_WORLD); + void* pp = nvshmem_ptr(buf.x_sorted, gpe); // self → own ptr (non-null) + if (pp == nullptr) { ok = false; break; } // same-host but no P2P → disable TMA + ok = make_get_tma_map(get_descs.src_x_desc[tp % gpus_per_node], pp, + hidden_dim, max_total_slots, kGetBoxRows); + } + if (ok) + ok = make_get_tma_map(get_descs.dst_staging_desc, buf.src_staging, + hidden_dim, remote_total_rows, kGetBoxRows); + get_descs.enabled = ok ? 1 : 0; + if (const char* e = std::getenv("LIGER_GET_TMA")) + if (e[0] == '0') get_descs.enabled = 0; + } + + // ---- Direct-to-peer Y store descriptors (FWD-only) ----------------------- + // Per-peer TMA-store maps over nvshmem_ptr(y_buf, peer) for SAME-HOST peers + // (self included), indexed by intra-host local rank. When a tile is_local, + // MLP2 stores Y straight into the destination's symmetric local_output via + // its peer descriptor (token-tile row), bypassing dst_staging + the put warp + // (which then only rubber-stamps). Same shape/layout as tma_store_y; only the + // base pointer differs (the peer's y_buf). For single-PE the only peer is + // self → desc[my_pe%gpus_per_node] == tma_store_y. + PeerYStoreDescs peer_y{}; + peer_y.enabled = 1; + for (int tp = 0; tp < num_pes; ++tp) { + bool same_host = (tp / gpus_per_node) == (my_pe / gpus_per_node); + if (!same_host) continue; // cross-host peers keep the dst_staging+putmem path + int gpe = nvshmem_team_translate_pe(team, tp, NVSHMEM_TEAM_WORLD); + void* py = nvshmem_ptr(buf.y_buf, gpe); // same-host → P2P ptr (non-null on NVLink) + LIGER_CHECK(py != nullptr, + "direct-to-peer Y store: nvshmem_ptr(y_buf, same-host pe ", tp, + ") is null (no P2P) — unsupported for this build"); + peer_y.desc[tp % gpus_per_node] = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(py)), + make_shape(static_cast(max_total_slots), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits2::SmemLayoutStore_1{}); + } + + // TMA descriptor for the final combine output [num_tokens, hidden_dim]. + // Box shape matches combine_tokens_tma's WG tile: 32 tokens × 256 hidden. + auto* ptr_combine_out = output_ptr; + auto tma_store_y_combine = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(ptr_combine_out), + make_shape(static_cast(num_tokens), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + CombineSmemLayoutSlot{}); + + // ---- Dynamic smem -------------------------------------------------------- + int sort_num_warps = kNumThreads / 32; + size_t smem_sort = sort_num_warps * num_experts * sizeof(int); + + size_t smem_mlp = sizeof(MoeSmem); + smem_mlp = (smem_mlp + sizeof(int) - 1) & ~(sizeof(int) - 1); + smem_mlp += num_pes * (experts_per_pe + 1) * sizeof(int); // remote_offsets + smem_mlp += num_pes * experts_per_pe * sizeof(int); // remote_counts + + // GET-via-TMA bounce: append a 128-aligned region (must match the kernel's + // get_bounce offset). Disable the TMA-get path if it would exceed the + // opt-in dynamic-smem cap (227 KiB) — that config just keeps warp getmem. + if (get_descs.enabled) { + size_t mlp_with_bounce = ((smem_mlp + 127) & ~((size_t)127)) + get_bounce_bytes; + if (mlp_with_bounce <= (size_t)227 * 1024) smem_mlp = mlp_with_bounce; + else get_descs.enabled = 0; + } + + // Combine phase reinterprets the same SMEM region as CombineSmem. + constexpr int kCombineNumWGs = + Config::kNumThreads / (kCombineWarpsPerWG * kCombineWarpSize); + size_t smem_combine = sizeof(CombineSmem); + + size_t smem_size = std::max({smem_sort, smem_mlp, smem_combine}); + + // ---- Launch fused kernel ----------------------------------------------- + auto kernel_fn = moe_fused_kernel; + + if (cudaError_t e = cudaFuncSetAttribute(kernel_fn, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaFuncSetAttribute(MaxDynamicSharedMemorySize=", smem_size, + ") failed: ", cudaGetErrorString(e)); + + // Reset per-call buffers in a single fused-memset kernel — one CTA per + // region. Replaces 11 cudaMemsetAsync calls that would otherwise + // serialize on the stream. + // + // Notes: + // - dispatch_done / mlp_done are already reset upstream in + // allocate_moe_buffers (the original "monotonic, never reset" design + // was retired — wait_until could exit on stale residual signal + // values in multi-PE shape-transition workloads). + // - all_expert_offsets stays unreset (overwritten by + // broadcast_expert_offsets on every call). + // - y_buf stays unreset (slots referenced by token_expert_slots[t,k] + // are always written by local MLP or the source PE's NVSHMEM + // put-back; padding slots are never read by combine_tokens_tma). + // - src_staging / dst_staging / z_buf intentionally NOT zero'd (see + // comment block below). + // + // Per-region reasons: + // - cta_counter: PeSync.barrier() does `% num_ctas`. Residual values + // from a prior call with different num_ctas misalign to the new + // modulus, hanging barrier() forever. + // - tile_expert_ids (0xff): sentinel-fill so unfetched slots produce + // expert_id = -1; matches moe_bwd.cu's 0xff fill of + // tile_expert_ids_{x,dy}. Without this, stale slot values from prior + // calls can match a valid expert id and confuse the MLP-side remote + // iterator at higher PE counts. + // - src/dst_ready/consumed: atomically accumulated; stale ready/ + // consumed values make the consumer-acquire wait fall through. + // - tile_expert_counts / cta_sums: atomicAdd'd during sort; must + // start at 0. + { + auto& scfg = get_symm_config(); + // L = MC · kCommNumStages, MC = num_blocks / NC (== gridDim.x / NC). + // MUST match the ring size in allocate_moe_buffers and the device-side + // comm/consumer stride — otherwise this memset under-covers the signal / + // tile_expert_ids arrays (slots [n_gemm/NC·CS, gridDim.x/NC·CS) left + // uninitialized) and the consumer reads stale src_ready / garbage expert + // ids for those slots (see the OOB note in allocate_moe_buffers). + int mc_local = num_blocks / Config::kNC; + int comm_l = mc_local * Config::kCommNumStages; + size_t sig_bytes = + static_cast(comm_l) * Config::kSubTiles * sizeof(int); + size_t eid_bytes = static_cast(comm_l) * sizeof(int); + int max_m_tiles = scfg.max_total_slots / kTileM; + int max_tile_counts = max_m_tiles * scfg.max_num_experts; + + FusedMemsetArgs ma{}; + fused_memset_add(ma, buf.phase_counter, + ((num_blocks + 1) / 2) * sizeof(int), 0x00); + fused_memset_add(ma, buf.cta_counter, + sizeof(int), 0x00); + fused_memset_add(ma, buf.cta_done, + num_blocks * sizeof(int), 0x00); + fused_memset_add(ma, buf.sorted_token_ids, + static_cast(max_total_slots) * sizeof(int), 0xff); + fused_memset_add(ma, buf.src_ready, sig_bytes, 0x00); + fused_memset_add(ma, buf.src_consumed, sig_bytes, 0x00); + fused_memset_add(ma, buf.dst_ready, sig_bytes, 0x00); + fused_memset_add(ma, buf.dst_consumed, sig_bytes, 0x00); + fused_memset_add(ma, buf.tile_expert_ids, eid_bytes, 0xff); + fused_memset_add(ma, buf.tile_expert_counts, + static_cast(max_tile_counts) * sizeof(int), 0x00); + fused_memset_add(ma, buf.cta_sums, + static_cast(num_blocks) * scfg.max_num_experts * sizeof(int), 0x00); + fused_memset_launch(ma, stream); + } + + // src_staging / dst_staging / z_buf intentionally NOT zero'd here: + // + // src_staging: comm warps fill the slots THIS CTA fetches; unfetched + // slots are gated by tile_expert_ids = -1 (set via the 0xff memset + // above) and RemoteIter::has_next() walks only per_cta_tiles, so the + // remote MLP never reads unfetched slots. Stale prev-call contents + // are inert under both bounds checks. + // + // dst_staging: write-then-read for routed slots only. Remote MLP + // writes Y for tiles RemoteIter returns; comm warps put back only + // those tiles. Padding/unused slots are never read at either end. + // + // z_buf: per-tile write→read within a single MLP launch. LocalIter is + // range-bounded to [local_e_start, local_e_end) experts via + // set_range; the remote pass is bounded by per_cta_tiles. Tiles + // outside the current call's active range are never written and + // never read — buffer-pool size growth from previous (larger) calls + // is invisible to the kernel. + // + // The tile_expert_ids 0xff fill above is what makes this safe — keep + // that. + + kernel_fn<<>>( + tma_load_x, tma_load_b, tma_load_c, tma_store_z, + tma_load_z, tma_load_a, tma_store_y, + tma_load_x_remote, tma_store_y_remote, + tma_store_y_combine, + mlp_dims, p, static_nsplit, get_descs, peer_y); + if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: kernel launch failed for grid=", grid.x, + " threads=", kNumThreads, " smem=", smem_size, ": ", + cudaGetErrorString(e)); + + // Post-launch global barrier: ensures all NVSHMEM operations (tile puts, + // signal ops) across all PEs are complete before any PE starts the next + // call. Without this, a fast PE could overwrite symmetric buffers + // (x_sorted, y_buf, tile signals) while a slow PE is still reading them. + // + // Stream-ordered variant — capturable into a CUDA graph. The host-side + // nvshmem_team_sync(team) would do an internal cudaStreamSynchronize and + // therefore fail with "operation not permitted when stream is capturing". + //if (num_pes > 1) { + // nvshmemx_barrier_on_stream(team, stream); + //} + + // Hand the persistent symmetric buffers back to the ABI wrapper so it can + // alias them as the *_out_symm views (shapes use the static MoeSymmConfig + // max; the views do NOT own the storage — SymmetricMemoryStack does). Y and + // the two sort outputs were written in place into the caller-owned buffers. + *a.x_sorted_out = buf.x_sorted; + *a.y_buf_out = y_buf_ptr; + *a.all_expert_offsets_out = buf.all_expert_offsets; + if (a.all_expert_counts_out != nullptr) + *a.all_expert_counts_out = buf.all_expert_counts; +} + +// ── Runtime dispatch table for tuned-config lookup ──────────────────── + +using MoeFwdFn = void (*)(const MoeFwdArgs&, int static_nsplit); +static constexpr int kFwdRuntimeNsplitSeed = 8; + +struct DispatchEntry { + int Compute; + int TileN1, TileK1, Stages1, EpiChunkN1; + int TileN2, TileK2, Stages2, EpiChunkN2; + int ZBufferSlots, CommNumStages; + int TileM; // tuned GEMM-tile field (legacy dispatch-table name) + int GemmTileM; // GEMM tile ∈ {64,128} + MoeFwdFn fn; +}; + +// 13-arg core: explicit GemmTileM. 12-arg shim sets GemmTileM = comm TileM. +#define LIGER_MOE_DISPATCH_ENTRY_G_C(Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) \ + { Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM, \ + &moe_fused_fwd_bf16 }, +#define LIGER_MOE_DISPATCH_ENTRY_C(Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) \ + LIGER_MOE_DISPATCH_ENTRY_G_C(Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, TileM) +#define LIGER_MOE_DISPATCH_ENTRY_G_SM90(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) \ + LIGER_MOE_DISPATCH_ENTRY_G_C(90, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) +#define LIGER_MOE_DISPATCH_ENTRY_SM90(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) \ + LIGER_MOE_DISPATCH_ENTRY_C(90, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) +#define LIGER_MOE_DISPATCH_ENTRY_G_SM100(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) \ + LIGER_MOE_DISPATCH_ENTRY_G_C(100, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) +#define LIGER_MOE_DISPATCH_ENTRY_SM100(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) \ + LIGER_MOE_DISPATCH_ENTRY_C(100, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) + +static const DispatchEntry kDispatchTable[] = { +#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 90 + LIGER_MOE_FWD_DISPATCH_CONFIGS_SM90(LIGER_MOE_DISPATCH_ENTRY_SM90, LIGER_MOE_DISPATCH_ENTRY_G_SM90) +#endif +#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 100 + LIGER_MOE_FWD_DISPATCH_CONFIGS_SM100(LIGER_MOE_DISPATCH_ENTRY_SM100, LIGER_MOE_DISPATCH_ENTRY_G_SM100) +#endif +}; +#undef LIGER_MOE_DISPATCH_ENTRY_SM100 +#undef LIGER_MOE_DISPATCH_ENTRY_G_SM100 +#undef LIGER_MOE_DISPATCH_ENTRY_SM90 +#undef LIGER_MOE_DISPATCH_ENTRY_G_SM90 +#undef LIGER_MOE_DISPATCH_ENTRY_C +#undef LIGER_MOE_DISPATCH_ENTRY_G_C +static constexpr int kNumDispatchEntries = + sizeof(kDispatchTable) / sizeof(kDispatchTable[0]); + +} // namespace liger + +// Joint fwd+bwd tuning table (auto-generated by benchmarks/tune_moe_fwd_bwd). +// We only consume the Fwd_* fields here; moe_bwd.cu consumes the same table's +// Bwd_* fields. Picking fwd and bwd from a single table guarantees the two +// directions are tuned as a pair (so their TileM agrees and their per-shape +// templates were measured together, not independently). +#include "moe_fwd_bwd_tuning_configs_sm100.cuh" + +namespace liger { + +// TunedConfig is the fwd-only projection of TunedConfigFwdBwd. The rest of +// the dispatch code (validity check, dispatch-entry match) was written +// against this struct and is reused unchanged. +struct TunedConfig { + int Compute; + int TK, TKE, D, I; + int TileN1, TileK1, Stages1, EpiChunkN1; + int TileN2, TileK2, Stages2, EpiChunkN2; + int ZBufferSlots, CommNumStages; + int TileM; // comm tile + int GemmTileM; // gemm tile (auto path: == TileM; FORCE_CONFIG may set 64) +}; + +// Project a TunedConfigFwdBwd row down to its fwd template params. The tuned +// table has no GemmTileM column, so the auto path keeps gemm == comm (TileM). +static TunedConfig project_fwd(const TunedConfigFwdBwd& c, int compute) { + return TunedConfig{ + compute, + c.TK, c.TKE, c.D, c.I, + c.Fwd_TileN1, c.Fwd_TileK1, c.Fwd_Stages1, c.Fwd_EpiChunkN1, + c.Fwd_TileN2, c.Fwd_TileK2, c.Fwd_Stages2, c.Fwd_EpiChunkN2, + c.Fwd_ZBufferSlots, c.Fwd_CommNumStages, + c.Fwd_TileM, c.Fwd_TileM, + }; +} + +// Find the dispatch table entry matching a TunedConfig's fields. +static const DispatchEntry* find_dispatch_entry(const TunedConfig& c) { + for (int i = 0; i < kNumDispatchEntries; ++i) { + const auto& e = kDispatchTable[i]; + if (e.Compute == c.Compute && + e.TileN1 == c.TileN1 && e.TileK1 == c.TileK1 && e.Stages1 == c.Stages1 && + e.EpiChunkN1 == c.EpiChunkN1 && + e.TileN2 == c.TileN2 && e.TileK2 == c.TileK2 && e.Stages2 == c.Stages2 && + e.EpiChunkN2 == c.EpiChunkN2 && + e.ZBufferSlots == c.ZBufferSlots && e.CommNumStages == c.CommNumStages && + e.TileM == c.TileM && e.GemmTileM == c.GemmTileM) { + return &e; + } + } + return nullptr; +} + +// Mirrors shape_valid() in tune_moe.cu: structural constraints the kernel +// requires on (D, I) for a given config. Configs that violate these for +// the runtime shape are skipped during nearest-neighbor lookup so we +// don't dispatch to a template instantiation the kernel can't run. +// +// Cooperative 2-WG consumer layout switches on TileM: +// - TileM=128 (M-split, Layout<_2,_1,_1>): each WG owns AtomTileM=64 +// rows × TileN cols. EpiChunkN must divide TileN. +// - TileM=64 (N-split, Layout<_1,_2,_1>): each WG owns TileM=64 rows +// × TileN/2 cols. EpiChunkN must divide TileN/2; TileN must be even. +// - Hopper uses ceil-div tile counts; SM100 retains exact divisibility. +// - store_buf is (AtomTileM=64, EpiChunkN) per WG in both modes. +static bool tuned_config_valid(int D, int I, const TunedConfig& c) { + if (c.TileM != 64 && c.TileM != 128) return false; + // The current Blackwell UMMA MLP1 consumer only supports the 128-row comm + // tile. Keep mirrored SM90 TileM=64 tuning rows out of SM100 auto-dispatch + // until that TMEM epilogue layout is implemented. + if (c.Compute == 100 && c.TileM != 128) return false; + if (c.Compute == 100) { + if (D % c.TileK1 != 0) return false; + if (I % c.TileN1 != 0) return false; + if (I % c.TileK2 != 0) return false; + if (D % c.TileN2 != 0) return false; + } + if (D % 8 != 0 || I % 8 != 0) return false; + // NC selection mirrors moe_fused_fwd_bf16: TileM=128 → NC=4, + // TileM=64 → NC=2. With the ticket-based GEMM iterator the only + // runtime requirement is NC | (gridDim.x · gridDim.y); runtime NS / + // CommNumStages no longer need to share any divisibility with NC. + const int nc = (c.TileM == 128) ? 4 : 2; + (void)nc; + // EpiChunkN must divide the per-WG N extent (WgTileN). + const int wg_tile_n1 = (c.TileM == 128) ? c.TileN1 : c.TileN1 / 2; + const int wg_tile_n2 = (c.TileM == 128) ? c.TileN2 : c.TileN2 / 2; + if (wg_tile_n1 <= 0 || wg_tile_n1 % c.EpiChunkN1 != 0) return false; + if (wg_tile_n2 <= 0 || wg_tile_n2 % c.EpiChunkN2 != 0) return false; + // N-split requires TileN to be evenly halvable (TileN/2 must be ≥ the + // WGMMA atom-N min, which is 8 for bf16 but 64 for our smem layout). + if (c.TileM == 64 && (c.TileN1 % 2 != 0 || c.TileN2 % 2 != 0)) return false; + + // SMEM cap (H100 = 228 KiB; budget 224 KiB inside MoeFusedConfig). + // Per-CTA, bf16 (2 B/elem). TileM scales the X/Z staging: + // MLP1 = 2 * (TileM·TK1·S1 + 2·TN1·TK1·S1 + 2·64·EC1) + // MLP2 = 2 * (TileM·TK2·S2 + TN2·TK2·S2 + 2·64·EC2) + // smem ≈ max(MLP1, MLP2) + int mlp1_bytes = 2 * (c.TileM * c.TileK1 * c.Stages1 + + 2 * c.TileN1 * c.TileK1 * c.Stages1 + + 2 * 64 * c.EpiChunkN1); + int mlp2_bytes = 2 * (c.TileM * c.TileK2 * c.Stages2 + + c.TileN2 * c.TileK2 * c.Stages2 + + 2 * 64 * c.EpiChunkN2); + int smem = mlp1_bytes > mlp2_bytes ? mlp1_bytes : mlp2_bytes; + constexpr int kSmemBudget = 224 * 1024; + if (smem > kSmemBudget) return false; + return true; +} + +static bool find_tuned_table(int compute, int n_pes, + const TunedConfigFwdBwd** table, + int* table_n) { + const TunedConfigFwdBwdTable* tables = + (n_pes <= 1) ? kTunedConfigTablesSingle : kTunedConfigTablesMulti; + const int tables_n = + (n_pes <= 1) ? kNumTunedConfigTablesSingle : kNumTunedConfigTablesMulti; + for (int i = 0; i < tables_n; ++i) { + if (tables[i].Compute == compute) { + *table = tables[i].configs; + *table_n = tables[i].count; + return *table_n > 0; + } + } + return false; +} + +// Nearest-neighbor lookup in log-space (TK, TKE, D, I), where TK = T*top_k +// is the total routed-token count (the actual M-axis the kernel sees) and +// TKE = TK / E_local (E_local = E / n_pes) is the per-LOCAL-expert avg +// K-range. Adding TKE to the distance lets the dispatcher distinguish +// shapes whose mlp3/mlp4 work distribution differs even when (TK, D, I) +// is the same — e.g. same total tokens, different expert count or PE +// count. +// Configs that violate structural constraints for the runtime (D, I) are +// skipped — the nearest tuned point is not always valid for the actual +// shape (e.g. tuned at I=4096 with TileN1=256, runtime I=768). +// Out-param API instead of pointer-into-static-array because each candidate +// TunedConfig is projected on the fly from kTunedConfigsFwdBwd (no stable +// address). Returns true when a valid config was found. +static bool find_nearest_tuned(int compute, int TK, int TKE, int D, int I, int n_pes, + TunedConfig* out) { + // World-size class: single-GPU (n_pes <= 1) and multi-GPU tables are tuned + // and stored separately, since the comm path (and thus the optimal NSplit/ + // CommNumStages/GEMM-tile) differs sharply between one PE and many. + const TunedConfigFwdBwd* table = nullptr; + int table_n = 0; + if (!find_tuned_table(compute, n_pes, &table, &table_n)) return false; + bool found = false; + double best_dist = 1e30; + double lTK = std::log((double)TK); + double lTKE = std::log((double)TKE); + double lD = std::log((double)D); + double lI = std::log((double)I); + for (int i = 0; i < table_n; ++i) { + TunedConfig c = project_fwd(table[i], compute); + if (!tuned_config_valid(D, I, c)) continue; + // Guard against TKE=0 in legacy rows (would yield -inf in log). + double dTK = std::log((double)c.TK) - lTK; + double dTKE = c.TKE > 0 + ? std::log((double)c.TKE) - lTKE + : 0.0; // legacy / placeholder rows contribute 0 to the TKE axis + double dD = std::log((double)c.D) - lD; + double dI = std::log((double)c.I) - lI; + double d = dTK*dTK + dTKE*dTKE + dD*dD + dI*dI; + if (d < best_dist) { best_dist = d; *out = c; found = true; } + } + return found; +} + +// ============================================================================ +// Auto-dispatch wrapper — pure tuned-table lookup +// ============================================================================ +// +// Picks the nearest tuned config (log-space distance over (T, D, I)) from +// moe_fwd_bwd_tuning_configs.cuh and dispatches to the matching template +// instantiation in kDispatchTable. Both tables are populated by the +// tune_moe_fwd_bwd sweep — +// shapes outside the swept range error out instead of falling back to a +// heuristic. Add the missing shape to the sweep (or to LIGER_MOE_TUNE_CONFIGS +// directly) to extend coverage. + +// Internal auto-dispatch: picks the nearest tuned config for the runtime shape +// (or honours LIGER_MOE_FORCE_CONFIG) and invokes the matching template +// instantiation via the dispatch table. Writes the chosen comm TileM to +// *chosen_tile_m. All shape info comes from the MoeFwdArgs bundle; the launcher +// writes the outputs in place. +void moe_fused_fwd_dispatch(const MoeFwdArgs& a, int* chosen_tile_m) { + const int T = a.num_tokens; + const int D = a.hidden_dim; + const int I = a.intermediate_dim; + const int top_k = a.top_k; + const int num_experts = a.num_experts; + const int TK = T * top_k; // total routed tokens — the M-axis the kernel sees + // TKE = TK / E_local where E_local = num_experts / n_pes (per-LOCAL-expert + // avg K-range). Matches the tuner's definition so the nearest-neighbor + // lookup picks rows recorded at comparable per-PE load. + const int n_pes = nvshmem_team_n_pes(a.team); + const int E_local = (n_pes > 0) ? std::max(1, num_experts / n_pes) : num_experts; + // The lookup operates in log space, so clamp the per-expert average to + // the smallest representable workload for tiny but valid token batches. + const int TKE = std::max(1, TK / E_local); + const int compute = moe_detect_compute_dispatch_key(a.device, "moe_fused_fwd_bf16_auto"); + + // ── Override path: LIGER_MOE_FORCE_CONFIG=TN1,TK1,S1,EC1,TN2,TK2,S2,EC2,ZB,CS,TM[,GemmTM] + // Lets a test suite sweep every compiled config without rebuilding. The + // optional final int is GemmTileM (comm/gemm tile decouple); omitted → == TM. + if (const char* s = std::getenv("LIGER_MOE_FORCE_CONFIG")) { + TunedConfig tc{}; + tc.Compute = compute; + tc.TileM = 128; // default when omitted (backward compat) + tc.GemmTileM = -1; // sentinel → default to TileM after parse + int n = std::sscanf(s, "%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d", + &tc.TileN1, &tc.TileK1, &tc.Stages1, &tc.EpiChunkN1, + &tc.TileN2, &tc.TileK2, &tc.Stages2, &tc.EpiChunkN2, + &tc.ZBufferSlots, &tc.CommNumStages, + &tc.TileM, &tc.GemmTileM); + LIGER_CHECK(n == 10 || n == 11 || n == 12, + "LIGER_MOE_FORCE_CONFIG must have 10, 11 or 12 comma-separated ints " + "(TN1,TK1,S1,EC1,TN2,TK2,S2,EC2,ZB,CS[,TM[,GemmTM]]); got '", s, "'"); + if (tc.GemmTileM <= 0) tc.GemmTileM = tc.TileM; // omitted → gemm == comm + const DispatchEntry* de = find_dispatch_entry(tc); + LIGER_CHECK(de != nullptr, + "LIGER_MOE_FORCE_CONFIG=", s, " is not in the compiled dispatch " + "table. Add it to LIGER_MOE_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); + de->fn(a, kFwdRuntimeNsplitSeed); + *chosen_tile_m = de->Compute == 100 + ? MoeCommGeometry<100>::TileM + : MoeCommGeometry<90>::TileM; + return; + } + + TunedConfig tc{}; + LIGER_CHECK(find_nearest_tuned(compute, TK, TKE, D, I, n_pes, &tc), + "moe_fused_fwd_bf16_auto: no valid tuned config for compute=", compute, + " shape (T=", T, + " K=", top_k, " TK=", TK, " TKE=", TKE, " D=", D, " I=", I, + "). Either moe_fwd_bwd_tuning_configs.cuh is empty (run tune_moe_fwd_bwd), " + "or every tuned entry violates structural constraints for this (D, I) " + "(tile/layout compatibility and shared-memory limits). " + "Add a compatible config to LIGER_MOE_TUNE_CONFIGS and re-tune."); + + const DispatchEntry* de = find_dispatch_entry(tc); + LIGER_CHECK(de != nullptr, + "moe_fused_fwd_bf16_auto: tuned config (Compute=", tc.Compute, + " TN1=", tc.TileN1, "/", tc.TileK1, "/", tc.Stages1, "/EC", tc.EpiChunkN1, + " TN2=", tc.TileN2, "/", tc.TileK2, "/", tc.Stages2, "/EC", tc.EpiChunkN2, + " ZB=", tc.ZBufferSlots, " CS=", tc.CommNumStages, + ") for shape (T=", T, " K=", top_k, " TK=", TK, " D=", D, " I=", I, + ") is not in the compiled dispatch table. " + "Add it to LIGER_MOE_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); + + de->fn(a, kFwdRuntimeNsplitSeed); + *chosen_tile_m = de->Compute == 100 + ? MoeCommGeometry<100>::TileM + : MoeCommGeometry<90>::TileM; +} + +// ───────────────────────────────────────────────────────────────────── +// Default (untuned) wrapper — single hardcoded config per TileM. Use +// this for the regular bench_moe / test_moe paths so they don't depend +// on the autotuner dispatch table. Returns the same 6-tuple as the auto +// wrapper. +// +// Each config below must have a matching explicit instantiation. The +// TileM=64 row is in LIGER_MOE_TUNE_CONFIGS (so LIGER_MOE_INSTANTIATE +// emits it automatically); the TileM=128 row is emitted as a direct +// `template` instantiation right below this function. +// ───────────────────────────────────────────────────────────────────── + +// Returns the compiled MoE config table as a list of 12-int rows +// (Compute, TileN1, TileK1, Stages1, EpiChunkN1, TileN2, TileK2, Stages2, +// EpiChunkN2, ZBufferSlots, CommNumStages, TileM). Kept for the tuner / a future +// sweep binding so they stay in sync with the dispatch table without manual +// duplication. Internal (std::vector) — not part of the flat ABI. +std::vector> +moe_list_compiled_configs() { + std::vector> out; + out.reserve(kNumDispatchEntries); + for (int i = 0; i < kNumDispatchEntries; ++i) { + const auto& e = kDispatchTable[i]; + out.emplace_back(e.Compute, e.TileN1, e.TileK1, e.Stages1, e.EpiChunkN1, + e.TileN2, e.TileK2, e.Stages2, e.EpiChunkN2, + e.ZBufferSlots, e.CommNumStages, e.TileM); + } + return out; +} + +} // namespace liger + +// ============================================================================ +// Flat extern "C" control ABI (declared in liger_cute/moe.h) +// ============================================================================ + +namespace { + +using liger_cute::detail::guarded; + +} // namespace + +extern "C" { + +liger_cute_status_t liger_cute_moe_get_symm_config(liger_cute_moe_symm_config_t* out) { + return guarded([&]() -> liger_cute_status_t { + LIGER_CHECK(out != nullptr, "moe_get_symm_config: out is null"); + const liger::MoeSymmConfig& cfg = liger::get_symm_config(); + out->max_total_slots = cfg.max_total_slots; + out->max_num_experts = cfg.max_num_experts; + out->hidden_dim = cfg.hidden_dim; + out->num_pes = cfg.num_pes; + out->experts_per_pe = cfg.experts_per_pe; + out->max_top_k = cfg.max_top_k; + out->initialized = cfg.initialized ? 1 : 0; + return LIGER_CUTE_OK; + }); +} + +liger_cute_status_t liger_cute_moe_configure_symmetric( + const int max_tokens, const int hidden_dim, const int max_num_experts, + const int max_top_k, const int num_pes, const int num_hosts, + const int gpus_per_host) { + return guarded([&]() -> liger_cute_status_t { + liger::moe_configure_symmetric(max_tokens, hidden_dim, max_num_experts, + max_top_k, num_pes, num_hosts, gpus_per_host); + return LIGER_CUTE_OK; + }); +} + +liger_cute_status_t liger_cute_moe_pop_fwd(void) { + return guarded([&]() -> liger_cute_status_t { + liger::moe_pop_fwd(); + return LIGER_CUTE_OK; + }); +} + +} // extern "C" diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_sm90.cuh new file mode 100644 index 00000000..81163fe7 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_sm90.cuh @@ -0,0 +1,200 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Fused MoE: MLP compute for local + remote tiles +// ═══════════════════════════════════════════════════════════════════ +// +// MLP-only device function — called by MLP warps (0,1,4-11). +// Comm warps (2-3) are handled separately at the kernel level to +// keep CommBuffers out of the MLP register scope. +// +// Execution flow (when remote tiles exist): +// a. mlp_fused_fwd with LocalMlpTileIterator (first half of column) +// b. NamedBarrier (MLP warps only) +// c. mlp_fused_fwd with RemoteMlpTileIterator +// d. NamedBarrier (MLP warps only) +// e. mlp_fused_fwd with LocalMlpTileIterator (second half of column) +// +// Splitting the local pass around the remote pass lets the comm warps +// overlap remote gets with the first local half and remote puts with +// the second local half. When there are no remote tiles, the kernel +// runs a single full local pass and returns. +// +// For remote tiles: +// - X input comes from src_staging (filled by comm get) +// - Y output goes to dst_staging (drained by comm put) +// - The RemoteMlpTileIterator acquire/release hooks synchronize +// with comm warps via StagePipe +// ═══════════════════════════════════════════════════════════════════ + +#include "mlp_sm90.cuh" + +namespace liger { + +using namespace cute; + +// ═══════════════════════════════════════════════════════════════════ +// SymmetricMemoryStack keys for the persistent fwd→bwd buffers. +// Centralized so set_size / put / pop / top sites all agree. +// ═══════════════════════════════════════════════════════════════════ +inline constexpr const char* kSymmKeyXSorted = "x_sorted"; +inline constexpr const char* kSymmKeyYBuf = "y_buf"; +inline constexpr const char* kSymmKeyAllExpertOffsets = "all_expert_offsets"; +inline constexpr const char* kSymmKeyAllExpertCounts = "all_expert_counts"; + +// ═══════════════════════════════════════════════════════════════════ +// MoE shared memory — union of MLP smem + comm smem +// ═══════════════════════════════════════════════════════════════════ + +template +struct MoeSmem { + // MLP shared memory (fused phase 1 + 2 with pipeline barriers). + MlpFusedSmem mlp; + + // Comm shared memory (total_tiles count + remote_offsets pointer). + CommSmem comm; +}; + +// ═══════════════════════════════════════════════════════════════════ +// moe_fused_fwd — MLP warps only (no CommBuffers in scope) +// ═══════════════════════════════════════════════════════════════════ + +template 1 means the GEMM steps + // through SubTiles sub-tiles of Traits1::TileM (=GemmTileM) rows per + // communication staging slot. 1 → unchanged single-tile path. + int SubTiles, + int Compute = 90, + typename LocalIter, typename RemoteIter, + typename TmaLoadX, typename TmaLoadW1, typename TmaStoreZ, + typename TmaLoadZ, typename TmaLoadW2, typename TmaStoreY> +__device__ __forceinline__ void moe_fused_fwd( + MoeSmem& smem, + LocalIter& local_iter, + // Remote-phase comm pointers (no CommBuffers — keeps MLP register scope clean) + int* src_ready, + int* src_consumed, + int* dst_ready, + int* dst_consumed, + // RemoteIter re-derives (expert, pe) and the per-tile is_local fence-scope + // predicate from remote_offsets (SMEM) using these — no per-slot HBM. + int experts_per_pe, + int num_pes, + int my_pe, + int gpus_per_node, + // MLP1 TMA descriptors + TmaLoadX const& tma_load_x, + TmaLoadW1 const& tma_load_b, + TmaLoadW1 const& tma_load_c, + TmaStoreZ const& tma_store_z, + // MLP2 TMA descriptors + TmaLoadZ const& tma_load_z, + TmaLoadW2 const& tma_load_a, + TmaStoreY const& tma_store_y, + // Remote-phase TMA descriptors (X loaded from staging, Y stored to staging) + TmaLoadX const& tma_load_x_remote, + TmaStoreY const& tma_store_y_remote, + // Invariant MLP dims — const ref into __grid_constant__ .param space + const MlpDims& dims, + // Per-phase params (change between local and remote) + int num_tokens, + int num_m_tiles, + // Flat-grid logical coordinates: col = flat_id / NSplit (this CTA's + // GEMM column), grid_x = n_gemm / NSplit (number of columns). Under + // the flat launch gridDim.x = num_blocks and gridDim.y = 1, so these + // replace every blockIdx.x / gridDim.x / gridDim.y read below. + int col, + int grid_x, + int split, // this CTA's N-split index = flat_id % runtime_nsplit + int runtime_nsplit, + int local_expert_start = 0, + // Direct-to-peer Y store: per-peer TMA descriptors over + // nvshmem_ptr(local_output, peer). Gated per-tile on is_local downstream. + const TmaStoreY* y_peer_desc = nullptr) { + + int warp_id = threadIdx.x / Traits1::WarpSize; + + // Construct the per-CTA x_barrier ONCE here so its target counter + // accumulates monotonically across all 3 mlp_fused_fwd calls + // (A1, B, A2). Each phase invokes mlp_fused_fwd which calls + // x_barrier.wait() inside its tile loop — they must share the + // same barrier object so target stays aligned with the shared + // device counter. Constructing the barrier inside mlp_fused_fwd + // would reset target=0 per call while the counter keeps growing, + // and the second mlp_fused_fwd would see counter > target on + // its very first poll and exit without actually waiting. + MlpFwdCtaBarrierT x_barrier(&dims.phase_counter[col], + runtime_nsplit); + + // UNIFIED PATH: the local MLP pass has been removed. The comm-side + // TileIterator now enumerates LOCAL experts (p=0) alongside remote ones, + // and the get warp stages local X from local symmetric memory (do_get + // is_local TMA path) exactly like remote tiles. So EVERY tile this column + // owns — local or remote — flows through the RemoteMlpTileIterator / + // src_staging→MLP→dst_staging pipeline below. local_iter is unused now. + // smem.comm.per_cta_tiles is this column's total tile count (stride + // grid_x); a column with no tiles returns immediately. + (void)local_iter; (void)tma_load_x; (void)tma_store_y; (void)num_tokens; + (void)num_m_tiles; (void)local_expert_start; + int total_remote = smem.comm.per_cta_tiles; + if (total_remote == 0) return; + + bool is_leader = (warp_id == 0 && (threadIdx.x % Traits1::WarpSize) == 0); + + // Build RemoteMlpTileIterator. The staging is a flat ring of + // L = MC · CommNumStages slots, where MC = (gridDim.x · gridDim.y) / NC, + // indexed by + // slot = (blockIdx.x + j · gridDim.x) mod L (MLP, ticket-based) + // slot = xc + (j mod CommNumStages) · MC + // (comm, NumStages cycle). + // Both sides resolve to the same slot for a given global tile T and + // cover the full L. No divisibility relationship is required between + // NC, NSplit and CommNumStages (runtime: NC must divide gridDim.x · + // gridDim.y so MC is integer). + // n_gemm = grid_x · runtime_nsplit — the GEMM/comm-active CTA count. + int n_gemm = grid_x * runtime_nsplit; + RemoteIter remote_iter; + remote_iter.init( + // remote_offsets in SMEM drives the embedded walker (null → no derivation). + smem.comm.remote_offsets, smem.comm.remote_counts, + experts_per_pe, num_pes, my_pe, gpus_per_node, + smem.comm.per_cta_tiles, + col, + n_gemm, + src_ready, src_consumed, + dst_ready, dst_consumed, + is_leader, + runtime_nsplit); + + // Remote MLP dimensions: TMA descriptors cover the full flat staging + // list = MC · CommNumStages · TileM rows. MC = gridDim.x / NC — MUST match + // the ring size in allocate_moe_buffers and the consumer's ring_len + // (RemoteMlpTileIterator mc = gridDim.x/NC); sizing on n_gemm here would + // leave the wrap-around slots [n_gemm/NC·CS, gridDim.x/NC·CS) outside the + // staging coverage (OOB) when gridDim.x/NC > n_gemm/NC. + int mc = gridDim.x / NC; + int remote_num_m_tiles = mc * CommNumStages; + + // NOTE: phase_counter is monotonic across phase A and phase B — + // the modulo-gridDim.y check in mlp_fused_fwd is invariant + // under any shared starting value. The host zeroes phase_counter + // before each kernel launch (see moe.cu); resetting it here would + // race across the gridDim.y CTAs in the column. + + // num_tokens = total staging rows = remote_num_m_tiles COMM tiles × CommTileM. + // CommTileM = Traits1::TileM (=GemmTileM) × SubTiles. num_m_tiles passed is the + // COMM tile count; mlp_fused_fwd scales it by SubTiles for the GEMM sub-tiles. + mlp_fused_fwd( + smem.mlp, remote_iter, + tma_load_x_remote, tma_load_b, tma_load_c, tma_store_z, + tma_load_z, tma_load_a, tma_store_y_remote, + dims, + remote_num_m_tiles * Traits1::TileM * SubTiles, remote_num_m_tiles, x_barrier, + /*local_expert_start=*/0, col, grid_x, split, runtime_nsplit, + /*y_peer_desc=*/y_peer_desc); +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_sm90.inc b/liger_cute_kernels/csrc/core/src/moe/moe_sm90.inc new file mode 100644 index 00000000..923ef834 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_sm90.inc @@ -0,0 +1,1774 @@ +// Fused MoE: sort+dispatch → MLP+combine — torch-free core (liger_cute). +// +// Ported from LigerCommKernels' moe.cu. The device kernels and the templated +// host launcher are upstream-verbatim (namespace liger, internal to the .so); +// only the boundary changed: the remaining public entry points here are flat +// extern "C" control/configuration APIs. Tensor-carrying MoE calls are exposed +// through TVM FFI and lower directly into the internal MoeFwdArgs/MoeBwdArgs +// launch bundles. torch::Tensor never appears here. +#define LIGER_CUTE_BUILDING 1 + +#include +#include // pulls in nvshmemi_device_state_d, required by the + // inline device helpers in . +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include "liger_cute/moe.h" +#include "liger_cute/check.h" +#include "liger_cute/detail/status.h" +#include "liger_cute/detail/comm_schedule.cuh" + +#include "moe_sm90.cuh" +#include "moe_symm_config_sm90.cuh" +#include "moe_launch.h" // MoeFwdArgs / MoeBwdArgs (shared with the tuner) +#include "moe_utils.cuh" +#include "moe_fwd_bwd_tune_configs_sm90.hpp" +#include "moe_dispatch_configs_sm90.cuh" +#include "tma_copy_atom_sm90.cuh" +#include "radix_sort.cuh" +#include "dispatch.cuh" +#include "combine.cuh" +#include "buffer_pool.cuh" // shim: liger:: aliases of liger_cute::detail infra +#include "fused_memset.cuh" + +using namespace cute; + +namespace liger { + +// 0 = build every dispatch family (developer fallback). CMake sets this to 90 +// for Hopper builds and 100 for Blackwell builds so normal native builds do not +// instantiate kernels for hardware that cannot run in that binary. +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif + +// init_comm_schedule lives in liger_cute::detail (comm_schedule.cuh); the host +// configure path below calls it through this alias. +using liger_cute::detail::init_comm_schedule; + +// ============================================================================ +// Global iteration counter for PeSync +// ============================================================================ +// +// PeSync's dispatch_done / mlp_done symmetric arrays must increase +// monotonically across calls (the wait-side uses NVSHMEM_CMP_GE). The +// iteration value is what each call SETs the signal to via +// nvshmemx_signal_op(SET). +// +// The counter lives in device memory (allocated by BufferPool as +// "moe_fwd_iter_counter") and is atomicAdded at the end of every +// kernel by CTA 0 thread 0 — see PeSync::advance_iteration_counter(). +// This is the form that survives CUDA graph capture: a host-side +// atomic would be read once on the host before the kernel launch and +// baked into the captured graph, so all replays would re-SET the same +// already-met value and the cross-PE wait would race; a GPU-resident +// counter incremented inside the captured kernel re-evaluates on +// every replay, giving each launch a fresh monotonic target. +// +// Buffer pool keys "moe_fwd_iter_counter" (fwd) and +// "moe_bwd_iter_counter" (bwd) name the two separate counters. + +// ============================================================================ +// MoeFusedConfig — bundles all tunable kernel parameters +// ============================================================================ +// +// TileM is fixed at 64 by the WGMMA hardware. +// NumThreads is fixed at 384 (10 MLP warps + 2 comm warps). +// All other parameters are exposed for autotuning. + +template < + typename Element_ = bfloat16_t, + int TileM_ = 128, + int TileN1_ = 128, + int TileK1_ = 64, + int Stages1_ = 4, + int EpiChunkN1_ = 64, + int TileN2_ = 256, + int TileK2_ = 64, + int Stages2_ = 4, + int EpiChunkN2_ = 64, + int ZBufferSlots_ = 4, + int CommNumStages_ = 2, + int NC_ = 4, + // GemmTileM_ decouples the GEMM compute tile (mlp1/mlp2 wgmma + TMA) from + // the COMM staging tile (TileM_). TileM_ is the comm tile — staging ring + // slots, per-tile semaphores, NC, sort/dispatch — and may be wider than the + // GEMM tile. GemmTileM_ ∈ {64,128}: when < TileM_, the GEMM waits on one comm + // semaphore once, then steps through kSubTiles = TileM_/GemmTileM_ sub-tiles + // of GemmTileM_ rows each (finer TMA + wgmma) within that one comm slot. + // Default GemmTileM_ = TileM_ → kSubTiles = 1 → bit-identical to the old + // single-TileM path. + int GemmTileM_ = TileM_, + int Compute_ = 90> +struct MoeFusedConfig { + using Element = Element_; + // GEMM tile (mlp1/mlp2): GemmTileM_ ∈ {64,128}, independent of the comm tile. + using Traits1 = Mlp1Traits< + Element, GemmTileM_, TileN1_, TileK1_, Stages1_, EpiChunkN1_>; + using Traits2 = Mlp2Traits; + + // No compile-time constraint between NC, NSplit and CommNumStages + // — the GEMM-side iterator is ticket-based (slot = T mod L, + // ticket = T / L) and synchronizes per-slot via raw atomics + // computed from the global tile index, so the only requirement + // (runtime) is that NC evenly divides gridDim.x · gridDim.y so MC + // is integer. CommNumStages is now a free buffering knob. + + // Compile-time smem budget. H100 dynamic-smem cap is 228 KiB; we leave + // ~4 KiB for the runtime-sized remote_offsets region appended by the + // launcher (see moe.cu: smem_mlp += num_pes * (experts_per_pe+1) * 4). + // Catches over-budget configs at link time instead of cudaErrorInvalidValue + // at launch. + static constexpr size_t kSmemBytes = sizeof(MoeSmem); + static constexpr size_t kSmemBudget = 224 * 1024; + static_assert(kSmemBytes <= kSmemBudget, + "MoE fused fwd kernel smem footprint exceeds 224 KiB budget — " + "reduce TileN/TileK/Stages for this config"); + + static constexpr int kNC = NC_; + static constexpr int kZBufferSlots = ZBufferSlots_; + static constexpr int kCommNumStages = CommNumStages_; + // kTileM is the COMM tile (staging ring slots, per-tile semaphores, NC, + // sort/dispatch granularity). + static constexpr int kTileM = TileM_; + static constexpr int kGemmTileM = GemmTileM_; // gemm tile ∈ {64,128} + static constexpr int kSubTiles = TileM_ / GemmTileM_; + static constexpr int kCompute = Compute_; + static constexpr int kNumThreads = 384; // fixed + + static_assert(TileM_ % GemmTileM_ == 0, + "CommTileM must be an integer multiple of GemmTileM"); + static_assert(TileM_ % NC_ == 0, + "CommTileM must be divisible by NC"); + static_assert(kSubTiles == 1 || kSubTiles == 2, + "CommTileM/GemmTileM must be 1 or 2"); + static_assert(GemmTileM_ == 64 || GemmTileM_ == 128, + "GemmTileM must be 64 or 128"); + + // Iterators operate on the COMM tile (kTileM), independent of the GEMM + // sub-tile size. + using LocalIter = LocalMlpTileIterator; + using RemoteIter = RemoteMlpTileIterator< + Element, kCommNumStages, kNC, kTileM, kSubTiles>; +}; + +// ============================================================================ +// MoeParams -- sort + dispatch + MLP + combine (single kernel) +// ============================================================================ + +template +struct MoeParams { + using Element = typename Config::Element; + + // Router outputs (read-only, from caller's tensors). + const int* expert_indices; // [T, K] + const Element* expert_weights; // [T, K] + + // Sort buffers. + SortBuffers sort; + + // Dispatch. + const Element* tokens; // [T, D] input (symmetric) + Element* dispatch_dst; // [total_slots, D] sorted output (symmetric) + + // Combine output. + Element* combine_output; // [T, D] final output + + // MLP + comm. + CommBuffers comm; + + // Dimensions. + int num_tokens; + int intermediate_dim; + int top_k; + int total_slots; + int num_m_tiles; // sort input tiles +}; + +// ============================================================================ +// MoeBuffers -- all device/symmetric allocations +// ============================================================================ + +template +struct MoeBuffers { + using Element = typename Config::Element; + + // MLP intermediate Z. + Element* z_buf; + int z_rows; + + // Sorted X (dispatch destination) -- symmetric, remote PEs get from here. + Element* x_sorted; + + // Output Y -- symmetric, local MLP writes here, remote PEs put here. + Element* y_buf; + + // Phase counters. + int* phase_counter; + + // Comm staging. + Element* src_staging; + Element* dst_staging; + int* src_ready; + int* src_consumed; + int* dst_ready; + int* dst_consumed; + int* tile_expert_ids; + + // All PEs' expert offsets -- symmetric. + int* all_expert_offsets; + int* all_expert_counts; + int* cta_counter; // [1] device memory — shared monotonic + + // Sort buffers. + int* tile_expert_counts; + int* sorted_token_ids; + int* token_expert_slots; + int* expert_offsets; // symmetric + int* expert_counts; // symmetric exact counts + int* cta_done; + int* cta_sums; + + // Per-tile expert IDs for MLP local iterator. + int* mlp_tile_expert_ids; +}; + +// MoeSymmConfig + get_symm_config() declaration live in the shared header so +// moe.cu and moe_bwd.cu can't disagree on the layout (see header). This TU +// owns the single instance. +MoeSymmConfig& get_symm_config() { + static MoeSymmConfig cfg; + return cfg; +} + +void moe_configure_symmetric( + int max_tokens, + int hidden_dim, + int max_num_experts, + int max_top_k, + int num_pes, + int num_hosts, + int gpus_per_host) { + static constexpr int kTileM = kMaxMoeCommTileM; + + // num_hosts × gpus_per_host must equal the team size so the comm + // schedule (g_dest_table / g_rank_table) addresses exactly the PEs + // in the team. Mismatch would silently map ranks to non-existent + // destinations (or skip some). + LIGER_CHECK(num_hosts * gpus_per_host == num_pes, + "moe_configure_symmetric: num_hosts (", num_hosts, ") * gpus_per_host (", + gpus_per_host, ") = ", num_hosts * gpus_per_host, + " must equal num_pes (", num_pes, ")"); + + // Populate g_dest_table / g_rank_table in device constant memory. + // Idempotent — safe to re-call with the same (N, M). + init_comm_schedule(num_hosts, gpus_per_host); + + auto& cfg = get_symm_config(); + cfg.max_total_slots = max_tokens * max_top_k + max_num_experts * kTileM; + cfg.max_num_experts = max_num_experts; + cfg.hidden_dim = hidden_dim; + cfg.num_pes = num_pes; + cfg.experts_per_pe = max_num_experts / num_pes; + cfg.max_top_k = max_top_k; + cfg.initialized = true; + + // Preset sizes for the symmetric stack used to hold X / Y / routing metadata + // across fwd → bwd. Element type is bf16 (= 2 bytes); kept in sync + // with allocate_moe_buffers below. + constexpr std::size_t kElemBytes = 2; // sizeof(__nv_bfloat16) + std::size_t symm_slot_bytes = + static_cast(cfg.max_total_slots) * cfg.hidden_dim * kElemBytes; + std::size_t all_off_bytes = + static_cast(cfg.num_pes) * (cfg.max_num_experts + 1) * sizeof(int); + std::size_t all_count_bytes = + static_cast(cfg.num_pes) * cfg.max_num_experts * sizeof(int); + + auto& stack = global_symmetric_stack(); + stack.set_size(kSymmKeyXSorted, symm_slot_bytes); + stack.set_size(kSymmKeyYBuf, symm_slot_bytes); + stack.set_size(kSymmKeyAllExpertOffsets, all_off_bytes); + stack.set_size(kSymmKeyAllExpertCounts, all_count_bytes); +} + +// Releases the forward metadata entries pushed onto the symmetric stack by the +// most recent fwd call. Must be called collectively +// on every PE in LIFO order vs. the fwd that produced them. After this, +// the buffers can be reused by a subsequent fwd. The torch::Tensor views +// returned by fwd remain valid pointers but will alias whatever lives in +// those slots after the next put. +void moe_pop_fwd() { + auto& stack = global_symmetric_stack(); + // LIFO order: bottom pushes first. Reverse the put order in + // allocate_moe_buffers (x_sorted, y_buf, all_expert_offsets, + // all_expert_counts). + stack.pop(kSymmKeyAllExpertCounts); + stack.pop(kSymmKeyAllExpertOffsets); + stack.pop(kSymmKeyYBuf); + stack.pop(kSymmKeyXSorted); +} + +template +static MoeBuffers allocate_moe_buffers( + BufferPool& pool, + cudaStream_t stream, + int grid_x, + int num_blocks, + int hidden_dim, + int intermediate_dim, + int num_tokens, + int top_k, + int num_experts, + int experts_per_pe, + int num_pes, + int total_slots, + int num_m_tiles, + // Caller-owned torch storage for the two sort outputs that bwd + // needs to consume. Passed in so we don't lease buffer-pool + // slots that the next fwd would overwrite. + int* token_expert_slots_ptr, + int* mlp_tile_expert_ids_ptr) { + + using Element = typename Config::Element; + static constexpr int kTileM = Config::kTileM; + static constexpr int kZBufferSlots = Config::kZBufferSlots; + + // Staging is a flat ring of L = MC · kCommNumStages slots, where + // MC = num_blocks / NC (num_blocks == launched gridDim.x). This MUST equal + // the device-side comm producer's stride (mlp_comms.cuh: MC = n_comm/NC, + // n_comm = (gridDim.x/NC)·NC) AND the MLP consumer's ring stride + // (RemoteMlpTileIterator::init: mc = gridDim.x/NC), so every tile's + // slot = g % L agrees across producer / consumer / allocation. Sizing this + // on grid_x·NSplit (= n_gemm) instead under-sizes the signal arrays whenever + // gridDim.x/NC > n_gemm/NC — i.e. num_sms % NSplit != 0, e.g. NSplit=8 on + // 132 SMs (33 vs 32) — and the gridDim.x-based ring then indexes past the + // end of src_ready / tile_expert_ids (OOB). num_blocks/NC == gridDim.x/NC + // == n_gemm/NC in the divisible (legacy) case, so this is a no-op there. + // Comm strides MC through NumStages per xc; MLP (ticket iterator) wraps mod L. + // Raw SM count — used as the device-independent upper bound on the comm + // staging row count (mc <= sm_count) so the symmetric staging size is a + // worst-case bound across configs (see staging sizing below). + int sm_count = 0; + { + int dev = 0; + cudaGetDevice(&dev); + cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev); + } + + MoeBuffers b; + + // Symmetric allocations use max sizes from MoeSymmConfig + // so all PEs allocate the same sizes regardless of per-call config. + auto& scfg = get_symm_config(); + LIGER_CHECK(scfg.initialized, "Call moe_configure_symmetric before moe_fused_fwd_bf16"); + + // Persistent symmetric tensors and routing metadata live in the symmetric + // stack so they can survive across fwd calls into bwd. + auto& sym_stack = global_symmetric_stack(); + + // Z buffer (device). + // size_t cast on b.z_rows: at NSplit=2 ZBuf=16 I=16384, the product + // z_rows · intermediate_dim · 2 = ~4.4 GB which overflows int32 if + // the multiply happens in int (b.z_rows and intermediate_dim are + // both int). Cast first so the entire product is in size_t. + int max_runtime_grid_x = (num_blocks + 1) / 2; // min candidate NS = 2 + int z_m_tiles = max_runtime_grid_x * kZBufferSlots; + b.z_rows = z_m_tiles * kTileM; + size_t z_bytes = static_cast(b.z_rows) * intermediate_dim * sizeof(Element); + b.z_buf = reinterpret_cast(pool.get_device("moe_z_buf", z_bytes)); + + // Sorted X -- symmetric, dispatch writes here, remote PEs get from here. + b.x_sorted = reinterpret_cast(sym_stack.put(kSymmKeyXSorted)); + + // Output Y -- symmetric, local MLP writes here, remote PEs put here. + b.y_buf = reinterpret_cast(sym_stack.put(kSymmKeyYBuf)); + + // Phase counters. + size_t counter_bytes = max_runtime_grid_x * sizeof(int); + b.phase_counter = reinterpret_cast( + pool.get_device("moe_phase_counter", counter_bytes)); + if (cudaError_t e = cudaMemsetAsync(b.phase_counter, 0, counter_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(phase_counter) failed: ", cudaGetErrorString(e)); + + // Flat staging list — L = grid_x * kCommNumStages slots, each + // holding TileM rows × hidden_dim columns. + // Size by the CONFIGURED MAX on every axis (scfg.hidden_dim, max_comm_stages, + // max_tile_m), NOT the per-call/per-config values: the symmetric pool caches + // by name and aborts on grow, so a small first config must not lock in a + // staging buffer a later larger config can't grow into. This covers all three + // independently-varying axes that the tuner sweeps: + // - hidden_dim: D=512 then D=8192 + // - CommNumStages: CS=2 then CS=8 (NS=16 config) + // - TileM: 64 then 128 + // The kernel indexes slots by the per-call (comm_l, tile_elems <= max), so an + // over-sized buffer is always safe (symmetric memory may be larger than + // needed, never smaller). + size_t staging_rows_ub = + (size_t)sm_count * scfg.max_comm_stages * scfg.max_tile_m; + size_t src_bytes = staging_rows_ub * scfg.hidden_dim * sizeof(Element); + size_t dst_bytes = staging_rows_ub * scfg.hidden_dim * sizeof(Element); + b.src_staging = reinterpret_cast(pool.get_symmetric("moe_src_staging", src_bytes)); + b.dst_staging = reinterpret_cast(pool.get_symmetric("moe_dst_staging", dst_bytes)); + + // Device buffers are cached by key across tuning/config changes, so size and + // clear them at a configuration-independent upper bound. Since + // kSubTiles <= NC, (num_blocks / NC) * stages * kSubTiles is bounded by + // sm_count * max_comm_stages. + size_t signal_l_ub = + static_cast(sm_count) * scfg.max_comm_stages; + size_t sig_bytes = signal_l_ub * sizeof(int); + b.src_ready = static_cast(pool.get_device("moe_src_ready", sig_bytes)); + b.src_consumed = static_cast(pool.get_device("moe_src_consumed", sig_bytes)); + b.dst_ready = static_cast(pool.get_device("moe_dst_ready", sig_bytes)); + b.dst_consumed = static_cast(pool.get_device("moe_dst_consumed", sig_bytes)); + if (cudaError_t e = cudaMemsetAsync(b.src_ready, 0, sig_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(src_ready) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(b.src_consumed, 0, sig_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(src_consumed) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(b.dst_ready, 0, sig_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(dst_ready) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(b.dst_consumed, 0, sig_bytes, stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(dst_consumed) failed: ", cudaGetErrorString(e)); + + // Per-parent-slot expert IDs use the same conservative upper bound. + size_t eid_bytes = signal_l_ub * sizeof(int); + b.tile_expert_ids = static_cast(pool.get_device("moe_tile_expert_ids", eid_bytes)); + + // All PEs' expert offsets (symmetric). Use max experts for stable sizing. + b.all_expert_offsets = static_cast(sym_stack.put(kSymmKeyAllExpertOffsets)); + b.all_expert_counts = static_cast(sym_stack.put(kSymmKeyAllExpertCounts)); + + // CTA counter for local all-CTA barriers around NVSHMEM collectives. + b.cta_counter = static_cast( + pool.get_device("moe_cta_counter", sizeof(int))); + + // Sort buffers -- use scfg max sizes to avoid reallocation across configs. + int max_sorted_slots = scfg.max_total_slots; + int max_m_tiles = scfg.max_total_slots / kTileM; + int max_tile_counts = max_m_tiles * scfg.max_num_experts; + int max_token_slots_ub = scfg.max_total_slots; + + (void)max_token_slots_ub; + b.tile_expert_counts = static_cast( + pool.get_device("moe_sort_tile_counts", max_tile_counts * sizeof(int))); + b.sorted_token_ids = static_cast( + pool.get_device("moe_sort_sorted_ids", max_sorted_slots * sizeof(int))); + // token_expert_slots: caller-owned torch tensor (passed in). Used by + // the bwd path after this call returns — must not alias a pool slot. + b.token_expert_slots = token_expert_slots_ptr; + b.expert_offsets = static_cast( + pool.get_symmetric("moe_sort_offsets", (scfg.max_num_experts + 1) * sizeof(int))); + b.expert_counts = static_cast( + pool.get_symmetric("moe_sort_counts", scfg.max_num_experts * sizeof(int))); + b.cta_done = static_cast( + pool.get_device("moe_sort_cta_done", num_blocks * sizeof(int))); + b.cta_sums = static_cast( + pool.get_device("moe_sort_cta_sums", num_blocks * scfg.max_num_experts * sizeof(int))); + + // Per-tile expert IDs for MLP — caller-owned torch tensor. + b.mlp_tile_expert_ids = mlp_tile_expert_ids_ptr; + + if (cudaError_t e = cudaMemsetAsync(b.cta_done, 0, num_blocks * sizeof(int), stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(cta_done) failed: ", cudaGetErrorString(e)); + if (cudaError_t e = cudaMemsetAsync(b.sorted_token_ids, 0xff, + max_sorted_slots * sizeof(int), stream); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaMemsetAsync(sorted_token_ids) failed: ", cudaGetErrorString(e)); + + return b; +} + +// ============================================================================ +// GET-via-TMA descriptor builder (raw CUtensorMap, FWD-only) +// ============================================================================ +// +// Box = [box_rows, kGetKChunk] over a row-major [n_rows, hidden_dim] bf16 +// tensor. SWIZZLE_NONE keeps the global bytes row-major, so the MLP's +// swizzled reader (tma_load_x_remote) sees correct data. Returns false on +// encode failure → caller disables the TMA-get path (falls back to getmem). +static inline bool make_get_tma_map(CUtensorMap& m, void* base, + int hidden_dim, int n_rows, int box_rows) { + uint64_t gdim[2] = { (uint64_t)hidden_dim, (uint64_t)n_rows }; + uint64_t gstride[1] = { (uint64_t)hidden_dim * sizeof(bfloat16_t) }; + uint32_t bdim[2] = { (uint32_t)liger::kGetKChunk, (uint32_t)box_rows }; + uint32_t estride[2] = { 1, 1 }; + CUresult r = cuTensorMapEncodeTiled( + &m, CU_TENSOR_MAP_DATA_TYPE_UINT16, 2, base, gdim, gstride, bdim, estride, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_L2_128B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + return r == CUDA_SUCCESS; +} + +// ============================================================================ +// Fused kernel: sort → dispatch → MLP + combine +// ============================================================================ + +template +__global__ void __launch_bounds__(Config::kNumThreads) +moe_fused_kernel( + __grid_constant__ TmaLoadX const tma_load_x, + __grid_constant__ TmaLoadW1 const tma_load_b, + __grid_constant__ TmaLoadW1 const tma_load_c, + __grid_constant__ TmaStoreZ const tma_store_z, + __grid_constant__ TmaLoadZ const tma_load_z, + __grid_constant__ TmaLoadW2 const tma_load_a, + __grid_constant__ TmaStoreY const tma_store_y, + __grid_constant__ TmaLoadX const tma_load_x_remote, + __grid_constant__ TmaStoreY const tma_store_y_remote, + __grid_constant__ TmaStoreYCombine const tma_store_y_combine, + __grid_constant__ MlpDims const mlp_dims, + __grid_constant__ MoeParams const p, + __grid_constant__ int const static_nsplit, + __grid_constant__ GetTmaDescs const get_descs, + __grid_constant__ PeerYStoreDescs const peer_y) { + + using Element = typename Config::Element; + using Traits1 = typename Config::Traits1; + using Traits2 = typename Config::Traits2; + using LocalIter = typename Config::LocalIter; + using RemoteIter = typename Config::RemoteIter; + constexpr int kTileM = Config::kTileM; + constexpr int kNumThreads = Config::kNumThreads; + constexpr int kZBufferSlots = Config::kZBufferSlots; + constexpr int kCommNumStages = Config::kCommNumStages; + constexpr int kNC = Config::kNC; + + // -- Sort ---------------------------------------------------------------- + { + extern __shared__ int sort_smem[]; + sort_tokens( + sort_smem, p.expert_indices, p.sort, + p.num_tokens, p.comm.num_experts, p.top_k, p.num_m_tiles); + } + __threadfence_system(); // device-scope (NVLink-only build; IB needs system scope) + __syncthreads(); + + // Target-based cross-CTA barrier — single instance per kernel call, + // reused by every PeSync::barrier / signal_dispatch / signal_mlp + // invocation below. Constructing a fresh barrier per call resets + // the per-CTA target to 0 while the shared device counter keeps + // growing → fast CTAs spin past the next-call target instantly. + SyncThreadsCtaCounterBarrier pe_barrier(p.comm.pe_sync.cta_counter, + (int)gridDim.x * (int)gridDim.y); + + // Cross-CTA barrier: sort writes sorted_token_ids across CTAs; + // dispatch and combine read them from other CTAs. + p.comm.pe_sync.barrier(pe_barrier); + + // -- Broadcast expert offsets to all PEs ---------------------------------- + broadcast_expert_offsets(p.comm); + __syncthreads(); + + // -- Dispatch ------------------------------------------------------------- + dispatch_tokens( + p.tokens, p.dispatch_dst, + p.sort.sorted_token_ids, + p.comm.hidden_dim, p.total_slots); + __syncthreads(); + + // H2 (IB coherence): dispatch_tokens wrote dispatch_dst (= x_sorted) with + // plain GPU stores; remote PEs RDMA-read it in do_get after wait_dispatch. + // Flush those stores to HBM so the remote NIC reads fresh data — the quiet + // below only orders this PE's OUTGOING nvshmem ops, not these local + // stores, and __syncthreads is device-scope. The offset-broadcast source + // is already covered by radix_sort's __threadfence_system. NVLink P2P + // reads are L2-coherent, so single-host never exercised this. + __threadfence_system(); // device-scope (NVLink-only build; IB needs system scope) + + p.comm.pe_sync.barrier_all(pe_barrier); + + // -- MLP (local + remote) ------------------------------------------------ + { + extern __shared__ char raw_smem[]; + auto& smem = *reinterpret_cast*>(raw_smem); + + size_t offsets_region = sizeof(MoeSmem); + offsets_region = (offsets_region + sizeof(int) - 1) & ~(sizeof(int) - 1); + size_t offsets_count = + (size_t)p.comm.num_pes() * (p.comm.experts_per_pe + 1); + if (threadIdx.x == 0) { + smem.comm.remote_offsets = + reinterpret_cast(raw_smem + offsets_region); + smem.comm.remote_counts = + smem.comm.remote_offsets + offsets_count; + } + __syncthreads(); + + // GET-via-TMA bounce region: placed after remote_offsets, 128-aligned + // (TMA dst smem must be 128B-aligned). MUST match the host smem_size + // computation (same sizeof(MoeSmem), num_pes, experts_per_pe). Only + // touched by the comm warps' do_get when get_descs.enabled. + char* get_bounce = nullptr; + if (get_descs.enabled) { + size_t boff = offsets_region + + (offsets_count + + (size_t)p.comm.num_pes() * p.comm.experts_per_pe) + * sizeof(int); + boff = (boff + 127) & ~((size_t)127); + get_bounce = raw_smem + boff; + } + + //if (threadIdx.x == 0) + // smem.comm.total_tiles = -1; + + // Flat-grid logical coordinates for this CTA. GEMM and comm gate to + // the first n_gemm = floor_NS CTAs (the comm↔MLP staging ring needs + // grid_x = n_gemm/static_nsplit and MC = n_gemm/NC consistent); the gap CTAs + // [n_gemm, num_blocks) fall through to the shared pe_barrier only. + int flat_id = blockIdx.x; + int n_gemm = p.comm.n_gemm; + int grid_x = n_gemm / static_nsplit; + int col = flat_id / static_nsplit; + int runtime_nsplit = static_nsplit; + + // gemm_active removed: comm warps self-gate on their precomputed tile + // count (smem.comm.total_tiles) so a CTA can do comm-only work without + // participating in GEMM (the point of the comm/GEMM grid relaxation). + // Run the prologue unconditionally so every launched CTA has total_tiles + // initialized; the MLP/GEMM path still gates on GEMM-grid membership + // (flat_id < n_gemm) below. + nvshmem_comm_prologue(smem.comm, p.comm, static_nsplit); + __syncthreads(); + + if (threadIdx.x == 0) { + int total_m_tiles = smem.comm.global_total; + int mlp_tiles = total_m_tiles; + if constexpr (Config::kSubTiles > 1) { + mlp_tiles = 0; + for (int pe = 0; pe < p.comm.num_pes(); ++pe) { + for (int e = 0; e < p.comm.experts_per_pe; ++e) { + int count = smem.comm.remote_counts[ + pe * p.comm.experts_per_pe + e]; + mlp_tiles += + (count + Config::kGemmTileM - 1) / + Config::kGemmTileM; + } + } + } + runtime_nsplit = select_runtime_nsplit_from_tiles( + mlp_tiles, + mlp_dims.num_n_tiles_1, + mlp_dims.num_n_tiles_2, + mlp_dims.num_k_tiles_1, + mlp_dims.num_k_tiles_2, + (int)gridDim.x, + static_nsplit); + n_gemm = ((int)gridDim.x / runtime_nsplit) * runtime_nsplit; + grid_x = n_gemm / runtime_nsplit; + col = flat_id / runtime_nsplit; + + int mlp_count = 0; + if (flat_id < n_gemm && col < mlp_tiles) + mlp_count = (mlp_tiles - 1 - col) / grid_x + 1; + smem.comm.per_cta_tiles = mlp_count; + smem.comm.runtime_nsplit = runtime_nsplit; + smem.comm.runtime_n_gemm = n_gemm; + smem.comm.runtime_grid_x = grid_x; + } + __syncthreads(); + runtime_nsplit = smem.comm.runtime_nsplit; + n_gemm = smem.comm.runtime_n_gemm; + grid_x = smem.comm.runtime_grid_x; + col = flat_id / runtime_nsplit; + + // Diverge: comm warps handle NVSHMEM, MLP warps handle compute. + // Separate scopes prevent CommBuffers from inflating MLP register budget. + // Fwd uses warps 1-3 (2 getters + 1 putter). + int warp_id = threadIdx.x / 32; + bool is_comm_warp = (warp_id == kCommGetWarp || warp_id == kCommPutWarp); + + if (is_comm_warp) { + // Comm warps: dedicated get/put — full CommBuffers access. Self-gate + // on the precomputed comm tile count — no GEMM coupling, so a + // comm-only CTA (no local GEMM work) still drives its comm tiles. + if (smem.comm.total_tiles > 0) { + nvshmem_comm_main< + Element, kTileM, kCommNumStages, kNC, Config::kSubTiles>( + smem.comm, p.comm, &get_descs, get_bounce); + } + } else if (flat_id < n_gemm) { + // MLP warps: no CommBuffers in scope. Gap CTAs (flat_id >= n_gemm) + // skip this entirely and fall through to the pe_barrier below. + int local_e_start = p.comm.my_pe() * p.comm.experts_per_pe; + int local_e_end = local_e_start + p.comm.experts_per_pe; + int local_offset = p.sort.expert_offsets[local_e_start]; + int local_end = p.sort.expert_offsets[local_e_end]; + int local_m_tiles = (local_end - local_offset) / kTileM; + int local_tile_off = local_offset / kTileM; + + LocalIter local_iter; + local_iter.init( + p.dispatch_dst + local_offset * p.comm.hidden_dim, + p.sort.tile_expert_ids + local_tile_off, + local_m_tiles, p.comm.hidden_dim, + /*col_stride=*/grid_x, /*start_col=*/col, local_tile_off); + + moe_fused_fwd( + smem, local_iter, + p.comm.src_ready, p.comm.src_consumed, + p.comm.dst_ready, p.comm.dst_consumed, + p.comm.experts_per_pe, p.comm.num_pes(), p.comm.my_pe(), + get_descs.gpus_per_node, + tma_load_x, tma_load_b, tma_load_c, tma_store_z, + tma_load_z, tma_load_a, tma_store_y, + tma_load_x_remote, tma_store_y_remote, + mlp_dims, + local_m_tiles * kTileM, local_m_tiles, + col, grid_x, /*split=*/flat_id % runtime_nsplit, runtime_nsplit, + local_e_start, + /*y_peer_desc=*/peer_y.enabled ? peer_y.desc : nullptr); + } + + __threadfence(); + __syncthreads(); + p.comm.pe_sync.barrier_all(pe_barrier); + } + + // -- Combine -------------------------------------------------------------- + // Reuse the kernel's dynamic SMEM region as CombineSmem. wait_mlp above + // is a cross-CTA + cross-PE barrier that ends in its own __syncthreads, so + // the CTA is already aligned and the MLP/Comm SMEM contents are dead by + // this point — no extra __syncthreads needed before the reinterpret. + { + constexpr int kCombineNumWGs = + kNumThreads / (kCombineWarpsPerWG * kCombineWarpSize); + extern __shared__ char raw_smem[]; + auto& combine_smem = + *reinterpret_cast*>(raw_smem); + combine_tokens_tma( + reinterpret_cast(p.comm.local_output), + p.expert_weights, + p.sort.token_expert_slots, + tma_store_y_combine, + combine_smem, + p.comm.hidden_dim, + p.num_tokens, + p.top_k); + } + +} + +// ============================================================================ +// Host wrapper +// ============================================================================ + +// MoeFwdArgs (the POD launch bundle) is defined in moe_launch.h, shared with the +// tuner. The runtime dispatch table holds a single function-pointer type +// (MoeFwdFn) while every tuned config is a distinct template instantiation. + +template < + typename Element_ = bfloat16_t, + int TileN1 = 128, + int TileK1 = 64, + int Stages1 = 4, + int TileN2 = 256, + int TileK2 = 64, + int Stages2 = 4, + int ZBufferSlots = 4, + int CommNumStages = 1, + int EpiChunkN1 = 64, // mlp1 epilogue chunk; must divide TileN1 + int EpiChunkN2 = 64, // mlp2 epilogue chunk; must divide TileN2 + int TileM = 128, // COMM tile (staging/semaphores/NC); 128 + int GemmTileM = TileM, // GEMM tile ∈ {64,128}; default == comm (no sub-tiling) + int Compute = 90> +void +moe_fused_fwd_bf16(const MoeFwdArgs& a, int static_nsplit) { + + // Unpack the POD argument bundle into the names the upstream body uses. + const int num_experts = a.num_experts; + const int top_k = a.top_k; + const nvshmem_team_t team = a.team; + cudaStream_t stream = a.stream; + + // NC = number of cooperating comm CTAs per (xc, yc) tile. With the + // ticket-based GEMM-side iterator, runtime NS and NC no longer need to + // share a divisibility relationship — the only runtime requirement + // is NC | (gridDim.x · gridDim.y) so MC is integer. + // + // Wrapper param order (10 "core" args + trailing EpiChunkN + TileM): + // Element, TN1, TK1, S1, TN2, TK2, S2, ZBuf, CS, + // EpiChunkN1, EpiChunkN2, TileM + // MoeFusedConfig interleaves TileM/EpiChunkN inside each phase: + // Element, TileM, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, + // ZBuf, CS, NC + // The remapping happens here. TileM is shared by MLP1 and MLP2. + // + // NC choice (verified via bench_mlp_comms BW sweep): + // - TileM=128: NC=4 gives ~2× the comm bandwidth of NC=2 because + // 4 cooperating comm CTAs split each tile, deepening the comm + // pipeline (more in-flight slots = grid.x/NC). + // - TileM=64: NC=2 already matches NC=4 throughput because tiles + // are half-size, so twice as many flow through the same pipeline. + // Both architectures use CommTileM=128 with NC=4. Keeping CommTileM uniform per + // architecture preserves + // all-to-all dispatch padding (x_sorted / expert_offsets, padded to the + // comm tile) across PEs. The 256/8 ratio keeps each SM100 comm CTA's row + // band and GET-bounce footprint unchanged. + using CommGeometry = MoeCommGeometry; + using Config = MoeFusedConfig< + Element_, /*CommTileM=*/CommGeometry::TileM, + TileN1, TileK1, Stages1, EpiChunkN1, + TileN2, TileK2, Stages2, EpiChunkN2, + ZBufferSlots, + CommNumStages * CommGeometry::RingStageScale + + CommGeometry::RingStageExtra, + /*NC=*/CommGeometry::NC, + /*GemmTileM=*/GemmTileM, /*Compute_=*/Compute>; + + using Element = typename Config::Element; + using Traits1 = typename Config::Traits1; + using Traits2 = typename Config::Traits2; + using TmaLoadAtom = TmaLoadAtomForCompute; + using TmaStoreAtom = TmaStoreAtomForCompute; + constexpr int kTileM = Config::kTileM; + constexpr int kNumThreads = Config::kNumThreads; + + // Iteration counter lives on the device and is incremented at the + // end of the kernel — see PeSync::advance_iteration_counter() and + // the comment block at the top of this file. + int num_tokens = a.num_tokens; + int hidden_dim = a.hidden_dim; + int experts_per_pe = a.experts_per_pe; // local experts (from weight matrices) + int intermediate_dim = a.intermediate_dim; + + // Upper bound for total_slots: T * K + E * TileM (TileM padding per expert). + // Actual value is computed by sort (expert_offsets[E]) inside the kernel. + int max_total_slots = num_tokens * top_k + num_experts * kTileM; + int num_m_tiles_in = (num_tokens + kTileM - 1) / kTileM; // sort input tiles + + if constexpr (Compute == 100) { + LIGER_CHECK(hidden_dim % Traits1::TileK == 0, + "hidden_dim (", hidden_dim, ") must be divisible by MLP1 TileK (", Traits1::TileK, ")"); + LIGER_CHECK(intermediate_dim % Traits1::TileN == 0, + "intermediate_dim (", intermediate_dim, ") must be divisible by MLP1 TileN (", Traits1::TileN, ")"); + LIGER_CHECK(intermediate_dim % Traits2::TileK == 0, + "intermediate_dim (", intermediate_dim, ") must be divisible by MLP2 TileK (", Traits2::TileK, ")"); + LIGER_CHECK(hidden_dim % Traits2::TileN == 0, + "hidden_dim (", hidden_dim, ") must be divisible by MLP2 TileN (", Traits2::TileN, ")"); + } + LIGER_CHECK(hidden_dim % 8 == 0, + "hidden_dim (", hidden_dim, ") must be divisible by 8 (int4 vectorization in dispatch/combine)"); + LIGER_CHECK(intermediate_dim % 8 == 0, + "intermediate_dim (", intermediate_dim, ") must be divisible by 8 (int4 vectorization)"); + LIGER_CHECK(top_k <= kCombineMaxTopK, + "top_k (", top_k, ") exceeds kCombineMaxTopK (", kCombineMaxTopK, + "). Bump the constant in combine.cuh."); + + int num_n_tiles_1 = (Compute == 90) + ? (intermediate_dim + Traits1::TileN - 1) / Traits1::TileN + : intermediate_dim / Traits1::TileN; + int num_n_tiles_2 = (Compute == 90) + ? (hidden_dim + Traits2::TileN - 1) / Traits2::TileN + : hidden_dim / Traits2::TileN; + + // Relaxed: runtime NS > num_n_tiles is allowed; surplus split CTAs run + // empty inner loops but still drive their TMA-load/store mbarriers so the + // pipeline stays balanced across the column. + + LIGER_CHECK(a.weight_expert_stride >= + static_cast(intermediate_dim) * hidden_dim, + "weight expert stride (", a.weight_expert_stride, + ") must cover at least one [I,D] expert matrix"); + LIGER_CHECK(a.weight_expert_stride % hidden_dim == 0, + "weight expert stride (", a.weight_expert_stride, + ") must contain complete hidden-dimension rows (D=", hidden_dim, ")"); + int weight_expert_stride_rows = + static_cast(a.weight_expert_stride / hidden_dim); + LIGER_CHECK(weight_expert_stride_rows % Traits1::TileN == 0, + "weight expert stride in rows (", weight_expert_stride_rows, + ") must be divisible by MLP1 TileN (", Traits1::TileN, ")"); + int total_n_rows_1 = + (experts_per_pe - 1) * weight_expert_stride_rows + intermediate_dim; + int total_n_rows_2 = experts_per_pe * hidden_dim; + + int num_pes = nvshmem_team_n_pes(team); + int my_pe = nvshmem_team_my_pe(team); + + // Grid — flat 1-D launch. static_nsplit only seeds n_gemm/grid_x before + // the kernel selects the runtime split from the actual routed tile count. + // + // GEMM and comm gate to the first n_gemm = floor_NS CTAs: the comm↔MLP + // staging ring couples producer and consumer (it needs grid_x = n_gemm/static_nsplit + // and MC = n_gemm/NC to be consistent), so they share one count. The + // sort/dispatch/combine grid-stride phases and the cross-CTA pe_barrier use + // ALL num_blocks launched CTAs. The runtime selector may choose a different + // NS inside the kernel and rewrites n_gemm/grid_x consistently for MLP. + int num_sms; + cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, a.device); + int floor_ns = num_sms - num_sms % static_nsplit; + int floor_nc = num_sms - num_sms % Config::kNC; + int n_gemm = floor_ns; // GEMM/comm-active CTAs + int grid_x = n_gemm / static_nsplit; // initial logical GEMM columns + int num_blocks = (floor_ns > floor_nc) ? floor_ns : floor_nc; // launched CTAs + dim3 grid(num_blocks); + + // stream comes from a.stream (the binding passes the current torch stream). + auto& pool = global_buffer_pool(); + + // The two sort outputs the bwd path consumes (token_expert_slots + // [max_total_slots], tile_expert_ids [max_total_slots/kTileM]) are + // caller-owned: the binding pre-allocated them at exactly these sizes and + // passed their device pointers in a. They must NOT alias pool slots that the + // next forward would overwrite. + int* token_expert_slots_ptr = a.token_expert_slots; + int* tile_expert_ids_ptr = a.tile_expert_ids; + + // Allocate all buffers. + auto buf = allocate_moe_buffers(pool, stream, grid_x, num_blocks, + hidden_dim, intermediate_dim, num_tokens, top_k, num_experts, + experts_per_pe, num_pes, + max_total_slots, num_m_tiles_in, + token_expert_slots_ptr, + tile_expert_ids_ptr); + if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: CUDA error before TMA descriptor setup: ", cudaGetErrorString(e)); + + // Final combine output [T, D] — caller-owned (binding-allocated) Y buffer. + Element* output_ptr = reinterpret_cast(a.Y); + Element* y_buf_ptr = buf.y_buf; + + auto& scfg = get_symm_config(); + + // ---- Build MoeParams ------------------------------------------------------- + MoeParams p; + p.expert_indices = a.expert_indices; + p.expert_weights = reinterpret_cast(a.expert_weights); + + p.sort.tile_expert_counts = buf.tile_expert_counts; + p.sort.sorted_token_ids = buf.sorted_token_ids; + p.sort.token_expert_slots = buf.token_expert_slots; + p.sort.expert_offsets = buf.expert_offsets; + p.sort.expert_counts = buf.expert_counts; + p.sort.tile_expert_ids = buf.mlp_tile_expert_ids; + p.sort.cta_done = buf.cta_done; + p.sort.cta_sums = buf.cta_sums; + + p.tokens = reinterpret_cast(a.X); + p.dispatch_dst = buf.x_sorted; + p.combine_output = output_ptr; + + p.comm.src_staging = buf.src_staging; + p.comm.dst_staging = buf.dst_staging; + p.comm.src_ready = buf.src_ready; + p.comm.src_consumed = buf.src_consumed; + p.comm.dst_ready = buf.dst_ready; + p.comm.dst_consumed = buf.dst_consumed; + p.comm.expert_offsets = buf.expert_offsets; + p.comm.expert_counts = buf.expert_counts; + p.comm.sorted_token_ids = buf.sorted_token_ids; + p.comm.local_tokens = buf.x_sorted; + p.comm.local_output = y_buf_ptr; + p.comm.tile_expert_ids = buf.tile_expert_ids; + p.comm.all_expert_offsets = buf.all_expert_offsets; + p.comm.all_expert_counts = buf.all_expert_counts; + p.comm.pe_sync.cta_counter = buf.cta_counter; + p.comm.pe_sync.num_pes = num_pes; + p.comm.pe_sync.my_pe = my_pe; + p.comm.pe_sync.team = team; + p.comm.hidden_dim = hidden_dim; + p.comm.num_experts = num_experts; + p.comm.experts_per_pe = experts_per_pe; + p.comm.n_gemm = n_gemm; // GEMM/comm-active CTA count (floor_NS) + + p.num_tokens = num_tokens; + p.intermediate_dim = intermediate_dim; + p.top_k = top_k; + p.total_slots = max_total_slots; + p.num_m_tiles = num_m_tiles_in; + + // ---- MLP invariant dims (passed as __grid_constant__) ---------------------- + MlpDims mlp_dims; + mlp_dims.hidden_dim = hidden_dim; + mlp_dims.intermediate_dim = intermediate_dim; + mlp_dims.num_experts = experts_per_pe; + mlp_dims.total_n_rows_1 = total_n_rows_1; + mlp_dims.total_n_rows_2 = total_n_rows_2; + mlp_dims.expert_n_stride_1 = weight_expert_stride_rows / Traits1::TileN; + mlp_dims.num_n_tiles_1 = num_n_tiles_1; + mlp_dims.num_n_tiles_2 = num_n_tiles_2; + mlp_dims.num_k_tiles_1 = (Compute == 90) + ? (hidden_dim + Traits1::TileK - 1) / Traits1::TileK + : hidden_dim / Traits1::TileK; + mlp_dims.num_k_tiles_2 = (Compute == 90) + ? (intermediate_dim + Traits2::TileK - 1) / Traits2::TileK + : intermediate_dim / Traits2::TileK; + // Row-tile extent of the symmetric local_output (y_buf) for the + // direct-to-peer Y store (num_y_m_tiles when storing at the token-tile row). + mlp_dims.y_buf_m_tiles = max_total_slots / kTileM; + mlp_dims.phase_counter = buf.phase_counter; + + // ---- MLP TMA descriptors ----------------------------------------------- + // X for MLP is x_sorted (dispatch destination, symmetric). + // All shape/stride values are cast to int64_t — see the long comment + // below the tma_load_c block for the rationale. B and C in particular + // are large: their row count = experts_per_pe · intermediate_dim, + // times hidden_dim → exceeds INT_MAX at production shapes. + auto tma_load_x = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.x_sorted)), + make_shape(static_cast(max_total_slots), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits1::SmemLayoutX_1{}); + + auto* ptr_B = reinterpret_cast(a.all_B); + auto* ptr_C = reinterpret_cast(a.all_C); + // Hopper keeps the expert as an explicit tensor mode. A tail TMA box can + // therefore go OOB only within that expert (loads zero-fill) instead of + // spilling into the next expert's flat rows. + auto tma_load_b = [&]() { + if constexpr (Compute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_B), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + a.weight_expert_stride)), + typename Traits1::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_B), + make_shape(static_cast(total_n_rows_1), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits1::SmemLayoutW_1{}); + } + }(); + auto tma_load_c = [&]() { + if constexpr (Compute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_C), + make_shape(static_cast(intermediate_dim), + static_cast(hidden_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(hidden_dim), + Int<1>{}, + a.weight_expert_stride)), + typename Traits1::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_C), + make_shape(static_cast(total_n_rows_1), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits1::SmemLayoutW_1{}); + } + }(); + + // New design's TMA store box = per-WG epi slot (AtomTileM, EpiChunkN). + // + // Shape / stride values are cast to int64_t so CUTE's layout math + // (cosize, byte-offset, total-elements) uses 64-bit arithmetic. The + // Z tensor in particular can exceed INT_MAX elements when + // grid_x · ZBuf · TileM · intermediate_dim > 2^31 — e.g. NSplit=2 + // ZBuf=16 at intermediate_dim=16384 wraps to a negative stride and + // IMAs the kernel. Casting one operand promotes the multiply to + // int64_t; we cast both to be safe across CUTE versions. + auto tma_store_z = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.z_buf), + make_shape(static_cast(buf.z_rows), + static_cast(intermediate_dim)), + make_stride(static_cast(intermediate_dim), Int<1>{})), + typename Traits1::SmemLayoutStoreSlot{}); + + auto tma_load_z = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.z_buf)), + make_shape(static_cast(buf.z_rows), + static_cast(intermediate_dim)), + make_stride(static_cast(intermediate_dim), Int<1>{})), + typename Traits2::SmemLayoutZ_1{}); + + auto* ptr_A = reinterpret_cast(a.all_A); + // Same expert isolation for the down projection, including K-tail loads. + auto tma_load_a = [&]() { + if constexpr (Compute == 90) { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_A), + make_shape(static_cast(hidden_dim), + static_cast(intermediate_dim), + static_cast(experts_per_pe)), + make_stride( + static_cast(intermediate_dim), + Int<1>{}, + static_cast(hidden_dim) * intermediate_dim)), + typename Traits2::SmemLayoutW_1{}); + } else { + return make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(ptr_A), + make_shape(static_cast(total_n_rows_2), + static_cast(intermediate_dim)), + make_stride(static_cast(intermediate_dim), Int<1>{})), + typename Traits2::SmemLayoutW_1{}); + } + }(); + + auto tma_store_y = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(y_buf_ptr), + make_shape(static_cast(max_total_slots), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits2::SmemLayoutStore_1{}); + + // Remote staging is a flat list of L = MC · kCommNumStages slots, + // each TileM rows tall. MC = num_blocks / NC (== gridDim.x / NC) — MUST + // match the ring size everywhere else (allocate_moe_buffers, the per-launch + // memset, the consumer's ring_len). Sizing this TMA descriptor on n_gemm + // would leave the wrap-around slots outside its row extent (TMA OOB) when + // num_blocks/NC > n_gemm/NC (num_sms % NSplit != 0). + int mc_remote = num_blocks / Config::kNC; + int remote_total_rows = mc_remote * Config::kCommNumStages * kTileM; + + auto tma_load_x_remote = make_tma_copy(TmaLoadAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(buf.src_staging)), + make_shape(static_cast(remote_total_rows), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits1::SmemLayoutX_1{}); + + auto tma_store_y_remote = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(buf.dst_staging), + make_shape(static_cast(remote_total_rows), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits2::SmemLayoutStore_1{}); + + // ---- GET-via-TMA descriptors (FWD-only, raw CUtensorMap) ----------------- + // SAME-HOST source maps over nvshmem_ptr(x_sorted, peer) + one dst map over + // local src_staging. TMA get is enabled only for same-host NVLink P2P peers + // (cross-host peers fall back to warp getmem); also needs divisible + // hidden_dim + smem budget (checked below). + // + // src_x_desc is indexed by INTRA-HOST local rank (pe % gpus_per_node), so the + // array only ever needs gpus_per_node (≤ kGetMaxPes) entries regardless of + // total PE count. Assumes host-contiguous, host-aligned ranks (true for a + // WORLD-contiguous team); the do_get / RemoteIter is_local predicate uses the + // same arithmetic. + GetTmaDescs get_descs{}; + get_descs.enabled = 0; + int gpus_per_node = nvshmem_team_n_pes(NVSHMEMX_TEAM_NODE); + LIGER_CHECK(gpus_per_node >= 1 && gpus_per_node <= liger::kGetMaxPes, + "gpus_per_node (", gpus_per_node, ") must be in [1, kGetMaxPes=", + liger::kGetMaxPes, "]"); + get_descs.gpus_per_node = gpus_per_node; + get_descs.my_pe = my_pe; + constexpr int kGetBoxRows = + Config::kTileM / (liger::kNumGetWarpsFwd * Config::kNC); + size_t get_bounce_bytes = (size_t) + liger::GetBounce::kTotalBytes; + { + bool ok = (hidden_dim % liger::kGetKChunk == 0); + for (int tp = 0; ok && tp < num_pes; ++tp) { + // Only same-host peers are P2P-reachable → only they get a TMA + // descriptor (indexed by local rank). Cross-host peers skip it and + // use getmem at runtime. + bool same_host = (tp / gpus_per_node) == (my_pe / gpus_per_node); + if (!same_host) continue; + int gpe = nvshmem_team_translate_pe(team, tp, NVSHMEM_TEAM_WORLD); + void* pp = nvshmem_ptr(buf.x_sorted, gpe); // self → own ptr (non-null) + if (pp == nullptr) { ok = false; break; } // same-host but no P2P → disable TMA + ok = make_get_tma_map(get_descs.src_x_desc[tp % gpus_per_node], pp, + hidden_dim, max_total_slots, kGetBoxRows); + } + if (ok) + ok = make_get_tma_map(get_descs.dst_staging_desc, buf.src_staging, + hidden_dim, remote_total_rows, kGetBoxRows); + get_descs.enabled = ok ? 1 : 0; + if (const char* e = std::getenv("LIGER_GET_TMA")) + if (e[0] == '0') get_descs.enabled = 0; + } + + // ---- Direct-to-peer Y store descriptors (FWD-only) ----------------------- + // Per-peer TMA-store maps over nvshmem_ptr(y_buf, peer) for SAME-HOST peers + // (self included), indexed by intra-host local rank. When a tile is_local, + // MLP2 stores Y straight into the destination's symmetric local_output via + // its peer descriptor (token-tile row), bypassing dst_staging + the put warp + // (which then only rubber-stamps). Same shape/layout as tma_store_y; only the + // base pointer differs (the peer's y_buf). For single-PE the only peer is + // self → desc[my_pe%gpus_per_node] == tma_store_y. + PeerYStoreDescs peer_y{}; + peer_y.enabled = 1; + for (int tp = 0; tp < num_pes; ++tp) { + bool same_host = (tp / gpus_per_node) == (my_pe / gpus_per_node); + if (!same_host) continue; // cross-host peers keep the dst_staging+putmem path + int gpe = nvshmem_team_translate_pe(team, tp, NVSHMEM_TEAM_WORLD); + void* py = nvshmem_ptr(buf.y_buf, gpe); // same-host → P2P ptr (non-null on NVLink) + LIGER_CHECK(py != nullptr, + "direct-to-peer Y store: nvshmem_ptr(y_buf, same-host pe ", tp, + ") is null (no P2P) — unsupported for this build"); + peer_y.desc[tp % gpus_per_node] = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(reinterpret_cast(py)), + make_shape(static_cast(max_total_slots), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + typename Traits2::SmemLayoutStore_1{}); + } + + // TMA descriptor for the final combine output [num_tokens, hidden_dim]. + // Box shape matches combine_tokens_tma's WG tile: 32 tokens × 256 hidden. + auto* ptr_combine_out = output_ptr; + auto tma_store_y_combine = make_tma_copy(TmaStoreAtom{}, + make_tensor(make_gmem_ptr(ptr_combine_out), + make_shape(static_cast(num_tokens), + static_cast(hidden_dim)), + make_stride(static_cast(hidden_dim), Int<1>{})), + CombineSmemLayoutSlot{}); + + // ---- Dynamic smem -------------------------------------------------------- + int sort_num_warps = kNumThreads / 32; + size_t smem_sort = sort_num_warps * num_experts * sizeof(int); + + size_t smem_mlp = sizeof(MoeSmem); + smem_mlp = (smem_mlp + sizeof(int) - 1) & ~(sizeof(int) - 1); + smem_mlp += num_pes * (experts_per_pe + 1) * sizeof(int); // remote_offsets + smem_mlp += num_pes * experts_per_pe * sizeof(int); // remote_counts + + // GET-via-TMA bounce: append a 128-aligned region (must match the kernel's + // get_bounce offset). Disable the TMA-get path if it would exceed the + // opt-in dynamic-smem cap (227 KiB) — that config just keeps warp getmem. + if (get_descs.enabled) { + size_t mlp_with_bounce = ((smem_mlp + 127) & ~((size_t)127)) + get_bounce_bytes; + if (mlp_with_bounce <= (size_t)227 * 1024) smem_mlp = mlp_with_bounce; + else get_descs.enabled = 0; + } + + // Combine phase reinterprets the same SMEM region as CombineSmem. + constexpr int kCombineNumWGs = + Config::kNumThreads / (kCombineWarpsPerWG * kCombineWarpSize); + size_t smem_combine = sizeof(CombineSmem); + + size_t smem_size = std::max({smem_sort, smem_mlp, smem_combine}); + + // ---- Launch fused kernel ----------------------------------------------- + auto kernel_fn = moe_fused_kernel; + + if (cudaError_t e = cudaFuncSetAttribute(kernel_fn, + cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: cudaFuncSetAttribute(MaxDynamicSharedMemorySize=", smem_size, + ") failed: ", cudaGetErrorString(e)); + + // Reset per-call buffers in a single fused-memset kernel — one CTA per + // region. Replaces 11 cudaMemsetAsync calls that would otherwise + // serialize on the stream. + // + // Notes: + // - dispatch_done / mlp_done are already reset upstream in + // allocate_moe_buffers (the original "monotonic, never reset" design + // was retired — wait_until could exit on stale residual signal + // values in multi-PE shape-transition workloads). + // - all_expert_offsets stays unreset (overwritten by + // broadcast_expert_offsets on every call). + // - y_buf stays unreset (slots referenced by token_expert_slots[t,k] + // are always written by local MLP or the source PE's NVSHMEM + // put-back; padding slots are never read by combine_tokens_tma). + // - src_staging / dst_staging / z_buf intentionally NOT zero'd (see + // comment block below). + // + // Per-region reasons: + // - cta_counter: PeSync.barrier() does `% num_ctas`. Residual values + // from a prior call with different num_ctas misalign to the new + // modulus, hanging barrier() forever. + // - tile_expert_ids (0xff): sentinel-fill so unfetched slots produce + // expert_id = -1; matches moe_bwd.cu's 0xff fill of + // tile_expert_ids_{x,dy}. Without this, stale slot values from prior + // calls can match a valid expert id and confuse the MLP-side remote + // iterator at higher PE counts. + // - src/dst_ready/consumed: atomically accumulated; stale ready/ + // consumed values make the consumer-acquire wait fall through. + // - tile_expert_counts / cta_sums: atomicAdd'd during sort; must + // start at 0. + { + auto& scfg = get_symm_config(); + // L = MC · kCommNumStages, MC = num_blocks / NC (== gridDim.x / NC). + // MUST match the ring size in allocate_moe_buffers and the device-side + // comm/consumer stride — otherwise this memset under-covers the signal / + // tile_expert_ids arrays (slots [n_gemm/NC·CS, gridDim.x/NC·CS) left + // uninitialized) and the consumer reads stale src_ready / garbage expert + // ids for those slots (see the OOB note in allocate_moe_buffers). + int mc_local = num_blocks / Config::kNC; + int comm_l = mc_local * Config::kCommNumStages; + size_t sig_bytes = + static_cast(comm_l) * Config::kSubTiles * sizeof(int); + size_t eid_bytes = static_cast(comm_l) * sizeof(int); + int max_m_tiles = scfg.max_total_slots / kTileM; + int max_tile_counts = max_m_tiles * scfg.max_num_experts; + + FusedMemsetArgs ma{}; + fused_memset_add(ma, buf.phase_counter, + ((num_blocks + 1) / 2) * sizeof(int), 0x00); + fused_memset_add(ma, buf.cta_counter, + sizeof(int), 0x00); + fused_memset_add(ma, buf.cta_done, + num_blocks * sizeof(int), 0x00); + fused_memset_add(ma, buf.sorted_token_ids, + static_cast(max_total_slots) * sizeof(int), 0xff); + fused_memset_add(ma, buf.src_ready, sig_bytes, 0x00); + fused_memset_add(ma, buf.src_consumed, sig_bytes, 0x00); + fused_memset_add(ma, buf.dst_ready, sig_bytes, 0x00); + fused_memset_add(ma, buf.dst_consumed, sig_bytes, 0x00); + fused_memset_add(ma, buf.tile_expert_ids, eid_bytes, 0xff); + fused_memset_add(ma, buf.tile_expert_counts, + static_cast(max_tile_counts) * sizeof(int), 0x00); + fused_memset_add(ma, buf.cta_sums, + static_cast(num_blocks) * scfg.max_num_experts * sizeof(int), 0x00); + fused_memset_launch(ma, stream); + } + + // src_staging / dst_staging / z_buf intentionally NOT zero'd here: + // + // src_staging: comm warps fill the slots THIS CTA fetches; unfetched + // slots are gated by tile_expert_ids = -1 (set via the 0xff memset + // above) and RemoteIter::has_next() walks only per_cta_tiles, so the + // remote MLP never reads unfetched slots. Stale prev-call contents + // are inert under both bounds checks. + // + // dst_staging: write-then-read for routed slots only. Remote MLP + // writes Y for tiles RemoteIter returns; comm warps put back only + // those tiles. Padding/unused slots are never read at either end. + // + // z_buf: per-tile write→read within a single MLP launch. LocalIter is + // range-bounded to [local_e_start, local_e_end) experts via + // set_range; the remote pass is bounded by per_cta_tiles. Tiles + // outside the current call's active range are never written and + // never read — buffer-pool size growth from previous (larger) calls + // is invisible to the kernel. + // + // The tile_expert_ids 0xff fill above is what makes this safe — keep + // that. + + kernel_fn<<>>( + tma_load_x, tma_load_b, tma_load_c, tma_store_z, + tma_load_z, tma_load_a, tma_store_y, + tma_load_x_remote, tma_store_y_remote, + tma_store_y_combine, + mlp_dims, p, static_nsplit, get_descs, peer_y); + if (cudaError_t e = cudaGetLastError(); e != cudaSuccess) + LIGER_FAIL_CUDA("moe_fwd: kernel launch failed for grid=", grid.x, + " threads=", kNumThreads, " smem=", smem_size, ": ", + cudaGetErrorString(e)); + + // Post-launch global barrier: ensures all NVSHMEM operations (tile puts, + // signal ops) across all PEs are complete before any PE starts the next + // call. Without this, a fast PE could overwrite symmetric buffers + // (x_sorted, y_buf, tile signals) while a slow PE is still reading them. + // + // Stream-ordered variant — capturable into a CUDA graph. The host-side + // nvshmem_team_sync(team) would do an internal cudaStreamSynchronize and + // therefore fail with "operation not permitted when stream is capturing". + //if (num_pes > 1) { + // nvshmemx_barrier_on_stream(team, stream); + //} + + // Hand the persistent symmetric buffers back to the ABI wrapper so it can + // alias them as the *_out_symm views (shapes use the static MoeSymmConfig + // max; the views do NOT own the storage — SymmetricMemoryStack does). Y and + // the two sort outputs were written in place into the caller-owned buffers. + *a.x_sorted_out = buf.x_sorted; + *a.y_buf_out = y_buf_ptr; + *a.all_expert_offsets_out = buf.all_expert_offsets; + if (a.all_expert_counts_out != nullptr) + *a.all_expert_counts_out = buf.all_expert_counts; +} + +// ── Runtime dispatch table for tuned-config lookup ──────────────────── + +using MoeFwdFn = void (*)(const MoeFwdArgs&, int static_nsplit); +static constexpr int kFwdRuntimeNsplitSeed = 8; + +struct DispatchEntry { + int Compute; + int TileN1, TileK1, Stages1, EpiChunkN1; + int TileN2, TileK2, Stages2, EpiChunkN2; + int ZBufferSlots, CommNumStages; + int TileM; // tuned GEMM-tile field (legacy dispatch-table name) + int GemmTileM; // GEMM tile ∈ {64,128} + MoeFwdFn fn; +}; + +// 13-arg core: explicit GemmTileM. 12-arg shim sets GemmTileM = comm TileM. +#define LIGER_MOE_DISPATCH_ENTRY_G_C(Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) \ + { Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM, \ + &moe_fused_fwd_bf16 }, +#define LIGER_MOE_DISPATCH_ENTRY_C(Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) \ + LIGER_MOE_DISPATCH_ENTRY_G_C(Compute, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, TileM) +#define LIGER_MOE_DISPATCH_ENTRY_G_SM90(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) \ + LIGER_MOE_DISPATCH_ENTRY_G_C(90, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) +#define LIGER_MOE_DISPATCH_ENTRY_SM90(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) \ + LIGER_MOE_DISPATCH_ENTRY_C(90, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) +#define LIGER_MOE_DISPATCH_ENTRY_G_SM100(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) \ + LIGER_MOE_DISPATCH_ENTRY_G_C(100, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM, GemmTileM) +#define LIGER_MOE_DISPATCH_ENTRY_SM100(TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) \ + LIGER_MOE_DISPATCH_ENTRY_C(100, TN1, TK1, S1, EC1, TN2, TK2, S2, EC2, ZBuf, CStages, TileM) + +static const DispatchEntry kDispatchTable[] = { +#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 90 + LIGER_MOE_FWD_DISPATCH_CONFIGS_SM90(LIGER_MOE_DISPATCH_ENTRY_SM90, LIGER_MOE_DISPATCH_ENTRY_G_SM90) +#endif +#if LIGER_CUTE_DISPATCH_COMPUTE == 0 || LIGER_CUTE_DISPATCH_COMPUTE == 100 + LIGER_MOE_FWD_DISPATCH_CONFIGS_SM100(LIGER_MOE_DISPATCH_ENTRY_SM100, LIGER_MOE_DISPATCH_ENTRY_G_SM100) +#endif +}; +#undef LIGER_MOE_DISPATCH_ENTRY_SM100 +#undef LIGER_MOE_DISPATCH_ENTRY_G_SM100 +#undef LIGER_MOE_DISPATCH_ENTRY_SM90 +#undef LIGER_MOE_DISPATCH_ENTRY_G_SM90 +#undef LIGER_MOE_DISPATCH_ENTRY_C +#undef LIGER_MOE_DISPATCH_ENTRY_G_C +static constexpr int kNumDispatchEntries = + sizeof(kDispatchTable) / sizeof(kDispatchTable[0]); + +} // namespace liger + +// Joint fwd+bwd tuning table (auto-generated by benchmarks/tune_moe_fwd_bwd). +// We only consume the Fwd_* fields here; moe_bwd.cu consumes the same table's +// Bwd_* fields. Picking fwd and bwd from a single table guarantees the two +// directions are tuned as a pair (so their TileM agrees and their per-shape +// templates were measured together, not independently). +#include "moe_fwd_bwd_tuning_configs_sm90.cuh" + +namespace liger { + +// TunedConfig is the fwd-only projection of TunedConfigFwdBwd. The rest of +// the dispatch code (validity check, dispatch-entry match) was written +// against this struct and is reused unchanged. +struct TunedConfig { + int Compute; + int TK, TKE, D, I; + int TileN1, TileK1, Stages1, EpiChunkN1; + int TileN2, TileK2, Stages2, EpiChunkN2; + int ZBufferSlots, CommNumStages; + int TileM; // comm tile + int GemmTileM; // gemm tile (auto path: == TileM; FORCE_CONFIG may set 64) +}; + +// Project a TunedConfigFwdBwd row down to its fwd template params. The tuned +// table has no GemmTileM column, so the auto path keeps gemm == comm (TileM). +static TunedConfig project_fwd(const TunedConfigFwdBwd& c, int compute) { + return TunedConfig{ + compute, + c.TK, c.TKE, c.D, c.I, + c.Fwd_TileN1, c.Fwd_TileK1, c.Fwd_Stages1, c.Fwd_EpiChunkN1, + c.Fwd_TileN2, c.Fwd_TileK2, c.Fwd_Stages2, c.Fwd_EpiChunkN2, + c.Fwd_ZBufferSlots, c.Fwd_CommNumStages, + c.Fwd_TileM, c.Fwd_TileM, + }; +} + +// Find the dispatch table entry matching a TunedConfig's fields. +static const DispatchEntry* find_dispatch_entry(const TunedConfig& c) { + for (int i = 0; i < kNumDispatchEntries; ++i) { + const auto& e = kDispatchTable[i]; + if (e.Compute == c.Compute && + e.TileN1 == c.TileN1 && e.TileK1 == c.TileK1 && e.Stages1 == c.Stages1 && + e.EpiChunkN1 == c.EpiChunkN1 && + e.TileN2 == c.TileN2 && e.TileK2 == c.TileK2 && e.Stages2 == c.Stages2 && + e.EpiChunkN2 == c.EpiChunkN2 && + e.ZBufferSlots == c.ZBufferSlots && e.CommNumStages == c.CommNumStages && + e.TileM == c.TileM && e.GemmTileM == c.GemmTileM) { + return &e; + } + } + return nullptr; +} + +// Mirrors shape_valid() in tune_moe.cu: structural constraints the kernel +// requires on (D, I) for a given config. Configs that violate these for +// the runtime shape are skipped during nearest-neighbor lookup so we +// don't dispatch to a template instantiation the kernel can't run. +// +// Cooperative 2-WG consumer layout switches on TileM: +// - TileM=128 (M-split, Layout<_2,_1,_1>): each WG owns AtomTileM=64 +// rows × TileN cols. EpiChunkN must divide TileN. +// - TileM=64 (N-split, Layout<_1,_2,_1>): each WG owns TileM=64 rows +// × TileN/2 cols. EpiChunkN must divide TileN/2; TileN must be even. +// - Hopper uses ceil-div tile counts; SM100 retains exact divisibility. +// - store_buf is (AtomTileM=64, EpiChunkN) per WG in both modes. +static bool tuned_config_valid(int D, int I, const TunedConfig& c) { + if (c.TileM != 64 && c.TileM != 128) return false; + // The current Blackwell UMMA MLP1 consumer only supports the 128-row comm + // tile. Keep mirrored SM90 TileM=64 tuning rows out of SM100 auto-dispatch + // until that TMEM epilogue layout is implemented. + if (c.Compute == 100 && c.TileM != 128) return false; + if (c.Compute == 100) { + if (D % c.TileK1 != 0) return false; + if (I % c.TileN1 != 0) return false; + if (I % c.TileK2 != 0) return false; + if (D % c.TileN2 != 0) return false; + } + if (D % 8 != 0 || I % 8 != 0) return false; + // NC selection mirrors moe_fused_fwd_bf16: TileM=128 → NC=4, + // TileM=64 → NC=2. With the ticket-based GEMM iterator the only + // runtime requirement is NC | (gridDim.x · gridDim.y); runtime NS / + // CommNumStages no longer need to share any divisibility with NC. + const int nc = (c.TileM == 128) ? 4 : 2; + (void)nc; + // EpiChunkN must divide the per-WG N extent (WgTileN). + const int wg_tile_n1 = (c.TileM == 128) ? c.TileN1 : c.TileN1 / 2; + const int wg_tile_n2 = (c.TileM == 128) ? c.TileN2 : c.TileN2 / 2; + if (wg_tile_n1 <= 0 || wg_tile_n1 % c.EpiChunkN1 != 0) return false; + if (wg_tile_n2 <= 0 || wg_tile_n2 % c.EpiChunkN2 != 0) return false; + // N-split requires TileN to be evenly halvable (TileN/2 must be ≥ the + // WGMMA atom-N min, which is 8 for bf16 but 64 for our smem layout). + if (c.TileM == 64 && (c.TileN1 % 2 != 0 || c.TileN2 % 2 != 0)) return false; + + // SMEM cap (H100 = 228 KiB; budget 224 KiB inside MoeFusedConfig). + // Per-CTA, bf16 (2 B/elem). TileM scales the X/Z staging: + // MLP1 = 2 * (TileM·TK1·S1 + 2·TN1·TK1·S1 + 2·64·EC1) + // MLP2 = 2 * (TileM·TK2·S2 + TN2·TK2·S2 + 2·64·EC2) + // smem ≈ max(MLP1, MLP2) + int mlp1_bytes = 2 * (c.TileM * c.TileK1 * c.Stages1 + + 2 * c.TileN1 * c.TileK1 * c.Stages1 + + 2 * 64 * c.EpiChunkN1); + int mlp2_bytes = 2 * (c.TileM * c.TileK2 * c.Stages2 + + c.TileN2 * c.TileK2 * c.Stages2 + + 2 * 64 * c.EpiChunkN2); + int smem = mlp1_bytes > mlp2_bytes ? mlp1_bytes : mlp2_bytes; + constexpr int kSmemBudget = 224 * 1024; + if (smem > kSmemBudget) return false; + return true; +} + +static bool find_tuned_table(int compute, int n_pes, + const TunedConfigFwdBwd** table, + int* table_n) { + const TunedConfigFwdBwdTable* tables = + (n_pes <= 1) ? kTunedConfigTablesSingle : kTunedConfigTablesMulti; + const int tables_n = + (n_pes <= 1) ? kNumTunedConfigTablesSingle : kNumTunedConfigTablesMulti; + for (int i = 0; i < tables_n; ++i) { + if (tables[i].Compute == compute) { + *table = tables[i].configs; + *table_n = tables[i].count; + return *table_n > 0; + } + } + return false; +} + +// Nearest-neighbor lookup in log-space (TK, TKE, D, I), where TK = T*top_k +// is the total routed-token count (the actual M-axis the kernel sees) and +// TKE = TK / E_local (E_local = E / n_pes) is the per-LOCAL-expert avg +// K-range. Adding TKE to the distance lets the dispatcher distinguish +// shapes whose mlp3/mlp4 work distribution differs even when (TK, D, I) +// is the same — e.g. same total tokens, different expert count or PE +// count. +// Configs that violate structural constraints for the runtime (D, I) are +// skipped — the nearest tuned point is not always valid for the actual +// shape (e.g. tuned at I=4096 with TileN1=256, runtime I=768). +// Out-param API instead of pointer-into-static-array because each candidate +// TunedConfig is projected on the fly from kTunedConfigsFwdBwd (no stable +// address). Returns true when a valid config was found. +static bool find_nearest_tuned(int compute, int TK, int TKE, int D, int I, int n_pes, + TunedConfig* out) { + // World-size class: single-GPU (n_pes <= 1) and multi-GPU tables are tuned + // and stored separately, since the comm path (and thus the optimal NSplit/ + // CommNumStages/GEMM-tile) differs sharply between one PE and many. + const TunedConfigFwdBwd* table = nullptr; + int table_n = 0; + if (!find_tuned_table(compute, n_pes, &table, &table_n)) return false; + bool found = false; + double best_dist = 1e30; + double lTK = std::log((double)TK); + double lTKE = std::log((double)TKE); + double lD = std::log((double)D); + double lI = std::log((double)I); + for (int i = 0; i < table_n; ++i) { + TunedConfig c = project_fwd(table[i], compute); + if (!tuned_config_valid(D, I, c)) continue; + // Guard against TKE=0 in legacy rows (would yield -inf in log). + double dTK = std::log((double)c.TK) - lTK; + double dTKE = c.TKE > 0 + ? std::log((double)c.TKE) - lTKE + : 0.0; // legacy / placeholder rows contribute 0 to the TKE axis + double dD = std::log((double)c.D) - lD; + double dI = std::log((double)c.I) - lI; + double d = dTK*dTK + dTKE*dTKE + dD*dD + dI*dI; + if (d < best_dist) { best_dist = d; *out = c; found = true; } + } + return found; +} + +// ============================================================================ +// Auto-dispatch wrapper — pure tuned-table lookup +// ============================================================================ +// +// Picks the nearest tuned config (log-space distance over (T, D, I)) from +// moe_fwd_bwd_tuning_configs.cuh and dispatches to the matching template +// instantiation in kDispatchTable. Both tables are populated by the +// tune_moe_fwd_bwd sweep — +// shapes outside the swept range error out instead of falling back to a +// heuristic. Add the missing shape to the sweep (or to LIGER_MOE_TUNE_CONFIGS +// directly) to extend coverage. + +// Internal auto-dispatch: picks the nearest tuned config for the runtime shape +// (or honours LIGER_MOE_FORCE_CONFIG) and invokes the matching template +// instantiation via the dispatch table. Writes the chosen comm TileM to +// *chosen_tile_m. All shape info comes from the MoeFwdArgs bundle; the launcher +// writes the outputs in place. +void moe_fused_fwd_dispatch(const MoeFwdArgs& a, int* chosen_tile_m) { + const int T = a.num_tokens; + const int D = a.hidden_dim; + const int I = a.intermediate_dim; + const int top_k = a.top_k; + const int num_experts = a.num_experts; + const int TK = T * top_k; // total routed tokens — the M-axis the kernel sees + // TKE = TK / E_local where E_local = num_experts / n_pes (per-LOCAL-expert + // avg K-range). Matches the tuner's definition so the nearest-neighbor + // lookup picks rows recorded at comparable per-PE load. + const int n_pes = nvshmem_team_n_pes(a.team); + const int E_local = (n_pes > 0) ? std::max(1, num_experts / n_pes) : num_experts; + // The lookup operates in log space, so clamp the per-expert average to + // the smallest representable workload for tiny but valid token batches. + const int TKE = std::max(1, TK / E_local); + const int compute = moe_detect_compute_dispatch_key(a.device, "moe_fused_fwd_bf16_auto"); + + // ── Override path: LIGER_MOE_FORCE_CONFIG=TN1,TK1,S1,EC1,TN2,TK2,S2,EC2,ZB,CS,TM[,GemmTM] + // Lets a test suite sweep every compiled config without rebuilding. The + // optional final int is GemmTileM (comm/gemm tile decouple); omitted → == TM. + if (const char* s = std::getenv("LIGER_MOE_FORCE_CONFIG")) { + TunedConfig tc{}; + tc.Compute = compute; + tc.TileM = 128; // default when omitted (backward compat) + tc.GemmTileM = -1; // sentinel → default to TileM after parse + int n = std::sscanf(s, "%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d,%d", + &tc.TileN1, &tc.TileK1, &tc.Stages1, &tc.EpiChunkN1, + &tc.TileN2, &tc.TileK2, &tc.Stages2, &tc.EpiChunkN2, + &tc.ZBufferSlots, &tc.CommNumStages, + &tc.TileM, &tc.GemmTileM); + LIGER_CHECK(n == 10 || n == 11 || n == 12, + "LIGER_MOE_FORCE_CONFIG must have 10, 11 or 12 comma-separated ints " + "(TN1,TK1,S1,EC1,TN2,TK2,S2,EC2,ZB,CS[,TM[,GemmTM]]); got '", s, "'"); + if (tc.GemmTileM <= 0) tc.GemmTileM = tc.TileM; // omitted → gemm == comm + const DispatchEntry* de = find_dispatch_entry(tc); + LIGER_CHECK(de != nullptr, + "LIGER_MOE_FORCE_CONFIG=", s, " is not in the compiled dispatch " + "table. Add it to LIGER_MOE_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); + de->fn(a, kFwdRuntimeNsplitSeed); + *chosen_tile_m = de->Compute == 100 + ? MoeCommGeometry<100>::TileM + : MoeCommGeometry<90>::TileM; + return; + } + + TunedConfig tc{}; + LIGER_CHECK(find_nearest_tuned(compute, TK, TKE, D, I, n_pes, &tc), + "moe_fused_fwd_bf16_auto: no valid tuned config for compute=", compute, + " shape (T=", T, + " K=", top_k, " TK=", TK, " TKE=", TKE, " D=", D, " I=", I, + "). Either moe_fwd_bwd_tuning_configs.cuh is empty (run tune_moe_fwd_bwd), " + "or every tuned entry violates structural constraints for this (D, I) " + "(tile/layout compatibility and shared-memory limits). " + "Add a compatible config to LIGER_MOE_TUNE_CONFIGS and re-tune."); + + const DispatchEntry* de = find_dispatch_entry(tc); + LIGER_CHECK(de != nullptr, + "moe_fused_fwd_bf16_auto: tuned config (Compute=", tc.Compute, + " TN1=", tc.TileN1, "/", tc.TileK1, "/", tc.Stages1, "/EC", tc.EpiChunkN1, + " TN2=", tc.TileN2, "/", tc.TileK2, "/", tc.Stages2, "/EC", tc.EpiChunkN2, + " ZB=", tc.ZBufferSlots, " CS=", tc.CommNumStages, + ") for shape (T=", T, " K=", top_k, " TK=", TK, " D=", D, " I=", I, + ") is not in the compiled dispatch table. " + "Add it to LIGER_MOE_TUNE_CONFIGS in moe_fwd_bwd_tune_configs.hpp."); + + de->fn(a, kFwdRuntimeNsplitSeed); + *chosen_tile_m = de->Compute == 100 + ? MoeCommGeometry<100>::TileM + : MoeCommGeometry<90>::TileM; +} + +// ───────────────────────────────────────────────────────────────────── +// Default (untuned) wrapper — single hardcoded config per TileM. Use +// this for the regular bench_moe / test_moe paths so they don't depend +// on the autotuner dispatch table. Returns the same 6-tuple as the auto +// wrapper. +// +// Each config below must have a matching explicit instantiation. The +// TileM=64 row is in LIGER_MOE_TUNE_CONFIGS (so LIGER_MOE_INSTANTIATE +// emits it automatically); the TileM=128 row is emitted as a direct +// `template` instantiation right below this function. +// ───────────────────────────────────────────────────────────────────── + +// Returns the compiled MoE config table as a list of 12-int rows +// (Compute, TileN1, TileK1, Stages1, EpiChunkN1, TileN2, TileK2, Stages2, +// EpiChunkN2, ZBufferSlots, CommNumStages, TileM). Kept for the tuner / a future +// sweep binding so they stay in sync with the dispatch table without manual +// duplication. Internal (std::vector) — not part of the flat ABI. +std::vector> +moe_list_compiled_configs() { + std::vector> out; + out.reserve(kNumDispatchEntries); + for (int i = 0; i < kNumDispatchEntries; ++i) { + const auto& e = kDispatchTable[i]; + out.emplace_back(e.Compute, e.TileN1, e.TileK1, e.Stages1, e.EpiChunkN1, + e.TileN2, e.TileK2, e.Stages2, e.EpiChunkN2, + e.ZBufferSlots, e.CommNumStages, e.TileM); + } + return out; +} + +} // namespace liger + +// ============================================================================ +// Flat extern "C" control ABI (declared in liger_cute/moe.h) +// ============================================================================ + +namespace { + +using liger_cute::detail::guarded; + +} // namespace + +extern "C" { + +liger_cute_status_t liger_cute_moe_get_symm_config(liger_cute_moe_symm_config_t* out) { + return guarded([&]() -> liger_cute_status_t { + LIGER_CHECK(out != nullptr, "moe_get_symm_config: out is null"); + const liger::MoeSymmConfig& cfg = liger::get_symm_config(); + out->max_total_slots = cfg.max_total_slots; + out->max_num_experts = cfg.max_num_experts; + out->hidden_dim = cfg.hidden_dim; + out->num_pes = cfg.num_pes; + out->experts_per_pe = cfg.experts_per_pe; + out->max_top_k = cfg.max_top_k; + out->initialized = cfg.initialized ? 1 : 0; + return LIGER_CUTE_OK; + }); +} + +liger_cute_status_t liger_cute_moe_configure_symmetric( + const int max_tokens, const int hidden_dim, const int max_num_experts, + const int max_top_k, const int num_pes, const int num_hosts, + const int gpus_per_host) { + return guarded([&]() -> liger_cute_status_t { + liger::moe_configure_symmetric(max_tokens, hidden_dim, max_num_experts, + max_top_k, num_pes, num_hosts, gpus_per_host); + return LIGER_CUTE_OK; + }); +} + +liger_cute_status_t liger_cute_moe_pop_fwd(void) { + return guarded([&]() -> liger_cute_status_t { + liger::moe_pop_fwd(); + return LIGER_CUTE_OK; + }); +} + +} // extern "C" diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_symm_config.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_symm_config.cuh index a8827991..ff64a889 100644 --- a/liger_cute_kernels/csrc/core/src/moe/moe_symm_config.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/moe_symm_config.cuh @@ -1,52 +1,11 @@ #pragma once -#include "moe_comm_config.cuh" +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -// ============================================================================ -// MoeSymmConfig — fixed parameters for symmetric-memory sizing. -// -// Single source of truth for the MoE symmetric-session config, shared by -// moe.cu (forward) and moe_bwd.cu (backward). BOTH translation units alias -// the one object returned by get_symm_config(), so the layout MUST be defined -// exactly once — here. Previously each .cu defined its own copy; when moe.cu -// gained max_comm_stages / max_tile_m (2026-06-03, ca1b943) and moe_bwd.cu's -// copy wasn't updated, the two TUs read members at different offsets and the -// bwd path saw a stale `initialized` flag (spurious "Call -// moe_configure_symmetric before moe_bwd_fwd_bf16"). A single definition makes -// that drift impossible. -// -// Requires nvshmem_team_t — include before this header. -// Must be set once (via moe_configure_symmetric) before the first kernel call; -// all PEs must agree. -// ============================================================================ - -namespace liger { - -struct MoeSymmConfig { - int max_total_slots; // upper bound across all configs - int max_num_experts; // max experts across all configs - int hidden_dim; // fixed across configs - int num_pes; - int experts_per_pe; // max_num_experts / num_pes - int max_top_k; // max top_k across configs (exposed via the flat ABI) - nvshmem_team_t team; // NVSHMEM team - // Worst-case comm-staging shape across every config the symmetric - // session may run. The symmetric staging pool (moe_src/dst_staging) is - // a single shared key sized ONCE; because get_symmetric aborts on grow, - // it must be reserved at the largest CommNumStages × TileM any config - // uses (the tuner sweeps NS=16/CS=8). - // Per-config sizing would - // let a small first config lock in a buffer a later large config can't - // grow into. Same reasoning as sizing by max hidden_dim, extended to - // the (CS, TileM) axes. The bwd path doesn't read these, but they MUST - // stay in the layout so its get_symm_config() view matches the fwd's. - int max_comm_stages = kMaxMoeCommStages; - int max_tile_m = kMaxMoeCommTileM; - bool initialized = false; -}; - -// Single instance lives in moe.cu (a function-local static). moe_bwd.cu reuses -// it through this same declaration. -MoeSymmConfig& get_symm_config(); - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "moe_symm_config_sm90.cuh" +#else +#include "moe_symm_config_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm100.cuh new file mode 100644 index 00000000..d558d161 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm100.cuh @@ -0,0 +1,52 @@ +#pragma once + +#include "moe_comm_config_sm100.cuh" + +// ============================================================================ +// MoeSymmConfig — fixed parameters for symmetric-memory sizing. +// +// Single source of truth for the MoE symmetric-session config, shared by +// moe.cu (forward) and moe_bwd.cu (backward). BOTH translation units alias +// the one object returned by get_symm_config(), so the layout MUST be defined +// exactly once — here. Previously each .cu defined its own copy; when moe.cu +// gained max_comm_stages / max_tile_m (2026-06-03, ca1b943) and moe_bwd.cu's +// copy wasn't updated, the two TUs read members at different offsets and the +// bwd path saw a stale `initialized` flag (spurious "Call +// moe_configure_symmetric before moe_bwd_fwd_bf16"). A single definition makes +// that drift impossible. +// +// Requires nvshmem_team_t — include before this header. +// Must be set once (via moe_configure_symmetric) before the first kernel call; +// all PEs must agree. +// ============================================================================ + +namespace liger { + +struct MoeSymmConfig { + int max_total_slots; // upper bound across all configs + int max_num_experts; // max experts across all configs + int hidden_dim; // fixed across configs + int num_pes; + int experts_per_pe; // max_num_experts / num_pes + int max_top_k; // max top_k across configs (exposed via the flat ABI) + nvshmem_team_t team; // NVSHMEM team + // Worst-case comm-staging shape across every config the symmetric + // session may run. The symmetric staging pool (moe_src/dst_staging) is + // a single shared key sized ONCE; because get_symmetric aborts on grow, + // it must be reserved at the largest CommNumStages × TileM any config + // uses (the tuner sweeps NS=16/CS=8). + // Per-config sizing would + // let a small first config lock in a buffer a later large config can't + // grow into. Same reasoning as sizing by max hidden_dim, extended to + // the (CS, TileM) axes. The bwd path doesn't read these, but they MUST + // stay in the layout so its get_symm_config() view matches the fwd's. + int max_comm_stages = kMaxMoeCommStages; + int max_tile_m = kMaxMoeCommTileM; + bool initialized = false; +}; + +// Single instance lives in moe.cu (a function-local static). moe_bwd.cu reuses +// it through this same declaration. +MoeSymmConfig& get_symm_config(); + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm90.cuh new file mode 100644 index 00000000..e1a033c5 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm90.cuh @@ -0,0 +1,52 @@ +#pragma once + +#include "moe_comm_config_sm90.cuh" + +// ============================================================================ +// MoeSymmConfig — fixed parameters for symmetric-memory sizing. +// +// Single source of truth for the MoE symmetric-session config, shared by +// moe.cu (forward) and moe_bwd.cu (backward). BOTH translation units alias +// the one object returned by get_symm_config(), so the layout MUST be defined +// exactly once — here. Previously each .cu defined its own copy; when moe.cu +// gained max_comm_stages / max_tile_m (2026-06-03, ca1b943) and moe_bwd.cu's +// copy wasn't updated, the two TUs read members at different offsets and the +// bwd path saw a stale `initialized` flag (spurious "Call +// moe_configure_symmetric before moe_bwd_fwd_bf16"). A single definition makes +// that drift impossible. +// +// Requires nvshmem_team_t — include before this header. +// Must be set once (via moe_configure_symmetric) before the first kernel call; +// all PEs must agree. +// ============================================================================ + +namespace liger { + +struct MoeSymmConfig { + int max_total_slots; // upper bound across all configs + int max_num_experts; // max experts across all configs + int hidden_dim; // fixed across configs + int num_pes; + int experts_per_pe; // max_num_experts / num_pes + int max_top_k; // max top_k across configs (exposed via the flat ABI) + nvshmem_team_t team; // NVSHMEM team + // Worst-case comm-staging shape across every config the symmetric + // session may run. The symmetric staging pool (moe_src/dst_staging) is + // a single shared key sized ONCE; because get_symmetric aborts on grow, + // it must be reserved at the largest CommNumStages × TileM any config + // uses (the tuner sweeps NS=16/CS=8). + // Per-config sizing would + // let a small first config lock in a buffer a later large config can't + // grow into. Same reasoning as sizing by max hidden_dim, extended to + // the (CS, TileM) axes. The bwd path doesn't read these, but they MUST + // stay in the layout so its get_symm_config() view matches the fwd's. + int max_comm_stages = kMaxMoeCommStages; + int max_tile_m = kMaxMoeCommTileM; + bool initialized = false; +}; + +// Single instance lives in moe.cu (a function-local static). moe_bwd.cu reuses +// it through this same declaration. +MoeSymmConfig& get_symm_config(); + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused.cuh b/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused.cuh index 6f1054d3..7d5e9848 100644 --- a/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused.cuh @@ -1,244 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// SiLU backward (fused tile): -// dV = dZ * SiLU(U) -// dU = dZ * V * grad_SiLU(U) -// ═══════════════════════════════════════════════════════════════════ -// -// Element-wise backward of Z = SiLU(U) * V. Uses all MLP warps -// (0, 1, 4–11 = 320 threads), excluding communication warps 2–3. -// Vectorized with uint4 (8 bf16 per load/store). -// -// num_elems is a runtime constant (tokens × intermediate_dim). +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -#include "math.cuh" - -namespace liger { - -static constexpr int kSiluBwdThreads = 320; // 10 warps × 32 - -template -__device__ __forceinline__ void silu_bwd_fused( - const Element* __restrict__ dZ, - const Element* __restrict__ U, - const Element* __restrict__ V, - Element* __restrict__ dU, - Element* __restrict__ dV, - int num_elems) { - - constexpr int kVecSize = sizeof(uint4) / sizeof(Element); - - // Dense thread ID across 10 participating warps (0,1,4-11). - // Warps 2-3 (comm) have already returned before reaching this call. - int warp_id = threadIdx.x / 32; - int lane = threadIdx.x % 32; - int dense_warp = (warp_id < 2) ? warp_id : warp_id - 2; - int tid = dense_warp * 32 + lane; - - auto* dZ_v = reinterpret_cast(dZ); - auto* U_v = reinterpret_cast(U); - auto* V_v = reinterpret_cast(V); - auto* dU_v = reinterpret_cast(dU); - auto* dV_v = reinterpret_cast(dV); - - int vec_count = num_elems / kVecSize; - - for (int i = tid; i < vec_count; i += kSiluBwdThreads) { - uint4 dz_vec = dZ_v[i]; - uint4 u_vec = U_v[i]; - uint4 v_vec = V_v[i]; - - auto* dz_e = reinterpret_cast(&dz_vec); - auto* u_e = reinterpret_cast(&u_vec); - auto* v_e = reinterpret_cast(&v_vec); - - uint4 du_out, dv_out; - auto* du_e = reinterpret_cast(&du_out); - auto* dv_e = reinterpret_cast(&dv_out); - - CUTE_UNROLL - for (int j = 0; j < kVecSize; ++j) { - float dz = static_cast(dz_e[j]); - float u = static_cast(u_e[j]); - float v = static_cast(v_e[j]); - float f = fast_sigmoid(u); - float df = f * (1.0f + u * (1.0f - f)); - - dv_e[j] = from_float(dz * u * f); - du_e[j] = from_float(dz * v * df); - } - - dV_v[i] = dv_out; - dU_v[i] = du_out; - } -} - -// ═══════════════════════════════════════════════════════════════════ -// N-tile-aware strided variant: processes a [num_rows, num_cols] -// sub-matrix within a row-major [num_rows, row_stride] buffer. -// Handles non-contiguous columns efficiently with full thread -// utilization across all rows × cols. -// ═══════════════════════════════════════════════════════════════════ - -template -__device__ __forceinline__ void silu_bwd_fused_strided( - const Element* __restrict__ dZ_base, - const Element* __restrict__ U_base, - const Element* __restrict__ V_base, - Element* __restrict__ dU_base, - Element* __restrict__ dV_base, - int num_rows, - int num_cols, - int row_stride) { - - constexpr int kVecSize = sizeof(uint4) / sizeof(Element); - - int warp_id = threadIdx.x / 32; - int lane = threadIdx.x % 32; - int dense_warp = (warp_id < 2) ? warp_id : warp_id - 2; - int tid = dense_warp * 32 + lane; - - // Total vectors across all rows - int cols_per_vec_row = num_cols / kVecSize; - int total_vecs = num_rows * cols_per_vec_row; - - for (int i = tid; i < total_vecs; i += kSiluBwdThreads) { - int r = i / cols_per_vec_row; - int c = i % cols_per_vec_row; - int off = r * row_stride / kVecSize + c; // offset in uint4 units - - auto* dZ_v = reinterpret_cast(dZ_base); - auto* U_v = reinterpret_cast(U_base); - auto* V_v = reinterpret_cast(V_base); - auto* dU_v = reinterpret_cast(dU_base); - auto* dV_v = reinterpret_cast(dV_base); - - uint4 dz_vec = dZ_v[off]; - uint4 u_vec = U_v[off]; - uint4 v_vec = V_v[off]; - - auto* dz_e = reinterpret_cast(&dz_vec); - auto* u_e = reinterpret_cast(&u_vec); - auto* v_e = reinterpret_cast(&v_vec); - - uint4 du_out, dv_out; - auto* du_e = reinterpret_cast(&du_out); - auto* dv_e = reinterpret_cast(&dv_out); - - CUTE_UNROLL - for (int j = 0; j < kVecSize; ++j) { - float dz = static_cast(dz_e[j]); - float u = static_cast(u_e[j]); - float v = static_cast(v_e[j]); - float f = fast_sigmoid(u); - float df = f * (1.0f + u * (1.0f - f)); - - dv_e[j] = from_float(dz * u * f); - du_e[j] = from_float(dz * v * df); - } - - dV_v[off] = dv_out; - dU_v[off] = du_out; - } -} - -// ═══════════════════════════════════════════════════════════════════ -// Pair-multiply variant: dU = dZ * U', dV = dZ * V' -// -// Used by the fused mlp_bwd Phase 1c. Reads U' = V·silu'(U) and -// V' = silu(U) (already stored by Phase 1a) plus dZ (just stored -// by Phase 1b'); writes dU/dV in place. -// -// **n-split**: caller passes (split_idx, num_splits) matching the -// blockIdx.y / gridDim.y partition that Phase 1b' used. Each CTA -// processes ONLY its own n-tile slice — the same slice it wrote in -// Phase 1b'. This keeps Phase 1c reads intra-CTA so an intra-CTA -// NamedBarrier between 1b' and 1c is sufficient. Cross-CTA dU/dV -// visibility for downstream phases is handled by the existing -// mlp_global_barrier. -// -// The elementwise pair phase uses the historical 9-warp work partition -// (warps 0, 4-11 = 288 threads). SM100 warp 3 enters the BWD MLP path for -// UMMA phases, but skips this elementwise helper. -// ═══════════════════════════════════════════════════════════════════ - -static constexpr int kSiluBwdPairThreads = 288; // 9 warps - -template -__device__ __forceinline__ void silu_bwd_pair_tile( - const Element* __restrict__ dZ_base, // gmem [num_tokens, I] - const Element* __restrict__ Up_base, // U' coefficient - const Element* __restrict__ Vp_base, // V' coefficient - Element* __restrict__ dU_base, // output (may alias dZ_base) - Element* __restrict__ dV_base, // output (may alias Vp_base) - int m, - int TileM, - int TileN, - int intermediate_dim, - int split_idx, // = blockIdx.y - int num_splits) { // = gridDim.y - - constexpr int kVecSize = sizeof(uint4) / sizeof(Element); - - // Dense thread ID across 9 participating warps (0, 4-11). Warps 1-2 are - // BWD comm warps; warp 3 is reserved for SM100 UMMA phases and does not - // contribute to this elementwise work mapping. - int warp_id = threadIdx.x / 32; - if (warp_id >= 1 && warp_id <= 3) return; - int lane = threadIdx.x % 32; - int dense_warp = (warp_id == 0) ? 0 : warp_id - 3; - int tid = dense_warp * 32 + lane; - - int num_n_tiles = (intermediate_dim + TileN - 1) / TileN; - int row_start = m * TileM; - int vec_per_tile_n = TileN / kVecSize; // vecs per row of one n-tile - int vec_row_stride = intermediate_dim / kVecSize; // vecs per row in full I-dim - int vecs_per_tile = TileM * vec_per_tile_n; // total vecs in one (TileM, TileN) tile - - auto* dZ_v = reinterpret_cast(dZ_base); - auto* Up_v = reinterpret_cast(Up_base); - auto* Vp_v = reinterpret_cast(Vp_base); - auto* dU_v = reinterpret_cast(dU_base); - auto* dV_v = reinterpret_cast(dV_base); - - for (int n = split_idx; n < num_n_tiles; n += num_splits) { - int col_vec_start = n * vec_per_tile_n; - int valid_col_vecs = min(vec_per_tile_n, vec_row_stride - col_vec_start); - - for (int i = tid; i < vecs_per_tile; i += kSiluBwdPairThreads) { - int local_row = i / vec_per_tile_n; - int local_col_vec = i % vec_per_tile_n; - if (local_col_vec >= valid_col_vecs) continue; - - size_t off = (size_t)(row_start + local_row) * vec_row_stride - + col_vec_start + local_col_vec; - - uint4 dz = dZ_v[off]; - uint4 up = Up_v[off]; - uint4 vp = Vp_v[off]; - - auto* dz_e = reinterpret_cast(&dz); - auto* up_e = reinterpret_cast(&up); - auto* vp_e = reinterpret_cast(&vp); - - uint4 du, dv; - auto* du_e = reinterpret_cast(&du); - auto* dv_e = reinterpret_cast(&dv); - - CUTE_UNROLL - for (int j = 0; j < kVecSize; ++j) { - float dz_f = static_cast(dz_e[j]); - float up_f = static_cast(up_e[j]); - float vp_f = static_cast(vp_e[j]); - du_e[j] = from_float(dz_f * up_f); - dv_e[j] = from_float(dz_f * vp_f); - } - - dU_v[off] = du; - dV_v[off] = dv; - } - } -} - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "silu_bwd_fused_sm90.cuh" +#else +#include "silu_bwd_fused_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm100.cuh new file mode 100644 index 00000000..2d1b4c42 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm100.cuh @@ -0,0 +1,244 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// SiLU backward (fused tile): +// dV = dZ * SiLU(U) +// dU = dZ * V * grad_SiLU(U) +// ═══════════════════════════════════════════════════════════════════ +// +// Element-wise backward of Z = SiLU(U) * V. Uses all MLP warps +// (0, 1, 4–11 = 320 threads), excluding communication warps 2–3. +// Vectorized with uint4 (8 bf16 per load/store). +// +// num_elems is a runtime constant (tokens × intermediate_dim). + +#include "math_sm100.cuh" + +namespace liger { + +static constexpr int kSiluBwdThreads = 320; // 10 warps × 32 + +template +__device__ __forceinline__ void silu_bwd_fused( + const Element* __restrict__ dZ, + const Element* __restrict__ U, + const Element* __restrict__ V, + Element* __restrict__ dU, + Element* __restrict__ dV, + int num_elems) { + + constexpr int kVecSize = sizeof(uint4) / sizeof(Element); + + // Dense thread ID across 10 participating warps (0,1,4-11). + // Warps 2-3 (comm) have already returned before reaching this call. + int warp_id = threadIdx.x / 32; + int lane = threadIdx.x % 32; + int dense_warp = (warp_id < 2) ? warp_id : warp_id - 2; + int tid = dense_warp * 32 + lane; + + auto* dZ_v = reinterpret_cast(dZ); + auto* U_v = reinterpret_cast(U); + auto* V_v = reinterpret_cast(V); + auto* dU_v = reinterpret_cast(dU); + auto* dV_v = reinterpret_cast(dV); + + int vec_count = num_elems / kVecSize; + + for (int i = tid; i < vec_count; i += kSiluBwdThreads) { + uint4 dz_vec = dZ_v[i]; + uint4 u_vec = U_v[i]; + uint4 v_vec = V_v[i]; + + auto* dz_e = reinterpret_cast(&dz_vec); + auto* u_e = reinterpret_cast(&u_vec); + auto* v_e = reinterpret_cast(&v_vec); + + uint4 du_out, dv_out; + auto* du_e = reinterpret_cast(&du_out); + auto* dv_e = reinterpret_cast(&dv_out); + + CUTE_UNROLL + for (int j = 0; j < kVecSize; ++j) { + float dz = static_cast(dz_e[j]); + float u = static_cast(u_e[j]); + float v = static_cast(v_e[j]); + float f = fast_sigmoid(u); + float df = f * (1.0f + u * (1.0f - f)); + + dv_e[j] = from_float(dz * u * f); + du_e[j] = from_float(dz * v * df); + } + + dV_v[i] = dv_out; + dU_v[i] = du_out; + } +} + +// ═══════════════════════════════════════════════════════════════════ +// N-tile-aware strided variant: processes a [num_rows, num_cols] +// sub-matrix within a row-major [num_rows, row_stride] buffer. +// Handles non-contiguous columns efficiently with full thread +// utilization across all rows × cols. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void silu_bwd_fused_strided( + const Element* __restrict__ dZ_base, + const Element* __restrict__ U_base, + const Element* __restrict__ V_base, + Element* __restrict__ dU_base, + Element* __restrict__ dV_base, + int num_rows, + int num_cols, + int row_stride) { + + constexpr int kVecSize = sizeof(uint4) / sizeof(Element); + + int warp_id = threadIdx.x / 32; + int lane = threadIdx.x % 32; + int dense_warp = (warp_id < 2) ? warp_id : warp_id - 2; + int tid = dense_warp * 32 + lane; + + // Total vectors across all rows + int cols_per_vec_row = num_cols / kVecSize; + int total_vecs = num_rows * cols_per_vec_row; + + for (int i = tid; i < total_vecs; i += kSiluBwdThreads) { + int r = i / cols_per_vec_row; + int c = i % cols_per_vec_row; + int off = r * row_stride / kVecSize + c; // offset in uint4 units + + auto* dZ_v = reinterpret_cast(dZ_base); + auto* U_v = reinterpret_cast(U_base); + auto* V_v = reinterpret_cast(V_base); + auto* dU_v = reinterpret_cast(dU_base); + auto* dV_v = reinterpret_cast(dV_base); + + uint4 dz_vec = dZ_v[off]; + uint4 u_vec = U_v[off]; + uint4 v_vec = V_v[off]; + + auto* dz_e = reinterpret_cast(&dz_vec); + auto* u_e = reinterpret_cast(&u_vec); + auto* v_e = reinterpret_cast(&v_vec); + + uint4 du_out, dv_out; + auto* du_e = reinterpret_cast(&du_out); + auto* dv_e = reinterpret_cast(&dv_out); + + CUTE_UNROLL + for (int j = 0; j < kVecSize; ++j) { + float dz = static_cast(dz_e[j]); + float u = static_cast(u_e[j]); + float v = static_cast(v_e[j]); + float f = fast_sigmoid(u); + float df = f * (1.0f + u * (1.0f - f)); + + dv_e[j] = from_float(dz * u * f); + du_e[j] = from_float(dz * v * df); + } + + dV_v[off] = dv_out; + dU_v[off] = du_out; + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Pair-multiply variant: dU = dZ * U', dV = dZ * V' +// +// Used by the fused mlp_bwd Phase 1c. Reads U' = V·silu'(U) and +// V' = silu(U) (already stored by Phase 1a) plus dZ (just stored +// by Phase 1b'); writes dU/dV in place. +// +// **n-split**: caller passes (split_idx, num_splits) matching the +// blockIdx.y / gridDim.y partition that Phase 1b' used. Each CTA +// processes ONLY its own n-tile slice — the same slice it wrote in +// Phase 1b'. This keeps Phase 1c reads intra-CTA so an intra-CTA +// NamedBarrier between 1b' and 1c is sufficient. Cross-CTA dU/dV +// visibility for downstream phases is handled by the existing +// mlp_global_barrier. +// +// The elementwise pair phase uses the historical 9-warp work partition +// (warps 0, 4-11 = 288 threads). SM100 warp 3 enters the BWD MLP path for +// UMMA phases, but skips this elementwise helper. +// ═══════════════════════════════════════════════════════════════════ + +static constexpr int kSiluBwdPairThreads = 288; // 9 warps + +template +__device__ __forceinline__ void silu_bwd_pair_tile( + const Element* __restrict__ dZ_base, // gmem [num_tokens, I] + const Element* __restrict__ Up_base, // U' coefficient + const Element* __restrict__ Vp_base, // V' coefficient + Element* __restrict__ dU_base, // output (may alias dZ_base) + Element* __restrict__ dV_base, // output (may alias Vp_base) + int m, + int TileM, + int TileN, + int intermediate_dim, + int split_idx, // = blockIdx.y + int num_splits) { // = gridDim.y + + constexpr int kVecSize = sizeof(uint4) / sizeof(Element); + + // Dense thread ID across 9 participating warps (0, 4-11). Warps 1-2 are + // BWD comm warps; warp 3 is reserved for SM100 UMMA phases and does not + // contribute to this elementwise work mapping. + int warp_id = threadIdx.x / 32; + if (warp_id >= 1 && warp_id <= 3) return; + int lane = threadIdx.x % 32; + int dense_warp = (warp_id == 0) ? 0 : warp_id - 3; + int tid = dense_warp * 32 + lane; + + int num_n_tiles = (intermediate_dim + TileN - 1) / TileN; + int row_start = m * TileM; + int vec_per_tile_n = TileN / kVecSize; // vecs per row of one n-tile + int vec_row_stride = intermediate_dim / kVecSize; // vecs per row in full I-dim + int vecs_per_tile = TileM * vec_per_tile_n; // total vecs in one (TileM, TileN) tile + + auto* dZ_v = reinterpret_cast(dZ_base); + auto* Up_v = reinterpret_cast(Up_base); + auto* Vp_v = reinterpret_cast(Vp_base); + auto* dU_v = reinterpret_cast(dU_base); + auto* dV_v = reinterpret_cast(dV_base); + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + int col_vec_start = n * vec_per_tile_n; + int valid_col_vecs = min(vec_per_tile_n, vec_row_stride - col_vec_start); + + for (int i = tid; i < vecs_per_tile; i += kSiluBwdPairThreads) { + int local_row = i / vec_per_tile_n; + int local_col_vec = i % vec_per_tile_n; + if (local_col_vec >= valid_col_vecs) continue; + + size_t off = (size_t)(row_start + local_row) * vec_row_stride + + col_vec_start + local_col_vec; + + uint4 dz = dZ_v[off]; + uint4 up = Up_v[off]; + uint4 vp = Vp_v[off]; + + auto* dz_e = reinterpret_cast(&dz); + auto* up_e = reinterpret_cast(&up); + auto* vp_e = reinterpret_cast(&vp); + + uint4 du, dv; + auto* du_e = reinterpret_cast(&du); + auto* dv_e = reinterpret_cast(&dv); + + CUTE_UNROLL + for (int j = 0; j < kVecSize; ++j) { + float dz_f = static_cast(dz_e[j]); + float up_f = static_cast(up_e[j]); + float vp_f = static_cast(vp_e[j]); + du_e[j] = from_float(dz_f * up_f); + dv_e[j] = from_float(dz_f * vp_f); + } + + dU_v[off] = du; + dV_v[off] = dv; + } + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm90.cuh new file mode 100644 index 00000000..567b647b --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm90.cuh @@ -0,0 +1,244 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// SiLU backward (fused tile): +// dV = dZ * SiLU(U) +// dU = dZ * V * grad_SiLU(U) +// ═══════════════════════════════════════════════════════════════════ +// +// Element-wise backward of Z = SiLU(U) * V. Uses all MLP warps +// (0, 1, 4–11 = 320 threads), excluding communication warps 2–3. +// Vectorized with uint4 (8 bf16 per load/store). +// +// num_elems is a runtime constant (tokens × intermediate_dim). + +#include "math_sm90.cuh" + +namespace liger { + +static constexpr int kSiluBwdThreads = 320; // 10 warps × 32 + +template +__device__ __forceinline__ void silu_bwd_fused( + const Element* __restrict__ dZ, + const Element* __restrict__ U, + const Element* __restrict__ V, + Element* __restrict__ dU, + Element* __restrict__ dV, + int num_elems) { + + constexpr int kVecSize = sizeof(uint4) / sizeof(Element); + + // Dense thread ID across 10 participating warps (0,1,4-11). + // Warps 2-3 (comm) have already returned before reaching this call. + int warp_id = threadIdx.x / 32; + int lane = threadIdx.x % 32; + int dense_warp = (warp_id < 2) ? warp_id : warp_id - 2; + int tid = dense_warp * 32 + lane; + + auto* dZ_v = reinterpret_cast(dZ); + auto* U_v = reinterpret_cast(U); + auto* V_v = reinterpret_cast(V); + auto* dU_v = reinterpret_cast(dU); + auto* dV_v = reinterpret_cast(dV); + + int vec_count = num_elems / kVecSize; + + for (int i = tid; i < vec_count; i += kSiluBwdThreads) { + uint4 dz_vec = dZ_v[i]; + uint4 u_vec = U_v[i]; + uint4 v_vec = V_v[i]; + + auto* dz_e = reinterpret_cast(&dz_vec); + auto* u_e = reinterpret_cast(&u_vec); + auto* v_e = reinterpret_cast(&v_vec); + + uint4 du_out, dv_out; + auto* du_e = reinterpret_cast(&du_out); + auto* dv_e = reinterpret_cast(&dv_out); + + CUTE_UNROLL + for (int j = 0; j < kVecSize; ++j) { + float dz = static_cast(dz_e[j]); + float u = static_cast(u_e[j]); + float v = static_cast(v_e[j]); + float f = fast_sigmoid(u); + float df = f * (1.0f + u * (1.0f - f)); + + dv_e[j] = from_float(dz * u * f); + du_e[j] = from_float(dz * v * df); + } + + dV_v[i] = dv_out; + dU_v[i] = du_out; + } +} + +// ═══════════════════════════════════════════════════════════════════ +// N-tile-aware strided variant: processes a [num_rows, num_cols] +// sub-matrix within a row-major [num_rows, row_stride] buffer. +// Handles non-contiguous columns efficiently with full thread +// utilization across all rows × cols. +// ═══════════════════════════════════════════════════════════════════ + +template +__device__ __forceinline__ void silu_bwd_fused_strided( + const Element* __restrict__ dZ_base, + const Element* __restrict__ U_base, + const Element* __restrict__ V_base, + Element* __restrict__ dU_base, + Element* __restrict__ dV_base, + int num_rows, + int num_cols, + int row_stride) { + + constexpr int kVecSize = sizeof(uint4) / sizeof(Element); + + int warp_id = threadIdx.x / 32; + int lane = threadIdx.x % 32; + int dense_warp = (warp_id < 2) ? warp_id : warp_id - 2; + int tid = dense_warp * 32 + lane; + + // Total vectors across all rows + int cols_per_vec_row = num_cols / kVecSize; + int total_vecs = num_rows * cols_per_vec_row; + + for (int i = tid; i < total_vecs; i += kSiluBwdThreads) { + int r = i / cols_per_vec_row; + int c = i % cols_per_vec_row; + int off = r * row_stride / kVecSize + c; // offset in uint4 units + + auto* dZ_v = reinterpret_cast(dZ_base); + auto* U_v = reinterpret_cast(U_base); + auto* V_v = reinterpret_cast(V_base); + auto* dU_v = reinterpret_cast(dU_base); + auto* dV_v = reinterpret_cast(dV_base); + + uint4 dz_vec = dZ_v[off]; + uint4 u_vec = U_v[off]; + uint4 v_vec = V_v[off]; + + auto* dz_e = reinterpret_cast(&dz_vec); + auto* u_e = reinterpret_cast(&u_vec); + auto* v_e = reinterpret_cast(&v_vec); + + uint4 du_out, dv_out; + auto* du_e = reinterpret_cast(&du_out); + auto* dv_e = reinterpret_cast(&dv_out); + + CUTE_UNROLL + for (int j = 0; j < kVecSize; ++j) { + float dz = static_cast(dz_e[j]); + float u = static_cast(u_e[j]); + float v = static_cast(v_e[j]); + float f = fast_sigmoid(u); + float df = f * (1.0f + u * (1.0f - f)); + + dv_e[j] = from_float(dz * u * f); + du_e[j] = from_float(dz * v * df); + } + + dV_v[off] = dv_out; + dU_v[off] = du_out; + } +} + +// ═══════════════════════════════════════════════════════════════════ +// Pair-multiply variant: dU = dZ * U', dV = dZ * V' +// +// Used by the fused mlp_bwd Phase 1c. Reads U' = V·silu'(U) and +// V' = silu(U) (already stored by Phase 1a) plus dZ (just stored +// by Phase 1b'); writes dU/dV in place. +// +// **n-split**: caller passes (split_idx, num_splits) matching the +// blockIdx.y / gridDim.y partition that Phase 1b' used. Each CTA +// processes ONLY its own n-tile slice — the same slice it wrote in +// Phase 1b'. This keeps Phase 1c reads intra-CTA so an intra-CTA +// NamedBarrier between 1b' and 1c is sufficient. Cross-CTA dU/dV +// visibility for downstream phases is handled by the existing +// mlp_global_barrier. +// +// The elementwise pair phase uses the historical 9-warp work partition +// (warps 0, 4-11 = 288 threads). SM100 warp 3 enters the BWD MLP path for +// UMMA phases, but skips this elementwise helper. +// ═══════════════════════════════════════════════════════════════════ + +static constexpr int kSiluBwdPairThreads = 288; // 9 warps + +template +__device__ __forceinline__ void silu_bwd_pair_tile( + const Element* __restrict__ dZ_base, // gmem [num_tokens, I] + const Element* __restrict__ Up_base, // U' coefficient + const Element* __restrict__ Vp_base, // V' coefficient + Element* __restrict__ dU_base, // output (may alias dZ_base) + Element* __restrict__ dV_base, // output (may alias Vp_base) + int m, + int TileM, + int TileN, + int intermediate_dim, + int split_idx, // = blockIdx.y + int num_splits) { // = gridDim.y + + constexpr int kVecSize = sizeof(uint4) / sizeof(Element); + + // Dense thread ID across 9 participating warps (0, 4-11). Warps 1-2 are + // BWD comm warps; warp 3 is reserved for SM100 UMMA phases and does not + // contribute to this elementwise work mapping. + int warp_id = threadIdx.x / 32; + if (warp_id >= 1 && warp_id <= 3) return; + int lane = threadIdx.x % 32; + int dense_warp = (warp_id == 0) ? 0 : warp_id - 3; + int tid = dense_warp * 32 + lane; + + int num_n_tiles = (intermediate_dim + TileN - 1) / TileN; + int row_start = m * TileM; + int vec_per_tile_n = TileN / kVecSize; // vecs per row of one n-tile + int vec_row_stride = intermediate_dim / kVecSize; // vecs per row in full I-dim + int vecs_per_tile = TileM * vec_per_tile_n; // total vecs in one (TileM, TileN) tile + + auto* dZ_v = reinterpret_cast(dZ_base); + auto* Up_v = reinterpret_cast(Up_base); + auto* Vp_v = reinterpret_cast(Vp_base); + auto* dU_v = reinterpret_cast(dU_base); + auto* dV_v = reinterpret_cast(dV_base); + + for (int n = split_idx; n < num_n_tiles; n += num_splits) { + int col_vec_start = n * vec_per_tile_n; + int valid_col_vecs = min(vec_per_tile_n, vec_row_stride - col_vec_start); + + for (int i = tid; i < vecs_per_tile; i += kSiluBwdPairThreads) { + int local_row = i / vec_per_tile_n; + int local_col_vec = i % vec_per_tile_n; + if (local_col_vec >= valid_col_vecs) continue; + + size_t off = (size_t)(row_start + local_row) * vec_row_stride + + col_vec_start + local_col_vec; + + uint4 dz = dZ_v[off]; + uint4 up = Up_v[off]; + uint4 vp = Vp_v[off]; + + auto* dz_e = reinterpret_cast(&dz); + auto* up_e = reinterpret_cast(&up); + auto* vp_e = reinterpret_cast(&vp); + + uint4 du, dv; + auto* du_e = reinterpret_cast(&du); + auto* dv_e = reinterpret_cast(&dv); + + CUTE_UNROLL + for (int j = 0; j < kVecSize; ++j) { + float dz_f = static_cast(dz_e[j]); + float up_f = static_cast(up_e[j]); + float vp_f = static_cast(vp_e[j]); + du_e[j] = from_float(dz_f * up_f); + dv_e[j] = from_float(dz_f * vp_f); + } + + dU_v[off] = du; + dV_v[off] = dv; + } + } +} + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/tile_iterator.cuh b/liger_cute_kernels/csrc/core/src/moe/tile_iterator.cuh index d73b8617..8e584946 100644 --- a/liger_cute_kernels/csrc/core/src/moe/tile_iterator.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/tile_iterator.cuh @@ -1,527 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// MLP Tile Iterators -// ═══════════════════════════════════════════════════════════════════ -// -// Provide a uniform interface for iterating over M-tiles in the -// fused MLP kernel. Each iterator yields MlpTileInfo per tile. -// -// Both iterators support staging acquire/release hooks: -// acquire_src / release_src — around phase 1 (X read) -// acquire_dst / release_dst — around phase 2 (Y write) -// For local tiles these are no-ops. For remote tiles they -// synchronize with the comm warps via StagePipe. -// -// ═══════════════════════════════════════════════════════════════════ +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -#include -#include -#include "mlp_comms.cuh" - -namespace liger { - -// ═══════════════════════════════════════════════════════════════════ -// MlpTileInfo — what the M-loop body needs per tile -// ═══════════════════════════════════════════════════════════════════ - -struct MlpTileInfo { - const void* x_ptr; // pointer to the start of this M-tile's X data - int expert; // expert ID for weight selection - int y_m; // staging slot — X-load coord AND staging Y-store coord - // ── Direct-to-peer Y store (cur_is_local tiles) ────────────────────── - int y_store_m; // token-tile index into the destination's local_output - // (= token_offset / TileM); the TMA store row. - int peer_rank; // destination peer's intra-host local rank — selects the - // per-peer TMA descriptor built over nvshmem_ptr(y_buf, pe). - bool is_local; // same-host (NVLink) → GEMM TMA-stores Y straight into the - // peer's symmetric local_output, bypassing dst_staging+put. - int valid_m64_subtiles; // number of 64-row output subtiles containing valid rows - int m_subtile; // GemmTileM row within the parent communication tile -}; - -// ═══════════════════════════════════════════════════════════════════ -// FlatMlpWorkItem — per-(m, n-chunk) work item for the flat scheduler -// ═══════════════════════════════════════════════════════════════════ -// -// The flat scheduler hands these out per stage per phase. A CTA loops -// while (iter.has_next()) { auto item = iter.next(); ... } -// processing each item with the per-(m, [n_start, n_end)) chunk -// producer/consumer. - -struct FlatMlpWorkItem { - int m; // global M-tile index (into expert_ids / arrival counters) - int expert; // expert ID after local_expert_start subtraction - int z_m; // Z-buffer M-slot index (double-buffered across stages) - int n_start; // chunk start n-tile (inclusive) - int n_end; // chunk end n-tile (exclusive); ≤ n_start + NS - int m_prev_wait; // global m whose mlp2_arrival must reach num_n_tiles_2 - // before this chunk can write Z[z_m]. -1 if no wait - // (first ZBufferSlots stages have no prior occupant). - int phase; // 0 = MLP1, 1 = MLP2 -}; - -// ═══════════════════════════════════════════════════════════════════ -// LocalMlpTileIterator -// ═══════════════════════════════════════════════════════════════════ - -template -struct LocalMlpTileIterator { - const Element* x_base; - const int* expert_ids; - int num_m_tiles; - int hidden_dim; - int tile_offset; // global tile offset (for TMA coordinates) - - int m; - int m_end; // exclusive upper bound for the current sub-range - int col_stride; // logical grid_x — column stride between this CTA's tiles - // (flat-grid launch: NOT gridDim.x, which is num_blocks) - - __device__ void init(const Element* x_base_, const int* expert_ids_, - int num_m_tiles_, int hidden_dim_, - int col_stride_, int start_col_, - int tile_offset_ = 0) { - x_base = x_base_; - expert_ids = expert_ids_; - num_m_tiles = num_m_tiles_; - hidden_dim = hidden_dim_; - tile_offset = tile_offset_; - // Logical column / stride passed explicitly (flat-grid launch) instead - // of reading blockIdx.x / gridDim.x. The fused driver overrides m via - // set_range() before iterating; standalone callers use m = start_col. - col_stride = col_stride_; - m = start_col_; - m_end = num_m_tiles; - } - - // Restrict iteration to [m_start, m_end_). Used by moe_fused_fwd to - // split the local M-loop around the remote phase so per-tile compute - // overlaps with remote get/put traffic. - __device__ void set_range(int m_start, int m_end_) { - m = m_start; - m_end = m_end_; - } - - __device__ bool has_next() const { - return m < m_end; - } - - // y_m is the global M-tile index (for TMA coordinates into full buffer). - __device__ MlpTileInfo next() { - MlpTileInfo info; - info.x_ptr = x_base + (size_t)m * TileM * hidden_dim; - info.expert = expert_ids[m]; - info.y_m = tile_offset + m; - // Local (non-fused) iteration never uses the direct-to-peer store path. - info.y_store_m = info.y_m; - info.peer_rank = 0; - info.is_local = false; - info.valid_m64_subtiles = TileM / 64; - info.m_subtile = 0; - - m += col_stride; - return info; - } - - // No-op staging hooks for local tiles. - __device__ void acquire_src() {} - __device__ void release_src(int /*lane*/) {} - __device__ void acquire_dst() {} - __device__ void release_dst(int /*lane*/) {} -}; - -// ═══════════════════════════════════════════════════════════════════ -// LocalFlatMlpTileIterator — chunked-flat scheduler work-item iterator -// ═══════════════════════════════════════════════════════════════════ -// -// 1D iteration interface. The iterator's internal state machine walks -// stages and phases (MLP1 then MLP2 per stage); each next() returns a -// FlatMlpWorkItem with m / n_start / n_end / expert / z_m / m_prev_wait -// / phase already populated. has_next() is mutating — it advances the -// state machine past phase / stage boundaries until either it finds -// the next work item or signals end-of-iteration. -// -// Grid is 2D: dim3(num_m_per_stage, NC). Total CTAs = num_m_per_stage·NC. -// Per-CTA flat index = blockIdx.x + blockIdx.y · gridDim.x; stride -// over the per-phase work pool = gridDim.x · gridDim.y. -// -// Usage: -// iter.init(expert_ids, num_m_tiles, num_m_per_stage, -// num_n_tiles_1, NS1, num_n_tiles_2, NS2); -// while (iter.has_next()) { -// FlatMlpWorkItem item = iter.next(); -// if (item.phase == 0) { /* MLP1 */ } else { /* MLP2 */ } -// } - -template -struct LocalFlatMlpTileIterator { - const int* expert_ids; - int num_m_tiles; - int num_m_per_stage; - int num_stages; - int local_expert_start; - - // Phase parameters captured at init. - int num_n_tiles_1; - int NS1; - int num_n_chunks_1; - int num_n_tiles_2; - int NS2; - int num_n_chunks_2; - - // CTA identity (decoded once at init). - int flat_cta_idx; // blockIdx.x + blockIdx.y · gridDim.x - int total_ctas; // gridDim.x · gridDim.y - - // State machine. - int stage; // current stage index (-1 before first has_next) - int stage_m_start; - int stage_m_count; - bool in_mlp1; // true → MLP1 active, false → MLP2 active - int total_work_phase; - int work_idx; - - __device__ void init(const int* expert_ids_, - int num_m_tiles_, - int num_m_per_stage_, - int num_n_tiles_1_, - int NS1_, - int num_n_tiles_2_, - int NS2_, - int local_expert_start_ = 0) { - expert_ids = expert_ids_; - num_m_tiles = num_m_tiles_; - num_m_per_stage = num_m_per_stage_; - num_stages = (num_m_tiles_ + num_m_per_stage_ - 1) / num_m_per_stage_; - local_expert_start = local_expert_start_; - num_n_tiles_1 = num_n_tiles_1_; - NS1 = NS1_; - num_n_chunks_1 = (num_n_tiles_1_ + NS1_ - 1) / NS1_; - num_n_tiles_2 = num_n_tiles_2_; - NS2 = NS2_; - num_n_chunks_2 = (num_n_tiles_2_ + NS2_ - 1) / NS2_; - flat_cta_idx = blockIdx.x + blockIdx.y * gridDim.x; - total_ctas = gridDim.x * gridDim.y; - - // Pre-state: first has_next() will advance into (stage=0, MLP1). - stage = -1; - stage_m_start = 0; - stage_m_count = 0; - in_mlp1 = false; - total_work_phase = 0; - work_idx = flat_cta_idx; - } - - // Advance to the next (stage, phase) that has any work for this CTA. - // Returns true if there is a work item to emit at the current state, - // false when iteration is exhausted. - __device__ bool has_next() { - while (work_idx >= total_work_phase) { - if (in_mlp1) { - // Finished MLP1 in this stage — start MLP2 (same stage). - in_mlp1 = false; - total_work_phase = stage_m_count * num_n_chunks_2; - work_idx = flat_cta_idx; - } else { - // Finished MLP2 (or pre-state) — advance to next stage's MLP1. - ++stage; - if (stage >= num_stages) return false; - stage_m_start = stage * num_m_per_stage; - int rem = num_m_tiles - stage_m_start; - stage_m_count = (rem < num_m_per_stage) ? rem : num_m_per_stage; - in_mlp1 = true; - total_work_phase = stage_m_count * num_n_chunks_1; - work_idx = flat_cta_idx; - } - } - return true; - } - - __device__ FlatMlpWorkItem next() { - FlatMlpWorkItem item; - int num_n_chunks_phase = in_mlp1 ? num_n_chunks_1 : num_n_chunks_2; - int num_n_tiles_phase = in_mlp1 ? num_n_tiles_1 : num_n_tiles_2; - int NS_phase = in_mlp1 ? NS1 : NS2; - int m_in_stage = work_idx / num_n_chunks_phase; - int chunk_idx = work_idx % num_n_chunks_phase; - item.m = stage_m_start + m_in_stage; - item.expert = expert_ids[item.m] - local_expert_start; - item.z_m = (stage % ZBufferSlots) * num_m_per_stage + m_in_stage; - item.n_start = chunk_idx * NS_phase; - int n_end = item.n_start + NS_phase; - item.n_end = (n_end < num_n_tiles_phase) ? n_end : num_n_tiles_phase; - if (in_mlp1 && stage >= ZBufferSlots) { - item.m_prev_wait = (stage - ZBufferSlots) * num_m_per_stage + m_in_stage; - } else { - item.m_prev_wait = -1; - } - item.phase = in_mlp1 ? 0 : 1; - work_idx += total_ctas; - return item; - } - - // No-op staging hooks for local (standalone) iteration; RemoteFlat - // variant will wire these to NVSHMEM src/dst pipes. - __device__ void acquire_src() {} - __device__ void release_src(int /*lane*/) {} - __device__ void acquire_dst() {} - __device__ void release_dst(int /*lane*/) {} -}; - -// ═══════════════════════════════════════════════════════════════════ -// RemoteMlpTileIterator — ticket-based slot synchronization -// ═══════════════════════════════════════════════════════════════════ -// -// Iterates over remote tiles prefetched by comm warps. Staging area -// is a flat ring of L = MC · NumStages slots, where MC = (gridDim.x · -// gridDim.y) / NC is the M-tile count produced per comm pipeline stage. -// -// Producer/consumer counts per slot (unchanged from comm side): -// src: 2·NC get warps produce, runtime NS GEMM CTAs consume. -// dst: runtime NS GEMM CTAs produce, NC put warps consume. -// -// ── Ring walk ────────────────────────────────────────────────────── -// Each column blockIdx.x walks its tile sequence with stride gridDim.x -// and wraps the slot index modulo L: -// T_j = blockIdx.x + j * gridDim.x (global tile index for j-th iter) -// slot = T_j mod L (physical staging slot) -// ticket = T_j / L (pass count over the ring) -// -// Comm writes tile T to the same slot (= T mod L), so producer and -// consumer always agree on the slot. There is no per-stage-id pipe -// state: the ticket carries enough information to compute the wait -// threshold directly. -// -// acquire_src(): wait for src_ready[slot] >= (ticket+1) · 2NC -// release_src(): atomic increment src_consumed[slot] by 1 (per CTA) -// acquire_dst(): wait for dst_consumed[slot] >= ticket · NC -// (first visit per slot — ticket == 0 — has no -// outstanding put to drain, skip the spin) -// release_dst(): atomic increment dst_ready[slot] by 1 (per CTA) -// -// With runtime NS CTAs cooperating per column, each contributes one -// atomic per release, so src_consumed and dst_ready advance by -// runtime NS per consumed/produced tile — exactly what the comm-side -// StagePipe expects. -// -// ── Decoupling from comm-side constants ──────────────────────────── -// Because the ticket comes from the global tile index (not a -// per-CTA cyclic counter), NumStages, runtime NS and NC need not share -// any divisibility relationship. The only correctness requirement is -// that NC evenly divides gridDim.x · gridDim.y so MC is integer -// (host-side concern, not a template constraint). - -template -struct RemoteMlpTileIterator { - // Slot's producer/consumer counts (compile-time, used for ticket math). - // FWD steady-state GET uses kNumGetWarpsFwd active warps/CTA (the TMA path - // saturates with one), so the MLP waits for kNumGetWarpsFwd·NC src signals. - static_assert(NC % MSubTiles == 0, - "communication CTAs must divide evenly across GEMM M-subtiles"); - static constexpr int RowNC = NC / MSubTiles; - static constexpr int NumProducersSrc = kNumGetWarpsFwd * RowNC; - static constexpr int NumConsumersDst = kNumPutWarpsFwd * RowNC; - - // (expert, pe) are re-derived in SMEM, NOT read from HBM: the MLP's j-th tile - // is global flat tile T_j = m_base + j·col_stride — the same flat position the - // comm placed at slot=T_j%L, ticket=T_j/L (the StagePipe ticket handshake - // guarantees the rendezvous). So an embedded TileIterator init(m_base, - // col_stride) walked in lockstep with idx reproduces what the comm wrote, - // with zero per-tile HBM load and fewer registers than holding HBM pointers. - GemmTileIterator tit; - bool derive_on; // false → no remote_offsets (test harness) → local - int gpus_per_node; // same-host divisor for is_local - int iter_my_pe; // this PE (team space) for is_local - int cur_is_local; // cached per-tile fence-scope predicate - - int total_tiles; - int m_base; // logical column index (flat_id / runtime NS) — column's - // tile-sequence start - int col_stride; // logical grid_x (= n_gemm / runtime NS) — tile-sequence - // stride; NOT gridDim.x under the flat-grid launch - int ring_len; // physical communication slots: L = MC · NumStages - int num_splits; - - int* src_ready; - int* src_consumed; - int* dst_ready; - int* dst_consumed; - - bool is_leader; // per-warp leader (lane 0) — does the atomic poll - - int idx; // iteration counter within this column - int cur_slot; // cached slot for current acquire/release cycle - int cur_data_slot; - int cur_ticket; // cached ticket for current acquire/release cycle - int cur_m_subtile; - int cur_parent_tile; - - // remote_offsets (SMEM) drives the embedded walker; pass nullptr to disable - // derivation (e.g. the standalone handshake test, which only exercises the - // ticket/slot protocol). experts_per_pe/num_pes/my_pe/gpus_per_node feed the - // walker and the same-host is_local predicate. - __device__ void init(const int* remote_offsets, - const int* remote_counts, - int experts_per_pe, - int num_pes, - int my_pe, - int gpus_per_node_, - int total_tiles_, - int m_base_, - int n_gemm_, - int* src_ready_, int* src_consumed_, - int* dst_ready_, int* dst_consumed_, - bool is_leader_, - int runtime_nsplit_) { - total_tiles = total_tiles_; - m_base = m_base_; - num_splits = runtime_nsplit_; - col_stride = n_gemm_ / num_splits; - idx = 0; - - // L = MC · NumStages, MC = n_gemm / NC. - int mc = gridDim.x / NC; - ring_len = mc * NumStages; - - src_ready = src_ready_; - src_consumed = src_consumed_; - dst_ready = dst_ready_; - dst_consumed = dst_consumed_; - is_leader = is_leader_; - cur_slot = 0; - cur_ticket = 0; - - gpus_per_node = gpus_per_node_; - iter_my_pe = my_pe; - cur_is_local = 1; - derive_on = (remote_offsets != nullptr); - if (derive_on) { - tit.remote_offsets = remote_offsets; - tit.offsets_stride = experts_per_pe + 1; - tit.remote_counts = remote_counts; - tit.counts_stride = experts_per_pe; - tit.experts_per_pe = experts_per_pe; - tit.num_pes = num_pes; - tit.my_pe = my_pe; - tit.init(m_base_, col_stride); - } - } - - __device__ bool has_next() const { - return idx < total_tiles; - } - - // Compute and cache (slot, ticket) for the current iteration. - // Called by acquire_src() — the rest of the per-tile API (next, - // release_src, acquire_dst, release_dst) reuses the cached values. - __device__ void compute_slot() { - int parent_t; - if (derive_on) { - parent_t = tit.current_parent_tile(); - cur_m_subtile = tit.current_m_subtile(); - } else { - int virtual_t = m_base + idx * col_stride; - parent_t = virtual_t / MSubTiles; - cur_m_subtile = virtual_t % MSubTiles; - } - cur_parent_tile = parent_t; - cur_data_slot = parent_t % ring_len; - cur_slot = cur_data_slot * MSubTiles + cur_m_subtile; - cur_ticket = parent_t / ring_len; - } - - // Re-derive (expert, pe) for this tile from SMEM offsets — the embedded - // walker advances in lockstep with idx and reproduces what the comm wrote at - // this slot (proven: same global flat position). No HBM load. - __device__ MlpTileInfo next() { - MlpTileInfo info; - info.x_ptr = nullptr; // not used — TMA loads by coordinate - info.y_m = cur_data_slot; - if (derive_on) { - TileInfo ti = tit.next(); - info.expert = ti.expert; - bool loc = (gpus_per_node > 0) - ? ((ti.pe / gpus_per_node) == (iter_my_pe / gpus_per_node)) - : true; - cur_is_local = (int)loc; - info.is_local = loc; - // Destination peer's intra-host local rank → per-peer y_buf descriptor. - info.peer_rank = (gpus_per_node > 0) ? (ti.pe % gpus_per_node) : 0; - // Direct-store row = this tile's token offset in the peer's local_output. - info.y_store_m = ti.token_offset / TileM; - info.valid_m64_subtiles = (ti.valid_rows + 63) / 64; - info.m_subtile = ti.m_subtile; - } else { - // Test harness (no SMEM offsets): no derivation, no direct store. - info.expert = 0; - cur_is_local = 1; - info.is_local = false; - info.peer_rank = 0; - info.y_store_m = cur_slot; - info.valid_m64_subtiles = TileM / 64; - info.m_subtile = 0; - } - idx++; - return info; - } - - // Wait for comm's get to finish writing this tile to its slot. - // Producer (comm) signals ready += 2·NC per tile-write at the slot, - // so after ticket K the wait threshold is (K+1)·2·NC. - __device__ void acquire_src() { - compute_slot(); - int target = (cur_ticket + 1) * NumProducersSrc; - if (is_leader) { - // Relaxed spin (ld.relaxed.gpu — cheapest read, no RMW, no per-iter - // fence); the __threadfence below establishes the acquire ordering - // once for all lanes on exit. atomic_ref ctor is zero-cost/hoisted. - cuda::atomic_ref r(src_ready[cur_slot]); - while (r.load(cuda::memory_order_relaxed) < target) {} - } - __syncwarp(); - // Device-scope: observing NIC-written src. The NIC's inbound-write - // visibility is not fixable by a fence (needs GDRCopy), so system scope - // here only costs bandwidth. Device fence covers the GPU-internal view. - __threadfence(); - } - - // Signal that this CTA has consumed the tile at cur_slot. Per - // tile (column-of-runtime-NS CTAs), src_consumed advances by runtime NS. - __device__ void release_src(int lane) { - __threadfence(); - if (lane == 0) { - atomicAdd(&src_consumed[cur_slot], 1); - } - } - - // Wait for the previous put from cur_slot to finish (so we can - // reuse the slot for a fresh Y write). ticket==0 → no prior put. - __device__ void acquire_dst() { - if (cur_ticket >= 1) { - int target = cur_ticket * NumConsumersDst; - if (is_leader) { - cuda::atomic_ref r(dst_consumed[cur_slot]); - while (r.load(cuda::memory_order_relaxed) < target) {} - } - __syncwarp(); - __threadfence(); - } - } - - // Signal that Y is ready in cur_slot for the comm put. Per tile, - // dst_ready advances by runtime NS. - __device__ void release_dst(int lane) { - // Same-host tiles store Y directly into the peer's symmetric - // local_output, so the put warp has no staged data to read. Cross-host - // tiles use dst_staging + putmem and require a system-scope fence before - // the put warp observes dst_ready. - if (!cur_is_local) __threadfence_system(); - if (lane == 0) { - atomicAdd(&dst_ready[cur_slot], 1); - } - } -}; - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "tile_iterator_sm90.cuh" +#else +#include "tile_iterator_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd.cuh b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd.cuh index 79e44ca0..91b494c8 100644 --- a/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd.cuh @@ -1,359 +1,11 @@ #pragma once -// ═══════════════════════════════════════════════════════════════════ -// Backward MLP tile iterators -// ═══════════════════════════════════════════════════════════════════ -// -// Both iterators mirror their forward counterparts in tile_iterator.cuh -// and add acquire_dy/release_dy hooks for the bwd-only dY pipe. Kept in -// a separate header from the forward iterators because they depend on -// mlp_comms_bwd.cuh (CommBuffersBwd), which forward code does not need. +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -#include "tile_iterator.cuh" -#include "mlp_comms_bwd.cuh" - -namespace liger { - -// ═══════════════════════════════════════════════════════════════════ -// LocalMlpTileIteratorBwd -// ═══════════════════════════════════════════════════════════════════ -// -// Backward local tile iterator. Mirrors LocalMlpTileIterator (including -// set_range / m_end sub-range support) and adds acquire_dy/release_dy -// hooks (no-ops for local tiles). -// -// Tiles are visited in sort order (expert-contiguous). The caller can -// detect expert boundaries from tile.expert if needed. - -template -struct LocalMlpTileIteratorBwd { - const Element* x_base; - const int* expert_ids; - int num_m_tiles; - int hidden_dim; - int tile_offset; - int col_stride; // flat-grid column stride (= grid_x); gridDim.x fallback - - int m; - int m_end; // exclusive upper bound for the current sub-range - - // Flat-grid launch: start_col_ = flat_id / NSplit (this CTA's logical - // column), col_stride_ = grid_x (= n_gemm / NSplit). Defaults (-1) fall - // back to blockIdx.x / gridDim.x for legacy 2-D-grid standalone callers. - __device__ void init(const Element* x_base_, const int* expert_ids_, - int num_m_tiles_, int hidden_dim_, - int tile_offset_ = 0, - int start_col_ = -1, int col_stride_ = -1) { - x_base = x_base_; - expert_ids = expert_ids_; - num_m_tiles = num_m_tiles_; - hidden_dim = hidden_dim_; - tile_offset = tile_offset_; - col_stride = (col_stride_ >= 0) ? col_stride_ : (int)gridDim.x; - m = (start_col_ >= 0) ? start_col_ : (int)blockIdx.x; - m_end = num_m_tiles; - } - - // Restrict iteration to [m_start, m_end_). Mirrors fwd - // LocalMlpTileIterator::set_range — kept for symmetry even though - // moe_fused_bwd currently only does a single full local pass. - __device__ void set_range(int m_start, int m_end_) { - m = m_start; - m_end = m_end_; - } - - __device__ bool has_next() const { - return m < m_end; - } - - __device__ MlpTileInfo next() { - MlpTileInfo info; - info.x_ptr = x_base + (size_t)m * TileM * hidden_dim; - info.expert = expert_ids[m]; - info.y_m = tile_offset + m; - info.valid_m64_subtiles = TileM / 64; - info.m_subtile = 0; - - m += col_stride; - return info; - } - - // All staging hooks are no-ops for local tiles. - __device__ void acquire_src() {} - __device__ void release_src(int /*lane*/) {} - __device__ void acquire_dy() {} - __device__ void release_dy(int /*lane*/) {} - __device__ void acquire_dst() {} - __device__ void release_dst(int /*lane*/) {} -}; - -// ═══════════════════════════════════════════════════════════════════ -// RemoteMlpTileIteratorBwd — ticket-based slot synchronization -// ═══════════════════════════════════════════════════════════════════ -// -// Mirrors the FWD RemoteMlpTileIterator's ticket-based design. Three -// pipes (X, dY, dX) all key on the same per-iteration (slot, ticket) -// pair, computed once in acquire_src() and reused by acquire_dy / -// acquire_dst / release_dst / release_src / release_dy. The release -// hooks for src and dy are deferred past Phase 2 mlp4 and mlp3 -// respectively (callers in mlp_bwd.cuh) — both reuse the same cached -// (cur_slot, cur_ticket) so deferral does not require any extra state. -// -// Ring shape: -// L = MC · NumStages -// MC = (gridDim.x · gridDim.y) / NC -// T_j = m_base + j · gridDim.x (column's j-th global tile) -// cur_slot = T_j mod L -// cur_ticket = T_j / L -// -// Producer/consumer counts per slot: -// src (X): 2·NC get warps produce, N_SPLIT GEMM CTAs consume. -// dy: 2·NC get warps produce, N_SPLIT GEMM CTAs consume. -// dst (dX): N_SPLIT GEMM CTAs produce, NC put warps consume. -// -// Wait thresholds (atomic poll by leader, fence-syncwarp for non-leaders): -// acquire_src(): x_src_ready[slot] >= (ticket+1) · 2NC -// acquire_dy(): dy_src_ready[slot] >= (ticket+1) · 2NC -// acquire_dst(): dst_consumed[slot] >= ticket · NC -// (first visit per slot — ticket == 0 — has no -// outstanding put to drain, skip the spin) -// -// Release increments (per consuming/producing CTA): -// release_src(): atomicAdd(&x_src_consumed[slot], 1) -// release_dy(): atomicAdd(&dy_src_consumed[slot], 1) -// release_dst(): atomicAdd(&dst_ready[slot], 1) -// -// Decoupling: since the ticket comes from the global tile index, no -// divisibility relationship between NumStages, N_SPLIT and NC is -// required. The only host-side correctness requirement is that NC -// evenly divides gridDim.x · gridDim.y so MC is integer. - -template -struct RemoteMlpTileIteratorBwd { - // Slot's producer/consumer counts. The unified backward comm loop always - // uses one GET warp and one PUT warp per CTA; do_get_bwd_tma selects TMA or - // getmem per tile without changing warp ownership. - int NumProducersSrc; // get warps × NC - int NumConsumersDst; // put warps × NC - - const int* tile_expert_ids; // [L] — comm writes expert id at slot index - const int* tile_valid_rows; // [L] — comm writes exact valid row count - int total_tiles; - int m_base; // logical column (flat_id / runtime NS) — tile-seq start - int col_stride; // grid_x = n_gemm / runtime NS (column stride) - int ring_len; // L = MC · NumStages, set at init - int num_splits; - - int* x_src_ready; - int* x_src_consumed; - int* dy_src_ready; - int* dy_src_consumed; - int* dst_ready; - int* dst_consumed; - - bool is_leader; // per-warp leader (lane 0) — does the atomic poll - - int idx; // iteration counter within this column - int cur_slot; // cached slot for current acquire/release cycle - int cur_ticket; // cached ticket for current acquire/release cycle - - // Flat-grid launch: m_base_ = flat_id / NSplit (logical column), - // n_gemm_ = floor_NS (= grid_x · NSplit). col_stride = n_gemm_ / N_SPLIT, - // MC = n_gemm_ / NC. n_gemm_ < 0 falls back to gridDim.x · gridDim.y for - // legacy 2-D-grid standalone callers. - __device__ void init(const int* tile_expert_ids_, - const int* tile_valid_rows_, - int total_tiles_, - int m_base_, - int n_gemm_, - int* x_src_ready_, int* x_src_consumed_, - int* dy_src_ready_, int* dy_src_consumed_, - int* dst_ready_, int* dst_consumed_, - bool is_leader_, - int runtime_nsplit_, - bool tma_enabled_ = false) { - (void)tma_enabled_; - NumProducersSrc = kNumGetWarpsBwd * NC; - NumConsumersDst = kNumPutWarpsBwd * NC; - tile_expert_ids = tile_expert_ids_; - tile_valid_rows = tile_valid_rows_; - total_tiles = total_tiles_; - m_base = m_base_; - num_splits = runtime_nsplit_; - idx = 0; - - int n_gemm = (n_gemm_ >= 0) ? n_gemm_ - : ((int)gridDim.x * (int)gridDim.y); - col_stride = n_gemm / num_splits; - // L = MC · NumStages, MC = launched NC-complete comm columns. - int mc = (int)gridDim.x / NC; - ring_len = mc * NumStages; - - x_src_ready = x_src_ready_; - x_src_consumed = x_src_consumed_; - dy_src_ready = dy_src_ready_; - dy_src_consumed= dy_src_consumed_; - dst_ready = dst_ready_; - dst_consumed = dst_consumed_; - is_leader = is_leader_; - cur_slot = 0; - cur_ticket = 0; - } - - __device__ bool has_next() const { - return idx < total_tiles; - } - - // Compute and cache (slot, ticket) for the current iteration. - // Called by acquire_src() — the rest of the per-tile API (next, - // acquire_dy, release_*, acquire_dst) reuses the cached values. - __device__ void compute_slot() { - int T = m_base + idx * col_stride; - cur_slot = T % ring_len; - cur_ticket = T / ring_len; - } - - // Called AFTER acquire_src(), so do_get_bwd's __threadfence has - // made tile_expert_ids[slot] visible. - __device__ MlpTileInfo next() { - MlpTileInfo info; - info.x_ptr = nullptr; // not used — TMA loads by coordinate - info.expert = tile_expert_ids[cur_slot]; - info.y_m = cur_slot; - info.valid_m64_subtiles = - (tile_valid_rows[cur_slot] + 63) / 64; - info.m_subtile = 0; - idx++; - return info; - } - - // Wait for comm warps to finish X get into staging. - // Producer (comm) signals ready += 2·NC per tile-write at the slot, - // so after ticket K the wait threshold is (K+1)·2NC. - __device__ void acquire_src() { - compute_slot(); - int target = (cur_ticket + 1) * NumProducersSrc; - if (is_leader) { - while (atomicAdd(&x_src_ready[cur_slot], 0) < target) {} - } - __syncwarp(); - __threadfence(); // observe NIC-written X (inbound visibility needs GDRCopy, not a fence) - } - - // Free X slot for comm warps to reuse (deferred past Phase 2 mlp4). - // Per tile (column-of-N_SPLIT CTAs), x_src_consumed advances by N_SPLIT. - __device__ void release_src(int lane) { - __threadfence(); - if (lane == 0) { - atomicAdd(&x_src_consumed[cur_slot], 1); - } - } - - // Sub-batch-grouped variant: release an EXPLICIT slot rather than the - // cached cur_slot. Used when Phase 1 acquires S sub-batches before the - // grouped Phase 2 — by release time cur_slot holds the LAST sub-batch's - // slot, so the earlier sub-batches must be freed by their saved slot. - // release_src/release_dy only ever index x_src_consumed[slot] / - // dy_src_consumed[slot] (the ticket is unused), so the saved slot alone - // is sufficient. Caller emits ONE __threadfence (shared across the batch - // of releases) before invoking these per saved slot. - __device__ void release_src_slot(int slot, int lane) { - if (lane == 0) { - atomicAdd(&x_src_consumed[slot], 1); - } - } - - // Wait for comm warps to finish dY get into staging. - __device__ void acquire_dy() { - int target = (cur_ticket + 1) * NumProducersSrc; - if (is_leader) { - while (atomicAdd(&dy_src_ready[cur_slot], 0) < target) {} - } - __syncwarp(); - __threadfence(); // observe NIC-written dY (inbound visibility needs GDRCopy, not a fence) - } - - // Free dY slot for comm warps to reuse (deferred past Phase 2 mlp3). - __device__ void release_dy(int lane) { - __threadfence(); - if (lane == 0) { - atomicAdd(&dy_src_consumed[cur_slot], 1); - } - } - - // Sub-batch-grouped variant (see release_src_slot). - __device__ void release_dy_slot(int slot, int lane) { - if (lane == 0) { - atomicAdd(&dy_src_consumed[slot], 1); - } - } - - // Wait for previous dX put to finish (staging slot free). - // ticket==0 → no prior put. - __device__ void acquire_dst() { - if (cur_ticket >= 1) { - int target = cur_ticket * NumConsumersDst; - if (is_leader) { - while (atomicAdd(&dst_consumed[cur_slot], 0) < target) {} - } - __syncwarp(); - __threadfence(); - } - } - - // Signal comm warps that dX is ready to be put. Per tile, dst_ready - // advances by N_SPLIT. - __device__ void release_dst(int lane) { - __threadfence_system(); - if (lane == 0) { - atomicAdd(&dst_ready[cur_slot], 1); - } - } -}; - -// ═══════════════════════════════════════════════════════════════════ -// FusedMlpTileIteratorBwd -// ═══════════════════════════════════════════════════════════════════ -// -// Thin wrapper over a single RemoteMlpTileIteratorBwd. Every BWD tile — -// local or remote — flows through the comm get→staging→MLP→put path (the -// separate materialized local pass was removed), so there is only the remote -// staging-ring sub-iterator and all accessors forward straight to it. The -// wrapper exists only to keep a stable iterator type at the moe_fused_bwd / -// mlp_fused_bwd_dual call sites. - -template -struct FusedMlpTileIteratorBwd { - RemoteMlpTileIteratorBwd remote; - - __device__ void init_remote(const int* tile_expert_ids, - const int* tile_valid_rows, - int total_tiles, - int m_base, - int n_gemm, - int* x_src_ready, int* x_src_consumed, - int* dy_src_ready, int* dy_src_consumed, - int* dst_ready, int* dst_consumed, - bool is_leader, - int runtime_nsplit, - bool tma_enabled = false) { - remote.init(tile_expert_ids, tile_valid_rows, - total_tiles, m_base, n_gemm, - x_src_ready, x_src_consumed, dy_src_ready, dy_src_consumed, - dst_ready, dst_consumed, is_leader, runtime_nsplit, tma_enabled); - } - - // All accessors forward to the remote sub-iterator (local pass removed). - __device__ bool has_next() const { return remote.has_next(); } - __device__ MlpTileInfo next() { return remote.next(); } - __device__ void acquire_src() { remote.acquire_src(); } - __device__ void release_src(int lane) { remote.release_src(lane); } - __device__ void release_src_slot(int slot, int lane) { remote.release_src_slot(slot, lane); } - __device__ void acquire_dy() { remote.acquire_dy(); } - __device__ void release_dy(int lane) { remote.release_dy(lane); } - __device__ void release_dy_slot(int slot, int lane) { remote.release_dy_slot(slot, lane); } - __device__ void acquire_dst() { remote.acquire_dst(); } - __device__ void release_dst(int lane) { remote.release_dst(lane); } -}; - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "tile_iterator_bwd_sm90.cuh" +#else +#include "tile_iterator_bwd_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm100.cuh new file mode 100644 index 00000000..ba021c16 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm100.cuh @@ -0,0 +1,359 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Backward MLP tile iterators +// ═══════════════════════════════════════════════════════════════════ +// +// Both iterators mirror their forward counterparts in tile_iterator.cuh +// and add acquire_dy/release_dy hooks for the bwd-only dY pipe. Kept in +// a separate header from the forward iterators because they depend on +// mlp_comms_bwd.cuh (CommBuffersBwd), which forward code does not need. + +#include "tile_iterator_sm100.cuh" +#include "mlp_comms_bwd_sm100.cuh" + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// LocalMlpTileIteratorBwd +// ═══════════════════════════════════════════════════════════════════ +// +// Backward local tile iterator. Mirrors LocalMlpTileIterator (including +// set_range / m_end sub-range support) and adds acquire_dy/release_dy +// hooks (no-ops for local tiles). +// +// Tiles are visited in sort order (expert-contiguous). The caller can +// detect expert boundaries from tile.expert if needed. + +template +struct LocalMlpTileIteratorBwd { + const Element* x_base; + const int* expert_ids; + int num_m_tiles; + int hidden_dim; + int tile_offset; + int col_stride; // flat-grid column stride (= grid_x); gridDim.x fallback + + int m; + int m_end; // exclusive upper bound for the current sub-range + + // Flat-grid launch: start_col_ = flat_id / NSplit (this CTA's logical + // column), col_stride_ = grid_x (= n_gemm / NSplit). Defaults (-1) fall + // back to blockIdx.x / gridDim.x for legacy 2-D-grid standalone callers. + __device__ void init(const Element* x_base_, const int* expert_ids_, + int num_m_tiles_, int hidden_dim_, + int tile_offset_ = 0, + int start_col_ = -1, int col_stride_ = -1) { + x_base = x_base_; + expert_ids = expert_ids_; + num_m_tiles = num_m_tiles_; + hidden_dim = hidden_dim_; + tile_offset = tile_offset_; + col_stride = (col_stride_ >= 0) ? col_stride_ : (int)gridDim.x; + m = (start_col_ >= 0) ? start_col_ : (int)blockIdx.x; + m_end = num_m_tiles; + } + + // Restrict iteration to [m_start, m_end_). Mirrors fwd + // LocalMlpTileIterator::set_range — kept for symmetry even though + // moe_fused_bwd currently only does a single full local pass. + __device__ void set_range(int m_start, int m_end_) { + m = m_start; + m_end = m_end_; + } + + __device__ bool has_next() const { + return m < m_end; + } + + __device__ MlpTileInfo next() { + MlpTileInfo info; + info.x_ptr = x_base + (size_t)m * TileM * hidden_dim; + info.expert = expert_ids[m]; + info.y_m = tile_offset + m; + info.valid_m64_subtiles = TileM / 64; + info.m_subtile = 0; + + m += col_stride; + return info; + } + + // All staging hooks are no-ops for local tiles. + __device__ void acquire_src() {} + __device__ void release_src(int /*lane*/) {} + __device__ void acquire_dy() {} + __device__ void release_dy(int /*lane*/) {} + __device__ void acquire_dst() {} + __device__ void release_dst(int /*lane*/) {} +}; + +// ═══════════════════════════════════════════════════════════════════ +// RemoteMlpTileIteratorBwd — ticket-based slot synchronization +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors the FWD RemoteMlpTileIterator's ticket-based design. Three +// pipes (X, dY, dX) all key on the same per-iteration (slot, ticket) +// pair, computed once in acquire_src() and reused by acquire_dy / +// acquire_dst / release_dst / release_src / release_dy. The release +// hooks for src and dy are deferred past Phase 2 mlp4 and mlp3 +// respectively (callers in mlp_bwd.cuh) — both reuse the same cached +// (cur_slot, cur_ticket) so deferral does not require any extra state. +// +// Ring shape: +// L = MC · NumStages +// MC = (gridDim.x · gridDim.y) / NC +// T_j = m_base + j · gridDim.x (column's j-th global tile) +// cur_slot = T_j mod L +// cur_ticket = T_j / L +// +// Producer/consumer counts per slot: +// src (X): 2·NC get warps produce, N_SPLIT GEMM CTAs consume. +// dy: 2·NC get warps produce, N_SPLIT GEMM CTAs consume. +// dst (dX): N_SPLIT GEMM CTAs produce, NC put warps consume. +// +// Wait thresholds (atomic poll by leader, fence-syncwarp for non-leaders): +// acquire_src(): x_src_ready[slot] >= (ticket+1) · 2NC +// acquire_dy(): dy_src_ready[slot] >= (ticket+1) · 2NC +// acquire_dst(): dst_consumed[slot] >= ticket · NC +// (first visit per slot — ticket == 0 — has no +// outstanding put to drain, skip the spin) +// +// Release increments (per consuming/producing CTA): +// release_src(): atomicAdd(&x_src_consumed[slot], 1) +// release_dy(): atomicAdd(&dy_src_consumed[slot], 1) +// release_dst(): atomicAdd(&dst_ready[slot], 1) +// +// Decoupling: since the ticket comes from the global tile index, no +// divisibility relationship between NumStages, N_SPLIT and NC is +// required. The only host-side correctness requirement is that NC +// evenly divides gridDim.x · gridDim.y so MC is integer. + +template +struct RemoteMlpTileIteratorBwd { + // Slot's producer/consumer counts. The unified backward comm loop always + // uses one GET warp and one PUT warp per CTA; do_get_bwd_tma selects TMA or + // getmem per tile without changing warp ownership. + int NumProducersSrc; // get warps × NC + int NumConsumersDst; // put warps × NC + + const int* tile_expert_ids; // [L] — comm writes expert id at slot index + const int* tile_valid_rows; // [L] — comm writes exact valid row count + int total_tiles; + int m_base; // logical column (flat_id / runtime NS) — tile-seq start + int col_stride; // grid_x = n_gemm / runtime NS (column stride) + int ring_len; // L = MC · NumStages, set at init + int num_splits; + + int* x_src_ready; + int* x_src_consumed; + int* dy_src_ready; + int* dy_src_consumed; + int* dst_ready; + int* dst_consumed; + + bool is_leader; // per-warp leader (lane 0) — does the atomic poll + + int idx; // iteration counter within this column + int cur_slot; // cached slot for current acquire/release cycle + int cur_ticket; // cached ticket for current acquire/release cycle + + // Flat-grid launch: m_base_ = flat_id / NSplit (logical column), + // n_gemm_ = floor_NS (= grid_x · NSplit). col_stride = n_gemm_ / N_SPLIT, + // MC = n_gemm_ / NC. n_gemm_ < 0 falls back to gridDim.x · gridDim.y for + // legacy 2-D-grid standalone callers. + __device__ void init(const int* tile_expert_ids_, + const int* tile_valid_rows_, + int total_tiles_, + int m_base_, + int n_gemm_, + int* x_src_ready_, int* x_src_consumed_, + int* dy_src_ready_, int* dy_src_consumed_, + int* dst_ready_, int* dst_consumed_, + bool is_leader_, + int runtime_nsplit_, + bool tma_enabled_ = false) { + (void)tma_enabled_; + NumProducersSrc = kNumGetWarpsBwd * NC; + NumConsumersDst = kNumPutWarpsBwd * NC; + tile_expert_ids = tile_expert_ids_; + tile_valid_rows = tile_valid_rows_; + total_tiles = total_tiles_; + m_base = m_base_; + num_splits = runtime_nsplit_; + idx = 0; + + int n_gemm = (n_gemm_ >= 0) ? n_gemm_ + : ((int)gridDim.x * (int)gridDim.y); + col_stride = n_gemm / num_splits; + // L = MC · NumStages, MC = launched NC-complete comm columns. + int mc = (int)gridDim.x / NC; + ring_len = mc * NumStages; + + x_src_ready = x_src_ready_; + x_src_consumed = x_src_consumed_; + dy_src_ready = dy_src_ready_; + dy_src_consumed= dy_src_consumed_; + dst_ready = dst_ready_; + dst_consumed = dst_consumed_; + is_leader = is_leader_; + cur_slot = 0; + cur_ticket = 0; + } + + __device__ bool has_next() const { + return idx < total_tiles; + } + + // Compute and cache (slot, ticket) for the current iteration. + // Called by acquire_src() — the rest of the per-tile API (next, + // acquire_dy, release_*, acquire_dst) reuses the cached values. + __device__ void compute_slot() { + int T = m_base + idx * col_stride; + cur_slot = T % ring_len; + cur_ticket = T / ring_len; + } + + // Called AFTER acquire_src(), so do_get_bwd's __threadfence has + // made tile_expert_ids[slot] visible. + __device__ MlpTileInfo next() { + MlpTileInfo info; + info.x_ptr = nullptr; // not used — TMA loads by coordinate + info.expert = tile_expert_ids[cur_slot]; + info.y_m = cur_slot; + info.valid_m64_subtiles = + (tile_valid_rows[cur_slot] + 63) / 64; + info.m_subtile = 0; + idx++; + return info; + } + + // Wait for comm warps to finish X get into staging. + // Producer (comm) signals ready += 2·NC per tile-write at the slot, + // so after ticket K the wait threshold is (K+1)·2NC. + __device__ void acquire_src() { + compute_slot(); + int target = (cur_ticket + 1) * NumProducersSrc; + if (is_leader) { + while (atomicAdd(&x_src_ready[cur_slot], 0) < target) {} + } + __syncwarp(); + __threadfence(); // observe NIC-written X (inbound visibility needs GDRCopy, not a fence) + } + + // Free X slot for comm warps to reuse (deferred past Phase 2 mlp4). + // Per tile (column-of-N_SPLIT CTAs), x_src_consumed advances by N_SPLIT. + __device__ void release_src(int lane) { + __threadfence(); + if (lane == 0) { + atomicAdd(&x_src_consumed[cur_slot], 1); + } + } + + // Sub-batch-grouped variant: release an EXPLICIT slot rather than the + // cached cur_slot. Used when Phase 1 acquires S sub-batches before the + // grouped Phase 2 — by release time cur_slot holds the LAST sub-batch's + // slot, so the earlier sub-batches must be freed by their saved slot. + // release_src/release_dy only ever index x_src_consumed[slot] / + // dy_src_consumed[slot] (the ticket is unused), so the saved slot alone + // is sufficient. Caller emits ONE __threadfence (shared across the batch + // of releases) before invoking these per saved slot. + __device__ void release_src_slot(int slot, int lane) { + if (lane == 0) { + atomicAdd(&x_src_consumed[slot], 1); + } + } + + // Wait for comm warps to finish dY get into staging. + __device__ void acquire_dy() { + int target = (cur_ticket + 1) * NumProducersSrc; + if (is_leader) { + while (atomicAdd(&dy_src_ready[cur_slot], 0) < target) {} + } + __syncwarp(); + __threadfence(); // observe NIC-written dY (inbound visibility needs GDRCopy, not a fence) + } + + // Free dY slot for comm warps to reuse (deferred past Phase 2 mlp3). + __device__ void release_dy(int lane) { + __threadfence(); + if (lane == 0) { + atomicAdd(&dy_src_consumed[cur_slot], 1); + } + } + + // Sub-batch-grouped variant (see release_src_slot). + __device__ void release_dy_slot(int slot, int lane) { + if (lane == 0) { + atomicAdd(&dy_src_consumed[slot], 1); + } + } + + // Wait for previous dX put to finish (staging slot free). + // ticket==0 → no prior put. + __device__ void acquire_dst() { + if (cur_ticket >= 1) { + int target = cur_ticket * NumConsumersDst; + if (is_leader) { + while (atomicAdd(&dst_consumed[cur_slot], 0) < target) {} + } + __syncwarp(); + __threadfence(); + } + } + + // Signal comm warps that dX is ready to be put. Per tile, dst_ready + // advances by N_SPLIT. + __device__ void release_dst(int lane) { + __threadfence_system(); + if (lane == 0) { + atomicAdd(&dst_ready[cur_slot], 1); + } + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// FusedMlpTileIteratorBwd +// ═══════════════════════════════════════════════════════════════════ +// +// Thin wrapper over a single RemoteMlpTileIteratorBwd. Every BWD tile — +// local or remote — flows through the comm get→staging→MLP→put path (the +// separate materialized local pass was removed), so there is only the remote +// staging-ring sub-iterator and all accessors forward straight to it. The +// wrapper exists only to keep a stable iterator type at the moe_fused_bwd / +// mlp_fused_bwd_dual call sites. + +template +struct FusedMlpTileIteratorBwd { + RemoteMlpTileIteratorBwd remote; + + __device__ void init_remote(const int* tile_expert_ids, + const int* tile_valid_rows, + int total_tiles, + int m_base, + int n_gemm, + int* x_src_ready, int* x_src_consumed, + int* dy_src_ready, int* dy_src_consumed, + int* dst_ready, int* dst_consumed, + bool is_leader, + int runtime_nsplit, + bool tma_enabled = false) { + remote.init(tile_expert_ids, tile_valid_rows, + total_tiles, m_base, n_gemm, + x_src_ready, x_src_consumed, dy_src_ready, dy_src_consumed, + dst_ready, dst_consumed, is_leader, runtime_nsplit, tma_enabled); + } + + // All accessors forward to the remote sub-iterator (local pass removed). + __device__ bool has_next() const { return remote.has_next(); } + __device__ MlpTileInfo next() { return remote.next(); } + __device__ void acquire_src() { remote.acquire_src(); } + __device__ void release_src(int lane) { remote.release_src(lane); } + __device__ void release_src_slot(int slot, int lane) { remote.release_src_slot(slot, lane); } + __device__ void acquire_dy() { remote.acquire_dy(); } + __device__ void release_dy(int lane) { remote.release_dy(lane); } + __device__ void release_dy_slot(int slot, int lane) { remote.release_dy_slot(slot, lane); } + __device__ void acquire_dst() { remote.acquire_dst(); } + __device__ void release_dst(int lane) { remote.release_dst(lane); } +}; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm90.cuh new file mode 100644 index 00000000..9be9aa3a --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm90.cuh @@ -0,0 +1,359 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// Backward MLP tile iterators +// ═══════════════════════════════════════════════════════════════════ +// +// Both iterators mirror their forward counterparts in tile_iterator.cuh +// and add acquire_dy/release_dy hooks for the bwd-only dY pipe. Kept in +// a separate header from the forward iterators because they depend on +// mlp_comms_bwd.cuh (CommBuffersBwd), which forward code does not need. + +#include "tile_iterator_sm90.cuh" +#include "mlp_comms_bwd_sm90.cuh" + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// LocalMlpTileIteratorBwd +// ═══════════════════════════════════════════════════════════════════ +// +// Backward local tile iterator. Mirrors LocalMlpTileIterator (including +// set_range / m_end sub-range support) and adds acquire_dy/release_dy +// hooks (no-ops for local tiles). +// +// Tiles are visited in sort order (expert-contiguous). The caller can +// detect expert boundaries from tile.expert if needed. + +template +struct LocalMlpTileIteratorBwd { + const Element* x_base; + const int* expert_ids; + int num_m_tiles; + int hidden_dim; + int tile_offset; + int col_stride; // flat-grid column stride (= grid_x); gridDim.x fallback + + int m; + int m_end; // exclusive upper bound for the current sub-range + + // Flat-grid launch: start_col_ = flat_id / NSplit (this CTA's logical + // column), col_stride_ = grid_x (= n_gemm / NSplit). Defaults (-1) fall + // back to blockIdx.x / gridDim.x for legacy 2-D-grid standalone callers. + __device__ void init(const Element* x_base_, const int* expert_ids_, + int num_m_tiles_, int hidden_dim_, + int tile_offset_ = 0, + int start_col_ = -1, int col_stride_ = -1) { + x_base = x_base_; + expert_ids = expert_ids_; + num_m_tiles = num_m_tiles_; + hidden_dim = hidden_dim_; + tile_offset = tile_offset_; + col_stride = (col_stride_ >= 0) ? col_stride_ : (int)gridDim.x; + m = (start_col_ >= 0) ? start_col_ : (int)blockIdx.x; + m_end = num_m_tiles; + } + + // Restrict iteration to [m_start, m_end_). Mirrors fwd + // LocalMlpTileIterator::set_range — kept for symmetry even though + // moe_fused_bwd currently only does a single full local pass. + __device__ void set_range(int m_start, int m_end_) { + m = m_start; + m_end = m_end_; + } + + __device__ bool has_next() const { + return m < m_end; + } + + __device__ MlpTileInfo next() { + MlpTileInfo info; + info.x_ptr = x_base + (size_t)m * TileM * hidden_dim; + info.expert = expert_ids[m]; + info.y_m = tile_offset + m; + info.valid_m64_subtiles = TileM / 64; + info.m_subtile = 0; + + m += col_stride; + return info; + } + + // All staging hooks are no-ops for local tiles. + __device__ void acquire_src() {} + __device__ void release_src(int /*lane*/) {} + __device__ void acquire_dy() {} + __device__ void release_dy(int /*lane*/) {} + __device__ void acquire_dst() {} + __device__ void release_dst(int /*lane*/) {} +}; + +// ═══════════════════════════════════════════════════════════════════ +// RemoteMlpTileIteratorBwd — ticket-based slot synchronization +// ═══════════════════════════════════════════════════════════════════ +// +// Mirrors the FWD RemoteMlpTileIterator's ticket-based design. Three +// pipes (X, dY, dX) all key on the same per-iteration (slot, ticket) +// pair, computed once in acquire_src() and reused by acquire_dy / +// acquire_dst / release_dst / release_src / release_dy. The release +// hooks for src and dy are deferred past Phase 2 mlp4 and mlp3 +// respectively (callers in mlp_bwd.cuh) — both reuse the same cached +// (cur_slot, cur_ticket) so deferral does not require any extra state. +// +// Ring shape: +// L = MC · NumStages +// MC = (gridDim.x · gridDim.y) / NC +// T_j = m_base + j · gridDim.x (column's j-th global tile) +// cur_slot = T_j mod L +// cur_ticket = T_j / L +// +// Producer/consumer counts per slot: +// src (X): 2·NC get warps produce, N_SPLIT GEMM CTAs consume. +// dy: 2·NC get warps produce, N_SPLIT GEMM CTAs consume. +// dst (dX): N_SPLIT GEMM CTAs produce, NC put warps consume. +// +// Wait thresholds (atomic poll by leader, fence-syncwarp for non-leaders): +// acquire_src(): x_src_ready[slot] >= (ticket+1) · 2NC +// acquire_dy(): dy_src_ready[slot] >= (ticket+1) · 2NC +// acquire_dst(): dst_consumed[slot] >= ticket · NC +// (first visit per slot — ticket == 0 — has no +// outstanding put to drain, skip the spin) +// +// Release increments (per consuming/producing CTA): +// release_src(): atomicAdd(&x_src_consumed[slot], 1) +// release_dy(): atomicAdd(&dy_src_consumed[slot], 1) +// release_dst(): atomicAdd(&dst_ready[slot], 1) +// +// Decoupling: since the ticket comes from the global tile index, no +// divisibility relationship between NumStages, N_SPLIT and NC is +// required. The only host-side correctness requirement is that NC +// evenly divides gridDim.x · gridDim.y so MC is integer. + +template +struct RemoteMlpTileIteratorBwd { + // Slot's producer/consumer counts. The unified backward comm loop always + // uses one GET warp and one PUT warp per CTA; do_get_bwd_tma selects TMA or + // getmem per tile without changing warp ownership. + int NumProducersSrc; // get warps × NC + int NumConsumersDst; // put warps × NC + + const int* tile_expert_ids; // [L] — comm writes expert id at slot index + const int* tile_valid_rows; // [L] — comm writes exact valid row count + int total_tiles; + int m_base; // logical column (flat_id / runtime NS) — tile-seq start + int col_stride; // grid_x = n_gemm / runtime NS (column stride) + int ring_len; // L = MC · NumStages, set at init + int num_splits; + + int* x_src_ready; + int* x_src_consumed; + int* dy_src_ready; + int* dy_src_consumed; + int* dst_ready; + int* dst_consumed; + + bool is_leader; // per-warp leader (lane 0) — does the atomic poll + + int idx; // iteration counter within this column + int cur_slot; // cached slot for current acquire/release cycle + int cur_ticket; // cached ticket for current acquire/release cycle + + // Flat-grid launch: m_base_ = flat_id / NSplit (logical column), + // n_gemm_ = floor_NS (= grid_x · NSplit). col_stride = n_gemm_ / N_SPLIT, + // MC = n_gemm_ / NC. n_gemm_ < 0 falls back to gridDim.x · gridDim.y for + // legacy 2-D-grid standalone callers. + __device__ void init(const int* tile_expert_ids_, + const int* tile_valid_rows_, + int total_tiles_, + int m_base_, + int n_gemm_, + int* x_src_ready_, int* x_src_consumed_, + int* dy_src_ready_, int* dy_src_consumed_, + int* dst_ready_, int* dst_consumed_, + bool is_leader_, + int runtime_nsplit_, + bool tma_enabled_ = false) { + (void)tma_enabled_; + NumProducersSrc = kNumGetWarpsBwd * NC; + NumConsumersDst = kNumPutWarpsBwd * NC; + tile_expert_ids = tile_expert_ids_; + tile_valid_rows = tile_valid_rows_; + total_tiles = total_tiles_; + m_base = m_base_; + num_splits = runtime_nsplit_; + idx = 0; + + int n_gemm = (n_gemm_ >= 0) ? n_gemm_ + : ((int)gridDim.x * (int)gridDim.y); + col_stride = n_gemm / num_splits; + // L = MC · NumStages, MC = launched NC-complete comm columns. + int mc = (int)gridDim.x / NC; + ring_len = mc * NumStages; + + x_src_ready = x_src_ready_; + x_src_consumed = x_src_consumed_; + dy_src_ready = dy_src_ready_; + dy_src_consumed= dy_src_consumed_; + dst_ready = dst_ready_; + dst_consumed = dst_consumed_; + is_leader = is_leader_; + cur_slot = 0; + cur_ticket = 0; + } + + __device__ bool has_next() const { + return idx < total_tiles; + } + + // Compute and cache (slot, ticket) for the current iteration. + // Called by acquire_src() — the rest of the per-tile API (next, + // acquire_dy, release_*, acquire_dst) reuses the cached values. + __device__ void compute_slot() { + int T = m_base + idx * col_stride; + cur_slot = T % ring_len; + cur_ticket = T / ring_len; + } + + // Called AFTER acquire_src(), so do_get_bwd's __threadfence has + // made tile_expert_ids[slot] visible. + __device__ MlpTileInfo next() { + MlpTileInfo info; + info.x_ptr = nullptr; // not used — TMA loads by coordinate + info.expert = tile_expert_ids[cur_slot]; + info.y_m = cur_slot; + info.valid_m64_subtiles = + (tile_valid_rows[cur_slot] + 63) / 64; + info.m_subtile = 0; + idx++; + return info; + } + + // Wait for comm warps to finish X get into staging. + // Producer (comm) signals ready += 2·NC per tile-write at the slot, + // so after ticket K the wait threshold is (K+1)·2NC. + __device__ void acquire_src() { + compute_slot(); + int target = (cur_ticket + 1) * NumProducersSrc; + if (is_leader) { + while (atomicAdd(&x_src_ready[cur_slot], 0) < target) {} + } + __syncwarp(); + __threadfence(); // observe NIC-written X (inbound visibility needs GDRCopy, not a fence) + } + + // Free X slot for comm warps to reuse (deferred past Phase 2 mlp4). + // Per tile (column-of-N_SPLIT CTAs), x_src_consumed advances by N_SPLIT. + __device__ void release_src(int lane) { + __threadfence(); + if (lane == 0) { + atomicAdd(&x_src_consumed[cur_slot], 1); + } + } + + // Sub-batch-grouped variant: release an EXPLICIT slot rather than the + // cached cur_slot. Used when Phase 1 acquires S sub-batches before the + // grouped Phase 2 — by release time cur_slot holds the LAST sub-batch's + // slot, so the earlier sub-batches must be freed by their saved slot. + // release_src/release_dy only ever index x_src_consumed[slot] / + // dy_src_consumed[slot] (the ticket is unused), so the saved slot alone + // is sufficient. Caller emits ONE __threadfence (shared across the batch + // of releases) before invoking these per saved slot. + __device__ void release_src_slot(int slot, int lane) { + if (lane == 0) { + atomicAdd(&x_src_consumed[slot], 1); + } + } + + // Wait for comm warps to finish dY get into staging. + __device__ void acquire_dy() { + int target = (cur_ticket + 1) * NumProducersSrc; + if (is_leader) { + while (atomicAdd(&dy_src_ready[cur_slot], 0) < target) {} + } + __syncwarp(); + __threadfence(); // observe NIC-written dY (inbound visibility needs GDRCopy, not a fence) + } + + // Free dY slot for comm warps to reuse (deferred past Phase 2 mlp3). + __device__ void release_dy(int lane) { + __threadfence(); + if (lane == 0) { + atomicAdd(&dy_src_consumed[cur_slot], 1); + } + } + + // Sub-batch-grouped variant (see release_src_slot). + __device__ void release_dy_slot(int slot, int lane) { + if (lane == 0) { + atomicAdd(&dy_src_consumed[slot], 1); + } + } + + // Wait for previous dX put to finish (staging slot free). + // ticket==0 → no prior put. + __device__ void acquire_dst() { + if (cur_ticket >= 1) { + int target = cur_ticket * NumConsumersDst; + if (is_leader) { + while (atomicAdd(&dst_consumed[cur_slot], 0) < target) {} + } + __syncwarp(); + __threadfence(); + } + } + + // Signal comm warps that dX is ready to be put. Per tile, dst_ready + // advances by N_SPLIT. + __device__ void release_dst(int lane) { + __threadfence_system(); + if (lane == 0) { + atomicAdd(&dst_ready[cur_slot], 1); + } + } +}; + +// ═══════════════════════════════════════════════════════════════════ +// FusedMlpTileIteratorBwd +// ═══════════════════════════════════════════════════════════════════ +// +// Thin wrapper over a single RemoteMlpTileIteratorBwd. Every BWD tile — +// local or remote — flows through the comm get→staging→MLP→put path (the +// separate materialized local pass was removed), so there is only the remote +// staging-ring sub-iterator and all accessors forward straight to it. The +// wrapper exists only to keep a stable iterator type at the moe_fused_bwd / +// mlp_fused_bwd_dual call sites. + +template +struct FusedMlpTileIteratorBwd { + RemoteMlpTileIteratorBwd remote; + + __device__ void init_remote(const int* tile_expert_ids, + const int* tile_valid_rows, + int total_tiles, + int m_base, + int n_gemm, + int* x_src_ready, int* x_src_consumed, + int* dy_src_ready, int* dy_src_consumed, + int* dst_ready, int* dst_consumed, + bool is_leader, + int runtime_nsplit, + bool tma_enabled = false) { + remote.init(tile_expert_ids, tile_valid_rows, + total_tiles, m_base, n_gemm, + x_src_ready, x_src_consumed, dy_src_ready, dy_src_consumed, + dst_ready, dst_consumed, is_leader, runtime_nsplit, tma_enabled); + } + + // All accessors forward to the remote sub-iterator (local pass removed). + __device__ bool has_next() const { return remote.has_next(); } + __device__ MlpTileInfo next() { return remote.next(); } + __device__ void acquire_src() { remote.acquire_src(); } + __device__ void release_src(int lane) { remote.release_src(lane); } + __device__ void release_src_slot(int slot, int lane) { remote.release_src_slot(slot, lane); } + __device__ void acquire_dy() { remote.acquire_dy(); } + __device__ void release_dy(int lane) { remote.release_dy(lane); } + __device__ void release_dy_slot(int slot, int lane) { remote.release_dy_slot(slot, lane); } + __device__ void acquire_dst() { remote.acquire_dst(); } + __device__ void release_dst(int lane) { remote.release_dst(lane); } +}; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm100.cuh new file mode 100644 index 00000000..50a66b01 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm100.cuh @@ -0,0 +1,527 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Tile Iterators +// ═══════════════════════════════════════════════════════════════════ +// +// Provide a uniform interface for iterating over M-tiles in the +// fused MLP kernel. Each iterator yields MlpTileInfo per tile. +// +// Both iterators support staging acquire/release hooks: +// acquire_src / release_src — around phase 1 (X read) +// acquire_dst / release_dst — around phase 2 (Y write) +// For local tiles these are no-ops. For remote tiles they +// synchronize with the comm warps via StagePipe. +// +// ═══════════════════════════════════════════════════════════════════ + +#include +#include +#include "mlp_comms_sm100.cuh" + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// MlpTileInfo — what the M-loop body needs per tile +// ═══════════════════════════════════════════════════════════════════ + +struct MlpTileInfo { + const void* x_ptr; // pointer to the start of this M-tile's X data + int expert; // expert ID for weight selection + int y_m; // staging slot — X-load coord AND staging Y-store coord + // ── Direct-to-peer Y store (cur_is_local tiles) ────────────────────── + int y_store_m; // token-tile index into the destination's local_output + // (= token_offset / TileM); the TMA store row. + int peer_rank; // destination peer's intra-host local rank — selects the + // per-peer TMA descriptor built over nvshmem_ptr(y_buf, pe). + bool is_local; // same-host (NVLink) → GEMM TMA-stores Y straight into the + // peer's symmetric local_output, bypassing dst_staging+put. + int valid_m64_subtiles; // number of 64-row output subtiles containing valid rows + int m_subtile; // GemmTileM row within the parent communication tile +}; + +// ═══════════════════════════════════════════════════════════════════ +// FlatMlpWorkItem — per-(m, n-chunk) work item for the flat scheduler +// ═══════════════════════════════════════════════════════════════════ +// +// The flat scheduler hands these out per stage per phase. A CTA loops +// while (iter.has_next()) { auto item = iter.next(); ... } +// processing each item with the per-(m, [n_start, n_end)) chunk +// producer/consumer. + +struct FlatMlpWorkItem { + int m; // global M-tile index (into expert_ids / arrival counters) + int expert; // expert ID after local_expert_start subtraction + int z_m; // Z-buffer M-slot index (double-buffered across stages) + int n_start; // chunk start n-tile (inclusive) + int n_end; // chunk end n-tile (exclusive); ≤ n_start + NS + int m_prev_wait; // global m whose mlp2_arrival must reach num_n_tiles_2 + // before this chunk can write Z[z_m]. -1 if no wait + // (first ZBufferSlots stages have no prior occupant). + int phase; // 0 = MLP1, 1 = MLP2 +}; + +// ═══════════════════════════════════════════════════════════════════ +// LocalMlpTileIterator +// ═══════════════════════════════════════════════════════════════════ + +template +struct LocalMlpTileIterator { + const Element* x_base; + const int* expert_ids; + int num_m_tiles; + int hidden_dim; + int tile_offset; // global tile offset (for TMA coordinates) + + int m; + int m_end; // exclusive upper bound for the current sub-range + int col_stride; // logical grid_x — column stride between this CTA's tiles + // (flat-grid launch: NOT gridDim.x, which is num_blocks) + + __device__ void init(const Element* x_base_, const int* expert_ids_, + int num_m_tiles_, int hidden_dim_, + int col_stride_, int start_col_, + int tile_offset_ = 0) { + x_base = x_base_; + expert_ids = expert_ids_; + num_m_tiles = num_m_tiles_; + hidden_dim = hidden_dim_; + tile_offset = tile_offset_; + // Logical column / stride passed explicitly (flat-grid launch) instead + // of reading blockIdx.x / gridDim.x. The fused driver overrides m via + // set_range() before iterating; standalone callers use m = start_col. + col_stride = col_stride_; + m = start_col_; + m_end = num_m_tiles; + } + + // Restrict iteration to [m_start, m_end_). Used by moe_fused_fwd to + // split the local M-loop around the remote phase so per-tile compute + // overlaps with remote get/put traffic. + __device__ void set_range(int m_start, int m_end_) { + m = m_start; + m_end = m_end_; + } + + __device__ bool has_next() const { + return m < m_end; + } + + // y_m is the global M-tile index (for TMA coordinates into full buffer). + __device__ MlpTileInfo next() { + MlpTileInfo info; + info.x_ptr = x_base + (size_t)m * TileM * hidden_dim; + info.expert = expert_ids[m]; + info.y_m = tile_offset + m; + // Local (non-fused) iteration never uses the direct-to-peer store path. + info.y_store_m = info.y_m; + info.peer_rank = 0; + info.is_local = false; + info.valid_m64_subtiles = TileM / 64; + info.m_subtile = 0; + + m += col_stride; + return info; + } + + // No-op staging hooks for local tiles. + __device__ void acquire_src() {} + __device__ void release_src(int /*lane*/) {} + __device__ void acquire_dst() {} + __device__ void release_dst(int /*lane*/) {} +}; + +// ═══════════════════════════════════════════════════════════════════ +// LocalFlatMlpTileIterator — chunked-flat scheduler work-item iterator +// ═══════════════════════════════════════════════════════════════════ +// +// 1D iteration interface. The iterator's internal state machine walks +// stages and phases (MLP1 then MLP2 per stage); each next() returns a +// FlatMlpWorkItem with m / n_start / n_end / expert / z_m / m_prev_wait +// / phase already populated. has_next() is mutating — it advances the +// state machine past phase / stage boundaries until either it finds +// the next work item or signals end-of-iteration. +// +// Grid is 2D: dim3(num_m_per_stage, NC). Total CTAs = num_m_per_stage·NC. +// Per-CTA flat index = blockIdx.x + blockIdx.y · gridDim.x; stride +// over the per-phase work pool = gridDim.x · gridDim.y. +// +// Usage: +// iter.init(expert_ids, num_m_tiles, num_m_per_stage, +// num_n_tiles_1, NS1, num_n_tiles_2, NS2); +// while (iter.has_next()) { +// FlatMlpWorkItem item = iter.next(); +// if (item.phase == 0) { /* MLP1 */ } else { /* MLP2 */ } +// } + +template +struct LocalFlatMlpTileIterator { + const int* expert_ids; + int num_m_tiles; + int num_m_per_stage; + int num_stages; + int local_expert_start; + + // Phase parameters captured at init. + int num_n_tiles_1; + int NS1; + int num_n_chunks_1; + int num_n_tiles_2; + int NS2; + int num_n_chunks_2; + + // CTA identity (decoded once at init). + int flat_cta_idx; // blockIdx.x + blockIdx.y · gridDim.x + int total_ctas; // gridDim.x · gridDim.y + + // State machine. + int stage; // current stage index (-1 before first has_next) + int stage_m_start; + int stage_m_count; + bool in_mlp1; // true → MLP1 active, false → MLP2 active + int total_work_phase; + int work_idx; + + __device__ void init(const int* expert_ids_, + int num_m_tiles_, + int num_m_per_stage_, + int num_n_tiles_1_, + int NS1_, + int num_n_tiles_2_, + int NS2_, + int local_expert_start_ = 0) { + expert_ids = expert_ids_; + num_m_tiles = num_m_tiles_; + num_m_per_stage = num_m_per_stage_; + num_stages = (num_m_tiles_ + num_m_per_stage_ - 1) / num_m_per_stage_; + local_expert_start = local_expert_start_; + num_n_tiles_1 = num_n_tiles_1_; + NS1 = NS1_; + num_n_chunks_1 = (num_n_tiles_1_ + NS1_ - 1) / NS1_; + num_n_tiles_2 = num_n_tiles_2_; + NS2 = NS2_; + num_n_chunks_2 = (num_n_tiles_2_ + NS2_ - 1) / NS2_; + flat_cta_idx = blockIdx.x + blockIdx.y * gridDim.x; + total_ctas = gridDim.x * gridDim.y; + + // Pre-state: first has_next() will advance into (stage=0, MLP1). + stage = -1; + stage_m_start = 0; + stage_m_count = 0; + in_mlp1 = false; + total_work_phase = 0; + work_idx = flat_cta_idx; + } + + // Advance to the next (stage, phase) that has any work for this CTA. + // Returns true if there is a work item to emit at the current state, + // false when iteration is exhausted. + __device__ bool has_next() { + while (work_idx >= total_work_phase) { + if (in_mlp1) { + // Finished MLP1 in this stage — start MLP2 (same stage). + in_mlp1 = false; + total_work_phase = stage_m_count * num_n_chunks_2; + work_idx = flat_cta_idx; + } else { + // Finished MLP2 (or pre-state) — advance to next stage's MLP1. + ++stage; + if (stage >= num_stages) return false; + stage_m_start = stage * num_m_per_stage; + int rem = num_m_tiles - stage_m_start; + stage_m_count = (rem < num_m_per_stage) ? rem : num_m_per_stage; + in_mlp1 = true; + total_work_phase = stage_m_count * num_n_chunks_1; + work_idx = flat_cta_idx; + } + } + return true; + } + + __device__ FlatMlpWorkItem next() { + FlatMlpWorkItem item; + int num_n_chunks_phase = in_mlp1 ? num_n_chunks_1 : num_n_chunks_2; + int num_n_tiles_phase = in_mlp1 ? num_n_tiles_1 : num_n_tiles_2; + int NS_phase = in_mlp1 ? NS1 : NS2; + int m_in_stage = work_idx / num_n_chunks_phase; + int chunk_idx = work_idx % num_n_chunks_phase; + item.m = stage_m_start + m_in_stage; + item.expert = expert_ids[item.m] - local_expert_start; + item.z_m = (stage % ZBufferSlots) * num_m_per_stage + m_in_stage; + item.n_start = chunk_idx * NS_phase; + int n_end = item.n_start + NS_phase; + item.n_end = (n_end < num_n_tiles_phase) ? n_end : num_n_tiles_phase; + if (in_mlp1 && stage >= ZBufferSlots) { + item.m_prev_wait = (stage - ZBufferSlots) * num_m_per_stage + m_in_stage; + } else { + item.m_prev_wait = -1; + } + item.phase = in_mlp1 ? 0 : 1; + work_idx += total_ctas; + return item; + } + + // No-op staging hooks for local (standalone) iteration; RemoteFlat + // variant will wire these to NVSHMEM src/dst pipes. + __device__ void acquire_src() {} + __device__ void release_src(int /*lane*/) {} + __device__ void acquire_dst() {} + __device__ void release_dst(int /*lane*/) {} +}; + +// ═══════════════════════════════════════════════════════════════════ +// RemoteMlpTileIterator — ticket-based slot synchronization +// ═══════════════════════════════════════════════════════════════════ +// +// Iterates over remote tiles prefetched by comm warps. Staging area +// is a flat ring of L = MC · NumStages slots, where MC = (gridDim.x · +// gridDim.y) / NC is the M-tile count produced per comm pipeline stage. +// +// Producer/consumer counts per slot (unchanged from comm side): +// src: 2·NC get warps produce, runtime NS GEMM CTAs consume. +// dst: runtime NS GEMM CTAs produce, NC put warps consume. +// +// ── Ring walk ────────────────────────────────────────────────────── +// Each column blockIdx.x walks its tile sequence with stride gridDim.x +// and wraps the slot index modulo L: +// T_j = blockIdx.x + j * gridDim.x (global tile index for j-th iter) +// slot = T_j mod L (physical staging slot) +// ticket = T_j / L (pass count over the ring) +// +// Comm writes tile T to the same slot (= T mod L), so producer and +// consumer always agree on the slot. There is no per-stage-id pipe +// state: the ticket carries enough information to compute the wait +// threshold directly. +// +// acquire_src(): wait for src_ready[slot] >= (ticket+1) · 2NC +// release_src(): atomic increment src_consumed[slot] by 1 (per CTA) +// acquire_dst(): wait for dst_consumed[slot] >= ticket · NC +// (first visit per slot — ticket == 0 — has no +// outstanding put to drain, skip the spin) +// release_dst(): atomic increment dst_ready[slot] by 1 (per CTA) +// +// With runtime NS CTAs cooperating per column, each contributes one +// atomic per release, so src_consumed and dst_ready advance by +// runtime NS per consumed/produced tile — exactly what the comm-side +// StagePipe expects. +// +// ── Decoupling from comm-side constants ──────────────────────────── +// Because the ticket comes from the global tile index (not a +// per-CTA cyclic counter), NumStages, runtime NS and NC need not share +// any divisibility relationship. The only correctness requirement is +// that NC evenly divides gridDim.x · gridDim.y so MC is integer +// (host-side concern, not a template constraint). + +template +struct RemoteMlpTileIterator { + // Slot's producer/consumer counts (compile-time, used for ticket math). + // FWD steady-state GET uses kNumGetWarpsFwd active warps/CTA (the TMA path + // saturates with one), so the MLP waits for kNumGetWarpsFwd·NC src signals. + static_assert(NC % MSubTiles == 0, + "communication CTAs must divide evenly across GEMM M-subtiles"); + static constexpr int RowNC = NC / MSubTiles; + static constexpr int NumProducersSrc = kNumGetWarpsFwd * RowNC; + static constexpr int NumConsumersDst = kNumPutWarpsFwd * RowNC; + + // (expert, pe) are re-derived in SMEM, NOT read from HBM: the MLP's j-th tile + // is global flat tile T_j = m_base + j·col_stride — the same flat position the + // comm placed at slot=T_j%L, ticket=T_j/L (the StagePipe ticket handshake + // guarantees the rendezvous). So an embedded TileIterator init(m_base, + // col_stride) walked in lockstep with idx reproduces what the comm wrote, + // with zero per-tile HBM load and fewer registers than holding HBM pointers. + GemmTileIterator tit; + bool derive_on; // false → no remote_offsets (test harness) → local + int gpus_per_node; // same-host divisor for is_local + int iter_my_pe; // this PE (team space) for is_local + int cur_is_local; // cached per-tile fence-scope predicate + + int total_tiles; + int m_base; // logical column index (flat_id / runtime NS) — column's + // tile-sequence start + int col_stride; // logical grid_x (= n_gemm / runtime NS) — tile-sequence + // stride; NOT gridDim.x under the flat-grid launch + int ring_len; // physical communication slots: L = MC · NumStages + int num_splits; + + int* src_ready; + int* src_consumed; + int* dst_ready; + int* dst_consumed; + + bool is_leader; // per-warp leader (lane 0) — does the atomic poll + + int idx; // iteration counter within this column + int cur_slot; // cached slot for current acquire/release cycle + int cur_data_slot; + int cur_ticket; // cached ticket for current acquire/release cycle + int cur_m_subtile; + int cur_parent_tile; + + // remote_offsets (SMEM) drives the embedded walker; pass nullptr to disable + // derivation (e.g. the standalone handshake test, which only exercises the + // ticket/slot protocol). experts_per_pe/num_pes/my_pe/gpus_per_node feed the + // walker and the same-host is_local predicate. + __device__ void init(const int* remote_offsets, + const int* remote_counts, + int experts_per_pe, + int num_pes, + int my_pe, + int gpus_per_node_, + int total_tiles_, + int m_base_, + int n_gemm_, + int* src_ready_, int* src_consumed_, + int* dst_ready_, int* dst_consumed_, + bool is_leader_, + int runtime_nsplit_) { + total_tiles = total_tiles_; + m_base = m_base_; + num_splits = runtime_nsplit_; + col_stride = n_gemm_ / num_splits; + idx = 0; + + // L = MC · NumStages, MC = n_gemm / NC. + int mc = gridDim.x / NC; + ring_len = mc * NumStages; + + src_ready = src_ready_; + src_consumed = src_consumed_; + dst_ready = dst_ready_; + dst_consumed = dst_consumed_; + is_leader = is_leader_; + cur_slot = 0; + cur_ticket = 0; + + gpus_per_node = gpus_per_node_; + iter_my_pe = my_pe; + cur_is_local = 1; + derive_on = (remote_offsets != nullptr); + if (derive_on) { + tit.remote_offsets = remote_offsets; + tit.offsets_stride = experts_per_pe + 1; + tit.remote_counts = remote_counts; + tit.counts_stride = experts_per_pe; + tit.experts_per_pe = experts_per_pe; + tit.num_pes = num_pes; + tit.my_pe = my_pe; + tit.init(m_base_, col_stride); + } + } + + __device__ bool has_next() const { + return idx < total_tiles; + } + + // Compute and cache (slot, ticket) for the current iteration. + // Called by acquire_src() — the rest of the per-tile API (next, + // release_src, acquire_dst, release_dst) reuses the cached values. + __device__ void compute_slot() { + int parent_t; + if (derive_on) { + parent_t = tit.current_parent_tile(); + cur_m_subtile = tit.current_m_subtile(); + } else { + int virtual_t = m_base + idx * col_stride; + parent_t = virtual_t / MSubTiles; + cur_m_subtile = virtual_t % MSubTiles; + } + cur_parent_tile = parent_t; + cur_data_slot = parent_t % ring_len; + cur_slot = cur_data_slot * MSubTiles + cur_m_subtile; + cur_ticket = parent_t / ring_len; + } + + // Re-derive (expert, pe) for this tile from SMEM offsets — the embedded + // walker advances in lockstep with idx and reproduces what the comm wrote at + // this slot (proven: same global flat position). No HBM load. + __device__ MlpTileInfo next() { + MlpTileInfo info; + info.x_ptr = nullptr; // not used — TMA loads by coordinate + info.y_m = cur_data_slot; + if (derive_on) { + TileInfo ti = tit.next(); + info.expert = ti.expert; + bool loc = (gpus_per_node > 0) + ? ((ti.pe / gpus_per_node) == (iter_my_pe / gpus_per_node)) + : true; + cur_is_local = (int)loc; + info.is_local = loc; + // Destination peer's intra-host local rank → per-peer y_buf descriptor. + info.peer_rank = (gpus_per_node > 0) ? (ti.pe % gpus_per_node) : 0; + // Direct-store row = this tile's token offset in the peer's local_output. + info.y_store_m = ti.token_offset / TileM; + info.valid_m64_subtiles = (ti.valid_rows + 63) / 64; + info.m_subtile = ti.m_subtile; + } else { + // Test harness (no SMEM offsets): no derivation, no direct store. + info.expert = 0; + cur_is_local = 1; + info.is_local = false; + info.peer_rank = 0; + info.y_store_m = cur_slot; + info.valid_m64_subtiles = TileM / 64; + info.m_subtile = 0; + } + idx++; + return info; + } + + // Wait for comm's get to finish writing this tile to its slot. + // Producer (comm) signals ready += 2·NC per tile-write at the slot, + // so after ticket K the wait threshold is (K+1)·2·NC. + __device__ void acquire_src() { + compute_slot(); + int target = (cur_ticket + 1) * NumProducersSrc; + if (is_leader) { + // Relaxed spin (ld.relaxed.gpu — cheapest read, no RMW, no per-iter + // fence); the __threadfence below establishes the acquire ordering + // once for all lanes on exit. atomic_ref ctor is zero-cost/hoisted. + cuda::atomic_ref r(src_ready[cur_slot]); + while (r.load(cuda::memory_order_relaxed) < target) {} + } + __syncwarp(); + // Device-scope: observing NIC-written src. The NIC's inbound-write + // visibility is not fixable by a fence (needs GDRCopy), so system scope + // here only costs bandwidth. Device fence covers the GPU-internal view. + __threadfence(); + } + + // Signal that this CTA has consumed the tile at cur_slot. Per + // tile (column-of-runtime-NS CTAs), src_consumed advances by runtime NS. + __device__ void release_src(int lane) { + __threadfence(); + if (lane == 0) { + atomicAdd(&src_consumed[cur_slot], 1); + } + } + + // Wait for the previous put from cur_slot to finish (so we can + // reuse the slot for a fresh Y write). ticket==0 → no prior put. + __device__ void acquire_dst() { + if (cur_ticket >= 1) { + int target = cur_ticket * NumConsumersDst; + if (is_leader) { + cuda::atomic_ref r(dst_consumed[cur_slot]); + while (r.load(cuda::memory_order_relaxed) < target) {} + } + __syncwarp(); + __threadfence(); + } + } + + // Signal that Y is ready in cur_slot for the comm put. Per tile, + // dst_ready advances by runtime NS. + __device__ void release_dst(int lane) { + // Same-host tiles store Y directly into the peer's symmetric + // local_output, so the put warp has no staged data to read. Cross-host + // tiles use dst_staging + putmem and require a system-scope fence before + // the put warp observes dst_ready. + if (!cur_is_local) __threadfence_system(); + if (lane == 0) { + atomicAdd(&dst_ready[cur_slot], 1); + } + } +}; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm90.cuh new file mode 100644 index 00000000..0c24d0f7 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm90.cuh @@ -0,0 +1,527 @@ +#pragma once + +// ═══════════════════════════════════════════════════════════════════ +// MLP Tile Iterators +// ═══════════════════════════════════════════════════════════════════ +// +// Provide a uniform interface for iterating over M-tiles in the +// fused MLP kernel. Each iterator yields MlpTileInfo per tile. +// +// Both iterators support staging acquire/release hooks: +// acquire_src / release_src — around phase 1 (X read) +// acquire_dst / release_dst — around phase 2 (Y write) +// For local tiles these are no-ops. For remote tiles they +// synchronize with the comm warps via StagePipe. +// +// ═══════════════════════════════════════════════════════════════════ + +#include +#include +#include "mlp_comms_sm90.cuh" + +namespace liger { + +// ═══════════════════════════════════════════════════════════════════ +// MlpTileInfo — what the M-loop body needs per tile +// ═══════════════════════════════════════════════════════════════════ + +struct MlpTileInfo { + const void* x_ptr; // pointer to the start of this M-tile's X data + int expert; // expert ID for weight selection + int y_m; // staging slot — X-load coord AND staging Y-store coord + // ── Direct-to-peer Y store (cur_is_local tiles) ────────────────────── + int y_store_m; // token-tile index into the destination's local_output + // (= token_offset / TileM); the TMA store row. + int peer_rank; // destination peer's intra-host local rank — selects the + // per-peer TMA descriptor built over nvshmem_ptr(y_buf, pe). + bool is_local; // same-host (NVLink) → GEMM TMA-stores Y straight into the + // peer's symmetric local_output, bypassing dst_staging+put. + int valid_m64_subtiles; // number of 64-row output subtiles containing valid rows + int m_subtile; // GemmTileM row within the parent communication tile +}; + +// ═══════════════════════════════════════════════════════════════════ +// FlatMlpWorkItem — per-(m, n-chunk) work item for the flat scheduler +// ═══════════════════════════════════════════════════════════════════ +// +// The flat scheduler hands these out per stage per phase. A CTA loops +// while (iter.has_next()) { auto item = iter.next(); ... } +// processing each item with the per-(m, [n_start, n_end)) chunk +// producer/consumer. + +struct FlatMlpWorkItem { + int m; // global M-tile index (into expert_ids / arrival counters) + int expert; // expert ID after local_expert_start subtraction + int z_m; // Z-buffer M-slot index (double-buffered across stages) + int n_start; // chunk start n-tile (inclusive) + int n_end; // chunk end n-tile (exclusive); ≤ n_start + NS + int m_prev_wait; // global m whose mlp2_arrival must reach num_n_tiles_2 + // before this chunk can write Z[z_m]. -1 if no wait + // (first ZBufferSlots stages have no prior occupant). + int phase; // 0 = MLP1, 1 = MLP2 +}; + +// ═══════════════════════════════════════════════════════════════════ +// LocalMlpTileIterator +// ═══════════════════════════════════════════════════════════════════ + +template +struct LocalMlpTileIterator { + const Element* x_base; + const int* expert_ids; + int num_m_tiles; + int hidden_dim; + int tile_offset; // global tile offset (for TMA coordinates) + + int m; + int m_end; // exclusive upper bound for the current sub-range + int col_stride; // logical grid_x — column stride between this CTA's tiles + // (flat-grid launch: NOT gridDim.x, which is num_blocks) + + __device__ void init(const Element* x_base_, const int* expert_ids_, + int num_m_tiles_, int hidden_dim_, + int col_stride_, int start_col_, + int tile_offset_ = 0) { + x_base = x_base_; + expert_ids = expert_ids_; + num_m_tiles = num_m_tiles_; + hidden_dim = hidden_dim_; + tile_offset = tile_offset_; + // Logical column / stride passed explicitly (flat-grid launch) instead + // of reading blockIdx.x / gridDim.x. The fused driver overrides m via + // set_range() before iterating; standalone callers use m = start_col. + col_stride = col_stride_; + m = start_col_; + m_end = num_m_tiles; + } + + // Restrict iteration to [m_start, m_end_). Used by moe_fused_fwd to + // split the local M-loop around the remote phase so per-tile compute + // overlaps with remote get/put traffic. + __device__ void set_range(int m_start, int m_end_) { + m = m_start; + m_end = m_end_; + } + + __device__ bool has_next() const { + return m < m_end; + } + + // y_m is the global M-tile index (for TMA coordinates into full buffer). + __device__ MlpTileInfo next() { + MlpTileInfo info; + info.x_ptr = x_base + (size_t)m * TileM * hidden_dim; + info.expert = expert_ids[m]; + info.y_m = tile_offset + m; + // Local (non-fused) iteration never uses the direct-to-peer store path. + info.y_store_m = info.y_m; + info.peer_rank = 0; + info.is_local = false; + info.valid_m64_subtiles = TileM / 64; + info.m_subtile = 0; + + m += col_stride; + return info; + } + + // No-op staging hooks for local tiles. + __device__ void acquire_src() {} + __device__ void release_src(int /*lane*/) {} + __device__ void acquire_dst() {} + __device__ void release_dst(int /*lane*/) {} +}; + +// ═══════════════════════════════════════════════════════════════════ +// LocalFlatMlpTileIterator — chunked-flat scheduler work-item iterator +// ═══════════════════════════════════════════════════════════════════ +// +// 1D iteration interface. The iterator's internal state machine walks +// stages and phases (MLP1 then MLP2 per stage); each next() returns a +// FlatMlpWorkItem with m / n_start / n_end / expert / z_m / m_prev_wait +// / phase already populated. has_next() is mutating — it advances the +// state machine past phase / stage boundaries until either it finds +// the next work item or signals end-of-iteration. +// +// Grid is 2D: dim3(num_m_per_stage, NC). Total CTAs = num_m_per_stage·NC. +// Per-CTA flat index = blockIdx.x + blockIdx.y · gridDim.x; stride +// over the per-phase work pool = gridDim.x · gridDim.y. +// +// Usage: +// iter.init(expert_ids, num_m_tiles, num_m_per_stage, +// num_n_tiles_1, NS1, num_n_tiles_2, NS2); +// while (iter.has_next()) { +// FlatMlpWorkItem item = iter.next(); +// if (item.phase == 0) { /* MLP1 */ } else { /* MLP2 */ } +// } + +template +struct LocalFlatMlpTileIterator { + const int* expert_ids; + int num_m_tiles; + int num_m_per_stage; + int num_stages; + int local_expert_start; + + // Phase parameters captured at init. + int num_n_tiles_1; + int NS1; + int num_n_chunks_1; + int num_n_tiles_2; + int NS2; + int num_n_chunks_2; + + // CTA identity (decoded once at init). + int flat_cta_idx; // blockIdx.x + blockIdx.y · gridDim.x + int total_ctas; // gridDim.x · gridDim.y + + // State machine. + int stage; // current stage index (-1 before first has_next) + int stage_m_start; + int stage_m_count; + bool in_mlp1; // true → MLP1 active, false → MLP2 active + int total_work_phase; + int work_idx; + + __device__ void init(const int* expert_ids_, + int num_m_tiles_, + int num_m_per_stage_, + int num_n_tiles_1_, + int NS1_, + int num_n_tiles_2_, + int NS2_, + int local_expert_start_ = 0) { + expert_ids = expert_ids_; + num_m_tiles = num_m_tiles_; + num_m_per_stage = num_m_per_stage_; + num_stages = (num_m_tiles_ + num_m_per_stage_ - 1) / num_m_per_stage_; + local_expert_start = local_expert_start_; + num_n_tiles_1 = num_n_tiles_1_; + NS1 = NS1_; + num_n_chunks_1 = (num_n_tiles_1_ + NS1_ - 1) / NS1_; + num_n_tiles_2 = num_n_tiles_2_; + NS2 = NS2_; + num_n_chunks_2 = (num_n_tiles_2_ + NS2_ - 1) / NS2_; + flat_cta_idx = blockIdx.x + blockIdx.y * gridDim.x; + total_ctas = gridDim.x * gridDim.y; + + // Pre-state: first has_next() will advance into (stage=0, MLP1). + stage = -1; + stage_m_start = 0; + stage_m_count = 0; + in_mlp1 = false; + total_work_phase = 0; + work_idx = flat_cta_idx; + } + + // Advance to the next (stage, phase) that has any work for this CTA. + // Returns true if there is a work item to emit at the current state, + // false when iteration is exhausted. + __device__ bool has_next() { + while (work_idx >= total_work_phase) { + if (in_mlp1) { + // Finished MLP1 in this stage — start MLP2 (same stage). + in_mlp1 = false; + total_work_phase = stage_m_count * num_n_chunks_2; + work_idx = flat_cta_idx; + } else { + // Finished MLP2 (or pre-state) — advance to next stage's MLP1. + ++stage; + if (stage >= num_stages) return false; + stage_m_start = stage * num_m_per_stage; + int rem = num_m_tiles - stage_m_start; + stage_m_count = (rem < num_m_per_stage) ? rem : num_m_per_stage; + in_mlp1 = true; + total_work_phase = stage_m_count * num_n_chunks_1; + work_idx = flat_cta_idx; + } + } + return true; + } + + __device__ FlatMlpWorkItem next() { + FlatMlpWorkItem item; + int num_n_chunks_phase = in_mlp1 ? num_n_chunks_1 : num_n_chunks_2; + int num_n_tiles_phase = in_mlp1 ? num_n_tiles_1 : num_n_tiles_2; + int NS_phase = in_mlp1 ? NS1 : NS2; + int m_in_stage = work_idx / num_n_chunks_phase; + int chunk_idx = work_idx % num_n_chunks_phase; + item.m = stage_m_start + m_in_stage; + item.expert = expert_ids[item.m] - local_expert_start; + item.z_m = (stage % ZBufferSlots) * num_m_per_stage + m_in_stage; + item.n_start = chunk_idx * NS_phase; + int n_end = item.n_start + NS_phase; + item.n_end = (n_end < num_n_tiles_phase) ? n_end : num_n_tiles_phase; + if (in_mlp1 && stage >= ZBufferSlots) { + item.m_prev_wait = (stage - ZBufferSlots) * num_m_per_stage + m_in_stage; + } else { + item.m_prev_wait = -1; + } + item.phase = in_mlp1 ? 0 : 1; + work_idx += total_ctas; + return item; + } + + // No-op staging hooks for local (standalone) iteration; RemoteFlat + // variant will wire these to NVSHMEM src/dst pipes. + __device__ void acquire_src() {} + __device__ void release_src(int /*lane*/) {} + __device__ void acquire_dst() {} + __device__ void release_dst(int /*lane*/) {} +}; + +// ═══════════════════════════════════════════════════════════════════ +// RemoteMlpTileIterator — ticket-based slot synchronization +// ═══════════════════════════════════════════════════════════════════ +// +// Iterates over remote tiles prefetched by comm warps. Staging area +// is a flat ring of L = MC · NumStages slots, where MC = (gridDim.x · +// gridDim.y) / NC is the M-tile count produced per comm pipeline stage. +// +// Producer/consumer counts per slot (unchanged from comm side): +// src: 2·NC get warps produce, runtime NS GEMM CTAs consume. +// dst: runtime NS GEMM CTAs produce, NC put warps consume. +// +// ── Ring walk ────────────────────────────────────────────────────── +// Each column blockIdx.x walks its tile sequence with stride gridDim.x +// and wraps the slot index modulo L: +// T_j = blockIdx.x + j * gridDim.x (global tile index for j-th iter) +// slot = T_j mod L (physical staging slot) +// ticket = T_j / L (pass count over the ring) +// +// Comm writes tile T to the same slot (= T mod L), so producer and +// consumer always agree on the slot. There is no per-stage-id pipe +// state: the ticket carries enough information to compute the wait +// threshold directly. +// +// acquire_src(): wait for src_ready[slot] >= (ticket+1) · 2NC +// release_src(): atomic increment src_consumed[slot] by 1 (per CTA) +// acquire_dst(): wait for dst_consumed[slot] >= ticket · NC +// (first visit per slot — ticket == 0 — has no +// outstanding put to drain, skip the spin) +// release_dst(): atomic increment dst_ready[slot] by 1 (per CTA) +// +// With runtime NS CTAs cooperating per column, each contributes one +// atomic per release, so src_consumed and dst_ready advance by +// runtime NS per consumed/produced tile — exactly what the comm-side +// StagePipe expects. +// +// ── Decoupling from comm-side constants ──────────────────────────── +// Because the ticket comes from the global tile index (not a +// per-CTA cyclic counter), NumStages, runtime NS and NC need not share +// any divisibility relationship. The only correctness requirement is +// that NC evenly divides gridDim.x · gridDim.y so MC is integer +// (host-side concern, not a template constraint). + +template +struct RemoteMlpTileIterator { + // Slot's producer/consumer counts (compile-time, used for ticket math). + // FWD steady-state GET uses kNumGetWarpsFwd active warps/CTA (the TMA path + // saturates with one), so the MLP waits for kNumGetWarpsFwd·NC src signals. + static_assert(NC % MSubTiles == 0, + "communication CTAs must divide evenly across GEMM M-subtiles"); + static constexpr int RowNC = NC / MSubTiles; + static constexpr int NumProducersSrc = kNumGetWarpsFwd * RowNC; + static constexpr int NumConsumersDst = kNumPutWarpsFwd * RowNC; + + // (expert, pe) are re-derived in SMEM, NOT read from HBM: the MLP's j-th tile + // is global flat tile T_j = m_base + j·col_stride — the same flat position the + // comm placed at slot=T_j%L, ticket=T_j/L (the StagePipe ticket handshake + // guarantees the rendezvous). So an embedded TileIterator init(m_base, + // col_stride) walked in lockstep with idx reproduces what the comm wrote, + // with zero per-tile HBM load and fewer registers than holding HBM pointers. + GemmTileIterator tit; + bool derive_on; // false → no remote_offsets (test harness) → local + int gpus_per_node; // same-host divisor for is_local + int iter_my_pe; // this PE (team space) for is_local + int cur_is_local; // cached per-tile fence-scope predicate + + int total_tiles; + int m_base; // logical column index (flat_id / runtime NS) — column's + // tile-sequence start + int col_stride; // logical grid_x (= n_gemm / runtime NS) — tile-sequence + // stride; NOT gridDim.x under the flat-grid launch + int ring_len; // physical communication slots: L = MC · NumStages + int num_splits; + + int* src_ready; + int* src_consumed; + int* dst_ready; + int* dst_consumed; + + bool is_leader; // per-warp leader (lane 0) — does the atomic poll + + int idx; // iteration counter within this column + int cur_slot; // cached slot for current acquire/release cycle + int cur_data_slot; + int cur_ticket; // cached ticket for current acquire/release cycle + int cur_m_subtile; + int cur_parent_tile; + + // remote_offsets (SMEM) drives the embedded walker; pass nullptr to disable + // derivation (e.g. the standalone handshake test, which only exercises the + // ticket/slot protocol). experts_per_pe/num_pes/my_pe/gpus_per_node feed the + // walker and the same-host is_local predicate. + __device__ void init(const int* remote_offsets, + const int* remote_counts, + int experts_per_pe, + int num_pes, + int my_pe, + int gpus_per_node_, + int total_tiles_, + int m_base_, + int n_gemm_, + int* src_ready_, int* src_consumed_, + int* dst_ready_, int* dst_consumed_, + bool is_leader_, + int runtime_nsplit_) { + total_tiles = total_tiles_; + m_base = m_base_; + num_splits = runtime_nsplit_; + col_stride = n_gemm_ / num_splits; + idx = 0; + + // L = MC · NumStages, MC = n_gemm / NC. + int mc = gridDim.x / NC; + ring_len = mc * NumStages; + + src_ready = src_ready_; + src_consumed = src_consumed_; + dst_ready = dst_ready_; + dst_consumed = dst_consumed_; + is_leader = is_leader_; + cur_slot = 0; + cur_ticket = 0; + + gpus_per_node = gpus_per_node_; + iter_my_pe = my_pe; + cur_is_local = 1; + derive_on = (remote_offsets != nullptr); + if (derive_on) { + tit.remote_offsets = remote_offsets; + tit.offsets_stride = experts_per_pe + 1; + tit.remote_counts = remote_counts; + tit.counts_stride = experts_per_pe; + tit.experts_per_pe = experts_per_pe; + tit.num_pes = num_pes; + tit.my_pe = my_pe; + tit.init(m_base_, col_stride); + } + } + + __device__ bool has_next() const { + return idx < total_tiles; + } + + // Compute and cache (slot, ticket) for the current iteration. + // Called by acquire_src() — the rest of the per-tile API (next, + // release_src, acquire_dst, release_dst) reuses the cached values. + __device__ void compute_slot() { + int parent_t; + if (derive_on) { + parent_t = tit.current_parent_tile(); + cur_m_subtile = tit.current_m_subtile(); + } else { + int virtual_t = m_base + idx * col_stride; + parent_t = virtual_t / MSubTiles; + cur_m_subtile = virtual_t % MSubTiles; + } + cur_parent_tile = parent_t; + cur_data_slot = parent_t % ring_len; + cur_slot = cur_data_slot * MSubTiles + cur_m_subtile; + cur_ticket = parent_t / ring_len; + } + + // Re-derive (expert, pe) for this tile from SMEM offsets — the embedded + // walker advances in lockstep with idx and reproduces what the comm wrote at + // this slot (proven: same global flat position). No HBM load. + __device__ MlpTileInfo next() { + MlpTileInfo info; + info.x_ptr = nullptr; // not used — TMA loads by coordinate + info.y_m = cur_data_slot; + if (derive_on) { + TileInfo ti = tit.next(); + info.expert = ti.expert; + bool loc = (gpus_per_node > 0) + ? ((ti.pe / gpus_per_node) == (iter_my_pe / gpus_per_node)) + : true; + cur_is_local = (int)loc; + info.is_local = loc; + // Destination peer's intra-host local rank → per-peer y_buf descriptor. + info.peer_rank = (gpus_per_node > 0) ? (ti.pe % gpus_per_node) : 0; + // Direct-store row = this tile's token offset in the peer's local_output. + info.y_store_m = ti.token_offset / TileM; + info.valid_m64_subtiles = (ti.valid_rows + 63) / 64; + info.m_subtile = ti.m_subtile; + } else { + // Test harness (no SMEM offsets): no derivation, no direct store. + info.expert = 0; + cur_is_local = 1; + info.is_local = false; + info.peer_rank = 0; + info.y_store_m = cur_slot; + info.valid_m64_subtiles = TileM / 64; + info.m_subtile = 0; + } + idx++; + return info; + } + + // Wait for comm's get to finish writing this tile to its slot. + // Producer (comm) signals ready += 2·NC per tile-write at the slot, + // so after ticket K the wait threshold is (K+1)·2·NC. + __device__ void acquire_src() { + compute_slot(); + int target = (cur_ticket + 1) * NumProducersSrc; + if (is_leader) { + // Relaxed spin (ld.relaxed.gpu — cheapest read, no RMW, no per-iter + // fence); the __threadfence below establishes the acquire ordering + // once for all lanes on exit. atomic_ref ctor is zero-cost/hoisted. + cuda::atomic_ref r(src_ready[cur_slot]); + while (r.load(cuda::memory_order_relaxed) < target) {} + } + __syncwarp(); + // Device-scope: observing NIC-written src. The NIC's inbound-write + // visibility is not fixable by a fence (needs GDRCopy), so system scope + // here only costs bandwidth. Device fence covers the GPU-internal view. + __threadfence(); + } + + // Signal that this CTA has consumed the tile at cur_slot. Per + // tile (column-of-runtime-NS CTAs), src_consumed advances by runtime NS. + __device__ void release_src(int lane) { + __threadfence(); + if (lane == 0) { + atomicAdd(&src_consumed[cur_slot], 1); + } + } + + // Wait for the previous put from cur_slot to finish (so we can + // reuse the slot for a fresh Y write). ticket==0 → no prior put. + __device__ void acquire_dst() { + if (cur_ticket >= 1) { + int target = cur_ticket * NumConsumersDst; + if (is_leader) { + cuda::atomic_ref r(dst_consumed[cur_slot]); + while (r.load(cuda::memory_order_relaxed) < target) {} + } + __syncwarp(); + __threadfence(); + } + } + + // Signal that Y is ready in cur_slot for the comm put. Per tile, + // dst_ready advances by runtime NS. + __device__ void release_dst(int lane) { + // Same-host tiles store Y directly into the peer's symmetric + // local_output, so the put warp has no staged data to read. Cross-host + // tiles use dst_staging + putmem and require a system-scope fence before + // the put warp observes dst_ready. + if (!cur_is_local) __threadfence_system(); + if (lane == 0) { + atomicAdd(&dst_ready[cur_slot], 1); + } + } +}; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom.cuh b/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom.cuh index 750ca602..ec9d0323 100644 --- a/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom.cuh @@ -1,35 +1,11 @@ #pragma once -#include +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -namespace liger { - -using namespace cute; - -template -struct TmaCopyAtomForCompute { - using Load = SM90_TMA_LOAD; - using Store = SM90_TMA_STORE; - using ReduceAdd = SM90_TMA_REDUCE_ADD; -}; - -template <> -struct TmaCopyAtomForCompute<100> { - // CUTLASS SM100 1-SM UMMA kernels still use the SM90 global TMA atom - // names. The SM100_TMA_2SM_* atoms are for 2-SM cluster TMA and require a - // different TiledCopy/thread layout. - using Load = SM90_TMA_LOAD; - using Store = SM90_TMA_STORE; - using ReduceAdd = SM90_TMA_REDUCE_ADD; -}; - -template -using TmaLoadAtomForCompute = typename TmaCopyAtomForCompute::Load; - -template -using TmaStoreAtomForCompute = typename TmaCopyAtomForCompute::Store; - -template -using TmaReduceAddAtomForCompute = typename TmaCopyAtomForCompute::ReduceAdd; - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "tma_copy_atom_sm90.cuh" +#else +#include "tma_copy_atom_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm100.cuh new file mode 100644 index 00000000..750ca602 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm100.cuh @@ -0,0 +1,35 @@ +#pragma once + +#include + +namespace liger { + +using namespace cute; + +template +struct TmaCopyAtomForCompute { + using Load = SM90_TMA_LOAD; + using Store = SM90_TMA_STORE; + using ReduceAdd = SM90_TMA_REDUCE_ADD; +}; + +template <> +struct TmaCopyAtomForCompute<100> { + // CUTLASS SM100 1-SM UMMA kernels still use the SM90 global TMA atom + // names. The SM100_TMA_2SM_* atoms are for 2-SM cluster TMA and require a + // different TiledCopy/thread layout. + using Load = SM90_TMA_LOAD; + using Store = SM90_TMA_STORE; + using ReduceAdd = SM90_TMA_REDUCE_ADD; +}; + +template +using TmaLoadAtomForCompute = typename TmaCopyAtomForCompute::Load; + +template +using TmaStoreAtomForCompute = typename TmaCopyAtomForCompute::Store; + +template +using TmaReduceAddAtomForCompute = typename TmaCopyAtomForCompute::ReduceAdd; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm90.cuh new file mode 100644 index 00000000..750ca602 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm90.cuh @@ -0,0 +1,35 @@ +#pragma once + +#include + +namespace liger { + +using namespace cute; + +template +struct TmaCopyAtomForCompute { + using Load = SM90_TMA_LOAD; + using Store = SM90_TMA_STORE; + using ReduceAdd = SM90_TMA_REDUCE_ADD; +}; + +template <> +struct TmaCopyAtomForCompute<100> { + // CUTLASS SM100 1-SM UMMA kernels still use the SM90 global TMA atom + // names. The SM100_TMA_2SM_* atoms are for 2-SM cluster TMA and require a + // different TiledCopy/thread layout. + using Load = SM90_TMA_LOAD; + using Store = SM90_TMA_STORE; + using ReduceAdd = SM90_TMA_REDUCE_ADD; +}; + +template +using TmaLoadAtomForCompute = typename TmaCopyAtomForCompute::Load; + +template +using TmaStoreAtomForCompute = typename TmaCopyAtomForCompute::Store; + +template +using TmaReduceAddAtomForCompute = typename TmaCopyAtomForCompute::ReduceAdd; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/tmem_load_op.cuh b/liger_cute_kernels/csrc/core/src/moe/tmem_load_op.cuh index 8ef7ee74..b88e849c 100644 --- a/liger_cute_kernels/csrc/core/src/moe/tmem_load_op.cuh +++ b/liger_cute_kernels/csrc/core/src/moe/tmem_load_op.cuh @@ -1,20 +1,11 @@ #pragma once -#include -#include +#ifndef LIGER_CUTE_DISPATCH_COMPUTE +#define LIGER_CUTE_DISPATCH_COMPUTE 0 +#endif -namespace liger { - -using namespace cute; - -template struct TmemLoadOpSelector; -template <> struct TmemLoadOpSelector<8> { using Op = SM100_TMEM_LOAD_32dp32b8x; }; -template <> struct TmemLoadOpSelector<16> { using Op = SM100_TMEM_LOAD_32dp32b16x; }; -template <> struct TmemLoadOpSelector<32> { using Op = SM100_TMEM_LOAD_32dp32b32x; }; -template <> struct TmemLoadOpSelector<64> { using Op = SM100_TMEM_LOAD_32dp32b64x; }; -template <> struct TmemLoadOpSelector<128> { using Op = SM100_TMEM_LOAD_32dp32b128x; }; - -template -using TmemLoadOp = typename TmemLoadOpSelector::Op; - -} // namespace liger +#if LIGER_CUTE_DISPATCH_COMPUTE == 90 +#include "tmem_load_op_sm90.cuh" +#else +#include "tmem_load_op_sm100.cuh" +#endif diff --git a/liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm100.cuh b/liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm100.cuh new file mode 100644 index 00000000..8ef7ee74 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm100.cuh @@ -0,0 +1,20 @@ +#pragma once + +#include +#include + +namespace liger { + +using namespace cute; + +template struct TmemLoadOpSelector; +template <> struct TmemLoadOpSelector<8> { using Op = SM100_TMEM_LOAD_32dp32b8x; }; +template <> struct TmemLoadOpSelector<16> { using Op = SM100_TMEM_LOAD_32dp32b16x; }; +template <> struct TmemLoadOpSelector<32> { using Op = SM100_TMEM_LOAD_32dp32b32x; }; +template <> struct TmemLoadOpSelector<64> { using Op = SM100_TMEM_LOAD_32dp32b64x; }; +template <> struct TmemLoadOpSelector<128> { using Op = SM100_TMEM_LOAD_32dp32b128x; }; + +template +using TmemLoadOp = typename TmemLoadOpSelector::Op; + +} // namespace liger diff --git a/liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm90.cuh b/liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm90.cuh new file mode 100644 index 00000000..8ef7ee74 --- /dev/null +++ b/liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm90.cuh @@ -0,0 +1,20 @@ +#pragma once + +#include +#include + +namespace liger { + +using namespace cute; + +template struct TmemLoadOpSelector; +template <> struct TmemLoadOpSelector<8> { using Op = SM100_TMEM_LOAD_32dp32b8x; }; +template <> struct TmemLoadOpSelector<16> { using Op = SM100_TMEM_LOAD_32dp32b16x; }; +template <> struct TmemLoadOpSelector<32> { using Op = SM100_TMEM_LOAD_32dp32b32x; }; +template <> struct TmemLoadOpSelector<64> { using Op = SM100_TMEM_LOAD_32dp32b64x; }; +template <> struct TmemLoadOpSelector<128> { using Op = SM100_TMEM_LOAD_32dp32b128x; }; + +template +using TmemLoadOp = typename TmemLoadOpSelector::Op; + +} // namespace liger diff --git a/liger_cute_kernels/test/test_moe_architecture_split.py b/liger_cute_kernels/test/test_moe_architecture_split.py new file mode 100644 index 00000000..a96f0807 --- /dev/null +++ b/liger_cute_kernels/test/test_moe_architecture_split.py @@ -0,0 +1,86 @@ +import re + +from pathlib import Path + +MOE_DIR = Path(__file__).resolve().parents[1] / "csrc" / "core" / "src" / "moe" + +DISJOINT_HEADERS = ( + "cta_barrier.cuh", + "math.cuh", + "mlp.cuh", + "mlp1_fused.cuh", + "mlp1_fused_act.cuh", + "mlp2_fused.cuh", + "mlp2_t.cuh", + "mlp2_t_fused.cuh", + "mlp3.cuh", + "mlp4.cuh", + "mlp5.cuh", + "mlp5_fused.cuh", + "mlp_bwd.cuh", + "mlp_comms.cuh", + "mlp_comms_bwd.cuh", + "models.cuh", + "moe.cuh", + "moe_bwd.cuh", + "moe_comm_config.cuh", + "moe_symm_config.cuh", + "silu_bwd_fused.cuh", + "tile_iterator.cuh", + "tile_iterator_bwd.cuh", + "tma_copy_atom.cuh", + "tmem_load_op.cuh", +) + + +def _architecture_name(name: str, architecture: str) -> str: + path = Path(name) + return f"{path.stem}_{architecture}{path.suffix}" + + +def test_moe_gemm_include_graph_is_architecture_local(): + split_names = set(DISJOINT_HEADERS) + split_names.update(("moe.cu", "moe_bwd.cu")) + + for architecture, opposite in (("sm90", "sm100"), ("sm100", "sm90")): + local_files = [_architecture_name(name, architecture) for name in DISJOINT_HEADERS] + local_files.extend( + ( + f"moe_{architecture}.inc", + f"moe_bwd_{architecture}.inc", + f"moe_fwd_bwd_tune_configs_{architecture}.hpp", + f"moe_fwd_bwd_tuning_configs_{architecture}.cuh", + ) + ) + for local_name in local_files: + text = (MOE_DIR / local_name).read_text() + includes = re.findall(r'#include\s+"([^"]+)"', text) + assert not any(f"_{opposite}." in include for include in includes), local_name + assert split_names.isdisjoint(includes), local_name + + +def test_moe_entrypoints_select_one_architecture(): + for name in DISJOINT_HEADERS: + text = (MOE_DIR / name).read_text() + assert f'#include "{_architecture_name(name, "sm90")}"' in text + assert f'#include "{_architecture_name(name, "sm100")}"' in text + + for name in ("moe.cu", "moe_bwd.cu"): + stem = Path(name).stem + text = (MOE_DIR / name).read_text() + assert f'#include "{stem}_sm90.inc"' in text + assert f'#include "{stem}_sm100.inc"' in text + + +def test_moe_communication_geometry_matches_mainline(): + for architecture in ("sm90", "sm100"): + text = (MOE_DIR / f"moe_comm_config_{architecture}.cuh").read_text() + for compute in (90, 100): + geometry = re.search( + rf"struct MoeCommGeometry<{compute}> \{{(?P.*?)\n\}};", + text, + re.DOTALL, + ) + assert geometry is not None + assert "static constexpr int TileM = 128;" in geometry["body"] + assert "static constexpr int NC = 4;" in geometry["body"]