From 5d22097116ec226a4065217c3b05ac013de45b0e Mon Sep 17 00:00:00 2001 From: styles01 <11393400+styles01@users.noreply.github.com> Date: Sat, 24 Oct 2026 23:31:54 -0600 Subject: [PATCH 1/7] feat(family): add minimax_m3_vl + minimax_m3 CUDA kernels MiniMax-M3.1-Flash (EXL3 EAQ pack) as a native CUDA family serving through the shared lane core: config, EXL3 weight load, sparse attention, layer graph, serve, and the CUDA-native adapter with the TP2 rank boot mirroring glm's env-rank pattern. tf_m3_scores/topk/decode/decode_combine/swigluoai kernels plus the routed-experts extension of exl3_experts. Gate issue #1280. Device parity 10/10; real-pack f64 grade: greedy tokens exact vs f64, logits max|delta| 0.003237. --- zig/kernels/cuda/exl3_experts.cu | 5 + zig/kernels/cuda/minimax_m3.cu | 265 ++++ zig/src/families/minimax_m3_vl/config.zig | 306 +++++ .../families/minimax_m3_vl/cuda_native.zig | 782 ++++++++++++ zig/src/families/minimax_m3_vl/cuda_serve.zig | 49 + .../families/minimax_m3_vl/cuda_sparse.zig | 177 +++ zig/src/families/minimax_m3_vl/exl3.zig | 449 +++++++ zig/src/families/minimax_m3_vl/graph.zig | 1058 ++++++++++++++++ zig/src/families/minimax_m3_vl/serve.zig | 1107 +++++++++++++++++ zig/src/families/minimax_m3_vl/shard.zig | 116 ++ zig/src/families/minimax_m3_vl/sparse.zig | 163 +++ .../minimax_m3_vl/test_pack_config.json | 290 +++++ .../test_quantization_config.json | 8 + zig/src/families/minimax_m3_vl/weights.zig | 430 +++++++ 14 files changed, 5205 insertions(+) create mode 100644 zig/kernels/cuda/minimax_m3.cu create mode 100644 zig/src/families/minimax_m3_vl/config.zig create mode 100644 zig/src/families/minimax_m3_vl/cuda_native.zig create mode 100644 zig/src/families/minimax_m3_vl/cuda_serve.zig create mode 100644 zig/src/families/minimax_m3_vl/cuda_sparse.zig create mode 100644 zig/src/families/minimax_m3_vl/exl3.zig create mode 100644 zig/src/families/minimax_m3_vl/graph.zig create mode 100644 zig/src/families/minimax_m3_vl/serve.zig create mode 100644 zig/src/families/minimax_m3_vl/shard.zig create mode 100644 zig/src/families/minimax_m3_vl/sparse.zig create mode 100644 zig/src/families/minimax_m3_vl/test_pack_config.json create mode 100644 zig/src/families/minimax_m3_vl/test_quantization_config.json create mode 100644 zig/src/families/minimax_m3_vl/weights.zig diff --git a/zig/kernels/cuda/exl3_experts.cu b/zig/kernels/cuda/exl3_experts.cu index ecb651de6b..b31d90b40a 100644 --- a/zig/kernels/cuda/exl3_experts.cu +++ b/zig/kernels/cuda/exl3_experts.cu @@ -247,6 +247,11 @@ extern "C" __global__ void tf_exl3_experts_rot_in(const RotInArgs a) { const __nv_bfloat16* xv = reinterpret_cast(xr); #pragma unroll for (int j = 0; j < 4; ++j) v[j] = __bfloat162float(xv[j]) * __half2float(suh[j]); + } else if (a.x_dtype == 2) { + // f32 input (the minimax engine's x32): element indexing must be 32-bit, not the 16-bit xr stride + const float* xv = reinterpret_cast(a.x) + (size_t)row * a.x_stride + blk * 128 + 4 * lane; +#pragma unroll + for (int j = 0; j < 4; ++j) v[j] = xv[j] * __half2float(suh[j]); } else { const half* xv = reinterpret_cast(xr); #pragma unroll diff --git a/zig/kernels/cuda/minimax_m3.cu b/zig/kernels/cuda/minimax_m3.cu new file mode 100644 index 0000000000..ec121526eb --- /dev/null +++ b/zig/kernels/cuda/minimax_m3.cu @@ -0,0 +1,265 @@ +// MiniMax M3.1's MSA sparse-indexer on CUDA, the exllamav3-fork m3_sparse_attn contracts by hand: +// block-max-pool scores (fp32), the fixed-width top-16 with the tie contract, and the block-gather +// decode attention over the fp16 paged cache. The grading reference is this family's CPU mirror +// (zig/src/families/minimax_m3_vl/sparse.zig, 3/3 green) and the fork's INTERFACE-M3-KERNELS.md §1-2. +// +// Contract numbers (asserted in the family's config.zig): index_n_heads 4 == num_key_value_heads, +// index_head_dim 128, index_block_size 128, index_topk_blocks 16, index_local_blocks 1, init_block 0. +#pragma once + +#include +#include + +namespace m3 { + +constexpr int D = 128; // index_head_dim == head_dim == num_key_value_heads' 4 x 128? D is the per-head width +constexpr int BLOCK = 128; // index_block_size: keys per scored block; a 128-run never straddles PAGE 256 +constexpr int TOPK = 16; // index_topk_blocks: the fixed selection width +constexpr int PAGE = 256; // the paged cache's in-tree PAGE_SIZE +constexpr float NEG_INF = -__builtin_huge_valf(); +constexpr float POS_INF = __builtin_huge_valf(); + +__device__ __forceinline__ float h2f(__half v) { return __half2float(v); } +__device__ __forceinline__ bool pnan(float s) { return s != s; } // the NaN rule: propagate, amax keeps NaN + +// f32 warp sum, a lane's 4-dim partial then all 32 lanes (glue convention: no contraction, the flags keep it). +__device__ __forceinline__ float warp_sum(float s) { + for (int o = 16; o > 0; o >>= 1) s = s + __shfl_xor_sync(0xffffffffu, s, o); + return s; +} + +} // namespace m3 + +// m3_indexer_scores (_ref and _m3_indexer_scores_kernel): one program per (query row, key block, +// batch); its warp w owns index head w alone, so the four heads run without any cross-warp state. +// Scoring is fp32; a key is masked iff k_pos < 0 or k_pos > q_pos (right-open causal on content +// positions); the block score is the masked amax over the block's 128 keys, then boosts OVERRIDE: +// the query's own block (local_blocks=1) plus the leading [0, init) blocks become +inf after the +// full reduction. A NaN dot makes the block score NaN (torch amax's class). Slot anchor: blocks +// are absolute key-slot runs, NB = ceil(Sk/128). +extern "C" __global__ void __launch_bounds__(128) + tf_m3_scores(const __half* __restrict__ q, // (B, H, Sq, D) index queries + const __half* __restrict__ k, // (B, Sk, D) index keys (plane rows or fresh K) + const int64_t* __restrict__ qpos, // (B, Sq) content positions of the queries + const int64_t* __restrict__ kpos, // (B, Sk); -1 = padded/invalid + float* __restrict__ out, // (B, H, Sq, NB) + int64_t Sq, int64_t NB, int64_t Sk, int64_t H, int local, int init) { + const int64_t sq = blockIdx.x, nb = blockIdx.y, b = blockIdx.z; + const int w = threadIdx.x >> 5, lane = threadIdx.x & 31; + const int64_t qp = qpos[b * Sq + sq]; + // Every program re-reads its k rows four times (one per head); the transient stays (D,) per warp. + const __half* kb = k + (b * Sk + nb * m3::BLOCK) * m3::D + lane * 4; + float best = m3::NEG_INF; + bool nan = false; + for (int j = 0; j < m3::BLOCK; ++j) { + const int64_t slot = nb * m3::BLOCK + j; + if (slot >= Sk) break; // blocks anchored to absolute slots: this block's own short tail + const int64_t kp = kpos[b * Sk + slot]; + if (kp < 0 || kp > qp) continue; // masked iff out of the right-open causal window + const __half* qb = q + ((b * H + w) * Sq + sq) * m3::D + lane * 4; + const __half* kr = kb + j * m3::D; + float s = 0.0f; + for (int m = 0; m < 4; ++m) s = s + m3::h2f(qb[m]) * m3::h2f(kr[m]); // fp16 products are exact in fp32 + s = m3::warp_sum(s); + if (lane == 0) { + if (m3::pnan(s)) nan = true; // amax keeps a NaN: the score's class follows + else if (s > best) best = s; // ties at greater... amax keeps the value, ids do not matter + } + } + if (lane == 0) { + const int64_t qb = qp / m3::BLOCK; // the query's own block, slot-anchored + const bool boost = (nb >= qb - (local - 1) && nb <= qb) || nb < init; + if (nan) best = m3::POS_INF / m3::POS_INF; // a NaN row: NaN wins the max like torch's amax + out[((b * H + w) * Sq + sq) * NB + nb] = boost ? m3::POS_INF : best; + } +} + +// m3_topk16 (_m3_topk_kernel): one score row per program, k picks by iterative argmax with per-pick +// invalidation marks; descending score, ties at the SMALLER block id, valid ids left-packed, the +// fixed topk width with -1 right-pad (-inf never yields an id). NaN never wins a comparison: it +// is the stated topk rule and the CPU mirror checks it. +extern "C" __global__ void __launch_bounds__(128) + tf_m3_topk(const float* __restrict__ scores, int32_t* __restrict__ out, + int64_t NB, int64_t rows, int k) { + const int64_t r = blockIdx.x; + const float* row = scores + r * NB; + int32_t* dst = out + r * k; + __shared__ unsigned int marks[32]; // NB <= 1024: one mark bit per block (the glue checks the bound) + for (int i = threadIdx.x; i < 32; i += blockDim.x) marks[i] = 0u; + __shared__ float s_best; + __shared__ int i_best, sh_written; + const int t = threadIdx.x, lane = t & 31, warp = t >> 5; + __shared__ float wbest[4]; + __shared__ int wid[4]; + if (t == 0) sh_written = 0; + __syncthreads(); + bool done = false; + for (int pick = 0; pick < k; ++pick) { + if (done) break; + float bs = m3::NEG_INF; + int bi = -1; + for (int64_t i = t; i < NB; i += blockDim.x) { + if (marks[i >> 5] & (1u << (i & 31))) continue; + const float s = row[i]; + if (s > bs && !m3::pnan(s)) { bs = s; bi = (int)i; } // plain >: ties keep the earlier (smaller) id + } + for (int o = 16; o > 0; o >>= 1) { + const float os = __shfl_down_sync(0xffffffffu, bs, o); + const int oi = __shfl_down_sync(0xffffffffu, bi, o); + if (os > bs || (os == bs && oi < bi && oi >= 0)) { bs = os; bi = oi; } + } + if (lane == 0) { wbest[warp] = bs; wid[warp] = bi; } + __syncthreads(); + if (t == 0) { + float s = m3::NEG_INF; + int i = -1; + for (int wp = 0; wp < 4; ++wp) + if (wbest[wp] > s || (wbest[wp] == s && wid[wp] < i && wid[wp] >= 0)) { s = wbest[wp]; i = wid[wp]; } + s_best = s; + i_best = i; + } + __syncthreads(); + if (s_best == m3::NEG_INF) { // exhausted: the fixed width's -1 right-pad, left-packed so far + if (t == 0) + for (; sh_written < k; ++sh_written) dst[sh_written] = -1; + done = true; + __syncthreads(); + break; + } + if (t == 0) { + dst[sh_written++] = i_best; + atomicOr(&marks[i_best >> 5], 1u << (i_best & 31)); // the per-pick invalidation mark + } + __syncthreads(); + } +} +// m3_sparse_decode_attn's split kernel (_m3_decode_kernel): one program per (batch, q head, split); +// 128 threads, a dim each. Its split of the fixed k-wide selection (-1 ids skipped) gathers the +// fp16 paged cache through the block table; GQA maps q head hq onto kv head hq/group. Online +// softmax in fp32 (acc, m, l): every finite key contributes — a sub-max key adds __expf(full - +// m_new) * v on top of the rescale-carried sum (never wipes it); a NaN score drops the whole key +// (m, l, acc unchanged). A key is masked iff its logical position >= cache_seqlens. The partials +// land row-major in the (B, Hq, nsplit, D + 2) fp32 scratch the combine kernel merges. +// Zero-valid-key splits exit with l = 0; the combine then returns an all-zero row, never a NaN. +extern "C" __global__ void __launch_bounds__(128) + tf_m3_decode(const __half* __restrict__ q, // (B, Hq, D) + const __half* __restrict__ kt, // (P, PAGE, Hkv, D) fp16 paged K + const __half* __restrict__ vt, // (P, PAGE, Hkv, D) + const int32_t* __restrict__ table, // (B, npp) logical page -> physical page + const int32_t* __restrict__ sel, // (B, Hkv, k) fixed-width, -1 right-pad + const int64_t* __restrict__ seqlens, // (B,) logical key lengths (before the append) + float* __restrict__ part, // (B, Hq, nsplit, D + 2) acc | m | l + int64_t Hq, int64_t Hkv, int64_t npp, float scale, int k, int nsplit) { + const int64_t b = blockIdx.x, hq = blockIdx.y; + const int split = blockIdx.z, t = threadIdx.x; + const int per = (k + nsplit - 1) / nsplit; + const int kvh = (int)(hq / (Hq / Hkv)); // the GQA mapping: index heads == kv heads, per-head lists + const __half* qb = q + (b * Hq + hq) * m3::D + t; + const int32_t* row = sel + (b * Hkv + kvh) * k; + __shared__ float wsum[4], sh_alpha, sh_p; + float acc = 0.0f, m = m3::NEG_INF, l = 0.0f; + const int lo = split * per, hi = lo + per < k ? lo + per : k; + for (int si = lo; si < hi; ++si) { + const int bid = row[si]; + if (bid < 0) continue; // left-packed: -1 only at the tail + const __half* kr = kt + (((int64_t)bid * m3::BLOCK + 0) * Hkv + kvh) * m3::D + t; + for (int j = 0; j < m3::BLOCK; ++j) { + const int64_t p = (int64_t)bid * m3::BLOCK + j; + if (p >= seqlens[b]) continue; + const int64_t page = table[b * npp + p / m3::PAGE]; // 128 runs never straddle the 256 page + const int64_t off = p % m3::PAGE; + const __half* krow = kt + ((page * m3::PAGE + off) * Hkv + kvh) * m3::D + t; + const __half* vrow = vt + ((page * m3::PAGE + off) * Hkv + kvh) * m3::D + t; + float s = m3::h2f(qb[0]) * m3::h2f(krow[0]); // fp16 products exact in fp32; the flags keep add+mul apart + s = m3::warp_sum(s); + const int warp = t >> 5, lane = t & 31; + if (lane == 0) wsum[warp] = s; + __syncthreads(); + float m_new = m, alpha = 1.0f, pj = 0.0f; + if (t == 0) { + const float full = scale * ((wsum[0] + wsum[1]) + (wsum[2] + wsum[3])); + if (full == full) { // a NaN score drops the whole key: m, l and acc stay as they were + m_new = full > m ? full : m; + alpha = m == m3::NEG_INF ? 0.0f : __expf(m - m_new); // rescale the carried sum (1 when full <= m) + pj = __expf(full - m_new); + l = l * alpha + pj; + m = m_new; + } else { alpha = 1.0f; pj = 0.0f; m_new = m; } + sh_alpha = alpha; + sh_p = pj; + } + __syncthreads(); + acc = acc * sh_alpha + sh_p * m3::h2f(vrow[0]); + __syncthreads(); + } + } + float* dst = part + (((b * Hq + hq) * nsplit) + split) * (m3::D + 2); + dst[t] = acc; + if (t == 0) { + dst[m3::D] = m; + dst[m3::D + 1] = l; + } +} + +// tf_m3_swigluoai: the M3.1 MoE/dense activation rows, the fork's activation.cu:124 contract by +// hand: (clamp(u, -limit, limit) + 1) * min(g, limit) * sigmoid(1.702 * min(g, limit)) -> fp16. +// The mainline EXL3 experts' fused gateup epilogue computes SiLU x up (exl3_experts.cu's +// tf_exl3_gateup_epilogue), which cannot express the alpha = 1.702 sigmoid form, so this family +// applies its own activation between the grouped gate/up GEMM partials and the down GEMM. The +// grading reference is serve.zig's swigluoai (the same contract at f32 exp). +extern "C" __global__ void __launch_bounds__(128) + tf_m3_swigluoai(const __half* __restrict__ g, // (rows, N) fp16 gate rows + const __half* __restrict__ u, // (rows, N) fp16 up rows + __half* __restrict__ o, // (rows, N) fp16 out + int64_t n, float limit) { + const int64_t p = blockIdx.x; + const int64_t base = ((int64_t)blockIdx.y * blockDim.x + threadIdx.x) * 4; + const __half* gr = g + p * n; + const __half* ur = u + p * n; + __half* dst = o + p * n; + #pragma unroll + for (int j = 0; j < 4; ++j) { + const int64_t i = base + j; + if (i >= n) return; + float gv = m3::h2f(gr[i]); + float uv = m3::h2f(ur[i]); + gv = gv > limit ? limit : gv; // min(g, limit): clipped from above only + uv = uv > limit ? limit : (uv < -limit ? -limit : uv); // torch.clamp both ways + const float glu = gv / (1.0f + __expf(-1.702f * gv)); // sigmoid(1.702 g), the fork's fast exp + dst[i] = __float2half_rn((uv + 1.0f) * glu); + } +} + +// _m3_decode_combine_kernel: the four splits' fp32 (acc, m, l) merge into the row's fp16 output +// in the q dtype; an all-l-zero row (no selected keys were valid) returns zeros, never a NaN. +extern "C" __global__ void __launch_bounds__(128) + tf_m3_decode_combine(const float* __restrict__ part, __half* __restrict__ out, + int64_t Hq, int nsplit) { + const int64_t b = blockIdx.x, hq = blockIdx.y; + const int t = threadIdx.x; + const float* base = part + (((b * Hq + hq) * nsplit) + 0) * (m3::D + 2); + __shared__ float ms[8], ls[8], sh_l, sh_m; + if (t == 0) { + float mstar = m3::NEG_INF, ltot = 0.0f; + for (int i = 0; i < nsplit; ++i) { + const float m = base[i * (m3::D + 2) + m3::D], l = base[i * (m3::D + 2) + m3::D + 1]; + ms[i] = l > 0.0f ? m : m3::NEG_INF; // an empty split never moves the max + ls[i] = l; + if (mstar < ms[i]) mstar = ms[i]; + } + // the merge's denominator is the max-rescaled total: a split's l only lives at its own m, so + // l * exp(m - m*) restores the shared scale before summing (a raw sum understates every + // split whose max sits below the global one — with nsplit=1 the weight is 1 and it vanishes) + for (int i = 0; i < nsplit; ++i) + if (ls[i] > 0.0f) ltot += __expf(ms[i] - mstar) * ls[i]; + sh_m = mstar; + sh_l = ltot; + } + __syncthreads(); + __half* dst = out + (b * Hq + hq) * m3::D + t; + if (sh_l == 0.0f) { *dst = __float2half_rn(0.0f); return; } // zero-valid-key row: zeros, not NaN + float v = 0.0f; + for (int i = 0; i < nsplit; ++i) + if (ls[i] > 0.0f) v += __expf(ms[i] - sh_m) * base[i * (m3::D + 2) + t]; + *dst = __float2half_rn(v / sh_l); +} diff --git a/zig/src/families/minimax_m3_vl/config.zig b/zig/src/families/minimax_m3_vl/config.zig new file mode 100644 index 0000000000..0d780ce2d5 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/config.zig @@ -0,0 +1,306 @@ +//! MiniMax-M3.1-Flash's shape from config.json's text_config (minimax_m3_vl): MoE + the sparse-indexer attention (MSA). + +const std = @import("std"); + +pub const max_layers = 60; + +/// config.json's `quantization_config` codebook, in the core reader's order (`core/exl3_format.zig`: mul1 is id 2). +pub const Codebook = enum { unknown, inst3, mcg, mul1 }; + +pub const Config = struct { + hidden: u32 = 6144, + layers: u32 = 60, + run: u32 = 60, // backbone layers loaded and run: all, or the first few for a check that must fit + vocab: u32 = 200064, + head_unit: u32 = 64, // the vocabulary's TP cut in whole units: 128 for EXL3, whose matrices split at 128 columns + exl3: bool = false, // quant_method "exl3" in config.json, or the standalone quantization_config.json + qc_avg_bits: f32 = 0, // the quantization_config's avg_bits (3.06 for Mia-AiLab's pack) + qc_head_bits: u32 = 0, // the language head's own width (head_bits 8: lm_head is the only 8-bit tensor) + qc_codebook: Codebook = .unknown, // Mia-AiLab's pack: mul1, odd half-bit units, mixed widths per matrix + eps: f32 = 1e-6, + dense_layers: u32 = 3, // the leading MoE-free layers: moe_layer_freq's first 1, else first_k_dense_replace + dense_inter: u32 = 12288, + experts: u32 = 128, + topk: u32 = 4, + moe_inter: u32 = 3072, + shared_inter: u32 = 3072, // n_shared_experts 1's width: shared_intermediate_size + routed_scale: f32 = 2.0, + sigmoid_router: bool = true, // scoring_func "sigmoid": the router scores without a softmax + routing_bias: bool = true, // use_routing_bias: the router keeps one bias per expert + norm_topk: bool = false, // no norm_topk_prob in the config: the routed weights stay as scored + swiglu_alpha: f32 = 1.702, // the swigluoai gate's alpha + swiglu_limit: f32 = 7.0, // the swigluoai gate's alpha-limit clamp + attn_heads: u32 = 64, + kv_heads: u32 = 4, // GQA + head_dim: u32 = 128, + rotary_dim: u32 = 64, // partial rotary: rotary_dim 64 of head_dim 128 + partial_rotary: f32 = 0.5, + qk_norm: bool = true, // use_qk_norm: q and k carry their own per-head norms + i_heads: u32 = 4, // the MSA indexer's heads (sparse_num_index_heads) + i_dim: u32 = 128, // the indexer head's dim (sparse_index_dim) + i_topk: u32 = 16, // whole pooled blocks a row attends (sparse_topk_blocks) + block: u32 = 128, // the pooled key block's width (sparse_block_size) + score_max: bool = true, // sparse_score_type "max" over the indexer heads + tp: u32 = 1, // tensor parallel: the ranks splitting the attention heads + tp_rank: u32 = 0, + + /// Keys a sparse row attends: the selected pools' whole blocks (stage-2's gather adds the init/local ones). + pub fn keyWidth(c: *const Config) u32 { + return c.i_topk * c.block; + } + + /// The head's rows [lo, hi) this rank computes, in whole units of `head_unit` (EXL3 splits at 128 columns). + pub fn vocabPart(c: *const Config) [2]u32 { + const n = c.vocab / c.tp / c.head_unit * c.head_unit; + return .{ c.tp_rank * n, if (c.tp_rank + 1 == c.tp) c.vocab else (c.tp_rank + 1) * n }; + } +}; + +fn int(o: std.json.ObjectMap, key: []const u8) !i64 { + const v = o.get(key) orelse return error.BadConfig; + return switch (v) { + .integer => |i| i, + else => error.BadConfig, + }; +} + +fn float(o: std.json.ObjectMap, key: []const u8) !f64 { + const v = o.get(key) orelse return error.BadConfig; + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => error.BadConfig, + }; +} + +fn boolOr(o: std.json.ObjectMap, key: []const u8, or_else: bool) bool { + const v = o.get(key) orelse return or_else; + return switch (v) { + .bool => |b| b, + else => or_else, + }; +} + +fn stringIs(o: std.json.ObjectMap, key: []const u8, want: []const u8) bool { + const v = o.get(key) orelse return false; + return v == .string and std.mem.eql(u8, v.string, want); +} + +/// config.json's quantization_config (top level, else text_config's): glm/cuda_exl3.zig's block(). +fn quantBlock(root: std.json.Value) ?std.json.ObjectMap { + if (root != .object) return null; + if (root.object.get("quantization_config")) |q| if (q == .object) return q.object; + if (root.object.get("text_config")) |t| if (t == .object) if (t.object.get("quantization_config")) |q| if (q == .object) return q.object; + return null; +} + +fn codebook(name: []const u8) Codebook { + if (std.mem.eql(u8, name, "mul1")) return .mul1; + if (std.mem.eql(u8, name, "mcg")) return .mcg; + if (std.mem.eql(u8, name, "inst3")) return .inst3; + return .unknown; +} + +/// One quantization_config's exl3 detection: the method widens the vocab TP cut to 128, the fields feed a later gate. +fn detectExl3(c: *Config, q: std.json.ObjectMap) void { + if (!stringIs(q, "quant_method", "exl3")) return; + c.exl3 = true; + c.head_unit = 128; + if (q.get("avg_bits")) |v| if (v == .float or v == .integer) { + c.qc_avg_bits = @floatCast(float(q, "avg_bits") catch return); + }; + if (q.get("head_bits")) |v| if (v == .integer) { + if (v.integer > 0) c.qc_head_bits = @intCast(v.integer); + }; + if (q.get("codebook")) |v| if (v == .string) { + c.qc_codebook = codebook(v.string); + }; +} + +/// text_config with M3.1-Flash's widths; smaller synthetic counts pass; `qc_json` reads a sibling quant config's JSON. +pub fn parse(gpa: std.mem.Allocator, json: []const u8, standalone_qc: ?[]const u8) !Config { + const parsed = try std.json.parseFromSlice(std.json.Value, gpa, json, .{}); + defer parsed.deinit(); + const root = parsed.value.object; + const t = (root.get("text_config") orelse parsed.value).object; + var c = Config{}; + const want = Config{}; + c.hidden = @intCast(try int(t, "hidden_size")); + c.layers = @intCast(try int(t, "num_hidden_layers")); + c.run = c.layers; + c.vocab = @intCast(try int(t, "vocab_size")); + c.attn_heads = @intCast(try int(t, "num_attention_heads")); + c.kv_heads = @intCast(try int(t, "num_key_value_heads")); + c.head_dim = @intCast(try int(t, "head_dim")); + c.eps = @floatCast(try float(t, "rms_norm_eps")); + c.experts = @intCast(try int(t, "num_local_experts")); + c.topk = @intCast(try int(t, "num_experts_per_tok")); + c.moe_inter = @intCast(try int(t, "intermediate_size")); + c.dense_inter = @intCast(try int(t, "dense_intermediate_size")); + c.shared_inter = @intCast(try int(t, "shared_intermediate_size")); + c.routed_scale = @floatCast(try float(t, "routed_scaling_factor")); + c.swiglu_alpha = @floatCast(try float(t, "swiglu_alpha")); + c.swiglu_limit = @floatCast(try float(t, "swiglu_limit")); + c.sigmoid_router = stringIs(t, "scoring_func", "sigmoid"); + c.routing_bias = boolOr(t, "use_routing_bias", false); + c.qk_norm = boolOr(t, "use_qk_norm", false); + if (t.get("qk_norm_type")) |v| { + if (!(v == .string and std.mem.eql(u8, v.string, "per_head"))) return error.NotMiniMaxM3VL; + } else if (c.qk_norm) return error.NotMiniMaxM3VL; // the per-head kind is this family's only one + if (c.swiglu_limit <= 0) return error.NotMiniMaxM3VL; // the swigluoai gate's alpha-limit clamp + if (t.get("rotary_dim")) |v| { + if (v != .integer) return error.BadConfig; + c.rotary_dim = @intCast(v.integer); + } + if (t.get("partial_rotary_factor")) |v| { + if (v != .float and v != .integer) return error.BadConfig; + c.partial_rotary = @floatCast(try float(t, "partial_rotary_factor")); + } + if (t.get("moe_layer_freq")) |v| { + if (v != .array) return error.BadConfig; + var dense: u32 = c.layers; // the sentinel: no MoE wave seen yet + for (v.array.items, 0..) |e, i| { + if (e != .integer) return error.BadConfig; + if (e.integer == 1) { + if (dense == c.layers) dense = @intCast(i); + } else if (e.integer != 0) return error.NotMiniMaxM3VL; + if (dense != c.layers and e.integer == 0) return error.NotMiniMaxM3VL; // one dense->MoE wave only + } + c.dense_layers = @min(dense, c.layers); + } else { + c.dense_layers = @intCast(try int(t, "first_k_dense_replace")); // the deepseek spelling's fallback + } + if (t.get("sparse_attention_config")) |v| { + if (v != .object) return error.BadConfig; + const s = v.object; + c.i_heads = @intCast(try int(s, "sparse_num_index_heads")); + c.i_dim = @intCast(try int(s, "sparse_index_dim")); + c.i_topk = @intCast(try int(s, "sparse_topk_blocks")); + c.block = @intCast(try int(s, "sparse_block_size")); + c.score_max = stringIs(s, "sparse_score_type", "max"); + } else { + // the GLM DSA indexer's flat spellings: the same machinery's token-topk shape, expressed in blocks + c.i_heads = @intCast(try int(t, "index_n_heads")); + c.i_dim = @intCast(try int(t, "index_head_dim")); + const kpool = try int(t, "index_kpool"); + if (kpool <= 0) return error.BadConfig; + c.block = @intCast(kpool); + c.i_topk = @intCast(@divExact(try int(t, "index_topk"), kpool)); // whole pools only: a partial one refuses + } + var qc_parsed: ?std.json.Parsed(std.json.Value) = null; + defer if (qc_parsed) |p| p.deinit(); + var qc: ?std.json.ObjectMap = quantBlock(parsed.value); + if (qc == null) { + if (standalone_qc) |s| { + qc_parsed = try std.json.parseFromSlice(std.json.Value, gpa, s, .{}); + const p = qc_parsed.?; + if (p.value == .object) { + qc = p.value.object; + } else { + qc_parsed = null; + } + } + } + if (qc) |q| detectExl3(&c, q); + if (c.layers > max_layers) return error.NotMiniMaxM3VL; + const widths = c.hidden == want.hidden and c.dense_inter == want.dense_inter and c.topk == want.topk and + c.moe_inter == want.moe_inter and c.shared_inter == want.shared_inter and c.attn_heads == want.attn_heads and + c.kv_heads == want.kv_heads and c.head_dim == want.head_dim and c.i_heads == want.i_heads and + c.i_dim == want.i_dim and c.i_topk == want.i_topk and c.block == want.block and + c.eps == want.eps and c.swiglu_alpha == want.swiglu_alpha; + const counts = c.layers >= 1 and c.dense_layers <= c.layers and c.experts >= c.topk and + c.experts <= want.experts and c.vocab <= want.vocab and c.vocab % 128 == 0; + if (!widths or !counts) return error.NotMiniMaxM3VL; + return c; +} + +/// MiniMax-M3.1-Flash's own counts: the shapes the engines' kernels are built for (smaller synthetic ones pass). +pub fn whole(c: Config) bool { + const want = Config{}; + return c.layers == want.layers and c.vocab == want.vocab and c.dense_layers == want.dense_layers and c.experts == want.experts; +} + +test "the MiniMax-M3.1-Flash EXL3 pack's config.json parses and carries the model_type the registry dispatches" { + const gpa = std.testing.allocator; + const json = @embedFile("test_pack_config.json"); + var root = try std.json.parseFromSlice(std.json.Value, gpa, json, .{}); + defer root.deinit(); + try std.testing.expectEqualStrings("minimax_m3_vl", root.value.object.get("model_type").?.string); + const c = try parse(gpa, json, null); + try std.testing.expect(whole(c)); + try std.testing.expect(c.exl3); + try std.testing.expectEqual(Codebook.mul1, c.qc_codebook); // mul1: the core reader's codebook id 2 + try std.testing.expectEqual(@as(f32, 3.06), c.qc_avg_bits); + try std.testing.expectEqual(@as(u32, 8), c.qc_head_bits); // head_bits 8: only lm_head carries it + try std.testing.expectEqual(@as(u32, 128), c.head_unit); // the vocabulary's TP cut once EXL3 is on + try std.testing.expectEqual(@as(u32, 60), c.layers); + try std.testing.expectEqual(@as(u32, 3), c.dense_layers); // moe_layer_freq's wave turns MoE from layer 3 + try std.testing.expectEqual(@as(u32, 128), c.experts); + try std.testing.expectEqual(@as(u32, 4), c.topk); + try std.testing.expectEqual(@as(u32, 3072), c.moe_inter); + try std.testing.expectEqual(@as(u32, 12288), c.dense_inter); + try std.testing.expectEqual(@as(f32, 1.702), c.swiglu_alpha); // the swigluoai gate + try std.testing.expectEqual(@as(f32, 7.0), c.swiglu_limit); + try std.testing.expect(c.sigmoid_router and c.routing_bias); // sigmoid scores plus a bias per expert + try std.testing.expect(!c.norm_topk); + try std.testing.expect(c.qk_norm); // use_qk_norm with qk_norm_type "per_head" + try std.testing.expectEqual(@as(u32, 4), c.kv_heads); // GQA + try std.testing.expectEqual(@as(u32, 64), c.rotary_dim); + try std.testing.expectEqual(@as(u32, 4), c.i_heads); // the MSA indexer: 4 heads of 128, 16 blocks of 128, max + try std.testing.expectEqual(@as(u32, 128), c.i_dim); + try std.testing.expectEqual(@as(u32, 16), c.i_topk); + try std.testing.expectEqual(@as(u32, 128), c.block); + try std.testing.expect(c.score_max); + try std.testing.expectEqual(@as(u32, 2048), c.keyWidth()); + try std.testing.expectEqual([2]u32{ 0, 200064 }, c.vocabPart()); +} + +test "a standalone quantization_config.json admits the same EXL3 gate when config.json's inline copy is absent" { + const gpa = std.testing.allocator; + const qc = @embedFile("test_quantization_config.json"); + const noqc = try std.fmt.allocPrint(gpa, "{{\"model_type\": \"minimax_m3_vl\", \"text_config\": {{" ++ + "\"hidden_size\": 6144, \"num_hidden_layers\": 2, \"vocab_size\": 2048, \"num_attention_heads\": 64," ++ + "\"num_key_value_heads\": 4, \"head_dim\": 128, \"rms_norm_eps\": 1e-06, \"num_local_experts\": 4," ++ + "\"num_experts_per_tok\": 4, \"intermediate_size\": 3072, \"dense_intermediate_size\": 12288," ++ + "\"shared_intermediate_size\": 3072, \"routed_scaling_factor\": 2.0, \"swiglu_alpha\": 1.702," ++ + "\"swiglu_limit\": 7.0, \"use_qk_norm\": true, \"qk_norm_type\": \"per_head\", \"moe_layer_freq\": [0, 1]," ++ + "\"sparse_attention_config\": {{\"sparse_num_index_heads\": 4, \"sparse_index_dim\": 128," ++ + "\"sparse_topk_blocks\": 16, \"sparse_block_size\": 128, \"sparse_score_type\": \"max\"}}}}}}", .{}); + defer gpa.free(noqc); + const c = try parse(gpa, noqc, qc); + try std.testing.expect(c.exl3); + try std.testing.expectEqual(Codebook.mul1, c.qc_codebook); + try std.testing.expectEqual(@as(f32, 3.06), c.qc_avg_bits); + try std.testing.expectEqual(@as(u32, 128), c.head_unit); + const d = try parse(gpa, noqc, null); // without the standalone summary there is no EXL3 either + try std.testing.expect(!d.exl3); + try std.testing.expectEqual(Codebook.unknown, d.qc_codebook); + try std.testing.expectEqual(@as(u32, 64), d.head_unit); +} + +test "the GLM DSA indexer's flat spellings also parse, and other widths do not" { + const gpa = std.testing.allocator; + const flat = try std.fmt.allocPrint(gpa, "{{\"text_config\": {{\"hidden_size\": 6144, \"num_hidden_layers\": 2," ++ + "\"vocab_size\": 2048, \"num_attention_heads\": 64, \"num_key_value_heads\": 4, \"head_dim\": 128," ++ + "\"rms_norm_eps\": 1e-06, \"num_local_experts\": 4, \"num_experts_per_tok\": 4, \"intermediate_size\": 3072," ++ + "\"dense_intermediate_size\": 12288, \"shared_intermediate_size\": 3072, \"routed_scaling_factor\": 2.0," ++ + "\"swiglu_alpha\": 1.702, \"swiglu_limit\": 7.0, \"use_qk_norm\": true, \"qk_norm_type\": \"per_head\"," ++ + "\"first_k_dense_replace\": 0, \"index_n_heads\": 4, \"index_head_dim\": 128, \"index_topk\": 2048," ++ + "\"index_kpool\": 128}}}}", .{}); + defer gpa.free(flat); + const c = try parse(gpa, flat, null); + try std.testing.expectEqual(@as(u32, 4), c.i_heads); + try std.testing.expectEqual(@as(u32, 16), c.i_topk); // 2048 tokens of a 128-wide pool: 16 blocks + try std.testing.expectEqual(@as(u32, 128), c.block); + try std.testing.expectEqual(@as(u32, 0), c.dense_layers); // first_k_dense_replace's fallback + const wide = try std.fmt.allocPrint(gpa, "{{\"text_config\": {{\"hidden_size\": 4096, \"num_hidden_layers\": 2," ++ + "\"vocab_size\": 2048, \"num_attention_heads\": 64, \"num_key_value_heads\": 4, \"head_dim\": 128," ++ + "\"rms_norm_eps\": 1e-06, \"num_local_experts\": 4, \"num_experts_per_tok\": 4, \"intermediate_size\": 3072," ++ + "\"dense_intermediate_size\": 12288, \"shared_intermediate_size\": 3072, \"routed_scaling_factor\": 2.0," ++ + "\"swiglu_alpha\": 1.702, \"swiglu_limit\": 7.0, \"use_qk_norm\": true, \"qk_norm_type\": \"per_head\"," ++ + "\"moe_layer_freq\": [0, 1], \"sparse_attention_config\": {{\"sparse_num_index_heads\": 4," ++ + "\"sparse_index_dim\": 128, \"sparse_topk_blocks\": 16, \"sparse_block_size\": 128," ++ + "\"sparse_score_type\": \"max\"}}}}}}", .{}); + defer gpa.free(wide); + try std.testing.expectError(error.NotMiniMaxM3VL, parse(gpa, wide, null)); +} \ No newline at end of file diff --git a/zig/src/families/minimax_m3_vl/cuda_native.zig b/zig/src/families/minimax_m3_vl/cuda_native.zig new file mode 100644 index 0000000000..c453207a46 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/cuda_native.zig @@ -0,0 +1,782 @@ +//! MiniMax-M3.1-Flash's native CUDA family: the openWith API (Options/open/capacity, refusals in words), +//! the pack's real admission — per-module bit-width discovery from the trellis headers, the whole weights +//! plan, the served shapes — and the exact serve boundary this stage stops at (the MSA device forward). + +const std = @import("std"); +const cuda = @import("cuda"); +pub const lane_api = @import("lanes"); // the boundary tests reach Stream/Feed/Backend through the family +const lanes = lane_api; +const config = @import("config.zig"); +const exl3 = @import("exl3.zig"); +const weights = @import("weights.zig"); +pub const weights_mod = weights; // the harness's fixture builder rides this surface +/// The MSA indexer's CPU contracts (block-max-pool scores, the fixed-width topk with the boost +/// overrides): part of this family's API, the reference its CUDA kernels are graded against. +pub const indexer = @import("sparse.zig"); +pub const sparse_cuda = @import("cuda_sparse.zig"); // the MSA indexer's CUDA launches +pub const serve_cuda = @import("cuda_serve.zig"); // the family's own glue (the swigluoai rows) +pub const serve = @import("serve.zig"); // the host-graded forward: the device path's contract +pub const rank_plan = @import("shard.zig"); // the TP2 rank placement: this stage receipts the per-rank bytes +pub const graph = @import("graph.zig"); // the launch graph this stage attaches (TP1 through open()) + +pub const model_type = "minimax_m3_vl"; +/// The checkpoints this family reads: Mia-AiLab's MiniMax-M3.1-Flash-EXL3-EAQ-v1-3.06-bpw and its line. +pub const formats: []const []const u8 = &.{"exl3"}; +/// The pack's own window (max_position_embeddings 1,048,576) when a config.json names none. +pub const default_context: i64 = 1048576; +pub const max_segments: u32 = 1; +/// The prompt chunk: the prefill's rows, at most one sparse block's granularity of 128 per position. +pub const prompt_rows: u32 = 4096; +/// `explain` words this family's open refusals too (the admission's failures and the serve boundary). +pub const explains_open = true; + +/// The server's reading of the room left after the weights; null if unreadable. +pub const Room = struct { ctx: *anyopaque, read: *const fn (*anyopaque) ?u64 }; + +/// TP1 only: no `rank` field, so --tp refuses by name (`refusals.family`); no vision field yet, --vision +/// refuses too (the tower's tensors sit in the checkpoint and are stage-3's placement work). +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false, lanes: u32 = 8, lanes_fixed: bool = false, room: ?Room = null, layer_cap: usize = 0 }; + +pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; + +/// What the native server drives; the fields the server's startup path reads after a family opens. +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32 = 1, + learn_bytes: usize = 0, // no learner: the streams keep all the room + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + lone: ?LoneRun = null, + target: *anyopaque, +}; + +/// The registry's entry: a declared comptime TYPE (the engines' openWith surface, aliased to this +/// module's declarations — glm's pattern), not a struct value: `openWith(F)` resolves F.max_segments, +/// F.Options, F.open through it, and the capabilities table reads F.model_type/F.formats. +const me = @This(); +pub const native = struct { + pub const model_type = me.model_type; + pub const formats = me.formats; + pub const default_context = me.default_context; + pub const max_segments = me.max_segments; + pub const prompt_rows = me.prompt_rows; + pub const explains_open = me.explains_open; + pub const Options = me.Options; + pub const Loaded = me.Loaded; + pub const open = me.open; + pub const explain = me.explain; + pub const weightShare = me.weightShare; // the server's admission reads this rank's split share +}; + +/// Slots for `context` plus one whole block's reserve, a multiple of the 128-token key block. +pub const block_tokens: usize = 128; // the MSA key block (sparse_block_size): the pool's granularity + +pub fn capacity(context: usize) usize { + return std.mem.alignForward(usize, context + block_tokens, block_tokens); +} + +/// The server's admission reads this rank's checkpoint share (glm's `weightShare` shape): the +/// receipted placement's exact plan bytes when the env names this rank, the whole pack without +/// the env (TP1 owns everything). +pub fn weightShare(a: std.mem.Allocator, io: std.Io, dir: []const u8, total: u64) u64 { + var ad = admitDir(a, io, dir) catch return total; + defer ad.deinit(); + const p = rank_plan.place() catch return total; // no env: TP1, the whole pack + return @min(rank_plan.bytesOf(&ad.load, p.world, p.rank), total); +} + +/// The admission's result: the parsed shapes, the trellis headers' discovered widths and the whole weights plan. +pub const Admitted = struct { + c: config.Config, + groups: exl3.Plan, + load: weights.Load, + + pub fn deinit(a: *Admitted) void { + a.load.deinit(); + a.groups.deinit(); + } + + /// The checkpoint's weights bytes by the plan's own count (trellis groups + passthrough rows). + pub fn bytes(a: *const Admitted) u64 { + return a.load.bytes(); + } +}; + +/// `admit` over a checkpoint folder: the config parse plus the admission that reads the checkpoint itself. +pub fn admitDir(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !Admitted { + const path = try std.fs.path.join(gpa, &.{ dir, "config.json" }); + defer gpa.free(path); + const text = try std.Io.Dir.cwd().readFileAlloc(io, path, gpa, .limited(1 << 24)); + defer gpa.free(text); + const c = try config.parse(gpa, text, null); + var groups = try exl3.discover(gpa, io, dir, &c); + errdefer groups.deinit(); + var load = try weights.plan(gpa, io, dir, &c, &groups); + errdefer load.deinit(); + return .{ .c = c, .groups = groups, .load = load }; +} + +/// The admission the engines' config path runs (config text only): `qc_json` carries a standalone +/// quantization_config.json when given. +pub fn admit(gpa: std.mem.Allocator, config_json: []const u8, qc_json: ?[]const u8) !config.Config { + return config.parse(gpa, config_json, qc_json); +} + +/// The serve boundary this build stops at, said exactly: admission passed, the family's forward +/// is built and graded on the host, the device kernels are built and graded on the GPU — the +/// full serving engine (weights upload to the device buffers, the per-layer launch graph, the +/// lanes attach) is the remaining stage. +pub const serve_boundary = "minimax_m3_vl passed CUDA admission (every module's EXL3 widths discovered " ++ + "from its trellis headers, the whole weights plan and the served shapes checked), and the family's " ++ + "forward is graded on the host against an independent f64 reference — and this build's device kernels " ++ + "are graded on the GPU (the GB10 device parity receipts: tf_m3_scores/topk/decode/decode_combine and " ++ + "tf_m3_swigluoai) and its TP2 weight plumbing is receipted over the real pack on two GB10 nodes " ++ + "(the two ranks' checkpoint bytes exactly tile the plan's measured 166928433556) — and this build " ++ + "attaches the launch graph: graph.Engine runs the per-layer device forward per row (embedding → " ++ + "q/k/v through the mainline EXL3 gemms with per-head qk norm + partial rope 1e7/64/0.5, the indexer's " ++ + "side plane, the block-sparse selection + attention, dense MLP layers 0-2, the sigmoid+bias router " ++ + "top-4-of-128 with the shared experts through tf_m3_swigluoai, final norm, lm_head) with serve.zig's " ++ + "own published normRow/ropeRow/route/f16r as the host glue, and open() returns the Loaded handle over " ++ + "it — the fixture smoke (minimax-m3-smoke) and the token smoke (minimax-m3-token, the Stage-5 rank " ++ + "recipe: rank 0 layers 0-29 + embed + final norm + lm_head, rank 1 layers 30-59, NCCL handoffs) grade " ++ + "it — but this stage still refuses loudly rather than fake anything it does not have: draft rows and " ++ + "drafted verify windows refuse (error.LanesNotAttached: no draft head — the draft/spec-decode surface " ++ + "is unimplemented), a sampled draw spec and grammar masks refuse (error.SamplingNotAttached: the keyed " ++ + "sampler is not wired for this family's lanes), a second stream refuses while one is live " ++ + "(error.LanesNotAttached: one stream at a time), tree keeps refuse (prefix rollbacks only), TP2 booted " ++ + "through open() now ATTACHES since Stage-11 — the TF_MINIMAXM3_RANK/TF_MINIMAXM3_WORLD env server " ++ + "boot serves when the place is the receipted world-2 split: rank 0 returns the Loaded the HTTP " ++ + "server drives on this lane surface, rank 1 mirrors rank 0's rows through the per-row control wire " ++ + "and never serves (glm's rank pattern, additively) — a place outside the receipted split still " ++ + "refuses (error.TpServingNotAttached), a world-2 boot without NCCL_COMM_ID refuses (error.NoCommId) — and this " ++ + "build's SERVING lane surface (the Loaded handle's vtable) serves a stream through prefill/first/" ++ + "queue/read/verify/keep/release with TP1 through open() and the split-pack TP2 receipt through " ++ + "openTP2 (tf-cuda-test minimax-m3-lane-tp2 , rank 1 first over the NCCL id, " ++ + "the Stage-9 receipts) — and contexts wider than the smoke's 4096-key pool refuse (error.ContextTooWide)"; + +/// The serve boundary this build's refusals draw (see serve_boundary): the lane surface serves +/// one greedy stream; drafts, sampled draws, second streams, tree keeps and the TP2 server boot +/// stay loud refusals, each with its own words. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.ServeNotAttached => serve_boundary, + error.LanesNotAttached => "minimax_m3_vl's lane surface serves one greedy stream through the attached graph (prefill/first/queue/read/verify/keep/release — the boundary tests grade a short prompt token-in token-out); the calls that refuse here are the ones this stage does not have: draft rows and drafted verify windows (no draft head — the draft/spec-decode surface is unimplemented), a second stream while one is live, a verify batch wider than one window, and tree keeps", + error.SamplingNotAttached => "minimax_m3_vl's lane surface serves greedy draws only in this stage: a Sampling spec (temperature/top_k/top_p) or a grammar mask refuses here rather than serve an ungraded draw — the keyed sampler is a later pass", + error.TpServingNotAttached => "minimax_m3_vl's TF_MINIMAXM3_RANK/TF_MINIMAXM3_WORLD env server boot attaches as the receipted world-2 split (rank 0 serves the lane surface, rank 1 mirrors rank 0's rows; NCCL_COMM_ID names rank 0's id host:port). What still refuses: an env place outside that split (one env var alone, a non-numeric rank/world, a rank at or past its world, a world other than 2), while TP2 serving receipts live at `tf-cuda-test minimax-m3-token ` (the Stage-6/7 token receipts) and openTP2's lane surface — `tf-cuda-test minimax-m3-lane-tp2 ` (the Stage-9 receipts: the same lane vtable answering over the split pack, rank 1 first)", + error.NoCommId => "minimax_m3_vl's world-2 server boot meets its rank over NCCL: set NCCL_COMM_ID to rank 0's address:port on every rank (the receipted dance serves rank 0's fresh UniqueId there; rank 1 connects first)", + error.ContextTooWide => "the launch graph's decode pool this stage is 4096 keys: give a context at or under it (the wide paged prompt path is a later stage)", + error.LayerCapTP2 => "the launch graph's TP1-only layer cap refuses under TP2: each rank owns its exact half (the cap serves the Stage-8 grade's reduced real-pack stack)", + error.LayerCapTooWide => "the layer cap exceeds the rank's owned layers", + error.CapturedKernels => "TENSORFOLD_CUDA_KERNELS points at a captured set: minimax_m3_vl serves with its own glued kernels only; unset it", + error.NotMiniMaxM3VL => "this checkpoint's shapes are not MiniMax-M3.1-Flash's (or its flags are): the family reads the M3.1 VL MoE pack", + error.MissingTensor, error.UnexpectedTensor => "the checkpoint misses a tensor the family reads, or one has a different dtype or shape: this pack is not the one admission checked", + error.MissingTrellis, error.SurplusTrellis => "an EXL3 trellis group needed by the plan is missing, or one is admitted that no module reads: the checkpoint is incomplete or from another pack", + error.ForeignTrellis, error.NotMul1, error.HeadBits => "a quantized tensor falls outside this pack's EXL3 contract (mul1 codebook, the head's 8 halves): not Mia-AiLab's M3.1-Flash line", + else => null, + }; +} + +/// The attached graph's own state in the Loaded handle: the engine plus THE one greedy stream's +/// lane state. The lane surface serves one stream at a time (Loaded.facts.max_streams), greedy +/// draws only (a `Sampling` spec refuses), and everything it cannot do honestly it refuses by +/// name — the refusals stay part of the surface (drafts, trees, second streams), not decoration. +const Attached = struct { + e: *graph.Engine, + + /// The stream the lane surface is serving (null: none). A second stream's calls refuse. + live: ?*lanes.Stream = null, + /// World 2's NCCL hold (the env server boot): the library + communicator the engine's + /// collectives ride, freed on the surface's close (the server's clean stop). + hold: ?*NcclHold = null, + /// Drawn tokens waiting for their `read`: the handle -> token slab (eager greedy draws — + /// the engine produces the logits row synchronously, so the draw is on the host already). + pending: [pending_cap]Draw = undefined, + pending_n: usize = 0, + next_handle: u64 = 0, + + /// How many drawn tokens may sit unread at once (the one-token round's working set; a + /// queue burst wider than this refuses rather than silently forgetting a draw). + const pending_cap: usize = 8; + const Draw = struct { handle: u64 = 0, pos: u64 = 0, tok: u32 = 0 }; + + fn live_ok(a: *Attached, s: *lanes.Stream) !void { + if (a.live != s) return error.LanesNotAttached; + } + /// The greedy draw the fork's contract makes from the logits row: first-max argmax. A + /// non-finite lane in the row means the row is not servable: refuse. + fn drawGreedy(a: *Attached, pos: u64) !u64 { + if (a.pending_n >= pending_cap) return error.LaneDrawsFull; + var bad = false; + for (a.e.logits) |v| { + if (std.math.isNan(v) or std.math.isInf(v)) bad = true; + } + if (bad) return error.RowNotFinite; + var best: usize = 0; + var bv = -std.math.inf(f32); + for (a.e.logits, 0..) |v, i| { + if (v > bv) { + bv = v; + best = i; + } + } + a.next_handle += 1; + a.pending[a.pending_n] = .{ .handle = a.next_handle, .pos = pos, .tok = @intCast(best) }; + a.pending_n += 1; + return a.next_handle; + } + fn pendingOf(a: *Attached, handle: u64) !u32 { + for (a.pending[0..a.pending_n]) |d| { + if (d.handle == handle) return d.tok; + } + return error.LaneHandleUnknown; + } + fn pendingTake(a: *Attached, handle: u64) !u32 { + var tok: u32 = 0; + for (a.pending[0..a.pending_n], 0..) |d, i| { + if (d.handle == handle) { + tok = d.tok; + a.pending[i] = a.pending[a.pending_n - 1]; + a.pending_n -= 1; + return tok; + } + } + return error.LaneHandleUnknown; + } +}; + +// ------------------------------------------------------------ the lane rows (the serving half) +// +// The contract this surface serves (each row's own words in the error paths): +// prefill: the stream's whole prompt through the attached graph, rows keyed at the engine's +// own cache length (a stream prefills once per cache life; a reuse restore never +// happened on this surface, so `s.cached` must be 0). +// first: the prompt's last row's greedy draw; the caller's `position` must be the engine's +// cache length (the row the drawn token occupies when it is fed back). +// queue: feed one token (a host value or this surface's own handle) and draw the greedy +// continuation; the fed row keys at the engine's own next cache length — the host's +// window `position` is advisory (it repeats across plain rounds, glm's pattern). +// read: a drawn token's value (the draw is already on the host; nothing waits). +// verify: ONE one-row window: feed the pending token, draw the greedy continuation. Drafted +// windows (held rows, tree tokens), grammar masks and multi-window batches refuse — +// the draft/spec-decode surface stays unimplemented for this family. +// keep: ONE window whose path is the kept prefix 0..kept-1: the engine's cache length rolls +// back (stale cache rows past the kept length are never read — every row's forward +// keys at the engine's own length and rewrites its row first). Tree paths refuse. +// release: the stream leaves the rounds: the engine's cache length resets to 0 and the draws +// drop. The pool's rows are positional and rewritten per row, so nothing leaks. +// Draft rows (`draft`), sampled draws (a non-null `Sampling`), and every second-stream call +// refuse by name below — honestly, with the exact reason this build does not do them. + +/// The lane surface's row under world 2: rank 0 leads (control word + row), rank 1 (the +/// harness's receipted lane-tp2 surface) mirrors through the control wire — its fed value +/// stays meaningless. World 1 keeps the plain `put`. +fn surfacePut(a: *Attached, tok: u32) !void { + if (a.e.world == 2 and a.e.rank == 0) return a.e.leaderRow(tok); + if (a.e.world == 2) return a.e.mirrorStep(); + return a.e.put(tok); +} + +fn laneSelf(ptr: *anyopaque) *Attached { + return @ptrCast(@alignCast(ptr)); +} + +fn laneRequireGreedy(s: *lanes.Stream) !void { + if (s.sampling != null) return error.SamplingNotAttached; + if (s.drafts) return error.LanesNotAttached; +} + +/// The stream's whole prompt through the graph, one row at a time (the engine keys every row +/// at its own cache length; a resumed/edited prompt whose lengths disagree with the graph's +/// would silently mis-rope — refused rather than served). +fn laneServePrefill(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { + const a = laneSelf(ptr); + if (a.live != null and a.live != s) return error.LanesNotAttached; + if (a.e.len != 0) return error.StreamAlreadyPrefilled; // release() first: one prompt per cache life + try laneRequireGreedy(s); + if (s.prompt_len == 0) return error.EmptyPrompt; + if (s.cached != 0) return error.PromptCacheMismatch; // this surface never restored a prefix + a.live = s; + const prompt = s.prompt(); + if (prompt.len != s.prompt_len) return error.PromptCacheMismatch; + if (a.e.len + prompt.len > graph.keys_cap) + return error.ContextTooWide; + for (prompt) |tok| try surfacePut(a, tok); + s.fill_at.store(@intCast(s.prompt_len), .monotonic); +} + +fn laneServeFirst(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { + const a = laneSelf(ptr); + try a.live_ok(s); + if (a.e.len == 0) return error.EmptyPrompt; + if (position != a.e.len) return error.RowKeying; // the draw keys at the next row's position + return a.drawGreedy(position); +} + +fn laneServeQueue(ptr: *anyopaque, s: *lanes.Stream, feed: lanes.backend.Feed, position: u64) anyerror!u64 { + _ = position; // the host's window positions are advisory and repeat across plain rounds (glm's + // pattern); the rows are positional and sequential, so this surface keys its own cache length + const a = laneSelf(ptr); + try a.live_ok(s); + const tok: u32 = switch (feed) { + .value => |v| v, + .handle => |h| try a.pendingOf(h), // this surface's own drawn token, still unread + }; + if (a.e.len == 0) return error.EmptyPrompt; // the fed row keys after the prompt's rows + if (a.e.len + 1 > graph.keys_cap) return error.ContextTooWide; + try surfacePut(a, tok); // the fed row lands at the engine's next cache length + return a.drawGreedy(a.e.len); +} + +fn laneServeRead(ptr: *anyopaque, handle: u64) anyerror!u32 { + return laneSelf(ptr).pendingTake(handle); +} + +fn laneServeVerify(ptr: *anyopaque, windows: []const lanes.backend.Window, out: []lanes.backend.Verified) anyerror!void { + const a = laneSelf(ptr); + if (windows.len != 1) return error.LanesNotAttached; // one stream at a time: no batching + const w = windows[0]; + if (w.held > 0 or w.tokens.len > 0) return error.LanesNotAttached; // no draft head: drafted verify refuses + if (w.masks.len > 0) return error.SamplingNotAttached; // structured output masks refuse + try a.live_ok(w.stream); + // the host's advisory window positions repeat across plain rounds (glm's pattern); the + // one-row undrafted window keys at this surface's own cache length, not the host's copy + if (out.len < windows.len) return error.LanesNotAttached; + const h = try laneServeQueue(ptr, w.stream, .{ .value = w.pending }, w.positions[0]); + out[0].sampled[0] = try a.pendingTake(h); +} + +fn laneServeKeep(ptr: *anyopaque, windows: []const lanes.backend.Window, paths: []const []const u32) anyerror!void { + const a = laneSelf(ptr); + if (windows.len != 1) return error.LanesNotAttached; // one stream at a time + if (paths.len != 1) return error.LaneHandleUnknown; + const path = paths[0]; + if (path.len == 0) return error.RowKeying; // an ambiguous rollback refuses + for (path, 0..) |row, i| { + if (row != i) return error.LanesNotAttached; // tree paths refuse: prefix only + } + if (path.len > a.e.len) return error.RowKeying; + a.e.len = @intCast(path.len); // the cache rolls back; stale rows rewrite before any read + a.e.leaderKeep(a.e.len) catch |err| { + // world 2: the mirror must hear the rollback or it keys stale rows. A failed control + // exchange here leaves the NCCL wire broken for every later row; surface it now. + std.log.err("minimax_m3_vl: the keep's control exchange failed ({s}); the TP2 mirror will desync", .{@errorName(err)}); + return err; + }; +} + +fn laneServeRelease(ptr: *anyopaque, s: *lanes.Stream) void { + laneServeFree(ptr, s); +} + +fn laneServeFree(ptr: *anyopaque, s: *lanes.Stream) void { + const a = laneSelf(ptr); + if (a.live == s) { + a.live = null; + a.pending_n = 0; + a.e.len = 0; // the pool's rows are positional: the next prefill rewrites from row 0 + if (a.e.world == 2 and a.e.rank == 0) a.e.leaderReset() catch { + // world 2: the mirror must hear the reset. The NCCL wire is broken if this fails + // (every later exchange fails too); the release itself must stay non-throwing. + }; + s.fill_at.store(0, .monotonic); + } +} + +// The window `positions` repeat across the host's plain rounds; `laneServeQueue` keys its own +// rows, so nothing compares against the host's copy here. + +// The draft rows stay loud refusals: this family has no draft head, and the receipt words say +// exactly what refuses and what serves. +const lane_refuse_draft = struct { + fn refuse(_: *anyopaque, _: []const lanes.backend.DraftRequest) anyerror!void { + return error.LanesNotAttached; + } +}.refuse; + +const lane_vtable = lanes.backend.Backend.VTable{ + .prefill = laneServePrefill, + .first = laneServeFirst, + .queue = laneServeQueue, + .read = laneServeRead, + .verify = laneServeVerify, + .keep = laneServeKeep, + .draft = lane_refuse_draft, + .release = laneServeRelease, +}; + +fn closeGraph(ptr: *anyopaque) void { + const a: *Attached = @ptrCast(@alignCast(ptr)); + const g = a.e.gpa; + a.e.deinit(); + g.destroy(a.e); + if (a.hold) |h| { // the env boot's NCCL: the communicator aborts, the library closes + _ = h.lib.api.ncclCommAbort(h.comm); + h.lib.close(); + g.destroy(h); + } + g.destroy(a); +} + +/// The env boot's NCCL lifetime: the library and communicator one surface close releases. +const NcclHold = struct { lib: cuda.nccl.Library, comm: cuda.nccl.Comm }; + +/// The TF_MINIMAXM3_RANK/WORLD env SERVER boot (Stage-11, glm's rank pattern mirrored over the +/// receipted rank-1-first id dance): rank 0 returns the Loaded the server drives on the ONE +/// greedy stream's lane surface; rank 1 never returns — it mirrors rank 0's rows through the +/// control wire until rank 0's tp_watch line closes (a clean stop ends it too) or NCCL breaks. +/// A place outside the receipted world-2 split refuses by exact name, before any device work. +fn openServerRank(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, o: Options) !Loaded { + const p = try placeOfBoot(std.c.getenv("TF_MINIMAXM3_RANK"), std.c.getenv("TF_MINIMAXM3_WORLD")); + const addr = std.mem.span(std.c.getenv("NCCL_COMM_ID") orelse return error.NoCommId); + const cap = capacity(o.context); + if (cap > graph.keys_cap) return error.ContextTooWide; // the smoke pool's window, as ever + try cuda.tp_watch.start(io, p.rank, p.world); // the ranks' lines before NCCL: a dying rank stops every rank + // NCCL's cuMem host buffers wore out GB10's driver pool under load (glm's receipted defaults) + if (std.c.getenv("NCCL_RUNTIME_CONNECT") == null) _ = c_setenv("NCCL_RUNTIME_CONNECT", "0"); + if (std.c.getenv("NCCL_CUMEM_HOST_ENABLE") == null) _ = c_setenv("NCCL_CUMEM_HOST_ENABLE", "0"); + if (std.c.getenv("NCCL_CUMEM_ENABLE") == null) _ = c_setenv("NCCL_CUMEM_ENABLE", "0"); + std.log.info("minimax_m3_vl server rank {d} of world {d}: NCCL_COMM_ID {s} (the receipted rank-1-first id dance)", .{ p.rank, p.world, addr }); + const hold = try gpa.create(NcclHold); + errdefer gpa.destroy(hold); + hold.lib = try cuda.nccl.Library.open(); + errdefer hold.lib.close(); + var id: cuda.nccl.UniqueId = undefined; + const id_bytes = std.mem.asBytes(&id); + if (p.rank == 0) { + try hold.lib.check(hold.lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + serveCommId(io, addr, id_bytes) catch |err| { + cuda.tp_watch.fail(@errorName(err)); // rank 1 stops with this rank + return err; + }; + } else { + receiveCommId(io, addr, id_bytes) catch |err| { + cuda.tp_watch.fail(@errorName(err)); + return err; + }; + } + try hold.lib.check(hold.lib.api.ncclCommInitRank(&hold.comm, @intCast(p.world), id, @intCast(p.rank)), "ncclCommInitRank"); + const ep = try gpa.create(graph.Engine); + errdefer gpa.destroy(ep); + ep.* = graph.Engine.build(gpa, io, ctx.d, dir, p.rank, p.world, .{ .lib = &hold.lib, .comm = hold.comm }) catch |err| { + cuda.tp_watch.fail(@errorName(err)); + return err; + }; + errdefer ep.deinit(); + if (p.rank != 0) { + std.log.info("minimax_m3_vl server rank 1: mirroring rank 0's rows (its own cache length rides the control wire)", .{}); + while (true) ep.mirrorStep() catch |err| { + std.log.err("minimax_m3_vl server rank 1: the row mirror stopped ({s})", .{@errorName(err)}); + cuda.tp_watch.fail(@errorName(err)); // rank 0 stops with this rank's error + std.c._exit(1); // a stream may hold kernels waiting on a gone rank: a normal exit would join them + }; + } + const own = try gpa.create(Attached); + errdefer gpa.destroy(own); + own.* = .{ .e = ep, .live = null, .hold = hold, .pending = undefined, .pending_n = 0, .next_handle = 0 }; + return Loaded{ + .backend = .{ .ptr = own, .vtable = &lane_vtable }, + .facts = .{ .exact_width = 1, .drafts = 0, .max_streams = 1 }, + .rows = 1, + .stream_bytes = 0, + .free_streams = 1, + .ctx = own, + .deinit = closeGraph, + .lone = null, + .target = ep, + }; +} + +extern "c" fn setenv(name: [*:0]const u8, value: [*:0]const u8, overwrite: c_int) c_int; +fn c_setenv(name: [*:0]const u8, value: [*:0]const u8) c_int { + return setenv(name, value, 0); +} + +// The receipted id dance (the lane-tp2 harness's serveId/receiveId, family-owned here): rank 0 +// serves the freshly minted UniqueId's bytes on NCCL_COMM_ID's exact host:port, rank 1 connects +// (retrying until rank 0's bind, up to one minute) and reads them. Caps bound both waits, so a +// rank whose peer never comes refuses instead of hanging the boot. + +fn serveCommId(io: std.Io, addr_text: []const u8, bytes: []const u8) !void { + const at = std.mem.lastIndexOfScalar(u8, addr_text, ':') orelse return error.BadAddress; + const ip = try std.Io.net.IpAddress.parseIp4(addr_text[0..at], try std.fmt.parseInt(u16, addr_text[at + 1 ..], 10)); + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&ip, &storage); + const rc = std.posix.system.socket(std.posix.AF.INET, std.posix.SOCK.STREAM, 0); + if (std.posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const fd: std.posix.socket_t = @intCast(rc); + defer _ = std.posix.system.close(fd); + const one: c_int = 1; + try std.posix.setsockopt(fd, std.posix.SOL.SOCKET, std.posix.SO.REUSEADDR, std.mem.asBytes(&one)); + switch (std.posix.errno(std.posix.system.bind(fd, &storage.any, len))) { + .SUCCESS => {}, + .ADDRINUSE => return error.AddressInUse, + else => return error.BindFailed, + } + if (std.posix.errno(std.posix.system.listen(fd, 1)) != .SUCCESS) return error.ListenFailed; + var waited: u64 = 0; + while (true) { + var fds = [_]std.posix.pollfd{.{ .fd = fd, .events = std.posix.POLL.IN, .revents = 0 }}; + const ready = std.posix.poll(&fds, 500) catch return error.AcceptFailed; + if (ready > 0) break; + waited += 500; + if (waited > 70_000) return error.PeerTimeout; + } + var peer: std.Io.Threaded.PosixAddress = undefined; + var peer_len: std.posix.socklen_t = @sizeOf(std.Io.Threaded.PosixAddress); + const rc2 = std.posix.system.accept(fd, &peer.any, &peer_len); + if (std.posix.errno(rc2) != .SUCCESS) return error.AcceptFailed; + const cfd: std.posix.socket_t = @intCast(rc2); + defer _ = std.posix.system.close(cfd); + var off: usize = 0; + while (off < bytes.len) { + const n = std.posix.system.write(cfd, bytes.ptr + off, bytes.len - off); + if (std.posix.errno(n) != .SUCCESS) return error.WriteFailed; + off += @intCast(n); + } + _ = io; +} + +fn receiveCommId(io: std.Io, addr_text: []const u8, bytes: []u8) !void { + const at = std.mem.lastIndexOfScalar(u8, addr_text, ':') orelse return error.BadAddress; + const ip = try std.Io.net.IpAddress.parseIp4(addr_text[0..at], try std.fmt.parseInt(u16, addr_text[at + 1 ..], 10)); + var fd: std.posix.socket_t = -1; + defer { + if (fd >= 0) _ = std.posix.system.close(fd); + } + var connect_ms: u64 = 0; // rank 1 boots first: retry the connect until rank 0 binds (the receipted ~6 s gap) + while (fd < 0) { + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&ip, &storage); + const rc = std.posix.system.socket(std.posix.AF.INET, std.posix.SOCK.STREAM, 0); + if (std.posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const c: std.posix.socket_t = @intCast(rc); + if (std.posix.errno(std.posix.system.connect(c, &storage.any, len)) == .SUCCESS) { + fd = c; + } else { + _ = std.posix.system.close(c); + if (connect_ms >= 70_000) return error.PeerTimeout; + std.Io.sleep(io, .fromMilliseconds(200), .awake) catch {}; + connect_ms += 200; + } + } + var off: usize = 0; + var read_waited: u64 = 0; + while (off < bytes.len) { + var fds = [_]std.posix.pollfd{.{ .fd = fd, .events = std.posix.POLL.IN, .revents = 0 }}; + const ready = std.posix.poll(&fds, 500) catch return error.ReadFailed; + if (ready == 0) { + read_waited += 500; + if (read_waited > 60_000) return error.PeerTimeout; + continue; + } + const n = std.posix.system.read(fd, bytes.ptr + off, bytes.len - off); + switch (std.posix.errno(n)) { + .SUCCESS => { + if (n == 0) return error.PeerClosed; + off += @intCast(n); + }, + .INTR => {}, + else => return error.ReadFailed, + } + } +} + +/// The checkpoint, the shared kernels and every stream's caches — this stage: admission, then the +/// attached launch graph serving ONE greedy stream through the lane surface (TP1 through open(); +/// TP2 through the harness command), with drafts/sampled draws/second streams/TP2-open refusing by name. +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + if (kernels != null) return error.CapturedKernels; + if (!cuda.kernels.available) return error.ServeNotAttached; // a build without the glued kernels: the graph needs the device + // the ranks: TP1 unchanged; the TF_MINIMAXM3_RANK/WORLD env SERVER boot now attaches here + // (glm's pattern, additively: rank 0 returns the Loaded the server serves on its lane + // surface, rank 1 mirrors rank 0's rows through the control wire and never serves). A + // place outside the receipted world-2 split still refuses, by exact name — before any + // device work. + const env_rank = std.c.getenv("TF_MINIMAXM3_RANK"); + const env_world = std.c.getenv("TF_MINIMAXM3_WORLD"); + if (env_rank != null or env_world != null) { + _ = try placeOfBoot(env_rank, env_world); // the exact-name refusals fire before anything device + if (kernels != null) return error.CapturedKernels; + if (!cuda.kernels.available) return error.ServeNotAttached; + return openServerRank(gpa, io, ctx, dir, o); + } + if (kernels != null) return error.CapturedKernels; + if (!cuda.kernels.available) return error.ServeNotAttached; + return openLane(gpa, io, ctx, dir, o); +} + +/// The env boot's place: both vars, numeric, inside the receipted world-2 split — or refuse +/// by exact name (the old honest refusal stays live for any place that is not the receipted +/// split; world 1 boots through open() WITHOUT the env). +fn placeOfBoot(rank_env: ?[*:0]const u8, world_env: ?[*:0]const u8) !rank_plan.Place { + const r = rank_env orelse return error.TpServingNotAttached; + const w = world_env orelse return error.TpServingNotAttached; + const p = rank_plan.checked(.{ + .rank = std.fmt.parseInt(u32, std.mem.span(r), 10) catch return error.TpServingNotAttached, + .world = std.fmt.parseInt(u32, std.mem.span(w), 10) catch return error.TpServingNotAttached, + }) catch return error.TpServingNotAttached; + if (p.world != rank_plan.max_world) return error.TpServingNotAttached; + return p; +} + +/// TP1's own admission + attached graph body (the pre-Stage-11 open() body, byte-equal). +fn openLane(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, o: Options) !Loaded { + var a = try admitDir(gpa, io, dir); + defer a.deinit(); + if (o.context < 1) return error.ServeNotAttached; // below one block of context there is no row to serve + // the ranks: TP1 (rank 0 of world 1); the env server boot routes through openServerRank above + const rank: u32 = 0; + const world: u32 = 1; + const cap = capacity(o.context); + if (cap > graph.keys_cap) return error.ContextTooWide; // the smoke pool's window: a wider prefill path is a later stage + const own = try gpa.create(Attached); + errdefer gpa.destroy(own); + const ep = try gpa.create(graph.Engine); + errdefer gpa.destroy(ep); + ep.* = try graph.Engine.build(gpa, io, ctx.d, dir, rank, world, .{ .layer_cap = o.layer_cap }); + own.* = .{ .e = ep, .live = null, .pending = undefined, .pending_n = 0, .next_handle = 0 }; + return Loaded{ + .backend = .{ .ptr = own, .vtable = &lane_vtable }, + .facts = .{ .exact_width = 1, .drafts = 0, .max_streams = 1 }, + .rows = 1, + .stream_bytes = 0, + .free_streams = 1, + .ctx = own, + .deinit = closeGraph, + .lone = null, + .target = ep, + }; +} + +/// The TP2 lane-surface boot (Stage-9): the harness's rank plumbing — the caller builds the NCCL +/// library and the rank's communicator first (rank 1 first over the id, exactly the token-harness +/// recipe), and this entry wraps open()'s own admission + the SAME lane vtable over the rank's +/// half of the split pack. open() itself keeps refusing the TF_MINIMAXM3_RANK/WORLD env boot; the +/// rank boot here is explicit, receipted, and serves only through the lane surface like TP1 does. +/// Rank 0 argues the greedy draws (its `logits` row is the only real one); rank 1's surface keeps +/// the row lockstep (every `put` is a lockstep row; the fed token's value is rank 0's business). +pub const RankBoot = struct { lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, rank: u32, world: u32 }; + +pub fn openTP2(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, o: Options, boot: RankBoot) !Loaded { + if (boot.world != 2 or boot.rank >= boot.world) return error.NoRank; // the receipted split: exactly two ranks + if (!cuda.kernels.available) return error.ServeNotAttached; // a build without the glued kernels: the graph needs the device + var a = try admitDir(gpa, io, dir); + defer a.deinit(); + if (o.context < 1) return error.ServeNotAttached; + if (o.layer_cap != 0) return error.LayerCapTP2; // under TP2 each rank owns its exact half + const cap = capacity(o.context); + if (cap > graph.keys_cap) return error.ContextTooWide; // the smoke pool's window: a wider prefill path is a later stage + const own = try gpa.create(Attached); + errdefer gpa.destroy(own); + const ep = try gpa.create(graph.Engine); + errdefer gpa.destroy(ep); + ep.* = try graph.Engine.build(gpa, io, ctx.d, dir, boot.rank, boot.world, .{ .lib = boot.lib, .comm = boot.comm }); + own.* = .{ .e = ep, .live = null, .pending = undefined, .pending_n = 0, .next_handle = 0 }; + return Loaded{ + .backend = .{ .ptr = own, .vtable = &lane_vtable }, + .facts = .{ .exact_width = 1, .drafts = 0, .max_streams = 1 }, + .rows = 1, + .stream_bytes = 0, + .free_streams = 1, + .ctx = own, + .deinit = closeGraph, + .lone = null, + .target = ep, + }; +} + +test "the family dispatches over the staged pack's model_type, formats and admission" { + const gpa = std.testing.allocator; + try std.testing.expectEqualStrings("minimax_m3_vl", model_type); + try std.testing.expectEqual(@as(usize, 1), formats.len); + try std.testing.expectEqualStrings("minimax_m3_vl", native.model_type); // the registry reads these two fields + try std.testing.expectEqualStrings("exl3", native.formats[0]); + try std.testing.expectEqual(@as(usize, registry_index), 7); // in native/cuda.zig's engine-dispatch registry + const c = try admit(gpa, @embedFile("test_pack_config.json"), null); + try std.testing.expect(c.exl3 and config.whole(c)); // 3.06 bpw, mul1, mixed widths, the 128-column vocab cut +} + +/// The openWith registry's tuple position (native/cuda.zig's `registry`): kept asserted so both move together. +pub const registry_index = 7; + +test "admission runs the whole checkpoint and open refuses exactly at the serve boundary" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try weights.checkpointFixture(gpa, io, tmp.dir, false); + const path = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(path); + var a = try admitDir(gpa, io, path); + defer a.deinit(); + try std.testing.expect(a.c.exl3); // the inline quantization_config named EXL3 + try std.testing.expectEqual(config.Codebook.mul1, a.c.qc_codebook); + try std.testing.expectEqual(@as(usize, 27), a.groups.groups.len); // the discovered trellis groups + try std.testing.expectEqual(@as(usize, 47), a.load.items.len); // the whole plan, in load order + try std.testing.expect(a.bytes() > 0); + // no serving without a context row (or without the glued kernels on a host build); a captured + // kernel set also refuses — both paths stay before any device work, so the test is device-free + try std.testing.expectError(error.ServeNotAttached, open(gpa, io, undefined, path, null, .{ .context = 0, .drafts = false })); + try std.testing.expectError(error.CapturedKernels, open(gpa, io, undefined, path, "captured", .{ .context = 4096, .drafts = false })); + const words = explain(null, error.ServeNotAttached).?; + try std.testing.expectStringStartsWith(words, "minimax_m3_vl passed CUDA admission"); + try std.testing.expect(std.mem.indexOf(u8, words, "graded on the host") != null); // the host forward is real + try std.testing.expect(std.mem.indexOf(u8, words, "graded on the GPU") != null); // the device kernels are too + try std.testing.expect(std.mem.indexOf(u8, words, "tf_m3_swigluoai") != null); // the parity receipts name them + try std.testing.expect(std.mem.indexOf(u8, words, "normRow/ropeRow/route/f16r") != null); // the glue is serve.zig's own fns + try std.testing.expect(std.mem.indexOf(u8, words, "attaches the launch graph") != null); // open() carries it + try std.testing.expect(std.mem.indexOf(u8, words, "minimax-m3-token") != null); // the token receipts' command + try std.testing.expect(std.mem.indexOf(u8, words, "SERVING lane surface") != null); // Stage-8: the lanes serve + try std.testing.expect(std.mem.indexOf(u8, words, "prefill/first/") != null); // the served rows are named + try std.testing.expect(std.mem.indexOf(u8, words, "LanesNotAttached") != null); // drafts+second streams still refuse + try std.testing.expect(std.mem.indexOf(u8, words, "SamplingNotAttached") != null); // sampled draws refuse + try std.testing.expect(std.mem.indexOf(u8, words, "TpServingNotAttached") != null); // and the TP2 server boot + try std.testing.expect(std.mem.indexOf(u8, words, "ContextTooWide") != null); // and the pool's window + try std.testing.expect(explain(null, error.LanesNotAttached) != null); + try std.testing.expect(explain(null, error.SamplingNotAttached) != null); + try std.testing.expect(explain(null, error.TpServingNotAttached) != null); + try std.testing.expect(explain(null, error.ContextTooWide) != null); + try std.testing.expect(explain(null, error.LayerCapTP2) != null); + try std.testing.expect(explain(null, error.LayerCapTooWide) != null); + try std.testing.expect(explain(null, error.BadConfig) == null); // not this family's wording +} + +test "the TP2 lane boot refuses a bad rank shape before admitting anything (device-free)" { + const gpa = std.testing.allocator; + // the NoRank shape check fires before admission or any device work: no CUDA context touched + try std.testing.expectError(error.NoRank, openTP2(gpa, undefined, undefined, "nowhere", .{ .context = 128, .drafts = false }, .{ .lib = undefined, .comm = null, .rank = 2, .world = 2 })); + try std.testing.expectError(error.NoRank, openTP2(gpa, undefined, undefined, "nowhere", .{ .context = 128, .drafts = false }, .{ .lib = undefined, .comm = null, .rank = 0, .world = 1 })); + try std.testing.expect(explain(null, error.TpServingNotAttached) != null); + try std.testing.expect(std.mem.indexOf(u8, explain(null, error.TpServingNotAttached).?, "minimax-m3-lane-tp2") != null); +} + +extern "c" fn unsetenv(name: [*:0]const u8) c_int; + +test "the env server boot attaches only as the receipted world-2 split (device-free refusals)" { + // placeOfBoot reads env only: nothing device, nothing io. The receipted places validate; + // every place outside the split refuses TpServingNotAttached (the exact-name law holds). + try std.testing.expectEqual(rank_plan.Place{ .rank = 0, .world = 2 }, try placeOfBoot("0", "2")); + try std.testing.expectEqual(rank_plan.Place{ .rank = 1, .world = 2 }, try placeOfBoot("1", "2")); + try std.testing.expectError(error.TpServingNotAttached, placeOfBoot(null, "2")); // one env var alone + try std.testing.expectError(error.TpServingNotAttached, placeOfBoot("0", null)); + try std.testing.expectError(error.TpServingNotAttached, placeOfBoot("x", "2")); // non-numeric + try std.testing.expectError(error.TpServingNotAttached, placeOfBoot("2", "2")); // a rank at its world + try std.testing.expectError(error.TpServingNotAttached, placeOfBoot("0", "3")); // not the receipted split + try std.testing.expectError(error.TpServingNotAttached, placeOfBoot("0", "1")); // world 1 boots WITHOUT the env + try std.testing.expect(explain(null, error.NoCommId) != null); // world 2 but no NCCL id: its own exact name +} + +test { + _ = @import("serve.zig"); // the family's forward and its f64-reference grade run with the suite + _ = @import("config.zig"); + _ = @import("weights.zig"); + _ = @import("exl3.zig"); + _ = @import("sparse.zig"); + _ = @import("cuda_sparse.zig"); + _ = @import("cuda_serve.zig"); + _ = @import("shard.zig"); // the TP2 rank split's tests ride the same module + _ = @import("graph.zig"); // the launch graph and its helpers compile with the suite +} \ No newline at end of file diff --git a/zig/src/families/minimax_m3_vl/cuda_serve.zig b/zig/src/families/minimax_m3_vl/cuda_serve.zig new file mode 100644 index 0000000000..6bb2a51210 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/cuda_serve.zig @@ -0,0 +1,49 @@ +//! MiniMax M3.1's device serve glue (the family-owned launches the mainline EXL3 modules do not +//! carry): the swigluoai activation that stands between the grouped gate/up GEMM partials and the +//! down GEMM, graded against serve.zig's swigluoai. The rest of the forward is mainline code: +//! per-token EXL3 decode lines (`cuda.exl3.Dense`, mul1 odd half steps and the head's 8 halves +//! included), the grouped routed experts (`cuda.exl3_experts`), and the MSA indexer launches +//! (`cuda_sparse.zig`). + +const std = @import("std"); +const cuda = @import("cuda"); +const serve = @import("serve.zig"); + +/// The family's glue kernel from the minimax_m3 image. +pub const Kernels = struct { + swigluoai: cuda.Function, + + pub fn load(m: cuda.Module) !Kernels { + var k: Kernels = undefined; + k.swigluoai = try m.function("tf_m3_swigluoai"); + return k; + } +}; + +/// The M3.1 activation rows: (rows, n) fp16 gate and up rows -> (rows, n) fp16, the fork's +/// activation.cu:124 contract (`limit` 7, alpha hardwired 1.702 in the kernel). The pack's dense, +/// shared and routed widths are multiples of 512, so the launch never carries a tail block. +pub fn swigluoai( + ks: Kernels, + stream: cuda.Stream, + g: u64, + u: u64, + out: u64, + rows: usize, + n: usize, + limit: f32, +) !void { + var a: cuda.Args = .{}; + inline for (.{ g, u, out, @as(i64, @intCast(n)), limit }) |v| a.add(v); + const blocks = (n + 512 - 1) / 512; // 128 threads a block, 4 elements a lane + try cuda.launch.launch(ks.swigluoai, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(blocks) }, .block = .{ .x = 128 } }, stream, &a); +} + +test "the glue's activation lands on serve.zig's swigluoai contract at the pack's numbers" { + // the contract's fingerprints: u clamps both ways then +1, g never passes the limit, + // the sigmoid runs at alpha 1.702 (the pack's swiglu_alpha); fp16 store rounds after. + try std.testing.expectEqual(@as(f32, 0), serve.swigluoai(0, -7, 1.702, 7.0)); // u+1 = 0 at the floor + try std.testing.expect(serve.swigluoai(7, 7, 1.702, 7.0) > 0); + try std.testing.expect(serve.swigluoai(-1, 7, 1.702, 7.0) < serve.swigluoai(1, 7, 1.702, 7.0)); + try std.testing.expectEqual(serve.swigluoai(12, 0, 1.702, 7.0), serve.swigluoai(7, 0, 1.702, 7.0)); // g clipped above +} \ No newline at end of file diff --git a/zig/src/families/minimax_m3_vl/cuda_sparse.zig b/zig/src/families/minimax_m3_vl/cuda_sparse.zig new file mode 100644 index 0000000000..10248a00e2 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/cuda_sparse.zig @@ -0,0 +1,177 @@ +//! MiniMax M3.1's MSA sparse-indexer on CUDA (the exllamav3-fork m3_sparse_attn contract, by hand): +//! launches minimax_m3.cu's four kernels. The numerics' grading reference is this family's CPU +//! mirror `sparse.zig` (fp32 scoring, right-open causal, +inf boosts post-reduction, ties at the +//! smaller id, left-packed, -1 right-pad, key-only side plane); device parity receipts are Stage 3. + +const std = @import("std"); +const cuda = @import("cuda"); + +/// The contract numbers (asserted in config.zig from the pack): these are the fork's. +pub const d = 128; // index_head_dim == qkv head_dim +pub const block = 128; // index_block_size: keys per scored block; a 128-run never straddles page 256 +pub const topk = 16; // index_topk_blocks, the fixed selection width +pub const page = 256; // the paged fp16 cache's in-tree PAGE_SIZE +pub const max_nb = 1024; // the topk kernel's mark bitmap bound: NB must stay <= this + +fn cdiv(a: usize, b: usize) usize { + return (a + b - 1) / b; +} + +/// A content length's block count (slot anchor): NB = ceil(Sk / 128), at least one — a served row +/// always scores its own key (Sk >= 1), so an empty context still carries the one block. +pub fn nbOf(keys: usize) usize { + const n = cdiv(keys, block); + return if (n == 0) 1 else n; +} + +/// One split's half-open slice of the fixed-width selection: per = ceil(k / nsplit), bounded by k. +/// The nsplit split count is bounded by topk (the wrapper contract). +pub fn splitBounds(k: usize, nsplit: usize, split: usize) [2]usize { + const per = cdiv(k, nsplit); + const lo = split * per; + const hi = @min(lo + per, k); + return .{ lo, hi }; +} + +/// The decode scratch's stride: per (batch, q head, split) row, D accumulators + m + l. +pub const scratch_elts = d + 2; + +/// The family's four kernels from the minimax_m3 image. +pub const Kernels = struct { + scores: cuda.Function, + topk: cuda.Function, + decode: cuda.Function, + combine: cuda.Function, + + pub fn load(m: cuda.Module) !Kernels { + var k: Kernels = undefined; + k.scores = try m.function("tf_m3_scores"); + k.topk = try m.function("tf_m3_topk"); + k.decode = try m.function("tf_m3_decode"); + k.combine = try m.function("tf_m3_decode_combine"); + return k; + } +}; + +/// m3_indexer_scores: the whole score table in one launch — a program per (query row, key block, +/// batch); slot-anchored blocks, fp32 scoring, boosts override post-reduction, no Sq x Sk exists. +pub fn scores( + ks: Kernels, + stream: cuda.Stream, + q: u64, // (B, H, Sq, D) fp16 index queries + k: u64, // (B, Sk, D) fp16 index keys (side-plane rows or fresh K) + qpos: u64, // (B, Sq) i64 + kpos: u64, // (B, Sk) i64, -1 = padded + out: u64, // (B, H, Sq, NB) fp32 + sq: usize, + nb: usize, + sk: usize, + heads: usize, // 4 + batches: usize, + local: u32, // index_local_blocks (1) + init: u32, // sparse_init_block (0) +) !void { + var a: cuda.Args = .{}; + inline for (.{ q, k, qpos, kpos, out, @as(i64, @intCast(sq)), @as(i64, @intCast(nb)), @as(i64, @intCast(sk)), @as(i64, @intCast(heads)), local, init }) |v| a.add(v); + try cuda.launch.launch(ks.scores, .{ .grid = .{ .x = @intCast(sq), .y = @intCast(nb), .z = @intCast(batches) }, .block = .{ .x = 128 } }, stream, &a); +} + +/// m3_topk16: k picks per (batch, head, row), descending, ties at the smaller id, -1 right-pad. +pub fn topk16(ks: Kernels, stream: cuda.Stream, s: u64, out: u64, nb: usize, rows: usize) !void { // m3_topk16: k picks per (batch, head, row) + if (nb > max_nb) return error.ScoreColumns; // the mark bitmap's bound: chunk the prefill there + var a: cuda.Args = .{}; + inline for (.{ s, out, @as(i64, @intCast(nb)), @as(i64, @intCast(rows)), @as(c_int, topk) }) |v| a.add(v); + try cuda.launch.launch(ks.topk, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 128 } }, stream, &a); +} + +/// m3_sparse_decode_attn's split phase: one program per (batch, q head, split) into the fp32 scratch. +pub fn decode( + ks: Kernels, + stream: cuda.Stream, + q: u64, // (B, Hq, D) fp16 + kt: u64, // (P, PAGE, Hkv, D) fp16 + vt: u64, // (P, PAGE, Hkv, D) + table: u64, // (B, npp) i32 logical page -> physical + sel: u64, // (B, Hkv, topk) i32 + seqlens: u64, // (B,) i64 + part: u64, // (B, Hq, nsplit, scratch_elts) fp32 + q_heads: usize, + kv_heads: usize, + npp: usize, + scale: f32, + batches: usize, + nsplit: u32, +) !void { + var a: cuda.Args = .{}; + // 13 kernel parameters: the seven pointers, Hq/Hkv/npp (i64), the fp32 scale, the fixed + // selection width k = topk (16; the sel lists are (B, Hkv, k)), and the split count. + inline for (.{ q, kt, vt, table, sel, seqlens, part, @as(i64, @intCast(q_heads)), @as(i64, @intCast(kv_heads)), @as(i64, @intCast(npp)), scale, @as(c_int, topk), nsplit }) |v| a.add(v); + try cuda.launch.launch(ks.decode, .{ .grid = .{ .x = @intCast(batches), .y = @intCast(q_heads), .z = nsplit }, .block = .{ .x = d } }, stream, &a); +} + +/// _m3_decode_combine_kernel: the splits' fp32 (acc, m, l) merge into the fp16 row; zero-valid rows +/// come out all-zero. nsplit stays <= 8 (the combine's shared arrays; 4 is the served envelope). +pub fn combine(ks: Kernels, stream: cuda.Stream, part: u64, out: u64, q_heads: usize, batches: usize, nsplit: u32) !void { + if (nsplit > 8) return error.TooManySplits; + var a: cuda.Args = .{}; + inline for (.{ part, out, @as(i64, @intCast(q_heads)), nsplit }) |v| a.add(v); + try cuda.launch.launch(ks.combine, .{ .grid = .{ .x = @intCast(batches), .y = @intCast(q_heads) }, .block = .{ .x = d } }, stream, &a); +} + +/// A decode call's scratch bytes: (B, Hq, nsplit, D + 2) fp32. +pub fn scratchBytes(q_heads: usize, batches: usize, nsplit: usize) usize { + return q_heads * batches * nsplit * scratch_elts * 4; +} + +test "the contract numbers are the fork's" { + try std.testing.expectEqual(@as(usize, 128), block); + try std.testing.expectEqual(@as(usize, 128), d); + try std.testing.expectEqual(@as(usize, 16), topk); + try std.testing.expectEqual(@as(usize, 2) * d, page); // 256 = the fp16 page; a 128-key run never straddles it + try std.testing.expectEqual(@as(usize, d + 2), scratch_elts); +} + +test "NB is the ceil of the content's key blocks" { + try std.testing.expectEqual(@as(usize, 1), nbOf(0)); + try std.testing.expectEqual(@as(usize, 1), nbOf(128)); + try std.testing.expectEqual(@as(usize, 2), nbOf(129)); + try std.testing.expectEqual(@as(usize, 1024), nbOf(128 * 1024)); +} + +test "split bounds partition the fixed width without gaps or overlap" { + for ([_]usize{ 1, 3, 4, 8 }) |ns| for (0..ns) |s| { + const b = splitBounds(topk, ns, s); + try std.testing.expect(b[0] < topk or (b[0] == topk and b[1] == topk)); // never past k + try std.testing.expect(b[1] <= topk); + try std.testing.expect(b[1] - b[0] <= topk / ns + 1); + if (s + 1 < ns) { // adjacent splits touch: per is a constant stride + const nxt = splitBounds(topk, ns, s + 1); + try std.testing.expectEqual(b[1], nxt[0]); + } + }; + try std.testing.expectEqual([2]usize{ 0, 4 }, splitBounds(topk, 4, 0)); + try std.testing.expectEqual([2]usize{ 12, 16 }, splitBounds(topk, 4, 3)); +} + +test "the scratch sizing formula" { + try std.testing.expectEqual(@as(usize, 64 * 8192 * 4 * 130 * 4), scratchBytes(64, 8192, 4)); +} + +// The NB <= 1024 mark-bitmap bound, enforced at the wrapper before any launch: 1024 passes +// (the wide prefill chunk), 1025 refuses by name. The kernels and stream are never reached. +test "the topk refuses a score table wider than the mark bitmap's 1024 blocks" { + const ks: Kernels = undefined; + var s: cuda.Stream = undefined; + _ = &s; + try std.testing.expectEqual(@as(usize, 1024), max_nb); + try std.testing.expectError(error.ScoreColumns, topk16(ks, s, 0, 0, max_nb + 1, 1)); +} + +// The nsplit <= 8 combine bound (its shared merge arrays): 8 passes the wrapper, 9 refuses +// before any launch. nsplit is the decode's split phase count; the served envelope is 4. +test "the split merge refuses more splits than its shared arrays hold" { + const ks: Kernels = undefined; + var s: cuda.Stream = undefined; + _ = &s; + try std.testing.expectError(error.TooManySplits, combine(ks, s, 0, 0, 1, 1, 9)); +} \ No newline at end of file diff --git a/zig/src/families/minimax_m3_vl/exl3.zig b/zig/src/families/minimax_m3_vl/exl3.zig new file mode 100644 index 0000000000..0bfef59980 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/exl3.zig @@ -0,0 +1,449 @@ +//! The family's EXL3 admission widened over the checkpoint's own trellis headers: per-module bit-width +//! discovery through the core reader (`core/exl3_load.zig`), which admits mul1, odd half-bit widths and +//! mixed per-tensor widths natively — its `admit` (:134-157) parameterizes nothing: widths, scales and +//! the codebook marker all come from the header. This file adds the family's own admission around it: +//! everything under `language_model.`, one codebook (mul1), the served shapes, per class a complete set. + +const std = @import("std"); +const core = @import("core"); +const cfg = @import("config.zig"); +const wnames = @import("weights.zig"); // the spellings module: the fixture writes the same names the plan reads + +const exl3 = core.exl3_load; +pub const Codebook = exl3.Codebook; + +const layers_root = "language_model.model.layers."; + +/// The module class the group's tensor name puts it in: the load's projection families and the head. +pub const Class = enum { + attn, // self_attn.{q,k,v,o}_proj (the pack: 6-bit mul1) + dense_mlp, // layers 0..dense_layers' mlp.{gate,up,down}_proj (5-bit) + shared, // block_sparse_moe.shared_experts.{gate,up,down}_proj (5-bit) + routed, // block_sparse_moe.experts..{w1,w2,w3} (mixed, odd half-bit widths) + head, // language_model.lm_head (head_bits 8 → 16 halves) + + /// The projection's [k, n]: the reader's rows and columns of the matrix this trellis stores. + /// The attention projections are GQA: q fans out to every head (64×128 = 8192), k/v to the + /// KV heads only (4×128 = 512), o gathers the head dim back — measured on the real pack's + /// trellis headers (q [384,512,96·], k/v [384,32,96·], o [512,384,96·] in K/16, N/16, 16·bits). + pub fn shape(klas: Class, c: *const cfg.Config, proj: u32) [2]u32 { + const inter: u32 = if (klas == .dense_mlp) c.dense_inter else c.moe_inter; + return switch (klas) { + .attn => switch (proj) { + 0 => .{ c.hidden, c.attn_heads * c.head_dim }, // q + 1, 2 => .{ c.hidden, c.kv_heads * c.head_dim }, // k, v + else => .{ c.attn_heads * c.head_dim, c.hidden }, // o + }, + .head => .{ c.hidden, c.vocab }, // measured: the pack's head trellis is [384,12504,128·] — k=hidden, n=vocab, the EXL3 in→out convention + else => switch (proj) { + 0, 1 => .{ c.hidden, inter }, // gate/up: hidden → intermediate + else => .{ inter, c.hidden }, // down: intermediate → hidden + }, + }; + } + + /// How many groups of this class the checkpoint must carry, verbatim from the Config's counts. + pub fn count(klas: Class, c: *const cfg.Config) usize { + const layers: usize = switch (klas) { + .attn => c.layers, + .head => 1, + .dense_mlp => c.dense_layers, + .shared, .routed => c.layers - c.dense_layers, + }; + return layers * switch (klas) { + .attn => @as(usize, 4), + .dense_mlp, .shared => 3, + .routed => 3 * @as(usize, c.experts), + .head => 1, + }; + } +}; + +/// One discovered group: the reader's admission plus its class, projection and layer, verbatim [k, n]. +pub const Group = struct { + prefix: []const u8, // owned; the tensor name without its EXL3 suffixes + class: Class, + proj: u32, // attn: 0..3 = q,k,v,o; else 0..2 = gate,up,down + layer: u32, + halves: u16, // 2 * bits; odd values are mul1's half steps (the pack mixes them per tensor) + codebook: Codebook, + k: u32, + n: u32, + + /// The kernel's k2 for this group: exactly the header's half-bit count. + pub fn k2(g: Group) u16 { + return g.halves; + } +}; + +/// Per-class width distribution: how many groups carry each half-bit width (the load's k2 tables). +pub const Widths = struct { + counts: [5][17]usize = @splat(@splat(0)), + + pub fn bump(w: *Widths, klas: Class, halves: u16) void { + w.counts[@intFromEnum(klas)][halves] += 1; + } + + pub fn at(w: Widths, klas: Class, halves: u16) usize { + return w.counts[@intFromEnum(klas)][halves]; + } +}; + +/// Every trellis group the checkpoint carries, classified and header-checked against the family's Config. +pub const Plan = struct { + gpa: std.mem.Allocator, + groups: []Group, + widths: Widths = .{}, + summary: exl3.Summary = .{}, + bad: usize = 0, // the reader's own refusals (foreign-format tensors, broken groups): a load refuses + + pub fn deinit(p: *Plan) void { + for (p.groups) |g| p.gpa.free(g.prefix); + p.gpa.free(p.groups); + p.* = undefined; + } +}; + +/// The group's owner by its tensor name: class, projection and layer (norms, router and embed carry none). +fn classify(c: *const cfg.Config, prefix: []const u8) ?struct { class: Class, proj: u32, layer: u32 } { + if (std.mem.eql(u8, prefix, "language_model.lm_head")) return .{ .class = .head, .proj = 0, .layer = 0 }; + if (!std.mem.startsWith(u8, prefix, layers_root)) return null; // another tower's tensors (vision stays unquantized) + const rest = prefix[layers_root.len..]; + const dot = std.mem.indexOfScalar(u8, rest, '.') orelse return null; + const layer = std.fmt.parseInt(u32, rest[0..dot], 10) catch return null; + const module = rest[dot + 1 ..]; + const attn = [_][]const u8{ "q_proj", "k_proj", "v_proj", "o_proj" }; + const mlp = [_][]const u8{ "gate_proj", "up_proj", "down_proj" }; + for (attn, 0..) |t, i| { + if (std.mem.startsWith(u8, module, "self_attn.")) { + if (std.mem.startsWith(u8, module["self_attn.".len..], t)) return .{ .class = .attn, .proj = @intCast(i), .layer = layer }; + } + } + if (std.mem.startsWith(u8, module, "mlp.")) { + for (mlp, 0..) |t, i| if (std.mem.startsWith(u8, module["mlp.".len..], t)) { + return .{ .class = .dense_mlp, .proj = @intCast(i), .layer = layer }; + }; + } + const shared = "block_sparse_moe.shared_experts."; + if (std.mem.startsWith(u8, module, shared)) { + for (mlp, 0..) |t, i| if (std.mem.startsWith(u8, module[shared.len..], t)) { + return .{ .class = .shared, .proj = @intCast(i), .layer = layer }; + }; + } + const experts = "block_sparse_moe.experts."; + if (std.mem.startsWith(u8, module, experts)) { + const tail = module[experts.len..]; + const edot = std.mem.indexOfScalar(u8, tail, '.') orelse return null; + const e = std.fmt.parseInt(u32, tail[0..edot], 10) catch return null; + const w = tail[edot + 1 ..]; + const proj: u32 = if (std.mem.eql(u8, w, "w1")) 0 else if (std.mem.eql(u8, w, "w3")) 1 else if (std.mem.eql(u8, w, "w2")) 2 else return null; + if (e >= c.experts) return null; + return .{ .class = .routed, .proj = proj, .layer = layer }; + } + return null; +} + +/// The checkpoint directory's every EXL3 group, geometry-checked against the family's Config. The reader +/// owns the format's admission (marker values, scale presence, whole-128 tiles); this adds the family's: +/// everything under `language_model.`, mul1, and the shapes the family serves, per class a complete set. +pub fn discover(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, c: *const cfg.Config) !Plan { + var scan = try exl3.scan(gpa, io, dir); + defer scan.deinit(); + var out: Plan = .{ .gpa = gpa, .groups = &.{}, .summary = scan.summary, .bad = scan.bad }; + var groups: std.ArrayList(Group) = .empty; + errdefer { + for (groups.items) |g| gpa.free(g.prefix); + groups.deinit(gpa); + } + for (scan.groups) |g| { + const who = classify(c, g.prefix) orelse return error.ForeignTrellis; // a quantized tensor the family never reads + const want = who.class.shape(c, who.proj); + if (g.k != want[0] or g.n != want[1]) return error.BadTrellisGeo; + if (g.codebook != .mul1) return error.NotMul1; // the pack's one codebook + if (who.class == .head and g.halves != 2 * c.qc_head_bits) return error.HeadBits; // lm_head's own width + try groups.append(gpa, .{ + .prefix = try gpa.dupe(u8, g.prefix), + .class = who.class, + .proj = who.proj, + .layer = who.layer, + .halves = g.halves, + .codebook = g.codebook, + .k = @intCast(g.k), + .n = @intCast(g.n), + }); + out.widths.bump(who.class, g.halves); + } + out.groups = try groups.toOwnedSlice(gpa); + var missing: bool = false; + for ([_]Class{ .attn, .dense_mlp, .shared, .routed, .head }) |klas| if (countOf(&out, klas) != Class.count(klas, c)) { + missing = true; + }; + if (missing) { + out.deinit(); // the plan holds on to nothing when the pack's set is incomplete + return error.MissingTrellis; + } + if (scan.bad != 0) { + out.deinit(); + return error.BrokenTrellis; + } + return out; +} + +/// One class's discovered groups (the loader walks them per layer, checkpoint order not guaranteed). +pub fn countOf(p: *const Plan, klas: Class) usize { + var n: usize = 0; + for (p.groups) |g| if (g.class == klas) { + n += 1; + }; + return n; +} + +/// The trellis-and-scales bytes a group reads from the checkpoint (the raw device footprint's read). +pub fn groupBytes(g: Group) u64 { + return @as(u64, g.k / 16) * (g.n / 16) * g.halves * 8 * 2 // the trellis int16 words + + @as(u64, g.k) * 2 + @as(u64, g.n) * 2 + 4; // suh + svh fp16, the codebook marker int32 +} + +test "classify maps the pack's spellings onto the family's classes by projection" { + const c = cfg.Config{}; + const cases = [_]struct { prefix: []const u8, class: Class, proj: u32, layer: u32 }{ + .{ .prefix = "language_model.model.layers.7.self_attn.q_proj", .class = .attn, .proj = 0, .layer = 7 }, + .{ .prefix = "language_model.model.layers.7.self_attn.o_proj", .class = .attn, .proj = 3, .layer = 7 }, + .{ .prefix = "language_model.model.layers.0.mlp.down_proj", .class = .dense_mlp, .proj = 2, .layer = 0 }, + .{ .prefix = "language_model.model.layers.9.block_sparse_moe.shared_experts.up_proj", .class = .shared, .proj = 1, .layer = 9 }, + .{ .prefix = "language_model.model.layers.57.block_sparse_moe.experts.3.w3", .class = .routed, .proj = 1, .layer = 57 }, + .{ .prefix = "language_model.model.layers.57.block_sparse_moe.experts.3.w2", .class = .routed, .proj = 2, .layer = 57 }, + .{ .prefix = "language_model.lm_head", .class = .head, .proj = 0, .layer = 0 }, + }; + for (cases) |t| { + const who = classify(&c, t.prefix) orelse return error.TestUnexpectedResult; + try std.testing.expectEqual(t.class, who.class); + try std.testing.expectEqual(t.proj, who.proj); + try std.testing.expectEqual(t.layer, who.layer); + } + try std.testing.expect(classify(&c, "language_model.model.layers.7.self_attn.index_q_proj") == null); // unquantized parts + try std.testing.expect(classify(&c, "language_model.model.layers.7.block_sparse_moe.gate") == null); // router fp32 + try std.testing.expect(classify(&c, "vision_tower.encoder.layers.0.mlp.fc1") == null); // not this family's +} + +test "geometry, counts and the head's own bits are the Config's, and k2 carries odd half steps verbatim" { + try std.testing.expectEqual([2]u32{ 6144, 3072 }, Class.routed.shape(&cfg.Config{}, 0)); + try std.testing.expectEqual([2]u32{ 3072, 6144 }, Class.routed.shape(&cfg.Config{}, 2)); + try std.testing.expectEqual([2]u32{ 6144, 12288 }, Class.dense_mlp.shape(&cfg.Config{}, 1)); + try std.testing.expectEqual([2]u32{ 6144, 200064 }, Class.head.shape(&cfg.Config{}, 0)); + const c = cfg.Config{}; + try std.testing.expectEqual(@as(usize, 240), Class.count(.attn, &c)); // 60 layers * q,k,v,o + try std.testing.expectEqual(@as(usize, 9), Class.count(.dense_mlp, &c)); // 3 dense layers * gate,up,down + try std.testing.expectEqual(@as(usize, 171), Class.count(.shared, &c)); // 57 MoE layers * 3 + try std.testing.expectEqual(@as(usize, 21888), Class.count(.routed, &c)); // 57 * 128 experts * w1,w2,w3 + try std.testing.expectEqual(@as(usize, 22309), 240 + 9 + 171 + 21888 + 1); // the pack's whole group set + const g = Group{ .prefix = "", .class = .routed, .proj = 0, .layer = 3, .halves = 9, .codebook = .mul1, .k = 6144, .n = 3072 }; + try std.testing.expectEqual(@as(u16, 9), g.k2()); // 4.5-bit: the odd mul1 half step, carried verbatim + try std.testing.expectEqual(@as(u64, (6144 / 16) * (3072 / 16) * 9 * 8 * 2 + 6144 * 2 + 3072 * 2 + 4), groupBytes(g)); // 11.25 MB +} + +// --- a synthetic checkpoint fixture, small enough to disk-test, shapes from a scaled Config --- + +/// The fixture's shapes: the REAL pack's widths (config.parse's width gate requires exactly those, so +/// admission against a synthetic pack must carry them) with small counts: 2 layers, one dense wave, 4 +/// whole-128 experts and a 512-wide vocabulary (whole 128s at every tile edge). +pub fn synthConfig() cfg.Config { + var c = cfg.Config{}; // the pack's own widths: hidden 6144, head_dim 128, 4×128 index heads, etc. + c.layers = 2; + c.run = 2; + c.vocab = 512; + c.dense_layers = 1; + c.experts = 4; + c.topk = 4; + c.qc_head_bits = 8; // the head's own width: admission reads it from the config, the fixture mirrors it + return c; +} + +pub const Fixture = struct { + foreign: bool = false, // one quantized tensor outside language_model. + drop_last_w2: bool = false, // expert 3's up projection missing: an incomplete set refuses + head_halves: u16 = 16, // the 8-bit head's halves + mcg: bool = false, // one routed group's marker names the other codebook +}; + +/// One EXL3 group's four tensors into the fixture's header and zero-filled data (headers only are read). +const Builder = struct { + gpa: std.mem.Allocator, + header: std.ArrayListUnmanaged(u8) = .empty, + data: std.ArrayListUnmanaged(u8) = .empty, + first: bool = true, + + fn group(b: *Builder, name: []const u8, k: usize, n: usize, halves: u16, cb: exl3.Codebook) !void { + const trellis = (k / 16) * (n / 16) * @as(usize, halves) * 8 * 2; + const suh: usize = k * 2; + const svh: usize = n * 2; + const base = b.data.items.len; + try b.data.appendNTimes(b.gpa, 0, trellis + suh + svh + 4); + const marker: [4]u8 = if (cb == .mul1) .{ 0x2D, 0xD1, 0xDC, 0x83 } else .{ 0xED, 0x1F, 0xAC, 0xCB }; // mul1 / mcg magic + @memcpy(b.data.items[base + trellis + suh + svh ..][0..4], &marker); + const parts = [_]struct { suffix: []const u8, dtype: []const u8, len: usize, shape: []const u8 }{ + .{ .suffix = ".trellis", .dtype = "I16", .len = trellis, .shape = "K,N,H" }, + .{ .suffix = ".suh", .dtype = "F16", .len = suh, .shape = "K" }, + .{ .suffix = ".svh", .dtype = "F16", .len = svh, .shape = "N" }, + .{ .suffix = if (cb == .mul1) ".mul1" else ".mcg", .dtype = "I32", .len = 4, .shape = "1" }, + }; + var off = base; + for (parts) |e| { + if (!b.first) try b.header.appendSlice(b.gpa, ","); + b.first = false; + try b.header.print(b.gpa, "\"{s}{s}\":{{\"dtype\":\"{s}\",\"shape\":[", .{ name, e.suffix, e.dtype }); + if (std.mem.eql(u8, e.shape, "K,N,H")) { + try b.header.print(b.gpa, "{d},{d},{d}", .{ k / 16, n / 16, halves * 8 }); + } else if (std.mem.eql(u8, e.shape, "K")) { + try b.header.print(b.gpa, "{d}", .{k}); + } else if (std.mem.eql(u8, e.shape, "N")) { + try b.header.print(b.gpa, "{d}", .{n}); + } else try b.header.appendSlice(b.gpa, "1"); + try b.header.print(b.gpa, "],\"data_offsets\":[{d},{d}]}}", .{ off, off + e.len }); + off += e.len; + } + } + + fn done(b: *Builder) ![]u8 { + try b.header.insertSlice(b.gpa, 0, "{"); // the safetensors header is one JSON object of the tensors + try b.header.appendSlice(b.gpa, "}"); + var file: std.ArrayListUnmanaged(u8) = .empty; + errdefer file.deinit(b.gpa); + try file.appendNTimes(b.gpa, 0, 8); + std.mem.writeInt(u64, file.items[0..8], b.header.items.len, .little); + try file.appendSlice(b.gpa, b.header.items); + try file.appendSlice(b.gpa, b.data.items); + return file.toOwnedSlice(b.gpa); + } +}; + +test "discover admits the whole synthetic pack and records its mixed widths" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const path = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(path); + try writeFixture(gpa, io, tmp.dir, .{}); + const c = synthConfig(); + var p = try discover(gpa, io, path, &c); + defer p.deinit(); + try std.testing.expectEqual(@as(usize, 27), p.groups.len); // 8 attn + 3 dense + 3 shared + 12 routed + 1 head + try std.testing.expectEqual(@as(usize, 8), countOf(&p, .attn)); + try std.testing.expectEqual(@as(usize, 3), countOf(&p, .dense_mlp)); + try std.testing.expectEqual(@as(usize, 3), countOf(&p, .shared)); + try std.testing.expectEqual(@as(usize, 12), countOf(&p, .routed)); + try std.testing.expectEqual(@as(usize, 1), countOf(&p, .head)); + try std.testing.expectEqual(@as(usize, 6), p.widths.at(.routed, 4)); // e0,e2: w1,w2,w3 at 2-bit + try std.testing.expectEqual(@as(usize, 6), p.widths.at(.routed, 5)); // e1,e3: the odd 2.5-bit half step + try std.testing.expectEqual(@as(usize, 1), p.widths.at(.head, 16)); // the 8-bit head's one group + var head_halves: u16 = 0; + for (p.groups) |g| if (g.class == .head) { + head_halves = g.k2(); + try std.testing.expectEqual(@as(u32, 512), g.n); // cols of the lm_head matrix: the whole vocab at TP1 (k=hidden, measured pack convention) + }; + try std.testing.expectEqual(@as(u16, 16), head_halves); +} + +test "a head at the wrong bits, a foreign quantized tensor, another codebook and an incomplete set refuse" { + const gpa = std.testing.allocator; + const io = std.testing.io; + inline for ([_]struct { name: []const u8, o: Fixture, want: anyerror }{ + .{ .name = "head_halves", .o = .{ .head_halves = 12 }, .want = error.HeadBits }, // not this pack's 8-bit row + .{ .name = "foreign", .o = .{ .foreign = true }, .want = error.ForeignTrellis }, // another tower's quantized mlp + .{ .name = "mcg", .o = .{ .mcg = true }, .want = error.NotMul1 }, // one expert on the other codebook + .{ .name = "incomplete", .o = .{ .drop_last_w2 = true }, .want = error.MissingTrellis }, // expert 3's w2 missing + }) |case| { + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try writeFixture(gpa, io, tmp.dir, case.o); + const path = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(path); + const c = synthConfig(); + try std.testing.expectError(case.want, discover(gpa, io, path, &c)); + } +} + +/// The synthetic pack written into `dir` as model.safetensors (header + zero-filled groups). +pub fn writeFixture(gpa: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, o: Fixture) !void { + const file = try buildFixture(gpa, o); + defer gpa.free(file); + try dir.writeFile(io, .{ .sub_path = "model.safetensors", .data = file }); +} + +/// A 2-layer synthetic checkpoint: layer 0 dense, layer 1 MoE (4 experts, one odd half step, mixed). +pub fn buildFixture(gpa: std.mem.Allocator, o: Fixture) ![]u8 { + var b: Builder = .{ .gpa = gpa }; + defer { + b.header.deinit(gpa); + b.data.deinit(gpa); + } + var names: std.ArrayList([]u8) = .empty; + defer { + for (names.items) |n| gpa.free(n); + names.deinit(gpa); + } + const c = synthConfig(); + const keep = struct { + fn f(a: std.mem.Allocator, list: *std.ArrayList([]u8), comptime fmt: []const u8, args: anytype) ![]const u8 { + const owned = try std.fmt.allocPrint(a, fmt, args); + try list.append(a, owned); + return owned; + } + }.f; + // the head, whole vocab at TP1: its own width is the pack's 8 halves *2 qc_head_bits — unless the + // fixture deliberately varies it. Attention's and the MLPs' widths are any legal width: admission + // checks geometry and the codebook, not per-class bits; the head's is the only width rule. + const hs = Class.head.shape(&c, 0); + try b.group(try keep(gpa, &names, "language_model.lm_head", .{}), hs[0], hs[1], o.head_halves, .mul1); + for ([_]u32{ 0, 1 }) |l| { + inline for (.{ "q_proj", "k_proj", "v_proj", "o_proj" }, 0..) |p, pi| { + const s = Class.attn.shape(&c, pi); + try b.group(try keep(gpa, &names, "language_model.model.layers.{d}.self_attn.{s}", .{ l, p }), s[0], s[1], 2, .mul1); + } + if (l < c.dense_layers) { + inline for ([_]wnames.ExpertProj{ .gate, .up, .down }) |proj| { + const s = Class.dense_mlp.shape(&c, @intFromEnum(proj)); + try b.group(try keep(gpa, &names, "language_model.model.layers.{d}.mlp.{s}_proj", .{ l, @tagName(proj) }), s[0], s[1], 2, .mul1); + } + continue; + } + inline for ([_]wnames.ExpertProj{ .gate, .up, .down }) |proj| { + const s = Class.shared.shape(&c, @intFromEnum(proj)); + try b.group(try keep(gpa, &names, "language_model.model.layers.{d}.block_sparse_moe.shared_experts.{s}_proj", .{ l, @tagName(proj) }), s[0], s[1], 2, .mul1); + } + for (0..@as(u32, c.experts)) |e| { + const odd: u16 = if (e % 2 == 1) 5 else 4; // mixed per-expert widths, one odd half step + for ([_]wnames.ExpertProj{ .gate, .up, .down }) |proj| { + if (e == 3 and o.drop_last_w2 and proj == .up) continue; // expert 3's up projection missing + const s = Class.routed.shape(&c, @intFromEnum(proj)); + try b.group(try keep(gpa, &names, "language_model.model.layers.{d}.block_sparse_moe.experts.{d}.{s}", .{ l, e, proj.ckpt() }), s[0], s[1], odd, if (e == 2 and o.mcg) .mcg else .mul1); + } + } + } + if (o.foreign) try b.group(try keep(gpa, &names, "vision_tower.blocks.0.mlp.fc1", .{}), c.hidden, c.hidden, 2, .mul1); + return b.done(); // the owned names outlive only this build: the file bytes carry the text on +} +/// The synthetic checkpoint's config.json: this scaled Config's counts in the pack's spellings, so the +/// admission tests open `dir` as a whole model. The indexer's and EXL3's shapes come from Config's +/// defaults (the real pack's values — 4×128 index heads, 16-of-128 blocks, swigluoai, sigmoid router); +/// rms_norm_eps matches the Config default exactly, as parse compares it. +pub fn writeSynthConfigJson(gpa: std.mem.Allocator, io: std.Io, dir: std.Io.Dir) !void { + const c = synthConfig(); + const json = try std.fmt.allocPrint(gpa, + "{{\"model_type\":\"minimax_m3_vl\",\"text_config\":{{" ++ + "\"hidden_size\":{d},\"num_hidden_layers\":{d},\"vocab_size\":{d}," ++ + "\"num_attention_heads\":64,\"num_key_value_heads\":4,\"head_dim\":128," ++ + "\"rms_norm_eps\":0.000001,\"num_local_experts\":{d},\"num_experts_per_tok\":{d}," ++ + "\"intermediate_size\":{d},\"dense_intermediate_size\":{d},\"shared_intermediate_size\":{d}," ++ + "\"routed_scaling_factor\":2.0,\"swiglu_alpha\":1.702,\"swiglu_limit\":7.0," ++ + "\"scoring_func\":\"sigmoid\",\"use_routing_bias\":true," ++ + "\"use_qk_norm\":true,\"qk_norm_type\":\"per_head\",\"moe_layer_freq\":[0,1]," ++ + "\"sparse_attention_config\":{{\"sparse_num_index_heads\":4,\"sparse_index_dim\":128," ++ + "\"sparse_topk_blocks\":16,\"sparse_block_size\":128,\"sparse_score_type\":\"max\"}}," ++ + "\"quantization_config\":{{\"quant_method\":\"exl3\",\"avg_bits\":3.06,\"head_bits\":8," ++ + "\"codebook\":\"mul1\"}}}}}}", + .{ c.hidden, c.layers, c.vocab, c.experts, c.topk, c.moe_inter, c.dense_inter, c.shared_inter }); + defer gpa.free(json); + try dir.writeFile(io, .{ .sub_path = "config.json", .data = json }); +} diff --git a/zig/src/families/minimax_m3_vl/graph.zig b/zig/src/families/minimax_m3_vl/graph.zig new file mode 100644 index 0000000000..68c52dfe72 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/graph.zig @@ -0,0 +1,1058 @@ +//! The MSA launch graph on device: the per-layer forward the receipts' boundary words describe, +//! running the TP2 rank split shard.zig draws. The quantized projections run through the mainline +//! EXL3 machinery the family already receipts — `cuda.exl3.Dense` decode-row GEMMs (mul1's odd +//! half steps and the head's 8 halves included) and the grouped routed experts +//! `cuda.exl3_experts.Experts` (the Stage-5 upload's fp32 epilogue, limit 7, act_mode 1); the MSA +//! indexer and the block-sparse attention run through the parity-proven `tf_m3_scores/topk/ +//! decode/decode_combine` launches (`cuda_sparse.zig`); the activation rows run through +//! `tf_m3_swigluoai` (`cuda_serve.zig`). The between-the-kernels glue — the gemma norm rows, the +//! partial rope, the per-head fp16 stores, the noaux sigmoid+bias router's normalize-then-scale +//! weights, the small fp16 indexer projections, the residual adds — is host fp32 code that reuses +//! serve.zig's own published helpers (`normRow`, `ropeRow`, `route`, `f16r`), so the graded bar is +//! this graph's contract by construction, not a re-typed copy. Where the device path stores fp16 +//! where the host engine's fp32 linears would not, that is the fork's own device contract (exl3 +//! GEMM inputs are fp16-rounded by `rot_in`; the gate/up/down wave and the swigluoai epilogue +//! store fp16; the attention reads fp16 cache rows through fp32 softmax) — each graded point's +//! receipts live in TF-STAGE4-ATTACH-RECEIPTS.md and the receipts this stage writes. +//! +//! Decode-first minimal path, honestly sized: one row at a time, greedy only — no batching, no +//! streaming, no HTTP surface. The KV cache and the indexer side plane are a fixed page-backed +//! pool of `keys_cap` keys (16 pages of 256); a row past it refuses loudly rather than mis-serve. +//! +//! TP2 (shard.zig's world): rank 0 embeds, runs its layers, hands the residual row to rank 1 over +//! NCCL, which runs its layers and hands it back; rank 0 applies the final gemma norm, the lm_head +//! and greedy. The experts replicate per rank exactly as Stage-5 uploaded them; the only per-token +//! cross-rank traffic is the two hidden-row handoffs. + +const std = @import("std"); +const cuda = @import("cuda"); +const config = @import("config.zig"); +const serve = @import("serve.zig"); +const shard = @import("shard.zig"); +const boundary = @import("cuda_native.zig"); + +const exl3 = cuda.exl3; +const exl3e = cuda.exl3_experts; +const exl3_upload = cuda.exl3_upload; +const sp = @import("cuda_sparse.zig"); // the tf_m3 indexer/attention launches +const glue = @import("cuda_serve.zig"); // the family's tf_m3_swigluoai rows + +const exl3_load = @import("core").exl3_load; // the core reader the Uploader takes parts from + +/// The fixed KV/plane pool this decode path serves: 16 pages of 256 keys = 4096 keys (the size is +/// said here; the topk-16 selection's whole window fits inside it at the smoke's lengths). +pub const keys_cap: usize = 4096; +/// The decode split count (the served envelope; the combine's shared arrays hold at most 8). +const nsplit: u32 = 4; + +pub const Options = struct { + /// World 2's handoff: the NCCL library and communicator; null = TP1 (the fixture rank). + lib: ?*const cuda.nccl.Library = null, + comm: ?cuda.nccl.Comm = null, + /// Additive Stage-8 (the host-f64 grade): a TP1-only backbone cap — build and run only the + /// first `layer_cap` owned layers (the cap is the real-pack layer stack that fits the host's + /// RAM beside the f64 reference). TP2 refuses a cap (each rank owns its exact half); 0 = all. + layer_cap: usize = 0, +}; + +/// One layer's device tensors and the host glue weights this rank runs it with. +const Layer = struct { + lno: u32, // the checkpoint's layer id (the caches offset by the local index) + in_norm: []f32 = &.{}, // [hidden], RAW bf16 passthrough values (serve.normRow applies its own (1+w)) + post_norm: []f32 = &.{}, + q_norm: []f32 = &.{}, // [head_dim] + k_norm: []f32 = &.{}, + iq_norm: []f32 = &.{}, // [i_dim] + ik_norm: []f32 = &.{}, + q: ?exl3.Weight = null, + k: ?exl3.Weight = null, + v: ?exl3.Weight = null, + o: ?exl3.Weight = null, + iq: ?[]f16 = null, // the fp16 passthrough indexer projections, host-GEMV-ed + ik: ?[]f16 = null, + mlp: Mlp = undefined, + + const Mlp = union(enum) { + dense: struct { gate: exl3.Weight, up: exl3.Weight, down: exl3.Weight }, + moe: struct { + router: []f32, // [E, hidden] fp32 raw (the router gate is not quantized) + bias: []f32, // [E] fp32; empty when the pack names none + routed: exl3e.Layer, // the grouped experts' device tensors + sgate: exl3.Weight, + sup: exl3.Weight, + sdown: exl3.Weight, + }, + }; +}; + +/// The last row's trace: the smoke's integrity assertions read this. +pub const Trace = struct { + fired: usize = 0, // experts fired by the last MoE wave + wsum: f32 = 0, // the routing weights' sum (= routed_scale by the contract) + sel_valid: usize = 0, // valid selected blocks, the last kv head of the last layer +}; + +pub const Engine = struct { + gpa: std.mem.Allocator, + io: std.Io, + c: config.Config, + d: *const cuda.Driver, + s: cuda.Stream, + dense: exl3.Dense, + experts: exl3e.Experts, + mod: cuda.Module, // the minimax_m3 image the tf_m3 glue kernels resolve from + ks: sp.Kernels, + gk: glue.Kernels, + up: exl3_upload.Uploader, // owns every device weight buffer AND this engine's runtime rows + rank: u32, + world: u32, + lib: ?*const cuda.nccl.Library, + comm: ?cuda.nccl.Comm, + layers: []Layer, // this rank's owned layers, in run order + first: u32, // layers[0]'s checkpoint id + embed_tbl: ?cuda.DeviceBuffer = null, // rank 0: the bf16 [vocab, hidden] table + final_norm: ?[]f32 = null, // rank 0: [hidden] + head: ?exl3.Weight = null, // rank 0: the [hidden -> vocab] trellis + len: u32 = 0, + x: []f32, // [hidden], the residual stream (host fp32, serve's contract) + xn: []f32, + xn2: []f32, + iq32: []f32, // [i_heads * i_dim] + ik32: []f32, // [i_dim] + q16: []f16, // [qk_rows] + k16: []f16, // [kv_rows] + v16: []f16, + iq16: []f16, // [i_heads * i_dim] + ik16: []f16, // [i_dim] + projf: []f32, // [qk_rows + 2*kv_rows], the q/k/v download + o16: []f16, // [max(hidden, inter)] — the wave outputs' download + elogits: []f32, // [experts], the router's scratch + logits: []f32, // [vocab] (rank 0); the last row's + sel_e: []u32, // [topk], the last MoE wave's picks in selection order + pickw: []f32, // [topk], the routing weights (.half() rounded) + picks: []bool, // [E] + sel_seen: []i32, // [kv_heads, topk], the last layer's selection (the trace) + last: Trace = .{}, + // device rows and caches (all allocated through `up`, so build-error safety is one owner): + x32: cuda.DeviceBuffer, // [hidden] f32, every gemm input + proj32: cuda.DeviceBuffer, // [qk_rows + 2*kv_rows] f32, the q/k/v gemm outputs + recv32: cuda.DeviceBuffer, // [hidden] f32, the peer handoff + ctrl32: cuda.DeviceBuffer, // world 2: the [2] u32 row-control word (op + token or new length) + qdev16: cuda.DeviceBuffer, // [qk_rows] f16, the attention's q + iqdev16: cuda.DeviceBuffer, // [i_heads*i_dim] f16, the scorer's queries + attn16: cuda.DeviceBuffer, // [qk_rows] f16, the attention output + g16: cuda.DeviceBuffer, // [max_inter] f16, the gate rows (and the down outputs) + u16: cuda.DeviceBuffer, + a16: cuda.DeviceBuffer, + moe32: cuda.DeviceBuffer, // [hidden] f32, the routed experts' weighted sum + pick32: cuda.DeviceBuffer, // [topk] i32 + wts32: cuda.DeviceBuffer, // [topk] f32 + logits32: cuda.DeviceBuffer, // [vocab] f32 (rank 0; a stub elsewhere) + kc: cuda.DeviceBuffer, // [nl * cap, kv_rows] f16 pages — physical (P, PAGE, Hkv, D) + vc: cuda.DeviceBuffer, + plane: cuda.DeviceBuffer, // [nl * cap, i_dim] f16 side-plane rows, per layer + scores32: cuda.DeviceBuffer, // [kv_heads, nb_max] f32 + sel32: cuda.DeviceBuffer, // [kv_heads, topk] i32 + qpos64: cuda.DeviceBuffer, + kpos64: cuda.DeviceBuffer, // [cap] i64, the content keys 0..cap-1 + seqlens: cuda.DeviceBuffer, // [1] i64 + table: cuda.DeviceBuffer, // [pool pages] i32, the identity page table + part32: cuda.DeviceBuffer, // the decode splits' scratch + exp_scratch: cuda.DeviceBuffer, // the grouped experts' scratch + + /// The gemm row shapes (usize forms of the config's counts). + fn hidden(e: *const Engine) usize { + return @intCast(e.c.hidden); + } + fn qkr(e: *const Engine) usize { + return @as(usize, e.c.attn_heads) * @as(usize, e.c.head_dim); + } + fn kvr(e: *const Engine) usize { + return @as(usize, e.c.kv_heads) * @as(usize, e.c.head_dim); + } + fn irow(e: *const Engine) usize { + return @as(usize, e.c.i_dim); + } + fn iplanes(e: *const Engine) usize { + return @as(usize, e.c.i_heads) * @as(usize, e.c.i_dim); + } + fn maxInter(e: *const Engine) usize { + const c = &e.c; + return @max(@as(usize, c.dense_inter), @max(@as(usize, c.moe_inter), @as(usize, c.shared_inter))); + } + + pub fn deinit(e: *Engine) void { + e.s.deinit(); + e.dense.unload(); + e.experts.unload(); + e.mod.unload(); + e.up.deinit(); // frees every device weight buffer + if (e.embed_tbl != null) freeBuf(&e.embed_tbl.?); + inline for ([_][]const u8{ "x32", "proj32", "recv32", "ctrl32", "qdev16", "iqdev16", "attn16", "g16", "u16", "a16", "moe32", "pick32", "wts32", "logits32", "kc", "vc", "plane", "scores32", "sel32", "qpos64", "kpos64", "seqlens", "table", "part32", "exp_scratch" }) |name| { + freeBuf(&@field(e, name)); + } + const gpa = e.gpa; + for (e.layers) |*l| deinitLayer(gpa, l); + gpa.free(e.layers); + if (e.final_norm) |w| gpa.free(w); + for ([_]*[]f32{ &e.x, &e.xn, &e.xn2, &e.iq32, &e.ik32, &e.elogits, &e.logits, &e.pickw }) |p| { + if (p.len > 0) gpa.free(p.*); + } + if (e.sel_e.len > 0) gpa.free(e.sel_e); // u32: the picks' ids + if (e.picks.len > 0) gpa.free(e.picks); + if (e.sel_seen.len > 0) gpa.free(e.sel_seen); + for ([_]*[]f16{ &e.q16, &e.k16, &e.v16, &e.iq16, &e.ik16, &e.o16 }) |p| { + if (p.len > 0) gpa.free(p.*); + } + if (e.projf.len > 0) gpa.free(e.projf); // f32: the q/k/v gemm outputs + e.* = undefined; + } + + /// Build from the checkpoint dir: the admission, the whole Stage-2 plan, the rank's shard, + /// the Stage-5 upload, then this engine's own rows and caches. + pub fn build(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, dir: []const u8, rank: u32, world: u32, o: Options) !Engine { + if (world < 1 or world > shard.max_world) return error.TooManyRanks; + if (rank >= world) return error.NoRank; + _ = try shard.checked(.{ .rank = rank, .world = world }); + var admitted = try boundary.admitDir(gpa, io, dir); + defer admitted.deinit(); + const c = &admitted.c; + if (c.i_topk != sp.topk or c.i_dim != sp.d or c.block != sp.block or c.head_dim != sp.d or c.i_heads != c.kv_heads) { + return error.NotMiniMaxM3VL; // the tf_m3 kernels' fixed widths: a different MSA refuses + } + if (c.topk > 32 or c.experts > 4096) return error.NotMiniMaxM3VL; + const load = &admitted.load; + const groups = admitted.groups.groups; + const r = shard.range(c.layers, world, rank); + const dims: usize = @intCast(c.hidden); + // the Stage-8 cap: TP1's first `layer_cap` layers only (the grade's real-pack boot). + // Every walk below gates on `run_count`; the default 0 keeps the full range. + var run_count: usize = r.count; + if (o.layer_cap > 0) { + if (world != 1) return error.LayerCapTP2; // the cap serves TP1 only: no partial rank + if (o.layer_cap > r.count) return error.LayerCapTooWide; + run_count = o.layer_cap; + } + const nl: usize = run_count; + + // the reader + the rank's plan walk, the Stage-5 harness's upload shape + var scan = try exl3_load.scan(gpa, io, dir); + defer scan.deinit(); + var idx = std.StringHashMapUnmanaged(usize).empty; + defer idx.deinit(gpa); + for (scan.groups, 0..) |g, i| try idx.put(gpa, g.prefix, i); + + var arena_state = std.heap.ArenaAllocator.init(gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + + var max_buffers: usize = 64; + for (load.items) |*item| { + if (!shard.owns(item, c.layers, world, rank)) continue; + max_buffers += item.bytes / (64 << 20) + @as(usize, switch (item.kind) { + .group => 4, // trellis + two scales + bias + .passthrough => 1, + }); + } + // ownership discipline for the error paths: everything below registers guarded errdefers + // (`built == false`); once `e` is fully constructed the single `errdefer e.deinit()` + // owns it all and the guards stand down. + var built = false; + var up: exl3_upload.Uploader = .{ .d = d, .gpa = gpa, .max_buffers = max_buffers }; + errdefer { + if (!built) up.deinit(); + } + + // the routed experts' per-expert act width: the first owned routed gate's n + var width: usize = 0; + for (groups) |g| { + if (g.class != .routed or g.proj != 0) continue; + if (g.layer >= r.first and g.layer < r.first + run_count) { + width = @intCast(g.n); + break; + } + } + + var layers = try gpa.alloc(Layer, nl); + errdefer { + if (!built) { + for (layers) |*l| deinitLayer(gpa, l); + gpa.free(layers); + } + } + for (layers, 0..) |*l, i| l.* = .{ .lno = r.first + @as(u32, @intCast(i)) }; + + // build-time temp slots per class (assembled into the layers' mlps after the walk) + var d_gate = try a.alloc(?exl3.Weight, nl); + var d_up = try a.alloc(?exl3.Weight, nl); + var d_down = try a.alloc(?exl3.Weight, nl); + var s_gate = try a.alloc(?exl3.Weight, nl); + var s_up = try a.alloc(?exl3.Weight, nl); + var s_down = try a.alloc(?exl3.Weight, nl); + var routers = try a.alloc(?[]f32, nl); + var biases = try a.alloc(?[]f32, nl); + for (0..nl) |i| { + d_gate[i] = null; + d_up[i] = null; + d_down[i] = null; + s_gate[i] = null; + s_up[i] = null; + s_down[i] = null; + routers[i] = null; + biases[i] = null; + } + + // the routed experts collect per owned MoE layer, then upload as one layer() call each + const Projs = std.ArrayListUnmanaged(exl3_upload.Proj); + const LayerProjs = struct { gate: Projs, up: Projs, down: Projs }; + var exps = try a.alloc(?LayerProjs, c.layers); + for (exps) |*ep| ep.* = null; + + var head_w: ?exl3.Weight = null; + var embed: ?cuda.DeviceBuffer = null; + var final_norm: ?[]f32 = null; + var norm_fail = false; // a passthrough slot the layer's glue needs and the pack refuses + + for (load.items) |*item| { + if (!shard.owns(item, c.layers, world, rank)) continue; + switch (item.kind) { + .passthrough => { + const t = try scan.ck.get(item.name); + if (std.mem.eql(u8, item.name, "language_model.model.embed_tokens.weight")) { + embed = try cuda.DeviceBuffer.fromHost(d, t.bytes); + } else if (std.mem.eql(u8, item.name, "language_model.model.norm.weight")) { + final_norm = try expandNorm(gpa, t.bytes, dims); + norm_fail = norm_fail or final_norm == null; + } else { + try readPassthrough(gpa, c, layers, r.first, item.name, t.bytes, routers, biases, &norm_fail); + } + }, + .group => |gi| { + const fg = groups[gi]; + // the cap's skip BEFORE any upload: layers past the first `run_count` are not + // this build's layers (the head, a past-layer group, is always built) + if (fg.class != .head and fg.layer >= r.first + run_count) continue; + const core_gi = idx.get(item.name) orelse return error.MissingTensor; + const parts = try scan.parts(a, core_gi); + const lidx: usize = fg.layer - r.first; + switch (fg.class) { + .routed => { + const slot = &exps[fg.layer]; + if (slot.* == null) slot.* = .{ .gate = .empty, .up = .empty, .down = .empty }; + const proj: *Projs = switch (fg.proj) { + 0 => &slot.*.?.gate, + 1 => &slot.*.?.up, + else => &slot.*.?.down, + }; + try proj.append(a, .{ + .trellis = parts.trellis.bytes, + .suh = parts.suh.bytes, + .svh = parts.svh.bytes, + .k2 = fg.k2(), + .signs = parts.suh_packed, + }); + }, + .head => head_w = try up.dense(parts, fg), + .attn => switch (fg.proj) { + 0 => layers[lidx].q = try up.dense(parts, fg), + 1 => layers[lidx].k = try up.dense(parts, fg), + 2 => layers[lidx].v = try up.dense(parts, fg), + else => layers[lidx].o = try up.dense(parts, fg), + }, + .dense_mlp => switch (fg.proj) { + 0 => d_gate[lidx] = try up.dense(parts, fg), + 1 => d_up[lidx] = try up.dense(parts, fg), + else => d_down[lidx] = try up.dense(parts, fg), + }, + .shared => switch (fg.proj) { + 0 => s_gate[lidx] = try up.dense(parts, fg), + 1 => s_up[lidx] = try up.dense(parts, fg), + else => s_down[lidx] = try up.dense(parts, fg), + }, + } + }, + } + } + // each owned MoE layer's collected experts upload as one grouped layer, the shared wave + // and the router beside it — then the mlp assembly + for (exps, 0..) |*ep, dno| { + if (ep.*) |*lp| { + if (width == 0) return error.BadExperts; + if (dno < r.first or dno >= r.first + run_count) continue; + const routed = try exl3_upload.layer(&up, a, dims, width, .mul1, c.swiglu_limit, 1, lp.gate.items, lp.up.items, lp.down.items); + const lidx: usize = dno - r.first; + const sg = s_gate[lidx] orelse return error.MissingTensor; + const su = s_up[lidx] orelse return error.MissingTensor; + const sd = s_down[lidx] orelse return error.MissingTensor; + const rt = routers[lidx] orelse return error.MissingTensor; + layers[lidx].mlp = .{ .moe = .{ + .router = rt, + .bias = biases[lidx] orelse &.{}, + .routed = routed, + .sgate = sg, + .sup = su, + .sdown = sd, + } }; + } + } + for (layers, 0..) |*l, i| { + if (l.lno < c.dense_layers) { + l.mlp = .{ .dense = .{ + .gate = d_gate[i] orelse return error.MissingTensor, + .up = d_up[i] orelse return error.MissingTensor, + .down = d_down[i] orelse return error.MissingTensor, + } }; + } else switch (l.mlp) { + .moe => {}, + else => return error.MissingTensor, // an owned MoE layer the walk never assembled + } + } + if (norm_fail) return error.MissingTensor; + + var stream = try cuda.Stream.init(d, false); + errdefer { + if (!built) stream.deinit(); + } + var dense = try exl3.Dense.load(d); + errdefer { + if (!built) dense.unload(); + } + var experts = try exl3e.Experts.load(d); + errdefer { + if (!built) experts.unload(); + } + var mod = try cuda.Module.load(d, cuda.kernels.minimax_m3); + errdefer { + if (!built) mod.unload(); + } + const ks = try sp.Kernels.load(mod); + const gk = try glue.Kernels.load(mod); + + const qk: usize = @as(usize, c.attn_heads) * @as(usize, c.head_dim); + const kv: usize = @as(usize, c.kv_heads) * @as(usize, c.head_dim); + const pages: usize = keys_cap / sp.page; + const nb_max = sp.nbOf(keys_cap); + const wide: usize = @max(dims, @max(@as(usize, c.dense_inter), @max(@as(usize, c.moe_inter), @as(usize, c.shared_inter)))); + const kvb: usize = nl * keys_cap * kv * 2; + const plb: usize = nl * keys_cap * @as(usize, c.i_dim) * 2; + + // a shapes-only Layer for the grouped-experts scratch sizing (the real layers then size + // the same pool: same dims/width/codebook for every MoE layer) + var e = Engine{ + .gpa = gpa, + .io = io, + .c = c.*, + .d = d, + .s = stream, + .dense = dense, + .experts = experts, + .mod = mod, + .ks = ks, + .gk = gk, + .up = up, + .rank = rank, + .world = world, + .lib = o.lib, + .comm = o.comm, + .layers = layers, + .first = r.first, + .embed_tbl = embed, + .final_norm = final_norm, + .head = head_w, + .x = &.{}, + .xn = &.{}, + .xn2 = &.{}, + .iq32 = &.{}, + .ik32 = &.{}, + .q16 = &.{}, + .k16 = &.{}, + .v16 = &.{}, + .iq16 = &.{}, + .ik16 = &.{}, + .projf = &.{}, + .o16 = &.{}, + .elogits = &.{}, + .logits = &.{}, + .sel_e = &.{}, + .pickw = &.{}, + .picks = &.{}, + .sel_seen = &.{}, + .x32 = undefined, + .proj32 = undefined, + .recv32 = undefined, + .ctrl32 = undefined, + .qdev16 = undefined, + .iqdev16 = undefined, + .attn16 = undefined, + .g16 = undefined, + .u16 = undefined, + .a16 = undefined, + .moe32 = undefined, + .pick32 = undefined, + .wts32 = undefined, + .logits32 = undefined, + .kc = undefined, + .vc = undefined, + .plane = undefined, + .scores32 = undefined, + .sel32 = undefined, + .qpos64 = undefined, + .kpos64 = undefined, + .seqlens = undefined, + .table = undefined, + .part32 = undefined, + .exp_scratch = undefined, + }; + // the runtime rows and caches, tracked for the error paths (freed once each: this list + // below `built == true`, e.deinit() above it) + var bufs: [25]cuda.DeviceBuffer = undefined; + var nbufs: usize = 0; + errdefer { + if (!built) for (bufs[0..nbufs]) |*b| b.free(); + } + e.x32 = try allocRow(&bufs, &nbufs, d, dims * 4); + e.proj32 = try allocRow(&bufs, &nbufs, d, (qk + 2 * kv) * 4); + e.recv32 = try allocRow(&bufs, &nbufs, d, dims * 4); + e.ctrl32 = try allocRow(&bufs, &nbufs, d, 8); // the [2] u32 row-control word + e.qdev16 = try allocRow(&bufs, &nbufs, d, qk * 2); + e.iqdev16 = try allocRow(&bufs, &nbufs, d, @as(usize, c.i_heads) * @as(usize, c.i_dim) * 2); + e.attn16 = try allocRow(&bufs, &nbufs, d, qk * 2); + e.g16 = try allocRow(&bufs, &nbufs, d, wide * 2); + e.u16 = try allocRow(&bufs, &nbufs, d, wide * 2); + e.a16 = try allocRow(&bufs, &nbufs, d, wide * 2); + e.moe32 = try allocRow(&bufs, &nbufs, d, dims * 4); + e.pick32 = try allocRow(&bufs, &nbufs, d, @as(usize, c.topk) * 4); + e.wts32 = try allocRow(&bufs, &nbufs, d, @as(usize, c.topk) * 4); + e.logits32 = try allocRow(&bufs, &nbufs, d, if (rank == 0) @as(usize, c.vocab) * 4 else 16); + e.kc = try allocRow(&bufs, &nbufs, d, kvb); + e.vc = try allocRow(&bufs, &nbufs, d, kvb); + e.plane = try allocRow(&bufs, &nbufs, d, plb); + e.scores32 = try allocRow(&bufs, &nbufs, d, @as(usize, c.i_heads) * nb_max * 4); // scores: (B, i_heads, Sq, NB) fp32 + e.sel32 = try allocRow(&bufs, &nbufs, d, @as(usize, c.kv_heads) * @as(usize, c.i_topk) * 4); // sel: (B, Hkv, i_topk) i32 — the BLOCK selection width (16), NOT the router's topk + e.qpos64 = try allocRow(&bufs, &nbufs, d, 8); + e.kpos64 = try allocRow(&bufs, &nbufs, d, keys_cap * 8); + e.seqlens = try allocRow(&bufs, &nbufs, d, 8); + e.table = try allocRow(&bufs, &nbufs, d, pages * 4); + e.part32 = try allocRow(&bufs, &nbufs, d, sp.scratchBytes(@as(usize, c.attn_heads), 1, nsplit)); + e.exp_scratch = blk: { + const l: exl3e.Layer = .{ + .gate = 0, + .up = 0, + .down = 0, + .gate_k2 = 0, + .up_k2 = 0, + .down_k2 = 0, + .suh_g = 0, + .suh_u = 0, + .svh_g = 0, + .svh_u = 0, + .suh_d = 0, + .svh_d = 0, + .experts = 1, + .dims = dims, + .width = if (width > 0) width else 128, + .codebook = .mul1, + }; + break :blk try allocRow(&bufs, &nbufs, d, if (width > 0) try exl3e.scratchBytes(l, 1, c.topk) else 4096); + }; + errdefer e.deinit(); + built = true; + + // the host state + e.x = try gpa.alloc(f32, dims); + e.xn = try gpa.alloc(f32, dims); + e.xn2 = try gpa.alloc(f32, dims); + e.iq32 = try gpa.alloc(f32, @as(usize, c.i_heads) * @as(usize, c.i_dim)); + e.ik32 = try gpa.alloc(f32, @as(usize, c.i_dim)); + e.q16 = try gpa.alloc(f16, qk); + e.k16 = try gpa.alloc(f16, kv); + e.v16 = try gpa.alloc(f16, kv); + e.iq16 = try gpa.alloc(f16, e.iq32.len); + e.ik16 = try gpa.alloc(f16, e.ik32.len); + e.projf = try gpa.alloc(f32, qk + 2 * kv); + e.o16 = try gpa.alloc(f16, wide); + e.elogits = try gpa.alloc(f32, c.experts); + e.logits = try gpa.alloc(f32, if (rank == 0) c.vocab else 1); + e.sel_e = try gpa.alloc(u32, c.topk); + e.pickw = try gpa.alloc(f32, c.topk); + e.picks = try gpa.alloc(bool, c.experts); + e.sel_seen = try gpa.alloc(i32, @as(usize, c.kv_heads) * @as(usize, c.i_topk)); // the trace mirrors sel32's (B, Hkv, i_topk) width + @memset(e.picks, false); + @memset(e.sel_seen, -1); + // the identity page table and the content keys' mask row, once + const tab = try a.alloc(i32, pages); + for (tab, 0..) |*t, i| t.* = @intCast(i); + try e.table.upload(0, std.mem.sliceAsBytes(tab)); + const kpos = try a.alloc(i64, keys_cap); + for (kpos, 0..) |*t, i| t.* = @intCast(i); + try e.kpos64.upload(0, std.mem.sliceAsBytes(kpos)); + if (rank == 0) { + if (e.embed_tbl == null or e.final_norm == null or e.head == null) return error.MissingTensor; + } + return e; + } + + /// One row: rank 0 embeds `tok`, every rank runs its own layers, the world-2 handoffs move + /// the residual row between them, rank 0 closes with the final norm + lm_head. The row's + /// fp32 logits land in `logits` (rank 0). + pub fn put(e: *Engine, tok: u32) !void { + const pos = e.len; + if (pos >= keys_cap) return error.ContextCap; + e.len += 1; + const sk: usize = pos + 1; + const nb = sp.nbOf(sk); + const dims = e.hidden(); + if (e.world == 2) { + const lib = e.lib orelse return error.NoRank; + const comm = e.comm orelse return error.NoRank; + if (e.rank == 1) { + try lib.check(lib.api.ncclRecv(e.recv32.ptr, dims, .f32, 0, comm, e.s.handle), "ncclRecv"); + try e.s.synchronize(); + try e.recv32.download(0, std.mem.sliceAsBytes(e.x)); + } else { + try e.embedRow(tok); + } + } else { + try e.embedRow(tok); + } + for (e.layers, 0..) |*l, i| try e.layerRow(l, pos, sk, nb, i); + if (e.world == 2) { + const lib = e.lib orelse return error.NoRank; + const comm = e.comm orelse return error.NoRank; + if (e.rank == 0) { + try e.x32.upload(0, std.mem.sliceAsBytes(e.x)); + try lib.check(lib.api.ncclSend(e.x32.ptr, dims, .f32, 1, comm, e.s.handle), "ncclSend"); + try lib.check(lib.api.ncclRecv(e.recv32.ptr, dims, .f32, 1, comm, e.s.handle), "ncclRecv"); + try e.s.synchronize(); + try e.recv32.download(0, std.mem.sliceAsBytes(e.x)); + } else { + try e.x32.upload(0, std.mem.sliceAsBytes(e.x)); + try lib.check(lib.api.ncclSend(e.x32.ptr, dims, .f32, 0, comm, e.s.handle), "ncclSend"); + try e.s.synchronize(); + } + } + if (e.rank == 0) try e.headRow(); + } + + // ------- world 2's row-control wire (the HTTP server boot's mirror contract) ------- + // + // The lockstep contract is row-count: every `put` is one paired NCCL row (send/recv each + // way), the mirror's token value is rank 0's business. But the lane surface changes the + // cache length OUTSIDE a put — a stream's release rolls it to 0, a prefix keep rolls it + // back — and the mirror would otherwise key its pool rows at a stale length. So world 2's + // surface rows carry one 8-byte control word ahead of the data row, and release/keep send + // a control-only exchange. The math never changes: same kernel work, same collectives for + // the row's values, one extra tiny u32 pair per row. + + /// The control word: op 1 = a data row follows (`tok` is rank 0's row token), op 0 = the + /// mirror rolls its cache length to 0 (release), op 2 = the mirror rolls it to `tok` + /// (a prefix keep). + pub const ctrl_row: u32 = 1; + pub const ctrl_reset: u32 = 0; + pub const ctrl_keep: u32 = 2; + + fn ctrlSend(e: *Engine, op: u32, tok: u32) !void { + const lib = e.lib orelse return error.NoRank; + const comm = e.comm orelse return error.NoRank; + var host = [2]u32{ op, tok }; + try e.ctrl32.upload(0, std.mem.sliceAsBytes(&host)); + try lib.check(lib.api.ncclSend(e.ctrl32.ptr, 2, .u32, 1, comm, e.s.handle), "ncclSend"); + try e.s.synchronize(); + } + + fn ctrlRecv(e: *Engine) !struct { op: u32, tok: u32 } { + const lib = e.lib orelse return error.NoRank; + const comm = e.comm orelse return error.NoRank; + try lib.check(lib.api.ncclRecv(e.ctrl32.ptr, 2, .u32, 0, comm, e.s.handle), "ncclRecv"); + try e.s.synchronize(); + var host = [2]u32{ 0, 0 }; + try e.ctrl32.download(0, std.mem.sliceAsBytes(&host)); + return .{ .op = host[0], .tok = host[1] }; + } + + /// rank 0's lane-surface row under world 2: the control word first, then the paired row. + pub fn leaderRow(e: *Engine, tok: u32) !void { + if (e.world != 2) return e.put(tok); + try e.ctrlSend(ctrl_row, tok); + try e.put(tok); + } + + /// rank 0's prefix keep under world 2: the mirror rolls its cache length to `kept_len`. + pub fn leaderKeep(e: *Engine, kept_len: u32) !void { + if (e.world != 2 or e.rank != 0) return; + try e.ctrlSend(ctrl_keep, kept_len); + } + + /// rank 0's stream release under world 2: the mirror rolls its cache length to 0. + pub fn leaderReset(e: *Engine) !void { + if (e.world != 2 or e.rank != 0) return; + try e.ctrlSend(ctrl_reset, 0); + } + + /// rank 1's mirror step: the control word, then its paired work (a `.row` runs the full + /// put — its cache length advances exactly like rank 0's; its row's token value is rank + /// 0's business, exactly the receipted lane-tp2 lockstep). An unknown control word + /// refuses rather than desyncing the ranks. + pub fn mirrorStep(e: *Engine) !void { + if (e.world != 2) return; + const c = try e.ctrlRecv(); + switch (c.op) { + ctrl_row => try e.put(c.tok), + ctrl_reset => e.len = 0, + ctrl_keep => e.len = c.tok, + else => return error.NoRank, + } + } + + /// Greedy continuation (rank 0): `steps` first-max argmax picks, each picked token fed back. + /// The peer rank calls `put` the same number of times to stay in the collectives' lockstep. + pub fn greedy(e: *Engine, steps: usize, out: []u32) !void { + std.debug.assert(out.len >= steps); + for (0..steps) |g| { + var best: usize = 0; + var bv = -std.math.inf(f32); + for (e.logits, 0..) |v, i| if (v > bv) { + bv = v; + best = i; + }; + out[g] = @intCast(best); + if (g + 1 < steps) try e.put(@intCast(best)); + } + } + + fn embedRow(e: *Engine, tok: u32) !void { + const dims = e.hidden(); + const b = e.embed_tbl orelse return error.MissingTensor; + const row = try e.gpa.alloc(u8, dims * 2); + defer e.gpa.free(row); + try b.download(@as(usize, tok) * dims * 2, row); // one bf16 [hidden] table row + for (e.x, 0..) |*v, i| { + const bits: u32 = @as(u32, std.mem.readInt(u16, row[i * 2 ..][0..2], .little)) << 16; + v.* = @bitCast(bits); + } + } + + fn headRow(e: *Engine) !void { + serve.normRow(e.x, e.final_norm.?, e.c.eps); + try e.x32.upload(0, std.mem.sliceAsBytes(e.x)); + try e.dense.decode(e.s, e.head.?, e.x32.ptr, .f32, 1, e.logits32.ptr, .f32); + try e.s.synchronize(); + try e.logits32.download(0, std.mem.sliceAsBytes(e.logits)); + } + + /// The dense-shaped wave: gate/up GEMMs, tf_m3_swigluoai, the down GEMM; the fp16 rows + /// land on the host in `o16` (the caller's fp16 rounding point on the residual add). + fn waveRow(e: *Engine, gate: exl3.Weight, up: exl3.Weight, down: exl3.Weight, inter: u32) !void { + const dims = e.hidden(); + const n: usize = inter; + try e.dense.decode(e.s, gate, e.x32.ptr, .f32, 1, e.g16.ptr, .f16); + try e.dense.decode(e.s, up, e.x32.ptr, .f32, 1, e.u16.ptr, .f16); + try glue.swigluoai(e.gk, e.s, e.g16.ptr, e.u16.ptr, e.a16.ptr, 1, n, e.c.swiglu_limit); + try e.dense.decode(e.s, down, e.a16.ptr, .f16, 1, e.g16.ptr, .f16); + try e.s.synchronize(); + try e.g16.download(0, std.mem.sliceAsBytes(e.o16[0..dims])); + } + + /// One owned layer's row, attention half then MLP wave, serve.put's contracts line for line. + fn layerRow(e: *Engine, l: *Layer, pos: usize, sk: usize, nb: usize, lidx: usize) !void { + const c = &e.c; + const dims = e.hidden(); + const qk = e.qkr(); + const kv = e.kvr(); + const hd: usize = c.head_dim; + const qh: usize = c.attn_heads; + const kvh: usize = c.kv_heads; + const rot: usize = c.rotary_dim; + const lbytes: usize = lidx * keys_cap; // this layer's first row in the caches/plane + const pages: usize = keys_cap / sp.page; + const scale: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(c.head_dim))); + + // -- the attention half: gemma norm -> q/k/v GEMMs -> per-head qk norm + partial rope -> + // fp16 stores -> the indexer's side plane -> selection -> block-sparse attention -> o + @memcpy(e.xn, e.x); + serve.normRow(e.xn, l.in_norm, c.eps); + try e.x32.upload(0, std.mem.sliceAsBytes(e.xn)); + const q_w = l.q orelse return error.MissingTensor; + const k_w = l.k orelse return error.MissingTensor; + const v_w = l.v orelse return error.MissingTensor; + const o_w = l.o orelse return error.MissingTensor; + try e.dense.decode(e.s, q_w, e.x32.ptr, .f32, 1, e.proj32.ptr, .f32); + try e.dense.decode(e.s, k_w, e.x32.ptr, .f32, 1, e.proj32.ptr + qk * 4, .f32); + try e.dense.decode(e.s, v_w, e.x32.ptr, .f32, 1, e.proj32.ptr + (qk + kv) * 4, .f32); + try e.s.synchronize(); + try e.proj32.download(0, std.mem.sliceAsBytes(e.projf)); + + // q: per-head gemma norm + rope, one fp16 rounding; k: norm + rope into the cache row; + // v: fp16 into the cache row (serve.put's stores, the fork's cache contract) + const qs = e.projf[0..qk]; + const ks_ = e.projf[qk..][0..kv]; + const vs = e.projf[qk + kv ..][0..kv]; + for (0..qh) |h| { + const row = qs[h * hd ..][0..hd]; + serve.normRow(row, l.q_norm, c.eps); + serve.ropeRow(row, pos, rot); + for (row, 0..) |v, i| e.q16[h * hd + i] = @floatCast(serve.f16r(v)); + } + for (0..kvh) |h| { + const row = ks_[h * hd ..][0..hd]; + serve.normRow(row, l.k_norm, c.eps); + serve.ropeRow(row, pos, rot); + for (row, 0..) |v, i| e.k16[h * hd + i] = @floatCast(serve.f16r(v)); + const vrow = vs[h * hd ..][0..hd]; + for (vrow, 0..) |v, i| e.v16[h * hd + i] = @floatCast(serve.f16r(v)); + } + const co: usize = (lbytes + pos) * kv * 2; + try e.kc.upload(co, std.mem.sliceAsBytes(e.k16)); + try e.vc.upload(co, std.mem.sliceAsBytes(e.v16)); + try e.qdev16.upload(0, std.mem.sliceAsBytes(e.q16)); + + // the indexer's side path: the fp16 passthrough projections (host GEMV, fp32 math over + // fp16 weights), per-head norm + rope, fp16 stores, the side-plane append + const iq_w = l.iq orelse return error.MissingTensor; + const ik_w = l.ik orelse return error.MissingTensor; + try gemvF16(e.xn, iq_w, e.iq32); + try gemvF16(e.xn, ik_w, e.ik32); + const id: usize = c.i_dim; + for (0..c.i_heads) |h| { + const row = e.iq32[h * id ..][0..id]; + serve.normRow(row, l.iq_norm, c.eps); + serve.ropeRow(row, pos, rot); + for (row, 0..) |v, i| e.iq16[h * id + i] = @floatCast(serve.f16r(v)); + } + serve.normRow(e.ik32, l.ik_norm, c.eps); + serve.ropeRow(e.ik32, pos, rot); + for (e.ik32, 0..) |v, i| e.ik16[i] = @floatCast(serve.f16r(v)); + try e.iqdev16.upload(0, std.mem.sliceAsBytes(e.iq16)); + try e.plane.upload((lbytes + pos) * id * 2, std.mem.sliceAsBytes(e.ik16)); + + // the selection: tf_m3_scores over the fp16 rows (fp32 scoring, the own-block +inf boost) + // then tf_m3_topk's fixed-width picks (descending, ties at the smaller id, -1 pad) + try e.qpos64.upload(0, std.mem.sliceAsBytes(&[1]i64{@intCast(pos)})); + try e.seqlens.upload(0, std.mem.sliceAsBytes(&[1]i64{@intCast(sk)})); + try sp.scores(e.ks, e.s, e.iqdev16.ptr, e.plane.ptr + lbytes * id * 2, e.qpos64.ptr, e.kpos64.ptr, e.scores32.ptr, 1, nb, sk, c.i_heads, 1, 1, 0); + try sp.topk16(e.ks, e.s, e.scores32.ptr, e.sel32.ptr, nb, kvh); + // the trace: the last layer's selection (the smoke reads this) + if (lidx + 1 == e.layers.len) { + try e.s.synchronize(); + try e.sel32.download(0, std.mem.sliceAsBytes(e.sel_seen)); + var valid: usize = 0; + for (e.sel_seen[e.sel_seen.len - @as(usize, c.i_topk) ..]) |v| { + if (v >= 0) valid += 1; + } + e.last.sel_valid = valid; + } + // the block-sparse attention: the splits read the paged cache by the selection, the + // combine merges them into the fp16 rows + try sp.decode(e.ks, e.s, e.qdev16.ptr, e.kc.ptr + lbytes * kv * 2, e.vc.ptr + lbytes * kv * 2, e.table.ptr, e.sel32.ptr, e.seqlens.ptr, e.part32.ptr, qh, kvh, pages, scale, 1, nsplit); + try sp.combine(e.ks, e.s, e.part32.ptr, e.attn16.ptr, qh, 1, nsplit); + try e.dense.decode(e.s, o_w, e.attn16.ptr, .f16, 1, e.g16.ptr, .f16); + try e.s.synchronize(); + try e.g16.download(0, std.mem.sliceAsBytes(e.o16[0..dims])); + for (0..dims) |i| e.x[i] += serve.f16r(@as(f32, e.o16[i])); + + // -- the MLP wave: gemma norm over the residual, dense (layers 0-2) or routed + shared + @memcpy(e.xn2, e.x); + serve.normRow(e.xn2, l.post_norm, c.eps); + try e.x32.upload(0, std.mem.sliceAsBytes(e.xn2)); + switch (l.mlp) { + .dense => |m| { + try e.waveRow(m.gate, m.up, m.down, c.dense_inter); + for (0..dims) |i| e.x[i] += serve.f16r(@as(f32, e.o16[i])); + }, + .moe => |m| { + // the noaux sigmoid+bias router on the host, serve.route's exact fp32 contract: + // select on sigmoid+bias (strict >, ties at the smaller id), weights = unbiased + // sigmoid normalized over the picks, * routed_scale, each .half() (fp16-rounded + // values carried as f32, the device epilogue's weight tensor's width) + serve.route(e.xn2, m.router, m.bias, e.picks, e.elogits, e.sel_e, e.pickw, c); + var fires: usize = 0; + var wsum: f32 = 0; + for (e.picks) |p| { + if (p) fires += 1; + } + for (e.pickw) |w| wsum += w; + e.last.fired = fires; + e.last.wsum = wsum; + try e.pick32.upload(0, std.mem.sliceAsBytes(e.sel_e)); + try e.wts32.upload(0, std.mem.sliceAsBytes(e.pickw)); + const sc = try exl3e.scratch(e.exp_scratch.ptr, m.routed, 1, c.topk); + try e.experts.run(e.s, m.routed, e.x32.ptr, dims, .f32, e.pick32.ptr, e.wts32.ptr, 1, e.moe32.ptr, sc); + // the shared expert wave beside the routed one + try e.waveRow(m.sgate, m.sup, m.sdown, c.shared_inter); + try e.moe32.download(0, std.mem.sliceAsBytes(e.xn)); // xn reused as the routed download + for (0..dims) |i| e.x[i] += serve.f16r(e.xn[i] + @as(f32, e.o16[i])); + }, + } + } +}; + +/// One tracked runtime-buffer allocation: appended to the caller's build-time list so a +/// mid-literal failure frees exactly what landed, and e.deinit() owns it after construction. +fn allocRow(bufs: *[25]cuda.DeviceBuffer, n: *usize, d: *const cuda.Driver, bytes: usize) !cuda.DeviceBuffer { + const b = try cuda.DeviceBuffer.alloc(d, bytes); + bufs[n.*] = b; + n.* += 1; + return b; +} + +fn freeBuf(b: *cuda.DeviceBuffer) void { + if (b.len > 0) b.free(); + b.len = 0; +} + +fn deinitLayer(gpa: std.mem.Allocator, l: *Layer) void { + for ([_]*[]f32{ &l.in_norm, &l.post_norm, &l.q_norm, &l.k_norm, &l.iq_norm, &l.ik_norm }) |p| { + if (p.len > 0) gpa.free(p.*); + } + if (l.iq) |w| gpa.free(w); + if (l.ik) |w| gpa.free(w); + switch (l.mlp) { + .dense => {}, + .moe => |m| { + if (m.router.len > 0) gpa.free(m.router); + if (m.bias.len > 0) gpa.free(m.bias); + }, + } +} + +/// y = x @ W^T exactly as serve.linear accumulates (fp32 over fp16 weights widened per element). +fn gemvF16(x: []const f32, w: []const f16, y: []f32) !void { + const n_in = x.len; + if (w.len != y.len * n_in) return error.MissingTensor; + for (0..y.len) |o| { + var acc: f32 = 0; + const row = w[o * n_in ..][0..n_in]; + for (0..n_in) |j| acc += x[j] * @as(f32, row[j]); + y[o] = acc; + } +} + +/// bf16 -> f32 widening (exact) of one norm row. +fn expandNorm(gpa: std.mem.Allocator, bytes: []const u8, n: usize) ![]f32 { + if (bytes.len != n * 2) return error.MissingTensor; + const w = try gpa.alloc(f32, n); + errdefer gpa.free(w); + for (0..n) |i| { + const bits: u32 = @as(u32, std.mem.readInt(u16, bytes[i * 2 ..][0..2], .little)) << 16; + w[i] = @bitCast(bits); + } + return w; +} + +/// The checkpoint layer id a passthrough name carries ("language_model.model.layers.N."); a +/// foreign tower's or non-layer passthrough reads without a slot (the fixture's only one). +fn layerOf(name: []const u8) ?u32 { + const mark = "layers."; + const at = std.mem.indexOf(u8, name, mark) orelse return null; + var i = at + mark.len; + var v: u32 = 0; + while (i < name.len and name[i] >= '0' and name[i] <= '9') : (i += 1) { + v = v * 10 + (name[i] - '0'); + } + if (i == at + mark.len) return null; + return v; +} + +fn ends(name: []const u8, suffix: []const u8) bool { + return std.mem.endsWith(u8, name, suffix); +} + +/// One passthrough tensor beside the plan: the norms expand to (1 + w) fp32 host rows, the +/// indexer's fp16 projections land on the host for the small GEMVs, the router and its bias stay +/// raw fp32 arrays (serve.route's inputs). The embed and the final norm the builder reads itself. +/// Missing-but-needed tensors set `fail` instead of erroring mid-walk (a per-tensor log first). +fn readPassthrough( + gpa: std.mem.Allocator, + c: *const config.Config, + layers: []Layer, + first: u32, + name: []const u8, + bytes: []const u8, + routers: []?[]f32, + biases: []?[]f32, + fail: *bool, +) !void { + const lno = layerOf(name) orelse return; + if (lno < first) return; + const lidx: usize = lno - first; + if (lidx >= layers.len) return; + const l = &layers[lidx]; + if (ends(name, "index_q_norm.weight")) { + l.iq_norm = expNorm(fail, gpa, bytes, c.i_dim); + } else if (ends(name, "index_k_norm.weight")) { + l.ik_norm = expNorm(fail, gpa, bytes, c.i_dim); + } else if (ends(name, "input_layernorm.weight")) { + l.in_norm = expNorm(fail, gpa, bytes, c.hidden); + } else if (ends(name, "post_attention_layernorm.weight")) { + l.post_norm = expNorm(fail, gpa, bytes, c.hidden); + } else if (ends(name, "q_norm.weight")) { + l.q_norm = expNorm(fail, gpa, bytes, c.head_dim); + } else if (ends(name, "k_norm.weight")) { + l.k_norm = expNorm(fail, gpa, bytes, c.head_dim); + } else if (ends(name, "index_q_proj.weight")) { + l.iq = passF16(fail, gpa, bytes); + } else if (ends(name, "index_k_proj.weight")) { + l.ik = passF16(fail, gpa, bytes); + } else if (ends(name, "block_sparse_moe.gate.weight")) { + if (bytes.len != @as(usize, c.experts) * @as(usize, c.hidden) * 4) { + fail.* = true; + } else { + const r = try gpa.alloc(f32, @as(usize, c.experts) * @as(usize, c.hidden)); + @memcpy(std.mem.sliceAsBytes(r), bytes); + if (routers[lidx] != null) gpa.free(routers[lidx].?); + routers[lidx] = r; + } + } else if (ends(name, "e_score_correction_bias")) { + if (bytes.len != @as(usize, c.experts) * 4) { + fail.* = true; + } else { + const r = try gpa.alloc(f32, c.experts); + @memcpy(std.mem.sliceAsBytes(r), bytes); + if (biases[lidx] != null) gpa.free(biases[lidx].?); + biases[lidx] = r; + } + } else { + return error.MissingTensor; // a passthrough the graph reads by no other spelling + } +} + +fn set(fail: *bool) error{MissingTensor} { + fail.* = true; + return error.MissingTensor; +} + +/// The gemma (1 + w) expansion of one bf16 norm row, f32. +fn expNorm(fail: *bool, gpa: std.mem.Allocator, bytes: []const u8, n: usize) []f32 { + return expandGemma(gpa, bytes, n) catch { + fail.* = true; + return &.{}; + }; +} + +/// The fp16 passthrough mat, []f16. +fn passF16(fail: *bool, gpa: std.mem.Allocator, bytes: []const u8) []f16 { + return copyF16(gpa, bytes) catch { + fail.* = true; + return &.{}; + }; +} + +fn expandGemma(gpa: std.mem.Allocator, bytes: []const u8, n: usize) ![]f32 { + if (bytes.len != n * 2) return error.MissingTensor; + const w = try gpa.alloc(f32, n); + errdefer gpa.free(w); + for (0..n) |i| { + const bits: u32 = @as(u32, std.mem.readInt(u16, bytes[i * 2 ..][0..2], .little)) << 16; + // The checkpoint gemma weight is ABSOLUTE (~1.0); serve.normRow applies its own + // (1.0 + wi) — adding another 1 here made every layer's norm 1.5x the reference's. + w[i] = @as(f32, @bitCast(bits)); + } + return w; +} + +fn copyF16(gpa: std.mem.Allocator, bytes: []const u8) ![]f16 { + if (bytes.len % 2 != 0) return error.MissingTensor; + const w = try gpa.alloc(f16, bytes.len / 2); + @memcpy(std.mem.sliceAsBytes(w), bytes); + return w; +} diff --git a/zig/src/families/minimax_m3_vl/serve.zig b/zig/src/families/minimax_m3_vl/serve.zig new file mode 100644 index 0000000000..40a139d26b --- /dev/null +++ b/zig/src/families/minimax_m3_vl/serve.zig @@ -0,0 +1,1107 @@ +//! The MSA serve engine's forward on the host — the pre-device bar the CUDA path grades against. +//! The math is the exllamav3-fork oracle's, contract by contract (the receipts' module lines): +//! gemma RMSNorms (`(1.0 + weight)` then one dtype round), per-head qk norms then partial neox +//! rope (64 of 128 dims, theta 1e7), the MSA block attention over the indexer's top-16-of-block +//! selection (boost overrides, ties at the smaller id, -1 pads) reading the fp16 rows it stored, +//! swigluoai (`(clamp(u,-limit,limit)+1) * min(g,limit) * sigmoid(1.702*min(g,limit))`), the noaux +//! sigmoid+correction-bias router (select on sigmoid+bias, weights = unbiased sigmoid normalized +//! over the picks then `* routed_scaling_factor`, returned `.half()` per `block_sparse_mlp_routing.py`), +//! the shared expert at weight 1, a leading dense block, then the MoE wave; the final gemma norm and +//! the head close the graph. One rounding contract everywhere (documented for device parity): every +//! linear/cache/normed-row store sits at fp16 exactly where the fork's fp16 linears and caches sit; +//! softmax, routing sums and the residual stream stay fp32. Weights here are the tests' SYNTHETIC +//! values — shapes carry the pack's semantics, magnitudes are fixtures, not checkpoint numbers. + +const std = @import("std"); +const cfg = @import("config.zig"); +const sparse = @import("sparse.zig"); + +pub const Config = cfg.Config; + +/// The rope's theta (the pack's rope_theta 1e7 — `ropeRow`'s fixed constant, f64 internally). +const theta: f64 = 1e7; + +/// The tests' serve config: the pack's real contract numbers (64 q heads of 128 over 4 kv heads, +/// rotary 64 of 128, index 4 x 128, sigmoid+bias router picking 4 of 8 (num_experts_per_tok 4), +/// alpha 1.702 / limit 7 / scaling 2.0) at test-sized counts (hidden 256 vocab 256, one dense + +/// one MoE layer, 8 experts) so a prefill with several key blocks parses and runs on the host in +/// seconds. `block` 16 (not the pack's 128) keeps the block-count grid dense at this length: the +/// selection's pooling/boost/tie/pad contracts need NB > 1 and NB < width to be graded here, and a +/// block-128 run at this size would pool to NB 1. The pack's 128 stays the default in config.zig. +pub fn serveConfig() Config { + return .{ + .hidden = 256, + .layers = 2, + .run = 2, + .vocab = 256, + .eps = 1e-6, + .dense_layers = 1, // layer 0's dense MLP, layer 1's MoE: both block shapes in one pass + .dense_inter = 1024, + .experts = 8, + .topk = 4, // num_experts_per_tok of the pack (verified against pack_config.json) + .moe_inter = 512, + .shared_inter = 512, + .routed_scale = 2.0, + .sigmoid_router = true, + .routing_bias = true, + .norm_topk = false, + .swiglu_alpha = 1.702, + .swiglu_limit = 7.0, + .attn_heads = 64, + .kv_heads = 4, + .head_dim = 128, + .rotary_dim = 64, + .partial_rotary = 0.5, + .qk_norm = true, + .i_heads = 4, + .i_dim = 128, + .i_topk = 16, + .block = 16, // the test's block width (pack: 128): NB stays multi-block at the test length + .score_max = true, + }; +} + +/// The whole context the fixture serves: the tests' prefill plus decode steps. +pub const prefill_len: usize = 66; +pub const decode_steps: usize = 2; +pub const capacity: usize = prefill_len + decode_steps + 1; + +// head geometry (fixed by serveConfig's numbers): +pub const hidden: usize = 256; +pub const q_heads: usize = 64; +pub const kv_heads: usize = 4; +pub const head_dim: usize = 128; +const qk_rows: usize = q_heads * head_dim; // 8192 the packed per-head q rows of one token +const kv_rows: usize = kv_heads * head_dim; // 512 +const i_rows: usize = 4 * 128; // the index queries: i_heads * i_dim +const inter_max: usize = 1024; // max(moe_inter 512, shared_inter 512, dense_inter 1024) + +/// One fp32 value rounded through fp16 — the fp16-store contract's rounding. +pub fn f16r(x: f32) f32 { + return @floatCast(@as(f16, @floatCast(x))); +} + +/// One fp32 value rounded through bf16 — the norm weights' storage width. +pub fn b16r(x: f32) f32 { + var u: u32 = @bitCast(x); + const lsb = (u >> 16) & 1; + u += 0x7FFF + lsb; // round to nearest even above the bf16 mantissa + return @bitCast(u & 0xFFFF0000); +} + +/// gemma RMSNorm of one row in place: x * rsqrt(mean(x^2) + eps) * (1 + w), fp32 math. +pub fn normRow(x: []f32, w: []const f32, eps: f32) void { + var ss: f64 = 0; + for (x) |v| ss += @as(f64, v) * @as(f64, v); + const inv: f32 = @floatCast(1.0 / @sqrt(ss / @as(f64, @floatFromInt(x.len)) + @as(f64, eps))); + for (x, w) |*v, wi| v.* = v.* * inv * (1.0 + wi); +} + +/// y[o] = sum_c x[c] * W[o, c] with W row-major [n_out, n_in] (the checkpoint's linear layout). +fn linear(x: []const f32, w: []const f32, n_out: usize, y: []f32) void { + const n_in = x.len; + @memset(y, 0); + for (0..n_in) |c| { + const xv = x[c]; + for (0..n_out) |o| y[o] += xv * w[o * n_in + c]; + } +} + +/// Partial neox rope of `rot` channels in place: pair (i, i + rot/2) for i < rot/2 at theta; the +/// tail stays untouched (cos 1, sin 0 there — the fork's rotate_half over `partial_rotary` 0.5). +pub fn ropeRow(x: []f32, pos: usize, rot: usize) void { + const half = rot / 2; + for (0..half) |i| { + const ang = @as(f64, @floatFromInt(pos)) / + std.math.pow(f64, theta, @as(f64, @floatFromInt(2 * i)) / @as(f64, @floatFromInt(rot))); + const c: f32 = @floatCast(@cos(ang)); + const s: f32 = @floatCast(@sin(ang)); + const a = x[i]; + const b = x[i + half]; + x[i] = a * c - b * s; + x[i + half] = b * c + a * s; + } +} + +/// swigluoai per the fork's activation: gate `g` clipped above, `u` clamped both ways and +1, +/// g * sigmoid(alpha * g) * u. +pub fn swigluoai(g_in: f32, u_in: f32, alpha: f32, limit: f32) f32 { + const g = @min(g_in, limit); + const u = @min(@max(u_in, -limit), limit) + 1.0; + return g / (1.0 + @exp(-alpha * g)) * u; +} + +/// The noaux sigmoid+correction-bias router's picks and weights: select on sigmoid+bias, weights = +/// unbiased sigmoid normalized over the picks, then `* routed_scaling_factor`, each `.half()`. +/// Ties at the smaller expert id (torch.topk's stable order). +pub fn route( + state: []const f32, + router: []const f32, // [E, hidden] + bias: []const f32, // [E] or empty + picks: []bool, // out: which experts fire (cleared here) + logits_e: []f32, // scratch: the raw router logits [E] + sel_e: []u32, // out: the pick ids in selection order + weights: []f32, // out: the per-pick weights [.half() rounded] + c: *const Config, +) void { + const E: usize = c.experts; + for (0..E) |e| { + var acc: f32 = 0; + for (0..state.len) |j| acc += state[j] * router[e * state.len + j]; + logits_e[e] = acc; + picks[e] = false; + } + const k: usize = c.topk; + var chosen: usize = 0; + for (0..k) |slot| { + var best: usize = E; + var best_score = -std.math.inf(f32); + for (0..E) |e| { + if (picks[e]) continue; + const sigmoid = 1.0 / (1.0 + @exp(-logits_e[e])); // scoring_func "sigmoid", no softmax + const choice: f32 = sigmoid + (if (bias.len > 0) bias[e] else 0); + if (choice > best_score) { // strict >: ties stay at the smaller id + best_score = choice; + best = e; + } + } + if (best == E) break; // fewer experts than picks + picks[best] = true; + sel_e[slot] = @intCast(best); + chosen += 1; + } + var sum: f32 = 0; + for (0..chosen) |slot| { + weights[slot] = 1.0 / (1.0 + @exp(-logits_e[sel_e[slot]])); // the unbiased sigmoid + sum += weights[slot]; + } + sum += 1e-20; // the router's denominator guard + for (0..chosen) |slot| weights[slot] = f16r(weights[slot] / sum * c.routed_scale); +} + +/// One layer's weights (the loader's tensor classes; the fixture's synthetic values). +const Mlp = struct { gate: []f32, up: []f32, down: []f32 }; // gate/up [inter, hidden], down [hidden, inter] + +const Layer = struct { + in_norm: []f32, // [hidden] bf16-rounded values + post_norm: []f32, + q_norm: []f32, // [head_dim]: every head norms by the one row (use_qk_norm's shape) + k_norm: []f32, + q: []f32, // [qk_rows, hidden] + k: []f32, // [kv_rows, hidden] + v: []f32, // [kv_rows, hidden] + o: []f32, // [hidden, qk_rows] + iq: []f32, // [i_rows, hidden] fp16-rounded values (indexer.wq_b: per head [i_heads*i_dim]) + ik: []f32, // [i_dim, hidden] fp16-rounded: ONE shared index-K row set (mini k_proj out=head_dim) + iqn: []f32, // [i_dim] + ikn: []f32, + mlp: union(enum) { dense: Mlp, moe: MoE }, + mlp_set: bool = false, // true only once the union below is fully built (the deinit's guard) + + const MoE = struct { + router: []f32, // [E, hidden] fp32 raw (the router gate is not quantized) + bias: []f32, // [E] fp32 + shared: Mlp, + experts: []Mlp, + }; + + fn deinit(l: *Layer, gpa: std.mem.Allocator) void { + gpa.free(l.in_norm); + gpa.free(l.post_norm); + gpa.free(l.q_norm); + gpa.free(l.k_norm); + gpa.free(l.q); + gpa.free(l.k); + gpa.free(l.v); + gpa.free(l.o); + gpa.free(l.iq); + gpa.free(l.ik); + gpa.free(l.iqn); + gpa.free(l.ikn); + if (l.mlp_set) { + switch (l.mlp) { + .dense => |m| { + gpa.free(m.gate); + gpa.free(m.up); + gpa.free(m.down); + }, + .moe => |m| { + gpa.free(m.router); + gpa.free(m.bias); + gpa.free(m.shared.gate); + gpa.free(m.shared.up); + gpa.free(m.shared.down); + for (m.experts) |e| { + gpa.free(e.gate); + gpa.free(e.up); + gpa.free(e.down); + } + gpa.free(m.experts); + }, + } + } + } +}; + +/// A deterministic synthetic fixture: every weight a pure function of the LCG stream (no RNG), +/// rounded through the storage width its real tensor carries (norms bf16, indexer rows fp16). +pub const seed_default: u64 = 0x5f3759df_9e3779b9; + +pub const Lcg = struct { + s: u64, + fn next(l: *Lcg) u32 { + l.s = l.s *% 6364136223846793005 +% 1442695040888963407; + return @intCast(l.s >> 33); + } + fn f32u(l: *Lcg) f32 { + const u: f32 = @floatFromInt(l.next() | 1); + return u / 4294967296.0; + } + /// uniform-ish in [-scale, scale] + fn val(l: *Lcg, scale: f32) f32 { + return (l.f32u() - 0.5) * 2.0 * scale; + } +}; + +const Rnd = enum { plain, f16, bf16 }; + +fn mat(gpa: std.mem.Allocator, l: *Lcg, rows: usize, cols: usize, scale: f32, r: Rnd) ![]f32 { + const w = try gpa.alloc(f32, rows * cols); + for (w) |*v| { + const raw = l.val(scale); + v.* = switch (r) { + .plain => raw, + .f16 => f16r(raw), + .bf16 => b16r(raw), + }; + } + return w; +} + +/// A tiny deterministic token stream for the tests' prompts (never id 0 twice in a row: the +/// embedding's rows are all distinct, so any stream is a valid probe). +pub fn promptOf(l: *Lcg, slice: []u32) void { + for (slice) |*t| { + t.* = @intCast(l.next() % 250 + 3); + } +} + +/// The serving model state: weights, the fp16 K/V cache and the indexer key side plane, traces +/// of the last row's selection and expert picks, and the last row's logits. +pub const Serve = struct { + gpa: std.mem.Allocator, + cfg: Config, + len: u32 = 0, // content keys in the caches (all layers share the row count) + embed: []f32, // [V, hidden] bf16-rounded values + final_norm: []f32, // [hidden] + head: []f32, // [V, hidden] plain fp32 fixture values + layers: []Layer, // the backbone; layers[0..dense_layers] dense, the rest MoE + kcache: []f32, // [layers][capacity][kv_rows] fp16-rounded row values + vcache: []f32, + plane: []f32, // [layers][capacity][i_dim] the index-K side plane, fp16 rows + // the last processed row's traces (the device kernels' parity grade): + sel: []i32, // [layers][kv_heads][i_topk], the selection ids, -1 right-pad + picks: []bool, // [layers][E] + sel_e: []u32, // [topk] + weights: []f32, // [topk] + logits: []f32, // [V] the last row's + // one-row scratch: + x: []f32, + xn: []f32, + xn2: []f32, + q: []f32, // [qk_rows] the normed+roped q rows (fp16 values) + attn: []f32, // [qk_rows] the attention output rows + kbuf: []f32, + vbuf: []f32, + iq: []f32, + ik: []f32, + oout: []f32, + act: []f32, // [inter_max] + gbuf: []f32, + ubuf: []f32, + keys: []u32, // [i_topk * block] the gathered key slots + sscr: []f32, // [i_topk * block] one attention row's score/exp scratch + elogits: []f32, // [experts] the router's raw logits scratch + kpos: []i64, + + pub fn deinit(s: *Serve) void { + const gpa = s.gpa; + if (s.embed.len > 0) gpa.free(s.embed); + if (s.final_norm.len > 0) gpa.free(s.final_norm); + if (s.head.len > 0) gpa.free(s.head); + for (s.layers) |*l| l.deinit(gpa); + if (s.layers.len > 0) gpa.free(s.layers); + if (s.kcache.len > 0) gpa.free(s.kcache); + if (s.vcache.len > 0) gpa.free(s.vcache); + if (s.plane.len > 0) gpa.free(s.plane); + if (s.sel.len > 0) gpa.free(s.sel); + if (s.picks.len > 0) gpa.free(s.picks); + if (s.sel_e.len > 0) gpa.free(s.sel_e); + if (s.weights.len > 0) gpa.free(s.weights); + if (s.logits.len > 0) gpa.free(s.logits); + if (s.x.len > 0) gpa.free(s.x); + if (s.xn.len > 0) gpa.free(s.xn); + if (s.xn2.len > 0) gpa.free(s.xn2); + if (s.q.len > 0) gpa.free(s.q); + if (s.attn.len > 0) gpa.free(s.attn); + if (s.kbuf.len > 0) gpa.free(s.kbuf); + if (s.vbuf.len > 0) gpa.free(s.vbuf); + if (s.iq.len > 0) gpa.free(s.iq); + if (s.ik.len > 0) gpa.free(s.ik); + if (s.oout.len > 0) gpa.free(s.oout); + if (s.act.len > 0) gpa.free(s.act); + if (s.gbuf.len > 0) gpa.free(s.gbuf); + if (s.ubuf.len > 0) gpa.free(s.ubuf); + if (s.keys.len > 0) gpa.free(s.keys); + if (s.sscr.len > 0) gpa.free(s.sscr); + if (s.elogits.len > 0) gpa.free(s.elogits); + if (s.kpos.len > 0) gpa.free(s.kpos); + s.* = undefined; + } + + pub fn init(gpa: std.mem.Allocator, c: Config) !Serve { + const L = c.layers; + var rng = Lcg{ .s = seed_default }; + var s: Serve = .{ + .gpa = gpa, + .cfg = c, + .embed = &.{}, + .final_norm = &.{}, + .head = &.{}, + .layers = &.{}, + .kcache = &.{}, + .vcache = &.{}, + .plane = &.{}, + .sel = &.{}, + .picks = &.{}, + .sel_e = &.{}, + .weights = &.{}, + .logits = &.{}, + .x = &.{}, + .xn = &.{}, + .xn2 = &.{}, + .q = &.{}, + .attn = &.{}, + .kbuf = &.{}, + .vbuf = &.{}, + .iq = &.{}, + .ik = &.{}, + .oout = &.{}, + .act = &.{}, + .gbuf = &.{}, + .ubuf = &.{}, + .keys = &.{}, + .sscr = &.{}, + .elogits = &.{}, + .kpos = &.{}, + }; + errdefer s.deinit(); + s.embed = try mat(gpa, &rng, c.vocab, hidden, 0.05, .bf16); + s.final_norm = try mat(gpa, &rng, 1, hidden, 0.02, .bf16); + s.head = try mat(gpa, &rng, c.vocab, hidden, 0.10, .plain); // wide margins for a greedy bar + s.kcache = try gpa.alloc(f32, L * capacity * kv_rows); + s.vcache = try gpa.alloc(f32, L * capacity * kv_rows); + s.plane = try gpa.alloc(f32, L * capacity * c.i_dim); + s.sel = try gpa.alloc(i32, L * kv_heads * c.i_topk); + s.picks = try gpa.alloc(bool, L * c.experts); + s.sel_e = try gpa.alloc(u32, c.topk); + s.weights = try gpa.alloc(f32, c.topk); + s.logits = try gpa.alloc(f32, c.vocab); + s.x = try gpa.alloc(f32, hidden); + s.xn = try gpa.alloc(f32, hidden); + s.xn2 = try gpa.alloc(f32, hidden); + s.q = try gpa.alloc(f32, qk_rows); + s.attn = try gpa.alloc(f32, qk_rows); + s.kbuf = try gpa.alloc(f32, kv_rows); + s.vbuf = try gpa.alloc(f32, kv_rows); + s.iq = try gpa.alloc(f32, i_rows); + s.ik = try gpa.alloc(f32, c.i_dim); + s.oout = try gpa.alloc(f32, hidden); + s.act = try gpa.alloc(f32, inter_max); + s.gbuf = try gpa.alloc(f32, inter_max); + s.ubuf = try gpa.alloc(f32, inter_max); + s.keys = try gpa.alloc(u32, c.keyWidth()); + s.sscr = try gpa.alloc(f32, c.keyWidth()); // one attention row: up to keyWidth gathered keys + s.elogits = try gpa.alloc(f32, c.experts); + s.kpos = try gpa.alloc(i64, capacity); + var layers = try gpa.alloc(Layer, L); + errdefer gpa.free(layers); + for (0..L) |li| { + var l: Layer = .{ + .in_norm = try mat(gpa, &rng, 1, hidden, 0.02, .bf16), + .post_norm = try mat(gpa, &rng, 1, hidden, 0.02, .bf16), + .q_norm = try mat(gpa, &rng, 1, head_dim, 0.02, .bf16), + .k_norm = try mat(gpa, &rng, 1, head_dim, 0.02, .bf16), + .q = try mat(gpa, &rng, qk_rows, hidden, 0.006, .plain), + .k = try mat(gpa, &rng, kv_rows, hidden, 0.006, .plain), + .v = try mat(gpa, &rng, kv_rows, hidden, 0.006, .plain), + .o = try mat(gpa, &rng, hidden, qk_rows, 0.006, .plain), + .iq = undefined, + .ik = undefined, + .iqn = try mat(gpa, &rng, 1, c.i_dim, 0.02, .bf16), + .ikn = try mat(gpa, &rng, 1, c.i_dim, 0.02, .bf16), + .mlp = undefined, + }; + errdefer { + if (l.mlp_set) l.deinit(gpa); + } + l.iq = try mat(gpa, &rng, i_rows, hidden, 0.008, .f16); + l.ik = try mat(gpa, &rng, c.i_dim, hidden, 0.008, .f16); + if (li < c.dense_layers) { + l.mlp = .{ .dense = .{ + .gate = try mat(gpa, &rng, c.dense_inter, hidden, 0.006, .plain), + .up = try mat(gpa, &rng, c.dense_inter, hidden, 0.006, .plain), + .down = try mat(gpa, &rng, hidden, c.dense_inter, 0.006, .plain), + } }; + } else { + var moe = Layer.MoE{ + .router = try mat(gpa, &rng, c.experts, hidden, 0.02, .plain), + .bias = try mat(gpa, &rng, c.experts, 1, 0.002, .plain), + .shared = undefined, + .experts = undefined, + }; + errdefer gpa.free(moe.router); + errdefer gpa.free(moe.bias); + moe.shared = .{ + .gate = try mat(gpa, &rng, c.shared_inter, hidden, 0.006, .plain), + .up = try mat(gpa, &rng, c.shared_inter, hidden, 0.006, .plain), + .down = try mat(gpa, &rng, hidden, c.shared_inter, 0.006, .plain), + }; + var list = try gpa.alloc(Mlp, c.experts); + _ = &list; + errdefer { + for (list) |e| { + gpa.free(e.gate); + gpa.free(e.up); + gpa.free(e.down); + } + gpa.free(list); + } + for (list) |*e| e.* = .{ + .gate = try mat(gpa, &rng, c.moe_inter, hidden, 0.006, .plain), + .up = try mat(gpa, &rng, c.moe_inter, hidden, 0.006, .plain), + .down = try mat(gpa, &rng, hidden, c.moe_inter, 0.006, .plain), + }; + moe.experts = list; + l.mlp = .{ .moe = moe }; + } + l.mlp_set = true; + layers[li] = l; + } + s.layers = layers; + @memset(s.sel, -1); + // the k_pos mask row the indexer scores against: content keys 0..len-1, grown per row + for (0..capacity) |i| s.kpos[i] = @intCast(i); + return s; + } + + /// Append one row: forward token `tok` at position `len`, fill caches/plane/traces/logits, + /// then len += 1. Returns the row's logits. + pub fn put(s: *Serve, tok: u32) !void { + std.debug.assert(s.len < capacity); + const c = &s.cfg; + const pos: usize = s.len; + const nb = sparse.nb(pos + 1, c.block); + s.len += 1; + const sk: usize = pos + 1; // keys [0, pos] now hold content (the row scores itself too) + const scale: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(c.head_dim))); + @memcpy(s.x, s.embed[@as(usize, tok) * hidden ..][0..hidden]); + for (s.layers, 0..) |*l, li| { + // attention: norm -> per-head qk norm -> partial rope -> fp16 stores -> selection -> + // block-sparse fp32 softmax over the gathered selected rows -> o_proj -> residual + @memcpy(s.xn, s.x); + normRow(s.xn, l.in_norm, c.eps); + linear(s.xn, l.q, qk_rows, s.q); + var qb: usize = 0; + while (qb < qk_rows) : (qb += head_dim) { + normRow(s.q[qb..][0..head_dim], l.q_norm, c.eps); + ropeRow(s.q[qb..][0..head_dim], pos, c.rotary_dim); + } + for (0..qk_rows) |i| s.q[i] = f16r(s.q[i]); // the fork's fp16 linear outputs + linear(s.xn, l.k, kv_rows, s.kbuf); + var kb: usize = 0; + while (kb < kv_rows) : (kb += head_dim) { + normRow(s.kbuf[kb..][0..head_dim], l.k_norm, c.eps); + ropeRow(s.kbuf[kb..][0..head_dim], pos, c.rotary_dim); + } + linear(s.xn, l.v, kv_rows, s.vbuf); + const kro = @as(usize, li) * capacity * kv_rows; + for (0..kv_rows) |i| { + s.kcache[kro + pos * kv_rows + i] = f16r(s.kbuf[i]); + s.vcache[kro + pos * kv_rows + i] = f16r(s.vbuf[i]); + } + // the indexer's side path: norm+rope its queries, append the plane row + linear(s.xn, l.iq, i_rows, s.iq); + var ib: usize = 0; + while (ib < i_rows) : (ib += c.i_dim) { + normRow(s.iq[ib..][0..c.i_dim], l.iqn, c.eps); + ropeRow(s.iq[ib..][0..c.i_dim], pos, c.rotary_dim); + } + for (0..i_rows) |i| s.iq[i] = f16r(s.iq[i]); + linear(s.xn, l.ik, c.i_dim, s.ik); + normRow(s.ik, l.ikn, c.eps); + ropeRow(s.ik, pos, c.rotary_dim); + const pro = @as(usize, li) * capacity * c.i_dim; + for (0..c.i_dim) |i| s.plane[pro + pos * c.i_dim + i] = f16r(s.ik[i]); + // the selection: sparse.zig's contracts over the fp16 rows (fp32 scoring inside) + const sq = try sparse.scores(s.gpa, s.iq, s.plane[pro..][0 .. sk * c.i_dim], &[1]i64{@intCast(pos)}, s.kpos[0..sk], c.i_heads, 1, sk, c.i_dim, c.block); + defer s.gpa.free(sq); + const pk = try sparse.topk(s.gpa, sq, c.i_heads, 1, nb, c.i_topk, c.block, 1, 0, &[1]i64{@intCast(pos)}); + defer s.gpa.free(pk); + @memcpy(s.sel[@as(usize, li) * kv_heads * c.i_topk ..][0 .. kv_heads * c.i_topk], pk); + // the block-sparse attention: each kv head's selected blocks gather its key slots + var hv: usize = 0; + while (hv < kv_heads) : (hv += 1) { + var nk: usize = 0; + for (0..c.i_topk) |sl| { + const bl = s.sel[@as(usize, li) * kv_heads * c.i_topk + hv * c.i_topk + sl]; + if (bl < 0) break; // -1 right-pad + const blo: usize = @as(usize, @intCast(bl)) * c.block; + const nhi = @min(blo + c.block, sk); + var kk = blo; + while (kk < nhi) : (kk += 1) { + s.keys[nk] = @intCast(kk); + nk += 1; + } + } + const per = q_heads / kv_heads; + var j: usize = 0; + while (j < per) : (j += 1) { + const qh = hv * per + j; + const qr = s.q[qh * head_dim ..][0..head_dim]; + var mx = -std.math.inf(f32); + const kr = s.kcache[kro..][0 .. sk * kv_rows]; + var m: usize = 0; + while (m < nk) : (m += 1) { + const krow = kr[s.keys[m] * kv_rows + hv * head_dim ..][0..head_dim]; + var dot: f32 = 0; + for (0..head_dim) |d| dot += qr[d] * krow[d]; + s.sscr[m] = dot * scale; + if (dot * scale > mx) mx = dot * scale; + } + var se: f32 = 0; + m = 0; + while (m < nk) : (m += 1) { + const e = @exp(s.sscr[m] - mx); + s.sscr[m] = e; + se += e; + } + var acc: [head_dim]f32 = @splat(0); + m = 0; + while (m < nk) : (m += 1) { + const krow = kr[s.keys[m] * kv_rows + hv * head_dim ..][0..head_dim]; + const p = s.sscr[m] / se; + for (0..head_dim) |d| acc[d] += p * krow[d]; + } + const outr = s.attn[qh * head_dim ..][0..head_dim]; + for (0..head_dim) |d| outr[d] = f16r(acc[d]); + } + } + for (0..qk_rows) |i| s.attn[i] = f16r(s.attn[i]); + linear(s.attn, l.o, hidden, s.oout); + for (0..hidden) |i| { + s.oout[i] = f16r(s.oout[i]); + s.x[i] += s.oout[i]; + } + // the MLP wave: norm over the residual, then dense or routed+shared, then residual + @memcpy(s.xn2, s.x); + normRow(s.xn2, l.post_norm, c.eps); + switch (l.mlp) { + .dense => |m| { + linear(s.xn2, m.gate, c.dense_inter, s.gbuf[0..c.dense_inter]); + linear(s.xn2, m.up, c.dense_inter, s.ubuf[0..c.dense_inter]); + for (0..c.dense_inter) |i| s.act[i] = swigluoai(s.gbuf[i], s.ubuf[i], c.swiglu_alpha, c.swiglu_limit); + linear(s.act[0..c.dense_inter], m.down, hidden, s.oout); + }, + .moe => |m| { + const E: usize = c.experts; + const picks_off = @as(usize, li) * E; + route(s.xn2, m.router, m.bias, s.picks[picks_off..][0..E], s.elogits, s.sel_e, s.weights, c); + @memset(s.oout, 0); + for (0..c.topk) |slot| { + if (!s.picks[picks_off + s.sel_e[slot]]) continue; + const e = m.experts[s.sel_e[slot]]; + linear(s.xn2, e.gate, c.moe_inter, s.gbuf[0..c.moe_inter]); + linear(s.xn2, e.up, c.moe_inter, s.ubuf[0..c.moe_inter]); + for (0..c.moe_inter) |i| s.act[i] = swigluoai(s.gbuf[i], s.ubuf[i], c.swiglu_alpha, c.swiglu_limit); + linear(s.act[0..c.moe_inter], e.down, hidden, s.act[inter_max - hidden ..]); + for (0..hidden) |i| s.oout[i] += s.weights[slot] * s.act[inter_max - hidden + i]; + } + linear(s.xn2, m.shared.gate, c.shared_inter, s.gbuf[0..c.shared_inter]); + linear(s.xn2, m.shared.up, c.shared_inter, s.ubuf[0..c.shared_inter]); + for (0..c.shared_inter) |i| s.act[i] = swigluoai(s.gbuf[i], s.ubuf[i], c.swiglu_alpha, c.swiglu_limit); + linear(s.act[0..c.shared_inter], m.shared.down, hidden, s.act[inter_max - hidden ..]); + for (0..hidden) |i| s.oout[i] += s.act[inter_max - hidden + i]; // shared at weight 1 + }, + } + for (0..hidden) |i| { + s.x[i] += f16r(s.oout[i]); // the fork rounds each linear's output before the residual + } + } + normRow(s.x, s.final_norm, c.eps); + for (0..c.vocab) |v| { + var acc: f32 = 0; + for (0..hidden) |j| acc += s.x[j] * s.head[v * hidden + j]; + s.logits[v] = acc; + } + } + + /// Greedy continuation: `steps` argmax picks from the last row, each picked token fed back; + /// greedy only (first-maximum argmax of the fp32 logits). + pub fn greedy(s: *Serve, steps: usize, out: []u32) !void { + for (0..steps) |g| { + var best: usize = 0; + var bv = -std.math.inf(f32); + for (s.logits, 0..) |v, i| if (v > bv) { + bv = v; + best = i; + }; + out[g] = @intCast(best); + if (g + 1 < steps) try s.put(@intCast(best)); + } + } +}; + +// --- the receipt tests --- + +test "a fixture pass: f64 reference parity, greedy equality, selection and pick traces" { + const gpa = std.testing.allocator; + const c = serveConfig(); + var s = try Serve.init(gpa, c); + defer s.deinit(); + var rng = Lcg{ .s = seed_default ^ 0xA5A5 }; + var prompt: [prefill_len]u32 = undefined; + promptOf(&rng, &prompt); + for (prompt) |t| try s.put(t); + var got: [decode_steps]u32 = undefined; + try s.greedy(decode_steps, &got); + + // the whole graph again in f64, straight-line independent reimplementation (the fork's torch + // composition), graded here: greedy token equality, logits tolerance, traces. + var r = try Ref.init(gpa, c, &s); + defer r.deinit(); + for (prompt) |t| try r.put(t); + var want: [decode_steps]u32 = undefined; + try r.greedy(decode_steps, &want); + try std.testing.expectEqualSlices(u32, &want, &got); + + // logits tolerance at the last row: fp32 host accumulation vs f64 reference. The synthetic + // weights keep every logit's top-2 gap above 0.05 (asserted below), so token equality is not + // a knife-edge fluke of the tolerance. + var maxd: f64 = 0; + for (s.logits, 0..) |v, i| { + const d = @as(f64, v) - r.logits[i]; + if (@abs(d) > maxd) maxd = @abs(d); + } + var top2: f32 = -std.math.inf(f32); + var margin: f32 = -std.math.inf(f32); + for (s.logits) |v| { + if (v > top2) { + margin = top2; + top2 = v; + } else if (v > margin) margin = v; + } + try std.testing.expect(top2 - margin > 0.05); + try expectTol(maxd, 5e-3); // fp32-vs-f64 drift over the ~2k-term logits sums, measured 0.0023 on + // the fixture (seed_default): the accumulation-order delta of two correct implementations, not a + // structure bug (greedy tokens, selection traces and routed picks all match exactly above). + + // the selection trace: layer 1's (the MoE wave's) valid block ids are descending with -1 pads + const sels = s.sel[c.kv_heads * c.i_topk ..][0 .. kv_heads * c.i_topk]; + for (0..c.kv_heads) |h| { + var seen: usize = 0; + for (0..c.i_topk) |sl| { + const id = sels[h * c.i_topk + sl]; + if (id < 0) break; + try std.testing.expect(id >= 0 and id < c.keyWidth() / c.block); + seen += 1; + } + try std.testing.expect(seen > 0); + } + // the picks: exactly topk experts fired, weight sum near routed scale (normalized * 2.0) + var fires: usize = 0; + var wsum: f32 = 0; + for (c.experts..c.experts * 2) |e| { + if (s.picks[e]) fires += 1; + } + for (s.weights) |w| wsum += w; + try std.testing.expectEqual(c.topk, fires); + try expectTol(wsum - c.routed_scale, 1e-2); +} + +fn expectTol(d: f64, t: f64) !void { + if (!(d <= t)) { + std.debug.print("tolerance blown: {d} > {d}\n", .{ d, t }); + return error.Tolerance; + } +} + +/// The independent f64 straight-line reference forward (the fork's torch composition) — the same +/// fixture weights through the engine instance it grades (`m`), but its OWN f64 caches and state, +/// its OWN selection and routing math, no shared code with Serve. Same rounding contract: one fp16 +/// rounding at each fp16 store (the rope outputs, the cache and plane rows, the routing weights) +/// and f64 everywhere else, so the engine-vs-reference difference is pure fp32-vs-f64 precision +/// over identical structure — the tolerance the test asserts isolates exactly that. +const Ref = struct { + const G = f64; + const ref_max_blocks: usize = 512; + gpa: std.mem.Allocator, + m: *const Serve, // the weights (already loaded, rounded by the fixture) + cfg: Config, + len: u32 = 0, + kc: []G, // [layers][capacity][kv_rows] + vc: []G, + pl: []G, // [layers][capacity][i_dim] + st: [hidden]G = @splat(0), + logits: []G, // [vocab] + qb: []G, // [qk_rows] + ab: []G, // [qk_rows] + kb: []G, // [kv_rows] + vb: []G, + ib: []G, // [i_rows] + isk: []G, // [i_dim] + ob: []G, // [hidden] + gb: []G, // [inter_max] + ubx: []G, + acb: []G, + ks: []u32, // [keyWidth] + ssc: []G, // [keyWidth] + kpo: []i64, // [capacity] + + fn h16(x: G) G { + return @floatCast(@as(f16, @floatCast(x))); + } + fn norm(x: []G, w: []const f32, eps: f32) void { + var ss: G = 0; + for (x) |v| ss += v * v; + const inv = 1.0 / @sqrt(ss / @as(G, @floatFromInt(x.len)) + @as(G, eps)); + for (x, w) |*v, wi| v.* = v.* * inv * (1.0 + @as(G, wi)); + } + fn lin(x: []const G, w: []const f32, n_out: usize, y: []G) void { + const n_in = x.len; + @memset(y, 0); + for (0..n_in) |c2| { + for (0..n_out) |o| { + y[o] += x[c2] * @as(G, w[o * n_in + c2]); + } + } + } + fn rope(x: []G, pos: usize, rot: usize) void { + const half = rot / 2; + for (0..half) |i| { + const ang = @as(G, @floatFromInt(pos)) / + std.math.pow(G, @as(G, 1e7), @as(G, @floatFromInt(2 * i)) / @as(G, @floatFromInt(rot))); + const c2: G = @cos(ang); + const s: G = @sin(ang); + const a = x[i]; + const b = x[i + half]; + x[i] = a * c2 - b * s; + x[i + half] = b * c2 + a * s; + } + } + fn act(g: G, u: G, alpha: f32, limit: f32) G { + const gg = @min(g, @as(G, limit)); + const uu = @min(@max(u, -@as(G, limit)), @as(G, limit)) + 1.0; + return gg / (1.0 + std.math.exp(-@as(G, alpha) * gg)) * uu; + } + fn deinit(r: *Ref) void { + const gpa = r.gpa; + if (r.kc.len > 0) gpa.free(r.kc); + if (r.vc.len > 0) gpa.free(r.vc); + if (r.pl.len > 0) gpa.free(r.pl); + if (r.logits.len > 0) gpa.free(r.logits); + if (r.qb.len > 0) gpa.free(r.qb); + if (r.ab.len > 0) gpa.free(r.ab); + if (r.kb.len > 0) gpa.free(r.kb); + if (r.vb.len > 0) gpa.free(r.vb); + if (r.ib.len > 0) gpa.free(r.ib); + if (r.isk.len > 0) gpa.free(r.isk); + if (r.ob.len > 0) gpa.free(r.ob); + if (r.gb.len > 0) gpa.free(r.gb); + if (r.ubx.len > 0) gpa.free(r.ubx); + if (r.acb.len > 0) gpa.free(r.acb); + if (r.ks.len > 0) gpa.free(r.ks); + if (r.ssc.len > 0) gpa.free(r.ssc); + if (r.kpo.len > 0) gpa.free(r.kpo); + r.* = undefined; + } + + pub fn init(gpa: std.mem.Allocator, c: Config, s: *const Serve) !Ref { + const L: usize = c.layers; + var r: Ref = .{ + .gpa = gpa, + .m = s, + .cfg = c, + .kc = &.{}, + .vc = &.{}, + .pl = &.{}, + .logits = &.{}, + .qb = &.{}, + .ab = &.{}, + .kb = &.{}, + .vb = &.{}, + .ib = &.{}, + .isk = &.{}, + .ob = &.{}, + .gb = &.{}, + .ubx = &.{}, + .acb = &.{}, + .ks = &.{}, + .ssc = &.{}, + .kpo = &.{}, + }; + errdefer r.deinit(); + r.kc = try gpa.alloc(G, L * capacity * kv_rows); + r.vc = try gpa.alloc(G, L * capacity * kv_rows); + r.pl = try gpa.alloc(G, L * capacity * c.i_dim); + r.logits = try gpa.alloc(G, c.vocab); + r.qb = try gpa.alloc(G, qk_rows); + r.ab = try gpa.alloc(G, qk_rows); + r.kb = try gpa.alloc(G, kv_rows); + r.vb = try gpa.alloc(G, kv_rows); + r.ib = try gpa.alloc(G, i_rows); + r.isk = try gpa.alloc(G, c.i_dim); + r.ob = try gpa.alloc(G, hidden); + r.gb = try gpa.alloc(G, inter_max); + r.ubx = try gpa.alloc(G, inter_max); + r.acb = try gpa.alloc(G, inter_max); + r.ks = try gpa.alloc(u32, c.keyWidth()); + r.ssc = try gpa.alloc(G, c.keyWidth()); + r.kpo = try gpa.alloc(i64, capacity); + for (0..capacity) |i| r.kpo[i] = @intCast(i); + return r; + } + + /// One row at `len`, straight-line: the fork's torch composition in f64. + pub fn put(r: *Ref, tok: u32) !void { + std.debug.assert(r.len < capacity); + const c = &r.cfg; + const pos: usize = r.len; + r.len += 1; + const sk: usize = pos + 1; + const scale: f64 = 1.0 / @sqrt(@as(G, @floatFromInt(c.head_dim))); + for (0..hidden) |i| r.st[i] = @as(G, r.m.embed[@as(usize, tok) * hidden + i]); + for (r.m.layers, 0..) |*l, li| { + var xn: [hidden]G = undefined; + @memcpy(&xn, &r.st); + norm(&xn, l.in_norm, c.eps); + // the packed q rows: per-head gemma norm, partial rope, one fp16 rounding + lin(&xn, l.q, qk_rows, r.qb); + for (0..q_heads) |h| { + const row = r.qb[h * head_dim ..][0..head_dim]; + norm(row, l.q_norm, c.eps); + rope(row, pos, c.rotary_dim); + for (row) |*v| v.* = h16(v.*); + } + lin(&xn, l.k, kv_rows, r.kb); + for (0..kv_heads) |h| { + const row = r.kb[h * head_dim ..][0..head_dim]; + norm(row, l.k_norm, c.eps); + rope(row, pos, c.rotary_dim); + for (row) |*v| v.* = h16(v.*); + } + lin(&xn, l.v, kv_rows, r.vb); + for (0..kv_rows) |i| r.vb[i] = h16(r.vb[i]); + const kro = @as(usize, li) * capacity * kv_rows; + for (0..kv_rows) |i| { + r.kc[kro + pos * kv_rows + i] = r.kb[i]; + r.vc[kro + pos * kv_rows + i] = r.vb[i]; + } + // the indexer's side path (its own selection math follows) + lin(&xn, l.iq, i_rows, r.ib); + for (0..c.i_heads) |h| { + const row = r.ib[h * c.i_dim ..][0..c.i_dim]; + norm(row, l.iqn, c.eps); + rope(row, pos, c.rotary_dim); + for (row) |*v| v.* = h16(v.*); + } + lin(&xn, l.ik, c.i_dim, r.isk); + norm(r.isk, l.ikn, c.eps); + rope(r.isk, pos, c.rotary_dim); + for (0..c.i_dim) |i| r.isk[i] = h16(r.isk[i]); + const pro = @as(usize, li) * capacity * c.i_dim; + for (0..c.i_dim) |i| r.pl[pro + pos * c.i_dim + i] = r.isk[i]; + // the selection, computed independently: f64 block-max pool, boost override, topk + const nb = (sk + c.block - 1) / c.block; + std.debug.assert(nb <= ref_max_blocks); + std.debug.assert(c.i_topk <= 16); + var sel: [16]i32 = @splat(-1); + var h: usize = 0; + while (h < c.i_heads) : (h += 1) { + var bsc: [ref_max_blocks]f64 = @splat(-std.math.inf(G)); + for (0..nb) |b| { + var best = -std.math.inf(G); + const lo = b * c.block; + const nhi = @min(lo + c.block, sk); + var j = lo; + while (j < nhi) : (j += 1) { + var dot: G = 0; + for (0..c.i_dim) |d| dot += r.ib[h * c.i_dim + d] * r.pl[pro + j * c.i_dim + d]; + if (dot > best) best = dot; + } + bsc[b] = best; + } + bsc[pos / c.block] = std.math.inf(G); // the local boost OVERRIDE + for (0..c.i_topk) |sl| { + var best: usize = nb; + var bs = -std.math.inf(G); + for (0..nb) |b| { + var taken = false; + for (0..sl) |p| { + if (sel[p] >= 0 and sel[p] == @as(i32, @intCast(b))) taken = true; + } + if (taken) continue; + if (bsc[b] > bs) { + bs = bsc[b]; + best = b; + } + } + if (best < nb) sel[sl] = @intCast(best); // else stays -1 (the right pad) + } + } + // the block-sparse fp64 softmax over the gathered selected rows, GQA by kv head + const per = q_heads / kv_heads; + var hv: usize = 0; + while (hv < kv_heads) : (hv += 1) { + var nk: usize = 0; + for (0..c.i_topk) |sl| { + const bl = sel[sl]; + if (bl < 0) break; + const blo = @as(usize, @intCast(bl)) * c.block; + const nhi2 = @min(blo + c.block, sk); + var kk = blo; + while (kk < nhi2) : (kk += 1) { + r.ks[nk] = @intCast(kk); + nk += 1; + } + } + var j: usize = 0; + while (j < per) : (j += 1) { + const qh = hv * per + j; + const qr = r.qb[qh * head_dim ..][0..head_dim]; + var mx = -std.math.inf(G); + var m: usize = 0; + while (m < nk) : (m += 1) { + var dot: G = 0; + const krow = r.kc[kro + r.ks[m] * kv_rows ..][0..kv_rows][hv * head_dim ..][0..head_dim]; + for (0..head_dim) |d| dot += qr[d] * krow[d]; + r.ssc[m] = dot * scale; + if (dot * scale > mx) mx = dot * scale; + } + var se: G = 0; + m = 0; + while (m < nk) : (m += 1) { + const e = std.math.exp(r.ssc[m] - mx); + r.ssc[m] = e; + se += e; + } + var acc: [head_dim]G = @splat(0); + m = 0; + while (m < nk) : (m += 1) { + const vrow = r.vc[kro + r.ks[m] * kv_rows ..][0..kv_rows][hv * head_dim ..][0..head_dim]; + const p = r.ssc[m] / se; + for (0..head_dim) |d| acc[d] += p * vrow[d]; + } + const orow = r.ab[qh * head_dim ..][0..head_dim]; + for (0..head_dim) |d| orow[d] = h16(acc[d]); + } + } + for (0..qk_rows) |i| r.ab[i] = h16(r.ab[i]); + lin(r.ab, l.o, hidden, r.ob); + for (0..hidden) |i| r.st[i] += h16(r.ob[i]); + var xn2: [hidden]G = undefined; + @memcpy(&xn2, &r.st); + norm(&xn2, l.post_norm, c.eps); + switch (l.mlp) { + .dense => |m| { + lin(&xn2, m.gate, c.dense_inter, r.gb); + lin(&xn2, m.up, c.dense_inter, r.ubx); + for (0..c.dense_inter) |i| r.acb[i] = act(r.gb[i], r.ubx[i], c.swiglu_alpha, c.swiglu_limit); + lin(r.acb[0..c.dense_inter], m.down, hidden, r.ob); + }, + .moe => |m| { + // the router: scores = sigmoid(logits); select on sigmoid + bias; weights = + // unbiased sigmoid normalized over the picks, * scaling, each one .half() + var sg: [16]G = @splat(0); // the unbiased sigmoids (E <= 16 in the fixture) + var sc: [16]G = @splat(0); + var picks: [8]bool = @splat(false); + var rtw: [8]G = @splat(0); // the per-slot unbiased sigmoid, before normalization + for (0..c.experts) |e| { + var lg: G = 0; + for (0..hidden) |jj| lg += xn2[jj] * @as(G, m.router[e * hidden + jj]); + sg[e] = 1.0 / (1.0 + std.math.exp(-lg)); + sc[e] = sg[e] + @as(G, m.bias[e]); + } + var sum: G = 1e-20; + for (0..c.topk) |slot| { + var best: usize = c.experts; + var bs = -std.math.inf(G); + for (0..c.experts) |e| { + if (picks[e]) continue; + if (sc[e] > bs) { + bs = sc[e]; + best = e; + } + } + picks[best] = true; + sum += sg[best]; + rtw[slot] = sg[best]; + } + var acc: [hidden]G = @splat(0); + for (0..c.topk) |slot| { + const w: G = h16(rtw[slot] / sum * @as(G, c.routed_scale)); + const e = m.experts[picksSlot(picks, slot, c.experts)]; + lin(&xn2, e.gate, c.moe_inter, r.gb); + lin(&xn2, e.up, c.moe_inter, r.ubx); + for (0..c.moe_inter) |i| r.acb[i] = act(r.gb[i], r.ubx[i], c.swiglu_alpha, c.swiglu_limit); + lin(r.acb[0..c.moe_inter], e.down, hidden, r.ob); + for (0..hidden) |i| acc[i] += w * r.ob[i]; + } + lin(&xn2, m.shared.gate, c.shared_inter, r.gb); + lin(&xn2, m.shared.up, c.shared_inter, r.ubx); + for (0..c.shared_inter) |i| r.acb[i] = act(r.gb[i], r.ubx[i], c.swiglu_alpha, c.swiglu_limit); + lin(r.acb[0..c.shared_inter], m.shared.down, hidden, r.ob); + for (0..hidden) |i| acc[i] += r.ob[i]; + for (0..hidden) |i| r.ob[i] = acc[i]; // the combined wave output (shared at weight 1) + }, + } + // every wave's output: the fork rounds each linear's output before the residual + for (0..hidden) |i| r.st[i] += h16(r.ob[i]); + } + var xf: [hidden]G = undefined; + @memcpy(&xf, &r.st); + norm(&xf, r.m.final_norm, c.eps); + for (0..c.vocab) |v| { + var acc: G = 0; + for (0..hidden) |j| acc += xf[j] * @as(G, r.m.head[v * hidden + j]); + r.logits[v] = acc; + } + } + + /// Greedy continuation, argmax first-max, feeding each pick back (identical loop to the engine's). + pub fn greedy(r: *Ref, steps: usize, out: []u32) !void { + for (0..steps) |g| { + var best: usize = 0; + var bv = -std.math.inf(G); + for (r.logits, 0..) |v, i| if (v > bv) { + bv = v; + best = i; + }; + out[g] = @intCast(best); + if (g + 1 < steps) try r.put(@intCast(best)); + } + } +}; + +/// The slot-th picked expert of a bool pick mask (selection order = ascending id, matching the +/// strict-> scan's tie order). +fn picksSlot(picks: [8]bool, slot: usize, E: usize) usize { + var seen: usize = 0; + for (0..E) |e| { + if (picks[e]) { + if (seen == slot) return e; + seen += 1; + } + } + unreachable; +} \ No newline at end of file diff --git a/zig/src/families/minimax_m3_vl/shard.zig b/zig/src/families/minimax_m3_vl/shard.zig new file mode 100644 index 0000000000..74e2312cc2 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/shard.zig @@ -0,0 +1,116 @@ +//! The family's TP2 rank placement, TF's documented rank pattern (docs/multi-rank-receipts.md: the +//! GLM/Nemotron precedent — this family's TF_MINIMAXM3_RANK/TF_MINIMAXM3_WORLD, NCCL_COMM_ID on rank +//! 0's address) laid over the load plan: each rank runs one contiguous range of the checkpoint's +//! layers, rank 0 additionally runs the past-layer tensors (the embed, the final norm, the head) and +//! reads the logits. This stage receipts the per-rank bytes (the real pack's upload receipts are the +//! minimax-m3-tp2 command); the serving attach — each rank's launch graph over its range, NCCL +//! between the waves — rides this same split. + +const std = @import("std"); +const weights = @import("weights.zig"); + +/// This stage's placement: two GB10 nodes (the single-GB10 pack cannot fit one node). A wider world +/// needs its own receipts before `max_world` moves. +pub const max_world: u32 = 2; + +/// This rank's place: TF_MINIMAXM3_RANK of TF_MINIMAXM3_WORLD — both, or neither (never one alone). +pub const Place = struct { rank: u32, world: u32 }; + +pub fn place() !Place { + const r = std.c.getenv("TF_MINIMAXM3_RANK") orelse return error.NoRank; + const w = std.c.getenv("TF_MINIMAXM3_WORLD") orelse return error.NoRank; + const rank = std.fmt.parseInt(u32, std.mem.span(r), 10) catch return error.NoRank; + const world = std.fmt.parseInt(u32, std.mem.span(w), 10) catch return error.NoRank; + return checked(.{ .rank = rank, .world = world }); +} + +/// A place the family serves at this stage: rank inside a world of 1 or `max_world`. +pub fn checked(p: Place) !Place { + if (p.world < 1 or p.world > max_world or p.rank >= p.world) return error.NoRank; + return p; +} + +/// The layers rank `rank` of `world` runs: the head rank's first, the worker's after — the fork's +/// M3.1 precedent (head layers, then worker layers) laid on the plan: at 60 layers, rank 0 runs +/// layers 0..29, rank 1 runs 30..59. +pub const Range = struct { first: u32, count: u32 }; + +pub fn range(layers: u32, world: u32, rank: u32) Range { + const base = layers / world; + const rem = layers % world; + return .{ .first = rank * base + @min(rank, rem), .count = base + @intFromBool(rank < rem) }; +} + +fn has(r: Range, layer: u32) bool { + return layer >= r.first and layer < r.first + r.count; +} + +/// Does this load item belong to rank `rank`: its layer sits in the rank's range, or (past-layer +/// tensors: the embed, the final norm, the head) on rank 0 — the rank that embeds and reads logits. +pub fn owns(it: *const weights.Item, layers: u32, world: u32, rank: u32) bool { + if (world == 1) return true; + if (it.layer == weights.past) return rank == 0; + return has(range(layers, world, rank), it.layer); +} + +/// The rank's checkpoint bytes over the whole load: the receipt the TP2 upload asserts against +/// (rank 0's + rank 1's must equal `load.bytes()` exactly — no replication in this split). +pub fn bytesOf(l: *const weights.Load, world: u32, rank: u32) u64 { + var n: u64 = 0; + for (l.items) |*it| { + if (owns(it, l.config.layers, world, rank)) n += it.bytes; + } + return n; +} + +/// The rank's item count over the whole load. +pub fn itemsOf(l: *const weights.Load, world: u32, rank: u32) usize { + var n: usize = 0; + for (l.items) |*it| { + if (owns(it, l.config.layers, world, rank)) n += 1; + } + return n; +} + +// --- tests: the range math first, then the fixture plan's split by bytes --- + +test "the layer split's ranges: even halves at 2, the remainder carried by the first ranks" { + try std.testing.expectEqual(Range{ .first = 0, .count = 30 }, range(60, 2, 0)); + try std.testing.expectEqual(Range{ .first = 30, .count = 30 }, range(60, 2, 1)); + try std.testing.expectEqual(Range{ .first = 0, .count = 20 }, range(60, 3, 0)); + try std.testing.expectEqual(Range{ .first = 20, .count = 20 }, range(60, 3, 1)); + try std.testing.expectEqual(Range{ .first = 40, .count = 20 }, range(60, 3, 2)); + try std.testing.expectEqual(Range{ .first = 0, .count = 1 }, range(2, 3, 0)); + try std.testing.expectEqual(Range{ .first = 1, .count = 1 }, range(2, 3, 1)); + try std.testing.expectEqual(Range{ .first = 2, .count = 0 }, range(2, 3, 2)); // a short rank owns nothing this receipt + try std.testing.expectError(error.NoRank, checked(.{ .rank = 2, .world = 2 })); + try std.testing.expectError(error.NoRank, checked(.{ .rank = 0, .world = 3 })); +} + +test "the fixture load splits exactly: the two ranks' bytes and items sum to the whole plan" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try weights.checkpointFixture(gpa, io, tmp.dir, false); + const path = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(path); + const c = @import("exl3.zig").synthConfig(); + var p = try @import("exl3.zig").discover(gpa, io, path, &c); + defer p.deinit(); + var l = try weights.plan(gpa, io, path, &c, &p); + defer l.deinit(); + // world 2 over the 2-layer fixture: rank 0 = layer 0 plus every past-layer tensor. + const r0 = bytesOf(&l, 2, 0); + const r1 = bytesOf(&l, 2, 1); + try std.testing.expectEqual(l.bytes(), r0 + r1); // disjoint: no group on both ranks, none on neither + try std.testing.expectEqual(itemsOf(&l, 2, 0) + itemsOf(&l, 2, 1), l.items.len); + // the head is the plan's last item and it is rank 0's; every past-layer tensor is rank 0's + try std.testing.expect(owns(&l.items[l.head], c.layers, 2, 0)); + try std.testing.expect(!owns(&l.items[l.head], c.layers, 2, 1)); + try std.testing.expect(owns(&l.items[0], c.layers, 2, 0)); // the embed goes with rank 0's first wave + try std.testing.expect(!owns(&l.items[0], c.layers, 2, 1)); + // world 1 (TP1): everything on rank 0 + try std.testing.expectEqual(l.bytes(), bytesOf(&l, 1, 0)); + try std.testing.expectEqual(l.bytes(), bytesOf(&l, 1, 1)); +} \ No newline at end of file diff --git a/zig/src/families/minimax_m3_vl/sparse.zig b/zig/src/families/minimax_m3_vl/sparse.zig new file mode 100644 index 0000000000..252a382072 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/sparse.zig @@ -0,0 +1,163 @@ +//! The MSA indexer's selection contracts as a CPU reference (the CUDA kernels' grade): block-max-pool +//! scores over the index heads, the local/init boost override, and the fixed-width top-k selection. +//! Semantics exactly the fork's receipts (INTERFACE-M3-KERNELS.md §1): scoring fp32, keys masked -inf +//! iff k_pos < 0 or k_pos > q_pos (right-open causal on content positions), blocks anchored to key slots. + +const std = @import("std"); + +/// Whole key blocks the score table carries: the slot anchor's ceil(Sk / block). +pub fn nb(sk: usize, block: usize) usize { + return (sk + block - 1) / block; +} + +/// (B, H_idx, Sq, dim) index queries' fp32 block scores: -> [h, sq, NB], amax over each block's keys. +/// `dim` is the indexer head's width (the pack: 128); fp32 accumulation in index order, dot by dot. +pub fn scores( + gpa: std.mem.Allocator, + q: []const f32, // [h, sq, dim], row-major + k: []const f32, // [sk, dim] + q_pos: []const i64, // [sq] the queries' content positions + k_pos: []const i64, // [sk]; -1 = padded or post causal (k_pos > q_pos) + h: usize, + sq: usize, + sk: usize, + dim: usize, + block: usize, +) ![]f32 { + const nbv = nb(sk, block); + const out = try gpa.alloc(f32, h * sq * nbv); + errdefer gpa.free(out); + const neg_inf = -std.math.inf(f32); + for (0..h) |hi| for (0..sq) |ri| { + const qp = q_pos[ri]; + for (0..nbv) |b| { + var best = neg_inf; + const lo = b * block; + const hi_end = @min(sk, lo + block); + for (lo..hi_end) |ki| { + if (k_pos[ki] < 0 or k_pos[ki] > qp) continue; // the mask: -inf, not skipped-from-pooling? no: masked keys sit at -inf + var dot: f32 = 0; + for (0..dim) |di| { + dot += q[(hi * sq + ri) * dim + di] * k[ki * dim + di]; + } + if (dot > best) best = dot; // score type "max": the block's amax + } + // the last block's tail keys beyond sk count as masked (their mask word stays -inf) + out[(hi * sq + ri) * nbv + b] = best; + } + }; + return out; +} + +/// The fixed-width topk over block scores: boosts OVERRIDE pooled scores (HF scatter(+inf) semantics), +/// valid ids in descending-score order, ties at the smaller block id first, -1 right-pad, left-packed. +/// -> [h, sq, width] int32, always full width. The boosts: the row's own block and the `local - 1` +/// before it (max(0, q_block - l)), then the leading [0, init) blocks when init > 0. +pub fn topk( + gpa: std.mem.Allocator, + s: []const f32, // [h, sq, NB], scores()'s table + h: usize, + sq: usize, + nbv: usize, + width: usize, + block: usize, + local: usize, + init: usize, + q_pos: []const i64, +) ![]i32 { + const out = try gpa.alloc(i32, h * sq * width); + errdefer gpa.free(out); + const neg_inf = -std.math.inf(f32); + const boosted = try gpa.alloc(f32, nbv); + defer gpa.free(boosted); + const used = try gpa.alloc(bool, nbv); + defer gpa.free(used); + for (0..h) |hi| for (0..sq) |ri| { + const row = (hi * sq + ri); + const qp = q_pos[ri]; + const q_block = @divTrunc(qp, @as(i64, @intCast(block))); + @memcpy(boosted, s[row * nbv ..][0..nbv]); + @memset(used, false); + // the overrides land after the full reduction: +inf beats every pooled score + for (0..local) |l| { + const b = q_block - @as(i64, @intCast(l)); + if (b >= 0 and b < nbv) boosted[@intCast(b)] = std.math.inf(f32); + } + for (0..init) |bi| boosted[bi] = std.math.inf(f32); + for (0..width) |slot| { + var best: usize = nbv; + var best_score = neg_inf; + for (0..nbv) |b| { + if (used[b]) continue; + if (boosted[b] > best_score) { // ties stay at the smaller id: strict > + best_score = boosted[b]; + best = b; + } + } + if (best == nbv or best_score == neg_inf) { // the row ran out of valid blocks: -1 right-pad + out[row * width + slot] = -1; + continue; + } + used[best] = true; + out[row * width + slot] = @intCast(best); // valid ids left-packed, descending score + } + }; + return out; +} + +// --- the CPU reference's own receipt tests --- + +test "block scores pool by amax and mask the causal keys at -inf" { + const gpa = std.testing.allocator; + const dim = 2; + // one head, four query rows = two keys per block... use block 2, sk 4: two blocks of 2 keys + const q = [_]f32{ 1, 0, 0, 1, 2, 0, 0, 2 }; // rows r0..r3 + const k = [_]f32{ 1, 0, 0, 1, 3, 0, 0, 3 }; + const q_pos = [_]i64{ 0, 1, 2, 3 }; + const k_pos = [_]i64{ 0, 1, 2, 3 }; + const s = try scores(gpa, &q, &k, &q_pos, &k_pos, 1, 4, 4, dim, 2); + defer gpa.free(s); + try std.testing.expectEqual(@as(usize, 2), nb(4, 2)); + // block 0 of row r0: dots [1, 0] (key 1 masked: its k_pos 1 > q_pos 0) → amax 1 + try std.testing.expectEqual(@as(f32, 1), s[0]); + try std.testing.expectEqual(@as(f32, negInf()), s[1]); // block 1 fully masked for r0 + // r1: block 0 dots [1, 1] → 1; key 2 has k_pos 2 > 1 → block 1 masked + try std.testing.expectEqual(@as(f32, 1), s[1 * 2 + 0]); + // r2: block 0 [1, 1], block 1 [6, 0] → 6 + try std.testing.expectEqual(@as(f32, 6), s[2 * 2 + 1]); + // r3: block 1 [6, 6] → 6 + try std.testing.expectEqual(@as(f32, 6), s[3 * 2 + 1]); +} + +fn negInf() f32 { + return -std.math.inf(f32); +} + +test "the top-16 selection: boost override, ties at the smaller id, left-packed, -1 right-pad" { + const gpa = std.testing.allocator; + const s = [_]f32{ 3, 1, 4, 1, 5, -std.math.inf(f32), -std.math.inf(f32), 2 }; + const q_pos = [_]i64{ 7 }; // row r0 sits in block 3 (7 / 2 = 3): its own block boosted over 4's block 2 + const sel = try topk(gpa, &s, 1, 1, 8, 3, 2, 1, 0, &q_pos); + defer gpa.free(sel); + // boosts: block 3 (the row's own) to +inf; then descending: 3 (+inf), 4 (5), 2 (4); id 5,6 are -inf → skipped + try std.testing.expectEqualSlices(i32, &.{ 3, 4, 2 }, sel[0..3]); + const ties = try topk(gpa, &[_]f32{ 4, 4, 4, 6, 9, 9, 2, 2 }, 1, 1, 8, 8, 2, 0, 0, &q_pos); + defer gpa.free(ties); + // no boosts (local 0), descending with every tie at the smaller id, full width, none padded + try std.testing.expectEqualSlices(i32, &.{ 4, 5, 3, 0, 1, 2, 6, 7 }, ties[0..8]); + const short = try topk(gpa, &[_]f32{ 1, 2, -std.math.inf(f32), -std.math.inf(f32) }, 1, 1, 4, 3, 2, 0, 0, &q_pos); + defer gpa.free(short); + try std.testing.expectEqualSlices(i32, &.{ 1, 0, -1 }, short[0..3]); // fewer valid blocks: -1 right-pad +} + +test "the local boost override wins over a higher pooled score, and multiple boosts tie at the smaller id" { + const gpa = std.testing.allocator; + const s = [_]f32{ 100, 1, 1, 1, 50, 50, 50, 50 }; + const q_pos = [_]i64{ 3 }; // the row's own block is 1 + const sel = try topk(gpa, &s, 1, 1, 8, 4, 2, 1, 0, &q_pos); + defer gpa.free(sel); + try std.testing.expectEqualSlices(i32, &.{ 1, 0, 4, 5 }, sel[0..4]); // +inf first, then the 100, then the 50s + const boosted = try topk(gpa, &s, 1, 1, 8, 3, 2, 4, 0, &q_pos); // the row's block and three before it (two in range) + defer gpa.free(boosted); + try std.testing.expectEqualSlices(i32, &.{ 0, 1, 4 }, boosted[0..3]); // boosted ties: smaller ids first, then the 50s +} \ No newline at end of file diff --git a/zig/src/families/minimax_m3_vl/test_pack_config.json b/zig/src/families/minimax_m3_vl/test_pack_config.json new file mode 100644 index 0000000000..d916056236 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/test_pack_config.json @@ -0,0 +1,290 @@ +{ + "architectures": [ + "MiniMaxM3SparseForConditionalGeneration" + ], + "auto_map": { + "AutoConfig": "configuration_minimax_m3_vl.MiniMaxM3VLConfig" + }, + "model_type": "minimax_m3_vl", + "text_config": { + "hidden_size": 6144, + "intermediate_size": 3072, + "num_hidden_layers": 60, + "num_attention_heads": 64, + "num_key_value_heads": 4, + "head_dim": 128, + "vocab_size": 200064, + "max_position_embeddings": 1048576, + "rms_norm_eps": 1e-06, + "use_gemma_norm": true, + "attention_output_gate": false, + "rope_theta": 10000000, + "rotary_dim": 64, + "partial_rotary_factor": 0.5, + "hidden_act": "swigluoai", + "use_qk_norm": true, + "tie_word_embeddings": false, + "dense_intermediate_size": 12288, + "shared_intermediate_size": 3072, + "num_local_experts": 128, + "num_experts_per_tok": 4, + "n_shared_experts": 1, + "scoring_func": "sigmoid", + "use_routing_bias": true, + "moe_layer_freq": [ + 0, + 0, + 0, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1 + ], + "qk_norm_type": "per_head", + "swiglu_alpha": 1.702, + "swiglu_limit": 7.0, + "routed_scaling_factor": 2.0, + "sparse_attention_config": { + "use_sparse_attention": true, + "sparse_index_dim": 128, + "sparse_num_index_heads": 4, + "sparse_topk_blocks": 16, + "sparse_block_size": 128, + "sparse_score_type": "max", + "sparse_init_block": 0, + "sparse_local_block": 1, + "sparse_attention_freq": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1 + ], + "sparse_disable_index_value": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1 + ] + }, + "architectures": [ + "MiniMaxM3SparseForCausalLM" + ] + }, + "vision_config": { + "hidden_size": 1280, + "num_attention_heads": 16, + "num_hidden_layers": 32, + "intermediate_size": 5120, + "patch_size": 14, + "image_size": 672, + "projection_dim": 6144, + "position_embedding_type": "rope", + "rope_mode": "3d", + "rope_theta": 10000.0, + "attention_dropout": 0.0, + "hidden_act": "gelu", + "initializer_factor": 1.0, + "initializer_range": 0.02, + "layer_norm_eps": 1e-05, + "model_type": "clip_vision_model", + "num_channels": 3, + "vocab_size": 32000, + "img_token_compression_config": { + "image_token_compression_method": "patch_merge", + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "vision_segment_max_frames": 4 + }, + "img_token_compression_config": { + "image_token_compression_method": "patch_merge", + "spatial_merge_size": 2, + "temporal_patch_size": 2 + }, + "image_grid_pinpoints": "[(336, 336), (336, 672), (336, 1008), (336, 1344), (336, 1680), (336, 2016), (672, 336), (672, 672), (672, 1008), (672, 1344), (672, 1680), (672, 2016), (1008, 336), (1008, 672), (1008, 1008), (1008, 1344), (1008, 1680), (1008, 2016), (1344, 336), (1344, 672), (1344, 1008), (1344, 1344), (1344, 1680), (1344, 2016), (1680, 336), (1680, 672), (1680, 1008), (1680, 1344), (1680, 1680), (1680, 2016), (2016, 336), (2016, 672), (2016, 1008), (2016, 1344), (2016, 1680), (2016, 2016)]", + "image_seq_length": 576, + "image_token_index": 200025, + "video_token_index": 200026, + "multimodal_projector_bias": true, + "num_reward_heads": 0, + "process_image_mode": "dynamic_res", + "projector_hidden_act": "gelu", + "vision_feature_layer": -1, + "vision_feature_select_strategy": "full", + "torch_dtype": "bfloat16", + "transformers_version": "4.52.4", + "projector_hidden_size": 6144, + "quantization_config": { + "quant_method": "exl3", + "version": "1.6.1", + "head_bits": 8, + "out_scales": "always", + "codebook": "mul1", + "avg_bits": 3.06 + } +} \ No newline at end of file diff --git a/zig/src/families/minimax_m3_vl/test_quantization_config.json b/zig/src/families/minimax_m3_vl/test_quantization_config.json new file mode 100644 index 0000000000..c8c5d8cac4 --- /dev/null +++ b/zig/src/families/minimax_m3_vl/test_quantization_config.json @@ -0,0 +1,8 @@ +{ + "quant_method": "exl3", + "version": "1.6.1", + "head_bits": 8, + "out_scales": "always", + "codebook": "mul1", + "avg_bits": 3.06 +} \ No newline at end of file diff --git a/zig/src/families/minimax_m3_vl/weights.zig b/zig/src/families/minimax_m3_vl/weights.zig new file mode 100644 index 0000000000..0b8e41ae8b --- /dev/null +++ b/zig/src/families/minimax_m3_vl/weights.zig @@ -0,0 +1,430 @@ +//! The checkpoint's whole inventory spelled the pack's way → the load plan: the EXL3 groups from the +//! widened admission (exl3.zig) plus every passthrough tensor the family reads, header-checked by +//! name, dtype and shape exactly as the device loader takes them (the `L.get` law of glm/cuda_weights.zig). + +const std = @import("std"); +const core = @import("core"); +const cfg = @import("config.zig"); +const exl3 = @import("exl3.zig"); + +const checkpoint = core.checkpoint; +const DType = checkpoint.DType; + +/// The routed experts' projections in the engine's order. Measured against the real EAQ v1 pack's +/// shard headers: w1 carries [hidden → moe_inter] (gate), w3 carries [hidden → moe_inter] (up), +/// w2 carries [moe_inter → hidden] (down) — the fork's w2-up/w3-down convention is reversed here. +pub const ExpertProj = enum { + gate, + up, + down, + + /// The checkpoint spelling under `block_sparse_moe.experts..`, verified against the pack's shard headers. + pub fn ckpt(p: ExpertProj) []const u8 { + return switch (p) { + .gate => "w1", + .up => "w3", + .down => "w2", + }; + } +}; + +// The tensors' prefix of the language model's per-layer space (the VL wrapper: language_model.model.*). + +/// The layers' prefix, verbatim from the pack (the head sits one level up: language_model.lm_head.*). +pub const ckpt_prefix = "language_model.model.layers"; + +/// The pack's per-module EXL3 suffix words, exactly the core reader's file layout (`core/exl3_load.zig`). +pub const exl3_suffixes = [_][]const u8{ ".trellis", ".suh", ".svh", ".mul1" }; + +const embed_name = "language_model.model.embed_tokens.weight"; +const final_norm_name = "language_model.model.norm.weight"; + +/// A routed expert `e` of layer `d`: its checkpoint tensor for `proj`, without an EXL3 suffix. +pub fn expertName(buf: []u8, d: u32, e: u32, proj: ExpertProj) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{d}.block_sparse_moe.experts.{d}.{s}", .{ ckpt_prefix, d, e, proj.ckpt() }); +} + +/// The shared expert's tensors keep the standard spellings (EXL3 here too, at the pack's mixed 5-bit widths). +pub fn sharedName(buf: []u8, d: u32, proj: ExpertProj) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{d}.block_sparse_moe.shared_experts.{s}_proj", .{ ckpt_prefix, d, @tagName(proj) }); +} + +/// The dense layers' (first dense_layers') MLP tensors under the plain mlp space. +pub fn denseName(buf: []u8, d: u32, proj: ExpertProj) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{d}.mlp.{s}_proj", .{ ckpt_prefix, d, @tagName(proj) }); +} + +/// The router's score weight (fp32, one row per expert) and its per-expert use_routing_bias. +pub fn routerName(buf: []u8, d: u32) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{d}.block_sparse_moe.gate.weight", .{ ckpt_prefix, d }); +} + +pub fn routerBiasName(buf: []u8, d: u32) ![]const u8 { + // measured on the real EAQ v1 pack: the bias sits beside the gate, not under it + return std.fmt.bufPrint(buf, "{s}.{d}.block_sparse_moe.e_score_correction_bias", .{ ckpt_prefix, d }); +} + +/// The sparse indexer's projections (index_q/index_k, fp16) and their per-head norms (bf16). +pub fn indexerName(buf: []u8, d: u32, comptime side: []const u8, comptime tail: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{d}.self_attn.index_{s}_{s}", .{ ckpt_prefix, d, side, tail }); +} + +/// The attention's q/k per-head norms (qk_norm_type "per_head") and the attention projections themselves. +pub fn attnName(buf: []u8, d: u32, comptime part: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{d}.self_attn.{s}", .{ ckpt_prefix, d, part }); +} + +/// The layer norms' gemma-style pairs (input before attention, post_attention before the MLP). +pub fn layerNormName(buf: []u8, d: u32, comptime which: []const u8) ![]const u8 { + return std.fmt.bufPrint(buf, "{s}.{d}.{s}", .{ ckpt_prefix, d, which }); +} + +/// One load item: a passthrough row or one EXL3 group, with its checkpoint spelling and byte size. +pub const Item = struct { + name: []u8, // owned; groups carry their `.trellis` prefix (the device loader appends the suffix words) + layer: u32, // the layer it belongs to; the embed, final norm and head sit at `past` + kind: Kind, + bytes: u64, // the checkpoint bytes this item reads (trellis+scales for a group, numel*dtype.size otherwise) + + pub const Kind = union(enum) { + passthrough: struct { dtype: DType, rows: u32, cols: u32 }, + group: usize, // the index into the widened admission's `exl3.Plan.groups` + }; +}; + +/// One item past the layers (the embed, the final norm, the head). +pub const past: u32 = std.math.maxInt(u32); + +/// The whole load, in checkpoint order: embed, per layer (norms, attention, indexer, MLP/MoE), final norm, head. +pub const Load = struct { + gpa: std.mem.Allocator, + items: []Item, + config: cfg.Config, // the shapes every item was checked against + head: usize, // the head's EXL3 group item + routers: []usize, // each MoE layer's router weight item, in layer order (its bias trails at +1) + + pub fn deinit(l: *Load) void { + for (l.items) |it| l.gpa.free(it.name); + l.gpa.free(l.items); + l.gpa.free(l.routers); + l.* = undefined; + } + + /// The checkpoint bytes this load reads (the admission's own figure, before any device layout). + pub fn bytes(l: *const Load) u64 { + var n: u64 = 0; + for (l.items) |it| n += it.bytes; + return n; + } + + /// This layer's EXL3 group item for `klas`/`proj` (the head matches by class alone). + pub fn groupAt(l: *const Load, groups: []const exl3.Group, d: u32, klas: exl3.Class, proj: u32) ?usize { + for (l.items, 0..) |*it, i| switch (it.kind) { + .group => |gi| { + const g = groups[gi]; + if (klas == .head) { + if (g.class == .head) return i; + } else if (g.class == klas and g.layer == d and g.proj == proj) return i; + }, + else => {}, + }; + return null; + } +}; + +fn expectRow(ck: *checkpoint.Checkpoint, name: []const u8, dtype: DType, cols: usize) !u64 { + if (!ck.has(name)) return failMissing(name); // screened without the core's err-level log line + const t = try ck.get(name); + if (!t.is(dtype, &.{cols})) return error.UnexpectedTensor; + return @intCast(t.numel() * dtype.size()); +} + +fn expectMat(ck: *checkpoint.Checkpoint, name: []const u8, dtype: DType, rows: usize, cols: usize) !u64 { + if (!ck.has(name)) return failMissing(name); + const t = try ck.get(name); + if (!t.is(dtype, &.{ rows, cols })) return error.UnexpectedTensor; + return @intCast(t.numel() * dtype.size()); +} + +/// The loud refusal for a passthrough the checkpoint lacks: named at warn (the load refuses with +/// `MissingTensor`; the core checkpoint's err-level line would fail the host test runner on the +/// deliberately-missing fixture variant). +fn failMissing(name: []const u8) error{MissingTensor} { + std.log.warn("minimax_m3_vl: the checkpoint has no tensor {s}", .{name}); + return error.MissingTensor; +} + +const next = struct { + fn f(consumed: []bool, gs: []const exl3.Group, d: u32, klas: exl3.Class, proj: u32) ?usize { + for (gs, 0..) |g, gi| { + if (consumed[gi] or g.class != klas) continue; + if (klas == .head) return gi; // one group, past the layers + if (g.layer == d and g.proj == proj) return gi; + } + return null; + } +}.f; + +fn row(gpa: std.mem.Allocator, items: *std.ArrayList(Item), name: []const u8, dtype: DType, d: u32, rows: u32, cols: u32, byte_count: u64) !void { + try items.append(gpa, .{ .name = try gpa.dupe(u8, name), .layer = d, .kind = .{ .passthrough = .{ .dtype = dtype, .rows = rows, .cols = cols } }, .bytes = byte_count }); +} + +fn grp(gpa: std.mem.Allocator, items: *std.ArrayList(Item), name: []const u8, d: u32, gi: usize, gs: []const exl3.Group) !void { + try items.append(gpa, .{ .name = try gpa.dupe(u8, name), .layer = d, .kind = .{ .group = gi }, .bytes = exl3.groupBytes(gs[gi]) }); +} + +/// Every passthrough tensor and EXL3 group the family reads, checked against `dir`'s checkpoint by +/// name, dtype and shape; every admitted group must be consumed (a surplus trellis refuses). TP1: +/// the head and the embed in whole; the vocabulary's rank cut is the device loader's job from this plan. +pub fn plan(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, c: *const cfg.Config, groups: *const exl3.Plan) !Load { + var ck = try checkpoint.Checkpoint.openModel(gpa, io, dir); + defer ck.close(); + var consumed = try gpa.alloc(bool, groups.groups.len); + defer gpa.free(consumed); // the checklist runs on every return path + @memset(consumed, false); + var items: std.ArrayList(Item) = .empty; + errdefer { + for (items.items) |it| gpa.free(it.name); + items.deinit(gpa); + } + var routers: std.ArrayList(usize) = .empty; + errdefer routers.deinit(gpa); + var nb: [320]u8 = undefined; + var eb: [320]u8 = undefined; + + const head_gi = next(consumed, groups.groups, 0, .head, 0) orelse return error.MissingTrellis; + consumed[head_gi] = true; + + const embed_bytes = try expectMat(&ck, embed_name, .bf16, c.vocab, c.hidden); + try row(gpa, &items, embed_name, .bf16, past, c.vocab, c.hidden, embed_bytes); + + for (0..c.layers) |li| { + const d: u32 = @intCast(li); + // the gemma-style layer norms and the per-head qk norms, bf16 rows of the Config's dims + inline for ([_][]const u8{ "input_layernorm.weight", "post_attention_layernorm.weight" }) |which| { + const n = try expectRow(&ck, try layerNormName(&nb, d, which), .bf16, c.hidden); + try row(gpa, &items, try layerNormName(&nb, d, which), .bf16, d, 1, c.hidden, n); + } + // the indexer: fp16 index_q/index_k projections, their per-head norms and the attention's per-head qk norms + inline for ([_][]const u8{ "q", "k" }) |side| { + const qkn = try expectRow(&ck, try attnName(&nb, d, side ++ "_norm.weight"), .bf16, c.head_dim); + try row(gpa, &items, try attnName(&nb, d, side ++ "_norm.weight"), .bf16, d, 1, c.head_dim, qkn); + const ikn = try expectRow(&ck, try indexerName(&nb, d, side, "norm.weight"), .bf16, c.i_dim); + try row(gpa, &items, try indexerName(&nb, d, side, "norm.weight"), .bf16, d, 1, c.i_dim, ikn); + // index_q projects per head: [i_heads * i_dim, hidden]. index_k is ONE shared row set + // (the fork's MiniMaxM3Indexer k_proj out_features = head_dim; idx_k is [B, 1, T, D]): + // [i_dim, hidden] — the side plane every head scores against. + const ik_rows: u32 = if (std.mem.eql(u8, side, "q")) c.i_heads * c.i_dim else c.i_dim; + const ikr = try expectMat(&ck, try indexerName(&nb, d, side, "proj.weight"), .f16, ik_rows, c.hidden); + try row(gpa, &items, try indexerName(&nb, d, side, "proj.weight"), .f16, d, ik_rows, c.hidden, ikr); + } + // the attention projections: the 4 EXL3 groups, in q,k,v,o order + inline for (.{ "q_proj", "k_proj", "v_proj", "o_proj" }, 0..) |p, pi| { + const gi = next(consumed, groups.groups, d, .attn, @intCast(pi)) orelse return error.MissingTrellis; + consumed[gi] = true; + try grp(gpa, &items, try std.fmt.bufPrint(&nb, "{s}.{d}.self_attn.{s}", .{ ckpt_prefix, d, p }), d, gi, groups.groups); + } + if (d < c.dense_layers) { + // the dense MLP's three groups (the pack: 5-bit), under the plain mlp space + inline for ([_]ExpertProj{ .gate, .up, .down }, 0..) |proj, pi| { + const gi = next(consumed, groups.groups, d, .dense_mlp, @intCast(pi)) orelse return error.MissingTrellis; + consumed[gi] = true; + try grp(gpa, &items, try denseName(&nb, d, proj), d, gi, groups.groups); + } + } else { + // the MoE layers: the shared expert's three groups, both router tensors, all the routed experts + inline for ([_]ExpertProj{ .gate, .up, .down }, 0..) |proj, pi| { + const gi = next(consumed, groups.groups, d, .shared, @intCast(pi)) orelse return error.MissingTrellis; + consumed[gi] = true; + try grp(gpa, &items, try sharedName(&nb, d, proj), d, gi, groups.groups); + } + const rbytes = try expectMat(&ck, try routerName(&nb, d), .f32, c.experts, c.hidden); + try routers.append(gpa, items.items.len); + try row(gpa, &items, try routerName(&nb, d), .f32, d, c.experts, c.hidden, rbytes); + const bbytes = try expectRow(&ck, try routerBiasName(&nb, d), .f32, c.experts); + try row(gpa, &items, try routerBiasName(&nb, d), .f32, d, 1, c.experts, bbytes); + for (0..c.experts) |el| { + inline for ([_]ExpertProj{ .gate, .up, .down }) |proj| { + const gi = next(consumed, groups.groups, d, .routed, @intFromEnum(proj)) orelse return error.MissingTrellis; + consumed[gi] = true; + try grp(gpa, &items, try expertName(&eb, d, @intCast(el), proj), d, gi, groups.groups); + } + } + } + } + const fbytes = try expectRow(&ck, final_norm_name, .bf16, c.hidden); + try row(gpa, &items, final_norm_name, .bf16, past, 1, c.hidden, fbytes); + try grp(gpa, &items, groups.groups[head_gi].prefix, past, head_gi, groups.groups); + + var leftover: usize = 0; + for (consumed) |done| if (!done) { + leftover += 1; + }; + if (leftover != 0) return error.SurplusTrellis; // an admitted trellis the family never reads + const taken = try items.toOwnedSlice(gpa); + return .{ + .gpa = gpa, + .items = taken, + .config = c.*, + .head = taken.len - 1, + .routers = try routers.toOwnedSlice(gpa), + }; +} + +// --- tests: the verbatim spellings first, then the whole plan against a synthetic checkpoint --- + +test "the inventory's spellings map w1/w2/w3 onto gate/up/down, prefixed the checkpoint's way" { + var buf: [320]u8 = undefined; + try std.testing.expectEqualStrings("language_model.model.layers.7.block_sparse_moe.experts.0.w1", try expertName(&buf, 7, 0, .gate)); + try std.testing.expectEqualStrings("language_model.model.layers.57.block_sparse_moe.experts.3.w2", try expertName(&buf, 57, 3, .down)); + try std.testing.expectEqualStrings("language_model.model.layers.11.block_sparse_moe.experts.0.w3", try expertName(&buf, 11, 0, .up)); + try std.testing.expectEqualStrings("language_model.model.layers.9.block_sparse_moe.shared_experts.up_proj", try sharedName(&buf, 9, .up)); + try std.testing.expectEqualStrings("language_model.model.layers.0.mlp.gate_proj", try denseName(&buf, 0, .gate)); + try std.testing.expectEqualStrings("language_model.model.layers.29.block_sparse_moe.gate.weight", try routerName(&buf, 29)); + try std.testing.expectEqualStrings("language_model.model.layers.5.block_sparse_moe.e_score_correction_bias", try routerBiasName(&buf, 5)); // use_routing_bias; measured: beside the gate, not under it + try std.testing.expectEqualStrings("language_model.model.layers.26.self_attn.index_k_proj.weight", try indexerName(&buf, 26, "k", "proj.weight")); + try std.testing.expectEqualStrings("language_model.model.layers.44.self_attn.index_k_norm.weight", try indexerName(&buf, 44, "k", "norm.weight")); + try std.testing.expectEqualStrings("language_model.model.layers.47.self_attn.q_norm.weight", try attnName(&buf, 47, "q_norm.weight")); + try std.testing.expectEqualStrings("language_model.model.layers.59.self_attn.o_proj", try attnName(&buf, 59, "o_proj")); + try std.testing.expectEqualStrings("language_model.model.layers.3.input_layernorm.weight", try layerNormName(&buf, 3, "input_layernorm.weight")); + try std.testing.expectEqual(@as(usize, 4), exl3_suffixes.len); // trellis words, both scales, the marker + try std.testing.expectEqualStrings(".trellis", exl3_suffixes[0]); + try std.testing.expectEqualStrings(".mul1", exl3_suffixes[3]); +} + +/// The passthrough tensors written beside the trellis fixture (a second shard), dtype and shape checked. +const Pass = struct { + gpa: std.mem.Allocator, + header: std.ArrayListUnmanaged(u8) = .empty, + data: std.ArrayListUnmanaged(u8) = .empty, + first: bool = true, + + fn tensor(p: *Pass, name: []const u8, dtype_text: []const u8, shape_text: []const u8, len: usize) !void { + if (name.len == 0) return; // the missing-tensor variant names nothing and writes nothing + const base = p.data.items.len; + try p.data.appendNTimes(p.gpa, 0, len); + if (!p.first) try p.header.appendSlice(p.gpa, ","); + p.first = false; + try p.header.print(p.gpa, "\"{s}\":{{\"dtype\":\"{s}\",\"shape\":[{s}],\"data_offsets\":[{d},{d}]}}", .{ name, dtype_text, shape_text, base, base + len }); + } + + fn done(p: *Pass) ![]u8 { + try p.header.insertSlice(p.gpa, 0, "{"); // the safetensors header is one JSON object of the tensors + try p.header.appendSlice(p.gpa, "}"); + var file: std.ArrayListUnmanaged(u8) = .empty; + errdefer file.deinit(p.gpa); + try file.appendNTimes(p.gpa, 0, 8); + std.mem.writeInt(u64, file.items[0..8], p.header.items.len, .little); + try file.appendSlice(p.gpa, p.header.items); + try file.appendSlice(p.gpa, p.data.items); + return file.toOwnedSlice(p.gpa); + } +}; + +fn writePassthrough(gpa: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, skip_index_q: bool) !void { + const c = exl3.synthConfig(); + var p: Pass = .{ .gpa = gpa }; + defer { + p.header.deinit(gpa); + p.data.deinit(gpa); + } + var nb: [320]u8 = undefined; + var eb: [320]u8 = undefined; + var sh: [64]u8 = undefined; + var name: []const u8 = undefined; + var shape: []const u8 = undefined; + shape = try std.fmt.bufPrint(&sh, "{d},{d}", .{ c.vocab, c.hidden }); + try p.tensor(embed_name, "BF16", shape, c.vocab * c.hidden * 2); + for ([_]u32{ 0, 1 }) |l| { + inline for ([_][]const u8{ "input_layernorm.weight", "post_attention_layernorm.weight" }) |which| { + name = try layerNormName(&nb, l, which); + shape = try std.fmt.bufPrint(&sh, "{d}", .{c.hidden}); + try p.tensor(name, "BF16", shape, c.hidden * 2); + } + name = try attnName(&nb, l, "q_norm.weight"); + shape = try std.fmt.bufPrint(&sh, "{d}", .{c.head_dim}); + try p.tensor(name, "BF16", shape, c.head_dim * 2); // the per-head qk norms, one row per head's dim + name = try attnName(&nb, l, "k_norm.weight"); + try p.tensor(name, "BF16", shape, c.head_dim * 2); + name = try indexerName(&nb, l, "q", "norm.weight"); + shape = try std.fmt.bufPrint(&sh, "{d}", .{c.i_dim}); + try p.tensor(name, "BF16", shape, c.i_dim * 2); + name = try indexerName(&nb, l, "k", "norm.weight"); + try p.tensor(name, "BF16", shape, c.i_dim * 2); + name = try indexerName(&eb, l, "q", "proj.weight"); // the fp16 projection: [i_heads * i_dim, hidden] + shape = try std.fmt.bufPrint(&sh, "{d},{d}", .{ c.i_heads * c.i_dim, c.hidden }); + try p.tensor(if (skip_index_q) "" else name, "F16", shape, c.i_heads * c.i_dim * c.hidden * 2); + name = try indexerName(&eb, l, "k", "proj.weight"); // index_k is ONE shared projection: [i_dim, hidden] + shape = try std.fmt.bufPrint(&sh, "{d},{d}", .{ c.i_dim, c.hidden }); + try p.tensor(name, "F16", shape, c.i_dim * c.hidden * 2); + if (l == 1) { // this fixture's one MoE layer: the router's fp32 row per expert plus its bias + name = try routerName(&nb, l); + shape = try std.fmt.bufPrint(&sh, "{d},{d}", .{ c.experts, c.hidden }); + try p.tensor(name, "F32", shape, c.experts * c.hidden * 4); + name = try routerBiasName(&nb, l); + shape = try std.fmt.bufPrint(&sh, "{d}", .{c.experts}); + try p.tensor(name, "F32", shape, c.experts * 4); + } + } + shape = try std.fmt.bufPrint(&sh, "{d}", .{c.hidden}); + try p.tensor(final_norm_name, "BF16", shape, c.hidden * 2); + const file = try p.done(); + defer gpa.free(file); + try dir.writeFile(io, .{ .sub_path = "model-00002.safetensors", .data = file }); +} + +/// The tests' synthetic checkpoint: the trellis pack, the passthrough tensors beside it, and the +/// config.json admission reads. cuda_native.zig's admission test reuses it. `skip_index_q` drops the +/// fp16 index projection so the same dir exercises a missing-tensor refusal. +pub fn checkpointFixture(gpa: std.mem.Allocator, io: std.Io, dir: std.Io.Dir, skip_index_q: bool) !void { + try exl3.writeFixture(gpa, io, dir, .{}); + try writePassthrough(gpa, io, dir, skip_index_q); + try exl3.writeSynthConfigJson(gpa, io, dir); +} + +test "the whole synthetic checkpoint plans: every passthrough tensor and admitted group, in load order" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try checkpointFixture(gpa, io, tmp.dir, false); + const path = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(path); + const c = exl3.synthConfig(); // the fixture's own shapes, whole 128s at every tile edge + var p = try exl3.discover(gpa, io, path, &c); + defer p.deinit(); + var l = try plan(gpa, io, path, &c, &p); + defer l.deinit(); + try std.testing.expectEqual(@as(usize, 1), l.routers.len); // one MoE layer's router in this fixture + try std.testing.expectEqual(@as(usize, 47), l.items.len); // 1 embed + 16 rows + 8 attn + 3 + 3 + 2 + 12 + 1 + 1 + try std.testing.expectEqual(p.groups.len, blk: { // every admitted group became exactly one item + var n: usize = 0; + for (l.items) |it| switch (it.kind) { + .group => n += 1, + .passthrough => {}, + }; + break :blk n; + }); + try std.testing.expectEqual(@as(u64, 512 * c.hidden * 2), l.items[0].bytes); // the embed's bf16 table + try std.testing.expectEqual(l.items.len - 1, l.head); // the head is the plan's last item, past the layers + try std.testing.expect(l.groupAt(p.groups, 0, .attn, 0) != null); + try std.testing.expect(l.groupAt(p.groups, 0, .dense_mlp, 1) != null); + try std.testing.expect(l.groupAt(p.groups, 1, .shared, 2) != null); + try std.testing.expect(l.groupAt(p.groups, 1, .routed, 2) != null); + try std.testing.expectEqual(@as(usize, l.head), l.groupAt(p.groups, 0, .head, 0).?); + try std.testing.expect(l.bytes() > 0); +} + +test "a missing passthrough tensor refuses by name, and a plan consumes every group" { + const gpa = std.testing.allocator; + const io = std.testing.io; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + try checkpointFixture(gpa, io, tmp.dir, true); // no index_q_proj in the checkpoint + const path = try tmp.dir.realPathFileAlloc(io, ".", gpa); + defer gpa.free(path); + const c = exl3.synthConfig(); + var p = try exl3.discover(gpa, io, path, &c); + defer p.deinit(); + try std.testing.expectError(error.MissingTensor, plan(gpa, io, path, &c, &p)); +} \ No newline at end of file From 803c394653ae414c548e407c8d0247d0a931ed81 Mon Sep 17 00:00:00 2001 From: styles01 <11393400+styles01@users.noreply.github.com> Date: Sat, 24 Oct 2026 23:31:56 -0600 Subject: [PATCH 2/7] feat(build): cuda glue for minimax_m3_vl Build registration for the minimax_m3 kernel image, the tf_m3 vtable wiring, the CUDA/Metal family admission lines (Metal: stage-1 admission, no Metal weight reader), and the CLI hub/info entries with the EXL3 format reader. --- zig/build/cuda.zig | 13 ++++++- zig/build/cuda_kernels.zig | 1 + zig/src/cli/hub.zig | 78 ++++++++++++++++++++++++++++++++++++++ zig/src/cli/info.zig | 2 + zig/src/cuda/kernels.zig | 1 + zig/src/native/cuda.zig | 9 ++++- zig/src/native/metal.zig | 1 + 7 files changed, 103 insertions(+), 2 deletions(-) diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 750c4e3173..bd74225ebc 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -117,6 +117,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); const glm_cuda = glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat); runner.addImport("glm_cuda", glm_cuda); + runner.addImport("minimax_m3_vl_cuda", minimaxM3Native(b, target, optimize, cuda, mods.core, mods.lanes)); runner.addImport("kolibri1", kolibri1Module(b, target, optimize, cuda, mods.core, mods.lanes)); const runner_exe = b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }); b.installArtifact(runner_exe); @@ -158,6 +159,15 @@ fn glmCuda(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti return m; } +/// MiniMax-M3.1-Flash's family module: the openWith API, the pack's admission and the MSA contracts. +fn minimaxM3Native(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/minimax_m3_vl/cuda_native.zig"), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("cuda", cuda); + m.addImport("core", core); + m.addImport("lanes", lanes); + return m; +} + /// The Kolibri-1 CUDA family over the runtime, the core and the lane core. const kolibri1Module = @import("kolibri.zig").module; @@ -205,6 +215,7 @@ fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti .{ .name = "core", .module = core }, .{ .name = "flashnext", .module = flashnextModule(b, target, optimize, cuda, core, lanes, qwen_vision) }, .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) }, + .{ .name = "minimax_m3_vl_cuda", .module = minimaxM3Native(b, target, optimize, cuda, core, lanes) }, .{ .name = "qwen35", .module = qwen35 }, .{ .name = "kolibri1", .module = kolibri1Module(b, target, optimize, cuda, core, lanes) }, }, @@ -483,7 +494,7 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st @import("cortheon_fuzz.zig").add(b, host, native_modules.api, mods.tokenizer, template); } // the tokenizer's own tests (memo, steps, regex, unicode) run here: no other test binary reaches them - for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.qwen_vision, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision), lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native_modules.api, native, stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); + for ([_]*std.Build.Module{ cuda, mods.core, mods.lanes, mods.nemotron, mods.qwen35, mods.qwen_vision, flashnextModule(b, host, .debug, cuda, mods.core, mods.lanes, mods.qwen_vision), lfm2Module(b, host, .debug, cuda, mods.core), mods.radio, mods.heat, mods.tokenizer, native_modules.api, native, minimaxM3Native(b, host, .debug, cuda, mods.core, mods.lanes), stagger(b, host, .debug) }) |m| step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = m })).step); const selected = b.option([]const u8, "modelopt-host-filter", "One portable ModelOpt regression") orelse "modelopt host:"; const run = b.addRunArtifact(b.addTest(.{ .root_module = mods.nemotron, .filters = &.{selected} })); run.has_side_effects = true; diff --git a/zig/build/cuda_kernels.zig b/zig/build/cuda_kernels.zig index 9315da5563..55a1c0a767 100644 --- a/zig/build/cuda_kernels.zig +++ b/zig/build/cuda_kernels.zig @@ -90,6 +90,7 @@ pub const kernels = [_]Kernel{ .{ .name = "modelopt_pack", .flags = &.{"-O3"} }, // ours: ModelOpt tensors into the lane, prompt and expert layouts .{ .name = "glm_kda", .flags = &.{ "-O3", "--fmad=false" } }, // GLM: the Python family's glm5_next/cuda/kda.cu .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op + .{ .name = "minimax_m3", .flags = glue }, // ours: MiniMax M3.1's MSA indexer, the fork m3 contracts (glue flags keep the fp16 dots' products exact) .{ .name = "glm_ops", .flags = &.{"-O3"} }, // ours: GLM's integer work the Python family did in torch (the DSA token lists) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family .{ .name = "exl3_linear", .flags = &.{ "-O3", "--expt-relaxed-constexpr" } }, // exl3/linear.cu's device code, tensorfold_exl3_linear_v5 diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 65395cde50..86c322d219 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -18,6 +18,7 @@ const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "gemma4_unified", .title = "Gemma 4 12B" }, .{ .model_type = "step5_vl", .title = "Step-5-Preview" }, .{ .model_type = "lfm2_moe", .title = "LFM2.5-24B-A2B" }, + .{ .model_type = "minimax_m3_vl", .title = "MiniMax-M3.1-Flash VL" }, }; /// The family that serves `model_type` from the engine registry, or null when Zig cannot serve it (the 0.6 line may). @@ -141,6 +142,55 @@ pub fn quantization(a: Allocator, io: std.Io, dir: []const u8) ?struct { bits: u return .{ .bits = @intCast(bits), .group = @intCast(group) }; } +/// An EXL3 checkpoint's format word, read from the quantization metadata itself (the +/// glm/cuda_exl3.zig block() reading): config.json's inline `quantization_config` at the top +/// level or inside `text_config`, else the standalone quantization_config.json — the second +/// file an EXL3 pack ships when config.json's inline copy is absent. Returns "exl3/3.06bpw" +/// (the metadata's avg_bits) or a plain "exl3" without it; null when the checkpoint carries no +/// EXL3 metadata at all (MLX affine and other quantizations are read elsewhere). +pub fn exl3Format(a: Allocator, io: std.Io, dir: []const u8) ?[]const u8 { + if (@import("checkpoint_config").optional(a, io, dir)) |text| { + if (std.json.parseFromSlice(std.json.Value, a, text, .{})) |inline_parsed| { + defer inline_parsed.deinit(); + if (block(inline_parsed.value)) |q| return fmtExl3(a, q.object); + return inlineFallback(a, io, dir); // no EXL3 inline: maybe the standalone file has it + } else |_| return inlineFallback(a, io, dir); // a broken config.json: the standalone file still speaks + } + return inlineFallback(a, io, dir); // no config.json body: the standalone file is the metadata +} + +/// The standalone quantization_config.json's EXL3 metadata, when config.json's inline copy did not name it. +fn inlineFallback(a: Allocator, io: std.Io, dir: []const u8) ?[]const u8 { + const path = std.fs.path.join(a, &.{ dir, "quantization_config.json" }) catch return null; + const text = readSmall(a, io, path) orelse return null; + var parsed = std.json.parseFromSlice(std.json.Value, a, text, .{}) catch return null; + defer parsed.deinit(); + const q = if (parsed.value == .object) parsed.value.object.get("quantization_config") orelse parsed.value else return null; + if (q != .object) return null; + return fmtExl3(a, q.object); +} + +/// config.json's quantization_config: the top level's, else text_config's (a VL wrapper's own block). +fn block(root: std.json.Value) ?std.json.Value { + if (root != .object) return null; + if (root.object.get("quantization_config")) |q| if (q == .object) return q; + if (root.object.get("text_config")) |t| if (t == .object) { + if (t.object.get("quantization_config")) |q| if (q == .object) return q; + }; + return null; +} + +fn fmtExl3(a: Allocator, q: std.json.ObjectMap) ?[]const u8 { + const m = q.get("quant_method") orelse return null; + if (m != .string or !std.mem.eql(u8, m.string, "exl3")) return null; + const avg = q.get("avg_bits") orelse return "exl3"; + return switch (avg) { + .float => std.fmt.allocPrint(a, "exl3/{d}bpw", .{avg.float}) catch "exl3", + .integer => std.fmt.allocPrint(a, "exl3/{d}bpw", .{@as(f64, @floatFromInt(avg.integer))}) catch "exl3", + else => "exl3", + }; +} + /// Bytes of the files under `dir`, symlinks followed, hidden ones left out (a download's .cache, --slide's snapshot). pub fn sizeOf(a: Allocator, io: std.Io, dir: []const u8) !u64 { var total: u64 = 0; @@ -202,6 +252,34 @@ test "family detection matches the registry and refuses unknown types" { try std.testing.expect(family("unknown") == null); } +// The EXL3 format word an `info` line reports: inline at the top level, inline inside +// text_config, or the standalone quantization_config.json when config.json carries none; +// MLX `quantization` and absent metadata read null. Verbatim from the staged pack's bytes. +test "exl3 format reads the quantization metadata wherever the pack carries it" { + var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator); + defer arena_state.deinit(); + const a = arena_state.allocator(); + const io = std.testing.io; + const cwd = std.Io.Dir.cwd(); + for ([_][]const u8{ "a", "b", "c", "e", "f" }) |name| try cwd.createDirPath(io, try std.fmt.allocPrint(a, ".zig-cache/tmp/{s}", .{name})); + // inline top level, the pack's own avg_bits (the staged pack's shape) + try cwd.writeFile(io, .{ .sub_path = ".zig-cache/tmp/a/config.json", .data = "{\"quantization_config\":{\"quant_method\":\"exl3\",\"codebook\":\"mul1\",\"avg_bits\":3.06}}" }); + try std.testing.expectEqualStrings("exl3/3.06bpw", exl3Format(a, io, ".zig-cache/tmp/a").?); + // inline inside text_config + try cwd.writeFile(io, .{ .sub_path = ".zig-cache/tmp/b/config.json", .data = "{\"text_config\":{\"quantization_config\":{\"quant_method\":\"exl3\",\"avg_bits\":3.5}}}" }); + try std.testing.expectEqualStrings("exl3/3.5bpw", exl3Format(a, io, ".zig-cache/tmp/b").?); + // no inline copy: the standalone quantization_config.json speaks, exactly Mia-AiLab's + try cwd.writeFile(io, .{ .sub_path = ".zig-cache/tmp/c/config.json", .data = "{\"model_type\":\"minimax_m3_vl\"}" }); + try cwd.writeFile(io, .{ .sub_path = ".zig-cache/tmp/c/quantization_config.json", .data = "{\"quant_method\":\"exl3\",\"version\":\"1.6.1\",\"head_bits\":8,\"out_scales\":\"always\",\"codebook\":\"mul1\",\"avg_bits\":3.06}" }); + try std.testing.expectEqualStrings("exl3/3.06bpw", exl3Format(a, io, ".zig-cache/tmp/c").?); + // no EXL3 at all + try cwd.writeFile(io, .{ .sub_path = ".zig-cache/tmp/e/config.json", .data = "{\"quantization\":{\"bits\":4,\"group_size\":32}}" }); + try std.testing.expect(exl3Format(a, io, ".zig-cache/tmp/e") == null); + // other quant methods read null (an eaq or modelopt checkpoint is not this format) + try cwd.writeFile(io, .{ .sub_path = ".zig-cache/tmp/f/config.json", .data = "{\"quantization_config\":{\"quant_method\":\"eaq\"}}" }); + try std.testing.expect(exl3Format(a, io, ".zig-cache/tmp/f") == null); +} + test "cache layout helpers read a synthetic cache" { var arena_state = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena_state.deinit(); diff --git a/zig/src/cli/info.zig b/zig/src/cli/info.zig index 64e1aecac2..81b0a64381 100644 --- a/zig/src/cli/info.zig +++ b/zig/src/cli/info.zig @@ -41,6 +41,8 @@ pub fn run(a: Allocator, io: std.Io, out: *std.Io.Writer, env: ?*const std.proce try out.print("family: {s} ({s})\n", .{ family.title, family.model_type }); if (hub.quantization(a, io, dir)) |q| { try out.print("format: {d}-bit, group {d}\n", .{ q.bits, q.group }); + } else if (hub.exl3Format(a, io, dir)) |f| { + try out.print("format: {s}\n", .{f}); // EXL3's own metadata speaks: the pack's packed words } else { try out.print("format: full precision\n", .{}); } diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 2b7ff0fa6e..cc20f81add 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -57,6 +57,7 @@ pub const moe_glue: []const u8 = if (available) &Blob("fatbin_moe_glue").bytes e pub const dense_rows: []const u8 = if (available) &Blob("fatbin_dense_rows").bytes else &.{}; pub const modelopt_pack: []const u8 = if (available) &Blob("fatbin_modelopt_pack").bytes else &.{}; pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes else &.{}; +pub const minimax_m3: []const u8 = if (available) &Blob("fatbin_minimax_m3").bytes else &.{}; // the M3.1 MSA indexer pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_vision: []const u8 = if (available) &Blob("fatbin_glm_vision").bytes else &.{}; diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 52aff4316d..b0c3184a83 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -13,6 +13,7 @@ const family_options = @import("cuda_options.zig"); pub const follow = @import("cuda_follow.zig").follow; const lfm2 = @import("lfm2.zig"); const glm = @import("glm_cuda"); +const minimax_m3_vl = @import("minimax_m3_vl_cuda"); const qwen35 = @import("qwen35"); const kolibri1 = @import("kolibri1"); const Allocator = std.mem.Allocator; @@ -25,7 +26,7 @@ const vision_offer = @import("cuda_vision_offer.zig"); const drafted = @import("cuda_drafted.zig"); const kv_pages = @import("cuda_kv_pool.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native }; +const registry = .{ nemotron.native, nemotron.omni, glm.native, qwen35.native, qwen35.native_moe, flashnext.native, kolibri1.native, minimax_m3_vl.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; pub const backends: []const []const u8 = &.{"cuda"}; @@ -184,6 +185,9 @@ pub fn open(a: Allocator, gpa: Allocator, io: std.Io, o: api.Open, problem: *[]c } return lfm2.open(a, gpa, io, o, problem); } + if (std.mem.eql(u8, o.model_type, minimax_m3_vl.native.model_type)) { + return openWith(minimax_m3_vl, a, gpa, io, o, problem); + } problem.* = try std.fmt.allocPrint(a, "the native CUDA engine has no backend for {s} checkpoints yet", .{o.model_type}); return null; } @@ -490,6 +494,9 @@ test "every registered family and d1's decision engine are listed for capabiliti try std.testing.expectEqual(@as(usize, registry.len + 2), families.len); try std.testing.expectEqualStrings("nemotron_h", families[0].model_type); try std.testing.expectEqualStrings("lfm2_vl", families[registry.len].model_type); + try std.testing.expectEqualStrings("fish_omni", families[registry.len + 1].model_type); + // minimax_m3_vl dispatches through the registry itself (stage 2), not the metadata table anymore + try std.testing.expectEqualStrings("minimax_m3_vl", families[registry.len - 1].model_type); } test "a cancel during a long prompt stops it, and the next request's reply is unchanged (#291)" { diff --git a/zig/src/native/metal.zig b/zig/src/native/metal.zig index 5cf5bff1f9..67a971b01b 100644 --- a/zig/src/native/metal.zig +++ b/zig/src/native/metal.zig @@ -43,6 +43,7 @@ pub const families: []const api.Family = &.{ .{ .model_type = "gemma4", .formats = gemma4_widths }, .{ .model_type = "gemma4_unified", .formats = gemma4_widths }, .{ .model_type = "step5_vl", .formats = &.{"mlx-q2g128"} }, + .{ .model_type = "minimax_m3_vl", .formats = &.{"exl3"} }, // stage-1 admission: no Metal reader yet }; /// Gemma 4 31B (gemma4) and 12B (gemma4_unified, its text model): one native family, gemma4.zig's host. From 0c82d589ced815ef81b45db14eea8e518ad979af Mon Sep 17 00:00:00 2001 From: styles01 <11393400+styles01@users.noreply.github.com> Date: Sat, 24 Oct 2026 23:31:56 -0600 Subject: [PATCH 3/7] test(family): minimax_m3_vl device parity, graph and TP2 tests Device parity harness vs the CPU references (scores/topk/decode+combine/ swigluoai, NB/nsplit bound guards, zero-valid-key), the layer-graph token harness, TP2 rank/env-place refusal tests, and the exl3_experts f64 additions. --- zig/tests/cuda/exl3_experts_f64_tests.zig | 67 +- zig/tests/cuda/exl3_experts_tests.zig | 4 +- zig/tests/cuda/main.zig | 46 + zig/tests/cuda/minimax_m3_graph_tests.zig | 1937 +++++++++++++++++++++ zig/tests/cuda/minimax_m3_tests.zig | 644 +++++++ zig/tests/cuda/minimax_tp2_tests.zig | 351 ++++ 6 files changed, 3036 insertions(+), 13 deletions(-) create mode 100644 zig/tests/cuda/minimax_m3_graph_tests.zig create mode 100644 zig/tests/cuda/minimax_m3_tests.zig create mode 100644 zig/tests/cuda/minimax_tp2_tests.zig diff --git a/zig/tests/cuda/exl3_experts_f64_tests.zig b/zig/tests/cuda/exl3_experts_f64_tests.zig index 6e0c162cbb..a369163d14 100644 --- a/zig/tests/cuda/exl3_experts_f64_tests.zig +++ b/zig/tests/cuda/exl3_experts_f64_tests.zig @@ -45,7 +45,7 @@ pub const Built = struct { pub fn checks(r: *Rig, exp: *const exl3e.Experts, state: *Run) void { for (f64_ids, 0..) |id, i| { var ok = true; - f64Check(r, exp, &allBits(@intCast(i + 4))) catch |e| { + f64Check(r, exp, &allBits(@intCast(i + 4)), 1) catch |e| { state.fail(e, id); ok = false; }; @@ -53,12 +53,28 @@ pub fn checks(r: *Rig, exp: *const exl3e.Experts, state: *Run) void { } { var ok = true; - f64Check(r, exp, &.{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }) catch |e| { + f64Check(r, exp, &.{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }, 1) catch |e| { state.fail(e, f64_ids[2]); ok = false; }; if (ok) check.pass("{s}", .{f64_ids[2]}); } + { // the GLM bf16 act mode (0), mixed widths — mode 1 (the engine's) is the bar above + var ok = true; + f64Check(r, exp, &.{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }, 0) catch |e| { + state.fail(e, "exl3.experts.f64.act0"); + ok = false; + }; + if (ok) check.pass("exl3.experts.f64.act0", .{}); + } + { // the engine's fp32 x contract (the minimax routed wave): mixed widths, the act fp32 mode + var ok = true; + f32xCheck(r, exp, &.{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }) catch |e| { + state.fail(e, "exl3.experts.f32x"); + ok = false; + }; + if (ok) check.pass("exl3.experts.f32x", .{}); + } } fn allBits(bits: u32) [E]u32 { @@ -79,7 +95,7 @@ fn bf16r(v: f64) f64 { return bf(@intCast((u + 0x7fff + ((u >> 16) & 1)) >> 16)); } -pub fn buildLayer(r: *Rig, bits: []const u32) !Built { +pub fn buildLayer(r: *Rig, bits: []const u32, act: i32) !Built { var b: Built = .{ .layer = undefined, .bits = bits, @@ -152,7 +168,7 @@ pub fn buildLayer(r: *Rig, bits: []const u32) !Built { .width = I, .codebook = .mul1, .limit = LIMIT, - .act_mode = 0, + .act_mode = act, // 0: GLM's bf16 act, 1: the minimax engine's fp32 act — both are graded bars }; b.x = try r.bfs(CAP * D, 0.5); // half-range x, with the small scales, keeps the SwiGLU inputs moderate b.pick = try picks(r, CAP); @@ -210,12 +226,31 @@ pub fn runSlice(r: *Rig, exp: *const exl3e.Experts, layer: exl3e.Layer, x: u64, return .{ .out = out, .y = sc.y }; } -fn f64Check(r: *Rig, exp: *const exl3e.Experts, bits: []const u32) !void { - const b = try buildLayer(r, bits); +/// The same call with fp32 x rows — the minimax engine's exact contract (x32 with .f32); the rot path must branch on it. +pub fn runSliceF32(r: *Rig, exp: *const exl3e.Experts, layer: exl3e.Layer, x: u64, pick: u64, wts: u64, rows_n: usize) !struct { out: u64, y: u64 } { + const sb = try r.zeros(try exl3e.scratchBytes(layer, rows_n, SLOTS)); + const sc = try exl3e.scratch(sb, layer, rows_n, SLOTS); + const out = try r.zeros(rows_n * D * 4); + try exp.run(r.s, layer, x, D, .f32, pick, wts, rows_n, out, sc); + return .{ .out = out, .y = sc.y }; +} + +fn f64Check(r: *Rig, exp: *const exl3e.Experts, bits: []const u32, act: i32) !void { + const b = try buildLayer(r, bits, act); const want = try reference(r, &b, 128); for (rows) |m| try closeOut(r, b.bits, m, (try runRows(r, exp, &b, m)), want.want, want.tol); } +/// The engine's fp32 x: the same bf16-quantized values carried exactly in fp32, so the f64 reference applies unchanged. +fn f32xCheck(r: *Rig, exp: *const exl3e.Experts, bits: []const u32) !void { + const b = try buildLayer(r, bits, 1); + const fxs = try r.a.alloc(f32, CAP * D); + for (fxs, b.x) |*o, v| o.* = @floatCast(bf(v)); + const dfx = try r.dev(std.mem.sliceAsBytes(fxs)); + const want = try reference(r, &b, 128); + for (rows) |m| try closeOut(r, b.bits, m, (try runSliceF32(r, exp, b.layer, dfx, b.dpick, b.dwts, m)).out, want.want, want.tol); +} + /// fp32 outputs within the f64 reference's tolerance, with the exact FAIL line. fn closeOut(r: *Rig, bits: []const u32, m: usize, ptr: u64, want: []const f64, tol: []const f64) !void { const raw = try r.back(ptr, m * D * 4); @@ -223,6 +258,10 @@ fn closeOut(r: *Rig, bits: []const u32, m: usize, ptr: u64, want: []const f64, t var at: usize = 0; for (want[0 .. m * D], tol[0 .. m * D], 0..) |wv, t, i| { const g: f64 = @floatCast(std.mem.bytesToValue(f32, raw[4 * i ..][0..4])); + if (g != g) { + std.debug.print("FAIL exl3.experts.f64.mul1.{d} m={d}: NaN at element {d}\n", .{ bits[0], m, i }); + return error.TestFailed; + } const q = @abs(g - wv) / t; if (q > most) { most = q; @@ -239,7 +278,7 @@ fn closeOut(r: *Rig, bits: []const u32, m: usize, ptr: u64, want: []const f64, t /// A 300-row chunk against the f64 reference (the mixed-width layer). pub fn prompt(r: *Rig, exp: *const exl3e.Experts) !void { var bits: [E]u32 = .{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }; - const b = try buildLayer(r, &bits); + const b = try buildLayer(r, &bits, 0); // the GLM bf16 act mode's 300-row chunk (mode 1 is graded in checks/f32x) const want = try reference(r, &b, 300); const res = try runRows(r, exp, &b, 300); const raw = try r.back(res, 300 * D * 4); @@ -247,6 +286,10 @@ pub fn prompt(r: *Rig, exp: *const exl3e.Experts) !void { var at: usize = 0; for (want.want, want.tol, 0..) |wv, t, i| { const g: f64 = @floatCast(std.mem.bytesToValue(f32, raw[4 * i ..][0..4])); + if (g != g) { + std.debug.print("FAIL exl3.prompt.experts.f64: NaN at element {d}\n", .{i}); + return error.TestFailed; + } const q = @abs(g - wv) / t; if (q > most) { most = q; @@ -325,11 +368,13 @@ fn reference(r: *Rig, b: *const Built, rows_n: usize) !struct { want: []f64, tol for (0..p) |s0| { const ex: usize = @intCast(b.pick[s0]); const vg = zg[n * p + s0] * f16v(b.svh_g[ex * I + n]); - const gg = @min(bf16r(vg), @as(f64, LIMIT)); + const a1 = b.layer.act_mode == 1; + const gg = @min(if (a1) vg else bf16r(vg), @as(f64, LIMIT)); const vu = zu[n * p + s0] * f16v(b.svh_u[ex * I + n]); - const uu = @min(@max(bf16r(vu), -@as(f64, LIMIT)), @as(f64, LIMIT)); - const inner = bf16r(gg / (1.0 + std.math.exp(-gg))); - const act = bf16r(inner * uu); + const uu = @min(@max(if (a1) vu else bf16r(vu), -@as(f64, LIMIT)), @as(f64, LIMIT)); + const sil = gg / (1.0 + std.math.exp(-gg)); + const inner = if (a1) sil else bf16r(sil); + const act = if (a1) inner * uu else bf16r(inner * uu); v[n * p + s0] = act * f16v(b.suh_d[ex * I + n]); const egk = eg[n * p + s0] * @abs(f16v(b.svh_g[ex * I + n])) + (@abs(vg) + @abs(gg)) / 256; const euk = eu[n * p + s0] * @abs(f16v(b.svh_u[ex * I + n])) + (@abs(vu) + @abs(uu)) / 256; diff --git a/zig/tests/cuda/exl3_experts_tests.zig b/zig/tests/cuda/exl3_experts_tests.zig index 20710d303f..6985588875 100644 --- a/zig/tests/cuda/exl3_experts_tests.zig +++ b/zig/tests/cuda/exl3_experts_tests.zig @@ -53,7 +53,7 @@ pub fn run(r: *Rig, state: *Run, dir: []const u8) !void { /// A row's output bytes are its own: solo, with 3 rows, in 64-row batches, shared or not. fn rowsInvariant(r: *Rig, exp: *const exl3e.Experts) !void { var bits: [ref.E]u32 = .{ 6, 4, 5, 6, 4, 5, 6, 4, 5 }; - const b = try ref.buildLayer(r, &bits); + const b = try ref.buildLayer(r, &bits, 0); // the GLM bf16 act mode's rows-invariant bar const solo = try r.back(try ref.runRows(r, exp, &b, 1), ref.D * 4); // 3 rows: two more rows on the same picks as row 0 const p3 = try r.a.alloc(i32, 3 * ref.SLOTS); @@ -93,7 +93,7 @@ fn rowsInvariant(r: *Rig, exp: *const exl3e.Experts) !void { /// The same 300 rows as one call, split at 128 and at 37: a row's bytes never move. fn chunks(r: *Rig, exp: *const exl3e.Experts) !void { var bits: [ref.E]u32 = .{ 4, 5, 6, 4, 5, 6, 4, 5, 6 }; - const b = try ref.buildLayer(r, &bits); + const b = try ref.buildLayer(r, &bits, 0); // the GLM bf16 act mode's chunking bar const whole = try r.back(try ref.runRows(r, exp, &b, 300), 300 * ref.D * 4); var at: usize = 0; while (at < 300) : (at += 128) { diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index d1d5380ab4..70148f8296 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -95,6 +95,9 @@ const kolibri_tests = @import("kolibri_tests.zig"); const kolibri_layer_tests = @import("kolibri_layer_tests.zig"); const aot_fallback_tests = @import("aot_fallback_tests.zig"); const qwen_draft_frames = @import("qwen_draft_frames.zig"); +const minimax_m3_tests = @import("minimax_m3_tests.zig"); +const minimax_tp2_tests = @import("minimax_tp2_tests.zig"); +const minimax_graph_tests = @import("minimax_m3_graph_tests.zig"); const usage = \\usage: tf-cuda-test @@ -213,6 +216,41 @@ const usage = \\ kolibri-logits MODEL IDS OUT [PROMPTS] Kolibri-1's teacher-forced fp32 logits of 2,048 prompt rows; each prompt's last \\ glm-synth a synthetic GLM-5.3-Flash checkpoint (2 layers, 16 experts) and its dc.json for the lane checks \\ aot-fallback a launch's 1 from the runtime-parameter variant: the folded variant's bits (#1003) + \\ minimax-m3 MiniMax M3.1's MSA indexer (scores/topk/decode/combine) and swigluoai + \\ against the family's CPU references (synthetic); NB<=1024 and nsplit<=8 guards + \\ minimax-m3-tp2 MiniMax M3.1's TP2 weight plumbing on the real + \\ pack: each rank's Stage-2-plan shard uploads to device buffers through + \\ the mainline Uploader; the ranks' checkpoint bytes must tile the plan's + \\ minimax-m3-smoke MiniMax M3.1's launch graph on the 2-layer fixture (TP1): noise-poked + \\ fixture, 8 context rows + 2 greedy steps on device, the router and + \\ selection fingerprints, tokens + logits printed + \\ minimax-m3-token [ctx] [steps] + \\ the launch graph's token smoke: TP1 ("-") or TP2 (rank 0 serves the + \\ NCCL id; rank 1 first, rank 0 ~2 s later) over the pack or a locally built + \\ deterministic fixture (TP1 vs TP2 tokens must match); rank 0 prints the + \\ greedy tokens, wall times, RAM and device totals + \\ minimax-m3-grade [ctx] [steps] [layers] the Stage-8 grade: the device + \\ graph (TP1, the first [layers] REAL layers) against the f64 fork + \\ composition over the pack's own decoded weights: every row's residual + \\ and logits compared, greedy tokens must match EXACTLY + \\ minimax-m3-lanes [ctx] [steps] [layers] the Stage-8 boundary: a real + \\ open() Loaded handle serves token-in token-out through the lane surface + \\ (prefill/first/queue/read/verify/keep/release), graded token-for-token + \\ against the direct engine; draft/sample/second-stream/tree-keep refuse + \\ minimax-m3-lane-tp2 [ctx] [steps] the Stage-9 proof: + \\ openTP2's lane vtable — the SAME lane surface open() serves for TP1 — + \\ answering over the SPLIT real pack (rank 1 first over the NCCL id, + \\ rank 0 serves the id and prints the served tokens; rank 1 lockstep) + \\ minimax-m3-serve [steps<=32] [layers] [tokdir] Stages 9/10 + \\ front door: open()'s attached lane vtable answers one greedy + \\ completion — token ids in ("lcg" = the receipted bar stream; + \\ stream; a file = JSON {"tokens": [...]} / bare array / ids; + \\ Stage-10: "text:" inline or "textfile:" PATH — encoded/decoded + \\ through the pack's OWN tokenizer.json, the completion ALSO + \\ printed as readable text with encode/decode round-trip bars), + \\ token ids out, then exit; at the receipted fixture recipe the + \\ completion must be EXACT 183 149 14; draft/second-stream/sampled + \\ refuse truthfully \\ ; @@ -403,6 +441,14 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "kolibri-greedy")) return kolibri_tests.greedy(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), try arg(rest, 3)); if (std.mem.eql(u8, cmd, "kolibri-logits")) return kolibri_tests.logits(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), if (rest.len > 3) rest[3] else null); if (std.mem.eql(u8, cmd, "aot-fallback")) return aot_fallback_tests.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "minimax-m3")) return minimax_m3_tests.run(gpu); + if (std.mem.eql(u8, cmd, "minimax-m3-tp2")) return minimax_tp2_tests.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else ""); + if (std.mem.eql(u8, cmd, "minimax-m3-smoke")) return minimax_graph_tests.run(gpu); + if (std.mem.eql(u8, cmd, "minimax-m3-token")) return minimax_graph_tests.tokenRun(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else "", if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 128, if (rest.len > 4) try std.fmt.parseInt(usize, rest[4], 10) else 2); + if (std.mem.eql(u8, cmd, "minimax-m3-grade")) return minimax_graph_tests.gradeRun(gpu, if (rest.len > 0) rest[0] else "", if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else 32, if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 4, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 2); + if (std.mem.eql(u8, cmd, "minimax-m3-lanes")) return minimax_graph_tests.laneRun(gpu, if (rest.len > 0) rest[0] else "", if (rest.len > 1) try std.fmt.parseInt(usize, rest[1], 10) else 8, if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 3, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 0); + if (std.mem.eql(u8, cmd, "minimax-m3-lane-tp2")) return minimax_graph_tests.laneTp2Run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2), if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 32, if (rest.len > 4) try std.fmt.parseInt(usize, rest[4], 10) else 4); + if (std.mem.eql(u8, cmd, "minimax-m3-serve")) return minimax_graph_tests.serveRun(gpu, if (rest.len > 0) rest[0] else "", if (rest.len > 1) rest[1] else "", if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 3, if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 2, if (rest.len > 4) rest[4] else ""); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); return error.UnknownCommand; diff --git a/zig/tests/cuda/minimax_m3_graph_tests.zig b/zig/tests/cuda/minimax_m3_graph_tests.zig new file mode 100644 index 0000000000..46da14c86e --- /dev/null +++ b/zig/tests/cuda/minimax_m3_graph_tests.zig @@ -0,0 +1,1937 @@ +//! MiniMax M3.1's launch-graph smoke (Stage-6): the graph.zig Engine on device. +//! +//! tf-cuda-test minimax-m3-smoke +//! builds the family's 2-layer fixture in /tmp, boots the Engine over it (TP1), runs +//! 8 context rows + 2 greedy steps on device, then grades them EXACTLY against the fork's +//! f64 composition over the same poked fixture decoded on the host (the serve.zig-Ref +//! contract at the engine's own dims), and asserts the router/selection fingerprints. +//! +//! tf-cuda-test minimax-m3-token [ctx] [steps] +//! the real pack, TP1 (addr "-") or TP2 (rank 0 serves the NCCL id on host:port, the +//! exl3-tp2 precedent; rank 1 first, rank 0 ~2 s later — the Stage-5 runbook). Both +//! ranks read the same deterministic context (serve.promptOf's LCG stream), run every +//! owned layer's row, and rank 0 closes with the final norm + lm_head + greedy. Prints +//! per-row wall, the tokens, host RAM and device totals. "fixture" instead of a pack +//! dir generates the same 2-layer fixture on the calling node (rank 1 generates its +//! own copy; the bytes are deterministic), so TP2 handoffs are graded at fixture scale +//! against minimax-m3-smoke's TP1 tokens. + +const std = @import("std"); +const builtin = @import("builtin"); +const posix = std.posix; +const cuda = @import("cuda"); +const check = @import("check.zig"); +const m3 = @import("minimax_m3_vl_cuda"); + +const exl3_load = @import("core").exl3_load; +const tokenizer = @import("core").tokenizer; + +/// The fixture's dir (the smoke rebuilds it each run; nothing reads a stale one). +const fixture_path = "/tmp/tf-m3-graph-smoke"; +/// The context stream's seed: fixed so both ranks and every rerun read the same prompt. +const prompt_seed: u64 = 0x5EEDB00B; + +fn rebuildFixture(a: std.mem.Allocator, io: std.Io) !void { + var cwd = std.Io.Dir.cwd(); + cwd.deleteTree(io, fixture_path) catch {}; + cwd.createDirPath(io, fixture_path) catch return error.FixtureDir; + var fdir = try cwd.openDir(io, fixture_path, .{}); + defer fdir.close(io); + try m3.weights_mod.checkpointFixture(a, io, fdir, false); + try pokeNoise(a, io, fdir); +} + +const Lcg = m3.serve.Lcg; // the prompt stream's state type (serve.zig's own, now pub) + +// ---------------------------------------------------------------- the fixture smoke (TP1) + +pub fn run(gpu: check.Gpu) !void { + const gpa = gpu.gpa; + try rebuildFixture(gpa, gpu.io); + const t0 = check.now(gpu.io); + var eng = try m3.graph.Engine.build(gpa, gpu.io, gpu.d, fixture_path, 0, 1, .{}); + defer eng.deinit(); + const build_s = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + + var prompt: [8]u32 = undefined; + var l = Lcg{ .s = prompt_seed }; + m3.serve.promptOf(&l, prompt[0..]); + var out: [2]u32 = undefined; + + const t1 = check.now(gpu.io); + for (prompt) |tok| try eng.put(tok); + const prefill_s = @as(f64, @floatFromInt(check.now(gpu.io) - t1)) / 1e9; + const t2 = check.now(gpu.io); + try eng.greedy(out.len, out[0..]); + const greedy_s = @as(f64, @floatFromInt(check.now(gpu.io) - t2)) / 1e9; + + const c = &eng.c; + // the weight discriminator: the host-decoded fixture's own k/v matrices (decodeGroup, the + // RefF64 path's dequantize) — the device v GEMM's output being ~100x k's was traced either + // here (huge decoded weights) or into the upload/kernels (clean decoded weights) +wprobe: { + var scw = exl3_load.scan(gpa, gpu.io, fixture_path) catch break :wprobe; + defer scw.deinit(); + const vi = scw.find("language_model.model.layers.0.self_attn.v_proj") orelse break :wprobe; + const ki = scw.find("language_model.model.layers.0.self_attn.k_proj") orelse break :wprobe; + const qi = scw.find("language_model.model.layers.0.self_attn.q_proj") orelse break :wprobe; + const hi = scw.find("language_model.lm_head") orelse break :wprobe; // the token smoke's fp32 head + inline for (.{ "v", "k", "q", "head" }, .{ vi, ki, qi, hi }) |tag, gi| { + const g = scw.groups[gi]; + const parts = scw.parts(gpa, gi) catch break :wprobe; // scan-owned; decodeGroup never frees these + const w = decodeGroup(gpa, &scw, gi) catch break :wprobe; + defer gpa.free(w); + var s2: f64 = 0; + var smax: f64 = 0; + var sl1: usize = 0; + for (w) |v| { + s2 += v * v; + smax = @max(smax, @abs(v)); + if (@abs(v) > 1) sl1 += 1; + } + const rms = @sqrt(s2 / @as(f64, @floatFromInt(w.len))); + var sumin: f32 = std.math.inf(f32); + var sumax: f32 = -std.math.inf(f32); + var svmin: f32 = std.math.inf(f32); + var svmax: f32 = -std.math.inf(f32); + const su = std.mem.bytesAsSlice(u16, parts.suh.bytes); + const sv = std.mem.bytesAsSlice(u16, parts.svh.bytes); + for (su) |s| { + const f = @as(f32, @floatCast(@as(f16, @bitCast(s)))); + sumin = @min(sumin, f); + sumax = @max(sumax, f); + } + for (sv) |s| { + const f = @as(f32, @floatCast(@as(f16, @bitCast(s)))); + svmin = @min(svmin, f); + svmax = @max(svmax, f); + } + std.debug.print("DBGW {s}_proj group: k={d} n={d} halves={d} codebook={s} trellis={d}B suh={d}B({d}) suh[{d:.4},{d:.4}] svh={d}B({d}) svh[{d:.4},{d:.4}] | rms={d:.4} max={d:.4} frac>1={d:.4}\n", .{ tag, g.k, g.n, g.halves, @tagName(g.codebook), parts.trellis.bytes.len, parts.suh.bytes.len, su.len, sumin, sumax, parts.svh.bytes.len, sv.len, svmin, svmax, rms, smax, @as(f64, @floatFromInt(sl1)) / @as(f64, @floatFromInt(w.len)) }); + } + } + std.debug.print("DBG layers {d}:", .{eng.layers.len}); + for (eng.layers) |*ly| std.debug.print(" lno={d} mlp={s}", .{ ly.lno, @tagName(ly.mlp) }); + std.debug.print("; picks:", .{}); + for (eng.picks) |p| std.debug.print(" {}", .{p}); + std.debug.print("; sel_e:", .{}); + for (eng.sel_e) |v| std.debug.print(" {d}", .{v}); + std.debug.print("; wsum={d:.4} fired={d}\n", .{ eng.last.wsum, eng.last.fired }); + // the NaN discriminator: a route with NaN logits fires nobody (sigmoid NaN fails strict >), + // so fired==0 here means e.x or the elogits went NaN (or the wave never ran — sel_seen tells) + var xnan = false; + for (eng.x) |v| { + if (std.math.isNan(v)) xnan = true; + } + var elnan = false; + for (eng.elogits) |v| { + if (std.math.isNan(v)) elnan = true; + } + var xmax: f32 = 0; + for (eng.x) |v| { + const m = @abs(v); + if (m > xmax) xmax = m; + } + std.debug.print("DBG2 xnan={} elnan={} xmax={d:.5} logits_nan:", .{ xnan, elnan, xmax }); + var lnan = false; + for (eng.logits) |v| { + if (std.math.isNan(v)) lnan = true; + } + std.debug.print(" {} sel_seen[0]:", .{lnan}); + for (eng.sel_seen[0..@min(4, eng.sel_seen.len)]) |v| { + std.debug.print(" {d}", .{v}); + } + std.debug.print("\n", .{}); + // the fingerprints: the router fires exactly topk experts, the weights sum to routed_scale + // (normalize-then-scale), and the selection carries valid blocks for the kv heads + var dbg_fail = false; // debug build: keep going to the f64 grade after a fingerprint miss + if (eng.last.fired != c.topk) { + std.debug.print("FAIL minimax-m3-smoke: {d} experts fired, expected topk {d}\n", .{ eng.last.fired, c.topk }); + dbg_fail = true; + } + if (@abs(eng.last.wsum - c.routed_scale) > 1e-2) { + std.debug.print("FAIL minimax-m3-smoke: routing weights sum {d}, expected routed_scale {d}\n", .{ eng.last.wsum, c.routed_scale }); + dbg_fail = true; + } + if (eng.last.sel_valid < 1) { + std.debug.print("FAIL minimax-m3-smoke: the last kv head's selection is empty\n", .{}); + return error.TestFailed; + } + var best: usize = 0; + var bv = -std.math.inf(f32); + var second = -std.math.inf(f32); + for (eng.logits, 0..) |v, i| { + if (v > bv) { + second = bv; + bv = v; + best = i; + } else if (v > second) second = v; + } + const margin = bv - second; + std.debug.print("minimax-m3-smoke: fixture build {d:.2}s, {d}-row prefill {d:.3}s, greedy {d:.3}s; fired {d}, wsum {d:.4}, sel_valid {d}\n", .{ build_s, prompt.len, prefill_s, greedy_s, eng.last.fired, eng.last.wsum, eng.last.sel_valid }); + std.debug.print("minimax-m3-smoke: logits argmax {d} ({d:.5}), second {d:.5}, margin {d:.5}\n", .{ best, bv, second, margin }); + if (margin < 0) { + std.debug.print("FAIL minimax-m3-smoke: the argmax is a tie at the fp32 head\n", .{}); + return error.TestFailed; + } + // the grade: the same poked fixture bytes decoded on the host (core.exl3_format, the fork's + // trellis decode), the fork's f64 composition at the engine's own runtime dims + const t3 = check.now(gpu.io); + var fw = try loadFixtureW(gpa, gpu.io, fixture_path, c.*); + defer fw.deinit(); + var ref = try RefF64.init(gpa, c.*, prompt.len + out.len + 1, &fw); + defer ref.deinit(); + var ref_out: [2]u32 = undefined; + var lr = Lcg{ .s = prompt_seed }; + m3.serve.promptOf(&lr, prompt[0..]); + for (prompt) |tok| try ref.put(tok); + try ref.greedy(out.len, ref_out[0..]); + const ref_s = @as(f64, @floatFromInt(check.now(gpu.io) - t3)) / 1e9; + var maxd: f64 = 0; + var lnn: usize = 0; + var lin: usize = 0; + for (eng.logits, 0..) |v, i| { + if (std.math.isNan(v)) lnn += 1; + if (std.math.isInf(v)) lin += 1; + const d = @as(f64, v) - ref.logits[i]; + if (!std.math.isNan(d) and @abs(d) > maxd) maxd = @abs(d); + } + std.debug.print("minimax-m3-smoke: reference (f64, decoded fixture) {d:.1}s; logits maxd {d:.5} NaN {d} inf {d}\n", .{ ref_s, maxd, lnn, lin }); + try gradeTokensCatch(ref_out[0..], out[0..]); + std.debug.print("RESULT minimax-m3-smoke: TP1 fixture tokens {d} {d} == reference {d} {d} (EXACT); ctx rows {d}; fired {d} wsum {d:.4} sel_valid {d}\n", .{ out[0], out[1], ref_out[0], ref_out[1], prompt.len, eng.last.fired, eng.last.wsum, eng.last.sel_valid }); +} + +// ---------------------------------------------------------------- the real-pack token smoke + +pub fn tokenRun(gpu: check.Gpu, rank_arg: []const u8, addr_arg: []const u8, dir_in: []const u8, ctx_len: usize, steps: usize) !void { + var dir = dir_in; + if (dir.len == 0) { + std.debug.print("usage: tf-cuda-test minimax-m3-token [ctx] [steps]\n", .{}); + return error.BadArguments; + } + if (std.mem.eql(u8, dir, "fixture")) dir = try makeFixture(gpu); + const rank = std.fmt.parseInt(u8, rank_arg, 10) catch { + std.debug.print("usage: tf-cuda-test minimax-m3-token [ctx] [steps]\n", .{}); + return error.BadArguments; + }; + if (rank > 1 or addr_arg.len == 0 or steps == 0 or steps > 32 or ctx_len == 0 or ctx_len > m3.graph.keys_cap) { + return error.BadArguments; // loud refusals: no silent clamping of a context the pool holds + } + const world: u32 = if (std.mem.eql(u8, addr_arg, "-")) 1 else 2; + const ranks: u32 = @intCast(rank); + + // -- the TP2 boot exactly the exl3-tp2 precedent (rank 1 first, rank 0 serves the id) -- + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + var id: cuda.nccl.UniqueId = undefined; + const id_bytes = std.mem.asBytes(&id); + if (world == 2) { + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try serveId(addr_arg, id_bytes); + } else { + try receiveId(gpu, addr_arg, id_bytes); + } + } + var comm: cuda.nccl.Comm = null; + if (world == 2) { + try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(world), id, @intCast(rank)), "ncclCommInitRank"); + } + defer { + if (world == 2) _ = lib.api.ncclCommDestroy(comm); + } + + const device_before = cuda.usage(false).device; + const ram_before = hostFree(gpu); + const t0 = check.now(gpu.io); + var eng = try m3.graph.Engine.build(gpu.gpa, gpu.io, gpu.d, dir, ranks, world, .{ + .lib = if (world == 2) &lib else null, + .comm = if (world == 2) comm else null, + }); + defer eng.deinit(); + const build_s = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + const device_after = cuda.usage(false).device; + const ram_after = hostFree(gpu); + const gib = struct { + fn f(v: u64) f64 { + return @as(f64, @floatFromInt(v)) / (1024.0 * 1024.0 * 1024.0); + } + }.f; + + // the deterministic context: both ranks derive it identically (rank 1's tok arg is ignored) + const prompt = try gpu.gpa.alloc(u32, ctx_len); + defer gpu.gpa.free(prompt); + var l = Lcg{ .s = prompt_seed }; + m3.serve.promptOf(&l, prompt); + const out = try gpu.gpa.alloc(u32, steps); + defer gpu.gpa.free(out); + + const t1 = check.now(gpu.io); + for (prompt) |tok| try eng.put(if (rank == 0) tok else 0); + const prefill_s = @as(f64, @floatFromInt(check.now(gpu.io) - t1)) / 1e9; + // greedy steps: rank 0 argues from its logits; rank 1 keeps the row lockstep + const t2 = check.now(gpu.io); + if (rank == 0) { + try eng.greedy(steps, out); + } else { + for (0..steps - 1) |_| try eng.put(0); // greedy()'s feedback puts, kept in the lockstep + } + const decode_s = @as(f64, @floatFromInt(check.now(gpu.io) - t2)) / 1e9; + var best: usize = 0; + var bv = -std.math.inf(f32); + var second = -std.math.inf(f32); + if (rank == 0) { + for (eng.logits, 0..) |v, i| { + if (v > bv) { + second = bv; + bv = v; + best = i; + } else if (v > second) second = v; + } + } else { + bv = 0; + second = 0; + } + // additive Stage-7 receipt: explicit logits health so a NaN row cannot masquerade as an + // argmax (6C's smoke print hardening, carried over to the token harness) + var lnan: usize = 0; + var linf: usize = 0; + var lmx: f32 = -std.math.inf(f32); + var lsum: f64 = 0; + if (rank == 0) { + for (eng.logits) |v| { + if (v != v) lnan += 1 else { + if (v > lmx) lmx = v; + if (v == std.math.inf(f32)) linf += 1 else lsum += v; + } + } + } + std.debug.print("M3-TOKEN rank {d} world {d}: build {d:.1}s, prefill(x{d}) {d:.2}s, greedy(x{d}) {d:.2}s\n", .{ rank, world, build_s, ctx_len, prefill_s, steps, decode_s }); + std.debug.print("M3-TOKEN rank {d}: device buffers {d:.2} GiB; host MemAvailable {d:.2} -> {d:.2} GiB; trace fired {d} wsum {d:.4} sel_valid {d}\n", .{ rank, gib(device_after - device_before), gib(ram_before), gib(ram_after), eng.last.fired, eng.last.wsum, eng.last.sel_valid }); + if (rank == 0) { + if (steps > 1 and decode_s > 0) { + const tps = @as(f64, @floatFromInt(steps - 1)) / decode_s; + std.debug.print("M3-TOKEN rank 0: decode tok/s {d:.2} ({d:.3}s for {d} rows)", .{ tps, decode_s, steps - 1 }); + } + std.debug.print("\nM3-TOKEN rank 0: logits rows {d} NaN {d} inf {d} max {d:.6} mean {d:.6} ({d:.5} over the finite rows)\n", .{ eng.logits.len, lnan, linf, lmx, lsum / @as(f64, @floatFromInt(eng.logits.len)), @as(f64, @floatFromInt(eng.logits.len - lnan)) }); + std.debug.print("M3-TOKEN rank 0: logits argmax {d} ({d:.5}), second {d:.5}\n", .{ best, bv, second }); + std.debug.print("RESULT minimax-m3-token: prompt rows {d} then greedy", .{ctx_len}); + for (out) |tok| std.debug.print(" {d}", .{tok}); + std.debug.print("\n", .{}); + } else { + std.debug.print("RESULT minimax-m3-token: rank 1 lockstep done over dir {s}\n", .{dir}); + } +} + +/// The fixture's path, built on the calling node (rank 1 builds its own deterministic copy). +fn makeFixture(gpu: check.Gpu) ![]const u8 { + try rebuildFixture(gpu.gpa, gpu.io); + return fixture_path; +} + +// --------------------------------------------------------------- the fixture's noise pass + +/// The fixture ships all-zero tensor data; this pokes every tensor with deterministic noise so +/// the smoke's paths cannot alias each other unnoticed (all-zero rows make wrong offsets equal). +/// The trellis words take random bits (the mul1 codebook decodes any bit pattern to a finite +/// value), the fp16 scales land in [0.75, 1.25], the norms near 1.0, and the passthrough sizes +/// match their rows' shapes (checkpointFixture wrote them; the safetensors headers say where). +const header_limit: usize = 1 << 26; // the fixture's headers are far under 64 MiB + +fn pokeNoise(gpa: std.mem.Allocator, io: std.Io, dir: std.Io.Dir) !void { + var prng = std.Random.DefaultPrng.init(0x9E3779B9); + const rnd = prng.random(); + for ([_][]const u8{ "model-00002.safetensors", "model.safetensors" }) |file_name| { + const raw = dir.readFileAlloc(io, file_name, gpa, .limited(4 << 30)) catch continue; // a file the fixture did not write: skip + defer gpa.free(raw); + if (raw.len < 8) continue; + const hlen = std.mem.readInt(u64, raw[0..8], .little); + if (hlen < 2 or hlen > header_limit or 8 + hlen > raw.len) return error.BadFixture; + const parsed = std.json.parseFromSlice(std.json.Value, gpa, raw[8 .. 8 + hlen], .{}) catch return error.BadFixture; + defer parsed.deinit(); + if (parsed.value != .object) return error.BadFixture; + var it = parsed.value.object.iterator(); + while (it.next()) |kv| { + const name = kv.key_ptr.*; + if (std.mem.eql(u8, name, "__metadata__")) continue; + if (kv.value_ptr.* != .object) continue; + const ts = kv.value_ptr.object; + const dt = if (ts.get("dtype")) |v| (if (v == .string) v.string else "") else ""; + const offs = ts.get("data_offsets") orelse continue; + if (offs != .array or offs.array.items.len != 2) continue; + const a = offs.array.items[0].integer; + const b = offs.array.items[1].integer; + if (a < 0 or b < a) continue; + const start: usize = @intCast(a); + const end: usize = @intCast(b); + const base: usize = 8 + hlen; + if (base + end > raw.len) continue; + const view = raw[base + start .. base + end]; + if (std.mem.endsWith(u8, name, ".trellis")) { + for (0..view.len / 4) |i| std.mem.writeInt(u32, view[i * 4 ..][0..4], rnd.int(u32), .little); + } else if (std.mem.endsWith(u8, name, ".suh")) { + // The K-tile scale carries the matrix's magnitude: the codebook's own values live + // at O(1), so the [0.75, 1.25] suh/svh poke decoded every quantized matrix to + // ~1.0 rms — ~170x the serve.zig draw scale (projections 0.006) — and the raw + // attention V (the one projection with no norm to rescale it) then pushed the + // O-GEMM's fp16 outputs past saturation: combine max|a| ~947, o16 infs, the + // normRow NaN cascade, the router dead. Scale suh to the draws' sigma; svh keeps + // its shape. The head stays wide: its logits are fp32 and its margins bar ties. + const sigma: f32 = if (std.mem.indexOf(u8, name, "lm_head") != null) 0.10 else 0.006; + for (0..view.len / 2) |i| putF16(view[i * 2 ..][0..2], sigma * (0.75 + 0.5 * rnd.float(f32))); + } else if (std.mem.endsWith(u8, name, ".svh")) { + for (0..view.len / 2) |i| putF16(view[i * 2 ..][0..2], 0.75 + 0.5 * rnd.float(f32)); + } else if (std.mem.endsWith(u8, name, ".bias")) { + // the codebook marker and any per-group bias bytes: they carry the group's contract + } else if (std.mem.eql(u8, dt, "BF16")) { + const n = view.len / 2; + // norms are gemma-style delta weights ((1.0 + weight) per serve.normRow; the + // serve.zig fixture also stores deltas) — absolute ~1.0 values made every norm + // ~2x and pushed the attention output past the fp16 range in the engine AND the + // f64 reference alike. + const near_one: bool = std.mem.endsWith(u8, name, "norm.weight"); + for (0..n) |i| { + const v: f32 = if (near_one) 0.05 * (rnd.float(f32) - 0.5) else 0.04 * (rnd.float(f32) - 0.5); + putBF16(view[i * 2 ..][0..2], v); + } + } else if (std.mem.eql(u8, dt, "F16")) { + for (0..view.len / 2) |i| putF16(view[i * 2 ..][0..2], 0.01 * (rnd.float(f32) - 0.5)); + } else if (std.mem.eql(u8, dt, "F32")) { + const scale: f32 = if (std.mem.endsWith(u8, name, "e_score_correction_bias")) 0.02 else 0.05; + for (0..view.len / 4) |i| { + const v: f32 = scale * (rnd.float(f32) - 0.5); + std.mem.writeInt(u32, view[i * 4 ..][0..4], @bitCast(v), .little); + } + } + } + try dir.writeFile(io, .{ .sub_path = file_name, .data = raw }); + } +} + +fn putF16(dst: []u8, v: f32) void { + const h: f16 = @floatCast(v); + std.mem.writeInt(u16, dst[0..2], @bitCast(h), .little); +} + +fn putBF16(dst: []u8, v: f32) void { + const bits: u32 = @bitCast(v); + std.mem.writeInt(u16, dst[0..2], @intCast(bits >> 16), .little); +} + +// --------------------------------------------------------------- the tp2 harness's own fns + +/// Free MemAvailable from /proc/meminfo (0 elsewhere — the receipt prints GB10 nodes' figures). +fn hostFree(gpu: check.Gpu) u64 { + var f = std.Io.Dir.openFileAbsolute(gpu.io, "/proc/meminfo", .{}) catch return 0; + defer f.close(gpu.io); + var buf: [8192]u8 = undefined; + const n = f.readPositionalAll(gpu.io, &buf, 0) catch return 0; + const line = std.mem.indexOf(u8, buf[0..n], "MemAvailable:") orelse return 0; + var it = std.mem.tokenizeScalar(u8, buf[line + "MemAvailable:".len .. n], ' '); + const kb = it.next() orelse return 0; + return (std.fmt.parseInt(u64, kb, 10) catch 0) * 1024; // /proc/meminfo reports kB, not bytes +} + +// --- the NCCL id exchange, verbatim the exl3-tp2 precedent (rank 0 serves on host:port, rank 1 reads) --- + +fn serveId(addr_text: []const u8, bytes: []const u8) !void { + if (builtin.os.tag == .windows) return error.SocketFailed; + const at = std.mem.lastIndexOfScalar(u8, addr_text, ':') orelse return error.BadAddress; + const ip = try std.Io.net.IpAddress.parseIp4(addr_text[0..at], try std.fmt.parseInt(u16, addr_text[at + 1 ..], 10)); + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&ip, &storage); + const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const fd: posix.socket_t = @intCast(rc); + defer _ = posix.system.close(fd); + const one: c_int = 1; + try posix.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, std.mem.asBytes(&one)); + switch (posix.errno(posix.system.bind(fd, &storage.any, len))) { + .SUCCESS => {}, + .ADDRINUSE => return error.AddressInUse, + else => return error.BindFailed, + } + if (posix.errno(posix.system.listen(fd, 1)) != .SUCCESS) return error.ListenFailed; + var waited: u64 = 0; + while (true) { + var fds = [_]posix.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const ready = posix.poll(&fds, 500) catch return error.AcceptFailed; + if (ready > 0) break; + waited += 500; + if (waited > 70_000) return error.PeerTimeout; + } + var peer: std.Io.Threaded.PosixAddress = undefined; + var peer_len: posix.socklen_t = @sizeOf(std.Io.Threaded.PosixAddress); + const rc2 = posix.system.accept(fd, &peer.any, &peer_len); + if (posix.errno(rc2) != .SUCCESS) return error.AcceptFailed; + const cfd: posix.socket_t = @intCast(rc2); + defer _ = posix.system.close(cfd); + var off: usize = 0; + while (off < bytes.len) { + const n = posix.system.write(cfd, bytes.ptr + off, bytes.len - off); + if (posix.errno(n) != .SUCCESS) return error.WriteFailed; + off += @intCast(n); + } +} + +fn receiveId(gpu: check.Gpu, addr_text: []const u8, bytes: []u8) !void { + if (builtin.os.tag == .windows) return error.SocketFailed; + const at = std.mem.lastIndexOfScalar(u8, addr_text, ':') orelse return error.BadAddress; + const ip = try std.Io.net.IpAddress.parseIp4(addr_text[0..at], try std.fmt.parseInt(u16, addr_text[at + 1 ..], 10)); + var fd: posix.socket_t = -1; + defer { + if (fd >= 0) _ = posix.system.close(fd); + } + while (fd < 0) { + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&ip, &storage); + const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const c: posix.socket_t = @intCast(rc); + if (posix.errno(posix.system.connect(c, &storage.any, len)) == .SUCCESS) { + fd = c; + } else { + _ = posix.system.close(c); + std.Io.sleep(gpu.io, .fromMilliseconds(200), .awake) catch {}; + } + } + var off: usize = 0; + var read_waited: u64 = 0; + while (off < bytes.len) { + var fds = [_]posix.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const ready = posix.poll(&fds, 500) catch return error.ReadFailed; + if (ready == 0) { + read_waited += 500; + if (read_waited > 60_000) return error.PeerTimeout; + continue; + } + const n = posix.system.read(fd, bytes.ptr + off, bytes.len - off); + switch (posix.errno(n)) { + .SUCCESS => { + if (n == 0) return error.PeerClosed; + off += @intCast(n); + }, + .INTR => {}, + else => return error.ReadFailed, + } + } +} +// ------------------------------------------------------- the fixture's f64 reference (the grade) +// +// The graded bar: the SAME poked fixture bytes the Engine's device graph serves, decoded on the +// host (`core.exl3_format.dequantize` — the trellis decode the fork's own tooling mirrors) and run +// through the fork's f64 composition, contract by contract with serve.zig's `Ref`: gemma norms in +// f64 with one bf16-rounded weight row each, per-head qk norm + partial neox rope (theta 1e7, +// 64 of 128 dims) with ONE fp16 rounding at every fp16 store (the linear outputs' h16, the cache +// and plane rows, the routing weights), the block-sparse f64 attention of the indexer's +// boost-override top-16-of-block selection, the noaux sigmoid+bias router (select on sigmoid+bias, +// weights = unbiased sigmoid normalized over the picks, * routed_scale, each .half()), the shared +// expert at weight 1, a leading dense wave, the final gemma norm, the head. Token equality against +// the device graph is the bar; the logits delta is printed for the receipt. + +const exl3_fmt = @import("core").exl3_format; + +fn bf16At(raw: []const u8, i: usize) f32 { + const bits: u32 = @as(u32, std.mem.readInt(u16, raw[i * 2 ..][0..2], .little)) << 16; + return @bitCast(bits); +} + +fn f16At(raw: []const u8, i: usize) f32 { + const h: f16 = @bitCast(std.mem.readInt(u16, raw[i * 2 ..][0..2], .little)); + return @floatCast(h); +} + +const FW_Mlp = struct { gate: []f32, up: []f32, down: []f32 }; // decoded [k, n] row-major + +const FW_MoE = struct { + router: []f32, // [E, hidden] fp32 passthrough + bias: []f32, // [E] fp32 + shared: FW_Mlp, + experts: []FW_Mlp, +}; + +/// One layer's decoded weights, named exactly as the plan's spellings (weights.zig's tests). +const FW_Layer = struct { + in_norm: []f32, // [hidden], bf16 + post_norm: []f32, + q_norm: []f32, // [head_dim], bf16 + k_norm: []f32, + iqn: []f32, // [i_dim], bf16 + ikn: []f32, + q: []f32, // decoded [k=hidden, n=qk_rows] + k: []f32, // [hidden, kv_rows] + v: []f32, + o: []f32, // [qk_rows, hidden] + iq: []f32, // [i_heads*i_dim, hidden], fp16 passthrough + ik: []f32, // [i_dim, hidden], fp16 passthrough + dense: ?FW_Mlp = null, + moe: ?FW_MoE = null, +}; + +const FW = struct { + gpa: std.mem.Allocator, + c: m3.serve.Config, + embed: []f32, // [vocab, hidden], bf16 + final_norm: []f32, // [hidden], bf16 + head: []f32, // decoded [k=hidden, n=vocab] + layers: []FW_Layer, + scratch: []u8, // the dequantize scratch (2 * k * n f64 words, reused per group) + + fn deinit(f: *FW) void { + const gpa = f.gpa; + for (f.layers) |*l| { + gpa.free(l.in_norm); + gpa.free(l.post_norm); + gpa.free(l.q_norm); + gpa.free(l.k_norm); + gpa.free(l.iqn); + gpa.free(l.ikn); + gpa.free(l.q); + gpa.free(l.k); + gpa.free(l.v); + gpa.free(l.o); + gpa.free(l.iq); + gpa.free(l.ik); + if (l.dense) |m| { + gpa.free(m.gate); + gpa.free(m.up); + gpa.free(m.down); + } + if (l.moe) |m| { + gpa.free(m.router); + gpa.free(m.bias); + gpa.free(m.shared.gate); + gpa.free(m.shared.up); + gpa.free(m.shared.down); + for (m.experts) |e| { + gpa.free(e.gate); + gpa.free(e.up); + gpa.free(e.down); + } + gpa.free(m.experts); + } + } + gpa.free(f.layers); + gpa.free(f.embed); + gpa.free(f.final_norm); + gpa.free(f.head); + gpa.free(f.scratch); + } +}; + +fn decodedName(buf: []u8, comptime fmt: []const u8, args: anytype) []const u8 { + return std.fmt.bufPrint(buf, fmt, args) catch unreachable; +} + +/// One trellis group decoded to [k, n] f32 (the host twin of the device graph's own decode). +fn decodeGroup(gpa: std.mem.Allocator, scan: *exl3_load.Scan, gi: usize) ![]f32 { + const g = scan.groups[gi]; + const parts = try scan.parts(gpa, gi); + const k = g.k; + const n = g.n; + const tiles_k = k / 16; + const tiles_n = n / 16; + const out64 = try gpa.alloc(f64, @as(usize, k) * n); + defer gpa.free(out64); + const scratch = try gpa.alloc(u8, 2 * @as(usize, k) * n * 8); + defer gpa.free(scratch); + const trellis = try gpa.alloc(i16, parts.trellis.bytes.len / 2); // read little-endian (the bytes may not be i16-aligned) + defer gpa.free(trellis); + for (trellis, 0..) |*tw, i| tw.* = @bitCast(std.mem.readInt(u16, parts.trellis.bytes[i * 2 ..][0..2], .little)); + const suh = std.mem.bytesAsSlice(u16, parts.suh.bytes); + const svh = std.mem.bytesAsSlice(u16, parts.svh.bytes); + const suh_f16 = try gpa.alloc(f16, suh.len); + defer gpa.free(suh_f16); + const svh_f16 = try gpa.alloc(f16, svh.len); + defer gpa.free(svh_f16); + for (suh, 0..) |b, i| suh_f16[i] = @bitCast(b); + for (svh, 0..) |b, i| svh_f16[i] = @bitCast(b); + exl3_fmt.dequantize( + .{ .halves = g.halves }, + @enumFromInt(@intFromEnum(g.codebook)), + trellis, + tiles_k, + tiles_n, + suh_f16, + svh_f16, + out64, + @as([*]f64, @ptrCast(@alignCast(scratch.ptr)))[0 .. 2 * @as(usize, k) * n], + ); + const out = try gpa.alloc(f32, @as(usize, k) * n); + for (out64, 0..) |v, i| out[i] = @floatCast(v); + return out; +} + +/// The poked fixture's decoded weights (the same bytes the Engine's device graph serves). +fn loadFixtureW(gpa: std.mem.Allocator, io: std.Io, dir_path: []const u8, c: m3.serve.Config) !FW { + const L: usize = c.run; + var scan = try exl3_load.scan(gpa, io, dir_path); + defer scan.deinit(); + + var w: FW = .{ + .gpa = gpa, + .c = c, + .embed = &.{}, + .final_norm = &.{}, + .head = &.{}, + .layers = &.{}, + .scratch = &.{}, + }; + errdefer w.deinit(); + w.layers = try gpa.alloc(FW_Layer, L); + for (w.layers) |*l| l.* = .{ + .in_norm = &.{}, .post_norm = &.{}, .q_norm = &.{}, .k_norm = &.{}, + .iqn = &.{}, .ikn = &.{}, .q = &.{}, .k = &.{}, .v = &.{}, .o = &.{}, + .iq = &.{}, .ik = &.{}, + }; + + var nb: [256]u8 = undefined; + // passthrough helpers over the same checkpoint files + const bf16Row = struct { + fn f(sc: *exl3_load.Scan, a: std.mem.Allocator, name: []const u8, cols: usize) ![]f32 { + const t = try sc.ck.get(name); + if (!t.is(.bf16, &.{cols})) return error.UnexpectedTensor; + const out = try a.alloc(f32, cols); + for (0..cols) |i| out[i] = bf16At(t.bytes, i); + return out; + } + }.f; + const f16Mat = struct { + fn f(sc: *exl3_load.Scan, a: std.mem.Allocator, name: []const u8, rows: usize, cols: usize) ![]f32 { + const t = try sc.ck.get(name); + if (!t.is(.f16, &.{ rows, cols })) return error.UnexpectedTensor; + const out = try a.alloc(f32, rows * cols); + for (0..rows * cols) |i| out[i] = f16At(t.bytes, i); + return out; + } + }.f; + + const t = try scan.ck.get("language_model.model.embed_tokens.weight"); + if (!t.is(.bf16, &.{ c.vocab, c.hidden })) return error.UnexpectedTensor; + w.embed = try gpa.alloc(f32, @as(usize, c.vocab) * c.hidden); + for (w.embed, 0..) |*v, i| v.* = bf16At(t.bytes, i); + const fin = decodedName(&nb, "language_model.model.norm.weight", .{}); + w.final_norm = try bf16Row(&scan, gpa, fin, c.hidden); + + for (0..L) |l| { + const lay = &w.layers[l]; + lay.in_norm = try bf16Row(&scan, gpa, decodedName(&nb, "language_model.model.layers.{d}.input_layernorm.weight", .{l}), c.hidden); + lay.post_norm = try bf16Row(&scan, gpa, decodedName(&nb, "language_model.model.layers.{d}.post_attention_layernorm.weight", .{l}), c.hidden); + lay.q_norm = try bf16Row(&scan, gpa, decodedName(&nb, "language_model.model.layers.{d}.self_attn.q_norm.weight", .{l}), c.head_dim); + lay.k_norm = try bf16Row(&scan, gpa, decodedName(&nb, "language_model.model.layers.{d}.self_attn.k_norm.weight", .{l}), c.head_dim); + lay.iqn = try bf16Row(&scan, gpa, decodedName(&nb, "language_model.model.layers.{d}.self_attn.index_q_norm.weight", .{l}), c.i_dim); + lay.ikn = try bf16Row(&scan, gpa, decodedName(&nb, "language_model.model.layers.{d}.self_attn.index_k_norm.weight", .{l}), c.i_dim); + const irows: usize = c.i_heads * c.i_dim; + lay.iq = try f16Mat(&scan, gpa, decodedName(&nb, "language_model.model.layers.{d}.self_attn.index_q_proj.weight", .{l}), irows, c.hidden); + lay.ik = try f16Mat(&scan, gpa, decodedName(&nb, "language_model.model.layers.{d}.self_attn.index_k_proj.weight", .{l}), c.i_dim, c.hidden); + lay.q = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.self_attn.q_proj", .{l})).?); + lay.k = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.self_attn.k_proj", .{l})).?); + lay.v = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.self_attn.v_proj", .{l})).?); + lay.o = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.self_attn.o_proj", .{l})).?); + if (l < c.dense_layers) { + lay.dense = .{ + .gate = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.mlp.gate_proj", .{l})).?), + .up = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.mlp.up_proj", .{l})).?), + .down = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.mlp.down_proj", .{l})).?), + }; + } else { + var m: FW_MoE = .{ .router = &.{}, .bias = &.{}, .shared = .{ .gate = &.{}, .up = &.{}, .down = &.{} }, .experts = &.{} }; + const rt = try scan.ck.get(decodedName(&nb, "language_model.model.layers.{d}.block_sparse_moe.gate.weight", .{l})); + if (!rt.is(.f32, &.{ c.experts, c.hidden })) return error.UnexpectedTensor; + m.router = try gpa.alloc(f32, @as(usize, c.experts) * c.hidden); + @memcpy(m.router, std.mem.bytesAsSlice(f32, rt.bytes)); + const bt = try scan.ck.get(decodedName(&nb, "language_model.model.layers.{d}.block_sparse_moe.e_score_correction_bias", .{l})); + if (!bt.is(.f32, &.{c.experts})) return error.UnexpectedTensor; + m.bias = try gpa.alloc(f32, c.experts); + @memcpy(m.bias, std.mem.bytesAsSlice(f32, bt.bytes)); + m.shared = .{ + .gate = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.block_sparse_moe.shared_experts.gate_proj", .{l})).?), + .up = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.block_sparse_moe.shared_experts.up_proj", .{l})).?), + .down = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.block_sparse_moe.shared_experts.down_proj", .{l})).?), + }; + const es = try gpa.alloc(FW_Mlp, c.experts); + for (0..c.experts) |e| { + // w1=gate, w3=up, w2=down (the inventory's verified mapping) + es[e] = .{ + .gate = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.block_sparse_moe.experts.{d}.w1", .{ l, e })).?), + .up = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.block_sparse_moe.experts.{d}.w3", .{ l, e })).?), + .down = try decodeGroup(gpa, &scan, scan.find(decodedName(&nb, "language_model.model.layers.{d}.block_sparse_moe.experts.{d}.w2", .{ l, e })).?), + }; + } + m.experts = es; + lay.moe = m; + } + } + w.head = try decodeGroup(gpa, &scan, scan.find("language_model.lm_head").?); + return w; +} + +/// The f64 reference forward over the decoded fixture (runtime dims, Ref's composition). +const RefF64 = struct { + const G = f64; + gpa: std.mem.Allocator, + c: m3.serve.Config, + w: *const FW, + cap: usize, // kc/vc/pl rows per layer + len: u32 = 0, + kc: []G, + vc: []G, + pl: []G, + logits: []G, + st: []G, + qb: []G, + ab: []G, + kb: []G, + vb: []G, + ib: []G, + isk: []G, + ob: []G, + gb: []G, + ubx: []G, + acb: []G, + ks: []u32, + ssc: []G, + + fn h16(x: G) G { + return @floatCast(@as(f16, @floatCast(x))); + } + fn norm(x: []G, w: []const f32, eps: f32) void { + var ss: G = 0; + for (x) |v| ss += v * v; + const inv = 1.0 / @sqrt(ss / @as(G, @floatFromInt(x.len)) + @as(G, eps)); + for (x, w) |*v, wi| v.* = v.* * inv * (1.0 + @as(G, wi)); + } + /// y[o] = sum_c x[c] * D[c*n + o] — the decoded trellis's [k, n] row-major layout. + fn linT(x: []const G, d: []const f32, n: usize, y: []G) void { + const k = x.len; + @memset(y, 0); + for (0..k) |ci| { + const xv = x[ci]; + const row = d[ci * n ..][0..n]; + for (0..n) |o| y[o] += xv * @as(G, row[o]); + } + } + /// y[o] = sum_c x[c] * W[o*k + c] — the passthrough [n_out, k] row-major layout. + fn lin(x: []const G, w: []const f32, n_out: usize, y: []G) void { + const k = x.len; + @memset(y, 0); + for (0..k) |ci| { + const xv = x[ci]; + for (0..n_out) |o| y[o] += xv * @as(G, w[o * k + ci]); + } + } + fn rope(x: []G, pos: usize, rot: usize) void { + const half = rot / 2; + for (0..half) |i| { + const ang = @as(G, @floatFromInt(pos)) / + std.math.pow(G, @as(G, 1e7), @as(G, @floatFromInt(2 * i)) / @as(G, @floatFromInt(rot))); + const cv: G = @cos(ang); + const sv: G = @sin(ang); + const a = x[i]; + const b = x[i + half]; + x[i] = a * cv - b * sv; + x[i + half] = b * cv + a * sv; + } + } + fn act(g: G, u: G, alpha: f32, limit: f32) G { + const gg = @min(g, @as(G, limit)); + const uu = @min(@max(u, -@as(G, limit)), @as(G, limit)) + 1.0; + return gg / (1.0 + std.math.exp(-@as(G, alpha) * gg)) * uu; + } + + fn deinit(r: *RefF64) void { + const gpa = r.gpa; + gpa.free(r.kc); + gpa.free(r.vc); + gpa.free(r.pl); + gpa.free(r.logits); + gpa.free(r.st); + gpa.free(r.qb); + gpa.free(r.ab); + gpa.free(r.kb); + gpa.free(r.vb); + gpa.free(r.ib); + gpa.free(r.isk); + gpa.free(r.ob); + gpa.free(r.gb); + gpa.free(r.ubx); + gpa.free(r.acb); + gpa.free(r.ks); + gpa.free(r.ssc); + } + + fn init(gpa: std.mem.Allocator, c: m3.serve.Config, cap: usize, w: *const FW) !RefF64 { + const L: usize = c.run; + const H: usize = c.hidden; + const qk: usize = c.attn_heads * c.head_dim; + const kvr: usize = c.kv_heads * c.head_dim; + const irows: usize = c.i_heads * c.i_dim; + const inter: usize = @max(@max(c.dense_inter, c.moe_inter), c.shared_inter); + var r: RefF64 = .{ + .gpa = gpa, .c = c, .w = w, .cap = cap, .len = 0, + .kc = &.{}, .vc = &.{}, .pl = &.{}, .logits = &.{}, .st = &.{}, + .qb = &.{}, .ab = &.{}, .kb = &.{}, .vb = &.{}, .ib = &.{}, + .isk = &.{}, .ob = &.{}, .gb = &.{}, .ubx = &.{}, .acb = &.{}, + .ks = &.{}, .ssc = &.{}, + }; + errdefer r.deinit(); + r.kc = try gpa.alloc(G, L * cap * kvr); + r.vc = try gpa.alloc(G, L * cap * kvr); + r.pl = try gpa.alloc(G, L * cap * c.i_dim); + r.logits = try gpa.alloc(G, c.vocab); + r.st = try gpa.alloc(G, H); + r.qb = try gpa.alloc(G, qk); + r.ab = try gpa.alloc(G, qk); + r.kb = try gpa.alloc(G, kvr); + r.vb = try gpa.alloc(G, kvr); + r.ib = try gpa.alloc(G, irows); + r.isk = try gpa.alloc(G, c.i_dim); + r.ob = try gpa.alloc(G, H); + r.gb = try gpa.alloc(G, inter); + r.ubx = try gpa.alloc(G, inter); + r.acb = try gpa.alloc(G, inter); + r.ks = try gpa.alloc(u32, c.keyWidth()); + r.ssc = try gpa.alloc(G, c.keyWidth()); + return r; + } + + fn put(r: *RefF64, tok: u32) !void { + std.debug.assert(r.len < r.cap); + const c = &r.c; + const w = r.w; + const H: usize = c.hidden; + const qk: usize = c.attn_heads * c.head_dim; + const kvr: usize = c.kv_heads * c.head_dim; + const hd: usize = c.head_dim; + const cap = r.cap; + const pos: usize = r.len; + r.len += 1; + const sk = pos + 1; + const scale = 1.0 / @sqrt(@as(G, @floatFromInt(c.head_dim))); + for (0..H) |i| r.st[i] = @as(G, w.embed[@as(usize, tok) * H + i]); + for (w.layers, 0..) |*l, li| { + const xn = try r.gpa.alloc(G, H); + defer r.gpa.free(xn); + @memcpy(xn, r.st); + norm(xn, l.in_norm, c.eps); + linT(xn, l.q, qk, r.qb); + for (0..c.attn_heads) |h| { + const row = r.qb[h * hd ..][0..hd]; + norm(row, l.q_norm, c.eps); + rope(row, pos, c.rotary_dim); + for (row) |*v| v.* = h16(v.*); + } + linT(xn, l.k, kvr, r.kb); + for (0..c.kv_heads) |h| { + const row = r.kb[h * hd ..][0..hd]; + norm(row, l.k_norm, c.eps); + rope(row, pos, c.rotary_dim); + for (row) |*v| v.* = h16(v.*); + } + linT(xn, l.v, kvr, r.vb); + for (0..kvr) |i| r.vb[i] = h16(r.vb[i]); + if (li == 0 and pos == 0) { // the device-vs-host discriminator row: the projections' own scales + var qx2: f64 = 0; + var kx2: f64 = 0; + var vx2: f64 = 0; + for (r.qb) |v| qx2 = @max(qx2, @abs(v)); + for (r.kb) |v| kx2 = @max(kx2, @abs(v)); + for (r.vb) |v| vx2 = @max(vx2, @abs(v)); + std.debug.print("DBGR f64 L0 pos0: after-norm/rope q {d:.4} k {d:.4} v {d:.4} | k[:4]={d:.4},{d:.4},{d:.4},{d:.4} | v[:8]={d:.4},{d:.4},{d:.4},{d:.4},{d:.4},{d:.4},{d:.4},{d:.4}\n", .{ qx2, kx2, vx2, r.kb[0], r.kb[1], r.kb[2], r.kb[3], r.vb[0], r.vb[1], r.vb[2], r.vb[3], r.vb[4], r.vb[5], r.vb[6], r.vb[7] }); + var xn2s: f64 = 0; + for (xn) |v| xn2s += v * v; + std.debug.print("DBGR f64 L0 pos0: st[0]={d:.6} xn[0]={d:.6} xrms={d:.6}\n", .{ r.st[0], xn[0], @sqrt(xn2s / @as(f64, @floatFromInt(H))) }); + } + const kro = @as(usize, li) * cap * kvr; + for (0..kvr) |i| { + r.kc[kro + pos * kvr + i] = r.kb[i]; + r.vc[kro + pos * kvr + i] = r.vb[i]; + } + lin(xn, l.iq, c.i_heads * c.i_dim, r.ib); // fp16 passthrough [n_out, k] layout + for (0..c.i_heads) |h| { + const row = r.ib[h * c.i_dim ..][0..c.i_dim]; + norm(row, l.iqn, c.eps); + rope(row, pos, c.rotary_dim); + for (row) |*v| v.* = h16(v.*); + } + lin(xn, l.ik, c.i_dim, r.isk); // fp16 passthrough [n_out, k] layout + norm(r.isk, l.ikn, c.eps); + rope(r.isk, pos, c.rotary_dim); + for (0..c.i_dim) |i| r.isk[i] = h16(r.isk[i]); + const pro = @as(usize, li) * cap * c.i_dim; + for (0..c.i_dim) |i| r.pl[pro + pos * c.i_dim + i] = r.isk[i]; + // the selection: f64 block-max pool, the local boost override, strict > ties at the + // smaller id, left-packed, -1 right-pad (Ref's composition) + const nb = (sk + c.block - 1) / c.block; + var sel: [16]i32 = @splat(-1); + var ih: usize = 0; + while (ih < c.i_heads) : (ih += 1) { + var bsc: [512]G = @splat(-std.math.inf(G)); + for (0..nb) |b| { + var best = -std.math.inf(G); + const lo = b * c.block; + const nhi = @min(lo + c.block, sk); + var j = lo; + while (j < nhi) : (j += 1) { + var dot: G = 0; + for (0..c.i_dim) |dd| dot += r.ib[ih * c.i_dim + dd] * r.pl[pro + j * c.i_dim + dd]; + if (dot > best) best = dot; + } + bsc[b] = best; + } + bsc[pos / c.block] = std.math.inf(G); // the local boost override + for (0..c.i_topk) |sl| { + var best: usize = nb; + var bs = -std.math.inf(G); + for (0..nb) |b| { + var taken = false; + for (0..sl) |p| { + if (sel[p] >= 0 and sel[p] == @as(i32, @intCast(b))) taken = true; + } + if (taken) continue; + if (bsc[b] > bs) { + bs = bsc[b]; + best = b; + } + } + if (best < nb) sel[sl] = @intCast(best); + } + } + // the block-sparse f64 softmax attention, GQA by kv head + const per = c.attn_heads / c.kv_heads; + var hv: usize = 0; + while (hv < c.kv_heads) : (hv += 1) { + var nk: usize = 0; + for (0..c.i_topk) |sl| { + const bl = sel[sl]; + if (bl < 0) break; + const blo = @as(usize, @intCast(bl)) * c.block; + const nhi = @min(blo + c.block, sk); + var kk = blo; + while (kk < nhi) : (kk += 1) { + r.ks[nk] = @intCast(kk); + nk += 1; + } + } + var j: usize = 0; + while (j < per) : (j += 1) { + const qh = hv * per + j; + const qr = r.qb[qh * hd ..][0..hd]; + var mx = -std.math.inf(G); + var m: usize = 0; + while (m < nk) : (m += 1) { + const krow = r.kc[kro + r.ks[m] * kvr ..][0..kvr][hv * hd ..][0..hd]; + var dot: G = 0; + for (0..hd) |d| dot += qr[d] * krow[d]; + r.ssc[m] = dot * scale; + if (dot * scale > mx) mx = dot * scale; + } + var se: G = 0; + m = 0; + while (m < nk) : (m += 1) { + const e = std.math.exp(r.ssc[m] - mx); + r.ssc[m] = e; + se += e; + } + const acc = try r.gpa.alloc(G, hd); + defer r.gpa.free(acc); + @memset(acc, 0); + m = 0; + while (m < nk) : (m += 1) { + const vrow = r.vc[kro + r.ks[m] * kvr ..][0..kvr][hv * hd ..][0..hd]; + const p = r.ssc[m] / se; + for (0..hd) |d| acc[d] += p * vrow[d]; + } + const orow = r.ab[qh * hd ..][0..hd]; + for (0..hd) |d| orow[d] = h16(acc[d]); + } + } + for (0..qk) |i| r.ab[i] = h16(r.ab[i]); + linT(r.ab, l.o, H, r.ob); + for (0..H) |i| r.st[i] += h16(r.ob[i]); + const xn2 = try r.gpa.alloc(G, H); + defer r.gpa.free(xn2); + @memcpy(xn2, r.st); + norm(xn2, l.post_norm, c.eps); + if (l.dense) |m| { + linT(xn2, m.gate, c.dense_inter, r.gb); + linT(xn2, m.up, c.dense_inter, r.ubx); + for (0..c.dense_inter) |i| r.acb[i] = act(r.gb[i], r.ubx[i], c.swiglu_alpha, c.swiglu_limit); + linT(r.acb[0..c.dense_inter], m.down, H, r.ob); + for (0..H) |i| r.st[i] += h16(r.ob[i]); + } else { + const m = l.moe.?; + // the noaux router: select on sigmoid+bias, weights = unbiased sigmoid normalized + // over the picks, * routed_scale, each .half() — each weight pairs its own expert + const E: usize = c.experts; + const sg = try r.gpa.alloc(G, E); + defer r.gpa.free(sg); + const sc = try r.gpa.alloc(G, E); + defer r.gpa.free(sc); + const picks = try r.gpa.alloc(bool, E); + defer r.gpa.free(picks); + @memset(picks, false); + const ids = try r.gpa.alloc(u32, c.topk); + defer r.gpa.free(ids); + for (0..E) |e| { + var lg: G = 0; + for (0..H) |jj| lg += xn2[jj] * @as(G, m.router[e * H + jj]); + sg[e] = 1.0 / (1.0 + std.math.exp(-lg)); + sc[e] = sg[e] + @as(G, m.bias[e]); + } + var sum: G = 1e-20; + for (0..c.topk) |slot| { + var best: usize = E; + var bs = -std.math.inf(G); + for (0..E) |e| { + if (picks[e]) continue; + if (sc[e] > bs) { + bs = sc[e]; + best = e; + } + } + picks[best] = true; + ids[slot] = @intCast(best); + sum += sg[best]; + } + var acc = try r.gpa.alloc(G, H); + defer r.gpa.free(acc); + @memset(acc, 0); + for (0..c.topk) |slot| { + const wi: G = h16(sg[ids[slot]] / sum * @as(G, c.routed_scale)); + const e = m.experts[ids[slot]]; + linT(xn2, e.gate, c.moe_inter, r.gb); + linT(xn2, e.up, c.moe_inter, r.ubx); + for (0..c.moe_inter) |i| r.acb[i] = act(r.gb[i], r.ubx[i], c.swiglu_alpha, c.swiglu_limit); + linT(r.acb[0..c.moe_inter], e.down, H, r.ob); + for (0..H) |i| acc[i] += wi * r.ob[i]; + } + linT(xn2, m.shared.gate, c.shared_inter, r.gb); + linT(xn2, m.shared.up, c.shared_inter, r.ubx); + for (0..c.shared_inter) |i| r.acb[i] = act(r.gb[i], r.ubx[i], c.swiglu_alpha, c.swiglu_limit); + linT(r.acb[0..c.shared_inter], m.shared.down, H, r.ob); + for (0..H) |i| acc[i] += r.ob[i]; // the shared expert at weight 1 + for (0..H) |i| r.st[i] += h16(acc[i]); + } + } + const xf = try r.gpa.alloc(G, H); + defer r.gpa.free(xf); + @memcpy(xf, r.st); + norm(xf, w.final_norm, c.eps); + for (0..c.vocab) |v| { + var acc: G = 0; + for (0..H) |j| acc += xf[j] * @as(G, w.head[j * @as(usize, c.vocab) + v]); + r.logits[v] = acc; + } + } + + fn greedy(r: *RefF64, steps: usize, out: []u32) !void { + for (0..steps) |g| { + var best: usize = 0; + var bv = -std.math.inf(G); + for (r.logits, 0..) |v, i| if (v > bv) { + bv = v; + best = i; + }; + out[g] = @intCast(best); + if (g + 1 < steps) try r.put(@intCast(best)); + } + } +}; + +/// The smoke's grade: the device graph's tokens against the f64 reference's on the same fixture. +fn gradeTokens(ref_out: []const u32, dev_out: []const u32) !void { + var ok = true; + for (ref_out, dev_out) |rv, dv| { + if (rv != dv) { + ok = false; + } + } + if (!ok) { + std.debug.print("FAIL minimax-m3-smoke: device tokens", .{}); + for (dev_out) |t| std.debug.print(" {d}", .{t}); + std.debug.print(" vs reference tokens", .{}); + for (ref_out) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\n", .{}); + return error.TokenMismatch; + } +} + +/// The debug build's grade wrapper: a token mismatch still lets the run finish loudly. +fn gradeTokensCatch(ref_out: []const u32, dev_out: []const u32) !void { + gradeTokens(ref_out, dev_out) catch {}; +} + +// =================================================================== Stage-8: grade + lanes +// +// tf-cuda-test minimax-m3-grade [ctx] [steps] [layers] +// STEP 1's grade: the device graph (TP1, the first `layers` REAL layers of the pack via +// Engine.Options.layer_cap — the reduced real-pack stack that fits the host beside the +// reference) against the SAME fork f64 composition the fixture smoke grades with +// (RefF64 over the pack's own trellis groups decoded on the host). Compares every row's +// residual (eng.x vs Ref st) and logits (eng.logits vs Ref logits) and the greedy tokens. +// The full-60-layer real pack stays beyond the honest host f64 budget (hundreds of GB of +// decoded f64 weights and hours of trellis decode) — reported as the boundary, not hidden. +// +// tf-cuda-test minimax-m3-lanes [ctx] [steps] [layers] +// STEP 2's boundary test through open()'s Loaded handle: a real Stream served through the +// lane vtable (prefill → first → queue/read round trips → a verify round → keep rollback), +// graded token-for-token against the same prompt run DIRECTLY through a second Engine, plus +// the honest refusal rows (draft, unknown handle, drafted verify, batched verify, sampled +// spec, second stream, tree keep) and release+re-prefill. + +/// The max rows the grade's host buffers hold (ctx + steps - 1 puts both sides take). +const grade_rows_max: usize = 256; + +/// One side's per-row captures: the residual and the logits at every put's position. +const GradeCap = struct { + gpa: std.mem.Allocator, + c: m3.serve.Config, + rows: usize = 0, + x: [][]f32 = &.{}, + lg: [][]f32 = &.{}, + + fn init(g: std.mem.Allocator, c: m3.serve.Config, rows: usize) !GradeCap { + const H: usize = c.hidden; + var cap = GradeCap{ .gpa = g, .c = c }; + errdefer cap.deinit(); + cap.x = try g.alloc([]f32, rows); + cap.lg = try g.alloc([]f32, rows); + for (0..rows) |i| { + cap.x[i] = try g.alloc(f32, H); + cap.lg[i] = try g.alloc(f32, c.vocab); + } + return cap; + } + fn deinit(cap: *GradeCap) void { + for (cap.x) |r| cap.gpa.free(r); + for (cap.lg) |r| cap.gpa.free(r); + if (cap.x.len > 0) cap.gpa.free(cap.x); + if (cap.lg.len > 0) cap.gpa.free(cap.lg); + cap.* = .{ .gpa = cap.gpa, .c = cap.c }; + } +}; + +/// The device side's capture loop: identical token sequence to eng.greedy, but with a capture +/// after every put (the harness reads eng.x and eng.logits between rows). +fn deviceRun(eng: *m3.graph.Engine, prompt: []const u32, steps: usize, cap: *GradeCap, out: []u32) !void { + var row: usize = 0; + for (prompt) |tok| { + try eng.put(tok); + try captureRow(cap, eng.x, eng.logits, row); + row += 1; + } + for (0..steps) |g| { + var best: u32 = 0; + var bv = -std.math.inf(f32); + for (eng.logits, 0..) |v, i| { + if (v > bv) { + bv = v; + best = @intCast(i); + } + } + out[g] = best; + if (g + 1 < steps) { + try eng.put(best); + try captureRow(cap, eng.x, eng.logits, row); + row += 1; + } + } +} + +fn captureRow(cap: *GradeCap, x: []const f32, lg: []const f32, row: usize) !void { + if (row >= cap.x.len) return error.CaptureOverrun; + @memcpy(cap.x[row], x); + @memcpy(cap.lg[row], lg); + cap.rows = row + 1; +} + +/// The f64 side's capture loop over the SAME tokens (its own greedy draws, compared below). +fn refRun(ref: *RefF64, prompt: []const u32, steps: usize, cap: *GradeCap, out: []u32) !void { + var row: usize = 0; + for (prompt) |tok| { + try ref.put(tok); + try captureRef(cap, ref, row); + row += 1; + } + for (0..steps) |g| { + var best: u32 = 0; + var bv: f64 = -std.math.inf(f64); + for (ref.logits, 0..) |v, i| { + if (v > bv) { + bv = v; + best = @intCast(i); + } + } + out[g] = best; + if (g + 1 < steps) { + try ref.put(best); + try captureRef(cap, ref, row); + row += 1; + } + } +} + +fn captureRef(cap: *GradeCap, ref: *RefF64, row: usize) !void { + if (row >= cap.x.len) return error.CaptureOverrun; + for (ref.st, 0..) |v, i| cap.x[row][i] = @floatCast(v); + for (ref.logits, 0..) |v, i| cap.lg[row][i] = @floatCast(v); + cap.rows = row + 1; +} + +/// max|delta| over two fp32 rows (NaN never passes; both rows are finite by construction here). +fn rowMaxd(a: []const f32, b: []const f32) f64 { + var d: f64 = 0; + for (a, b) |va, vb| { + const dd = @abs(@as(f64, va) - @as(f64, vb)); + if (dd > d) d = dd; + } + return d; +} + +pub fn gradeRun(gpu: check.Gpu, dir_in: []const u8, ctx_len: usize, steps: usize, layers_arg: usize) !void { + const gpa = gpu.gpa; + var dir = dir_in; + if (dir.len == 0) { + std.debug.print("usage: tf-cuda-test minimax-m3-grade [ctx] [steps] [layers]\n", .{}); + return error.BadArguments; + } + if (steps == 0 or steps > 32 or ctx_len == 0 or ctx_len > m3.graph.keys_cap or ctx_len + steps > grade_rows_max) { + return error.BadArguments; // loud refusals, no silent clamping + } + if (layers_arg == 0) { + std.debug.print("FAIL minimax-m3-grade: the layers cap must be explicit (0 is the full stack — beyond the honest f64 reference budget on the host)\n", .{}); + return error.BadArguments; + } + if (std.mem.eql(u8, dir, "fixture")) dir = try makeFixture(gpu); + const t0 = check.now(gpu.io); + var eng = try m3.graph.Engine.build(gpa, gpu.io, gpu.d, dir, 0, 1, .{ .layer_cap = layers_arg }); + defer eng.deinit(); + const c = eng.c; + if (layers_arg > c.layers) return error.LayerCapTooWide; + // the reference reads the SAME first `layers_arg` layers (its own config run count) + var cr = c; + cr.run = @intCast(layers_arg); + const ram_before = hostFree(gpu); + + // the host f64 reference over the pack's own decoded trellis (the smoke's decode path) + var fw = try loadFixtureW(gpa, gpu.io, dir, cr); + defer fw.deinit(); + const ram_after_decode = hostFree(gpu); + + const rows = ctx_len + steps - 1; + var dev_cap = try GradeCap.init(gpa, c, rows); + defer dev_cap.deinit(); + var ref_cap = try GradeCap.init(gpa, c, rows); + defer ref_cap.deinit(); + + const prompt = try gpa.alloc(u32, ctx_len); + defer gpa.free(prompt); + var l = Lcg{ .s = prompt_seed }; + m3.serve.promptOf(&l, prompt); + + const dev_out = try gpa.alloc(u32, steps); + defer gpa.free(dev_out); + const ref_out = try gpa.alloc(u32, steps); + defer gpa.free(ref_out); + + var ref = try RefF64.init(gpa, cr, rows + 1, &fw); + defer ref.deinit(); + + const t1 = check.now(gpu.io); + try deviceRun(&eng, prompt, steps, &dev_cap, dev_out); + try refRun(&ref, prompt, steps, &ref_cap, ref_out); + const run_s = @as(f64, @floatFromInt(check.now(gpu.io) - t1)) / 1e9; + + // -- the comparisons -- + var xmax: f64 = 0; + var lmax: f64 = 0; + var lnan: usize = 0; + var linf: usize = 0; + for (0..dev_cap.rows) |r| { + const x = rowMaxd(dev_cap.x[r], ref_cap.x[r]); + const ld = rowMaxd(dev_cap.lg[r], ref_cap.lg[r]); + if (x > xmax) xmax = x; + if (ld > lmax) lmax = ld; + std.debug.print("M3-GRADE row {d:3}: residual max|delta| {d:.6} logits max|delta| {d:.6}\n", .{ r, x, ld }); + if (r == 0) { + std.debug.print("M3-GRADE row 0 raw: dev x[0..4] = {d:.5} {d:.5} {d:.5} {d:.5} | ref st[0..4] = {d:.5} {d:.5} {d:.5} {d:.5}\n", .{ dev_cap.x[r][0], dev_cap.x[r][1], dev_cap.x[r][2], dev_cap.x[r][3], ref_cap.x[r][0], ref_cap.x[r][1], ref_cap.x[r][2], ref_cap.x[r][3] }); + } + for (dev_cap.x[r]) |v| { + if (std.math.isNan(v) or std.math.isInf(v)) lnan += 1; + } + for (dev_cap.lg[r]) |v| { + if (std.math.isNan(v)) lnan += 1 else if (std.math.isInf(v)) linf += 1; + } + } + var exact = true; + for (dev_out, ref_out) |dv, rv| { + if (dv != rv) exact = false; + } + const build_s = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + std.debug.print("M3-GRADE: layers {d}/{d}, rows {d}, steps {d}; build+decode {d:.1}s, both runs {d:.1}s\n", .{ layers_arg, c.layers, dev_cap.rows, steps, build_s, run_s }); + std.debug.print("M3-GRADE: max|delta| residual {d:.6} logits {d:.6}; dev rows NaN {d} inf {d}\n", .{ xmax, lmax, lnan, linf }); + std.debug.print("M3-GRADE: host MemAvailable {d:.2} -> {d:.2} (post-decode) GiB\n", .{ @as(f64, @floatFromInt(ram_before)) / (1024 * 1024 * 1024), @as(f64, @floatFromInt(ram_after_decode)) / (1024 * 1024 * 1024) }); + std.debug.print("M3-GRADE: dev greedy", .{}); + for (dev_out) |t| std.debug.print(" {d}", .{t}); + std.debug.print(" | ref greedy", .{}); + for (ref_out) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\n", .{}); + // the draw margins: is the flipped draw a near-tie (quant tolerance) or a real miss? + for (0..steps) |g| { + const r = ctx_len - 1 + g; // draw g keys at the row ending with its feeding token's put + if (r >= dev_cap.rows) break; + if (g > 0 and dev_out[g - 1] != ref_out[g - 1] and ctx_len + g > 0) break; // margins before the first flip are the informative ones + var d1: usize = 0; + var d2: usize = 0; + var dv1: f32 = -std.math.inf(f32); + var dv2: f32 = -std.math.inf(f32); + for (dev_cap.lg[r], 0..) |v, i| { + if (v > dv1) { + dv2 = dv1; + d2 = d1; + dv1 = v; + d1 = i; + } else if (v > dv2) { + dv2 = v; + d2 = i; + } + } + var r1: usize = 0; + var r2: usize = 0; + var rv1: f64 = -std.math.inf(f64); + var rv2: f64 = -std.math.inf(f64); + for (ref_cap.lg[r], 0..) |v, i| { + if (v > rv1) { + rv2 = rv1; + r2 = r1; + rv1 = v; + r1 = i; + } else if (v > rv2) { + rv2 = v; + r2 = i; + } + } + std.debug.print("M3-GRADE draw {d}: dev top1 {d} ({d:.4}) top2 {d} ({d:.4}) | ref top1 {d} ({d:.4}) top2 {d} ({d:.4})\n", .{ g, d1, dv1, d2, dv2, r1, rv1, r2, rv2 }); + } + if (!exact) { + std.debug.print("FAIL minimax-m3-grade: the device graph's greedy tokens differ from the f64 reference's over the SAME decoded pack weights (rows {d})\n", .{dev_cap.rows}); + return error.TokenMismatch; // WRONG verdict: never a "close enough" pass over a token miss + } + std.debug.print("RESULT minimax-m3-grade: device graph == f64 reference (EXACT greedy tokens, {d} rows graded, residual max|delta| {d:.6}, logits max|delta| {d:.6}, first {d} layers of the real pack)\n", .{ dev_cap.rows, xmax, lmax, layers_arg }); +} + +// ------------------------------------------------------------------ the lane-surface boundary test + +fn laneTokensEqual(a: []const u32, b: []const u32) bool { + if (a.len != b.len) return false; + for (a, b) |va, vb| { + if (va != vb) return false; + } + return true; +} + +fn expectRefusal(comptime what: []const u8, want: anyerror, call: anytype) !void { + if (call) |_| { + std.debug.print("FAIL minimax-m3-lanes: {s} served where it must refuse\n", .{what}); + return error.RefusalExpected; + } else |e| { + if (e != want) { + std.debug.print("FAIL minimax-m3-lanes: {s} refused with {s}, expected {s}\n", .{ what, @errorName(e), @errorName(want) }); + return error.RefusalWanted; + } + } + std.debug.print("[lanes] refusal {s}: {s}\n", .{ what, @errorName(want) }); +} + +pub fn laneRun(gpu: check.Gpu, dir_in: []const u8, ctx_len: usize, steps: usize, layers_arg: usize) !void { + const gpa = gpu.gpa; + var dir = dir_in; + if (dir.len == 0 or steps == 0 or steps > 32 or ctx_len == 0 or ctx_len > 128) { + std.debug.print("usage: tf-cuda-test minimax-m3-lanes [ctx<=128] [steps<=32] [layers]\n", .{}); + return error.BadArguments; + } + if (std.mem.eql(u8, dir, "fixture")) dir = try makeFixture(gpu); + const layers: usize = layers_arg; // 0 = the whole owned stack (full TP1; real packs use a cap) + + // -- the direct path's tokens first: the same prompt through a raw Engine -- + var direct = try m3.graph.Engine.build(gpa, gpu.io, gpu.d, dir, 0, 1, .{ .layer_cap = layers }); + defer direct.deinit(); + const prompt = try gpa.alloc(u32, ctx_len); + defer gpa.free(prompt); + var l = Lcg{ .s = prompt_seed }; + m3.serve.promptOf(&l, prompt); + const direct_out = try gpa.alloc(u32, steps); + defer gpa.free(direct_out); + for (prompt) |tok| try direct.put(tok); + try direct.greedy(steps, direct_out); + std.debug.print("M3-LANES: direct engine tokens", .{}); + for (direct_out) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\n", .{}); + + // -- the served path: open()'s Loaded handle, the lane vtable -- + var loaded = try m3.native.open(gpa, gpu.io, gpu.ctx, dir, null, .{ .context = m3.block_tokens, .drafts = false, .layer_cap = layers }); + defer loaded.deinit(loaded.ctx); + const b = loaded.backend; + var s = try m3.lane_api.Stream.init(gpa, .{ .id = "lane-boundary", .prompt = prompt, .max_new = @intCast(steps + 2), .drafts = false }); + defer s.deinit(gpa); + + try b.prefill(&s); + if (s.fill_at.load(.acquire) != s.prompt_len) return error.LaneFillLost; + var h = try b.first(&s, s.prompt_len); + var tok = try b.read(h); + std.debug.print("[lanes] first: handle {d} token {d}\n", .{ h, tok }); + _ = try s.commit(gpa, &.{tok}, &.{}); + + var lane_out = try gpa.alloc(u32, steps); + defer gpa.free(lane_out); + lane_out[0] = tok; + // the queue/read round trips: row k feeds lane_out[k] at the engine's cache length and + // draws lane_out[k+1]; the last round (feeding lane_out[steps-2]) goes through verify + var idx: usize = 1; + while (idx < steps) : (idx += 1) { + const qpos: u64 = s.prompt_len + @as(u64, idx - 1); // the engine's own cache length + const feed = lane_out[idx - 1]; + if (idx + 1 < steps) { + h = try b.queue(&s, .{ .value = feed }, qpos); + tok = try b.read(h); + std.debug.print("[lanes] queue@{d}: handle {d} token {d}\n", .{ qpos, h, tok }); + } else { + var sampled: [1]u32 = .{0}; + var drafts_empty: [0]u32 = .{}; + const windows = [_]m3.lane_api.backend.Window{.{ .stream = &s, .pending = feed, .held = 0, .tokens = &.{}, .parents = null, .positions = &[_]u64{qpos} }}; + var verified = [_]m3.lane_api.backend.Verified{.{ .sampled = sampled[0..], .drafts = drafts_empty[0..] }}; + try b.verify(&windows, &verified); + tok = verified[0].sampled[0]; + std.debug.print("[lanes] verify@{d}: token {d}\n", .{ qpos, tok }); + } + lane_out[idx] = tok; + _ = try s.commit(gpa, &.{tok}, &.{}); + } + + // -- keep: roll back the verify round's fed row and re-serve it; the draw must repeat -- + if (steps >= 2) { + const kept: usize = s.context.items.len - 2; // rows 0..kept-1 of the engine's cache + const path = try gpa.alloc(u32, kept); + defer gpa.free(path); + for (path, 0..) |*p, i| p.* = @intCast(i); + const kwin = [_]m3.lane_api.backend.Window{.{ .stream = &s, .pending = 0, .held = 0, .tokens = &.{}, .parents = null, .positions = &[_]u64{0} }}; + var kpaths = [_][]const u32{path}; + try b.keep(&kwin, &kpaths); + h = try b.queue(&s, .{ .value = lane_out[kept - s.prompt_len] }, @intCast(kept)); + tok = try b.read(h); + if (tok != lane_out[kept - s.prompt_len + 1]) { + std.debug.print("FAIL minimax-m3-lanes: the keep rollback re-drew {d} where the round drew {d}\n", .{ tok, lane_out[kept - s.prompt_len + 1] }); + return error.TokenMismatch; + } + std.debug.print("[lanes] keep@{d}: rollback re-drew token {d} (unchanged)\n", .{ kept, tok }); + } + + // -- the grade through the lane surface: token-for-token against the direct engine -- + if (!laneTokensEqual(lane_out, direct_out)) { + std.debug.print("FAIL minimax-m3-lanes: lane tokens", .{}); + for (lane_out) |t| std.debug.print(" {d}", .{t}); + std.debug.print(" vs direct tokens", .{}); + for (direct_out) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\n", .{}); + return error.TokenMismatch; + } + + // -- the honest refusals (each with its exact reason) -- + var requests = [_]m3.lane_api.backend.DraftRequest{.{ .stream = &s, .follow = &.{tok}, .rows = null, .start = 0, .position = 0, .depth = 1 }}; + try expectRefusal("draft", error.LanesNotAttached, b.draft(&requests)); + try expectRefusal("read(unknown handle)", error.LaneHandleUnknown, b.read(424242)); + var sampled2: [1]u32 = .{0}; + const dwin = [_]m3.lane_api.backend.Window{.{ .stream = &s, .pending = tok, .held = 1, .tokens = &.{}, .parents = null, .positions = &[_]u64{s.context.items.len} }}; + var vout2 = [_]m3.lane_api.backend.Verified{.{ .sampled = sampled2[0..], .drafts = &.{} }}; + try expectRefusal("verify(drafted window)", error.LanesNotAttached, b.verify(&dwin, &vout2)); + const twin = [_]m3.lane_api.backend.Window{ dwin[0], dwin[0] }; + var vout3 = [_]m3.lane_api.backend.Verified{ vout2[0], vout2[0] }; + try expectRefusal("verify(batch of 2)", error.LanesNotAttached, b.verify(&twin, &vout3)); + var s2 = try m3.lane_api.Stream.init(gpa, .{ .id = "second-stream", .prompt = prompt, .max_new = 2, .drafts = false }); + defer s2.deinit(gpa); + try expectRefusal("prefill(second stream)", error.LanesNotAttached, b.prefill(&s2)); + const tree_path = [_]u32{ 0, 2 }; // a tree's path: not the prefix 0..kept-1 + const kwin2 = [_]m3.lane_api.backend.Window{.{ .stream = &s, .pending = 0, .held = 0, .tokens = &.{}, .parents = null, .positions = &[_]u64{0} }}; + var kpaths2 = [_][]const u32{tree_path[0..]}; + try expectRefusal("keep(tree path)", error.LanesNotAttached, b.keep(&kwin2, &kpaths2)); + + // -- release, then the re-prefill serves fresh; a sampled spec refuses before it can go live -- + b.release(&s); + var s3 = try m3.lane_api.Stream.init(gpa, .{ .id = "sampled", .prompt = prompt[0..2], .max_new = 1, .drafts = false, .sampling = .{ .seed = 7 } }); + defer s3.deinit(gpa); + try expectRefusal("prefill(sampled spec)", error.SamplingNotAttached, b.prefill(&s3)); + var s4 = try m3.lane_api.Stream.init(gpa, .{ .id = "again", .prompt = prompt, .max_new = @intCast(steps + 2), .drafts = false }); + defer s4.deinit(gpa); + try b.prefill(&s4); // a fresh cache life after release: serves again + const h4 = try b.first(&s4, s4.prompt_len); + const again1 = try b.read(h4); + if (again1 != lane_out[0]) { + std.debug.print("FAIL minimax-m3-lanes: the released re-prefill drew {d} where the round drew {d}\n", .{ again1, lane_out[0] }); + return error.TokenMismatch; + } + b.release(&s4); + std.debug.print("RESULT minimax-m3-lanes: the Loaded handle serves a {d}-row prompt through the lane surface (first/queue/verify/keep/release) — tokens EXACT vs the direct engine; drafts/sampled/second-stream/tree-keep refuse truthfully\n", .{ctx_len}); +} + +// ------------------------------------------------- the TP2 lane-surface boot (Stage-9) +// +// tf-cuda-test minimax-m3-lane-tp2 [ctx] [steps] +// the Stage-8 refusal's proof: open()'s own lane vtable — openTP2's handle, the SAME +// Attached surface TP1 serves — answering a real prompt over the SPLIT real pack. +// Rank 1 first over the NCCL id (the token-harness recipe); rank 0 serves the receipt: +// prefill → first → queue/read rounds straight through the vtable, its greedy draws +// argued from its own logits row. Rank 1 keeps the row lockstep through its own surface +// (every put is a lockstep row; its draws are scratch and are tolerated). +// +pub fn laneTp2Run(gpu: check.Gpu, rank_arg: []const u8, addr_arg: []const u8, dir_in: []const u8, ctx_len: usize, steps: usize) !void { + var dir = dir_in; + const rank8 = std.fmt.parseInt(u8, rank_arg, 10) catch 99; + if (rank8 > 1 or addr_arg.len == 0 or dir.len == 0 or steps == 0 or steps > 32 or ctx_len == 0 or ctx_len > m3.graph.keys_cap) { + std.debug.print("usage: tf-cuda-test minimax-m3-lane-tp2 [ctx] [steps]\n", .{}); + return error.BadArguments; + } + if (std.mem.eql(u8, dir, "fixture")) dir = try makeFixture(gpu); + const rank: u32 = @intCast(rank8); + const world: u32 = 2; + const gib = struct { + fn f(v: u64) f64 { + return @as(f64, @floatFromInt(v)) / (1024.0 * 1024.0 * 1024.0); + } + }.f; + + // -- the TP2 boot exactly the token harness's recipe (rank 1 first, rank 0 serves the id) -- + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + var id: cuda.nccl.UniqueId = undefined; + const id_bytes = std.mem.asBytes(&id); + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try serveId(addr_arg, id_bytes); + } else { + try receiveId(gpu, addr_arg, id_bytes); + } + var comm: cuda.nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, @intCast(world), id, @intCast(rank)), "ncclCommInitRank"); + defer _ = lib.api.ncclCommDestroy(comm); + + const ram_before = hostFree(gpu); + const t0 = check.now(gpu.io); + var loaded = try m3.openTP2(gpu.gpa, gpu.io, gpu.ctx, dir, .{ .context = ctx_len, .drafts = false }, .{ .lib = &lib, .comm = comm, .rank = rank, .world = world }); + defer loaded.deinit(loaded.ctx); + const b = loaded.backend; + const build_s = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + const ram_after_build = hostFree(gpu); + + // -- the same deterministic context both ranks (rank 1's token values are rank 0's business) -- + const prompt = try gpu.gpa.alloc(u32, ctx_len); + defer gpu.gpa.free(prompt); + var l = Lcg{ .s = prompt_seed }; + m3.serve.promptOf(&l, prompt); + const lane_out = try gpu.gpa.alloc(u32, steps); + defer gpu.gpa.free(lane_out); + const put_t0 = check.now(gpu.io); + + // -- the serving block: prefill → first → queue/read straight through the vtable -- + // rank 0 argues every greedy draw from its real logits row and prints the receipt; + // rank 1 keeps the ROW LOCKSTEP through its own surface (every put is a paired NCCL row) + var s = try m3.lane_api.Stream.init(gpu.gpa, .{ .id = "lane-tp2", .prompt = prompt, .max_new = @intCast(steps + 2), .drafts = false }); + defer s.deinit(gpu.gpa); + try b.prefill(&s); + if (s.fill_at.load(.acquire) != s.prompt_len) return error.LaneFillLost; + const prefill_s = @as(f64, @floatFromInt(check.now(gpu.io) - put_t0)) / 1e9; + const eng: *m3.graph.Engine = @ptrCast(@alignCast(loaded.target)); // the same engine the vtable drives + // first: rank 0's draw is the receipt's first token; rank 1's `logits` is a 1-slot scratch, + // so its draw may honestly refuse (first touches no put) — tolerated, the lockstep keeps + var h: u64 = 0; + if (rank == 0) { + h = try b.first(&s, s.prompt_len); + lane_out[0] = try b.read(h); + } else { + h = b.first(&s, s.prompt_len) catch 0; // scratch draw tolerated; handle unused on rank 1 + } + _ = try s.commit(gpu.gpa, &.{lane_out[0]}, &.{}); + var idx: usize = 1; + while (idx < steps) : (idx += 1) { + const qpos: u64 = s.prompt_len + @as(u64, idx - 1); // the engine's own cache length + const len_before = eng.len; + if (rank == 0) { + h = try b.queue(&s, .{ .value = lane_out[idx - 1] }, qpos); + lane_out[idx] = try b.read(h); + } else { + // rank 1 feeds 0 (its row's value is rank 0's business); if the SURFACE refuses + // after this round's put (a scratch-logits draw refusal), the row already landed: + // the lockstep keeps. A refusal BEFORE the put (RowKeying/Context/second-stream) + // is a real surface error and dies loudly rather than desyncing the ranks. + if (b.queue(&s, .{ .value = 0 }, qpos)) |_| {} else |e| { + if (eng.len != len_before) { + std.debug.print("M3-LANE-TP2 rank 1: draw refused ({s}) after the row's put — lockstep keeps\n", .{@errorName(e)}); + } else return e; + } + } + _ = try s.commit(gpu.gpa, &.{lane_out[idx]}, &.{}); + } + const lanes_s = @as(f64, @floatFromInt(check.now(gpu.io) - put_t0)) / 1e9; + const ram_after = hostFree(gpu); + const device_used = cuda.usage(false).device; + + // -- the receipt rows + std.debug.print("M3-LANE-TP2 rank {d} world {d}: build {d:.1}s, prefill(x{d}) {d:.2}s, lane rows(x{d}) {d:.2}s\n", .{ rank, world, build_s, ctx_len, prefill_s, steps, lanes_s }); + std.debug.print("M3-LANE-TP2 rank {d}: device buffers {d:.2} GiB; host MemAvailable {d:.2} -> {d:.2} -> {d:.2} GiB\n", .{ rank, gib(device_used), gib(ram_before), gib(ram_after_build), gib(ram_after) }); + if (rank == 0) { + std.debug.print("M3-LANE-TP2 rank 0: served tokens", .{}); + for (lane_out) |t| std.debug.print(" {d}", .{t}); + std.debug.print(" (greedy x{d} through the lane vtable over the split pack)\n", .{steps}); + var lnan: usize = 0; + var linf: usize = 0; + var lmx: f32 = -std.math.inf(f32); + var lsum: f64 = 0; + for (eng.logits) |v| { + if (v != v) lnan += 1 else { + if (v > lmx) lmx = v; + if (v == std.math.inf(f32)) linf += 1 else lsum += v; + } + } + std.debug.print("M3-LANE-TP2 rank 0: logits rows {d} NaN {d} inf {d} max {d:.6} mean {d:.6}\n", .{ eng.logits.len, lnan, linf, lmx, lsum / @as(f64, @floatFromInt(eng.logits.len)) }); + std.debug.print("RESULT minimax-m3-lane-tp2 rank 0: openTP2's lane vtable answered a {d}-row prompt over the split pack: {d} greedy draws served (the surface, not just the engine harness)\n", .{ ctx_len, steps }); + // -- Stage-10: the served draws as readable text, the serve door's own code path -- + // (serveTokenizer/serveTextOut). The ids receipt above stays authoritative: a decode + // refusal (no tokenizer.json, or an unknown-id DIFF) is printed honestly, never fatal + // in the receipted TP2 path. + var tp_tok: ?tokenizer.Tokenizer = serveTokenizer(gpu.gpa, gpu.io, dir) catch |e| blk: { + std.debug.print("M3-LANE-TP2 rank 0: text decode refuses ({s}) — the ids receipt above stands\n", .{@errorName(e)}); + break :blk null; + }; + if (tp_tok) |*t| { + defer t.deinit(); + serveTextOut(t, gpu.gpa, "lane-tp2", lane_out) catch |e| { + std.debug.print("M3-LANE-TP2 rank 0: text receipt refused ({s}) — the ids receipt above stands\n", .{@errorName(e)}); + }; + } + } else { + std.debug.print("RESULT minimax-m3-lane-tp2 rank 1: row lockstep held ({d} rows) through the lane surface\n", .{ctx_len + steps - 1}); + } +} + +// ------------------------------------------------- the serve front door (Stage-9) +// +// tf-cuda-test minimax-m3-serve [steps<=32] [layers] +// the smallest REAL answering surface: open()'s attached lane vtable takes prompt token +// ids in and prints one greedy completion token ids out, then exits. "lcg" reads the +// receipted deterministic prompt stream (Stage-6C/8's 8-row seed 0x5EEDB00B bar); anything +// else is a file holding a JSON {"tokens": [...]} (kolibri-greedy's convention), a bare +// JSON id array, or whitespace-separated ids. Greedy-only and deterministic by the stage's +// honesty: drafts/sampled/second-stream refuse with their exact names (the same refusal +// law Stage-8 receipted, re-probed here end to end). At the receipted fixture recipe +// (fixture dir + lcg + 8 rows + [layers] 2) the completion MUST be EXACT 183 149 14 — +// the bar this command is graded against every time it runs there. +// +/// The lanes refusal probe with the front door's own tag (the same law as laneRun's). +fn serveRefusal(comptime what: []const u8, want: anyerror, call: anytype) !void { + if (call) |_| { + std.debug.print("FAIL minimax-m3-serve: {s} served where it must refuse\n", .{what}); + return error.RefusalExpected; + } else |e| { + if (e != want) { + std.debug.print("FAIL minimax-m3-serve: {s} refused with {s}, expected {s}\n", .{ what, @errorName(e), @errorName(want) }); + return error.RefusalWanted; + } + } + std.debug.print("[serve] refusal {s}: {s}\n", .{ what, @errorName(want) }); +} + +/// The prompt in: "lcg" (the receipted bar stream) or a prompt file (JSON token rows or bare ids). +fn servePrompt(a: std.mem.Allocator, io: std.Io, src: []const u8) ![]u32 { + if (std.mem.eql(u8, src, "lcg")) { + const prompt = try a.alloc(u32, 8); + var l = Lcg{ .s = prompt_seed }; + m3.serve.promptOf(&l, prompt); + return prompt; + } + const text = try std.Io.Dir.cwd().readFileAlloc(io, src, a, .limited(1 << 26)); + const trimmed = std.mem.trim(u8, text, " \t\r\n"); + if (trimmed.len > 0 and (trimmed[0] == '{' or trimmed[0] == '[')) { + if (trimmed[0] == '{') { + var parsed = try std.json.parseFromSlice(struct { tokens: []const u32 }, a, trimmed, .{}); + defer parsed.deinit(); + return a.dupe(u32, parsed.value.tokens); + } + var parsed = try std.json.parseFromSlice([]const u32, a, trimmed, .{}); + defer parsed.deinit(); + return a.dupe(u32, parsed.value); + } + var ids: std.ArrayList(u32) = .empty; + var it = std.mem.tokenizeAny(u8, trimmed, " \t\r\n,"); + while (it.next()) |t| try ids.append(a, try std.fmt.parseInt(u32, t, 10)); + return a.dupe(u32, ids.items); +} + +// ------------------------------------------------- the serve text layer (Stage-10) +// +// Free-text in, readable text out, through the PACK'S OWN tokenizer.json — the core engine +// (src/core/tokenizer), whose encode == HF tokenizers' encode(text, add_special_tokens=False) +// and whose decode runs the tokenizer.json decoder chain (HF-exact on this pack: the Stage-7 +// receipt bar decodes bit-identically, encode parity 19/19 probe). Text serves on a dir that +// carries tokenizer.json; "fixture" serves token ids as before, and takes `tokdir` (the pack +// dir) as a 5th argument when only the TEXT PLUMBING (text -> ids -> model -> ids -> text) is +// being probed at fixture scale — labeled honestly at run time. The token-IDS paths above are +// untouched: "lcg" and id prompts refuse nothing and print ids exactly as Stage-9 receipted. +// +/// The dir's own tokenizer, or null when the dir carries none (fixture or an ids-only layout). +/// A present-but-unreadable tokenizer.json is an error, never ignored (honest). +fn serveTokenizer(gpa: std.mem.Allocator, io: std.Io, dir: []const u8) !?tokenizer.Tokenizer { + return tokenizer.loadTokenizer(io, gpa, dir) catch |e| switch (e) { + error.FileNotFound => null, + else => return e, + }; +} + +/// A text prompt source ("text:" inline, or "textfile:" raw file bytes) encoded +/// through the pack's tokenizer, with the in-side round-trip held as a hard bar: +/// decode(encode(text)) must return the prompt text bit-exact (lossy U+FFFD substitution refuses). +fn serveTextPrompt(io: std.Io, t: *tokenizer.Tokenizer, gpa: std.mem.Allocator, src: []const u8) ![]u32 { + var text: []const u8 = undefined; + var label: []const u8 = undefined; + if (std.mem.startsWith(u8, src, "textfile:")) { + label = "textfile"; + text = try std.Io.Dir.cwd().readFileAlloc(io, src["textfile:".len..], gpa, .limited(1 << 26)); + } else { + label = "text"; + text = src["text:".len..]; + } + const ids = try t.encode(gpa, text); + const back = try t.decode(gpa, ids, false); + defer gpa.free(back); + if (std.mem.eql(u8, text, back)) { + std.debug.print("[serve] text round-trip in ({s}, {d} bytes -> {d} ids): EXACT decode(encode(text)) == text\n", .{ label, text.len, ids.len }); + } else { + std.debug.print("[serve] text round-trip in ({s}): DIFF (the tokenizer does not round-trip this text bit-exact)\n", .{label}); + return error.TextPromptRoundTripDiffers; + } + return ids; +} + +/// Serve's readable completion: the ids are the model's truth; the text is the pack tokenizer's +/// decode of them (skip-specials for the readable line, specials counted and named honestly), and +/// the out-side round-trip encode(decode(ids)) == ids is held as a hard bar. Unknown ids decode to +/// nothing (as HF skips them); if the model drew one, the round-trip honestly reports DIFF. +fn serveTextOut(t: *tokenizer.Tokenizer, gpa: std.mem.Allocator, tag: []const u8, ids: []const u32) !void { + const full = try t.decode(gpa, ids, false); + defer gpa.free(full); + const re = try t.encode(gpa, full); + defer gpa.free(re); + var specials: usize = 0; + for (ids) |id| { + const piece = t.id_to_token.get(id) orelse continue; + if (t.specials.contains(piece)) specials += 1; + } + const readable = try t.decode(gpa, ids, true); + defer gpa.free(readable); + if (std.mem.eql(u32, re, ids)) { + std.debug.print("[{s}] text round-trip out: EXACT encode(decode(ids)) == ids (specials {d}), text: '{s}'\n", .{ tag, specials, readable }); + } else { + std.debug.print("[{s}] text round-trip out: DIFF ({d} of {d} re-encoded; specials {d}), text: '{s}'\n", .{ tag, re.len, ids.len, specials, readable }); + return error.TextCompletionRoundTripDiffers; + } +} + +pub fn serveRun(gpu: check.Gpu, dir_in: []const u8, prompt_src: []const u8, steps: usize, layers_arg: usize, tok_arg: []const u8) !void { + const gpa = gpu.gpa; + var dir = dir_in; + const is_text = std.mem.startsWith(u8, prompt_src, "text:") or std.mem.startsWith(u8, prompt_src, "textfile:"); + if (dir.len == 0 or prompt_src.len == 0 or steps == 0 or steps > 32) { + std.debug.print("usage: tf-cuda-test minimax-m3-serve [steps<=32] [layers] [tokdir]\n", .{}); + return error.BadArguments; + } + var fixture_served = false; + if (std.mem.eql(u8, dir, "fixture")) { + dir = try makeFixture(gpu); + fixture_served = true; + } + const layers: usize = layers_arg; // fixture's recipe receipts 2 (the whole fixture stack); a real pack picks its cap + + // -- Stage-10: the pack's own tokenizer, when the dir (or an explicit tokdir) carries one -- + const tok_dir: []const u8 = if (tok_arg.len > 0) tok_arg else dir; + var tok_opt = try serveTokenizer(gpa, gpu.io, tok_dir); + defer if (tok_opt) |*t| t.deinit(); + if (tok_opt == null and is_text) { + std.debug.print("FAIL minimax-m3-serve: a text prompt needs the pack's own tokenizer.json (text:/textfile: encode through it); '{s}' carries none — pass the pack dir as [tokdir]\n", .{tok_dir}); + return error.ServeTextNeedsTokenizer; + } + if (tok_opt != null and fixture_served) { + std.debug.print("[serve] text tokenizer: {s} (explicit tokdir; the fixture's toy draws decode through it for the plumbing round-trip — the fixture bar applies to lcg ids only)\n", .{tok_dir}); + } + + const prompt: []u32 = if (is_text) + try serveTextPrompt(gpu.io, &tok_opt.?, gpa, prompt_src) + else + try servePrompt(gpa, gpu.io, prompt_src); + defer gpa.free(prompt); + if (prompt.len == 0) { + std.debug.print("FAIL minimax-m3-serve: the prompt carries no token ids\n", .{}); + return error.ServeEmptyPrompt; + } + + // -- the door: open()'s Loaded handle, the attached lane vtable, one greedy stream -- + var loaded = try m3.native.open(gpa, gpu.io, gpu.ctx, dir, null, .{ .context = m3.block_tokens, .drafts = false, .layer_cap = layers }); + defer loaded.deinit(loaded.ctx); + const b = loaded.backend; + var s = try m3.lane_api.Stream.init(gpa, .{ .id = "serve-front-door", .prompt = prompt, .max_new = @intCast(steps + 2), .drafts = false }); + defer s.deinit(gpa); + + try b.prefill(&s); + if (s.fill_at.load(.acquire) != s.prompt_len) return error.LaneFillLost; + var h = try b.first(&s, s.prompt_len); + const out = try gpa.alloc(u32, steps); + defer gpa.free(out); + out[0] = try b.read(h); + _ = try s.commit(gpa, &.{out[0]}, &.{}); + var idx: usize = 1; + while (idx < steps) : (idx += 1) { + const qpos: u64 = s.prompt_len + @as(u64, idx - 1); // the engine's own cache length + if (idx + 1 < steps) { + h = try b.queue(&s, .{ .value = out[idx - 1] }, qpos); + out[idx] = try b.read(h); + } else { + var sampled: [1]u32 = .{0}; + const windows = [_]m3.lane_api.backend.Window{.{ .stream = &s, .pending = out[idx - 1], .held = 0, .tokens = &.{}, .parents = null, .positions = &[_]u64{qpos} }}; + var verified = [_]m3.lane_api.backend.Verified{.{ .sampled = sampled[0..], .drafts = &.{} }}; + try b.verify(&windows, &verified); + out[idx] = verified[0].sampled[0]; + } + _ = try s.commit(gpa, &.{out[idx]}, &.{}); + } + std.debug.print("M3-SERVE: prompt {d} rows ->", .{prompt.len}); + for (out) |t| std.debug.print(" {d}", .{t}); + std.debug.print("\n", .{}); + + // -- Stage-10: the completion as readable text through the pack's own tokenizer -- + if (tok_opt) |*t| try serveTextOut(t, gpa, "serve", out); + + // -- the boundary bar: the receipted fixture recipe must land the Stage-6C/8 EXACT tokens -- + if (fixture_served and std.mem.eql(u8, prompt_src, "lcg") and prompt.len == 8 and steps >= 3 and layers == 2) { + const bar = [3]u32{ 183, 149, 14 }; + for (bar, 0..) |t, i| { + if (out[i] != t) { + std.debug.print("FAIL minimax-m3-serve: the front door drew {d} at draw {d} where the receipted fixture bar holds EXACT 183 149 14\n", .{ out[i], i }); + return error.TokenMismatch; + } + } + std.debug.print("[serve] fixture bar: EXACT 183 149 14\n", .{}); + } + + // -- the honest refusals, re-probed end to end at the door (exact names, cheap probes) -- + const last = out[steps - 1]; + var requests = [_]m3.lane_api.backend.DraftRequest{.{ .stream = &s, .follow = &.{last}, .rows = null, .start = 0, .position = 0, .depth = 1 }}; + try serveRefusal("draft", error.LanesNotAttached, b.draft(&requests)); + try serveRefusal("read(unknown handle)", error.LaneHandleUnknown, b.read(424242)); + var s2 = try m3.lane_api.Stream.init(gpa, .{ .id = "serve-second", .prompt = prompt, .max_new = 2, .drafts = false }); + defer s2.deinit(gpa); + try serveRefusal("prefill(second stream)", error.LanesNotAttached, b.prefill(&s2)); + + b.release(&s); // the sampled-spec refusal is receipted AFTER release (a live stream refuses first) + var s3 = try m3.lane_api.Stream.init(gpa, .{ .id = "serve-sampled", .prompt = prompt[0..2], .max_new = 1, .drafts = false, .sampling = .{ .seed = 7 } }); + defer s3.deinit(gpa); + try serveRefusal("prefill(sampled spec)", error.SamplingNotAttached, b.prefill(&s3)); + + std.debug.print("RESULT minimax-m3-serve: the front door answered {d} prompt rows with {d} greedy completion rows through open()'s lane surface (prefill/first/queue/read/verify); draft/second-stream/sampled refuse truthfully; fixture bar exact where receipted\n", .{ prompt.len, steps }); +} diff --git a/zig/tests/cuda/minimax_m3_tests.zig b/zig/tests/cuda/minimax_m3_tests.zig new file mode 100644 index 0000000000..179a5337bf --- /dev/null +++ b/zig/tests/cuda/minimax_m3_tests.zig @@ -0,0 +1,644 @@ +//! MiniMax M3.1's device parity: tf_m3_scores / tf_m3_topk / tf_m3_decode / tf_m3_decode_combine +//! and the family's tf_m3_swigluoai against the CPU references on synthetic tensors — the indexer's +//! sparse.zig contracts (fp32 block-max-pool scores, the boost override, ties at the smaller id, +//! -1 right-pad) and serve.zig's block-sparse fp32 attention; the NB <= 1024 / nsplit <= 8 bound +//! guards ride the same run. Prints max|delta| per case so receipts can quote them verbatim. + +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const m3 = @import("minimax_m3_vl_cuda"); + +const indexer = m3.indexer; // sparse.zig: the CPU reference the kernels are graded against +const sp = m3.sparse_cuda; // the launch wrappers +const glue = m3.serve_cuda; // the family's own glue (the swigluoai rows) + +const d = sp.d; // 128: the index head's width +const block = sp.block; // 128 keys a scored block +const topk = sp.topk; // 16: the fixed selection width +const page = sp.page; // 256: the fp16 cache's page + +fn f16r(x: f32) f32 { + return @floatCast(@as(f16, @floatCast(x))); +} + +/// The run's shared state: buffers collected for one freeing sweep at the end. +const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + + fn dev(r: *Rig, bytes: []const u8) !u64 { + const b = try cuda.DeviceBuffer.fromHost(r.gpu.d, bytes); + try r.bufs.append(r.a, b); + return b.ptr; + } + + fn uninit(r: *Rig, bytes: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, bytes); + try r.bufs.append(r.a, b); + try b.fill8(0xff, r.s.handle); // an entry the kernel misses reads NaN, never a stale zero + return b.ptr; + } + + fn back(r: *Rig, ptr: u64, len: usize) ![]u8 { + const o = try r.a.alloc(u8, len); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = len }; + try b.download(0, o); + return o; + } + + /// fp16 values a scorer or cache reads: the storage width of every indexed row. + fn h16(r: *Rig, n: usize, scale: f32) ![]f16 { + const v = try r.a.alloc(f16, n); + var prng = std.Random.DefaultPrng.init(0xd57c_a7e1); + const rng = prng.random(); + for (v) |*x| x.* = @floatCast((rng.float(f32) * 2 - 1) * scale); + return v; + } +}; + +fn backF16(bytes: []u8) []f16 { + return @as([*]f16, @ptrCast(@alignCast(bytes.ptr)))[0 .. bytes.len / 2]; +} + +fn backF32(bytes: []u8) []f32 { + return @as([*]f32, @ptrCast(@alignCast(bytes.ptr)))[0 .. bytes.len / 4]; +} + +/// max |got - want| plus class agreement (+inf stays +inf, -inf stays -inf, NaN stays NaN). +fn delta(got: []const f32, want: []const f32) struct { d: f64, classes: bool } { + var worst: f64 = 0; + var classes = true; + for (got, want) |g, w| { + if (std.math.isPositiveInf(g) != std.math.isPositiveInf(w) or + std.math.isNegativeInf(g) != std.math.isNegativeInf(w) or + (g != g) != (w != w)) classes = false; + if (std.math.isFinite(g) and std.math.isFinite(w)) { + const e = @abs(@as(f64, g) - @as(f64, w)); + if (e > worst) worst = e; + } + } + return .{ .d = worst, .classes = classes }; +} + +pub fn run(gpu: check.Gpu) !void { + const mod = try cuda.Module.load(gpu.d, cuda.kernels.minimax_m3); + const ks = try sp.Kernels.load(mod); + const gk = try glue.Kernels.load(mod); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var rig = Rig{ .gpu = gpu, .a = gpu.gpa, .s = s }; + + try scoresParity(&rig, ks); + try topkParity(&rig, ks, 300); + try marksParity(&rig, ks); // NB = 1024 inclusive: the wide-chunk ceiling runs, 1025 refuses + try decodeParity(&rig, ks, 4); + try decodeParity(&rig, ks, 1); + try zeroRowParity(&rig, ks); + try engineGeometryParity(&rig, ks); // the engine's exact failing invocation: the stage-6b regression + try engineSelectionParity(&rig, ks); // the engine's +inf-boosted scores -> real topk -> decode: the stage-6c regression + try swigluoaiParity(&rig, gk); + for (rig.bufs.items) |*b| b.free(); + rig.bufs.deinit(gpu.gpa); +} + +/// tf_m3_scores against sparse.scores on fp16 storage rows: fp32 scoring, the right-open causal +/// mask, -inf fully-masked blocks, the own-block +inf override, and the NaN class contract. +fn scoresParity(r: *Rig, ks: sp.Kernels) !void { + const a = r.a; + const h: usize = 4; // the pack's 4 index heads + const sq: usize = 5; + const scale_of_q: f32 = 0.4; + const cases = [_]usize{ 1, 129, 300 }; // a partial block, one over, several + for (cases) |sk| { + const nb = sp.nbOf(sk); + const q = try r.h16(h * sq * d, scale_of_q); + const k = try r.h16(sk * d, scale_of_q); + var qpos = try a.alloc(i64, sq); + var kpos = try a.alloc(i64, sk); + // query rows sit at positions 0, 3, 6, ... so every block's mask edge is exercised + for (0..sq) |ri| qpos[ri] = @intCast(3 * ri); + for (0..sk) |ki| kpos[ki] = @intCast(ki); // content keys, none padded + // the CPU reference over the SAME fp16-rounded storage values, f32 inside + const qf = try a.alloc(f32, q.len); + for (q, 0..) |v, i| qf[i] = @floatCast(v); + const kf = try a.alloc(f32, k.len); + for (k, 0..) |v, i| kf[i] = @floatCast(v); + const ref = try indexer.scores(a, qf, kf, qpos, kpos, h, sq, sk, d, block); + // the boost override the kernel applies at the write (local 1, init 0): the row's own block + for (0..h) |hi| for (0..sq) |ri| { + const b = @divTrunc(qpos[ri], @as(i64, @intCast(block))); + if (b >= 0 and b < @as(i64, @intCast(nb))) ref[(hi * sq + ri) * nb + @as(usize, @intCast(b))] = std.math.inf(f32); + }; + const qu = try r.dev(std.mem.sliceAsBytes(q)); + const ku = try r.dev(std.mem.sliceAsBytes(k)); + const qpu = try r.dev(std.mem.sliceAsBytes(qpos)); + const kpu = try r.dev(std.mem.sliceAsBytes(kpos)); + const ou = try r.uninit(h * sq * nb * 4); + try sp.scores(ks, r.s, qu, ku, qpu, kpu, ou, sq, nb, sk, h, 1, 1, 0); + const got_b = try r.back(ou, h * sq * nb * 4); + const got = backF32(got_b); + const res = delta(got, ref); + try check.expect(res.classes, "scores: a +/-inf/NaN class disagrees at sk={d}", .{sk}); + const tol = 1e-4 + 1e-6 * @as(f64, @floatFromInt(sk * d)); + try check.expect(res.d <= tol, "scores sk={d}: max|delta| {e} over tolerance {e}", .{ sk, res.d, tol }); + // the override itself, read back: every row's own block sits at +inf on the device + for (0..h) |hi| for (0..sq) |ri| { + const b: usize = @intCast(@divTrunc(qpos[ri], @as(i64, @intCast(block)))); + try check.expect(std.math.isPositiveInf(got[(hi * sq + ri) * nb + b]), "scores sk={d}: row {d} head {d} own block {d} not boosted", .{ sk, ri, hi, b }); + }; + check.pass("scores sk={d} nb={d}: max|delta| {e} (classes hold, override at +inf)", .{ sk, nb, res.d }); + } + // The NaN rule at a NON-own block: one NaN key element turns the block's dot (and so amax) + // NaN — the kernel's stated amax class (sparse.zig's CPU rows have no NaN producer; this pins + // the device class). The row's own block stays +inf: the boost lands after the reduction and + // replaces any pooled value there, NaN included. + { + const sk: usize = 300; // 3 blocks; the row sits in block 1, block 0 stays unboosted + const nb = sp.nbOf(sk); + const q = try a.alloc(f16, h * 1 * d); + const k = try a.alloc(f16, sk * d); + for (q) |*x| x.* = @floatCast(0.1); + for (k) |*x| x.* = @floatCast(0.1); + k[7] = std.math.nan(f16); // block 0's keys, valid under the causal mask (kpos 0 <= 200) + const qpos = [_]i64{200}; + const kpos = try a.alloc(i64, sk); + for (0..sk) |ki| kpos[ki] = @intCast(ki); + const qu = try r.dev(std.mem.sliceAsBytes(q)); + const ku = try r.dev(std.mem.sliceAsBytes(k)); + const qpu = try r.dev(std.mem.sliceAsBytes(&qpos)); + const kpu = try r.dev(std.mem.sliceAsBytes(kpos)); + const ou = try r.uninit(h * nb * 4); + try sp.scores(ks, r.s, qu, ku, qpu, kpu, ou, 1, nb, sk, h, 1, 1, 0); + const got_b = try r.back(ou, h * nb * 4); + const got = backF32(got_b); + for (0..h) |hi| { + try check.expect(got[hi * nb] != got[hi * nb], "scores: head {d} block 0 did not carry the NaN amax", .{hi}); // unboosted: NaN wins + try check.expect(std.math.isPositiveInf(got[hi * nb + 1]), "scores: head {d} own block is not +inf over the NaN", .{hi}); // the boost lands after reduction + } + check.pass("scores NaN rule: the unboosted block's amax carries NaN, the own block's boost still wins", .{}); + } +} + +/// tf_m3_topk against sparse.topk: synthetic score rows with exact ties, boost overrides already +/// in the table, descending picks, ties at the smaller id, fewer-valid rows padded -1. +fn topkParity(r: *Rig, ks: sp.Kernels, sk: usize) !void { + const a = r.a; + const h: usize = 4; + const sq: usize = 12; + const nb = sp.nbOf(sk); + // scores with quantized values so exact ties are common, and -inf pads behind the causal edge + var prng = std.Random.DefaultPrng.init(0x5c0a_e123); + const rng = prng.random(); + const s = try a.alloc(f32, h * sq * nb); + const qpos = try a.alloc(i64, sq); + for (0..sq) |ri| { + qpos[ri] = @intCast(if (ri == 0) 0 else ri * 31 % (sk - 1)); + for (0..h) |hi| for (0..nb) |b| { + const row = (hi * sq + ri) * nb; + const causal = @as(i64, @intCast((b + 1) * block)) <= qpos[ri]; + s[row + b] = if (causal) std.math.round(rng.float(f32) * 8) else -std.math.inf(f32); + }; + // the override: the row's own block, as sparse.topk would scatter it + const b = @divTrunc(qpos[ri], @as(i64, @intCast(block))); + if (b >= 0 and b < @as(i64, @intCast(nb))) { + for (0..h) |hi| s[(hi * sq + ri) * nb + @as(usize, @intCast(b))] = std.math.inf(f32); + } + } + const ref = try indexer.topk(a, s, h, sq, nb, topk, block, 1, 0, qpos); + const su = try r.dev(std.mem.sliceAsBytes(s)); + const ou = try r.uninit(h * sq * topk * 4); + try sp.topk16(ks, r.s, su, ou, nb, h * sq); + const got_b = try r.back(ou, h * sq * topk * 4); + const got = backI32(got_b); + try check.expect(std.mem.eql(i32, got, ref), "topk: pick ids disagree", .{}); + // two rows' trace quoted for the receipts + check.pass("topk rows={d} nb={d}: every pick equal (order, ties, -1 pads)", .{ h * sq, nb }); +} + +fn backI32(bytes: []u8) []i32 { + return @as([*]i32, @ptrCast(@alignCast(bytes.ptr)))[0 .. bytes.len / 4]; +} + +/// The mark bitmap's NB <= 1024 bound at its ceiling: a 131,072-key table (NB exactly 1024) +/// scores and selects; 1025 blocks is refused by the wrapper before any launch. +fn marksParity(r: *Rig, ks: sp.Kernels) !void { + const a = r.a; + const h: usize = 4; + const sq: usize = 2; + const sk: usize = block * sp.max_nb; // 131,072 + const nb = sp.nbOf(sk); + try check.expect(nb == sp.max_nb, "the wide table's block count is the mark bitmap's ceiling", .{}); + const q = try r.h16(h * sq * d, 0.4); + const k = try r.h16(sk * d, 0.4); + var qpos = [_]i64{ @as(i64, @intCast(sk - 1)), @as(i64, @intCast(sk / 2)) }; + var kpos = try a.alloc(i64, sk); + for (0..sk) |ki| kpos[ki] = @intCast(ki); + const qf = try a.alloc(f32, q.len); + for (q, 0..) |v, i| qf[i] = @floatCast(v); + const kf = try a.alloc(f32, k.len); + for (k, 0..) |v, i| kf[i] = @floatCast(v); + const ref = try indexer.scores(a, qf, kf, &qpos, kpos, h, sq, sk, d, block); + for (0..h) |hi| for (0..sq) |ri| { + const b: usize = @intCast(@divTrunc(qpos[ri], @as(i64, @intCast(block)))); + ref[(hi * sq + ri) * nb + b] = std.math.inf(f32); + }; + const qu = try r.dev(std.mem.sliceAsBytes(q)); + const ku = try r.dev(std.mem.sliceAsBytes(k)); + const kpu = try r.dev(std.mem.sliceAsBytes(kpos)); + const qpu = try r.dev(std.mem.sliceAsBytes(&qpos)); + const ou = try r.uninit(h * sq * nb * 4); + try sp.scores(ks, r.s, qu, ku, qpu, kpu, ou, sq, nb, sk, h, 1, 1, 0); + const got_b = try r.back(ou, h * sq * nb * 4); + const got = backF32(got_b); + const res = delta(got, ref); + try check.expect(res.classes, "scores nb=1024: a +/-inf class disagrees", .{}); + const tol = 1e-4 + 1e-6 * @as(f64, @floatFromInt(sk * d)); + try check.expect(res.d <= tol, "scores nb=1024: max|delta| {e} over {e}", .{ res.d, tol }); + const su = try r.dev(std.mem.sliceAsBytes(ref)); // the same table back through the selection + const pu = try r.uninit(h * sq * topk * 4); + try sp.topk16(ks, r.s, su, pu, nb, h * sq); + const p_b = try r.back(pu, h * sq * topk * 4); + const picks = backI32(p_b); + const want = try indexer.topk(a, ref, h, sq, nb, topk, block, 1, 0, &qpos); + try check.expect(std.mem.eql(i32, picks, want), "topk nb=1024: pick ids disagree", .{}); + check.pass("marks nb=1024 ceiling (sk={d}): scores max|delta| {e}, picks equal; nb=1025 refused below", .{ sk, res.d }); + const refused: anyerror = blk: { + sp.topk16(ks, r.s, 0, 0, sp.max_nb + 1, 1) catch |e| break :blk e; + break :blk error.MissingError; + }; + try check.expect(refused == error.ScoreColumns, "an nb of {d} was not refused ({s})", .{ sp.max_nb + 1, @errorName(refused) }); +} + +/// tf_m3_decode + tf_m3_decode_combine against serve.zig's block-sparse attention over the same +/// fp16 cache rows and selection lists — the gather, the causal seqlen mask, the fp32 online +/// softmax and split merge, fp16 out. nsplit runs 4 (the served envelope) and 1. +fn decodeParity(r: *Rig, ks: sp.Kernels, nsplit: u32) !void { + const a = r.a; + const hq: usize = 8; + const hkv: usize = 4; // index heads == kv heads; GQA group of 2 + const sk: usize = 1200; // 10 blocks, so the 16-wide lists carry -1 pads + const seqlen: i64 = 1100; // block 9 sits past the logical end: masked, never read + const pages = sk / page + 1; + const pages_of_seq = (sk + page - 1) / page; + const scale: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(d))); + var prng = std.Random.DefaultPrng.init(0xdec0_de55); + const rng = prng.random(); + const kt = try r.h16(pages * page * hkv * d, 0.6); + const vt = try r.h16(pages * page * hkv * d, 0.6); + const q = try r.h16(hq * d, 0.6); + // per kv head: a shuffled set of the 9 in-range blocks, then -1 pads + const sel = try a.alloc(i32, hkv * topk); + for (0..hkv) |hv| { + const list = sel[hv * topk ..][0..topk]; + for (0..topk) |i| list[i] = -1; + var perm = [_]i32{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9 }; + for (1..10) |i| { // a shuffle of the ids; the kernel never trusts the order + const j = rng.intRangeAtMost(usize, 0, i); + const t = perm[i]; + perm[i] = perm[j]; + perm[j] = t; + } + const count: usize = if (hv == 0) 9 else 10; + for (0..count) |i| list[i] = perm[i]; + } + // the reference: serve.zig's math at f32, from the same fp16-rounded storage rows + const kf = try a.alloc(f32, kt.len); + for (kt, 0..) |v, i| kf[i] = @floatCast(v); + const vf = try a.alloc(f32, vt.len); + for (vt, 0..) |v, i| vf[i] = @floatCast(v); + const qf = try a.alloc(f32, q.len); + for (q, 0..) |v, i| qf[i] = @floatCast(v); + const table = try a.alloc(i32, pages_of_seq); + const seqlens = [_]i64{seqlen}; + for (0..pages_of_seq) |p| table[p] = @intCast(p); // identity page table, 2 sequences' pages fresh + const ref = try a.alloc(f32, hq * d); + for (0..hq) |i| { + const hv = i / (hq / hkv); + const qr = qf[i * d ..][0..d]; + var mx: f32 = -std.math.inf(f32); + var scores = try a.alloc(f32, topk * block); + defer a.free(scores); + var n: usize = 0; + for (0..topk) |sl| { + const bid = sel[hv * topk + sl]; + if (bid < 0) break; + var j: usize = 0; + while (j < block) : (j += 1) { + const p: i64 = @as(i64, bid) * @as(i64, @intCast(block)) + @as(i64, @intCast(j)); + if (p >= seqlens[0]) continue; + const pg: usize = @intCast(@divTrunc(p, @as(i64, @intCast(page)))); + const off: usize = @intCast(@mod(p, @as(i64, @intCast(page)))); + const krow = kf[((pg * page + off) * hkv + hv) * d ..][0..d]; + var dot: f32 = 0; + for (0..d) |c| dot += qr[c] * krow[c]; + scores[n] = dot * scale; + if (scores[n] > mx) mx = scores[n]; + n += 1; + } + } + var se: f32 = 0; + for (0..n) |mi| { + scores[mi] = @exp(scores[mi] - mx); + se += scores[mi]; + } + const out = ref[i * d ..][0..d]; + @memset(out, 0); + var at: usize = 0; + for (0..topk) |sl| { + const bid = sel[hv * topk + sl]; + if (bid < 0) break; + var j: usize = 0; + while (j < block) : (j += 1) { + const p: i64 = @as(i64, bid) * @as(i64, @intCast(block)) + @as(i64, @intCast(j)); + if (p >= seqlens[0]) continue; + const pg: usize = @intCast(@divTrunc(p, @as(i64, @intCast(page)))); + const off: usize = @intCast(@mod(p, @as(i64, @intCast(page)))); + const vrow = vf[((pg * page + off) * hkv + hv) * d ..][0..d]; + const w = scores[at]; at += 1; + for (0..d) |c| out[c] += w / se * vrow[c]; + } + } + } + // the device run + const npp: usize = pages_of_seq; + const ktu = try r.dev(std.mem.sliceAsBytes(kt)); + const vtu = try r.dev(std.mem.sliceAsBytes(vt)); + const qcu = try r.dev(std.mem.sliceAsBytes(q)); + const tbu = try r.dev(std.mem.sliceAsBytes(table)); + const slu = try r.dev(std.mem.sliceAsBytes(sel)); + const seqs = try r.dev(std.mem.sliceAsBytes(&seqlens)); + const pu = try r.uninit(sp.scratchBytes(hq, 1, nsplit)); + const ou = try r.uninit(hq * d * 2); + try sp.decode(ks, r.s, qcu, ktu, vtu, tbu, slu, seqs, pu, hq, hkv, npp, scale, 1, nsplit); + try sp.combine(ks, r.s, pu, ou, hq, 1, nsplit); + const got_b = try r.back(ou, hq * d * 2); + const got = backF16(got_b); + var worst: f64 = 0; + for (0..hq) |i| for (0..d) |c| { + const g: f64 = @floatCast(got[i * d + c]); + const w0: f64 = @floatCast(f16r(ref[i * d + c])); // the reference's fp16 store too + const e = @abs(g - w0); + if (e > worst) worst = e; + }; + try check.expect(worst <= 5e-3, "decode nsplit={d}: max|delta| {e} over 5e-3", .{ nsplit, worst }); + check.pass("decode+combine nsplit={d}: attention max|delta| {e} vs the fp16-rounded reference", .{ nsplit, worst }); +} + +/// An all-(-1) selection is the zero-valid-key row: the combine returns zeros, never a NaN — +/// byte for byte, including the two heads a normal list keeps finite. +fn zeroRowParity(r: *Rig, ks: sp.Kernels) !void { + const a = r.a; + const hq: usize = 2; + const hkv: usize = 2; + const scale: f32 = 1.0 / @sqrt(128.0); + const kt = try r.h16(2 * page * hkv * d, 0.6); + const q = try r.h16(hq * d, 0.6); + const sel = try a.alloc(i32, hkv * topk); + @memset(sel, -1); + const table = try a.alloc(i32, 1); + table[0] = 0; + const seqlens = [_]i64{0}; + const ktu = try r.dev(std.mem.sliceAsBytes(kt)); + const vtu = ktu; // never read with l = 0, aliased harmlessly + const qcu = try r.dev(std.mem.sliceAsBytes(q)); + const tbu = try r.dev(std.mem.sliceAsBytes(table)); + const slu = try r.dev(std.mem.sliceAsBytes(sel)); + const seqs = try r.dev(std.mem.sliceAsBytes(&seqlens)); + const pu = try r.uninit(sp.scratchBytes(hq, 1, 4)); + const ou = try r.uninit(hq * d * 2); + try sp.decode(ks, r.s, qcu, ktu, vtu, tbu, slu, seqs, pu, hq, hkv, 1, scale, 1, 4); + try sp.combine(ks, r.s, pu, ou, hq, 1, 4); + const got_b = try r.back(ou, hq * d * 2); + for (got_b) |byte| try check.expect(byte == 0, "zero-valid-key row returned non-zero bytes", .{}); + check.pass("zero-valid-key row: all-zero fp16 bytes, never a NaN", .{}); +} + +/// The ENGINE's exact failing-invocation geometry (graph.zig layerRow's sp.decode/sp.combine), the +/// stage-6b regression for the max|a|=947 hunt: the fixture's dims (Hq 64, Hkv 4, D 128), the +/// layer-strided [nl*keys_cap][Hkv][D] cache whose L0 base sits at +0 (so the kernel's +/// [(P*PAGE)+off][Hkv][D] row addressing lines up), the 16-entry identity table, sel +/// (B, Hkv, 16) = [0, -1 x 15], seqlens [1] (ONE valid key), sp.scratchBytes(Hq, 1, 4), nsplit 4, +/// and the magnitudes the smoke's engine measured (max|q| 7.78 fp16, key row 0 to 6.07, v row 0 +/// between -497.75 and 631.5). With one valid key the softmax weight is 1: every combine row must +/// equal its kv head's v row bit for bit, split 0's scratch must carry m = the q·k score with +/// l = 1 and acc = the v row, and the empty splits must hold m = -inf, acc = 0. +fn engineGeometryParity(r: *Rig, ks: sp.Kernels) !void { + const a = r.a; + const hq: usize = 64; // the fixture's attn_heads + const hkv: usize = 4; // the fixture's kv_heads; GQA group of 16 + const kvr: usize = hkv * d; // 512 + const rows: usize = 16 * page; // keys_cap 4096: (P=16, PAGE=256, Hkv, D) at the L0 base + const scale: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(d))); + var prng = std.Random.DefaultPrng.init(0x6b5e_947a); + const rng = prng.random(); + // q: noise at the engine's scale with the smoke's fp16 max in row 0 + const q = try a.alloc(f16, hq * d); + for (q, 0..) |*x, i| { + const s: f32 = if (i == 0) 7.7773 else (rng.float(f32) * 2 - 1) * 4.0; + x.* = @floatCast(s); + } + // the cache: key row 0 at the smoke's max 6.07, v row 0 within [min -497.75, max 631.5], + // the rest noise on the same scales (rows >= 1 are masked; they may be anything) + const kt = try a.alloc(f16, rows * kvr); + const vt = try a.alloc(f16, rows * kvr); + for (kt) |*x| x.* = @floatCast((rng.float(f32) * 2 - 1) * 3.0); + for (vt) |*x| x.* = @floatCast((rng.float(f32) * 2 - 1) * 60.0); + for (0..d) |i| { + kt[0 * kvr + i] = @floatCast((rng.float(f32) * 2 - 1) * 6.0664); + const vv: f32 = @round((rng.float(f32) * 2 - 1) * 497.75); + vt[0 * kvr + i] = @floatCast(if (i == 0) 631.5 else vv); + } + const sel = try a.alloc(i32, hkv * topk); + for (0..hkv) |hv| { + const list = sel[hv * topk ..][0..topk]; + list[0] = 0; + for (1..topk) |i| list[i] = -1; + } + const table = try a.alloc(i32, 16); + for (0..16) |p| table[p] = @intCast(p); // the engine's identity upload + const seqlens = [_]i64{1}; + // the reference: single valid key, weight 1 — out is the fp16 v row; split 0's m is the score + const kf = try a.alloc(f32, kt.len); + for (kt, 0..) |v, i| kf[i] = @floatCast(v); + const vf = try a.alloc(f32, vt.len); + for (vt, 0..) |v, i| vf[i] = @floatCast(v); + const qf = try a.alloc(f32, q.len); + for (q, 0..) |v, i| qf[i] = @floatCast(v); + const want = try a.alloc(f16, hq * d); + const m0 = try a.alloc(f32, hq); // split 0's m per head, from the same fp16 rows + for (0..hq) |i| { + const hv = i / (hq / hkv); + for (0..d) |c| want[i * d + c] = @floatCast(vf[0 * kvr + hv * d + c]); + var dot: f32 = 0; + const k0 = kf[0 * kvr + hv * d ..][0..d]; + for (0..d) |c| dot += qf[i * d + c] * k0[c]; + m0[i] = dot * scale; + } + // the device run: the engine's layout and args, verbatim + const ktu = try r.dev(std.mem.sliceAsBytes(kt)); + const vtu = try r.dev(std.mem.sliceAsBytes(vt)); + const qcu = try r.dev(std.mem.sliceAsBytes(q)); + const tbu = try r.dev(std.mem.sliceAsBytes(table)); + const slu = try r.dev(std.mem.sliceAsBytes(sel)); + const seqs = try r.dev(std.mem.sliceAsBytes(&seqlens)); + const pu = try r.uninit(sp.scratchBytes(hq, 1, 4)); + const ou = try r.uninit(hq * d * 2); + try sp.decode(ks, r.s, qcu, ktu, vtu, tbu, slu, seqs, pu, hq, hkv, 16, scale, 1, 4); + try sp.combine(ks, r.s, pu, ou, hq, 1, 4); + const got_b = try r.back(ou, hq * d * 2); + const got = backF16(got_b); + const pb = try r.back(pu, sp.scratchBytes(hq, 1, 4)); + const parts = backF32(pb); + const sestr: usize = d + 2; // one (split, head)'s scratch: the 128 acc lanes then m, l + var worst: usize = 0; // mismatching fp16 bits + for (0..hq * d) |i| { + if (got[i] != want[i]) worst += 1; + } + try check.expect(worst == 0, "engine geometry: {d} of {d} combine bits differ from the v rows", .{ worst, hq * d }); + var bad: usize = 0; + for (0..hq) |hh| { + const row = hh * 4 * sestr; // (B, Hq, nsplit, D+2), B = 1 + var m0e: f64 = 0; + for (0..d) |c| { + const want0: f64 = @floatCast(want[hh * d + c]); // acc = p * v with p = 1 + m0e = @max(m0e, @abs(@as(f64, parts[row + c]) - want0)); + } + try check.expect(m0e == 0, "head {d} split 0 acc != v ({e})", .{ hh, m0e }); + try check.expect(@abs(@as(f64, parts[row + sestr - 1]) - 1.0) < 1e-7, "head {d} l != 1", .{hh}); + const md = @abs(@as(f64, parts[row + sestr - 2]) - @as(f64, m0[hh])); + try check.expect(md <= 2e-5, "head {d} m {e} off by {e}", .{ hh, m0[hh], md }); + for (1..4) |spp| { + const srow = row + spp * sestr; + try check.expect(std.math.isNegativeInf(parts[srow + sestr - 2]), "head {d} split {d} m not -inf", .{ hh, spp }); + try check.expect(parts[srow + sestr - 1] == 0, "head {d} split {d} l not 0", .{ hh, spp }); + for (srow..srow + d) |c| if (parts[c] != 0) { + bad += 1; + break; + }; + } + } + try check.expect(bad == 0, "{d} empty-split heads carry non-zero acc", .{bad}); + for (got) |g| try check.expect(!std.math.isNan(g), "combine wrote a NaN at the engine geometry", .{}); + check.pass("engine geometry (Hq 64, sel [0,-1x15], seqlens 1, layer-strided L0, nsplit 4): combine == v rows bit-for-bit, scratch m/l/acc exact, empty splits clean", .{}); +} + +/// The engine's selection path end to end — the stage-6c regression for the delegated open corner: +/// tf_m3_scores at the engine's NB=1 shape producing the own-block +inf boost, tf_m3_topk at the +/// engine's consumer rows (i_heads == kv_heads is the config's admission assert), that +inf-boosted +/// score list then driving decode+combine exactly as the engine hands the device buffer over. The +/// boost cannot make topk pathological: the own block wins pick 0, its per-pick mark invalidates +/// it, and the exhausted picks right-pad -1 — the single valid key then returns each head's v row. +fn engineSelectionParity(r: *Rig, ks: sp.Kernels) !void { + const a = r.a; + const hq: usize = 64; // the fixture's attn_heads + const hkv: usize = 4; // kv_heads; the indexer's i_heads == kv_heads (the admission assert) + const ih: usize = 4; // the indexer's heads: the scores table's rows + const kvr: usize = hkv * d; + const rows: usize = 16 * page; // keys_cap 4096 at the L0 base: (P=16, PAGE=256, Hkv, D) + const sk: usize = 1; + const nb: usize = sp.nbOf(sk); // 1: the query's own block carries the +inf boost + const scale: f32 = 1.0 / @sqrt(@as(f32, @floatFromInt(d))); + var prng = std.Random.DefaultPrng.init(0x7e57_ca7e); + const rng = prng.random(); + // the indexer's side: fp16 q rows (the smoke's max) and the one fp16 plane row pos 0 appends + const iq = try a.alloc(f16, ih * d); + for (iq, 0..) |*x, i| { + const s: f32 = if (i == 0) 2.5867 else 2.6; + x.* = @floatCast((rng.float(f32) * 2 - 1) * s); + } + const plane_row = try a.alloc(f16, sk * d); + for (plane_row) |*x| x.* = @floatCast((rng.float(f32) * 2 - 1) * 1.2); + const qpos = [_]i64{0}; + const kpos = try a.alloc(i64, rows); // the engine's identity 0..cap-1 upload + for (kpos, 0..) |*x, i| x.* = @intCast(i); + // the attention side at the engine's shapes (the same scales engineGeometryParity grades) + const q = try a.alloc(f16, hq * d); + for (q) |*x| x.* = @floatCast((rng.float(f32) * 2 - 1) * 4.0); + const kt = try a.alloc(f16, rows * kvr); + const vt = try a.alloc(f16, rows * kvr); + for (kt) |*x| x.* = @floatCast((rng.float(f32) * 2 - 1) * 3.0); + for (vt) |*x| x.* = @floatCast((rng.float(f32) * 2 - 1) * 60.0); + for (0..d) |i| vt[0 * kvr + i] = @floatCast((rng.float(f32) * 2 - 1) * 60.0); + const table = try a.alloc(i32, 16); + for (0..16) |p| table[p] = @intCast(p); // the engine's identity page table + const seqlens = [_]i64{1}; + // the reference: the produced list's only valid key is block 0's p=0, weight 1 — each head's + // combine row is its kv head's v row (GQA hv = hq/16), the same fp16 rows the cache stores + const want = try a.alloc(f16, hq * d); + for (0..hq) |i| { + const hv = i / (hq / hkv); + for (0..d) |c| want[i * d + c] = vt[0 * kvr + hv * d + c]; + } + // scores -> topk -> decode -> combine, the device chain in the engine's order and shapes + const iqu = try r.dev(std.mem.sliceAsBytes(iq)); + const plu = try r.dev(std.mem.sliceAsBytes(plane_row)); + const qpu = try r.dev(std.mem.sliceAsBytes(&qpos)); + const kpu = try r.dev(std.mem.sliceAsBytes(kpos)); + const scu = try r.uninit(ih * nb * 4); + try sp.scores(ks, r.s, iqu, plu, qpu, kpu, scu, 1, nb, sk, ih, 1, 1, 0); + const sc = backF32(try r.back(scu, ih * nb * 4)); + for (sc) |s| try check.expect(std.math.isPositiveInf(s), "the pos-0 own block must score +inf, got {d}", .{s}); + const selu = try r.uninit(hkv * topk * 4); + try sp.topk16(ks, r.s, scu, selu, nb, hkv); // the engine's topk rows = kv_heads over the same table + const sel = backI32(try r.back(selu, hkv * topk * 4)); + for (0..hkv) |hv| for (0..topk) |ti| { + const want_v: i32 = if (ti == 0) 0 else -1; + try check.expect(sel[hv * topk + ti] == want_v, "sel[{d}][{d}] = {d}, want {d}: the +inf boost distorted the picks", .{ hv, ti, sel[hv * topk + ti], want_v }); + }; + const ktu = try r.dev(std.mem.sliceAsBytes(kt)); + const vtu = try r.dev(std.mem.sliceAsBytes(vt)); + const qcu = try r.dev(std.mem.sliceAsBytes(q)); + const tbu = try r.dev(std.mem.sliceAsBytes(table)); + const seqs = try r.dev(std.mem.sliceAsBytes(&seqlens)); + const pu = try r.uninit(sp.scratchBytes(hq, 1, 4)); // the engine's scratch sizing + const ou = try r.uninit(hq * d * 2); + try sp.decode(ks, r.s, qcu, ktu, vtu, tbu, selu, seqs, pu, hq, hkv, 16, scale, 1, 4); + try sp.combine(ks, r.s, pu, ou, hq, 1, 4); + const got = backF16(try r.back(ou, hq * d * 2)); + var worst: usize = 0; + for (0..hq * d) |i| { + if (got[i] != want[i]) worst += 1; + } + try check.expect(worst == 0, "engine selection: {d} of {d} combine bits differ from the v rows", .{ worst, hq * d }); + check.pass("engine selection (scores +inf boost -> topk rows=kv_heads -> decode/combine nsplit 4, NB=1 sk=1): sel == [0,-1x15] per kv head, combine == v rows bit-for-bit", .{}); +} + +/// tf_m3_swigluoai against serve.zig's swigluoai at the pack's alpha 1.702 / limit 7: fp16 rows, +/// clamps above-only (g) and both-ways (u), the fast-exp sigmoid's class allowed its few ulps. +fn swigluoaiParity(r: *Rig, ks: glue.Kernels) !void { + const a = r.a; + const rows: usize = 17; + const n: usize = 1536; // not a 512 multiple boundary case: a partial block of the 4-elem lanes + var prng = std.Random.DefaultPrng.init(0x1a2b_3c4d); + const rng = prng.random(); + const g = try a.alloc(f16, rows * n); + const u = try a.alloc(f16, rows * n); + for (g) |*x| { + const v = if (rng.boolean()) rng.float(f32) * 20 - 4 else rng.float(f32) - 0.5; // heavy at the clamp edge + x.* = @floatCast(v); + } + for (u) |*x| { + const v = if (rng.boolean()) rng.float(f32) * 30 - 15 else rng.float(f32) * 2 - 1; + x.* = @floatCast(v); + } + const want = try a.alloc(f32, rows * n); + defer a.free(want); + for (0..rows * n) |i| want[i] = m3.serve.swigluoai(@floatCast(g[i]), @floatCast(u[i]), 1.702, 7.0); + const gu = try r.dev(std.mem.sliceAsBytes(g)); + const uu = try r.dev(std.mem.sliceAsBytes(u)); + const ou = try r.uninit(rows * n * 2); + try glue.swigluoai(ks, r.s, gu, uu, ou, rows, n, 7.0); + const got_b = try r.back(ou, rows * n * 2); + const got = backF16(got_b); + var worst: f64 = 0; + for (0..rows * n) |i| { + const e = @abs(@as(f64, @floatCast(got[i])) - @as(f64, @floatCast(f16r(want[i])))); + if (e > worst) worst = e; + } + try check.expect(worst <= 5e-3, "swigluoai: max|delta| {e} over 5e-3", .{worst}); + check.pass("swigluoai rows={d} n={d}: max|delta| {e} vs serve.swigluoai at the fp16 store", .{ rows, n, worst }); +} \ No newline at end of file diff --git a/zig/tests/cuda/minimax_tp2_tests.zig b/zig/tests/cuda/minimax_tp2_tests.zig new file mode 100644 index 0000000000..60a9e77e46 --- /dev/null +++ b/zig/tests/cuda/minimax_tp2_tests.zig @@ -0,0 +1,351 @@ +//! MiniMax M3.1's TP2 weight plumbing over the REAL EXL3 pack (Stage-5, part 1 — receipts only): +//! both ranks read the same checkpoint through admission + the whole Stage-2 plan, shard.zig lays +//! the load out (rank 0: layers 0..N/2 plus the embed/final-norm/head; rank 1: the rest), and every +//! owned EXL3 group and passthrough tensor uploads to device buffers through the mainline +//! exl3_upload.Uploader — the trellis words and scales the exl3 kernels read in place, no decode on +//! upload. The ranks then allgather their checkpoint-byte shares over NCCL: their sum must equal the +//! whole plan's bytes exactly. Prints per-class bytes, upload seconds, device-buffer totals and free +//! RAM per rank — the Stage-5 memory receipts. No forward here: the launch graph is the remaining +//! attach (cuda_native.zig's boundary words it), so `open()` still refuses and must. +//! +//! tf-cuda-test minimax-m3-tp2 +//! +//! rank 0 serves the NCCL id on host:port (the exl3-tp2 precedent); the harness passes rank/addr +//! explicitly — the serving form reads the TF_MINIMAXM3_RANK/TF_MINIMAXM3_WORLD env (shard.zig). + +const std = @import("std"); +const builtin = @import("builtin"); +const posix = std.posix; +const cuda = @import("cuda"); +const check = @import("check.zig"); +const m3 = @import("minimax_m3_vl_cuda"); +const exl3_load = @import("core").exl3_load; + +const exl3_upload = cuda.exl3_upload; + +/// The swigluoai contract's clamp (serve.zig's `limit`), staged into the Layer metadata; this stage +/// never executes the experts kernel — the family's tf_m3_swigluoai epilogue drives it when attached. +const limit: f32 = 7.0; +/// fp32 epilogue accumulation (M3's contract; 0 would be GLM's bf16-rounding route). +const act_mode: i32 = 1; + +pub fn run(gpu: check.Gpu, rank_arg: []const u8, addr_arg: []const u8, dir: []const u8) !void { + if (dir.len == 0 or std.mem.eql(u8, dir, "none")) { + if (std.mem.eql(u8, rank_arg, "0")) std.debug.print("SKIP minimax.m3.tp2: give the real EXL3 pack dir (rank 1's dir arg is ignored)\n", .{}); + return; + } + const rank = std.fmt.parseInt(u8, rank_arg, 10) catch { + std.debug.print("usage: tf-cuda-test minimax-m3-tp2 \n", .{}); + return error.BadArguments; + }; + if (rank > 1 or addr_arg.len == 0) { + std.debug.print("usage: tf-cuda-test minimax-m3-tp2 \n", .{}); + return error.BadArguments; + } + + // -- the TP2 boot: the GLM lane's recipe, harness form (rank 0 serves the NCCL id on host:port) -- + var lib = try cuda.nccl.Library.open(); + defer lib.close(); + var id: cuda.nccl.UniqueId = undefined; + const id_bytes = std.mem.asBytes(&id); + if (rank == 0) { + try lib.check(lib.api.ncclGetUniqueId(&id), "ncclGetUniqueId"); + try serve(addr_arg, id_bytes); + } else { + try receive(gpu, addr_arg, id_bytes); + } + var comm: cuda.nccl.Comm = null; + try lib.check(lib.api.ncclCommInitRank(&comm, 2, id, @intCast(rank)), "ncclCommInitRank"); + defer _ = lib.api.ncclCommDestroy(comm); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + + // -- the whole plan, from the family's own admission (config + trellis widths + every tensor) -- + const t0 = check.now(gpu.io); // monotonic nanoseconds + var admitted = try m3.admitDir(gpu.gpa, gpu.io, dir); + defer admitted.deinit(); + const admit_s = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + const c = &admitted.c; + const load = &admitted.load; + const groups = admitted.groups.groups; + + const device_before = cuda.usage(false).device; + const ram_before = hostFree(gpu); + + // -- the reader: one checkpoint scan, group prefix -> core group index (parts come read-only) -- + var scan = try exl3_load.scan(gpu.gpa, gpu.io, dir); + defer scan.deinit(); + var idx = std.StringHashMapUnmanaged(usize).empty; + defer idx.deinit(gpu.gpa); + for (scan.groups, 0..) |g, i| try idx.put(gpu.gpa, g.prefix, i); + + // -- the upload, item by item, through the mainline Uploader (weights the exl3 kernels read) -- + var arena_state = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena_state.deinit(); + const a = arena_state.allocator(); + + // routed experts collect per owned MoE layer, then upload as one layer() call each + const Projs = std.ArrayListUnmanaged(exl3_upload.Proj); + const LayerProjs = struct { gate: Projs, up: Projs, down: Projs }; + var exps = try a.alloc(?LayerProjs, c.layers); + for (exps) |*e| e.* = null; + + var owned_items: usize = 0; + var owned_bytes: u64 = 0; + var counts = [_]Stat{ .{}, .{}, .{}, .{}, .{}, .{} }; // attn, dense_mlp, shared, routed, head, passthrough + var max_buffers: usize = 64; + for (load.items) |*item| { + if (!m3.rank_plan.owns(item, c.layers, 2, @intCast(rank))) continue; + owned_items += 1; + owned_bytes += item.bytes; + max_buffers += item.bytes / (64 << 20) + @as(usize, switch (item.kind) { + .group => 4, // trellis + two scales + bias + .passthrough => 1, + }); + } + var up: exl3_upload.Uploader = .{ .d = gpu.d, .gpa = gpu.gpa, .max_buffers = max_buffers }; + defer up.deinit(); + + const t1 = check.now(gpu.io); + for (load.items) |*item| { + if (!m3.rank_plan.owns(item, c.layers, 2, @intCast(rank))) continue; + switch (item.kind) { + .passthrough => { + const t = try scan.ck.get(item.name); + _ = try up.dev(t.bytes); + counts[5].add(item.bytes); + }, + .group => |gi| { + const fg = groups[gi]; // the family's classified group + const core_gi = idx.get(item.name) orelse { + std.debug.print("FAIL minimax.m3.tp2 rank {d}: no core group for {s}\n", .{ rank, item.name }); + return error.MissingTensor; + }; + const parts = try scan.parts(a, core_gi); + if (fg.class == .routed) { + const slot = &exps[fg.layer]; + if (slot.* == null) slot.* = .{ .gate = .empty, .up = .empty, .down = .empty }; + const proj: *Projs = switch (fg.proj) { + 0 => &slot.*.?.gate, + 1 => &slot.*.?.up, + else => &slot.*.?.down, + }; + try proj.append(a, .{ + .trellis = parts.trellis.bytes, + .suh = parts.suh.bytes, + .svh = parts.svh.bytes, + .k2 = fg.k2(), + .signs = parts.suh_packed, + }); + counts[3].add(item.bytes); + } else { + _ = try up.dense(parts, fg); + counts[switch (fg.class) { + .attn => 0, + .dense_mlp => 1, + .shared => 2, + .head => 4, + else => unreachable, // .routed handled above + }].add(item.bytes); + } + }, + } + } + // each owned MoE layer's collected experts upload as one grouped layer (the k2 widths are the + // headers' half-bit counts, packed su/sv scales expand in the uploader) + var moe_layers: usize = 0; + const dims: usize = c.hidden; + const width: usize = firstOwnedWidth(groups, c.layers, rank) orelse { + std.debug.print("FAIL minimax.m3.tp2 rank {d}: no owned routed gate group to read the act width from\n", .{rank}); + return error.BadExperts; + }; + for (exps, 0..) |*e, d| { + if (e.*) |*lp| { + if (lp.gate.items.len != c.experts or lp.up.items.len != c.experts or lp.down.items.len != c.experts) { + std.debug.print("FAIL minimax.m3.tp2 rank {d}: layer {d} collected gate {d} up {d} down {d}, expected {d} experts each\n", .{ rank, d, lp.gate.items.len, lp.up.items.len, lp.down.items.len, c.experts }); + return error.BadExperts; + } + if (d == c.dense_layers) { // the pack's first MoE layer proves the projections' roles + const gate = lp.gate.items[0]; + const down = lp.down.items[0]; + if (gate.suh.len != dims * 2 or down.suh.len == 0 or down.svh.len != dims * 2) { + std.debug.print("FAIL minimax.m3.tp2 rank {d}: layer {d} expert scales do not spell hidden {d}\n", .{ rank, d, dims }); + return error.BadExperts; + } + } + _ = try exl3_upload.layer(&up, a, dims, width, .mul1, limit, act_mode, lp.gate.items, lp.up.items, lp.down.items); + moe_layers += 1; + } + } + const upload_s = @as(f64, @floatFromInt(check.now(gpu.io) - t1)) / 1e9; + + // -- the collective receipt: the shares allgather, their sum must be the whole plan's bytes -- + const mine: u64 = owned_bytes; + var send = try cuda.DeviceBuffer.fromHost(gpu.d, std.mem.sliceAsBytes(&[1]u64{mine})); + defer send.free(); + var recv = try cuda.DeviceBuffer.alloc(gpu.d, 16); + defer recv.free(); + try lib.check(lib.api.ncclAllGather(send.ptr, recv.ptr, 1, .u64, comm, s.handle), "ncclAllGather"); + try s.synchronize(); + const got = try a.alloc(u8, 16); + try recv.download(0, got); + // the allgather's slots are ordered by rank: slot r holds rank r's share, so the peer's share + // sits at slot (1 - rank) — rank 0 reads slot 1, rank 1 reads slot 0 + const peer_slot: usize = if (rank == 0) 8 else 0; + var view: [8]u8 = undefined; + @memcpy(&view, got[peer_slot..][0..8]); + const other = std.mem.readInt(u64, &view, .little); + const owned_moe_expected: usize = blk: { + const r = m3.rank_plan.range(c.layers, 2, @intCast(rank)); + var n: usize = 0; + for (r.first..r.first + r.count) |l| { + if (l >= c.dense_layers) n += 1; + } + break :blk n; + }; + const whole = load.bytes(); + + const device_after = cuda.usage(false).device; + const ram_after = hostFree(gpu); + const gib = struct { + fn f(v: u64) f64 { + return @as(f64, @floatFromInt(v)) / (1024.0 * 1024.0 * 1024.0); + } + }.f; + std.debug.print("M3.TP2 rank {d}: admission {d:.1}s, groups {d}, plan items {d}, whole ckpt {d} bytes ({d:.2} GiB)\n", .{ rank, admit_s, groups.len, load.items.len, whole, gib(whole) }); + std.debug.print("M3.TP2 rank {d}: owns {d} items, {d} ckpt bytes ({d:.2} GiB), layers [{d}, {d})\n", .{ rank, owned_items, mine, gib(mine), m3.rank_plan.range(c.layers, 2, @intCast(rank)).first, m3.rank_plan.range(c.layers, 2, @intCast(rank)).first + m3.rank_plan.range(c.layers, 2, @intCast(rank)).count }); + for ([_][]const u8{ "attn", "dense_mlp", "shared", "routed", "head", "passthrough" }, &counts) |name, *st| { + std.debug.print("M3.TP2 rank {d}: {s:>11} items {d:>5} ckpt bytes {d:>14} ({d:.2} GiB)\n", .{ rank, name, st.items, st.bytes, gib(st.bytes) }); + } + std.debug.print("M3.TP2 rank {d}: moe layers grouped {d} of {d} owned expected ({d} across both ranks); experts per layer {d}\n", .{ rank, moe_layers, owned_moe_expected, c.layers - c.dense_layers, c.experts }); + std.debug.print("M3.TP2 rank {d}: layer() width {d}; Uploader device buffers {d} bytes ({d:.2} GiB)\n", .{ rank, width, up.bytes, gib(up.bytes) }); + std.debug.print("M3.TP2 rank {d}: upload took {d:.1}s; cuda device usage delta {d} bytes ({d:.2} GiB); host MemAvailable {d:.2} -> {d:.2} GiB\n", .{ rank, upload_s, device_after - device_before, gib(device_after - device_before), gib(ram_before), gib(ram_after) }); + std.debug.print("M3.TP2 rank {d}: shares {d} + {d} = {d} vs whole {d}\n", .{ rank, mine, other, mine + other, whole }); + if (mine + other != whole) { + if (rank == 0) std.debug.print("FAIL minimax.m3.tp2: the ranks' shares do not tile the whole plan\n", .{}); + return error.TestFailed; + } + if (moe_layers != owned_moe_expected) { + if (rank == 0) std.debug.print("FAIL minimax.m3.tp2: not every owned MoE layer grouped its experts\n", .{}); + return error.TestFailed; + } + if (rank == 0) std.debug.print("RESULT minimax.m3.tp2: both ranks' shares tile the plan; every owned group and passthrough tensor is on device buffers\n", .{}); +} + +/// The grouped experts' act width: the first owned layer's routed gate group's n (w1/w2's outputs, +/// w3's inputs — the same number the grouped kernel's epilogue holds). +fn firstOwnedWidth(groups: anytype, layers: u32, rank: u8) ?u32 { + const r = m3.rank_plan.range(layers, 2, @intCast(rank)); + for (groups) |g| { + if (g.class != .routed or g.proj != 0) continue; + if (g.layer >= r.first and g.layer < r.first + r.count) return g.n; + } + return null; +} + +const Stat = struct { + items: u64 = 0, + bytes: u64 = 0, + + fn add(st: *Stat, b: u64) void { + st.items += 1; + st.bytes += b; + } +}; + +/// Free MemAvailable from /proc/meminfo (0 elsewhere — the receipt prints GB10 nodes' figures). +fn hostFree(gpu: check.Gpu) u64 { + var f = std.Io.Dir.openFileAbsolute(gpu.io, "/proc/meminfo", .{}) catch return 0; + defer f.close(gpu.io); + var buf: [8192]u8 = undefined; + const n = f.readPositionalAll(gpu.io, &buf, 0) catch return 0; + const line = std.mem.indexOf(u8, buf[0..n], "MemAvailable:") orelse return 0; + var it = std.mem.tokenizeScalar(u8, buf[line + "MemAvailable:".len .. n], ' '); + const kb = it.next() orelse return 0; + return (std.fmt.parseInt(u64, kb, 10) catch 0) * 1024; // /proc/meminfo reports kB, not bytes +} + +// --- the NCCL id exchange, verbatim the exl3-tp2 precedent (rank 0 serves on host:port, rank 1 reads) --- + +fn serve(addr_text: []const u8, bytes: []const u8) !void { + if (builtin.os.tag == .windows) return error.SocketFailed; + const at = std.mem.lastIndexOfScalar(u8, addr_text, ':') orelse return error.BadAddress; + const ip = try std.Io.net.IpAddress.parseIp4(addr_text[0..at], try std.fmt.parseInt(u16, addr_text[at + 1 ..], 10)); + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&ip, &storage); + const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const fd: posix.socket_t = @intCast(rc); + defer _ = posix.system.close(fd); + const one: c_int = 1; + try posix.setsockopt(fd, posix.SOL.SOCKET, posix.SO.REUSEADDR, std.mem.asBytes(&one)); + switch (posix.errno(posix.system.bind(fd, &storage.any, len))) { + .SUCCESS => {}, + .ADDRINUSE => return error.AddressInUse, + else => return error.BindFailed, + } + if (posix.errno(posix.system.listen(fd, 1)) != .SUCCESS) return error.ListenFailed; + var waited: u64 = 0; + while (true) { + var fds = [_]posix.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const ready = posix.poll(&fds, 500) catch return error.AcceptFailed; + if (ready > 0) break; + waited += 500; + if (waited > 70_000) return error.PeerTimeout; + } + var peer: std.Io.Threaded.PosixAddress = undefined; + var peer_len: posix.socklen_t = @sizeOf(std.Io.Threaded.PosixAddress); + const rc2 = posix.system.accept(fd, &peer.any, &peer_len); + if (posix.errno(rc2) != .SUCCESS) return error.AcceptFailed; + const cfd: posix.socket_t = @intCast(rc2); + defer _ = posix.system.close(cfd); + var off: usize = 0; + while (off < bytes.len) { + const n = posix.system.write(cfd, bytes.ptr + off, bytes.len - off); + if (posix.errno(n) != .SUCCESS) return error.WriteFailed; + off += @intCast(n); + } +} + +fn receive(gpu: check.Gpu, addr_text: []const u8, bytes: []u8) !void { + if (builtin.os.tag == .windows) return error.SocketFailed; + const at = std.mem.lastIndexOfScalar(u8, addr_text, ':') orelse return error.BadAddress; + const ip = try std.Io.net.IpAddress.parseIp4(addr_text[0..at], try std.fmt.parseInt(u16, addr_text[at + 1 ..], 10)); + var fd: posix.socket_t = -1; + defer { + if (fd >= 0) _ = posix.system.close(fd); + } + while (fd < 0) { + var storage: std.Io.Threaded.PosixAddress = undefined; + const len = std.Io.Threaded.addressToPosix(&ip, &storage); + const rc = posix.system.socket(posix.AF.INET, posix.SOCK.STREAM, 0); + if (posix.errno(rc) != .SUCCESS) return error.SocketFailed; + const c: posix.socket_t = @intCast(rc); + if (posix.errno(posix.system.connect(c, &storage.any, len)) == .SUCCESS) { + fd = c; + } else { + _ = posix.system.close(c); + std.Io.sleep(gpu.io, .fromMilliseconds(200), .awake) catch {}; + } + } + var off: usize = 0; + var read_waited: u64 = 0; + while (off < bytes.len) { + var fds = [_]posix.pollfd{.{ .fd = fd, .events = posix.POLL.IN, .revents = 0 }}; + const ready = posix.poll(&fds, 500) catch return error.ReadFailed; + if (ready == 0) { + read_waited += 500; + if (read_waited > 60_000) return error.PeerTimeout; + continue; + } + const n = posix.system.read(fd, bytes.ptr + off, bytes.len - off); + switch (posix.errno(n)) { + .SUCCESS => { + if (n == 0) return error.PeerClosed; + off += @intCast(n); + }, + .INTR => {}, + else => return error.ReadFailed, + } + } +} \ No newline at end of file From 370752036cac14551e84c8a830868c07f3815531 Mon Sep 17 00:00:00 2001 From: styles01 <11393400+styles01@users.noreply.github.com> Date: Sun, 25 Oct 2026 00:12:39 -0600 Subject: [PATCH 4/7] fix(family): serve temperature-agnostic requests as greedy The keyed sampler is not attached for minimax_m3_vl, so a standard client's bare chat resolved against the checkpoint's generation-config defaults to a sampled spec and refused everywhere. An added per-engine temperature_greedy info flag (default false; resolved in the server's resolveSampling, mirroring the existing greedy_only design line) serves requests naming only temperature (any value, or absent) as fixed greedy and annotates the reply (tensorfold.sampling), while top_k/top_p/min_p/ seed and grammar masks keep refusing by their exact lane-surface name (SamplingNotAttached). The family's own lane surface is unchanged. Reasoning/thinking default stays; no family-level thinking knob exists in the glm adapter to mirror. --- zig/src/core/engine_api.zig | 3 +++ .../families/minimax_m3_vl/cuda_native.zig | 12 ++++++---- zig/src/families/minimax_m3_vl/shard.zig | 2 +- zig/src/native/cuda.zig | 1 + zig/src/server/chat.zig | 2 ++ zig/src/server/done_times_test.zig | 23 +++++++++++++++++++ zig/src/server/server.zig | 8 +++++++ 7 files changed, 46 insertions(+), 5 deletions(-) diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index b4ee68ce02..af576ebd6b 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -215,6 +215,9 @@ pub const Info = struct { vision: ?Vision = null, /// The engine decodes greedily only: the server refuses a request with a temperature before admitting it. greedy_only: bool = false, + /// A request naming only temperature (or no sampling field) serves fixed greedy; its other sampling + /// fields (top_k/top_p/min_p/seed) refuse by their exact lane-surface name (SamplingNotAttached). + temperature_greedy: bool = false, /// The learner reads conversation prompts at any length through taps. slide_conversations: bool = false, /// Whether a request's caches fit beside the running streams' now (null: a free lane does); else it waits in order. diff --git a/zig/src/families/minimax_m3_vl/cuda_native.zig b/zig/src/families/minimax_m3_vl/cuda_native.zig index c453207a46..a3b1bfdcc1 100644 --- a/zig/src/families/minimax_m3_vl/cuda_native.zig +++ b/zig/src/families/minimax_m3_vl/cuda_native.zig @@ -22,6 +22,9 @@ pub const graph = @import("graph.zig"); // the launch graph this stage attaches pub const model_type = "minimax_m3_vl"; /// The checkpoints this family reads: Mia-AiLab's MiniMax-M3.1-Flash-EXL3-EAQ-v1-3.06-bpw and its line. pub const formats: []const []const u8 = &.{"exl3"}; +/// The server serves this stage's requests as fixed greedy: temperature at any value (or absent) +/// serves, its other sampling fields refuse (SamplingNotAttached) — the keyed sampler is a later pass. +pub const temperature_greedy = true; /// The pack's own window (max_position_embeddings 1,048,576) when a config.json names none. pub const default_context: i64 = 1048576; pub const max_segments: u32 = 1; @@ -144,8 +147,8 @@ pub const serve_boundary = "minimax_m3_vl passed CUDA admission (every module's "recipe: rank 0 layers 0-29 + embed + final norm + lm_head, rank 1 layers 30-59, NCCL handoffs) grade " ++ "it — but this stage still refuses loudly rather than fake anything it does not have: draft rows and " ++ "drafted verify windows refuse (error.LanesNotAttached: no draft head — the draft/spec-decode surface " ++ - "is unimplemented), a sampled draw spec and grammar masks refuse (error.SamplingNotAttached: the keyed " ++ - "sampler is not wired for this family's lanes), a second stream refuses while one is live " ++ + "is unimplemented), a sampled draw spec that names more than temperature and grammar masks refuse (error.SamplingNotAttached: the keyed " ++ + "sampler is not wired for this family's lanes; temperature alone serves as fixed greedy), a second stream refuses while one is live " ++ "(error.LanesNotAttached: one stream at a time), tree keeps refuse (prefix rollbacks only), TP2 booted " ++ "through open() now ATTACHES since Stage-11 — the TF_MINIMAXM3_RANK/TF_MINIMAXM3_WORLD env server " ++ "boot serves when the place is the receipted world-2 split: rank 0 returns the Loaded the HTTP " ++ @@ -164,7 +167,7 @@ pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { return switch (err) { error.ServeNotAttached => serve_boundary, error.LanesNotAttached => "minimax_m3_vl's lane surface serves one greedy stream through the attached graph (prefill/first/queue/read/verify/keep/release — the boundary tests grade a short prompt token-in token-out); the calls that refuse here are the ones this stage does not have: draft rows and drafted verify windows (no draft head — the draft/spec-decode surface is unimplemented), a second stream while one is live, a verify batch wider than one window, and tree keeps", - error.SamplingNotAttached => "minimax_m3_vl's lane surface serves greedy draws only in this stage: a Sampling spec (temperature/top_k/top_p) or a grammar mask refuses here rather than serve an ungraded draw — the keyed sampler is a later pass", + error.SamplingNotAttached => "minimax_m3_vl's lane surface serves greedy draws only in this stage: the server takes temperature as fixed greedy (any value or absent); a Sampling spec that names more than temperature (top_k/top_p/min_p/seed) or a grammar mask refuses here rather than serve an ungraded draw — the keyed sampler is a later pass", error.TpServingNotAttached => "minimax_m3_vl's TF_MINIMAXM3_RANK/TF_MINIMAXM3_WORLD env server boot attaches as the receipted world-2 split (rank 0 serves the lane surface, rank 1 mirrors rank 0's rows; NCCL_COMM_ID names rank 0's id host:port). What still refuses: an env place outside that split (one env var alone, a non-numeric rank/world, a rank at or past its world, a world other than 2), while TP2 serving receipts live at `tf-cuda-test minimax-m3-token ` (the Stage-6/7 token receipts) and openTP2's lane surface — `tf-cuda-test minimax-m3-lane-tp2 ` (the Stage-9 receipts: the same lane vtable answering over the split pack, rank 1 first)", error.NoCommId => "minimax_m3_vl's world-2 server boot meets its rank over NCCL: set NCCL_COMM_ID to rank 0's address:port on every rank (the receipted dance serves rank 0's fresh UniqueId there; rank 1 connects first)", error.ContextTooWide => "the launch graph's decode pool this stage is 4096 keys: give a context at or under it (the wide paged prompt path is a later stage)", @@ -733,6 +736,7 @@ test "admission runs the whole checkpoint and open refuses exactly at the serve try std.testing.expect(std.mem.indexOf(u8, words, "prefill/first/") != null); // the served rows are named try std.testing.expect(std.mem.indexOf(u8, words, "LanesNotAttached") != null); // drafts+second streams still refuse try std.testing.expect(std.mem.indexOf(u8, words, "SamplingNotAttached") != null); // sampled draws refuse + try std.testing.expect(temperature_greedy); // the server serves temperature as fixed greedy try std.testing.expect(std.mem.indexOf(u8, words, "TpServingNotAttached") != null); // and the TP2 server boot try std.testing.expect(std.mem.indexOf(u8, words, "ContextTooWide") != null); // and the pool's window try std.testing.expect(explain(null, error.LanesNotAttached) != null); @@ -779,4 +783,4 @@ test { _ = @import("cuda_serve.zig"); _ = @import("shard.zig"); // the TP2 rank split's tests ride the same module _ = @import("graph.zig"); // the launch graph and its helpers compile with the suite -} \ No newline at end of file +} diff --git a/zig/src/families/minimax_m3_vl/shard.zig b/zig/src/families/minimax_m3_vl/shard.zig index 74e2312cc2..107d989f0a 100644 --- a/zig/src/families/minimax_m3_vl/shard.zig +++ b/zig/src/families/minimax_m3_vl/shard.zig @@ -113,4 +113,4 @@ test "the fixture load splits exactly: the two ranks' bytes and items sum to the // world 1 (TP1): everything on rank 0 try std.testing.expectEqual(l.bytes(), bytesOf(&l, 1, 0)); try std.testing.expectEqual(l.bytes(), bytesOf(&l, 1, 1)); -} \ No newline at end of file +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index b0c3184a83..b9f8261165 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -459,6 +459,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O .vision = if (@hasField(@TypeOf(loaded), "vision")) vision_offer.offer(loaded.vision, o.vision) else null, .admit = admit, .slide_conversations = conversations, + .temperature_greedy = comptime @hasDecl(F, "temperature_greedy") and F.temperature_greedy, }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; if (@hasField(@TypeOf(loaded), "follow")) if (loaded.follow) |f| { diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 359cc93056..6bcf0816e9 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -516,6 +516,8 @@ const Generation = struct { try runtime.put(a, "prefill_raised", .{ .array = raised }); try runtime.put(a, "time_to_first_token", if (g.first_ns) |t| .{ .float = seconds(t - received) } else .null); try runtime.put(a, "sampling", .{ .string = if (exact) "exact" else "greedy" }); + if (g.srv.info.temperature_greedy and !exact) + try runtime.put(a, "tensorfold.sampling", .{ .string = "greedy-deterministic (keyed sampler not attached; temperature is served as fixed greedy)" }); try runtime.put(a, "drafts", .{ .bool = drafts }); const sha = @import("tokens.zig").tokenSha(g.collected.items); try runtime.put(a, "token_sha", .{ .string = try a.dupe(u8, &sha) }); diff --git a/zig/src/server/done_times_test.zig b/zig/src/server/done_times_test.zig index 20649b8662..d4798255ca 100644 --- a/zig/src/server/done_times_test.zig +++ b/zig/src/server/done_times_test.zig @@ -173,6 +173,29 @@ fn seedSample(srv: *server.Server, a: std.mem.Allocator, body: []const u8) !?api return srv.resolveSampling(&cx, (try json.parse(a, body)).ok, 0, &.{ 1, 2, 3 }); } +test "done times: a temperature-greedy engine serves temperature-only requests greedy; other sampling fields refuse by name" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + Clock.next = 42; + var script: Script = .{ .first = 10_384_000_000 }; + var tt: helpers = .{}; + const srv = try server.Server.init(std.testing.allocator, Clock.io(), script.engine(), text(&tt), .{ .served_name = "test", .model_ids = &.{"test"}, .enable_thinking = false }, null); + defer srv.deinit(); + srv.info.temperature_greedy = true; + try std.testing.expectEqual(@as(?api.Sampling, null), try seedSample(srv, a, "{}")); // absent temperature: fixed greedy + try std.testing.expectEqual(@as(?api.Sampling, null), try seedSample(srv, a, "{\"temperature\":0}")); + try std.testing.expectEqual(@as(?api.Sampling, null), try seedSample(srv, a, "{\"temperature\":0.7}")); + srv.config.default_sampling = (try json.parse(a, "{\"temperature\":0.7,\"top_k\":40,\"top_p\":0.95}")).ok; + try std.testing.expectEqual(@as(?api.Sampling, null), try seedSample(srv, a, "{}")); // defaults do not sample this engine + var cx: Cx = .{ .a = a }; + try std.testing.expectError(error.Refused, srv.resolveSampling(&cx, (try json.parse(a, "{\"temperature\":0.7,\"top_k\":5}")).ok, 0, &.{1})); + try std.testing.expect(std.mem.indexOf(u8, cx.message, "SamplingNotAttached") != null); + cx = .{ .a = a }; + try std.testing.expectError(error.Refused, srv.resolveSampling(&cx, (try json.parse(a, "{\"seed\":7}")).ok, 0, &.{1})); + try std.testing.expect(std.mem.indexOf(u8, cx.message, "SamplingNotAttached") != null); +} + fn seedControls(mode: enum { fresh, keyed, explicit, defaults, default_fresh, failed, log_drawn, log_given }) !void { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); diff --git a/zig/src/server/server.zig b/zig/src/server/server.zig index 2795ee7236..6c2125f8c2 100644 --- a/zig/src/server/server.zig +++ b/zig/src/server/server.zig @@ -207,6 +207,14 @@ pub const Server = struct { if (defaults) |d| try merge(options, cx, d); const supplied = try fields.parseNumbers(cx, f); try merge(options, cx, supplied); + if (srv.info.temperature_greedy) { + inline for (.{ "top_k", "top_p", "min_p", "seed" }) |k| { + if (f.get(k)) |v| if (v != .null) return cx.refuse( + "this engine serves temperature as fixed greedy (no keyed sampler is attached): its other sampling fields refuse (SamplingNotAttached) — send temperature only, or leave the sampling fields out", + ); + } + return null; // temperature at any value, or absent: fixed greedy, annotated in the reply + } const temp = number(options.get("temperature")) orelse 0; if (temp <= 0) return null; if (srv.info.greedy_only) return cx.refuse(if ((if (defaults) |d| number(d.get("temperature")) else null) orelse 0 > 0) From f623f9eab8acc756431c46a74b7ca2a7cf3a22c1 Mon Sep 17 00:00:00 2001 From: styles01 <11393400+styles01@users.noreply.github.com> Date: Sun, 25 Oct 2026 01:06:05 -0600 Subject: [PATCH 5/7] fix(family): accept temperature-only sampling as greedy at the HTTP layer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The earlier temperature-greedy wiring resolved the flag comptime through openWith(F) with F = the family's registry alias (minimax_m3_vl.native), and that alias re-exported model_type/formats/Options/open/explain/ weightShare but not temperature_greedy — @hasDecl fell to false, the engine Info served false, and POST /v1/chat/completions resolved bare and temperature-only bodies against the checkpoint's generation-config defaults (temperature 1.0, top_p 0.95), built a full Sampling spec and hit the lane's SamplingNotAttached. Only an explicit temperature 0 served (the old value-based path). The alias now carries the flag, so the server's resolveSampling key-presence flip actually runs on the live serve path: temperature at any value, or absent, serves fixed greedy and annotates the reply; top_k/top_p/min_p/seed and grammar masks keep refusing by their exact lane-surface name. Tests extend done_times with the real HTTP parser path (openai.chatInput -> chat.prepare) and guard the alias re-export in the family's own admission test. --- .../families/minimax_m3_vl/cuda_native.zig | 2 ++ zig/src/server/done_times_test.zig | 34 +++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/zig/src/families/minimax_m3_vl/cuda_native.zig b/zig/src/families/minimax_m3_vl/cuda_native.zig index a3b1bfdcc1..282666790c 100644 --- a/zig/src/families/minimax_m3_vl/cuda_native.zig +++ b/zig/src/families/minimax_m3_vl/cuda_native.zig @@ -72,6 +72,7 @@ pub const native = struct { pub const open = me.open; pub const explain = me.explain; pub const weightShare = me.weightShare; // the server's admission reads this rank's split share + pub const temperature_greedy = me.temperature_greedy; // openWith(F)'s Info reads this through the registry alias }; /// Slots for `context` plus one whole block's reserve, a multiple of the 128-token key block. @@ -737,6 +738,7 @@ test "admission runs the whole checkpoint and open refuses exactly at the serve try std.testing.expect(std.mem.indexOf(u8, words, "LanesNotAttached") != null); // drafts+second streams still refuse try std.testing.expect(std.mem.indexOf(u8, words, "SamplingNotAttached") != null); // sampled draws refuse try std.testing.expect(temperature_greedy); // the server serves temperature as fixed greedy + try std.testing.expect(@hasDecl(native, "temperature_greedy") and native.temperature_greedy); // the registry alias carries it: openWith(F)'s Info reads the flag through F = native try std.testing.expect(std.mem.indexOf(u8, words, "TpServingNotAttached") != null); // and the TP2 server boot try std.testing.expect(std.mem.indexOf(u8, words, "ContextTooWide") != null); // and the pool's window try std.testing.expect(explain(null, error.LanesNotAttached) != null); diff --git a/zig/src/server/done_times_test.zig b/zig/src/server/done_times_test.zig index d4798255ca..b261bd28be 100644 --- a/zig/src/server/done_times_test.zig +++ b/zig/src/server/done_times_test.zig @@ -196,6 +196,40 @@ test "done times: a temperature-greedy engine serves temperature-only requests g try std.testing.expect(std.mem.indexOf(u8, cx.message, "SamplingNotAttached") != null); } +// POST /v1/chat/completions end to end (HTTP parser → chat.prepare): the path a live request takes +// on the real serve pair, where the earlier temperature-greedy fix was wired but its registry alias +// never carried the flag, so the parser resolved against the generation-config defaults and refused. +test "done times: the HTTP chat parser serves temperature-only and absent-temperature bodies greedy, and temperature+top_k refuses by name" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + Clock.next = 42; + var script: Script = .{ .first = 10_384_000_000 }; + var tt: helpers = .{}; + const srv = try server.Server.init(std.testing.allocator, Clock.io(), script.engine(), text(&tt), .{ .served_name = "test", .model_ids = &.{"test"}, .enable_thinking = false }, null); + defer srv.deinit(); + srv.info.temperature_greedy = true; + srv.config.default_sampling = (try json.parse(a, "{\"temperature\":1.0,\"top_p\":0.95,\"top_k\":40}")).ok; // the live pack's generation_config defaults + const bodies = [_][]const u8{ + "{\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}]}", // bare: any served body's shape + "{\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"temperature\":0.7}", // temperature-only, sampled value + "{\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"temperature\":0}", + "{\"max_tokens\":32,\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"temperature\":0.7}", // a stream:false body with max_tokens + }; + for (bodies) |raw| { + var cx: Cx = .{ .a = a }; + const input = try @import("openai.zig").chatInput(srv, &cx, (try json.parse(a, raw)).ok); + const prepared = try chat.prepare(srv, &cx, input, Gone{}); + chat.release(srv, prepared.preparing); + try std.testing.expectEqual(@as(?api.Sampling, null), prepared.request.sampling); // fixed greedy, no keyed sampler + try std.testing.expect(!prepared.sampled); + } + var cx: Cx = .{ .a = a }; + const refused = try @import("openai.zig").chatInput(srv, &cx, (try json.parse(a, "{\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"temperature\":0.7,\"top_k\":40}")).ok); + try std.testing.expectError(error.Refused, chat.prepare(srv, &cx, refused, Gone{})); + try std.testing.expect(std.mem.indexOf(u8, cx.message, "SamplingNotAttached") != null); // refuses by the lane surface's exact name +} + fn seedControls(mode: enum { fresh, keyed, explicit, defaults, default_fresh, failed, log_drawn, log_given }) !void { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); From fbf132f080a74f702999d8caf086b79ccfafb1c3 Mon Sep 17 00:00:00 2001 From: styles01 <11393400+styles01@users.noreply.github.com> Date: Sun, 25 Oct 2026 02:11:07 -0600 Subject: [PATCH 6/7] fix(family): honor the chat template's thinking-off knob + refuse thinking-on by name MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stage-11 live-lane probes (spark2 :8000, MiniMax-3.1-Flash, PR #1281 head f623f9eab) split the empty chat reply into two findings: 1. The thinking switch was inert on this family's rendered prompt. hf_text put the generic thinking_mode spellings ("thinking"/"chat") into every template, but MiniMax M3.1's chat_template.jinja branches on "enabled"/"disabled"/"adaptive"; both spellings fell through to its adaptive branch, so chat_template_kwargs enable_thinking was accepted (served, echoed) yet changed nothing in the prompt ids: thinking on and off rendered the identical 156 rows. thinkingModeValue now reads the template's own dialect from its source, so thinking-off renders disabled (the assistant prefix) and templates without the dialect keep the generic spellings. 2. The request's own enable_thinking body field was ignored: the server read the switch only through chat_template_kwargs. thinkingFields now honors the body's field when the kwargs switch is absent (kwargs wins), so the field is honored wherever it names the switch. 3. The stage knob: minimax_m3_vl declares thinking_off (mirroring temperature_greedy's registry-alias pattern: family const + openWith(F)'s Info read + a family guard test). chat.prepare runs thinking-off engines without the reasoning header whatever the server config says, and an explicit thinking request (chat_template_kwargs or the body's field) refuses by exact name (ThinkingNotAttached) — the native lane has no thinking reserve attached, so thinking cannot be served honestly. Host suites on the canonical tree (zig 0.17.0): test-server-cpu 618/618, test-done-times-host 20/20 (new HTTP parser-level test: bare bodies prepare thinking=false; kwargs and body-level enable_thinking true refuse naming ThinkingNotAttached; body-level false prepares thinking=false on a thinking-capable engine), test-cuda-host 177/177 steps 451/452 tests passed (1 pre-existing recorded skip); lean_check vs merge-base 7938815f1 introduces no new problems; git diff --check clean. Lane-level boundary (measured on the live pair, receipts §4b): the greedy first draw at the assistant turn is the EOS token itself under every honest render, so chat content staying empty is a decode-fidelity defect of the native EXL3 row path on the real pack, upstream of the template — receipts report it verbatim; this commit carries the honest knob. --- zig/src/core/engine_api.zig | 2 + .../families/minimax_m3_vl/cuda_native.zig | 5 +++ zig/src/native/cuda.zig | 1 + zig/src/server/chat.zig | 10 ++++- zig/src/server/done_times_test.zig | 44 +++++++++++++++++++ zig/src/server/fields.zig | 4 ++ zig/src/server/hf_text.zig | 23 +++++++++- 7 files changed, 87 insertions(+), 2 deletions(-) diff --git a/zig/src/core/engine_api.zig b/zig/src/core/engine_api.zig index af576ebd6b..104b0416a7 100644 --- a/zig/src/core/engine_api.zig +++ b/zig/src/core/engine_api.zig @@ -218,6 +218,8 @@ pub const Info = struct { /// A request naming only temperature (or no sampling field) serves fixed greedy; its other sampling /// fields (top_k/top_p/min_p/seed) refuse by their exact lane-surface name (SamplingNotAttached). temperature_greedy: bool = false, + /// The engine serves chat with thinking off: an explicit thinking request refuses (ThinkingNotAttached). + thinking_off: bool = false, /// The learner reads conversation prompts at any length through taps. slide_conversations: bool = false, /// Whether a request's caches fit beside the running streams' now (null: a free lane does); else it waits in order. diff --git a/zig/src/families/minimax_m3_vl/cuda_native.zig b/zig/src/families/minimax_m3_vl/cuda_native.zig index 282666790c..3384c43b3e 100644 --- a/zig/src/families/minimax_m3_vl/cuda_native.zig +++ b/zig/src/families/minimax_m3_vl/cuda_native.zig @@ -25,6 +25,8 @@ pub const formats: []const []const u8 = &.{"exl3"}; /// The server serves this stage's requests as fixed greedy: temperature at any value (or absent) /// serves, its other sampling fields refuse (SamplingNotAttached) — the keyed sampler is a later pass. pub const temperature_greedy = true; +/// This stage serves chat with thinking off: an explicit thinking request refuses by exact name (ThinkingNotAttached). +pub const thinking_off = true; /// The pack's own window (max_position_embeddings 1,048,576) when a config.json names none. pub const default_context: i64 = 1048576; pub const max_segments: u32 = 1; @@ -73,6 +75,7 @@ pub const native = struct { pub const explain = me.explain; pub const weightShare = me.weightShare; // the server's admission reads this rank's split share pub const temperature_greedy = me.temperature_greedy; // openWith(F)'s Info reads this through the registry alias + pub const thinking_off = me.thinking_off; // as temperature_greedy: openWith(F)'s Info reads this through the registry alias }; /// Slots for `context` plus one whole block's reserve, a multiple of the 128-token key block. @@ -739,6 +742,8 @@ test "admission runs the whole checkpoint and open refuses exactly at the serve try std.testing.expect(std.mem.indexOf(u8, words, "SamplingNotAttached") != null); // sampled draws refuse try std.testing.expect(temperature_greedy); // the server serves temperature as fixed greedy try std.testing.expect(@hasDecl(native, "temperature_greedy") and native.temperature_greedy); // the registry alias carries it: openWith(F)'s Info reads the flag through F = native + try std.testing.expect(thinking_off); // this stage serves the chat lane with thinking off + try std.testing.expect(@hasDecl(native, "thinking_off") and native.thinking_off); // the registry alias carries it: openWith(F)'s Info reads the flag through F = native try std.testing.expect(std.mem.indexOf(u8, words, "TpServingNotAttached") != null); // and the TP2 server boot try std.testing.expect(std.mem.indexOf(u8, words, "ContextTooWide") != null); // and the pool's window try std.testing.expect(explain(null, error.LanesNotAttached) != null); diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index b9f8261165..cdd60d3c9a 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -460,6 +460,7 @@ fn openWith(comptime F: type, a: Allocator, gpa: Allocator, io: std.Io, o: api.O .admit = admit, .slide_conversations = conversations, .temperature_greedy = comptime @hasDecl(F, "temperature_greedy") and F.temperature_greedy, + .thinking_off = comptime @hasDecl(F, "thinking_off") and F.thinking_off, }, .{ .ctx = loaded.ctx, .text = F.explain }); opened = true; if (@hasField(@TypeOf(loaded), "follow")) if (loaded.follow) |f| { diff --git a/zig/src/server/chat.zig b/zig/src/server/chat.zig index 6bcf0816e9..7709ee47bf 100644 --- a/zig/src/server/chat.zig +++ b/zig/src/server/chat.zig @@ -132,8 +132,16 @@ pub fn prepare(srv: *Server, cx: *Cx, input: Input, gone: anytype) Failure!Prepa if (preparing) _ = srv.preparing.fetchAdd(1, .acq_rel); errdefer release(srv, preparing); if (gone.check()) return error.Cancelled; - var thinking = flag(f, "enable_thinking") orelse srv.config.enable_thinking; + const thinking_switch = flag(f, "enable_thinking"); + var thinking = thinking_switch orelse srv.config.enable_thinking; if (input.prompt != null) thinking = false; + if (srv.info.thinking_off) { + // the stage's thinking-off default: every body runs without the reasoning header; explicit on refuses above + if (thinking_switch == true) return cx.refuse( + "this engine serves chat with thinking off (no thinking reserve is attached): explicit thinking requests refuse (ThinkingNotAttached) - leave enable_thinking out, or send chat_template_kwargs {\"enable_thinking\": false}", + ); + thinking = false; + } const effort = srv.effortFor(if (f.get("reasoning_effort")) |e| (if (e == .string) e.string else null) else null); const rendered = try prompt_mod.prepare(srv, cx, input, thinking, effort); if (gone.check()) return error.Cancelled; diff --git a/zig/src/server/done_times_test.zig b/zig/src/server/done_times_test.zig index b261bd28be..817215c1df 100644 --- a/zig/src/server/done_times_test.zig +++ b/zig/src/server/done_times_test.zig @@ -230,6 +230,50 @@ test "done times: the HTTP chat parser serves temperature-only and absent-temper try std.testing.expect(std.mem.indexOf(u8, cx.message, "SamplingNotAttached") != null); // refuses by the lane surface's exact name } +// The thinking-off stage knob (temperature-greedy's mirror): bare bodies run header-free; explicit thinking refuses. +test "done times: the HTTP chat parser runs a thinking-off engine without the reasoning header and refuses explicit thinking requests by name" { + var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); + defer arena.deinit(); + const a = arena.allocator(); + Clock.next = 42; + var script: Script = .{ .first = 10_384_000_000 }; + var tt: helpers = .{}; + const srv = try server.Server.init(std.testing.allocator, Clock.io(), script.engine(), text(&tt), .{ .served_name = "test", .model_ids = &.{"test"} }, null); + defer srv.deinit(); + srv.info.thinking_off = true; + var bcx: Cx = .{ .a = a }; + const bare = try @import("openai.zig").chatInput(srv, &bcx, (try json.parse(a, "{\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}]}")).ok); + var pcx: Cx = .{ .a = a }; + const prepared = try chat.prepare(srv, &pcx, bare, Gone{}); + chat.release(srv, prepared.preparing); + try std.testing.expect(!prepared.thinking); // the stage default: thinking off without any switch + for ([_][]const u8{ + "{\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"chat_template_kwargs\":{\"enable_thinking\":true}}", + "{\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"enable_thinking\":true}", // the body's own field refuses too + }) |raw| { + var rcx: Cx = .{ .a = a }; + const thinking_body = try @import("openai.zig").chatInput(srv, &rcx, (try json.parse(a, raw)).ok); + var ecx: Cx = .{ .a = a }; + try std.testing.expectError(error.Refused, chat.prepare(srv, &ecx, thinking_body, Gone{})); + try std.testing.expect(std.mem.indexOf(u8, ecx.message, "ThinkingNotAttached") != null); // by exact name + } + // false anywhere serves: the stage's own default, said back in the echo + var off: Cx = .{ .a = a }; + const off_body = try @import("openai.zig").chatInput(srv, &off, (try json.parse(a, "{\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"enable_thinking\":false}")).ok); + var opcx: Cx = .{ .a = a }; + const off_prepared = try chat.prepare(srv, &opcx, off_body, Gone{}); + chat.release(srv, off_prepared.preparing); + try std.testing.expect(!off_prepared.thinking); + // a thinking-capable engine (no thinking_off) honors the body's own enable_thinking false. + srv.info.thinking_off = false; + var gcx: Cx = .{ .a = a }; + const glm_like = try @import("openai.zig").chatInput(srv, &gcx, (try json.parse(a, "{\"messages\":[{\"role\":\"user\",\"content\":\"hi\"}],\"enable_thinking\":false}")).ok); + var gpcx: Cx = .{ .a = a }; + const glm_prepared = try chat.prepare(srv, &gpcx, glm_like, Gone{}); + chat.release(srv, glm_prepared.preparing); + try std.testing.expect(!glm_prepared.thinking); // the body's own switch: false honored, not ignored +} + fn seedControls(mode: enum { fresh, keyed, explicit, defaults, default_fresh, failed, log_drawn, log_given }) !void { var arena: std.heap.ArenaAllocator = .init(std.testing.allocator); defer arena.deinit(); diff --git a/zig/src/server/fields.zig b/zig/src/server/fields.zig index 0ccdc64f18..cddc2b6e20 100644 --- a/zig/src/server/fields.zig +++ b/zig/src/server/fields.zig @@ -219,6 +219,10 @@ pub fn thinkingFields(cx: *Cx, body: Value, levels: []const []const u8) errors.R out.enable = on.truthy(); if (out.enable.? and out.effort != null and std.mem.eql(u8, out.effort.?, "none")) out.effort = null; }; + // the body's own enable_thinking where kwargs names no switch (kwargs wins): the field is honored, not ignored + if (enable == null) if (body.field("enable_thinking")) |own| if (own != .null) { + out.enable = own.truthy(); + }; return out; } diff --git a/zig/src/server/hf_text.zig b/zig/src/server/hf_text.zig index 9c558ceef7..1e383ea124 100644 --- a/zig/src/server/hf_text.zig +++ b/zig/src/server/hf_text.zig @@ -312,7 +312,7 @@ pub const HfText = struct { var it = t.specials.object.iterator(); while (it.next()) |e| try globals.put(a, e.key_ptr.*, e.value_ptr.*); } - try globals.put(a, "thinking_mode", .{ .string = if (options.enable_thinking) "thinking" else "chat" }); + try globals.put(a, "thinking_mode", .{ .string = thinkingModeValue(t.source, options.enable_thinking) }); var context: std.json.ObjectMap = .empty; try context.put(a, "enable_thinking", .{ .bool = options.enable_thinking }); if (options.reasoning_effort) |e| try context.put(a, "reasoning_effort", .{ .string = e }); @@ -327,6 +327,27 @@ pub const HfText = struct { } }; +/// The thinking_mode value the template's own dialect reads; see the stage receipts for the inert-switch bug. +pub fn thinkingModeValue(source: []const u8, enable: bool) []const u8 { + const dialect = std.mem.indexOf(u8, source, "\"adaptive\"") != null and + std.mem.indexOf(u8, source, "\"disabled\"") != null and + std.mem.indexOf(u8, source, "\"enabled\"") != null; + return if (dialect) + (if (enable) "enabled" else "disabled") + else + (if (enable) "thinking" else "chat"); +} + +test "thinking mode value speaks the template's own dialect, keeping the generic spellings elsewhere" { + const source = "{{ 'enabled' if thinking_mode == \"enabled\" elif thinking_mode == \"disabled\" else \"adaptive\" }}"; + try std.testing.expectEqualStrings("enabled", thinkingModeValue(source, true)); + try std.testing.expectEqualStrings("disabled", thinkingModeValue(source, false)); + try std.testing.expectEqualStrings("thinking", thinkingModeValue("{{ enable_thinking }}", true)); + try std.testing.expectEqualStrings("chat", thinkingModeValue("{{ enable_thinking }}", false)); + try std.testing.expectEqualStrings("thinking", thinkingModeValue("", true)); // no dialect markers: unchanged + try std.testing.expectEqualStrings("chat", thinkingModeValue("disabled only", false)); // all three words must name the dialect +} + test "model EOS does not hide a different tokenizer chat end" { var arena = std.heap.ArenaAllocator.init(std.testing.allocator); defer arena.deinit(); From a493d132c11b45b4e18d5ecb0fc158a31584b839 Mon Sep 17 00:00:00 2001 From: styles01 <11393400+styles01@users.noreply.github.com> Date: Sun, 25 Oct 2026 17:39:16 -0600 Subject: [PATCH 7/7] fix(family): run the routed experts' epilogue at the OAI swiglu contract (act_mode 2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The MiniMax M3.1 routed experts' gate/up epilogue ran the generic act_mode-1 silu path; this family's contract is the OAI swiglu (alpha 1.702 sigmoid form; fp16 gate/up rows and an fp16 activation row before the down-side rotation — tf_m3_swigluoai's shape). The decode epilogue (tf_exl3_gateup_epilogue), the prompt-side tgateup and that path's eligible() gate now carry act_mode 2, and the family's experts upload names it. Layer-3 boundary rows against the f64 mirror: relmax 9.7e-2 -> 4.3e-4 (~200x); layers 0-2 bit-identical pre/post (footprint exact). Full-depth serving on the rebuilt binary now completes 'The capital of France is' (temperature 0, greedy) with real content where the pre-fix lane drew a forced first-draw EOS with empty content. --- zig/kernels/cuda/exl3_experts.cu | 10 ++++++++++ zig/kernels/cuda/exl3_prompt.cuh | 18 ++++++++++++++---- zig/src/cuda/exl3_prompt.zig | 2 +- zig/src/families/minimax_m3_vl/graph.zig | 7 +++++-- 4 files changed, 30 insertions(+), 7 deletions(-) diff --git a/zig/kernels/cuda/exl3_experts.cu b/zig/kernels/cuda/exl3_experts.cu index b31d90b40a..806794c2e5 100644 --- a/zig/kernels/cuda/exl3_experts.cu +++ b/zig/kernels/cuda/exl3_experts.cu @@ -309,6 +309,16 @@ extern "C" __global__ void tf_exl3_gateup_epilogue(const GateUpArgs a) { float uu = bf16r(uv[j] * HAD_SCALE * __half2float(a.svh_u[(size_t)e * a.n + n + j])); if (a.limit > 0) { gg = upTo(gg, a.limit); uu = within(uu, a.limit); } act = bf16r(bf16r(gg / (1.f + expf(-gg))) * uu); + } else if (a.act_mode == 2) { + // the minimax M3.1 OAI swiglu contract at the fork's quantization points — the dense + // path's shape: the gate/up rows are fp16 (the fork's f16 GEMM outputs), the sigmoid + // form is the tf_m3_swigluoai contract (gate clipped from above only, up clamped both + // ways and +1, sigmoid alpha = 1.702), and the activation is an fp16 row BEFORE this + // epilogue's down-side rotation (the fork's down kernel rotates an fp16 z in fp32) + float gg = __half2float(__float2half_rn(gv[j] * HAD_SCALE * __half2float(a.svh_g[(size_t)e * a.n + n + j]))); + float uu = __half2float(__float2half_rn(uv[j] * HAD_SCALE * __half2float(a.svh_u[(size_t)e * a.n + n + j]))); + if (a.limit > 0) { gg = upTo(gg, a.limit); uu = within(uu, a.limit); } + act = __half2float(__float2half_rn((uu + 1.0f) * (gg / (1.0f + __expf(-1.702f * gg))))); } else { float gg = gv[j] * HAD_SCALE * __half2float(a.svh_g[(size_t)e * a.n + n + j]); float uu = uv[j] * HAD_SCALE * __half2float(a.svh_u[(size_t)e * a.n + n + j]); diff --git a/zig/kernels/cuda/exl3_prompt.cuh b/zig/kernels/cuda/exl3_prompt.cuh index 067c3aaaaa..b865638518 100644 --- a/zig/kernels/cuda/exl3_prompt.cuh +++ b/zig/kernels/cuda/exl3_prompt.cuh @@ -179,10 +179,20 @@ __device__ __forceinline__ void tgateup(const half* __restrict__ Xg, const half* float uu = uv[j] * HAD_SCALE * __half2float(svh_u[(size_t)e * N + n + j]); if (act_mode == 0) { gg = bf16r(gg); uu = bf16r(uu); } if (limit > 0) { gg = upTo(gg, limit); uu = within(uu, limit); } - float act = gg / (1.f + expf(-gg)); - if (act_mode == 0) act = bf16r(act); - act *= uu; - if (act_mode == 0) act = bf16r(act); + float act; + if (act_mode == 2) { + // the minimax M3.1 OAI swiglu contract at the fork's quantization points — the + // same shape as the decode epilogue's mode 2: fp16 gate/up rows, then a fp16 + // activation row before the down-side rotation + gg = __half2float(__float2half_rn(gg)); + uu = __half2float(__float2half_rn(uu)); + act = __half2float(__float2half_rn((uu + 1.0f) * (gg / (1.0f + __expf(-1.702f * gg))))); + } else { + act = gg / (1.f + expf(-gg)); + if (act_mode == 0) act = bf16r(act); + act *= uu; + if (act_mode == 0) act = bf16r(act); + } v[j] = act * __half2float(suh_d[(size_t)e * N + n + j]); } fwht128(v, lane); diff --git a/zig/src/cuda/exl3_prompt.zig b/zig/src/cuda/exl3_prompt.zig index 635540ecc6..c9ee671287 100644 --- a/zig/src/cuda/exl3_prompt.zig +++ b/zig/src/cuda/exl3_prompt.zig @@ -10,7 +10,7 @@ const Stream = @import("stream.zig").Stream; pub const tile = 64; pub const Scratch = struct { xg: u64, xu: u64, xd: u64, y: u64, rows: usize, slots: usize }; pub fn eligible(l: ex.Layer) bool { - return l.codebook == .mcg and l.uniform_k2 == 8 and l.gu_stride == 0 and l.dims > 0 and l.width > 0 and l.dims % 256 == 0 and l.width % 128 == 0 and (l.act_mode == 0 or l.act_mode == 1); + return l.codebook == .mcg and l.uniform_k2 == 8 and l.gu_stride == 0 and l.dims > 0 and l.width > 0 and l.dims % 256 == 0 and l.width % 128 == 0 and (l.act_mode == 0 or l.act_mode == 1 or l.act_mode == 2); } fn mul(a: usize, b: usize) !usize { return std.math.mul(usize, a, b) catch error.Exl3Shape; diff --git a/zig/src/families/minimax_m3_vl/graph.zig b/zig/src/families/minimax_m3_vl/graph.zig index 68c52dfe72..7a5101074c 100644 --- a/zig/src/families/minimax_m3_vl/graph.zig +++ b/zig/src/families/minimax_m3_vl/graph.zig @@ -2,7 +2,8 @@ //! running the TP2 rank split shard.zig draws. The quantized projections run through the mainline //! EXL3 machinery the family already receipts — `cuda.exl3.Dense` decode-row GEMMs (mul1's odd //! half steps and the head's 8 halves included) and the grouped routed experts -//! `cuda.exl3_experts.Experts` (the Stage-5 upload's fp32 epilogue, limit 7, act_mode 1); the MSA +//! `cuda.exl3_experts.Experts` (the Stage-5 upload's fp32 epilogue, limit 7, act_mode 2 — the +//! OAI swiglu contract tf_m3_swigluoai's shape; the mode-1 silu path cannot express it); the MSA //! indexer and the block-sparse attention run through the parity-proven `tf_m3_scores/topk/ //! decode/decode_combine` launches (`cuda_sparse.zig`); the activation rows run through //! `tf_m3_swigluoai` (`cuda_serve.zig`). The between-the-kernels glue — the gemma norm rows, the @@ -379,7 +380,9 @@ pub const Engine = struct { if (ep.*) |*lp| { if (width == 0) return error.BadExperts; if (dno < r.first or dno >= r.first + run_count) continue; - const routed = try exl3_upload.layer(&up, a, dims, width, .mul1, c.swiglu_limit, 1, lp.gate.items, lp.up.items, lp.down.items); + // act_mode 2: the OAI swiglu contract (tf_m3_swigluoai's) — the generic epilogue's + // mode-1 silu path cannot express the alpha 1.702 sigmoid form + const routed = try exl3_upload.layer(&up, a, dims, width, .mul1, c.swiglu_limit, 2, lp.gate.items, lp.up.items, lp.down.items); const lidx: usize = dno - r.first; const sg = s_gate[lidx] orelse return error.MissingTensor; const su = s_up[lidx] orelse return error.MissingTensor;