linkedin/Liger-Kernel
1
0
Fork 0
mirror of https://github.com/linkedin/Liger-Kernel.git synced 2026-10-12 02:17:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

Isolate SM90 and SM100 MoE kernel paths (#1412)

Browse Source 
## 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 <[email protected]>
Co-authored-by: Copilot <[email protected]>
This commit is contained in:
Jari Kolehmainen
2026-08-27 17:44:06 +00:00
committed by GitHub
co-authored by jkolehma Copilot
parent 38c6a455ea
commit df40cf97d4
88 changed files with 38712 additions and 18914 deletions
Download Patch File Download Diff File Expand all files Collapse all files
liger_cute_kernels/csrc/core/src/moe/cta_barrier.cuh
+8 -132
View File
@@ -1,135 +1,11 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda/atomic>
#include <cutlass/arch/barrier.h>
#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<kNumThreads, kBarrierId> 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 <int NumThreads, int BarrierId>
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<int, cuda::thread_scope_device> 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<int, cuda::thread_scope_device> 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
liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm100.cuh
+135
View File
@@ -0,0 +1,135 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda/atomic>
#include <cutlass/arch/barrier.h>
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<kNumThreads, kBarrierId> 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 <int NumThreads, int BarrierId>
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<int, cuda::thread_scope_device> 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<int, cuda::thread_scope_device> r(*counter);
r.fetch_add(1, cuda::memory_order_release);
while (r.load(cuda::memory_order_acquire) < target) {}
}
__syncthreads();
}
};
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/cta_barrier_sm90.cuh
+135
View File
@@ -0,0 +1,135 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda/atomic>
#include <cutlass/arch/barrier.h>
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<kNumThreads, kBarrierId> 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 <int NumThreads, int BarrierId>
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<int, cuda::thread_scope_device> 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<int, cuda::thread_scope_device> r(*counter);
r.fetch_add(1, cuda::memory_order_release);
while (r.load(cuda::memory_order_acquire) < target) {}
}
__syncthreads();
}
};
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/math.cuh
+8 -152
View File
@@ -1,155 +1,11 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cutlass/numeric_types.h>
#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 <int N, typename Tensor>
__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 <typename Element>
__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<cutlass::bfloat16_t>(float x) {
__nv_bfloat16 tmp = __float2bfloat16(x);
return reinterpret_cast<const cutlass::bfloat16_t&>(tmp);
}
template <>
__device__ __forceinline__ cutlass::half_t from_float<cutlass::half_t>(float x) {
__half tmp = __float2half(x);
return reinterpret_cast<const cutlass::half_t&>(tmp);
}
// ═══════════════════════════════════════════════════════════════════
// Element → float conversion (template specialization per type)
// ═══════════════════════════════════════════════════════════════════
template <typename Element>
__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>(cutlass::bfloat16_t v) {
return __bfloat162float(reinterpret_cast<const __nv_bfloat16&>(v));
}
template <>
__device__ __forceinline__ float to_float<cutlass::half_t>(cutlass::half_t v) {
return __half2float(reinterpret_cast<const __half&>(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 <typename T>
__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<cutlass::bfloat16_t>(const cutlass::bfloat16_t* p) {
__nv_bfloat16 raw = __ldg(reinterpret_cast<const __nv_bfloat16*>(p));
return reinterpret_cast<const cutlass::bfloat16_t&>(raw);
}
template <>
__device__ __forceinline__ cutlass::half_t ldg_elem<cutlass::half_t>(const cutlass::half_t* p) {
__half raw = __ldg(reinterpret_cast<const __half*>(p));
return reinterpret_cast<const cutlass::half_t&>(raw);
}
} // namespace liger
#if LIGER_CUTE_DISPATCH_COMPUTE == 90
#include "math_sm90.cuh"
#else
#include "math_sm100.cuh"
#endif
liger_cute_kernels/csrc/core/src/moe/math_sm100.cuh
+155
View File
@@ -0,0 +1,155 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cutlass/numeric_types.h>
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 <int N, typename Tensor>
__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 <typename Element>
__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<cutlass::bfloat16_t>(float x) {
__nv_bfloat16 tmp = __float2bfloat16(x);
return reinterpret_cast<const cutlass::bfloat16_t&>(tmp);
}
template <>
__device__ __forceinline__ cutlass::half_t from_float<cutlass::half_t>(float x) {
__half tmp = __float2half(x);
return reinterpret_cast<const cutlass::half_t&>(tmp);
}
// ═══════════════════════════════════════════════════════════════════
// Element → float conversion (template specialization per type)
// ═══════════════════════════════════════════════════════════════════
template <typename Element>
__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>(cutlass::bfloat16_t v) {
return __bfloat162float(reinterpret_cast<const __nv_bfloat16&>(v));
}
template <>
__device__ __forceinline__ float to_float<cutlass::half_t>(cutlass::half_t v) {
return __half2float(reinterpret_cast<const __half&>(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 <typename T>
__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<cutlass::bfloat16_t>(const cutlass::bfloat16_t* p) {
__nv_bfloat16 raw = __ldg(reinterpret_cast<const __nv_bfloat16*>(p));
return reinterpret_cast<const cutlass::bfloat16_t&>(raw);
}
template <>
__device__ __forceinline__ cutlass::half_t ldg_elem<cutlass::half_t>(const cutlass::half_t* p) {
__half raw = __ldg(reinterpret_cast<const __half*>(p));
return reinterpret_cast<const cutlass::half_t&>(raw);
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/math_sm90.cuh
+155
View File
@@ -0,0 +1,155 @@
#pragma once
#include <cuda_runtime.h>
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cutlass/numeric_types.h>
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 <int N, typename Tensor>
__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 <typename Element>
__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<cutlass::bfloat16_t>(float x) {
__nv_bfloat16 tmp = __float2bfloat16(x);
return reinterpret_cast<const cutlass::bfloat16_t&>(tmp);
}
template <>
__device__ __forceinline__ cutlass::half_t from_float<cutlass::half_t>(float x) {
__half tmp = __float2half(x);
return reinterpret_cast<const cutlass::half_t&>(tmp);
}
// ═══════════════════════════════════════════════════════════════════
// Element → float conversion (template specialization per type)
// ═══════════════════════════════════════════════════════════════════
template <typename Element>
__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>(cutlass::bfloat16_t v) {
return __bfloat162float(reinterpret_cast<const __nv_bfloat16&>(v));
}
template <>
__device__ __forceinline__ float to_float<cutlass::half_t>(cutlass::half_t v) {
return __half2float(reinterpret_cast<const __half&>(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 <typename T>
__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<cutlass::bfloat16_t>(const cutlass::bfloat16_t* p) {
__nv_bfloat16 raw = __ldg(reinterpret_cast<const __nv_bfloat16*>(p));
return reinterpret_cast<const cutlass::bfloat16_t&>(raw);
}
template <>
__device__ __forceinline__ cutlass::half_t ldg_elem<cutlass::half_t>(const cutlass::half_t* p) {
__half raw = __ldg(reinterpret_cast<const __half*>(p));
return reinterpret_cast<const cutlass::half_t&>(raw);
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/mlp.cuh
+8 -414
View File
@@ -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 <typename Traits1, typename Traits2, int Compute = 90>
struct MlpFusedSmem {
union {
Mlp1FusedSmem<Traits1> mlp1;
Mlp2FusedSmem<Traits2> mlp2;
};
// Pipeline barrier storage — persistent across phases.
// Both phases use a single fused pipe (one mbarrier per stage).
typename Mlp1MainloopPipelineFor<Traits1, Compute>::SharedStorage p1_pipe;
typename Mlp2MainloopPipelineFor<Traits2, Compute>::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 <int Compute>
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 <typename Traits1, typename Traits2,
int ZBufferSlots = 2,
// SubTiles = CommTileM/GemmTileM. Each communication tile (one
// acquire_src/release_src/acquire_dst/release_dst) is consumed as
// SubTiles GemmTileM-row regions: mlp1 produces all Z regions, then
// mlp2 drains them. The communication tile remains intact; only the
// TMA + wgmma granularity is finer. 1 → unchanged.
int SubTiles = 1,
int Compute = 90,
typename TileIter,
typename TmaLoadX, typename TmaLoadW1, typename TmaStoreZ,
typename TmaLoadZ, typename TmaLoadW2, typename TmaStoreY>
__device__ __forceinline__ void mlp_fused_fwd(
MlpFusedSmem<Traits1, Traits2, Compute>& 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<Compute> & 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<Traits1, Compute>;
using PipeState1 = typename Traits1::PipelineState;
using Pipeline2 = Mlp2MainloopPipelineFor<Traits2, Compute>;
using PipeState2 = typename Traits2::PipelineState;
int warp_id = threadIdx.x / Traits1::WarpSize;
// MLP-warp NamedBarrier arrival count — must equal MlpFwdCtaBarrierT<Compute>'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<Traits1>(smem.p1_pipe);
else
return mlp1_make_pipe<Traits1>(smem.p1_pipe);
}();
auto p2_pipe = [&]() {
if constexpr (Compute == 100)
return mlp2_make_pipe_umma<Traits2>(smem.p2_pipe);
else
return mlp2_make_pipe<Traits2>(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<Pipeline1>() : PipeState1{};
PipeState1 p1_cons_state;
PipeState2 p2_state = (warp_id == 0) ?
cutlass::make_producer_start_state<Pipeline2>() : 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<Traits1>(
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<Traits1, true>(
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<Traits1, 100>(
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<Traits1, 90>(
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<Traits2>(
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<Traits2, true>(
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<Traits2, 100>(
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<Traits2, 90>(
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
liger_cute_kernels/csrc/core/src/moe/mlp1_fused.cuh
+7 -930
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutX_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>>{}));
using SmemLayoutX = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}));
// 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<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using SmemLayoutStore_1 = SmemLayoutStoreSlot;
using SmemLayoutActStoreSlot = Layout<
Shape<Int<ActStoreTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
static constexpr int TmaTransBytesX =
static_cast<int>(size(SmemLayoutX_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp1MainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Expert weight pointers (device-side)
// ═══════════════════════════════════════════════════════════════════
template <typename Element>
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 <typename Traits>
struct Mlp1FusedSmem {
using Element = typename Traits::Element;
static constexpr int smem_X_size = cosize_v<typename Traits::SmemLayoutX>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, bool Expert3D = false, int MSubTiles = 1,
typename Pipeline,
typename TmaLoadX, typename TmaLoadW>
__device__ __forceinline__ void mlp1_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(hidden_dim)));
auto mB = [&]() {
if constexpr (Expert3D) {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(hidden_dim)));
}
}();
auto mC = [&]() {
if constexpr (Expert3D) {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mB,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(expert_or_n_offset + n, _));
}
}();
auto gC = [&]() {
if constexpr (Expert3D) {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreZ>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto acc_C = partition_fragment_C(tiled_mma,
Shape<Int<Traits::TileM>, Int<Traits::TileN>>{});
// Identity coords for per-thread acc → store_buf mapping.
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_z_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreZ>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::K>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_z_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
// 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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<ElementAccumulator>(shape(tTR_sD)).
auto tTR_rU = make_tensor<float>(shape(tTR_cChunk)); // f32 regs (U, reused)
auto tTR_rV = make_tensor<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged;
// pass mlp1_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, int MSubTiles = 1,
typename Pipeline, typename TmaStoreZ>
__device__ __forceinline__ void mlp1_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<Compute>::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
liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act.cuh
+7 -699
View File
@@ -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 <typename Traits>
struct Mlp1FusedActSmem {
using Element = typename Traits::Element;
static constexpr int smem_X_size = cosize_v<typename Traits::SmemLayoutX>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size =
cosize_v<typename Traits::SmemLayoutActStoreSlot>;
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 <typename Traits, bool Expert3D = false, typename Pipeline,
typename TmaLoadX, typename TmaLoadW>
__device__ __forceinline__ void mlp1_fused_act_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(hidden_dim)));
auto mB = [&]() {
if constexpr (Expert3D) {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(hidden_dim)));
}
}();
auto mC = [&]() {
if constexpr (Expert3D) {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(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<Traits::TileM>{}, Int<Traits::TileK>{}),
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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mB,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(expert_or_n_offset + n, _));
}
}();
auto gC = [&]() {
if constexpr (Expert3D) {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, typename Pipeline, typename TmaStore>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto acc_C = partition_fragment_C(tiled_mma,
Shape<Int<Traits::TileM>, Int<Traits::TileN>>{});
// Identity coords for per-thread acc → store_buf mapping.
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_z_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(v * sil_d);
sStoreV(store_m, chunk_n) =
static_cast<Element>(vprime);
sStoreZ(store_m, chunk_n) =
static_cast<Element>(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<Traits::ActStoreTileM>{},
Int<Traits::EpiChunkN>{});
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 <typename Traits, typename Pipeline, typename TmaStore>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::K>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_z_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
// 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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<float>(shape(tTR_cChunk)); // f32 regs (U)
auto tTR_rV = make_tensor<float>(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<Element>(uprime);
Element vprime_e = static_cast<Element>(vprime);
Element z_e = static_cast<Element>(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<StoreTileM>{}, Int<EpiChunkN>{}),
make_coord(row_tile, col_tile));
auto gV = local_tile(mV,
make_tile(Int<StoreTileM>{}, Int<EpiChunkN>{}),
make_coord(row_tile, col_tile));
auto gZ = local_tile(mZ,
make_tile(Int<StoreTileM>{}, Int<EpiChunkN>{}),
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 <Traits, 100> for the Blackwell path.
template <typename Traits, int Compute = 90, typename Pipeline, typename TmaStore>
__device__ __forceinline__ void mlp1_fused_act_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<Compute>::template run<Traits>(
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
liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm100.cuh
+703
View File
@@ -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 <typename Traits>
struct Mlp1FusedActSmem {
using Element = typename Traits::Element;
static constexpr int smem_X_size = cosize_v<typename Traits::SmemLayoutX>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size =
cosize_v<typename Traits::SmemLayoutActStoreSlot>;
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 <typename Traits, bool Expert3D = false, typename Pipeline,
typename TmaLoadX, typename TmaLoadW>
__device__ __forceinline__ void mlp1_fused_act_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(hidden_dim)));
auto mB = [&]() {
if constexpr (Expert3D) {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(hidden_dim)));
}
}();
auto mC = [&]() {
if constexpr (Expert3D) {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(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<Traits::TileM>{}, Int<Traits::TileK>{}),
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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mB,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(expert_or_n_offset + n, _));
}
}();
auto gC = [&]() {
if constexpr (Expert3D) {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, typename Pipeline, typename TmaStore>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto acc_C = partition_fragment_C(tiled_mma,
Shape<Int<Traits::TileM>, Int<Traits::TileN>>{});
// Identity coords for per-thread acc → store_buf mapping.
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_z_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(v * sil_d);
sStoreV(store_m, chunk_n) =
static_cast<Element>(vprime);
sStoreZ(store_m, chunk_n) =
static_cast<Element>(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<Traits::ActStoreTileM>{},
Int<Traits::EpiChunkN>{});
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 <typename Traits, typename Pipeline, typename TmaStore>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::K>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_z_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
// 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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<float>(shape(tTR_cChunk)); // f32 regs (U)
auto tTR_rV = make_tensor<float>(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<Element>(uprime);
Element vprime_e = static_cast<Element>(vprime);
Element z_e = static_cast<Element>(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<StoreTileM>{}, Int<EpiChunkN>{}),
make_coord(row_tile, col_tile));
auto gV = local_tile(mV,
make_tile(Int<StoreTileM>{}, Int<EpiChunkN>{}),
make_coord(row_tile, col_tile));
auto gZ = local_tile(mZ,
make_tile(Int<StoreTileM>{}, Int<EpiChunkN>{}),
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 <Traits, 100> for the Blackwell path.
template <typename Traits, int Compute = 90, typename Pipeline, typename TmaStore>
__device__ __forceinline__ void mlp1_fused_act_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<Compute>::template run<Traits>(
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
liger_cute_kernels/csrc/core/src/moe/mlp1_fused_act_sm90.cuh
+703
View File
@@ -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 <typename Traits>
struct Mlp1FusedActSmem {
using Element = typename Traits::Element;
static constexpr int smem_X_size = cosize_v<typename Traits::SmemLayoutX>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size =
cosize_v<typename Traits::SmemLayoutActStoreSlot>;
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 <typename Traits, bool Expert3D = false, typename Pipeline,
typename TmaLoadX, typename TmaLoadW>
__device__ __forceinline__ void mlp1_fused_act_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(hidden_dim)));
auto mB = [&]() {
if constexpr (Expert3D) {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(hidden_dim)));
}
}();
auto mC = [&]() {
if constexpr (Expert3D) {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(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<Traits::TileM>{}, Int<Traits::TileK>{}),
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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mB,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(expert_or_n_offset + n, _));
}
}();
auto gC = [&]() {
if constexpr (Expert3D) {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, typename Pipeline, typename TmaStore>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto acc_C = partition_fragment_C(tiled_mma,
Shape<Int<Traits::TileM>, Int<Traits::TileN>>{});
// Identity coords for per-thread acc → store_buf mapping.
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_z_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(v * sil_d);
sStoreV(store_m, chunk_n) =
static_cast<Element>(vprime);
sStoreZ(store_m, chunk_n) =
static_cast<Element>(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<Traits::ActStoreTileM>{},
Int<Traits::EpiChunkN>{});
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 <typename Traits, typename Pipeline, typename TmaStore>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::K>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_z_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
// 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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<float>(shape(tTR_cChunk)); // f32 regs (U)
auto tTR_rV = make_tensor<float>(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<Element>(uprime);
Element vprime_e = static_cast<Element>(vprime);
Element z_e = static_cast<Element>(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<StoreTileM>{}, Int<EpiChunkN>{}),
make_coord(row_tile, col_tile));
auto gV = local_tile(mV,
make_tile(Int<StoreTileM>{}, Int<EpiChunkN>{}),
make_coord(row_tile, col_tile));
auto gZ = local_tile(mZ,
make_tile(Int<StoreTileM>{}, Int<EpiChunkN>{}),
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 <Traits, 100> for the Blackwell path.
template <typename Traits, int Compute = 90, typename Pipeline, typename TmaStore>
__device__ __forceinline__ void mlp1_fused_act_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedActSmem<Traits>& 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<Compute>::template run<Traits>(
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
liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm100.cuh
+934
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutX_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>>{}));
using SmemLayoutX = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}));
// 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<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using SmemLayoutStore_1 = SmemLayoutStoreSlot;
using SmemLayoutActStoreSlot = Layout<
Shape<Int<ActStoreTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
static constexpr int TmaTransBytesX =
static_cast<int>(size(SmemLayoutX_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp1MainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Expert weight pointers (device-side)
// ═══════════════════════════════════════════════════════════════════
template <typename Element>
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 <typename Traits>
struct Mlp1FusedSmem {
using Element = typename Traits::Element;
static constexpr int smem_X_size = cosize_v<typename Traits::SmemLayoutX>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, bool Expert3D = false, int MSubTiles = 1,
typename Pipeline,
typename TmaLoadX, typename TmaLoadW>
__device__ __forceinline__ void mlp1_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(hidden_dim)));
auto mB = [&]() {
if constexpr (Expert3D) {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(hidden_dim)));
}
}();
auto mC = [&]() {
if constexpr (Expert3D) {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mB,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(expert_or_n_offset + n, _));
}
}();
auto gC = [&]() {
if constexpr (Expert3D) {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreZ>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto acc_C = partition_fragment_C(tiled_mma,
Shape<Int<Traits::TileM>, Int<Traits::TileN>>{});
// Identity coords for per-thread acc → store_buf mapping.
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_z_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreZ>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::K>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_z_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
// 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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<ElementAccumulator>(shape(tTR_sD)).
auto tTR_rU = make_tensor<float>(shape(tTR_cChunk)); // f32 regs (U, reused)
auto tTR_rV = make_tensor<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged;
// pass mlp1_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, int MSubTiles = 1,
typename Pipeline, typename TmaStoreZ>
__device__ __forceinline__ void mlp1_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<Compute>::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
liger_cute_kernels/csrc/core/src/moe/mlp1_fused_sm90.cuh
+934
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutX_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>>{}));
using SmemLayoutX = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}));
// 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<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using SmemLayoutStore_1 = SmemLayoutStoreSlot;
using SmemLayoutActStoreSlot = Layout<
Shape<Int<ActStoreTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
static constexpr int TmaTransBytesX =
static_cast<int>(size(SmemLayoutX_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp1MainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Expert weight pointers (device-side)
// ═══════════════════════════════════════════════════════════════════
template <typename Element>
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 <typename Traits>
struct Mlp1FusedSmem {
using Element = typename Traits::Element;
static constexpr int smem_X_size = cosize_v<typename Traits::SmemLayoutX>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, bool Expert3D = false, int MSubTiles = 1,
typename Pipeline,
typename TmaLoadX, typename TmaLoadW>
__device__ __forceinline__ void mlp1_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(hidden_dim)));
auto mB = [&]() {
if constexpr (Expert3D) {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(hidden_dim)));
}
}();
auto mC = [&]() {
if constexpr (Expert3D) {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mB,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(expert_or_n_offset + n, _));
}
}();
auto gC = [&]() {
if constexpr (Expert3D) {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreZ>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto acc_C = partition_fragment_C(tiled_mma,
Shape<Int<Traits::TileM>, Int<Traits::TileN>>{});
// Identity coords for per-thread acc → store_buf mapping.
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_z_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreZ>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::K>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_z_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
// 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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<ElementAccumulator>(shape(tTR_sD)).
auto tTR_rU = make_tensor<float>(shape(tTR_cChunk)); // f32 regs (U, reused)
auto tTR_rV = make_tensor<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged;
// pass mlp1_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, int MSubTiles = 1,
typename Pipeline, typename TmaStoreZ>
__device__ __forceinline__ void mlp1_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp1FusedSmem<Traits>& 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<Compute>::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
liger_cute_kernels/csrc/core/src/moe/mlp2_fused.cuh
+7 -826
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>>{}));
using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}));
// Per-WG store buffer slot: (AtomTileM, EpiChunkN).
using SmemLayoutStoreSlot = Layout<
Shape<Int<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using SmemLayoutStore_1 = SmemLayoutStoreSlot;
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
static constexpr int TmaTransBytesZ =
static_cast<int>(size(SmemLayoutZ_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp2MainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Fused MLP2 shared memory — no pipeline storage (managed externally)
// ═══════════════════════════════════════════════════════════════════
template <typename Traits>
struct Mlp2FusedSmem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, bool Expert3D = false, int MSubTiles = 1,
typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp2_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(intermediate_dim)));
auto mA = [&]() {
if constexpr (Expert3D) {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mA,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_y_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::K>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_y_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
// 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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<ElementAccumulator>(shape(tTR_sD)).
auto tTR_rAcc = make_tensor<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged;
// pass mlp2_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, int MSubTiles = 1,
typename Pipeline, typename TmaStoreY>
__device__ __forceinline__ void mlp2_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<Compute>::template run<Traits, MSubTiles>(
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
liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm100.cuh
+830
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>>{}));
using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}));
// Per-WG store buffer slot: (AtomTileM, EpiChunkN).
using SmemLayoutStoreSlot = Layout<
Shape<Int<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using SmemLayoutStore_1 = SmemLayoutStoreSlot;
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
static constexpr int TmaTransBytesZ =
static_cast<int>(size(SmemLayoutZ_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp2MainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Fused MLP2 shared memory — no pipeline storage (managed externally)
// ═══════════════════════════════════════════════════════════════════
template <typename Traits>
struct Mlp2FusedSmem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, bool Expert3D = false, int MSubTiles = 1,
typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp2_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(intermediate_dim)));
auto mA = [&]() {
if constexpr (Expert3D) {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mA,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_y_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::K>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_y_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
// 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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<ElementAccumulator>(shape(tTR_sD)).
auto tTR_rAcc = make_tensor<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged;
// pass mlp2_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, int MSubTiles = 1,
typename Pipeline, typename TmaStoreY>
__device__ __forceinline__ void mlp2_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<Compute>::template run<Traits, MSubTiles>(
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
liger_cute_kernels/csrc/core/src/moe/mlp2_fused_sm90.cuh
+830
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>>{}));
using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtom{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}));
// Per-WG store buffer slot: (AtomTileM, EpiChunkN).
using SmemLayoutStoreSlot = Layout<
Shape<Int<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using SmemLayoutStore_1 = SmemLayoutStoreSlot;
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
static constexpr int TmaTransBytesZ =
static_cast<int>(size(SmemLayoutZ_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp2MainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Fused MLP2 shared memory — no pipeline storage (managed externally)
// ═══════════════════════════════════════════════════════════════════
template <typename Traits>
struct Mlp2FusedSmem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, bool Expert3D = false, int MSubTiles = 1,
typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp2_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(intermediate_dim)));
auto mA = [&]() {
if constexpr (Expert3D) {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(total_n_rows),
static_cast<int64_t>(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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_n_offset));
} else {
return local_tile(mA,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_y_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, int MSubTiles, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::K>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_y_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
// 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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<ElementAccumulator>(shape(tTR_sD)).
auto tTR_rAcc = make_tensor<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged;
// pass mlp2_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, int MSubTiles = 1,
typename Pipeline, typename TmaStoreY>
__device__ __forceinline__ void mlp2_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2FusedSmem<Traits>& 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<Compute>::template run<Traits, MSubTiles>(
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
liger_cute_kernels/csrc/core/src/moe/mlp2_t.cuh
+9 -538
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom<Element>;
using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>>{}));
// 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<TileN>, Int<TileK>>{},
Step<_2, _1>{}));
using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{},
Step<_2, _1, _3>{}));
// Per-WG store buffer slot: (AtomTileM, EpiChunkN).
using SmemLayoutStoreSlot = Layout<
Shape<Int<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using SmemLayoutStore_1 = SmemLayoutStoreSlot;
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
// 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 <K,K>. The B
// operand descriptor is produced by make_umma_desc<Major::MN> 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<Major::MN> accepts (mma_traits_sm100.hpp).
using TiledMmaUmma = decltype(make_tiled_mma(
SM100_MMA_F16BF16_SS<Element, Element, ElementAccum,
TileM, TileN,
UMMA::Major::K, UMMA::Major::MN>{}));
static constexpr int TmaTransBytesZ =
static_cast<int>(size(SmemLayoutZ_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp2TMainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Shared Memory
// ═══════════════════════════════════════════════════════════════════
template <typename Traits>
struct Mlp2TSmem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits, typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp2_t_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(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<int64_t>(intermediate_dim),
static_cast<int64_t>(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<Traits::TileM>{}, Int<Traits::TileK>{}),
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<Traits::TileN>{}, Int<Traits::TileK>{}),
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 <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreY>
__device__ __forceinline__ void mlp2_t_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, typename TmaLoadZ, typename TmaLoadW, typename TmaStoreY>
__device__ __forceinline__ void mlp2_t_fwd(
Mlp2TSmem<Traits>& 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<Pipeline>() : PipeState{};
PipeState cons_state;
if (is_producer) {
mlp2_t_producer<Traits>(
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<Traits>(
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
liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused.cuh
+7 -675
View File
@@ -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 <typename Traits>
struct Mlp2TFusedSmem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, bool Expert3D = false, bool ThrottleTma = false,
typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp2_t_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(hidden_dim)));
auto mA = [&]() {
if constexpr (Expert3D) {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(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<Traits::TileM>{}, Int<Traits::TileK>{}),
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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_k_offset));
} else {
return local_tile(mA,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Pipeline>();
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 <int Compute>
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 <typename Traits, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_y_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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<K,MN>) — the ONLY difference vs the mlp2_fused
// Compute=100 consumer (SS<K,K>). The epilogue is byte-for-byte identical.
// ───────────────────────────────────────────────────────────────────
template <>
struct Mlp2TFusedConsumerImpl<100> {
template <typename Traits, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_y_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged; pass
// mlp2_t_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreY>
__device__ __forceinline__ void mlp2_t_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<Compute>::template run<Traits>(
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
liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm100.cuh
+679
View File
@@ -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 <typename Traits>
struct Mlp2TFusedSmem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, bool Expert3D = false, bool ThrottleTma = false,
typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp2_t_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(hidden_dim)));
auto mA = [&]() {
if constexpr (Expert3D) {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(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<Traits::TileM>{}, Int<Traits::TileK>{}),
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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_k_offset));
} else {
return local_tile(mA,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Pipeline>();
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 <int Compute>
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 <typename Traits, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_y_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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<K,MN>) — the ONLY difference vs the mlp2_fused
// Compute=100 consumer (SS<K,K>). The epilogue is byte-for-byte identical.
// ───────────────────────────────────────────────────────────────────
template <>
struct Mlp2TFusedConsumerImpl<100> {
template <typename Traits, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_y_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged; pass
// mlp2_t_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreY>
__device__ __forceinline__ void mlp2_t_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<Compute>::template run<Traits>(
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
liger_cute_kernels/csrc/core/src/moe/mlp2_t_fused_sm90.cuh
+679
View File
@@ -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 <typename Traits>
struct Mlp2TFusedSmem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, bool Expert3D = false, bool ThrottleTma = false,
typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp2_t_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(hidden_dim)));
auto mA = [&]() {
if constexpr (Expert3D) {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_a.get_tma_tensor(make_shape(
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(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<Traits::TileM>{}, Int<Traits::TileK>{}),
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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(n, _, expert_or_k_offset));
} else {
return local_tile(mA,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Pipeline>();
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 <int Compute>
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 <typename Traits, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_y_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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<K,MN>) — the ONLY difference vs the mlp2_fused
// Compute=100 consumer (SS<K,K>). The epilogue is byte-for-byte identical.
// ───────────────────────────────────────────────────────────────────
template <>
struct Mlp2TFusedConsumerImpl<100> {
template <typename Traits, typename Pipeline, typename TmaStoreY>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_y_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
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<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged; pass
// mlp2_t_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreY>
__device__ __forceinline__ void mlp2_t_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TFusedSmem<Traits>& 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<Compute>::template run<Traits>(
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
liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm100.cuh
+540
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom<Element>;
using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>>{}));
// 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<TileN>, Int<TileK>>{},
Step<_2, _1>{}));
using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{},
Step<_2, _1, _3>{}));
// Per-WG store buffer slot: (AtomTileM, EpiChunkN).
using SmemLayoutStoreSlot = Layout<
Shape<Int<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using SmemLayoutStore_1 = SmemLayoutStoreSlot;
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
// 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 <K,K>. The B
// operand descriptor is produced by make_umma_desc<Major::MN> 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<Major::MN> accepts (mma_traits_sm100.hpp).
using TiledMmaUmma = decltype(make_tiled_mma(
SM100_MMA_F16BF16_SS<Element, Element, ElementAccum,
TileM, TileN,
UMMA::Major::K, UMMA::Major::MN>{}));
static constexpr int TmaTransBytesZ =
static_cast<int>(size(SmemLayoutZ_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp2TMainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Shared Memory
// ═══════════════════════════════════════════════════════════════════
template <typename Traits>
struct Mlp2TSmem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits, typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp2_t_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(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<int64_t>(intermediate_dim),
static_cast<int64_t>(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<Traits::TileM>{}, Int<Traits::TileK>{}),
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<Traits::TileN>{}, Int<Traits::TileK>{}),
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 <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreY>
__device__ __forceinline__ void mlp2_t_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, typename TmaLoadZ, typename TmaLoadW, typename TmaStoreY>
__device__ __forceinline__ void mlp2_t_fwd(
Mlp2TSmem<Traits>& 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<Pipeline>() : PipeState{};
PipeState cons_state;
if (is_producer) {
mlp2_t_producer<Traits>(
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<Traits>(
pipe, cons_state,
smem, tma_store_y,
intermediate_dim,
num_m_tiles, num_n_tiles, num_k_tiles);
}
__syncthreads();
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/mlp2_t_sm90.cuh
+540
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom<Element>;
using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>>{}));
// 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<TileN>, Int<TileK>>{},
Step<_2, _1>{}));
using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{},
Step<_2, _1, _3>{}));
// Per-WG store buffer slot: (AtomTileM, EpiChunkN).
using SmemLayoutStoreSlot = Layout<
Shape<Int<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _1>>;
using SmemLayoutStore_1 = SmemLayoutStoreSlot;
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
// 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 <K,K>. The B
// operand descriptor is produced by make_umma_desc<Major::MN> 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<Major::MN> accepts (mma_traits_sm100.hpp).
using TiledMmaUmma = decltype(make_tiled_mma(
SM100_MMA_F16BF16_SS<Element, Element, ElementAccum,
TileM, TileN,
UMMA::Major::K, UMMA::Major::MN>{}));
static constexpr int TmaTransBytesZ =
static_cast<int>(size(SmemLayoutZ_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp2TMainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Shared Memory
// ═══════════════════════════════════════════════════════════════════
template <typename Traits>
struct Mlp2TSmem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStoreSlot>;
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 <typename Traits, typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp2_t_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TSmem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(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<int64_t>(intermediate_dim),
static_cast<int64_t>(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<Traits::TileM>{}, Int<Traits::TileK>{}),
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<Traits::TileN>{}, Int<Traits::TileK>{}),
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 <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreY>
__device__ __forceinline__ void mlp2_t_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp2TSmem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, typename TmaLoadZ, typename TmaLoadW, typename TmaStoreY>
__device__ __forceinline__ void mlp2_t_fwd(
Mlp2TSmem<Traits>& 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<Pipeline>() : PipeState{};
PipeState cons_state;
if (is_producer) {
mlp2_t_producer<Traits>(
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<Traits>(
pipe, cons_state,
smem, tma_store_y,
intermediate_dim,
num_m_tiles, num_n_tiles, num_k_tiles);
}
__syncthreads();
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/mlp3.cuh
+6 -1779
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp3_sm100.cuh
+1784
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp3_sm90.cuh
+1784
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp4.cuh
+6 -1834
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp4_sm100.cuh
+1839
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp4_sm90.cuh
+1839
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp5.cuh
+9 -597
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom<Element>;
// One-stage layouts (used for TMA descriptors).
using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtomW{},
Shape<Int<TileN>, Int<TileK>>{}, Step<_2, _1>{}));
// Multi-stage smem layouts.
using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}, Step<_2, _1, _3>{}));
// Per-WG store buffer: (AtomTileM, EpiChunkN). One slot per WG.
using SmemLayoutStore = Layout<
Shape<Int<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _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<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
static constexpr int TmaTransBytesZ =
static_cast<int>(size(SmemLayoutZ_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp5MainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Shared memory layout
// ═══════════════════════════════════════════════════════════════════
template <typename Traits>
struct Mlp5Smem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStore>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp5_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _));
auto gC_all = local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
make_coord(m, _));
auto tZgDU = cta_tma_du.partition_S(gDU);
auto gDV = local_tile(mDV,
make_tile(Int<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreDX>
__device__ __forceinline__ void mlp5_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, typename TmaLoadZ, typename TmaLoadW, typename TmaStoreDX>
__device__ __forceinline__ void mlp5_fwd(
Mlp5Smem<Traits>& 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<Traits>(smem.pipe_storage);
__syncthreads();
PipeState state = is_producer ?
cutlass::make_producer_start_state<Pipeline>() : PipeState{};
PipeState cons_state;
if (is_producer) {
mlp5_producer<Traits>(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<Traits>(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
liger_cute_kernels/csrc/core/src/moe/mlp5_fused.cuh
+7 -625
View File
@@ -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 <typename Traits, bool Expert3D = false, typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp5_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(intermediate_dim)));
auto mDV = tma_load_dv.get_tma_tensor(make_shape(
static_cast<int64_t>(num_tokens),
static_cast<int64_t>(intermediate_dim)));
auto mB = [&]() {
if constexpr (Expert3D) {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(total_k_cols)));
}
}();
auto mC = [&]() {
if constexpr (Expert3D) {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _, expert_or_k_offset));
} else {
return local_tile(mB,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _));
}
}();
auto gC_all = [&]() {
if constexpr (Expert3D) {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _, expert_or_k_offset));
} else {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
make_coord(m, _));
auto tZgDU = cta_tma_du.partition_S(gDU);
auto gDV = local_tile(mDV,
make_tile(Int<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, typename Pipeline, typename TmaStoreDX>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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<K), C in phase 2 (k≥K) — into the same W slot. The
// consumer runs one continuous k-loop over all 2·num_k_tiles stages into ONE
// TMEM accumulator, setting the per-instruction accumulate bit
// (UMMA::ScaleOut) to Zero on the very first MMA only (writes/clears the acc)
// and One on every subsequent MMA — INCLUDING the phase-1→phase-2 boundary.
// The bit is NOT reset when phase 2 starts, so the dU·B term is preserved.
//
// Epilogue (reused from mlp2_fused's single-accumulator UMMA epilogue): the
// two consumer warpgroups split TileN; each reads its N-half from TMEM→regs in
// EpiChunkN=64 chunks, casts to bf16, and stores 64-row (AtomTileM) tiles via
// the existing reg→SMEM→TMA path (host TMA atom unchanged).
// ───────────────────────────────────────────────────────────────────
template <>
struct Mlp5FusedConsumerImpl<100> {
template <typename Traits, typename Pipeline, typename TmaStoreDX>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::MN>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N)
auto tTR_rAcc = make_tensor<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged; pass
// mlp5_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreDX>
__device__ __forceinline__ void mlp5_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Compute>::template run<Traits>(
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
liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm100.cuh
+629
View File
@@ -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 <typename Traits, bool Expert3D = false, typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp5_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(intermediate_dim)));
auto mDV = tma_load_dv.get_tma_tensor(make_shape(
static_cast<int64_t>(num_tokens),
static_cast<int64_t>(intermediate_dim)));
auto mB = [&]() {
if constexpr (Expert3D) {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(total_k_cols)));
}
}();
auto mC = [&]() {
if constexpr (Expert3D) {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _, expert_or_k_offset));
} else {
return local_tile(mB,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _));
}
}();
auto gC_all = [&]() {
if constexpr (Expert3D) {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _, expert_or_k_offset));
} else {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
make_coord(m, _));
auto tZgDU = cta_tma_du.partition_S(gDU);
auto gDV = local_tile(mDV,
make_tile(Int<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, typename Pipeline, typename TmaStoreDX>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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<K), C in phase 2 (k≥K) — into the same W slot. The
// consumer runs one continuous k-loop over all 2·num_k_tiles stages into ONE
// TMEM accumulator, setting the per-instruction accumulate bit
// (UMMA::ScaleOut) to Zero on the very first MMA only (writes/clears the acc)
// and One on every subsequent MMA — INCLUDING the phase-1→phase-2 boundary.
// The bit is NOT reset when phase 2 starts, so the dU·B term is preserved.
//
// Epilogue (reused from mlp2_fused's single-accumulator UMMA epilogue): the
// two consumer warpgroups split TileN; each reads its N-half from TMEM→regs in
// EpiChunkN=64 chunks, casts to bf16, and stores 64-row (AtomTileM) tiles via
// the existing reg→SMEM→TMA path (host TMA atom unchanged).
// ───────────────────────────────────────────────────────────────────
template <>
struct Mlp5FusedConsumerImpl<100> {
template <typename Traits, typename Pipeline, typename TmaStoreDX>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::MN>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N)
auto tTR_rAcc = make_tensor<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged; pass
// mlp5_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreDX>
__device__ __forceinline__ void mlp5_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Compute>::template run<Traits>(
pipe, state, smem, tma_store_dx, m, hidden_dim,
num_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits);
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/mlp5_fused_sm90.cuh
+629
View File
@@ -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 <typename Traits, bool Expert3D = false, typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp5_fused_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<int64_t>(num_tokens),
static_cast<int64_t>(intermediate_dim)));
auto mDV = tma_load_dv.get_tma_tensor(make_shape(
static_cast<int64_t>(num_tokens),
static_cast<int64_t>(intermediate_dim)));
auto mB = [&]() {
if constexpr (Expert3D) {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_b.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(total_k_cols)));
}
}();
auto mC = [&]() {
if constexpr (Expert3D) {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(intermediate_dim),
static_cast<int64_t>(num_experts)));
} else {
return tma_load_c.get_tma_tensor(make_shape(
static_cast<int64_t>(hidden_dim),
static_cast<int64_t>(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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _, expert_or_k_offset));
} else {
return local_tile(mB,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _));
}
}();
auto gC_all = [&]() {
if constexpr (Expert3D) {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _, expert_or_k_offset));
} else {
return local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
make_coord(m, _));
auto tZgDU = cta_tma_du.partition_S(gDU);
auto gDV = local_tile(mDV,
make_tile(Int<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <int Compute>
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 <typename Traits, typename Pipeline, typename TmaStoreDX>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<int64_t>(num_m_tiles) * Traits::TileM,
static_cast<int64_t>(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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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<K), C in phase 2 (k≥K) — into the same W slot. The
// consumer runs one continuous k-loop over all 2·num_k_tiles stages into ONE
// TMEM accumulator, setting the per-instruction accumulate bit
// (UMMA::ScaleOut) to Zero on the very first MMA only (writes/clears the acc)
// and One on every subsequent MMA — INCLUDING the phase-1→phase-2 boundary.
// The bit is NOT reset when phase 2 starts, so the dU·B term is preserved.
//
// Epilogue (reused from mlp2_fused's single-accumulator UMMA epilogue): the
// two consumer warpgroups split TileN; each reads its N-half from TMEM→regs in
// EpiChunkN=64 chunks, casts to bf16, and stores 64-row (AtomTileM) tiles via
// the existing reg→SMEM→TMA path (host TMA atom unchanged).
// ───────────────────────────────────────────────────────────────────
template <>
struct Mlp5FusedConsumerImpl<100> {
template <typename Traits, typename Pipeline, typename TmaStoreDX>
static __device__ __forceinline__ void run(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Element, Element, float, TileM, TileN,
UMMA::Major::K, UMMA::Major::MN>{});
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<TileM>{}, Int<TileN>{}));
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<cute::_1, cute::_1, cute::_1>{});
auto acc_prod_state = cutlass::make_producer_start_state<AccPipe>();
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<int64_t>(num_m_tiles) * TileM,
static_cast<int64_t>(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<TileM>{}, Int<EpiChunkN>{});
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<EpiChunkN>{}, 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<TileM>{}, Int<EpiChunkN>{}));
auto tTR_cChunk = thr_t2r.partition_D(cChunk); // (Cpy,Cpy_M,Cpy_N)
auto tTR_rAcc = make_tensor<float>(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<Element>(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<AtomTileM>{}, Int<EpiChunkN>{}),
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<Traits>(...)) are unchanged; pass
// mlp5_fused_consumer<Traits, 100>(...) to select the Blackwell path.
// ───────────────────────────────────────────────────────────────────
template <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreDX>
__device__ __forceinline__ void mlp5_fused_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Compute>::template run<Traits>(
pipe, state, smem, tma_store_dx, m, hidden_dim,
num_m_tiles, num_n_tiles, num_k_tiles, split_idx, num_splits);
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/mlp5_sm100.cuh
+599
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom<Element>;
// One-stage layouts (used for TMA descriptors).
using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtomW{},
Shape<Int<TileN>, Int<TileK>>{}, Step<_2, _1>{}));
// Multi-stage smem layouts.
using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}, Step<_2, _1, _3>{}));
// Per-WG store buffer: (AtomTileM, EpiChunkN). One slot per WG.
using SmemLayoutStore = Layout<
Shape<Int<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _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<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
static constexpr int TmaTransBytesZ =
static_cast<int>(size(SmemLayoutZ_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp5MainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Shared memory layout
// ═══════════════════════════════════════════════════════════════════
template <typename Traits>
struct Mlp5Smem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStore>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp5_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _));
auto gC_all = local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
make_coord(m, _));
auto tZgDU = cta_tma_du.partition_S(gDU);
auto gDV = local_tile(mDV,
make_tile(Int<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreDX>
__device__ __forceinline__ void mlp5_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, typename TmaLoadZ, typename TmaLoadW, typename TmaStoreDX>
__device__ __forceinline__ void mlp5_fwd(
Mlp5Smem<Traits>& 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<Traits>(smem.pipe_storage);
__syncthreads();
PipeState state = is_producer ?
cutlass::make_producer_start_state<Pipeline>() : PipeState{};
PipeState cons_state;
if (is_producer) {
mlp5_producer<Traits>(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<Traits>(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
liger_cute_kernels/csrc/core/src/moe/mlp5_sm90.cuh
+599
View File
@@ -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 <cute/tensor.hpp>
#include <cute/algorithm/gemm.hpp>
#include <cute/atom/copy_traits_sm90_tma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
// 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 <cute/arch/mma_sm100_umma.hpp>
#include <cute/atom/mma_traits_sm100.hpp>
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#include <cutlass/pipeline/sm100_pipeline.hpp> // 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<Element, AtomN>::Atom;
using WgLayout = cute::conditional_t<
kMSplit,
Layout<Shape<_2, _1, _1>>,
Layout<Shape<_1, _2, _1>>>;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, WgLayout>;
using SmemLayoutAtomZ = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutAtomW = GMMA::Layout_MN_SW128_Atom<Element>;
// One-stage layouts (used for TMA descriptors).
using SmemLayoutZ_1 = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutW_1 = decltype(tile_to_shape(SmemLayoutAtomW{},
Shape<Int<TileN>, Int<TileK>>{}, Step<_2, _1>{}));
// Multi-stage smem layouts.
using SmemLayoutZ = decltype(tile_to_shape(SmemLayoutAtomZ{},
Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutW = decltype(tile_to_shape(SmemLayoutAtomW{},
Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}, Step<_2, _1, _3>{}));
// Per-WG store buffer: (AtomTileM, EpiChunkN). One slot per WG.
using SmemLayoutStore = Layout<
Shape<Int<AtomTileM>, Int<EpiChunkN>>,
Stride<Int<EpiChunkN>, _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<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
// 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::_1, cute::_1, cute::_1>,
cute::Shape<cute::_1, cute::_1, cute::_1>>;
// 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<AccStages>;
static constexpr int TmaTransBytesZ =
static_cast<int>(size(SmemLayoutZ_1{}) * sizeof(Element));
static constexpr int TmaTransBytesW =
static_cast<int>(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 <typename Traits, int Compute>
using Mlp5MainloopPipelineFor = cute::conditional_t<
Compute == 100,
typename Traits::MainloopPipelineUmma,
typename Traits::MainloopPipeline>;
// ═══════════════════════════════════════════════════════════════════
// Shared memory layout
// ═══════════════════════════════════════════════════════════════════
template <typename Traits>
struct Mlp5Smem {
using Element = typename Traits::Element;
static constexpr int smem_Z_size = cosize_v<typename Traits::SmemLayoutZ>;
static constexpr int smem_W_size = cosize_v<typename Traits::SmemLayoutW>;
static constexpr int smem_store_size = cosize_v<typename Traits::SmemLayoutStore>;
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 <typename Traits>
__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 <typename Traits>
__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 <typename Traits, typename Pipeline,
typename TmaLoadZ, typename TmaLoadW>
__device__ __forceinline__ void mlp5_producer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Traits::TileN>{}, Int<Traits::TileK>{}),
make_coord(_, _));
auto gC_all = local_tile(mC,
make_tile(Int<Traits::TileN>{}, Int<Traits::TileK>{}),
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<Traits::TileM>{}, Int<Traits::TileK>{}),
make_coord(m, _));
auto tZgDU = cta_tma_du.partition_S(gDU);
auto gDV = local_tile(mDV,
make_tile(Int<Traits::TileM>{}, Int<Traits::TileK>{}),
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 <typename Traits, int Compute = 90, typename Pipeline, typename TmaStoreDX>
__device__ __forceinline__ void mlp5_consumer(
Pipeline& pipe,
typename Traits::PipelineState& state,
Mlp5Smem<Traits>& 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<Traits::TileM>, Int<Traits::TileN>>{});
auto cC = make_identity_tensor(make_shape(Int<Traits::TileM>{}, Int<Traits::TileN>{}));
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<K_PIPE_MMAS>();
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<Element>(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<Traits::AtomTileM>{}, Int<Traits::EpiChunkN>{}),
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 <typename Traits, typename TmaLoadZ, typename TmaLoadW, typename TmaStoreDX>
__device__ __forceinline__ void mlp5_fwd(
Mlp5Smem<Traits>& 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<Traits>(smem.pipe_storage);
__syncthreads();
PipeState state = is_producer ?
cutlass::make_producer_start_state<Pipeline>() : PipeState{};
PipeState cons_state;
if (is_producer) {
mlp5_producer<Traits>(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<Traits>(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
liger_cute_kernels/csrc/core/src/moe/mlp_bwd.cuh
+9 -1883
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm100.cuh
+1885
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp_bwd_sm90.cuh
+1885
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp_comms.cuh
+9 -1341
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd.cuh
+9 -770
View File
@@ -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 <typename Element, int BoxRows>
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 <typename Element, int TileM, int K, int NC, int BoxRows>
__device__ __forceinline__ void do_get_bwd_tma(
StagePipe<K, kNumGetWarpsBwd * NC, 1>& 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<Element, BoxRows>;
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<Element*>(bounce_warp);
uint64_t* mbar = reinterpret_cast<uint64_t*>(
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 <typename Element, int TileM, int K, int NC>
__device__ __forceinline__ void do_get_bwd(
StagePipe<K, kNumGetWarpsPerCta * NC, 1>& 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<const int4*>(
static_cast<Element*>(p2p_remote) + offset);
int4* dst_v = reinterpret_cast<int4*>(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 <typename Element, int TileM, int K, int NC>
__device__ __forceinline__ void do_put_bwd(
StagePipe<K, 1, NC>& 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<Element*>(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 <int TileM, int NC = 2>
__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<TileM> 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<TileM> 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<Element*>(bufs.x_staging);
auto* dy_staging = static_cast<Element*>(bufs.dy_staging);
auto* dx_staging = static_cast<Element*>(bufs.dx_staging);
auto* remote_x_base = static_cast<Element*>(bufs.remote_x);
auto* remote_dy_base = static_cast<Element*>(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<K, kNumGetWarpsBwd * NC, 1> x_pipe;
x_pipe.init(bufs.x_src_ready + xc, bufs.x_src_consumed + xc,
is_leader, MC, kNumGetWarpsBwd * NC, runtime_nsplit);
StagePipe<K, kNumGetWarpsBwd * NC, 1> 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<Element, TileM, K, NC, BoxRows>(
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<Element, TileM, K, NC, BoxRows>(
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<K, 1, NC> 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<Element, TileM, K, NC>(
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
liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm100.cuh
+772
View File
@@ -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 <typename Element, int BoxRows>
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 <typename Element, int TileM, int K, int NC, int BoxRows>
__device__ __forceinline__ void do_get_bwd_tma(
StagePipe<K, kNumGetWarpsBwd * NC, 1>& 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<Element, BoxRows>;
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<Element*>(bounce_warp);
uint64_t* mbar = reinterpret_cast<uint64_t*>(
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 <typename Element, int TileM, int K, int NC>
__device__ __forceinline__ void do_get_bwd(
StagePipe<K, kNumGetWarpsPerCta * NC, 1>& 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<const int4*>(
static_cast<Element*>(p2p_remote) + offset);
int4* dst_v = reinterpret_cast<int4*>(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 <typename Element, int TileM, int K, int NC>
__device__ __forceinline__ void do_put_bwd(
StagePipe<K, 1, NC>& 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<Element*>(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 <int TileM, int NC = 2>
__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<TileM> 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<TileM> 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<Element*>(bufs.x_staging);
auto* dy_staging = static_cast<Element*>(bufs.dy_staging);
auto* dx_staging = static_cast<Element*>(bufs.dx_staging);
auto* remote_x_base = static_cast<Element*>(bufs.remote_x);
auto* remote_dy_base = static_cast<Element*>(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<K, kNumGetWarpsBwd * NC, 1> x_pipe;
x_pipe.init(bufs.x_src_ready + xc, bufs.x_src_consumed + xc,
is_leader, MC, kNumGetWarpsBwd * NC, runtime_nsplit);
StagePipe<K, kNumGetWarpsBwd * NC, 1> 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<Element, TileM, K, NC, BoxRows>(
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<Element, TileM, K, NC, BoxRows>(
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<K, 1, NC> 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<Element, TileM, K, NC>(
dst_pipe, bufs, info,
dx_staging, tile_elems, xc, MC, chunk_idx, ib_seq, lane,
kNumPutWarpsBwd, put_local_idx);
}
}
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/mlp_comms_bwd_sm90.cuh
+772
View File
@@ -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 <typename Element, int BoxRows>
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 <typename Element, int TileM, int K, int NC, int BoxRows>
__device__ __forceinline__ void do_get_bwd_tma(
StagePipe<K, kNumGetWarpsBwd * NC, 1>& 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<Element, BoxRows>;
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<Element*>(bounce_warp);
uint64_t* mbar = reinterpret_cast<uint64_t*>(
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 <typename Element, int TileM, int K, int NC>
__device__ __forceinline__ void do_get_bwd(
StagePipe<K, kNumGetWarpsPerCta * NC, 1>& 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<const int4*>(
static_cast<Element*>(p2p_remote) + offset);
int4* dst_v = reinterpret_cast<int4*>(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 <typename Element, int TileM, int K, int NC>
__device__ __forceinline__ void do_put_bwd(
StagePipe<K, 1, NC>& 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<Element*>(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 <int TileM, int NC = 2>
__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<TileM> 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<TileM> 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<Element*>(bufs.x_staging);
auto* dy_staging = static_cast<Element*>(bufs.dy_staging);
auto* dx_staging = static_cast<Element*>(bufs.dx_staging);
auto* remote_x_base = static_cast<Element*>(bufs.remote_x);
auto* remote_dy_base = static_cast<Element*>(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<K, kNumGetWarpsBwd * NC, 1> x_pipe;
x_pipe.init(bufs.x_src_ready + xc, bufs.x_src_consumed + xc,
is_leader, MC, kNumGetWarpsBwd * NC, runtime_nsplit);
StagePipe<K, kNumGetWarpsBwd * NC, 1> 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<Element, TileM, K, NC, BoxRows>(
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<Element, TileM, K, NC, BoxRows>(
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<K, 1, NC> 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<Element, TileM, K, NC>(
dst_pipe, bufs, info,
dx_staging, tile_elems, xc, MC, chunk_idx, ib_seq, lane,
kNumPutWarpsBwd, put_local_idx);
}
}
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm100.cuh
+1343
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp_comms_sm90.cuh
+1343
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/mlp_sm100.cuh
+417
View File
@@ -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 <typename Traits1, typename Traits2, int Compute = 90>
struct MlpFusedSmem {
union {
Mlp1FusedSmem<Traits1> mlp1;
Mlp2FusedSmem<Traits2> mlp2;
};
// Pipeline barrier storage — persistent across phases.
// Both phases use a single fused pipe (one mbarrier per stage).
typename Mlp1MainloopPipelineFor<Traits1, Compute>::SharedStorage p1_pipe;
typename Mlp2MainloopPipelineFor<Traits2, Compute>::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 <int Compute>
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 <typename Traits1, typename Traits2,
int ZBufferSlots = 2,
// SubTiles = CommTileM/GemmTileM. Each communication tile (one
// acquire_src/release_src/acquire_dst/release_dst) is consumed as
// SubTiles GemmTileM-row regions: mlp1 produces all Z regions, then
// mlp2 drains them. The communication tile remains intact; only the
// TMA + wgmma granularity is finer. 1 → unchanged.
int SubTiles = 1,
int Compute = 90,
typename TileIter,
typename TmaLoadX, typename TmaLoadW1, typename TmaStoreZ,
typename TmaLoadZ, typename TmaLoadW2, typename TmaStoreY>
__device__ __forceinline__ void mlp_fused_fwd(
MlpFusedSmem<Traits1, Traits2, Compute>& 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<Compute> & 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<Traits1, Compute>;
using PipeState1 = typename Traits1::PipelineState;
using Pipeline2 = Mlp2MainloopPipelineFor<Traits2, Compute>;
using PipeState2 = typename Traits2::PipelineState;
int warp_id = threadIdx.x / Traits1::WarpSize;
// MLP-warp NamedBarrier arrival count — must equal MlpFwdCtaBarrierT<Compute>'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<Traits1>(smem.p1_pipe);
else
return mlp1_make_pipe<Traits1>(smem.p1_pipe);
}();
auto p2_pipe = [&]() {
if constexpr (Compute == 100)
return mlp2_make_pipe_umma<Traits2>(smem.p2_pipe);
else
return mlp2_make_pipe<Traits2>(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<Pipeline1>() : PipeState1{};
PipeState1 p1_cons_state;
PipeState2 p2_state = (warp_id == 0) ?
cutlass::make_producer_start_state<Pipeline2>() : 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<Traits1>(
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<Traits1, true>(
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<Traits1, 100>(
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<Traits1, 90>(
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<Traits2>(
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<Traits2, true>(
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<Traits2, 100>(
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<Traits2, 90>(
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
liger_cute_kernels/csrc/core/src/moe/mlp_sm90.cuh
+417
View File
@@ -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 <typename Traits1, typename Traits2, int Compute = 90>
struct MlpFusedSmem {
union {
Mlp1FusedSmem<Traits1> mlp1;
Mlp2FusedSmem<Traits2> mlp2;
};
// Pipeline barrier storage — persistent across phases.
// Both phases use a single fused pipe (one mbarrier per stage).
typename Mlp1MainloopPipelineFor<Traits1, Compute>::SharedStorage p1_pipe;
typename Mlp2MainloopPipelineFor<Traits2, Compute>::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 <int Compute>
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 <typename Traits1, typename Traits2,
int ZBufferSlots = 2,
// SubTiles = CommTileM/GemmTileM. Each communication tile (one
// acquire_src/release_src/acquire_dst/release_dst) is consumed as
// SubTiles GemmTileM-row regions: mlp1 produces all Z regions, then
// mlp2 drains them. The communication tile remains intact; only the
// TMA + wgmma granularity is finer. 1 → unchanged.
int SubTiles = 1,
int Compute = 90,
typename TileIter,
typename TmaLoadX, typename TmaLoadW1, typename TmaStoreZ,
typename TmaLoadZ, typename TmaLoadW2, typename TmaStoreY>
__device__ __forceinline__ void mlp_fused_fwd(
MlpFusedSmem<Traits1, Traits2, Compute>& 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<Compute> & 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<Traits1, Compute>;
using PipeState1 = typename Traits1::PipelineState;
using Pipeline2 = Mlp2MainloopPipelineFor<Traits2, Compute>;
using PipeState2 = typename Traits2::PipelineState;
int warp_id = threadIdx.x / Traits1::WarpSize;
// MLP-warp NamedBarrier arrival count — must equal MlpFwdCtaBarrierT<Compute>'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<Traits1>(smem.p1_pipe);
else
return mlp1_make_pipe<Traits1>(smem.p1_pipe);
}();
auto p2_pipe = [&]() {
if constexpr (Compute == 100)
return mlp2_make_pipe_umma<Traits2>(smem.p2_pipe);
else
return mlp2_make_pipe<Traits2>(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<Pipeline1>() : PipeState1{};
PipeState1 p1_cons_state;
PipeState2 p2_state = (warp_id == 0) ?
cutlass::make_producer_start_state<Pipeline2>() : 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<Traits1>(
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<Traits1, true>(
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<Traits1, 100>(
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<Traits1, 90>(
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<Traits2>(
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<Traits2, true>(
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<Traits2, 100>(
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<Traits2, 90>(
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
liger_cute_kernels/csrc/core/src/moe/models.cuh
+8 -186
View File
@@ -1,189 +1,11 @@
#pragma once
// #ifndef NDEBUG
// #define NDEBUG
// #endif
#ifndef LIGER_CUTE_DISPATCH_COMPUTE
#define LIGER_CUTE_DISPATCH_COMPUTE 0
#endif
#include <cuda_runtime.h>
#include <cute/tensor.hpp>
#include <cute/atom/mma_atom.hpp>
#include <cute/atom/mma_traits_sm90_gmma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
#include <cutlass/numeric_types.h>
namespace liger {
using namespace cute;
// ═══════════════════════════════════════════════════════════════════
// GMMA atom selection
// ═══════════════════════════════════════════════════════════════════
template <typename Element, int N>
struct GmmaSelector;
static constexpr auto GmmaK = GMMA::Major::K;
static constexpr auto GmmaMN = GMMA::Major::MN;
template <> struct GmmaSelector<bfloat16_t, 8> { using Atom = SM90_64x8x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 16> { using Atom = SM90_64x16x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 32> { using Atom = SM90_64x32x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 64> { using Atom = SM90_64x64x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 128> { using Atom = SM90_64x128x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 192> { using Atom = SM90_64x192x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 256> { using Atom = SM90_64x256x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 8> { using Atom = SM90_64x8x16_F32F16F16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 16> { using Atom = SM90_64x16x16_F32F16F16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 32> { using Atom = SM90_64x32x16_F32F16F16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 64> { using Atom = SM90_64x64x16_F32F16F16_SS<GmmaK, GmmaK>; };
// ═══════════════════════════════════════════════════════════════════
// GMMA atom selection — A=K-major, B=MN-major (SS<K, MN>)
// Used when B operand has MN-contiguous smem (e.g. mlp2_t_fused).
// ═══════════════════════════════════════════════════════════════════
template <typename Element, int N>
struct GmmaSelectorKMN;
template <> struct GmmaSelectorKMN<bfloat16_t, 8> { using Atom = SM90_64x8x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 16> { using Atom = SM90_64x16x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 32> { using Atom = SM90_64x32x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 64> { using Atom = SM90_64x64x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 128> { using Atom = SM90_64x128x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 256> { using Atom = SM90_64x256x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 8> { using Atom = SM90_64x8x16_F32F16F16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 16> { using Atom = SM90_64x16x16_F32F16F16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 32> { using Atom = SM90_64x32x16_F32F16F16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 64> { using Atom = SM90_64x64x16_F32F16F16_SS<GmmaK, GmmaMN>; };
// Backward compat alias
template <typename Element, int N>
struct GmmaSelectorMN : GmmaSelectorKMN<Element, N> {};
// ═══════════════════════════════════════════════════════════════════
// GMMA atom selection — both MN-major (SS<MN, MN>)
// Used when both operands have MN-contiguous smem (e.g. mlp3).
// ═══════════════════════════════════════════════════════════════════
template <typename Element, int N>
struct GmmaSelectorMNMN;
template <> struct GmmaSelectorMNMN<bfloat16_t, 8> { using Atom = SM90_64x8x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 16> { using Atom = SM90_64x16x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 32> { using Atom = SM90_64x32x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 64> { using Atom = SM90_64x64x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 128> { using Atom = SM90_64x128x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 256> { using Atom = SM90_64x256x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 8> { using Atom = SM90_64x8x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 16> { using Atom = SM90_64x16x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 32> { using Atom = SM90_64x32x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 64> { using Atom = SM90_64x64x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
// ═══════════════════════════════════════════════════════════════════
// 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<Element, TileN>::Atom;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, Layout<Shape<_1, _1, _1>>>;
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutA_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutB_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileN>, Int<TileK>>{}));
using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutB = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}));
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
static constexpr int TmaTransactionBytesA = static_cast<int>(size(SmemLayoutA_1{}) * sizeof(Element));
static constexpr int TmaTransactionBytesB = static_cast<int>(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 <typename Traits>
struct RouterSmem {
using Element = typename Traits::Element;
static constexpr int smem_A_size = cosize_v<typename Traits::SmemLayoutA>;
static constexpr int smem_B_size = cosize_v<typename Traits::SmemLayoutB>;
// 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
liger_cute_kernels/csrc/core/src/moe/models_sm100.cuh
+189
View File
@@ -0,0 +1,189 @@
#pragma once
// #ifndef NDEBUG
// #define NDEBUG
// #endif
#include <cuda_runtime.h>
#include <cute/tensor.hpp>
#include <cute/atom/mma_atom.hpp>
#include <cute/atom/mma_traits_sm90_gmma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
#include <cutlass/numeric_types.h>
namespace liger {
using namespace cute;
// ═══════════════════════════════════════════════════════════════════
// GMMA atom selection
// ═══════════════════════════════════════════════════════════════════
template <typename Element, int N>
struct GmmaSelector;
static constexpr auto GmmaK = GMMA::Major::K;
static constexpr auto GmmaMN = GMMA::Major::MN;
template <> struct GmmaSelector<bfloat16_t, 8> { using Atom = SM90_64x8x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 16> { using Atom = SM90_64x16x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 32> { using Atom = SM90_64x32x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 64> { using Atom = SM90_64x64x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 128> { using Atom = SM90_64x128x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 192> { using Atom = SM90_64x192x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 256> { using Atom = SM90_64x256x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 8> { using Atom = SM90_64x8x16_F32F16F16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 16> { using Atom = SM90_64x16x16_F32F16F16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 32> { using Atom = SM90_64x32x16_F32F16F16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 64> { using Atom = SM90_64x64x16_F32F16F16_SS<GmmaK, GmmaK>; };
// ═══════════════════════════════════════════════════════════════════
// GMMA atom selection — A=K-major, B=MN-major (SS<K, MN>)
// Used when B operand has MN-contiguous smem (e.g. mlp2_t_fused).
// ═══════════════════════════════════════════════════════════════════
template <typename Element, int N>
struct GmmaSelectorKMN;
template <> struct GmmaSelectorKMN<bfloat16_t, 8> { using Atom = SM90_64x8x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 16> { using Atom = SM90_64x16x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 32> { using Atom = SM90_64x32x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 64> { using Atom = SM90_64x64x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 128> { using Atom = SM90_64x128x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 256> { using Atom = SM90_64x256x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 8> { using Atom = SM90_64x8x16_F32F16F16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 16> { using Atom = SM90_64x16x16_F32F16F16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 32> { using Atom = SM90_64x32x16_F32F16F16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 64> { using Atom = SM90_64x64x16_F32F16F16_SS<GmmaK, GmmaMN>; };
// Backward compat alias
template <typename Element, int N>
struct GmmaSelectorMN : GmmaSelectorKMN<Element, N> {};
// ═══════════════════════════════════════════════════════════════════
// GMMA atom selection — both MN-major (SS<MN, MN>)
// Used when both operands have MN-contiguous smem (e.g. mlp3).
// ═══════════════════════════════════════════════════════════════════
template <typename Element, int N>
struct GmmaSelectorMNMN;
template <> struct GmmaSelectorMNMN<bfloat16_t, 8> { using Atom = SM90_64x8x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 16> { using Atom = SM90_64x16x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 32> { using Atom = SM90_64x32x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 64> { using Atom = SM90_64x64x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 128> { using Atom = SM90_64x128x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 256> { using Atom = SM90_64x256x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 8> { using Atom = SM90_64x8x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 16> { using Atom = SM90_64x16x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 32> { using Atom = SM90_64x32x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 64> { using Atom = SM90_64x64x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
// ═══════════════════════════════════════════════════════════════════
// 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<Element, TileN>::Atom;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, Layout<Shape<_1, _1, _1>>>;
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutA_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutB_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileN>, Int<TileK>>{}));
using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutB = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}));
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
static constexpr int TmaTransactionBytesA = static_cast<int>(size(SmemLayoutA_1{}) * sizeof(Element));
static constexpr int TmaTransactionBytesB = static_cast<int>(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 <typename Traits>
struct RouterSmem {
using Element = typename Traits::Element;
static constexpr int smem_A_size = cosize_v<typename Traits::SmemLayoutA>;
static constexpr int smem_B_size = cosize_v<typename Traits::SmemLayoutB>;
// 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
liger_cute_kernels/csrc/core/src/moe/models_sm90.cuh
+189
View File
@@ -0,0 +1,189 @@
#pragma once
// #ifndef NDEBUG
// #define NDEBUG
// #endif
#include <cuda_runtime.h>
#include <cute/tensor.hpp>
#include <cute/atom/mma_atom.hpp>
#include <cute/atom/mma_traits_sm90_gmma.hpp>
#include <cutlass/pipeline/sm90_pipeline.hpp>
#include <cutlass/arch/barrier.h>
#include <cutlass/numeric_types.h>
namespace liger {
using namespace cute;
// ═══════════════════════════════════════════════════════════════════
// GMMA atom selection
// ═══════════════════════════════════════════════════════════════════
template <typename Element, int N>
struct GmmaSelector;
static constexpr auto GmmaK = GMMA::Major::K;
static constexpr auto GmmaMN = GMMA::Major::MN;
template <> struct GmmaSelector<bfloat16_t, 8> { using Atom = SM90_64x8x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 16> { using Atom = SM90_64x16x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 32> { using Atom = SM90_64x32x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 64> { using Atom = SM90_64x64x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 128> { using Atom = SM90_64x128x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 192> { using Atom = SM90_64x192x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<bfloat16_t, 256> { using Atom = SM90_64x256x16_F32BF16BF16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 8> { using Atom = SM90_64x8x16_F32F16F16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 16> { using Atom = SM90_64x16x16_F32F16F16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 32> { using Atom = SM90_64x32x16_F32F16F16_SS<GmmaK, GmmaK>; };
template <> struct GmmaSelector<half_t, 64> { using Atom = SM90_64x64x16_F32F16F16_SS<GmmaK, GmmaK>; };
// ═══════════════════════════════════════════════════════════════════
// GMMA atom selection — A=K-major, B=MN-major (SS<K, MN>)
// Used when B operand has MN-contiguous smem (e.g. mlp2_t_fused).
// ═══════════════════════════════════════════════════════════════════
template <typename Element, int N>
struct GmmaSelectorKMN;
template <> struct GmmaSelectorKMN<bfloat16_t, 8> { using Atom = SM90_64x8x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 16> { using Atom = SM90_64x16x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 32> { using Atom = SM90_64x32x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 64> { using Atom = SM90_64x64x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 128> { using Atom = SM90_64x128x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<bfloat16_t, 256> { using Atom = SM90_64x256x16_F32BF16BF16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 8> { using Atom = SM90_64x8x16_F32F16F16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 16> { using Atom = SM90_64x16x16_F32F16F16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 32> { using Atom = SM90_64x32x16_F32F16F16_SS<GmmaK, GmmaMN>; };
template <> struct GmmaSelectorKMN<half_t, 64> { using Atom = SM90_64x64x16_F32F16F16_SS<GmmaK, GmmaMN>; };
// Backward compat alias
template <typename Element, int N>
struct GmmaSelectorMN : GmmaSelectorKMN<Element, N> {};
// ═══════════════════════════════════════════════════════════════════
// GMMA atom selection — both MN-major (SS<MN, MN>)
// Used when both operands have MN-contiguous smem (e.g. mlp3).
// ═══════════════════════════════════════════════════════════════════
template <typename Element, int N>
struct GmmaSelectorMNMN;
template <> struct GmmaSelectorMNMN<bfloat16_t, 8> { using Atom = SM90_64x8x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 16> { using Atom = SM90_64x16x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 32> { using Atom = SM90_64x32x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 64> { using Atom = SM90_64x64x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 128> { using Atom = SM90_64x128x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<bfloat16_t, 256> { using Atom = SM90_64x256x16_F32BF16BF16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 8> { using Atom = SM90_64x8x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 16> { using Atom = SM90_64x16x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 32> { using Atom = SM90_64x32x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
template <> struct GmmaSelectorMNMN<half_t, 64> { using Atom = SM90_64x64x16_F32F16F16_SS<GmmaMN, GmmaMN>; };
// ═══════════════════════════════════════════════════════════════════
// 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<Element, TileN>::Atom;
using TiledMma = TiledMMA<MMA_Atom<GmmaAtom>, Layout<Shape<_1, _1, _1>>>;
using SmemLayoutAtom = GMMA::Layout_K_SW128_Atom<Element>;
using SmemLayoutA_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileM>, Int<TileK>>{}));
using SmemLayoutB_1 = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileN>, Int<TileK>>{}));
using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileM>, Int<TileK>, Int<Stages>>{}));
using SmemLayoutB = decltype(tile_to_shape(SmemLayoutAtom{}, Shape<Int<TileN>, Int<TileK>, Int<Stages>>{}));
using MainloopPipeline = cutlass::PipelineTmaAsync<Stages>;
using PipelineState = cutlass::PipelineState<Stages>;
static constexpr int TmaTransactionBytesA = static_cast<int>(size(SmemLayoutA_1{}) * sizeof(Element));
static constexpr int TmaTransactionBytesB = static_cast<int>(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 <typename Traits>
struct RouterSmem {
using Element = typename Traits::Element;
static constexpr int smem_A_size = cosize_v<typename Traits::SmemLayoutA>;
static constexpr int smem_B_size = cosize_v<typename Traits::SmemLayoutB>;
// 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
liger_cute_kernels/csrc/core/src/moe/moe.cu
+4 -1770
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/moe.cuh
+8 -197
View File
@@ -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 <typename Traits1, typename Traits2, int Compute = 90>
struct MoeSmem {
// MLP shared memory (fused phase 1 + 2 with pipeline barriers).
MlpFusedSmem<Traits1, Traits2, Compute> mlp;
// Comm shared memory (total_tiles count + remote_offsets pointer).
CommSmem comm;
};
// ═══════════════════════════════════════════════════════════════════
// moe_fused_fwd — MLP warps only (no CommBuffers in scope)
// ═══════════════════════════════════════════════════════════════════
template <typename Traits1, typename Traits2,
int ZBufferSlots,
int CommNumStages,
int NC,
// kSubTiles = CommTileM / GemmTileM. >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<Traits1, Traits2, Compute>& 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<Compute> 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<Traits1, Traits2, ZBufferSlots, SubTiles, Compute>(
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
liger_cute_kernels/csrc/core/src/moe/moe_bwd.cu
+4 -2087
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/moe_bwd.cuh
+8 -216
View File
@@ -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 <typename Traits1, typename Traits2T, typename Traits3,
typename Traits4, typename Traits5, int Compute = 90>
struct MoeBwdSmem {
MlpFusedBwdSmem<Traits1, Traits2T, Traits3, Traits4, Traits5, Compute> mlp;
CommSmem comm;
};
// ═══════════════════════════════════════════════════════════════════
// moe_fused_bwd — local phase → barrier → remote phase
// ═══════════════════════════════════════════════════════════════════
template <typename Traits1, typename Traits2T, typename Traits3,
typename Traits4, typename Traits5,
int CommNumStages, int NSplit2, int SubBatch,
// SubTiles = CommTileM / GemmTileM. >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<Traits1, Traits2T, Traits3, Traits4, Traits5, Compute>& 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<typename Traits1::Element>& 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<Compute> global_barrier(bufs.barrier_counter,
(int)gridDim.x * (int)gridDim.y);
MlpBwdCtaBarrierT<Compute> 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<typename Traits1::Element> 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
liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.cuh
+219
View File
@@ -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 <typename Traits1, typename Traits2T, typename Traits3,
typename Traits4, typename Traits5, int Compute = 90>
struct MoeBwdSmem {
MlpFusedBwdSmem<Traits1, Traits2T, Traits3, Traits4, Traits5, Compute> mlp;
CommSmem comm;
};
// ═══════════════════════════════════════════════════════════════════
// moe_fused_bwd — local phase → barrier → remote phase
// ═══════════════════════════════════════════════════════════════════
template <typename Traits1, typename Traits2T, typename Traits3,
typename Traits4, typename Traits5,
int CommNumStages, int NSplit2, int SubBatch,
// SubTiles = CommTileM / GemmTileM. >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<Traits1, Traits2T, Traits3, Traits4, Traits5, Compute>& 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<typename Traits1::Element>& 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<Compute> global_barrier(bufs.barrier_counter,
(int)gridDim.x * (int)gridDim.y);
MlpBwdCtaBarrierT<Compute> 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<typename Traits1::Element> 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
liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm100.inc
+2091
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.cuh
+219
View File
@@ -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 <typename Traits1, typename Traits2T, typename Traits3,
typename Traits4, typename Traits5, int Compute = 90>
struct MoeBwdSmem {
MlpFusedBwdSmem<Traits1, Traits2T, Traits3, Traits4, Traits5, Compute> mlp;
CommSmem comm;
};
// ═══════════════════════════════════════════════════════════════════
// moe_fused_bwd — local phase → barrier → remote phase
// ═══════════════════════════════════════════════════════════════════
template <typename Traits1, typename Traits2T, typename Traits3,
typename Traits4, typename Traits5,
int CommNumStages, int NSplit2, int SubBatch,
// SubTiles = CommTileM / GemmTileM. >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<Traits1, Traits2T, Traits3, Traits4, Traits5, Compute>& 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<typename Traits1::Element>& 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<Compute> global_barrier(bufs.barrier_counter,
(int)gridDim.x * (int)gridDim.y);
MlpBwdCtaBarrierT<Compute> 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<typename Traits1::Element> 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
liger_cute_kernels/csrc/core/src/moe/moe_bwd_sm90.inc
+2091
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/moe_comm_config.cuh
+8 -42
View File
@@ -1,45 +1,11 @@
#pragma once
namespace liger {
#ifndef LIGER_CUTE_DISPATCH_COMPUTE
#define LIGER_CUTE_DISPATCH_COMPUTE 0
#endif
template <int Compute>
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
liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm100.cuh
+45
View File
@@ -0,0 +1,45 @@
#pragma once
namespace liger {
template <int Compute>
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
liger_cute_kernels/csrc/core/src/moe/moe_comm_config_sm90.cuh
+45
View File
@@ -0,0 +1,45 @@
#pragma once
namespace liger {
template <int Compute>
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
liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm100.cuh
+1 -1
View File
@@ -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.
liger_cute_kernels/csrc/core/src/moe/moe_dispatch_configs_sm90.cuh
+1 -1
View File
@@ -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:
liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm100.hpp
+182
View File
@@ -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
liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tune_configs_sm90.hpp
+182
View File
@@ -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
liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm100.cuh
+29
View File
@@ -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
liger_cute_kernels/csrc/core/src/moe/moe_fwd_bwd_tuning_configs_sm90.cuh
+29
View File
@@ -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
liger_cute_kernels/csrc/core/src/moe/moe_sm100.cuh
+200
View File
@@ -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 <typename Traits1, typename Traits2, int Compute = 90>
struct MoeSmem {
// MLP shared memory (fused phase 1 + 2 with pipeline barriers).
MlpFusedSmem<Traits1, Traits2, Compute> mlp;
// Comm shared memory (total_tiles count + remote_offsets pointer).
CommSmem comm;
};
// ═══════════════════════════════════════════════════════════════════
// moe_fused_fwd — MLP warps only (no CommBuffers in scope)
// ═══════════════════════════════════════════════════════════════════
template <typename Traits1, typename Traits2,
int ZBufferSlots,
int CommNumStages,
int NC,
// kSubTiles = CommTileM / GemmTileM. >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<Traits1, Traits2, Compute>& 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<Compute> 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<Traits1, Traits2, ZBufferSlots, SubTiles, Compute>(
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
liger_cute_kernels/csrc/core/src/moe/moe_sm100.inc
+1774
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/moe_sm90.cuh
+200
View File
@@ -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 <typename Traits1, typename Traits2, int Compute = 90>
struct MoeSmem {
// MLP shared memory (fused phase 1 + 2 with pipeline barriers).
MlpFusedSmem<Traits1, Traits2, Compute> mlp;
// Comm shared memory (total_tiles count + remote_offsets pointer).
CommSmem comm;
};
// ═══════════════════════════════════════════════════════════════════
// moe_fused_fwd — MLP warps only (no CommBuffers in scope)
// ═══════════════════════════════════════════════════════════════════
template <typename Traits1, typename Traits2,
int ZBufferSlots,
int CommNumStages,
int NC,
// kSubTiles = CommTileM / GemmTileM. >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<Traits1, Traits2, Compute>& 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<Compute> 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<Traits1, Traits2, ZBufferSlots, SubTiles, Compute>(
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
liger_cute_kernels/csrc/core/src/moe/moe_sm90.inc
+1774
View File
File diff suppressed because it is too large Load Diff
liger_cute_kernels/csrc/core/src/moe/moe_symm_config.cuh
+8 -49
View File
@@ -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 <nvshmem.h> 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
liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm100.cuh
+52
View File
@@ -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 <nvshmem.h> 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
liger_cute_kernels/csrc/core/src/moe/moe_symm_config_sm90.cuh
+52
View File
@@ -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 <nvshmem.h> 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
liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused.cuh
+8 -241
View File
@@ -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 <typename Element>
__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<const uint4*>(dZ);
auto* U_v = reinterpret_cast<const uint4*>(U);
auto* V_v = reinterpret_cast<const uint4*>(V);
auto* dU_v = reinterpret_cast<uint4*>(dU);
auto* dV_v = reinterpret_cast<uint4*>(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<Element*>(&dz_vec);
auto* u_e = reinterpret_cast<Element*>(&u_vec);
auto* v_e = reinterpret_cast<Element*>(&v_vec);
uint4 du_out, dv_out;
auto* du_e = reinterpret_cast<Element*>(&du_out);
auto* dv_e = reinterpret_cast<Element*>(&dv_out);
CUTE_UNROLL
for (int j = 0; j < kVecSize; ++j) {
float dz = static_cast<float>(dz_e[j]);
float u = static_cast<float>(u_e[j]);
float v = static_cast<float>(v_e[j]);
float f = fast_sigmoid(u);
float df = f * (1.0f + u * (1.0f - f));
dv_e[j] = from_float<Element>(dz * u * f);
du_e[j] = from_float<Element>(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 <typename Element>
__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<const uint4*>(dZ_base);
auto* U_v = reinterpret_cast<const uint4*>(U_base);
auto* V_v = reinterpret_cast<const uint4*>(V_base);
auto* dU_v = reinterpret_cast<uint4*>(dU_base);
auto* dV_v = reinterpret_cast<uint4*>(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<Element*>(&dz_vec);
auto* u_e = reinterpret_cast<Element*>(&u_vec);
auto* v_e = reinterpret_cast<Element*>(&v_vec);
uint4 du_out, dv_out;
auto* du_e = reinterpret_cast<Element*>(&du_out);
auto* dv_e = reinterpret_cast<Element*>(&dv_out);
CUTE_UNROLL
for (int j = 0; j < kVecSize; ++j) {
float dz = static_cast<float>(dz_e[j]);
float u = static_cast<float>(u_e[j]);
float v = static_cast<float>(v_e[j]);
float f = fast_sigmoid(u);
float df = f * (1.0f + u * (1.0f - f));
dv_e[j] = from_float<Element>(dz * u * f);
du_e[j] = from_float<Element>(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 <typename Element>
__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<const uint4*>(dZ_base);
auto* Up_v = reinterpret_cast<const uint4*>(Up_base);
auto* Vp_v = reinterpret_cast<const uint4*>(Vp_base);
auto* dU_v = reinterpret_cast<uint4*>(dU_base);
auto* dV_v = reinterpret_cast<uint4*>(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<Element*>(&dz);
auto* up_e = reinterpret_cast<Element*>(&up);
auto* vp_e = reinterpret_cast<Element*>(&vp);
uint4 du, dv;
auto* du_e = reinterpret_cast<Element*>(&du);
auto* dv_e = reinterpret_cast<Element*>(&dv);
CUTE_UNROLL
for (int j = 0; j < kVecSize; ++j) {
float dz_f = static_cast<float>(dz_e[j]);
float up_f = static_cast<float>(up_e[j]);
float vp_f = static_cast<float>(vp_e[j]);
du_e[j] = from_float<Element>(dz_f * up_f);
dv_e[j] = from_float<Element>(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
liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm100.cuh
+244
View File
@@ -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 <typename Element>
__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<const uint4*>(dZ);
auto* U_v = reinterpret_cast<const uint4*>(U);
auto* V_v = reinterpret_cast<const uint4*>(V);
auto* dU_v = reinterpret_cast<uint4*>(dU);
auto* dV_v = reinterpret_cast<uint4*>(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<Element*>(&dz_vec);
auto* u_e = reinterpret_cast<Element*>(&u_vec);
auto* v_e = reinterpret_cast<Element*>(&v_vec);
uint4 du_out, dv_out;
auto* du_e = reinterpret_cast<Element*>(&du_out);
auto* dv_e = reinterpret_cast<Element*>(&dv_out);
CUTE_UNROLL
for (int j = 0; j < kVecSize; ++j) {
float dz = static_cast<float>(dz_e[j]);
float u = static_cast<float>(u_e[j]);
float v = static_cast<float>(v_e[j]);
float f = fast_sigmoid(u);
float df = f * (1.0f + u * (1.0f - f));
dv_e[j] = from_float<Element>(dz * u * f);
du_e[j] = from_float<Element>(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 <typename Element>
__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<const uint4*>(dZ_base);
auto* U_v = reinterpret_cast<const uint4*>(U_base);
auto* V_v = reinterpret_cast<const uint4*>(V_base);
auto* dU_v = reinterpret_cast<uint4*>(dU_base);
auto* dV_v = reinterpret_cast<uint4*>(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<Element*>(&dz_vec);
auto* u_e = reinterpret_cast<Element*>(&u_vec);
auto* v_e = reinterpret_cast<Element*>(&v_vec);
uint4 du_out, dv_out;
auto* du_e = reinterpret_cast<Element*>(&du_out);
auto* dv_e = reinterpret_cast<Element*>(&dv_out);
CUTE_UNROLL
for (int j = 0; j < kVecSize; ++j) {
float dz = static_cast<float>(dz_e[j]);
float u = static_cast<float>(u_e[j]);
float v = static_cast<float>(v_e[j]);
float f = fast_sigmoid(u);
float df = f * (1.0f + u * (1.0f - f));
dv_e[j] = from_float<Element>(dz * u * f);
du_e[j] = from_float<Element>(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 <typename Element>
__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<const uint4*>(dZ_base);
auto* Up_v = reinterpret_cast<const uint4*>(Up_base);
auto* Vp_v = reinterpret_cast<const uint4*>(Vp_base);
auto* dU_v = reinterpret_cast<uint4*>(dU_base);
auto* dV_v = reinterpret_cast<uint4*>(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<Element*>(&dz);
auto* up_e = reinterpret_cast<Element*>(&up);
auto* vp_e = reinterpret_cast<Element*>(&vp);
uint4 du, dv;
auto* du_e = reinterpret_cast<Element*>(&du);
auto* dv_e = reinterpret_cast<Element*>(&dv);
CUTE_UNROLL
for (int j = 0; j < kVecSize; ++j) {
float dz_f = static_cast<float>(dz_e[j]);
float up_f = static_cast<float>(up_e[j]);
float vp_f = static_cast<float>(vp_e[j]);
du_e[j] = from_float<Element>(dz_f * up_f);
dv_e[j] = from_float<Element>(dz_f * vp_f);
}
dU_v[off] = du;
dV_v[off] = dv;
}
}
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/silu_bwd_fused_sm90.cuh
+244
View File
@@ -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 <typename Element>
__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<const uint4*>(dZ);
auto* U_v = reinterpret_cast<const uint4*>(U);
auto* V_v = reinterpret_cast<const uint4*>(V);
auto* dU_v = reinterpret_cast<uint4*>(dU);
auto* dV_v = reinterpret_cast<uint4*>(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<Element*>(&dz_vec);
auto* u_e = reinterpret_cast<Element*>(&u_vec);
auto* v_e = reinterpret_cast<Element*>(&v_vec);
uint4 du_out, dv_out;
auto* du_e = reinterpret_cast<Element*>(&du_out);
auto* dv_e = reinterpret_cast<Element*>(&dv_out);
CUTE_UNROLL
for (int j = 0; j < kVecSize; ++j) {
float dz = static_cast<float>(dz_e[j]);
float u = static_cast<float>(u_e[j]);
float v = static_cast<float>(v_e[j]);
float f = fast_sigmoid(u);
float df = f * (1.0f + u * (1.0f - f));
dv_e[j] = from_float<Element>(dz * u * f);
du_e[j] = from_float<Element>(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 <typename Element>
__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<const uint4*>(dZ_base);
auto* U_v = reinterpret_cast<const uint4*>(U_base);
auto* V_v = reinterpret_cast<const uint4*>(V_base);
auto* dU_v = reinterpret_cast<uint4*>(dU_base);
auto* dV_v = reinterpret_cast<uint4*>(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<Element*>(&dz_vec);
auto* u_e = reinterpret_cast<Element*>(&u_vec);
auto* v_e = reinterpret_cast<Element*>(&v_vec);
uint4 du_out, dv_out;
auto* du_e = reinterpret_cast<Element*>(&du_out);
auto* dv_e = reinterpret_cast<Element*>(&dv_out);
CUTE_UNROLL
for (int j = 0; j < kVecSize; ++j) {
float dz = static_cast<float>(dz_e[j]);
float u = static_cast<float>(u_e[j]);
float v = static_cast<float>(v_e[j]);
float f = fast_sigmoid(u);
float df = f * (1.0f + u * (1.0f - f));
dv_e[j] = from_float<Element>(dz * u * f);
du_e[j] = from_float<Element>(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 <typename Element>
__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<const uint4*>(dZ_base);
auto* Up_v = reinterpret_cast<const uint4*>(Up_base);
auto* Vp_v = reinterpret_cast<const uint4*>(Vp_base);
auto* dU_v = reinterpret_cast<uint4*>(dU_base);
auto* dV_v = reinterpret_cast<uint4*>(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<Element*>(&dz);
auto* up_e = reinterpret_cast<Element*>(&up);
auto* vp_e = reinterpret_cast<Element*>(&vp);
uint4 du, dv;
auto* du_e = reinterpret_cast<Element*>(&du);
auto* dv_e = reinterpret_cast<Element*>(&dv);
CUTE_UNROLL
for (int j = 0; j < kVecSize; ++j) {
float dz_f = static_cast<float>(dz_e[j]);
float up_f = static_cast<float>(up_e[j]);
float vp_f = static_cast<float>(vp_e[j]);
du_e[j] = from_float<Element>(dz_f * up_f);
dv_e[j] = from_float<Element>(dz_f * vp_f);
}
dU_v[off] = du;
dV_v[off] = dv;
}
}
}
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/tile_iterator.cuh
+8 -524
View File
@@ -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 <cuda_runtime.h>
#include <cuda/atomic>
#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 <typename Element, int TileM = 128>
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 <typename Element, int TileM = 128, int ZBufferSlots = 2>
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 <typename Element, int NumStages, int NC = 2, int TileM = 128,
int MSubTiles = 1>
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<TileM, MSubTiles> 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<int, cuda::thread_scope_device> 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<int, cuda::thread_scope_device> 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
liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd.cuh
+8 -356
View File
@@ -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 <typename Element, int TileM = 128>
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 <typename Element, int NumStages, int NC = 2, int TileM = 128>
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 <typename Element, int NumStages, int NC = 2, int TileM = 128>
struct FusedMlpTileIteratorBwd {
RemoteMlpTileIteratorBwd<Element, NumStages, NC, TileM> 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
liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm100.cuh
+359
View File
@@ -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 <typename Element, int TileM = 128>
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 <typename Element, int NumStages, int NC = 2, int TileM = 128>
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 <typename Element, int NumStages, int NC = 2, int TileM = 128>
struct FusedMlpTileIteratorBwd {
RemoteMlpTileIteratorBwd<Element, NumStages, NC, TileM> 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
liger_cute_kernels/csrc/core/src/moe/tile_iterator_bwd_sm90.cuh
+359
View File
@@ -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 <typename Element, int TileM = 128>
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 <typename Element, int NumStages, int NC = 2, int TileM = 128>
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 <typename Element, int NumStages, int NC = 2, int TileM = 128>
struct FusedMlpTileIteratorBwd {
RemoteMlpTileIteratorBwd<Element, NumStages, NC, TileM> 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
liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm100.cuh
+527
View File
@@ -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 <cuda_runtime.h>
#include <cuda/atomic>
#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 <typename Element, int TileM = 128>
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 <typename Element, int TileM = 128, int ZBufferSlots = 2>
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 <typename Element, int NumStages, int NC = 2, int TileM = 128,
int MSubTiles = 1>
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<TileM, MSubTiles> 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<int, cuda::thread_scope_device> 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<int, cuda::thread_scope_device> 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
liger_cute_kernels/csrc/core/src/moe/tile_iterator_sm90.cuh
+527
View File
@@ -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 <cuda_runtime.h>
#include <cuda/atomic>
#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 <typename Element, int TileM = 128>
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 <typename Element, int TileM = 128, int ZBufferSlots = 2>
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 <typename Element, int NumStages, int NC = 2, int TileM = 128,
int MSubTiles = 1>
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<TileM, MSubTiles> 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<int, cuda::thread_scope_device> 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<int, cuda::thread_scope_device> 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
liger_cute_kernels/csrc/core/src/moe/tma_copy_atom.cuh
+8 -32
View File
@@ -1,35 +1,11 @@
#pragma once
#include <cute/atom/copy_traits_sm90_tma.hpp>
#ifndef LIGER_CUTE_DISPATCH_COMPUTE
#define LIGER_CUTE_DISPATCH_COMPUTE 0
#endif
namespace liger {
using namespace cute;
template <int Compute>
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 <int Compute>
using TmaLoadAtomForCompute = typename TmaCopyAtomForCompute<Compute>::Load;
template <int Compute>
using TmaStoreAtomForCompute = typename TmaCopyAtomForCompute<Compute>::Store;
template <int Compute>
using TmaReduceAddAtomForCompute = typename TmaCopyAtomForCompute<Compute>::ReduceAdd;
} // namespace liger
#if LIGER_CUTE_DISPATCH_COMPUTE == 90
#include "tma_copy_atom_sm90.cuh"
#else
#include "tma_copy_atom_sm100.cuh"
#endif
liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm100.cuh
+35
View File
@@ -0,0 +1,35 @@
#pragma once
#include <cute/atom/copy_traits_sm90_tma.hpp>
namespace liger {
using namespace cute;
template <int Compute>
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 <int Compute>
using TmaLoadAtomForCompute = typename TmaCopyAtomForCompute<Compute>::Load;
template <int Compute>
using TmaStoreAtomForCompute = typename TmaCopyAtomForCompute<Compute>::Store;
template <int Compute>
using TmaReduceAddAtomForCompute = typename TmaCopyAtomForCompute<Compute>::ReduceAdd;
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/tma_copy_atom_sm90.cuh
+35
View File
@@ -0,0 +1,35 @@
#pragma once
#include <cute/atom/copy_traits_sm90_tma.hpp>
namespace liger {
using namespace cute;
template <int Compute>
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 <int Compute>
using TmaLoadAtomForCompute = typename TmaCopyAtomForCompute<Compute>::Load;
template <int Compute>
using TmaStoreAtomForCompute = typename TmaCopyAtomForCompute<Compute>::Store;
template <int Compute>
using TmaReduceAddAtomForCompute = typename TmaCopyAtomForCompute<Compute>::ReduceAdd;
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/tmem_load_op.cuh
+8 -17
View File
@@ -1,20 +1,11 @@
#pragma once
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
#ifndef LIGER_CUTE_DISPATCH_COMPUTE
#define LIGER_CUTE_DISPATCH_COMPUTE 0
#endif
namespace liger {
using namespace cute;
template <int EpiChunkN> 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 <int EpiChunkN>
using TmemLoadOp = typename TmemLoadOpSelector<EpiChunkN>::Op;
} // namespace liger
#if LIGER_CUTE_DISPATCH_COMPUTE == 90
#include "tmem_load_op_sm90.cuh"
#else
#include "tmem_load_op_sm100.cuh"
#endif
liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm100.cuh
+20
View File
@@ -0,0 +1,20 @@
#pragma once
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
namespace liger {
using namespace cute;
template <int EpiChunkN> 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 <int EpiChunkN>
using TmemLoadOp = typename TmemLoadOpSelector<EpiChunkN>::Op;
} // namespace liger
liger_cute_kernels/csrc/core/src/moe/tmem_load_op_sm90.cuh
+20
View File
@@ -0,0 +1,20 @@
#pragma once
#include <cute/arch/copy_sm100.hpp>
#include <cute/atom/copy_traits_sm100.hpp>
namespace liger {
using namespace cute;
template <int EpiChunkN> 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 <int EpiChunkN>
using TmemLoadOp = typename TmemLoadOpSelector<EpiChunkN>::Op;
} // namespace liger
liger_cute_kernels/test/test_moe_architecture_split.py
+86
View File
@@ -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<body>.*?)\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"]