Niko1221/Strata
1
0 Fork 0
mirror of https://github.com/Niko1221/Strata.git synced 2026-11-09 03:59:45 -10:00
Code Issues Packages Projects Releases 1 Wiki Activity

Strata Q2 blob fused experts (moe_fused.cu CUDA and gfx11 kernels) and their parity test follow the compiled geometry

Browse Source 
Weight-row counts, K stages, code/scale row bytes, the scale staging copies, the H row stride, the down store stride and
the launch grids come from kMoeEmbd / kMoeFF.  The CUDA kernel stages a row's gate/up scales in 16-byte copies, so the hidden
width must be a multiple of 512 (static_assert).
This commit is contained in:
Niko1221 committed 2026-11-08 04:48:52 +01:00
1 parent 2351f2cccb
commit da7b1d022e
2 files changed
+52 -39
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/prefill/moe_fused.cu
+36 -26
View File
@@ -44,17 +44,23 @@ void ck(cudaError_t e, const char* what) {
constexpr int AB = 80; // activation bytes per 64 values: 64 codes, then {d0, c0, d1, c1}
constexpr int MAGIC = 0x4B400000; // the bits of 1.5 * 2^23
constexpr float MAGICF = 12582912.0f;
// The Strata Q2_0 blob (moe_mmq.cu strata_q2_kernel, kernels.cu blob_dequant_kernel): gate/up codes [1280][640 B]
// (rows interleaved: gate of feature f at row 2f, its up at 2f + 1), down codes [2560][160 B], gate/up scales
// [1280][40] fp16, down scales [2560][10] fp16.
constexpr size_t O_GU_CODES = 0, O_D_CODES = (size_t) 1280 * 640, O_GU_SC = O_D_CODES + (size_t) 2560 * 160,
O_D_SC = O_GU_SC + (size_t) 1280 * 40 * 2;
// The Strata Q2_0 blob (moe_mmq.cu strata_q2_kernel, kernels.cu blob_dequant_kernel): gate/up codes [2 F][E / 4 B]
// (rows interleaved: gate of feature f at row 2f, its up at 2f + 1), down codes [E][F / 4 B], gate/up scales
// [2 F][E / 64] fp16, down scales [E][F / 64] fp16; E, F the compiled geometry's (Flash-Next 2560, 640: [1280][640 B],
// [2560][160 B], [1280][40], [2560][10]).
constexpr int GU_ROWS = kMoeGuRows, D_ROWS = kMoeEmbd; // weight rows of the two products
constexpr int ROW_GU = kMoeEmbd / 4, ROW_D = kMoeFF / 4; // code bytes per weight row
constexpr int SC_GU = kMoeEmbd / 64, SC_D = kMoeFF / 64; // 64-weight blocks (scales) per row = stages along K
constexpr int H_ROW_BYTES = SC_D * AB; // bytes of a row of H (AB per 64 values)
static_assert(kMoeEmbd % 64 == 0 && kMoeFF % 64 == 0, "the Strata Q2 blob holds whole 64-weight blocks");
constexpr size_t O_GU_CODES = 0, O_D_CODES = (size_t) GU_ROWS * ROW_GU, O_GU_SC = O_D_CODES + (size_t) D_ROWS * ROW_D,
O_D_SC = O_GU_SC + (size_t) GU_ROWS * SC_GU * 2;
constexpr int THREADS = 512; // 16 warps: 4 along the weight rows (64 each) x 4 along the tile rows (16 each)
constexpr int WROWS = 256; // weight rows per work item: 128 features of gate/up, or 256 outputs of down
constexpr int STAGES = 4; // cp.async pipeline depth, one 64-weight block of K per stage
constexpr int STAGE_BYTES = WROWS * 16 + kTileRows * AB;
constexpr size_t smem_bytes(bool gu) {
return (size_t) WROWS * (gu ? 40 : 10) * 2 + STAGES * STAGE_BYTES + kTileRows * 4;
return (size_t) WROWS * (gu ? SC_GU : SC_D) * 2 + STAGES * STAGE_BYTES + kTileRows * 4;
}
// the stored position of value k of a 32-value half-block (see "The K order")
@@ -220,7 +226,7 @@ __device__ __forceinline__ void mma_s8(int (&d)[4], const uint32_t (&a)[4], uint
// One work item = (tile of up to 64 routed rows of expert e, a block of 256 weight rows); a persistent grid walks
// the batch's items (their count is on the device). GU: gate/up rows of 128 features against the rows' tokens'
// activations (`act` per token, 2560 values), SwiGLU, H to int8 per 32 features into `out` (per row, 640 values).
// activations (`act` per token, E values), SwiGLU, H to int8 per 32 features into `out` (per row, F values).
// !GU: down rows (256 outputs) against H (`act` per row), FP32 into `dm` per row.
// Warp (wf, wt): weight rows 64 wf.. as four m16 tiles (GU: a tile = 8 features, its gate rows as mma rows 0-7 and
// its up rows as 8-15, so a lane holds gate and up of one feature: SwiGLU in registers), tile rows 16 wt.. as two n8.
@@ -229,11 +235,14 @@ __global__ void __launch_bounds__(THREADS, 1)
expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, const int32_t* __restrict__ src,
uint8_t* __restrict__ out, float* __restrict__ dm) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750
constexpr int NKB = GU ? 40 : 10; // 64-weight blocks along K
constexpr int NFB = GU ? 5 : 10; // weight-row blocks per tile
constexpr int CODE_LD = GU ? 640 : 160; // code bytes per weight row
constexpr int NKB = GU ? SC_GU : SC_D; // 64-weight blocks along K
constexpr int NFB = (GU ? GU_ROWS : D_ROWS) / WROWS; // weight-row blocks per tile
static_assert(GU_ROWS % WROWS == 0 && D_ROWS % WROWS == 0, "the Strata Q2 expert kernel: weight rows are whole work items of 256");
static_assert((SC_GU * 2) % 16 == 0, "the Strata Q2 expert kernel stages a row's gate/up scales in 16-byte copies (hidden width a multiple of 512)");
constexpr int NSC = GU ? SC_GU * 2 / 16 : SC_D * 2 / 4; // scale copies per weight row (16 B for gate/up, 4 B for down)
constexpr int CODE_LD = GU ? ROW_GU : ROW_D; // code bytes per weight row
constexpr size_t O_CODES = GU ? O_GU_CODES : O_D_CODES, O_SC = GU ? O_GU_SC : O_D_SC;
constexpr int ACT_LD = NKB * AB; // bytes per activation row: 3200 (a token), 800 (a row's H)
constexpr int ACT_LD = NKB * AB; // bytes per activation row: 3200 (a token), 800 (a row's H) for Flash-Next
constexpr unsigned M2 = 0x03030303u;
extern __shared__ __align__(16) uint8_t smem[];
const __half* wsc = (const __half*) smem; // [256][NKB] the work item's weight scales
@@ -259,10 +268,10 @@ expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, c
const int r = row0 + min(tid, nrows - 1); // rows past the tile's end repeat its last one
srow[tid] = GU ? src[r] : r;
}
for (int c = tid; c < WROWS * 5; c += THREADS) { // the scales: 5 x 16 B (GU) or 5 x 4 B a row
const int r = c / 5, q = c % 5;
if (GU) cp16(smem + r * 80 + q * 16, blob + O_SC + (size_t) (rbase + r) * 80 + q * 16);
else cp4(smem + r * 20 + q * 4, blob + O_SC + (size_t) (rbase + r) * 20 + q * 4);
for (int c = tid; c < WROWS * NSC; c += THREADS) { // the scales: NSC x 16 B (GU) or NSC x 4 B a row (5 for Flash-Next)
const int r = c / NSC, q = c % NSC;
if (GU) cp16(smem + r * (SC_GU * 2) + q * 16, blob + O_SC + (size_t) (rbase + r) * (SC_GU * 2) + q * 16);
else cp4(smem + r * (SC_D * 2) + q * 4, blob + O_SC + (size_t) (rbase + r) * (SC_D * 2) + q * 4);
}
__syncthreads(); // srow
auto load = [&](int kb) {
@@ -356,7 +365,7 @@ expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, c
for (int o = 4; o < 32; o <<= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
const int r = nb0 + 8 * n + 2 * tig + c;
if (r < nrows) {
uint8_t* o = out + (size_t) (row0 + r) * (10 * AB) + (blk >> 1) * AB;
uint8_t* o = out + (size_t) (row0 + r) * H_ROW_BYTES + (blk >> 1) * AB;
const int hh = blk & 1;
#pragma unroll
for (int i = 0; i < 4; ++i) o[32 * hh + perm32(8 * i + g)] = (uint8_t) (int8_t) qv[i];
@@ -373,7 +382,7 @@ expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, c
#pragma unroll
for (int q = 0; q < 4; ++q) {
const int r = nb0 + 8 * n + 2 * tig + (q & 1);
if (r < nrows) dm[(size_t) (row0 + r) * 2560 + rbase + (q < 2 ? ra[i] : rb[i])] = acc[i][n][q];
if (r < nrows) dm[(size_t) (row0 + r) * kMoeEmbd + rbase + (q < 2 ? ra[i] : rb[i])] = acc[i][n][q];
}
}
}
@@ -417,9 +426,10 @@ __global__ void __launch_bounds__(W_THREADS)
expert_w11_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, const int32_t* __restrict__ src,
uint8_t* __restrict__ out, float* __restrict__ dm) {
#if STRATA_FUSED_W11
constexpr int NKB = GU ? 40 : 10; // 64-weight blocks along K
constexpr int NFB = (GU ? 1280 : 2560) / W_ROWS; // work items per tile
constexpr int CODE_LD = GU ? 640 : 160; // code bytes per weight row
constexpr int NKB = GU ? SC_GU : SC_D; // 64-weight blocks along K
constexpr int NFB = (GU ? GU_ROWS : D_ROWS) / W_ROWS; // work items per tile
static_assert(GU_ROWS % W_ROWS == 0 && D_ROWS % W_ROWS == 0, "the Strata Q2 gfx11 kernel: weight rows are whole work items of 128");
constexpr int CODE_LD = GU ? ROW_GU : ROW_D; // code bytes per weight row
constexpr size_t O_CODES = GU ? O_GU_CODES : O_D_CODES, O_SC = GU ? O_GU_SC : O_D_SC;
constexpr int ACT_LD = NKB * AB; // bytes per activation row
constexpr unsigned M2 = 0x03030303u;
@@ -555,7 +565,7 @@ expert_w11_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ ac
#pragma unroll 8
for (int j = 0; j < 32; ++j) am = fmaxf(am, fabsf(hs[r][32 * hh + j]));
const float inv = am > 0.0f ? 127.0f / am : 0.0f;
uint8_t* o = out + (size_t) (row0 + r) * (10 * AB) + (size_t) fb * AB;
uint8_t* o = out + (size_t) (row0 + r) * H_ROW_BYTES + (size_t) fb * AB;
int sum = 0;
uint32_t wd[8] = {};
#pragma unroll
@@ -578,7 +588,7 @@ expert_w11_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ ac
for (int nt = 0; nt < 2; ++nt) {
const int r = 32 * wn + 16 * nt + l16;
if (r < nrows) {
float* d = dm + (size_t) (row0 + r) * 2560 + rbase + 32 * wm + 16 * mt + hi;
float* d = dm + (size_t) (row0 + r) * kMoeEmbd + rbase + 32 * wm + 16 * mt + hi;
#pragma unroll
for (int i = 0; i < 8; ++i) d[2 * i] = acc[mt][nt][i];
}
@@ -735,13 +745,13 @@ void experts(const Batch& b, int n_expert, int64_t n, const void* scratch, const
// the most tiles the batch can have: every row in it, plus a partial tile per expert
const int64_t tiles = (n + kTileRows - 1) / kTileRows + (b.e1 - b.e0);
#if defined(__HIPCC__)
const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (1280 / W_ROWS), (int64_t) d.sms * wgp_blocks(d.occ_gu));
const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (2560 / W_ROWS), (int64_t) d.sms * wgp_blocks(d.occ_d));
const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (GU_ROWS / W_ROWS), (int64_t) d.sms * wgp_blocks(d.occ_gu));
const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (D_ROWS / W_ROWS), (int64_t) d.sms * wgp_blocks(d.occ_d));
expert_w11_kernel<true><<<g_gu, W_THREADS, 0, s>>>(b, tb, (const uint8_t*) xa, src, (uint8_t*) ha, nullptr);
expert_w11_kernel<false><<<g_d, W_THREADS, 0, s>>>(b, tb, (const uint8_t*) ha, src, nullptr, dm);
#else
const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * 5, (int64_t) d.sms * d.occ_gu);
const unsigned g_d = (unsigned) std::min<int64_t>(tiles * 10, (int64_t) d.sms * d.occ_d);
const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (GU_ROWS / WROWS), (int64_t) d.sms * d.occ_gu);
const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (D_ROWS / WROWS), (int64_t) d.sms * d.occ_d);
expert_kernel<true><<<g_gu, THREADS, smem_bytes(true), s>>>(b, tb, (const uint8_t*) xa, src, (uint8_t*) ha,
nullptr);
expert_kernel<false><<<g_d, THREADS, smem_bytes(false), s>>>(b, tb, (const uint8_t*) ha, src, nullptr, dm);
tests/cuda/prefill_fused_moe_test.cpp
+16 -13
View File
@@ -32,9 +32,12 @@ namespace {
namespace mmq = strata::prefill::mmq;
namespace fused = strata::prefill::fused;
constexpr int N = 2560, FF = 640, K = 10, GROUP = 16, Q2 = 42; // GGML_TYPE_Q2_0
constexpr size_t BLOB = 1382400, O_D_CODES = (size_t) 1280 * 640, O_GU_SC = O_D_CODES + (size_t) 2560 * 160,
O_D_SC = O_GU_SC + (size_t) 1280 * 40 * 2;
// The compiled geometry's widths and top-k (Flash-Next: 2560 x 640, top 10); the blob is the Strata Q2 layout at that shape
// (Flash-Next: 1,382,400 B): gate/up codes [2 FF][N / 4 B], down codes [N][FF / 4 B], gate/up scales [2 FF][N / 64] f16, down scales [N][FF / 64] f16.
constexpr int N = fused::kMoeEmbd, FF = fused::kMoeFF, FF2 = fused::kMoeGuRows, K = (int) strata::core::kCompiledGeometry.n_expert_used, GROUP = 16, Q2 = 42; // GGML_TYPE_Q2_0
constexpr int ROW_GU = N / 4, ROW_D = FF / 4, SC_GU = N / 64, SC_D = FF / 64;
constexpr size_t O_D_CODES = (size_t) FF2 * ROW_GU, O_GU_SC = O_D_CODES + (size_t) N * ROW_D, O_D_SC = O_GU_SC + (size_t) FF2 * SC_GU * 2,
BLOB = O_D_SC + (size_t) N * SC_D * 2;
void ck(cudaError_t e, const char* what) {
if (e != cudaSuccess) throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(e));
@@ -62,7 +65,7 @@ std::vector<uint8_t> make_blob(std::mt19937& rng) {
return b;
}
// the blob's weights as floats: gate/up [1280][2560] (interleaved rows), down [2560][640]
// the blob's weights as floats: gate/up [2 FF][N] (interleaved rows), down [N][FF]
void dequant(const std::vector<uint8_t>& b, std::vector<float>& gu, std::vector<float>& dn) {
auto row = [&](size_t codes, size_t ld, size_t scales, int nb, int r, float* out) {
for (int k = 0; k < nb * 64; ++k) {
@@ -72,10 +75,10 @@ void dequant(const std::vector<uint8_t>& b, std::vector<float>& gu, std::vector<
out[k] = ggml_fp16_to_fp32(h) * (float) (q - 1);
}
};
gu.resize((size_t) 1280 * N);
gu.resize((size_t) FF2 * N);
dn.resize((size_t) N * FF);
for (int r = 0; r < 1280; ++r) row(0, 640, O_GU_SC, 40, r, gu.data() + (size_t) r * N);
for (int r = 0; r < N; ++r) row(O_D_CODES, 160, O_D_SC, 10, r, dn.data() + (size_t) r * FF);
for (int r = 0; r < FF2; ++r) row(0, ROW_GU, O_GU_SC, SC_GU, r, gu.data() + (size_t) r * N);
for (int r = 0; r < N; ++r) row(O_D_CODES, ROW_D, O_D_SC, SC_D, r, dn.data() + (size_t) r * FF);
}
// activations: N(0, 1) with a few large values (the hidden state has outliers); token `zero` all zero
@@ -132,14 +135,14 @@ Routing sort_rows(const std::vector<int32_t>& ids, int E) {
struct MmqBufs {
Dev xq, gu, h, hq, dm, ident, bounds, grp_gu, grp_d;
MmqBufs(int64_t rows, int E)
: xq(mmq::q8_bytes(rows, N)), gu((size_t) rows * 1280 * 4), h((size_t) rows * FF * 4),
: xq(mmq::q8_bytes(rows, N)), gu((size_t) rows * FF2 * 4), h((size_t) rows * FF * 4),
hq(mmq::q8_bytes(rows, FF)),
dm((size_t) rows * N * 4), ident((size_t) rows * 4), bounds((size_t) (2 * (E + E / GROUP + 2)) * 4),
grp_gu(GROUP * mmq::matrix_bytes(Q2, 1280, N) + 4096), grp_d(GROUP * mmq::matrix_bytes(Q2, N, FF) + 4096) {}
grp_gu(GROUP * mmq::matrix_bytes(Q2, FF2, N) + 4096), grp_d(GROUP * mmq::matrix_bytes(Q2, N, FF) + 4096) {}
};
void run_mmq(mmq::Context& ctx, MmqBufs& b, const Routing& r, const float* x_dev, const int32_t* src_dev,
const std::vector<const uint8_t*>& blob, int64_t rows, cudaStream_t s) {
const size_t gub = mmq::matrix_bytes(Q2, 1280, N), db = mmq::matrix_bytes(Q2, N, FF);
const size_t gub = mmq::matrix_bytes(Q2, FF2, N), db = mmq::matrix_bytes(Q2, N, FF);
const int E = (int) r.cnt.size();
std::vector<int32_t> order;
for (int e = 0; e < E; ++e) if (r.cnt[(size_t) e] > 0) order.push_back(e);
@@ -166,11 +169,11 @@ void run_mmq(mmq::Context& ctx, MmqBufs& b, const Routing& r, const float* x_dev
ck(cudaMemsetAsync(b.grp_gu.as<uint8_t>() + (size_t) ngx * gub, 0, 4096, s), "tail");
ck(cudaMemsetAsync(b.grp_d.as<uint8_t>() + (size_t) ngx * db, 0, 4096, s), "tail");
mmq::Product gu;
gu.w = b.grp_gu.p; gu.type = Q2; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = gub;
gu.w = b.grp_gu.p; gu.type = Q2; gu.w_rows = FF2; gu.w_cols = N; gu.expert_bytes = gub;
gu.n = ngx; gu.xq = b.xq.p; gu.bounds = b.bounds.as<int32_t>() + j0; gu.ids = b.ident.as<int32_t>();
gu.total_rows = rows; gu.max_rows = maxr; gu.dst = b.gu.as<float>(); gu.ld_dst = 1280;
gu.total_rows = rows; gu.max_rows = maxr; gu.dst = b.gu.as<float>(); gu.ld_dst = FF2;
ctx.run(gu, s);
mmq::swiglu(b.gu.as<float>() + r0 * 1280, b.h.as<float>() + r0 * FF, nr, FF, true, s);
mmq::swiglu(b.gu.as<float>() + r0 * FF2, b.h.as<float>() + r0 * FF, nr, FF, true, s);
mmq::quantize(b.h.as<float>() + r0 * FF, nullptr, b.hq.p, Q2, FF, FF, nr, s);
mmq::Product dn;
dn.w = b.grp_d.p; dn.type = Q2; dn.w_rows = N; dn.w_cols = FF; dn.expert_bytes = db;