mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-08 13:33:42 +08:00
Daniel Koepping
91 lines
3.1 KiB
C++
91 lines
3.1 KiB
C++
#include "kittens.cuh"
|
|
|
|
using namespace kittens;
|
|
|
|
#ifndef ROUTER_M
|
|
#define ROUTER_M 16384
|
|
#endif
|
|
#ifndef ROUTER_K
|
|
#define ROUTER_K 2880
|
|
#endif
|
|
#ifndef ROUTER_E
|
|
#define ROUTER_E 32
|
|
#endif
|
|
|
|
constexpr int BLOCK_M = 64;
|
|
constexpr int BLOCK_K = 64;
|
|
constexpr int NUM_WARPS = 4;
|
|
constexpr int THREADS = NUM_WARPS * WARP_THREADS;
|
|
|
|
using G = kittens::group<NUM_WARPS>;
|
|
using XST = st_bf<BLOCK_M, BLOCK_K, st_16x32_s>;
|
|
using WST = st_bf<ROUTER_E, BLOCK_K, st_16x32_s>;
|
|
using XRT = rt_bf<16, BLOCK_K, row_l, rt_16x32_s>;
|
|
using WRT = rt_bf<ROUTER_E, BLOCK_K, row_l, rt_16x32_s>;
|
|
using CRT = rt_fl<16, ROUTER_E, col_l, rt_16x16_s>;
|
|
|
|
static_assert(ROUTER_M % BLOCK_M == 0, "ROUTER_M must be divisible by 64");
|
|
static_assert(ROUTER_K % BLOCK_K == 0, "ROUTER_K must be divisible by 64");
|
|
static_assert(ROUTER_E == 32, "the small-N tile is specialized for 32 experts");
|
|
|
|
extern "C" __global__ __launch_bounds__(THREADS, 4) void moe_router_mfma(
|
|
float *__restrict__ out, bf16 *__restrict__ x_ptr, bf16 *__restrict__ weight_ptr,
|
|
bf16 *__restrict__ bias) {
|
|
gl<bf16, 1, 1, ROUTER_M, ROUTER_K> X{x_ptr, nullptr, nullptr, nullptr, nullptr};
|
|
gl<bf16, 1, 1, ROUTER_E, ROUTER_K> W{weight_ptr, nullptr, nullptr, nullptr, nullptr};
|
|
|
|
__shared__ XST Xs[2];
|
|
__shared__ WST Ws[2];
|
|
|
|
XRT xr;
|
|
WRT wr;
|
|
CRT accum;
|
|
zero(accum);
|
|
|
|
const int block_m = __builtin_amdgcn_workgroup_id_x();
|
|
const int warp_m = warpid();
|
|
|
|
// Prefetch the next tile into alternate LDS storage without changing the MFMA accumulation order.
|
|
G::load(Xs[0], X, {0, 0, block_m, 0});
|
|
G::load(Ws[0], W, {0, 0, 0, 0});
|
|
asm volatile("s_waitcnt vmcnt(0)");
|
|
asm volatile("s_waitcnt lgkmcnt(0)");
|
|
__builtin_amdgcn_s_barrier();
|
|
#pragma unroll
|
|
for (int kk = 0; kk < ROUTER_K / BLOCK_K; kk++) {
|
|
const int cur = kk & 1;
|
|
const int next = cur ^ 1;
|
|
load(xr, subtile_inplace<16, BLOCK_K>(Xs[cur], {warp_m, 0}));
|
|
load(wr, subtile_inplace<ROUTER_E, BLOCK_K>(Ws[cur], {0, 0}));
|
|
asm volatile("s_waitcnt lgkmcnt(0)");
|
|
if (kk + 1 < ROUTER_K / BLOCK_K) {
|
|
G::load(Xs[next], X, {0, 0, block_m, kk + 1});
|
|
G::load(Ws[next], W, {0, 0, 0, kk + 1});
|
|
}
|
|
__builtin_amdgcn_s_setprio(1);
|
|
mma_ABt(accum, xr, wr, accum);
|
|
__builtin_amdgcn_s_setprio(0);
|
|
__builtin_amdgcn_sched_barrier(0);
|
|
if (kk + 1 < ROUTER_K / BLOCK_K) {
|
|
asm volatile("s_waitcnt vmcnt(0)");
|
|
asm volatile("s_waitcnt lgkmcnt(0)");
|
|
__builtin_amdgcn_s_barrier();
|
|
}
|
|
}
|
|
|
|
// A 16x16 MFMA accumulator is column-layout: each lane owns four consecutive rows
|
|
// at one column. Store all 64x32 FP32 results directly; no padded or undersized output ABI.
|
|
const int lane = laneid();
|
|
const int row0 = block_m * BLOCK_M + warp_m * 16 + 4 * (lane / 16);
|
|
const int lane_col = lane % 16;
|
|
#pragma unroll
|
|
for (int j = 0; j < ROUTER_E / 16; j++) {
|
|
const int col = j * 16 + lane_col;
|
|
const float b = (float)bias[col];
|
|
const float vals[4] = {accum.tiles[0][j].data[0].x, accum.tiles[0][j].data[0].y,
|
|
accum.tiles[0][j].data[1].x, accum.tiles[0][j].data[1].y};
|
|
#pragma unroll
|
|
for (int r = 0; r < 4; r++) out[(long long)(row0 + r) * ROUTER_E + col] = vals[r] + b;
|
|
}
|
|
}
|