diff --git a/zig/build/qwen_contract.zig b/zig/build/qwen_contract.zig index ba84730974..6bf2b4b859 100644 --- a/zig/build/qwen_contract.zig +++ b/zig/build/qwen_contract.zig @@ -14,7 +14,7 @@ pub fn host(b: *std.Build, core: *std.Build.Module, all: *std.Build.Step) void { pub fn images(b: *std.Build, runner: *std.Build.Step.Compile, controls: ?std.Build.LazyPath, h3: ?std.Build.LazyPath, kernels: anytype, files: []const ?std.Build.LazyPath) void { const install = @import("kernel_contract.zig").cuda(b, runner, controls, h3); for (kernels, files) |k, file| { - if (std.mem.eql(u8, k.name, "moe") or std.mem.eql(u8, k.name, "qwen36_mtp")) { + if (std.mem.eql(u8, k.name, "moe") or std.mem.eql(u8, k.name, "experts") or std.mem.eql(u8, k.name, "qwen36_mtp")) { if (file) |image| install.dependOn(&b.addInstallFile(image, b.fmt("fatbin/{s}.fatbin", .{k.name})).step); } } diff --git a/zig/exl3_prompt_host.zig b/zig/exl3_prompt_host.zig index fdfe6bcbcd..df7f7281eb 100644 --- a/zig/exl3_prompt_host.zig +++ b/zig/exl3_prompt_host.zig @@ -56,7 +56,7 @@ fn run(rows: usize, prompt: bool, shared: bool) !void { const s: cuda.Stream = .{ .d = &d, .handle = null }; var k: ex.Experts = .{ .mod = undefined, .group = function(&d, 1), .rot_in = function(&d, 2), .grouped = @splat(function(&d, 3)), .gateup = function(&d, 4), .down_combine = function(&d, 5) }; if (comptime @hasField(ex.Experts, "prompt")) k.prompt = .{ .rotate = function(&d, 11), .gateup = function(&d, 12), .down = function(&d, 13) }; - const ek: glm.experts.Kernels = .{ .nvfp4 = undefined, .router = .{ .plan_small = function(&d, 20), .plan_rank = function(&d, 21), .plan_offsets = function(&d, 22), .plan_scatter = function(&d, 23) } }; + const ek: glm.experts.Kernels = .{ .nvfp4 = undefined, .router = .{ .plan_small = function(&d, 20), .plan_rank = function(&d, 21), .plan_offsets = function(&d, 22), .plan_scatter = function(&d, 23), .plan_rows = function(&d, 24), .plan_rank_rows = function(&d, 25), .plan_scatter_rows = function(&d, 26) } }; const w: glm.exl3.Experts = .{ .gt = 100, .ut = 200, .dt = 300, .k2 = 400, .suh_g = 500, .suh_u = 600, .svh_g = 700, .svh_u = 800, .suh_d = 900, .svh_d = 1000, .shx = shared }; const sc: glm.exl3.Scratch = .{ .picks = 0x100000000, .xg = 0, .xu = 0, .z = 0 }; const plan: cuda.grouped.Plan = .{ .items = 10000, .members = 20000, .counts = 30000, .rank = 40000, .hist = 50000 }; diff --git a/zig/kernels/cuda/experts.cu b/zig/kernels/cuda/experts.cu index e41adf817e..ea6b1a3a73 100644 --- a/zig/kernels/cuda/experts.cu +++ b/zig/kernels/cuda/experts.cu @@ -228,8 +228,36 @@ __global__ void __launch_bounds__(PLAN_THREADS) } } +// plan_kernel over the first *rows rows' pairs (per a row): a captured bucket's padding rows get no items. __global__ void __launch_bounds__(PLAN_THREADS) - plan_rank(const int* __restrict__ picks, int P, int E, int* __restrict__ rank, int* __restrict__ hist) { + plan_rows(const int* __restrict__ picks, int P, int E, int T, int* __restrict__ members, + int* __restrict__ items, int* __restrict__ counts, const int* __restrict__ rows, int per) { + P = min(P, *rows * per); + __shared__ int pk[PSMALL]; + __shared__ int cnt[EMAX]; + __shared__ int off[EMAX]; + const int tid = threadIdx.x; + for (int e = tid; e < EMAX; e += PLAN_THREADS) cnt[e] = 0; + __syncthreads(); + for (int p = tid; p < P; p += PLAN_THREADS) { + const int e = picks[p]; + pk[p] = e; + atomicAdd(&cnt[e], 1); + } + __syncthreads(); + const int o = place_items(tid < E ? cnt[tid] : 0, E, T, items, counts); + if (tid < E) off[tid] = o; + __syncthreads(); + for (int p = tid; p < P; p += PLAN_THREADS) { + const int e = pk[p]; + int rank = 0; + for (int q = 0; q < p; ++q) rank += pk[q] == e; + members[off[e] + rank] = p; + } +} + +__device__ __forceinline__ void rank_pairs(const int* __restrict__ picks, int P, int E, int* __restrict__ rank, + int* __restrict__ hist) { __shared__ int cnt[EMAX]; const int tid = threadIdx.x, lane = tid & 31, warp = tid >> 5; for (int e = tid; e < E; e += PLAN_THREADS) cnt[e] = 0; @@ -251,6 +279,18 @@ __global__ void __launch_bounds__(PLAN_THREADS) for (int x = tid; x < E; x += PLAN_THREADS) hist[(size_t)blockIdx.x * E + x] = cnt[x]; } +__global__ void __launch_bounds__(PLAN_THREADS) + plan_rank(const int* __restrict__ picks, int P, int E, int* __restrict__ rank, int* __restrict__ hist) { + rank_pairs(picks, P, E, rank, hist); +} + +// plan_rank over the first *rows rows' pairs (per a row): a captured bucket's padding pairs count nowhere. +__global__ void __launch_bounds__(PLAN_THREADS) + plan_rank_rows(const int* __restrict__ picks, int P, int E, int* __restrict__ rank, int* __restrict__ hist, + const int* __restrict__ rows, int per) { + rank_pairs(picks, min(P, *rows * per), E, rank, hist); +} + __global__ void __launch_bounds__(PLAN_THREADS) plan_offsets(int nblk, int E, int T, int* __restrict__ hist, int* __restrict__ items, int* __restrict__ counts) { const int tid = threadIdx.x; @@ -273,6 +313,14 @@ __global__ void plan_scatter(const int* __restrict__ picks, int P, int E, const if (p < P) members[hist[(size_t)(p / PLAN_THREADS) * E + picks[p]] + rank[p]] = p; } +// plan_scatter over the first *rows rows' pairs (per a row). +__global__ void plan_scatter_rows(const int* __restrict__ picks, int P, int E, const int* __restrict__ rank, + const int* __restrict__ hist, int* __restrict__ members, const int* __restrict__ rows, + int per) { + const int p = blockIdx.x * blockDim.x + threadIdx.x; + if (p < min(P, *rows * per)) members[hist[(size_t)(p / PLAN_THREADS) * E + picks[p]] + rank[p]] = p; +} + } // namespace tf_experts // Decode form for groups of 64: up with relu^2 (epilogue 1), down to fp32 (epilogue 0). diff --git a/zig/kernels/cuda/moe.cu b/zig/kernels/cuda/moe.cu index 2adb06a220..89153a569c 100644 --- a/zig/kernels/cuda/moe.cu +++ b/zig/kernels/cuda/moe.cu @@ -14,10 +14,12 @@ constexpr int KC = 256; // inputs a shared-memory chunk holds } // namespace -// OUT[r][e] = fp32 x[r] . w[e] over D in index order (one fma chain a pair); grid (rows / 16, E / 16), block 256. +// OUT[r][e] = fp32 x[r] . w[e] in index order (an fma chain a pair), grid (R / 16, E / 16); skips rows >= P[2]. extern "C" __global__ void __launch_bounds__(256) tf_moe_router(const __nv_bfloat16* __restrict__ X, int x_stride, const __nv_bfloat16* __restrict__ W, float* __restrict__ OUT, - int R, int D, int E) { + int R, int D, int E, const int* __restrict__ P) { + if (P != nullptr) R = min(R, P[2]); + if (static_cast(blockIdx.x) * TILE >= R) return; __shared__ __nv_bfloat16 xs[TILE][KC + 2]; __shared__ __nv_bfloat16 ws[TILE][KC + 2]; const int row0 = blockIdx.x * TILE, e0 = blockIdx.y * TILE; @@ -47,10 +49,12 @@ __device__ __forceinline__ bool beats(float v, int i, float bv, int bi) { return } // namespace -// Row r by a warp: TOPK experts by logit (ties to the lower id), weights bf16(exp(l - top) / sum), the shared slot. +// Row r by a warp: TOPK experts by logit (ties to lower id), bf16(exp(l - top) / sum), shared slot; r < P[2]. extern "C" __global__ void __launch_bounds__(32) tf_moe_topk(const float* __restrict__ L, int* __restrict__ PICK, - float* __restrict__ WT, int NE, int TOPK) { + float* __restrict__ WT, int NE, int TOPK, + const int* __restrict__ P) { const int64_t r = blockIdx.x; + if (P != nullptr && r >= P[2]) return; const int lane = threadIdx.x; const int NL = NE + 1, SLOTS = TOPK + 1; // L[r][NE], the shared gate's logit, weighs slot TOPK by bf16(sigmoid(bf16(it))) float v[MAXC]; @@ -97,10 +101,12 @@ extern "C" __global__ void __launch_bounds__(32) tf_moe_topk(const float* __rest } } -// OUT[r][d] = bf16(sum over slots in order of fma(w[r][s], y[r][s][d])); y fp32 (decode) or bf16 (prompt rows). +// OUT[r][d] = bf16(sum over slots in order of fma(w[r][s], y[r][s][d])), y fp32 or bf16; rows >= P[2] skipped. extern "C" __global__ void __launch_bounds__(256) tf_moe_combine(const void* __restrict__ Y, int y_f32, const float* __restrict__ WT, - __nv_bfloat16* __restrict__ OUT, int S, int D) { + __nv_bfloat16* __restrict__ OUT, int S, int D, + const int* __restrict__ P) { const int64_t r = blockIdx.x; + if (P != nullptr && r >= P[2]) return; for (int d = blockIdx.y * 256 + threadIdx.x; d < D; d += gridDim.y * 256) { float acc = 0.0f; for (int s = 0; s < S; ++s) { diff --git a/zig/src/core/contract/qwen_served_records.zig b/zig/src/core/contract/qwen_served_records.zig index 9647f0d6a8..05de437fee 100644 --- a/zig/src/core/contract/qwen_served_records.zig +++ b/zig/src/core/contract/qwen_served_records.zig @@ -1,12 +1,15 @@ -//! The Qwen CUDA engine's MoE and MTP launches, read by kernel-contract-limits from every SM's image. +//! The Qwen CUDA engine's MoE, plan and MTP launches, read by kernel-contract-limits from every SM's image. const std = @import("std"); const registry = @import("registry.zig"); /// nvcc 13.0's usage, the compiler of the release fatbins. pub const records = [_]registry.CudaRecord{ - .{ .launch = .{ .name = "tf_moe_router", .symbol = "tf_moe_router", .image = "moe", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 40, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 } } }, + .{ .launch = .{ .name = "tf_moe_router", .symbol = "tf_moe_router", .image = "moe", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 }, .{ .regs = 48, .static = 16512, .bound = 256 } } }, .{ .launch = .{ .name = "tf_moe_topk", .symbol = "tf_moe_topk", .image = "moe", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 35, .static = 0, .bound = 32 }, .{ .regs = 35, .static = 0, .bound = 32 }, .{ .regs = 35, .static = 0, .bound = 32 }, .{ .regs = 34, .static = 0, .bound = 32 }, .{ .regs = 34, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "tf_moe_combine", .symbol = "tf_moe_combine", .image = "moe", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 40, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 }, .{ .regs = 38, .static = 0, .bound = 256 } } }, + .{ .launch = .{ .name = "experts plan_rows", .symbol = "_ZN10tf_experts9plan_rowsEPKiiiiPiS2_S2_S1_i", .image = "experts", .threads = 1024, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 12672, .bound = 1024 }, .{ .regs = 38, .static = 12672, .bound = 1024 }, .{ .regs = 38, .static = 12672, .bound = 1024 }, .{ .regs = 32, .static = 12672, .bound = 1024 }, .{ .regs = 32, .static = 12672, .bound = 1024 } } }, + .{ .launch = .{ .name = "experts plan_rank_rows", .symbol = "_ZN10tf_experts14plan_rank_rowsEPKiiiPiS2_S1_i", .image = "experts", .threads = 1024, .min_sm = 80 }, .usage = .{ .{ .regs = 32, .static = 4096, .bound = 1024 }, .{ .regs = 40, .static = 4096, .bound = 1024 }, .{ .regs = 40, .static = 4096, .bound = 1024 }, .{ .regs = 42, .static = 4096, .bound = 1024 }, .{ .regs = 42, .static = 4096, .bound = 1024 } } }, + .{ .launch = .{ .name = "experts plan_scatter_rows", .symbol = "_ZN10tf_experts17plan_scatter_rowsEPKiiiS1_S1_PiS1_i", .image = "experts", .threads = 256, .min_sm = 80 }, .usage = .{ .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 }, .{ .regs = 14, .static = 0, .bound = 0 } } }, .{ .launch = .{ .name = "tf_q36m_norm_rope", .symbol = "tf_q36m_norm_rope", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 36, .static = 0, .bound = 32 }, .{ .regs = 37, .static = 0, .bound = 32 }, .{ .regs = 39, .static = 0, .bound = 32 }, .{ .regs = 42, .static = 0, .bound = 32 }, .{ .regs = 42, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "tf_q36m_attn", .symbol = "tf_q36m_attn", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 50, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 51, .static = 0, .bound = 32 }, .{ .regs = 54, .static = 0, .bound = 32 }, .{ .regs = 54, .static = 0, .bound = 32 } } }, .{ .launch = .{ .name = "tf_q36m_attn_tree", .symbol = "tf_q36m_attn_tree", .image = "qwen36_mtp", .threads = 32, .min_sm = 80 }, .usage = .{ .{ .regs = 53, .static = 0, .bound = 32 }, .{ .regs = 54, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 57, .static = 0, .bound = 32 }, .{ .regs = 57, .static = 0, .bound = 32 } } }, @@ -17,7 +20,7 @@ pub const cuda_132 = [_]registry.Override{ .{ .name = "tf_q36m_attn", .usage = .{ .{ .regs = 50, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 51, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 }, .{ .regs = 52, .static = 0, .bound = 32 } } }, }; -test "the six served MoE and MTP launches are in the live registry and fit every SM by their records" { +test "the nine served MoE, plan and MTP launches are in the live registry and fit every SM by their records" { for (records) |record| { var found: usize = 0; for (registry.all_cuda) |active| if (std.mem.eql(u8, active.launch.symbol, record.launch.symbol)) { diff --git a/zig/src/cuda/grouped.zig b/zig/src/cuda/grouped.zig index 6b9cdb2b6c..08d95144c9 100644 --- a/zig/src/cuda/grouped.zig +++ b/zig/src/cuda/grouped.zig @@ -22,42 +22,68 @@ pub const Plan = struct { members: u64, items: u64, counts: u64, rank: u64, hist /// experts.cu's plan instances (cuobjdump -symbols of the experts fatbin). pub const symbols = struct { pub const plan = "_ZN10tf_experts11plan_kernelEPKiiiiPiS2_S2_"; + pub const plan_rows = "_ZN10tf_experts9plan_rowsEPKiiiiPiS2_S2_S1_i"; pub const rank = "_ZN10tf_experts9plan_rankEPKiiiPiS2_"; + pub const rank_rows = "_ZN10tf_experts14plan_rank_rowsEPKiiiPiS2_S1_i"; pub const offsets = "_ZN10tf_experts12plan_offsetsEiiiPiS0_S0_"; pub const scatter = "_ZN10tf_experts12plan_scatterEPKiiiS1_S1_Pi"; + pub const scatter_rows = "_ZN10tf_experts17plan_scatter_rowsEPKiiiS1_S1_PiS1_i"; }; pub const Router = struct { plan_small: Function, + plan_rows: Function, plan_rank: Function, plan_offsets: Function, plan_scatter: Function, + plan_rank_rows: Function, + plan_scatter_rows: Function, /// The plan kernels from a loaded experts module (kernels.experts). pub fn resolve(mod: Module) !Router { - return .{ .plan_small = try mod.function(symbols.plan), .plan_rank = try mod.function(symbols.rank), .plan_offsets = try mod.function(symbols.offsets), .plan_scatter = try mod.function(symbols.scatter) }; + return .{ .plan_small = try mod.function(symbols.plan), .plan_rows = try mod.function(symbols.plan_rows), .plan_rank = try mod.function(symbols.rank), .plan_offsets = try mod.function(symbols.offsets), .plan_scatter = try mod.function(symbols.scatter), .plan_rank_rows = try mod.function(symbols.rank_rows), .plan_scatter_rows = try mod.function(symbols.scatter_rows) }; } fn go(f: Function, s: Stream, grid: usize, block: u32, a: *launch_.Args) !void { try launch_.launch(f, .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = block } }, s, a); } + /// route over the first `rows.*` rows' pairs (`per` a row; `rows` a device int, 0: all), at any width. + pub fn routeRows(r: Router, s: Stream, picks: u64, pairs: usize, count: usize, tile: usize, p: Plan, rows: u64, per: usize) !void { + if (rows == 0) return r.route(s, picks, pairs, count, tile, p); + if (pairs > small) return r.wide(s, picks, pairs, count, tile, p, .{ .rows = rows, .per = per }); + var a: launch_.Args = .{}; + a.add(picks); + for ([_]usize{ pairs, count, tile }) |v| a.add(@as(c_int, @intCast(v))); + for ([_]u64{ p.members, p.items, p.counts, rows }) |v| a.add(v); + a.add(@as(c_int, @intCast(per))); + return go(r.plan_rows, s, 1, 1024, &a); + } + /// experts.route: `pairs` int32 expert ids at `picks` grouped into items of at most `tile` pairs (16 decode). pub fn route(r: Router, s: Stream, picks: u64, pairs: usize, count: usize, tile: usize, p: Plan) !void { + if (pairs > small) return r.wide(s, picks, pairs, count, tile, p, null); var a: launch_.Args = .{}; - if (pairs <= small) { - a.add(picks); - for ([_]usize{ pairs, count, tile }) |v| a.add(@as(c_int, @intCast(v))); - for ([_]u64{ p.members, p.items, p.counts }) |v| a.add(v); - return go(r.plan_small, s, 1, 1024, &a); - } + a.add(picks); + for ([_]usize{ pairs, count, tile }) |v| a.add(@as(c_int, @intCast(v))); + for ([_]u64{ p.members, p.items, p.counts }) |v| a.add(v); + return go(r.plan_small, s, 1, 1024, &a); + } + + /// The plan ranked in blocks of 1024 pairs; `live` (a device row count and pairs a row) caps the pairs planned. + fn wide(r: Router, s: Stream, picks: u64, pairs: usize, count: usize, tile: usize, p: Plan, live: ?struct { rows: u64, per: usize }) !void { const nblk = (pairs + 1023) / 1024; + var a: launch_.Args = .{}; a.add(picks); a.add(@as(c_int, @intCast(pairs))); a.add(@as(c_int, @intCast(count))); a.add(p.rank); a.add(p.hist); - try go(r.plan_rank, s, nblk, 1024, &a); + if (live) |l| { + a.add(l.rows); + a.add(@as(c_int, @intCast(l.per))); + } + try go(if (live != null) r.plan_rank_rows else r.plan_rank, s, nblk, 1024, &a); var b: launch_.Args = .{}; for ([_]usize{ nblk, count, tile }) |v| b.add(@as(c_int, @intCast(v))); for ([_]u64{ p.hist, p.items, p.counts }) |v| b.add(v); @@ -67,7 +93,11 @@ pub const Router = struct { c.add(@as(c_int, @intCast(pairs))); c.add(@as(c_int, @intCast(count))); for ([_]u64{ p.rank, p.hist, p.members }) |v| c.add(v); - try go(r.plan_scatter, s, (pairs + 255) / 256, 256, &c); + if (live) |l| { + c.add(l.rows); + c.add(@as(c_int, @intCast(l.per))); + } + try go(if (live != null) r.plan_scatter_rows else r.plan_scatter, s, (pairs + 255) / 256, 256, &c); } }; diff --git a/zig/src/cuda/moe.zig b/zig/src/cuda/moe.zig index d293ae770d..7ecb48b45b 100644 --- a/zig/src/cuda/moe.zig +++ b/zig/src/cuda/moe.zig @@ -22,20 +22,21 @@ pub const Fns = struct { return .{ .router_fn = try mod.function("tf_moe_router"), .topk_fn = try mod.function("tf_moe_topk"), .combine_fn = try mod.function("tf_moe_combine") }; } - /// logits[r][e] = fp32 x[r] . w[e]: `rows` bf16 rows `x_stride` apart against `count` rows of `dims`, gate last. - pub fn route(f: Fns, s: Stream, x: u64, x_stride: usize, rows: usize, w: u64, logits: u64, dims: usize, count: usize) !void { + /// logits[r][e] = fp32 x[r] . w[e], `count` rows of `dims` with the gate last; `params` word 2 (if set) caps rows. + pub fn route(f: Fns, s: Stream, x: u64, x_stride: usize, rows: usize, w: u64, logits: u64, dims: usize, count: usize, params: u64) !void { var a: launch_.Args = .{}; a.add(x); a.add(int(x_stride)); a.add(w); a.add(logits); for ([_]usize{ rows, dims, count }) |v| a.add(int(v)); + a.add(params); const cfg: launch_.Config = .{ .grid = .{ .x = @intCast((rows + 15) / 16), .y = @intCast((count + 15) / 16), .z = 1 }, .block = .{ .x = 256 } }; try launch_.launch(f.router_fn, cfg, s, &a); } /// Each row's top-k picks and bf16-rounded weights ([rows][top_k + 1]); slot top_k is the shared expert. - pub fn pick(f: Fns, s: Stream, logits: u64, picks: u64, wts: u64, rows: usize, experts: usize, top_k: usize) !void { + pub fn pick(f: Fns, s: Stream, logits: u64, picks: u64, wts: u64, rows: usize, experts: usize, top_k: usize, params: u64) !void { if (experts > max_experts or top_k > max_top_k) return error.RoutingTooWide; var a: launch_.Args = .{}; a.add(logits); @@ -43,11 +44,12 @@ pub const Fns = struct { a.add(wts); a.add(int(experts)); a.add(int(top_k)); + a.add(params); try launch_.launch(f.topk_fn, .{ .grid = .{ .x = @intCast(rows), .y = 1, .z = 1 }, .block = .{ .x = 32 } }, s, &a); } /// out[r] = bf16(sum of w[r][s] * y[r][s]) over the slots in order; y fp32 (decode rows) or bf16 (prompt rows). - pub fn combine(f: Fns, s: Stream, y: u64, y_f32: bool, wts: u64, out: u64, rows: usize, slots: usize, dims: usize) !void { + pub fn combine(f: Fns, s: Stream, y: u64, y_f32: bool, wts: u64, out: u64, rows: usize, slots: usize, dims: usize, params: u64) !void { var a: launch_.Args = .{}; a.add(y); a.add(@as(c_int, @intFromBool(y_f32))); @@ -55,6 +57,7 @@ pub const Fns = struct { a.add(out); a.add(int(slots)); a.add(int(dims)); + a.add(params); const blocks = @min((dims + 255) / 256, 16); try launch_.launch(f.combine_fn, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(blocks), .z = 1 }, .block = .{ .x = 256 } }, s, &a); } diff --git a/zig/src/families/qwen3_5/cuda_engine.zig b/zig/src/families/qwen3_5/cuda_engine.zig index ebe6a4d056..b63b8d9125 100644 --- a/zig/src/families/qwen3_5/cuda_engine.zig +++ b/zig/src/families/qwen3_5/cuda_engine.zig @@ -557,7 +557,7 @@ pub const Engine = struct { try e.run(e.fns.mlp_act, .{ (total * g.intermediate + 255) / 256, 1, 1 }, .{ 256, 1, 1 }, .{ e.gate, e.up, e.act, n, t_params, @as(c_int, @intCast(g.intermediate)) }); try e.proj(m.down, e.act, e.r, total, prompt); }, - .moe => |m| try e.routedRows(m.router, m.experts, e.x, e.r, total, prompt and total > most_window), + .moe => |m| try e.routedRows(m.router, m.experts, e.x, e.r, total, prompt and total > most_window, t_params), } if (shape.taps) for (e.taps[0..e.n_taps], 0..) |t, ti| if (t == i) { try e.run(e.fns.tap, .{ total, 1, 1 }, .{ 256, 1, 1 }, .{ e.h, e.r, t_segs, t_rows, @as(c_int, @intCast(ti)), @as(c_int, @intCast(g.hidden)), t_params }); diff --git a/zig/src/families/qwen3_5/cuda_mtp.zig b/zig/src/families/qwen3_5/cuda_mtp.zig index 2e330be4ec..c1565e2610 100644 --- a/zig/src/families/qwen3_5/cuda_mtp.zig +++ b/zig/src/families/qwen3_5/cuda_mtp.zig @@ -287,7 +287,7 @@ pub const Mtp = struct { } try d.proj(w.o, d.att, d.r, rows); try d.norm(d.h, d.r, w.post_norm, d.x, rows); - try e.routedRows(w.mlp.moe.router, w.mlp.moe.experts, d.x, d.r, rows, rows > most_window); + try e.routedRows(w.mlp.moe.router, w.mlp.moe.experts, d.x, d.r, rows, rows > most_window, 0); try d.norm(d.h, d.r, w.norm, d.x, rows); if (pass == .ring) try e.run(d.ring_fns.ring_put, .{ rows, 1, 1 }, .{ 256, 1, 1 }, .{ d.kb, d.vb, d.ringK(ctx), d.ringV(ctx), @as(c_int, @intCast(kw)), p0, cap }); } diff --git a/zig/src/families/qwen3_5/cuda_routed.zig b/zig/src/families/qwen3_5/cuda_routed.zig index 89788cfe26..608a30467d 100644 --- a/zig/src/families/qwen3_5/cuda_routed.zig +++ b/zig/src/families/qwen3_5/cuda_routed.zig @@ -34,8 +34,8 @@ pub const Scratch = struct { } }; -/// Routed experts over `total` normed rows at x into out; `wide` rows take the prefill kernels, else decode's. -pub fn rows(e: *E.Engine, router: u64, layer: cuda.experts.Layer, x: u64, out: u64, total: usize, wide: bool) !void { +/// Routed experts over `total` normed rows at x into out (wide: prefill kernels); `params` word 2 caps the rows. +pub fn rows(e: *E.Engine, router: u64, layer: cuda.experts.Layer, x: u64, out: u64, total: usize, wide: bool, params: u64) !void { const g = e.g; const s = e.routed; const slots = g.slots(); @@ -45,9 +45,9 @@ pub fn rows(e: *E.Engine, router: u64, layer: cuda.experts.Layer, x: u64, out: u const items = cuda.grouped.maxItems(pairs, count, tile); const plan: cuda.grouped.Plan = .{ .members = s.members, .items = s.items, .counts = s.counts, .rank = s.rank, .hist = s.hist }; const grouped: cuda.experts.Grouped = .{ .affine4 = &e.experts, .affine3 = &e.experts3 }; - try e.moe.route(e.stream, x, g.hidden, total, router, s.logits, g.hidden, count); - try e.moe.pick(e.stream, s.logits, s.picks, s.wts, total, g.experts, g.top_k); - try e.plan.route(e.stream, s.picks, pairs, count, tile, plan); + try e.moe.route(e.stream, x, g.hidden, total, router, s.logits, g.hidden, count, params); + try e.moe.pick(e.stream, s.logits, s.picks, s.wts, total, g.experts, g.top_k, params); + try e.plan.routeRows(e.stream, s.picks, pairs, count, tile, plan, if (params == 0) 0 else params + 2 * 4, slots); const rows_in: cuda.experts.Rows = .{ .x = x, .stride = g.hidden, .slots = slots }; const act_in: cuda.experts.Rows = .{ .x = s.act, .stride = g.intermediate }; if (wide) { @@ -57,5 +57,5 @@ pub fn rows(e: *E.Engine, router: u64, layer: cuda.experts.Layer, x: u64, out: u try grouped.decode(e.stream, .up, rows_in, layer, plan, s.act, items, -1); try grouped.decode(e.stream, .down_f32, act_in, layer, plan, s.y, items, -1); } - try e.moe.combine(e.stream, s.y, !wide, s.wts, out, total, slots, g.hidden); + try e.moe.combine(e.stream, s.y, !wide, s.wts, out, total, slots, g.hidden, params); } diff --git a/zig/tests/cuda/contract_qwen_served.zig b/zig/tests/cuda/contract_qwen_served.zig index 78271fc476..470ec4df56 100644 --- a/zig/tests/cuda/contract_qwen_served.zig +++ b/zig/tests/cuda/contract_qwen_served.zig @@ -35,9 +35,7 @@ pub fn run(gpu: check.Gpu) !void { var first: ?[]const u8 = null; for (0..5) |_| { switch (mode) { - .router => try fns.route(s, pointers[0], 2048 + @as(usize, if (variant) 16 else 0), rows, pointers[1], pointers[2], 2048, 257), - .topk => try fns.pick(s, pointers[0], pointers[1], pointers[2], rows, 256, 8), - .combine => try fns.combine(s, pointers[0], variant, pointers[1], pointers[2], rows, 9, 2048), + .router, .topk, .combine => try served(fns, s, mode, pointers, rows, variant, 0), else => try cuda.launch.launch(f, .{ .grid = .{ .x = fixture.grid[0], .y = fixture.grid[1], .z = fixture.grid[2] }, .block = .{ .x = fixture.threads } }, s, &args), } try s.synchronize(); @@ -53,7 +51,77 @@ pub fn run(gpu: check.Gpu) !void { if (first) |bytes| try check.expect(std.mem.eql(u8, bytes, output), "{s} nondeterministic bytes", .{ref.symbol(mode)}) else first = try a.dupe(u8, output); try check.expect(try rig.intact() and !rig.overlap() and std.mem.eql(u8, before, try rig.inputBytes()), "{s} changed an input or guard", .{ref.symbol(mode)}); } + if (@backingInt(mode) < 3) for ([_]usize{ 1, rows / 2 + 1, rows - 1 }) |live| if (live < rows) try liveRows(gpu, s, fns, mode, fixture, rows, variant, live); }; std.debug.print("PASS {s}: served geometry, FP64, five repeats, guards, inputs; finite corruption RED\n", .{ref.symbol(mode)}); } + try planRows(gpu, s); +} + +fn served(fns: cuda.moe.Fns, s: cuda.Stream, mode: ref.Mode, p: []const u64, rows: usize, variant: bool, params: u64) !void { + switch (mode) { + .router => try fns.route(s, p[0], 2048 + @as(usize, if (variant) 16 else 0), rows, p[1], p[2], 2048, 257, params), + .topk => try fns.pick(s, p[0], p[1], p[2], rows, 256, 8, params), + .combine => try fns.combine(s, p[0], variant, p[1], p[2], rows, 9, 2048, params), + else => unreachable, + } +} + +/// A bucket of `rows` with params word 2 at `live`: the live rows meet FP64, the padding rows' outputs stay unwritten. +fn liveRows(gpu: check.Gpu, s: cuda.Stream, fns: cuda.moe.Fns, mode: ref.Mode, fixture: ref.Fixture, rows: usize, variant: bool, live: usize) !void { + var rig = rig_.Rig{ .gpu = gpu, .s = s }; + defer rig.deinit(); + var pointers: [3]u64 = undefined; + for (fixture.buffers.items, &pointers) |b, *p| p.* = if (b.output) try rig.output(b.bytes.len) else try rig.input(b.bytes); + const params = [8]i32{ 0, 0, @intCast(live), 1, 1, 0, 0, 0 }; + try served(fns, s, mode, &pointers, rows, variant, try rig.input(std.mem.sliceAsBytes(¶ms))); + try s.synchronize(); + for (fixture.results.items) |result| { + const bytes = try rig.bytes(pointers[result.buffer], fixture.buffers.items[result.buffer].bytes.len); + const per = result.want.len / rows; + const width: usize = if (result.elem == .bf16) 2 else 4; + const head: ref.Result = .{ .buffer = result.buffer, .elem = result.elem, .want = result.want[0 .. live * per], .bound = result.bound[0 .. live * per] }; + try check.expect(head.accepts(bytes[0 .. live * per * width]), "{s} FP64 mismatch at {d} live rows of {d}", .{ ref.symbol(mode), live, rows }); + for (bytes[live * per * width ..]) |byte| try check.expect(byte == 0xff, "{s} wrote a padding row ({d} live of {d})", .{ ref.symbol(mode), live, rows }); + } + try check.expect(try rig.intact() and !rig.overlap(), "{s} broke a guard at {d} live rows", .{ ref.symbol(mode), live }); +} + +/// routeRows at n live rows equals route over those rows' pairs byte for byte, padding pairs never planned; buckets +/// past 113 rows (1,024 pairs) take the wide plan. +fn planRows(gpu: check.Gpu, s: cuda.Stream) !void { + var mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + defer mod.unload(); + const router = try cuda.grouped.Router.resolve(mod); + const per = 9; + const count = 257; + for ([_]usize{ 1, 16, 32, 64, 114, 128 }) |bucket| for (1..bucket + 1) |live| { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const pairs = bucket * per; + const picks = try a.alloc(i32, pairs); + for (picks, 0..) |*v, i| v.* = if (i % per == per - 1) 256 else @intCast((i * 37 + i / per * 11) % 256); + const params = [8]i32{ 0, 0, @intCast(live), 1, 1, 0, 0, 0 }; + const items = 3 * cuda.grouped.maxItems(pairs, count, 16) * 4; + var outs: [2][]const u8 = undefined; + for (0..2) |k| { + var rig = rig_.Rig{ .gpu = gpu, .s = s }; + defer rig.deinit(); + const dp = try rig.input(std.mem.sliceAsBytes(picks)); + const dlive = try rig.input(std.mem.sliceAsBytes(¶ms)) + 8; + const plan: cuda.grouped.Plan = .{ .members = try rig.output(pairs * 4), .items = try rig.output(items), .counts = try rig.output(8), .rank = try rig.scratch(pairs * 4), .hist = try rig.scratch((pairs + 1023) / 1024 * count * 4) }; + var first: ?[]const u8 = null; + for (0..5) |_| { + if (k == 0) try router.routeRows(s, dp, pairs, count, 16, plan, dlive, per) else try router.route(s, dp, live * per, count, 16, plan); + try s.synchronize(); + const out = try rig.outputBytes(); + if (first) |b| try check.expect(std.mem.eql(u8, b, out), "routeRows nondeterministic bytes", .{}) else first = try a.dupe(u8, out); + } + try check.expect(try rig.intact() and !rig.overlap(), "routeRows broke a guard", .{}); + outs[k] = first.?; + } + try check.expect(std.mem.eql(u8, outs[0], outs[1]), "routeRows at {d} live rows of {d} differs from route over their pairs", .{ live, bucket }); + }; + std.debug.print("PASS plan_rows: buckets 1, 16, 32, 64, 114, 128 at every live count (wide past 1,024 pairs) equal the plan over the live pairs, five repeats, guards\n", .{}); } diff --git a/zig/tests/cuda/contract_tests.zig b/zig/tests/cuda/contract_tests.zig index d5d52993cc..f0caee2116 100644 --- a/zig/tests/cuda/contract_tests.zig +++ b/zig/tests/cuda/contract_tests.zig @@ -41,7 +41,7 @@ const h3_image = struct { /// The embedded image a registration names (the controls' only when this binary was built with -Dnvcc). fn image(name: []const u8) []const u8 { if (std.mem.eql(u8, name, "h3")) return &h3_image.bytes; - inline for (.{ "norms", "moe", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "lane_gemv", "prompt16", "qmmf", "ck_experts", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "nemotron_mamba", "sample", "qwen35", "qwen35_keep", "qwen35_kv8", "qwen35_draft", "ck_lane" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); + inline for (.{ "norms", "moe", "experts", "qwen36_mtp", "nemotron_attention", "nemotron_route", "nemotron_latent", "qmm_group", "lane_gemv", "prompt16", "qmmf", "ck_experts", "rank_sum", "prefill_attention", "vit", "embed_rows", "route_rows", "flashnext_gdn", "flashnext_batch", "flashnext_hc", "flashnext_attn", "flashnext_mtp", "qwen_vision", "exl3_linear", "exl3_experts", "exl3_candidate", "experts3", "experts3_prefill", "drama3", "drama3_fp8", "drama3_codec", "glm_mm", "glm_dflash", "nemotron_mamba", "sample", "qwen35", "qwen35_keep", "qwen35_kv8", "qwen35_draft", "ck_lane" }) |n| if (std.mem.eql(u8, name, n)) return @field(kk, n); return if (options.controls) &controls.bytes else &.{}; } diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index 6f984ad7ad..f9d3a293de 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -147,7 +147,8 @@ const usage = \\ affine8 MLX affine 2 to 8-bit tiles and the 8-bit lane_gemv: widened codes, float64, rows (synthetic) \\ carveout [MiB] [card] GB10 display memory: round trips and bandwidth (SKIP without the card) \\ nvfp4 NVFP4 projections against recorded lane-matmul and prompt-GEMM fixture bytes - \\ moe softmax routing and SwiGLU experts against host references + \\ moe softmax routing and SwiGLU experts against host references, then moe-live + \\ moe-live a captured bucket's padding rows: never routed, planned or combined (synthetic) \\ experts3 affine-3 packing and outputs against affine-4 and fp64 \\ ggml Bolt's IQ2_XXS experts and Q8_0 / IQ2_XXS dense rows on real tensors against fp64 (tools/zig/bolt_ggml_fixture.py) \\ ggml-bench IQ2_XXS expert and Q8_0 lm_head decode GB/s (synthetic blocks) @@ -262,6 +263,7 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "glm-vision")) return glm_vision_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "moe")) return moe_tests.run(gpu); + if (std.mem.eql(u8, cmd, "moe-live")) return @import("moe_live_tests.zig").run(gpu); if (std.mem.eql(u8, cmd, "experts3")) return experts3_tests.run(gpu); if (std.mem.eql(u8, cmd, "ggml")) return ggml_tests.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "ggml-bench")) return ggml_tests.bench(gpu); diff --git a/zig/tests/cuda/moe_live_tests.zig b/zig/tests/cuda/moe_live_tests.zig new file mode 100644 index 0000000000..df79a0aea3 --- /dev/null +++ b/zig/tests/cuda/moe_live_tests.zig @@ -0,0 +1,219 @@ +//! The routed chain in a captured bucket: rows past the params' word 2 are padding, never routed, planned or combined. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const moe = @import("moe_tests.zig"); +const Gpu = check.Gpu; + +const routed = 64; +const count = routed + 1; +const top_k = 8; +const slots = top_k + 1; +const dims = 256; +const ni = 128; +const most = 128; +const unset: u8 = 0xab; +const stale_pick = 5; // an expert a padding row's stale pick names: planned, it would add that expert's pairs + +/// The chain's kernels, weights and scratch for buckets of up to `most` rows. +const Chain = struct { + gpu: Gpu, + s: cuda.Stream, + f: cuda.moe.Fns, + router: cuda.grouped.Router, + grouped: cuda.experts.Grouped, + layer: cuda.experts.Layer, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + w: u64 = 0, + logits: u64 = 0, + picks: u64 = 0, + wts: u64 = 0, + plan: cuda.grouped.Plan = undefined, + act: u64 = 0, + y: u64 = 0, + + fn buf(c: *Chain, len: usize) !u64 { + var b = try cuda.DeviceBuffer.alloc(c.gpu.d, len); + try b.fill8(0, null); + try c.bufs.append(c.gpu.gpa, b); + return b.ptr; + } + + fn put(c: *Chain, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(c.gpu.d, bytes); + try c.bufs.append(c.gpu.gpa, b); + return b.ptr; + } + + fn get(c: *Chain, ptr: u64, out: []u8) !void { + try c.s.synchronize(); + try c.gpu.d.check(c.gpu.d.api.cuMemcpyDtoH_v2(out.ptr, ptr, out.len), "cuMemcpyDtoH"); + } + + fn set(c: *Chain, ptr: u64, bytes: []const u8) !void { + try c.s.synchronize(); + try c.gpu.d.check(c.gpu.d.api.cuMemcpyHtoD_v2(ptr, bytes.ptr, bytes.len), "cuMemcpyHtoD"); + } + + fn fill(c: *Chain, ptr: u64, byte: u8, len: usize) !void { + try c.s.synchronize(); + try c.gpu.d.check(c.gpu.d.api.cuMemsetD8_v2(ptr, byte, len), "cuMemsetD8"); + } + + /// True when every byte from ptr + `from` to ptr + `end` still holds `byte`. + fn still(c: *Chain, a: std.mem.Allocator, ptr: u64, from: usize, end: usize, byte: u8) !bool { + const host = try a.alloc(u8, end - from); + defer a.free(host); + try c.get(ptr + from, host); + return std.mem.allEqual(u8, host, byte); + } + + /// The engine's decode chain over a bucket of `rows` rows at x into out (cuda_routed.rows, wide off). + fn run(c: *Chain, x: u64, out: u64, rows: usize, params: u64) !void { + const pairs = rows * slots; + const items = cuda.grouped.maxItems(pairs, count, 16); + try c.f.route(c.s, x, dims, rows, c.w, c.logits, dims, count, params); + try c.f.pick(c.s, c.logits, c.picks, c.wts, rows, routed, top_k, params); + try c.router.routeRows(c.s, c.picks, pairs, count, 16, c.plan, if (params == 0) 0 else params + 8, slots); + try c.grouped.decode(c.s, .up, .{ .x = x, .stride = dims, .slots = slots }, c.layer, c.plan, c.act, items, -1); + try c.grouped.decode(c.s, .down_f32, .{ .x = c.act, .stride = ni }, c.layer, c.plan, c.y, items, -1); + try c.f.combine(c.s, c.y, true, c.wts, out, rows, slots, dims, params); + } + + /// The pairs the last plan grouped: the items' counts summed, and the largest member. + fn planned(c: *Chain, a: std.mem.Allocator) !struct { pairs: usize, top: i32 } { + var counts: [2]i32 = undefined; + try c.get(c.plan.counts, std.mem.sliceAsBytes(&counts)); + const items = try a.alloc(i32, 3 * @as(usize, @intCast(counts[0]))); + defer a.free(items); + try c.get(c.plan.items, std.mem.sliceAsBytes(items)); + var pairs: usize = 0; + for (0..items.len / 3) |i| pairs += @intCast(items[3 * i + 2]); + const members = try a.alloc(i32, pairs); + defer a.free(members); + try c.get(c.plan.members, std.mem.sliceAsBytes(members)); + return .{ .pairs = pairs, .top = if (pairs > 0) std.mem.max(i32, members) else -1 }; + } +}; + +pub fn run(gpu: Gpu) !void { + const a = gpu.gpa; + var prng = std.Random.DefaultPrng.init(0x11fe); + const r = prng.random(); + var mods: [4]cuda.Module = undefined; + for (&mods, [_][]const u8{ cuda.kernels.moe, cuda.kernels.experts, cuda.kernels.experts_prefill, cuda.kernels.experts_pack }) |*m, img| m.* = try cuda.Module.load(gpu.d, img); + defer for (&mods) |*m| m.unload(); + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + const ex = try cuda.experts.Affine4Experts.resolve(mods[1], mods[2], mods[3], sms); + var c: Chain = .{ + .gpu = gpu, + .s = try cuda.Stream.init(gpu.d, true), + .f = try cuda.moe.Fns.resolve(mods[0]), + .router = try cuda.grouped.Router.resolve(mods[1]), + .grouped = .{ .affine4 = &ex }, + .layer = undefined, + }; + defer c.s.deinit(); + defer { + for (c.bufs.items) |*b| b.free(); + c.bufs.deinit(a); + } + // a bf16 router (routed experts, then the shared gate), affine-4 SwiGLU experts with the shared one last + const w = try a.alloc(u16, count * dims); + defer a.free(w); + for (w) |*v| v.* = moe.randBf16(r, 0.1); + c.w = try c.put(std.mem.sliceAsBytes(w)); + const up = try c.buf(cuda.experts.Affine4Experts.pairBytes(count, ni, dims)); + const down = try c.buf(cuda.experts.Affine4Experts.bytes(count, dims, ni)); + for ([_]struct { u64, usize, usize, usize, usize }{ .{ up, ni, dims, 2, 0 }, .{ up, ni, dims, 2, 1 }, .{ down, dims, ni, 1, 0 } }) |t| { + const tb = try moe.table(a, r, count, t[1], t[2]); + defer { + a.free(tb.words); + a.free(tb.scales); + a.free(tb.biases); + } + try ex.packInto(c.s, try c.put(std.mem.sliceAsBytes(tb.words)), try c.put(std.mem.sliceAsBytes(tb.scales)), try c.put(std.mem.sliceAsBytes(tb.biases)), t[0], count, t[1], t[2], t[3], t[4]); + try c.s.synchronize(); + } + c.layer = .{ .format = .affine4, .up = up, .down = down, .width = ni, .dims = dims, .experts = count, .swiglu = true }; + const pairs = most * slots; + c.logits = try c.buf(most * count * 4); + c.picks = try c.buf(pairs * 4); + c.wts = try c.buf(pairs * 4); + c.plan = .{ .members = try c.buf(pairs * 4), .items = try c.buf(3 * cuda.grouped.maxItems(pairs, count, 16) * 4), .counts = try c.buf(8), .rank = try c.buf(pairs * 4), .hist = try c.buf((pairs + 1023) / 1024 * count * 4) }; + c.act = try c.buf(pairs * ni * 2); + c.y = try c.buf(pairs * dims * 4); + // rows of normed activations; every row is non-zero, so a padding row routed would plan its own experts + const x = try a.alloc(u16, most * dims); + defer a.free(x); + for (x) |*v| v.* = moe.randBf16(r, 1.0); + const dx = try c.put(std.mem.sliceAsBytes(x)); + // each row alone, uncaptured and without params: the bits every live row of a bucket must equal + const want = try a.alloc(u8, most * dims * 2); + defer a.free(want); + const one = try c.buf(dims * 2); + for (0..most) |row| { + try c.run(dx + row * dims * 2, one, 1, 0); + try c.get(one, want[row * dims * 2 ..][0 .. dims * 2]); + } + const out = try c.buf(most * dims * 2); + const got = try a.alloc(u8, most * dims * 2); + defer a.free(got); + const poisoned = try a.alloc(u16, most * dims); + defer a.free(poisoned); + const stale = try a.alloc(i32, pairs); + defer a.free(stale); + @memset(stale, stale_pick); + var params = [8]i32{ 0, 0, 0, 0, 0, 0, 0, 0 }; + const dparams = try c.put(std.mem.sliceAsBytes(¶ms)); + // each bucket captured once, then replayed at live counts in this order (the count is read at each replay); + // nine pairs a row, so buckets past 113 rows take the wide plan, and 114 live rows are its first wide count + const sweeps = [_]struct { bucket: usize, live: []const usize }{ + .{ .bucket = 1, .live = &.{1} }, + .{ .bucket = 16, .live = &.{ 1, 2, 3, 7, 15, 16, 3, 14, 2 } }, + .{ .bucket = 32, .live = &.{ 17, 18, 31, 32, 3, 30, 17 } }, + .{ .bucket = 64, .live = &.{ 33, 1, 64, 40, 63, 9 } }, + .{ .bucket = 128, .live = &.{ 1, 9, 33, 64, 127, 128, 113, 114, 2, 120, 64 } }, + }; + for (sweeps) |sw| { + // capture over clean rows, so a replay's padding is whatever the replay leaves there + try c.set(dx, std.mem.sliceAsBytes(x)); + try cuda.graph.beginCapture(c.s, .thread_local); + c.run(dx, out, sw.bucket, dparams) catch |err| { + var broken = cuda.graph.endCapture(c.s) catch return err; + broken.deinit(); + return err; + }; + var graph = try cuda.graph.endCapture(c.s); + defer graph.deinit(); + var exec = try graph.instantiate(); + defer exec.deinit(); + for (sw.live) |live| { + // poisoned padding: rows from `live` up NaN, +huge and -huge in turn; stale picks name a real expert, + // stale weights and logits are NaN, and the experts' pair outputs hold a sentinel + @memcpy(poisoned, x); + for (live..most) |row| @memset(poisoned[row * dims ..][0..dims], switch (row % 3) { + 0 => 0x7fc0, + 1 => 0x7f7f, + else => 0xff7f, + }); + try c.set(dx, std.mem.sliceAsBytes(poisoned)); + try c.set(c.picks, std.mem.sliceAsBytes(stale)); + try c.fill(c.wts, 0xff, pairs * 4); + try c.fill(c.logits, 0xff, most * count * 4); + try c.fill(c.act, unset, pairs * ni * 2); + try c.fill(c.y, unset, pairs * dims * 4); + try c.fill(out, unset, most * dims * 2); + params[2] = @intCast(live); + try c.set(dparams, std.mem.sliceAsBytes(¶ms)); + try exec.launchOn(c.s); + const p = try c.planned(a); + try check.expect(p.pairs == live * slots and p.top < live * slots, "bucket {d} at {d} live rows planned {d} pairs (largest member {d}), want the live rows' {d}", .{ sw.bucket, live, p.pairs, p.top, live * slots }); + try check.expect(try c.still(a, c.act, live * slots * ni * 2, sw.bucket * slots * ni * 2, unset) and try c.still(a, c.y, live * slots * dims * 4, sw.bucket * slots * dims * 4, unset), "bucket {d} at {d} live rows: an expert wrote a padding pair", .{ sw.bucket, live }); + try c.get(out, got); + try check.sameBytes("live rows against one-row runs", got[0 .. live * dims * 2], want[0 .. live * dims * 2]); + for (got[live * dims * 2 ..]) |byte| try check.expect(byte == unset, "bucket {d} at {d} live rows wrote a padding row", .{ sw.bucket, live }); + } + check.pass("moe live rows: bucket {d} replayed at {any} live rows, padding poisoned, plans only their pairs, experts never see padding, equals one-row runs, skips padding", .{ sw.bucket, sw.live }); + } +} diff --git a/zig/tests/cuda/moe_tests.zig b/zig/tests/cuda/moe_tests.zig index 2bfd33a415..7b859fd2fe 100644 --- a/zig/tests/cuda/moe_tests.zig +++ b/zig/tests/cuda/moe_tests.zig @@ -17,7 +17,7 @@ fn rbf(x: f32) f32 { return bf16f(f32bf16(x)); } -fn randBf16(r: std.Random, scale: f32) u16 { +pub fn randBf16(r: std.Random, scale: f32) u16 { return f32bf16((r.float(f32) * 2.0 - 1.0) * scale); } @@ -50,7 +50,7 @@ fn router(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { defer dw.free(); var dl = try cuda.DeviceBuffer.alloc(gpu.d, rows * count * 4); defer dl.free(); - try f.route(s, dx.ptr, dims, rows, dw.ptr, dl.ptr, dims, count); + try f.route(s, dx.ptr, dims, rows, dw.ptr, dl.ptr, dims, count, 0); try s.synchronize(); const got = try check.download(gpu, dl); defer a.free(got); @@ -107,7 +107,7 @@ fn topk(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { defer dp.free(); var dw = try cuda.DeviceBuffer.alloc(gpu.d, rows * slots * 4); defer dw.free(); - try f.pick(s, dl.ptr, dp.ptr, dw.ptr, rows, experts, top_k); + try f.pick(s, dl.ptr, dp.ptr, dw.ptr, rows, experts, top_k, 0); try s.synchronize(); const got_ids = try check.download(gpu, dp); defer a.free(got_ids); @@ -148,7 +148,7 @@ fn combine(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { defer dw.free(); var dout = try cuda.DeviceBuffer.alloc(gpu.d, rows * dims * 2); defer dout.free(); - try f.combine(s, dy.ptr, true, dw.ptr, dout.ptr, rows, slots, dims); + try f.combine(s, dy.ptr, true, dw.ptr, dout.ptr, rows, slots, dims, 0); try s.synchronize(); const got = try check.download(gpu, dout); defer a.free(got); @@ -157,9 +157,9 @@ fn combine(gpu: Gpu, f: cuda.moe.Fns, s: cuda.Stream, r: std.Random) !void { } /// Random affine-4 experts (words, scales, biases) as the checkpoint stores them. -const Table = struct { words: []u32, scales: []u16, biases: []u16, count: usize, n: usize, k: usize }; +pub const Table = struct { words: []u32, scales: []u16, biases: []u16, count: usize, n: usize, k: usize }; -fn table(a: std.mem.Allocator, r: std.Random, count: usize, n: usize, k: usize) !Table { +pub fn table(a: std.mem.Allocator, r: std.Random, count: usize, n: usize, k: usize) !Table { const t: Table = .{ .words = try a.alloc(u32, count * n * k / 8), .scales = try a.alloc(u16, count * n * k / 64), .biases = try a.alloc(u16, count * n * k / 64), .count = count, .n = n, .k = k }; for (t.words) |*w| w.* = r.int(u32); for (t.scales) |*v| v.* = f32bf16(0.02 + r.float(f32) * 0.03); @@ -306,4 +306,5 @@ pub fn run(gpu: Gpu) !void { try topk(gpu, f, s, r); try combine(gpu, f, s, r); try pairExperts(gpu, s, r); + try @import("moe_live_tests.zig").run(gpu); }