141 lines
5.5 KiB
C++
141 lines
5.5 KiB
C++
// test-moe-perf: microbench for CPU MoE expert GEMM (IQ4_XS, AVX2 path).
|
|
// Exercises the exact production path (GGML_OP_MUL_MAT_ID / fused up-gate)
|
|
// with qwen4exp-like shapes across batch sizes and thread counts.
|
|
// Reports ms/call + effective GB/s of ACTIVE expert weights.
|
|
//
|
|
// Usage: test-moe-perf [--quick] (quick trims shapes/threads for CI)
|
|
|
|
#include "ggml.h"
|
|
|
|
#include <cmath>
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <string>
|
|
#include <time.h>
|
|
#include <vector>
|
|
|
|
static uint64_t rng_state = 0x9e3779b97f4a7c15ULL;
|
|
static float frand(float lo, float hi) {
|
|
rng_state ^= rng_state << 13;
|
|
rng_state ^= rng_state >> 7;
|
|
rng_state ^= rng_state << 17;
|
|
const double u = (double)(rng_state >> 11) * (1.0 / 9007199254740992.0);
|
|
return (float)(lo + u * (hi - lo));
|
|
}
|
|
|
|
static double now_ms(void) {
|
|
struct timespec ts;
|
|
clock_gettime(CLOCK_MONOTONIC, &ts);
|
|
return ts.tv_sec * 1000.0 + ts.tv_nsec / 1e6;
|
|
}
|
|
|
|
struct Shape { int K; int N; int E; int used; };
|
|
|
|
static void bench_one(const Shape & sh, int Ny, int nth, bool fused, bool quick) {
|
|
const int64_t K = sh.K, N = sh.N, E = sh.E, U = sh.used;
|
|
const size_t wrow = ggml_row_size(GGML_TYPE_IQ4_XS, K);
|
|
const size_t mem = (wrow * N * E) * (fused ? 2 : 1)
|
|
+ (size_t)K * U * Ny * 4 + (size_t)N * U * Ny * 4 + (size_t)U * Ny * 4
|
|
+ 256 * 1024 * 1024;
|
|
struct ggml_init_params ip = { mem, nullptr, false };
|
|
struct ggml_context * ctx = ggml_init(ip);
|
|
if (!ctx) { printf("ggml_init failed\n"); return; }
|
|
|
|
// experts: [K, N, E] IQ4_XS from random f32 (values irrelevant for perf)
|
|
std::vector<float> f32((size_t)K * N);
|
|
for (auto & x : f32) x = frand(-1.0f, 1.0f);
|
|
struct ggml_tensor * w1 = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ4_XS, K, N, E);
|
|
ggml_set_name(w1, "w1");
|
|
{
|
|
std::vector<float> imatrix(K, 1.0f);
|
|
char * dst = (char *)w1->data;
|
|
for (int e = 0; e < E; ++e) {
|
|
for (int64_t r = 0; r < N; ++r) {
|
|
ggml_quantize_chunk(GGML_TYPE_IQ4_XS, f32.data() + r * K,
|
|
dst + ((size_t)e * N + r) * wrow, 0, 1, K, imatrix.data(), nullptr);
|
|
}
|
|
}
|
|
}
|
|
struct ggml_tensor * w2 = nullptr;
|
|
if (fused) {
|
|
w2 = ggml_new_tensor_3d(ctx, GGML_TYPE_IQ4_XS, K, N, E);
|
|
ggml_set_name(w2, "w2");
|
|
memcpy(w2->data, w1->data, ggml_nbytes(w2));
|
|
}
|
|
// activations F32 [K, U, Ny] (3D like production: one row per expert slot)
|
|
struct ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, K, U, Ny);
|
|
ggml_set_name(b, "b");
|
|
for (int64_t i = 0; i < (int64_t)K * U * Ny; ++i) ((float *)b->data)[i] = frand(-1.0f, 1.0f);
|
|
// ids [used, Ny]: random expert per slot
|
|
struct ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, U, Ny);
|
|
ggml_set_name(ids, "ids");
|
|
for (int64_t i = 0; i < (int64_t)U * Ny; ++i) ((int32_t *)ids->data)[i] = (int32_t)(frand(0, 0.999f) * E);
|
|
|
|
struct ggml_tensor * out;
|
|
if (fused) {
|
|
out = ggml_moe_up_gate(ctx, w1, w2, b, ids, GGML_UNARY_OP_SILU);
|
|
} else {
|
|
out = ggml_mul_mat_id(ctx, w1, b, ids);
|
|
}
|
|
struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16, false);
|
|
ggml_build_forward_expand(gf, out);
|
|
|
|
// warmup (also ramps CPU frequency)
|
|
for (int i = 0; i < 8; ++i) ggml_graph_compute_with_ctx(ctx, gf, nth);
|
|
// timed: min over reps (robust to frequency/noise spikes)
|
|
double t0 = now_ms();
|
|
ggml_graph_compute_with_ctx(ctx, gf, nth);
|
|
double dt = now_ms() - t0;
|
|
int reps = dt > 0 ? (int)(1000.0 / dt) : 5;
|
|
if (reps < 8) reps = 8;
|
|
if (reps > 40) reps = 40;
|
|
if (quick && reps > 10) reps = 10;
|
|
double best = 1e30;
|
|
for (int i = 0; i < reps; ++i) {
|
|
t0 = now_ms();
|
|
ggml_graph_compute_with_ctx(ctx, gf, nth);
|
|
dt = now_ms() - t0;
|
|
if (dt < best) best = dt;
|
|
}
|
|
dt = best;
|
|
|
|
// active expert bytes per call (up+gate for fused)
|
|
const double bpp = 4.25 / 8.0;
|
|
double bytes = (double)Ny * U * K * N * bpp * (fused ? 2 : 1);
|
|
printf("%s K=%-5d N=%-5d E=%-3d U=%-2d Ny=%-3d T=%-2d %8.3f ms/call %7.1f GB/s\n",
|
|
fused ? "fused" : "plain", sh.K, sh.N, sh.E, sh.used, Ny, nth, dt, bytes / dt / 1e6);
|
|
ggml_free(ctx);
|
|
}
|
|
|
|
int main(int argc, char ** argv) {
|
|
bool quick = false;
|
|
int force_threads = 0;
|
|
for (int i = 1; i < argc; ++i) {
|
|
if (!strcmp(argv[i], "--quick")) quick = true;
|
|
else if ((!strcmp(argv[i], "-t") || !strcmp(argv[i], "--threads")) && i + 1 < argc) force_threads = atoi(argv[++i]);
|
|
else { printf("usage: %s [--quick] [-t N]\n", argv[0]); return 1; }
|
|
}
|
|
ggml_quantize_init(GGML_TYPE_IQ4_XS);
|
|
// E=512 mirrors qwen4exp (expert weights >> L3, no false cache hits)
|
|
const std::vector<Shape> shapes = quick ? std::vector<Shape>{{2048, 2048, 128, 10}}
|
|
: std::vector<Shape>{{2048, 1024, 512, 10}, {2048, 2048, 512, 10}, {4096, 2048, 128, 10}};
|
|
const std::vector<int> nys = quick ? std::vector<int>{1, 8, 32} : std::vector<int>{1, 2, 4, 8, 16, 32, 64};
|
|
std::vector<int> nths = quick ? std::vector<int>{16} : std::vector<int>{4, 8, 12, 16};
|
|
if (force_threads > 0) nths = std::vector<int>{force_threads};
|
|
for (auto sh : shapes) {
|
|
for (int Ny : nys) {
|
|
for (int nth : nths) {
|
|
bench_one(sh, Ny, nth, false, quick);
|
|
}
|
|
}
|
|
if (!quick) {
|
|
for (int Ny : nys) bench_one(sh, Ny, 16, true, quick);
|
|
} else {
|
|
for (int Ny : nys) bench_one(sh, Ny, nths[0], true, quick);
|
|
}
|
|
}
|
|
ggml_quantize_free();
|
|
return 0;
|
|
}
|