onwx 0.9.0: Qwen3.6 35B-A3B experts from the SSD (~7 GB of memory), images, faster MTP; changelog shown with updates
5123724 verified 19 days ago
43.2 kB
// glue_cpu: the onwx all-NPU prefill's glue ops (xd_gen pfx_glue.cc / pfx_glue.py) on the CPU, on the same buffers
// (the mapped BOs) in the same DDR layouts, with the NPU kernels' arithmetic reproduced step by step (bf16 hi / lo
// splits, 3-product fp32 multiplies, the division-free sigmoid, the FWHT's butterfly order, the bfp16ebs8 rounding),
// so the next NPU GEMM reads the same X tiles.  AVX-512 (Zen 5); a persistent thread pool (pinned: the least
// performant core class when there are 2+, else every physical core; Windows: EcoQoS).
//
//   cg_info(int out[8])     [0] AVX-512 usable [1] efficiency classes (mask) [2] chosen cores [3] physical cores
//   cg_init(nthreads, mode) the pool (0: one a chosen core); mode bits: 1 pin, 2 EcoQoS (Windows). Returns threads.
//   cg_run(op, M, H, I, ND, KO, NO, ZB, a0, a1, a2, a3, dbg)
//       op: 0 n1 (h P X) 1 r1n (h Y P X) 2 r3n (h Y P X) 3 r3 (h Y P) 4 act (Y P X) 5 dno (O Z P X) 6 rotp (O P X)
//       7 rotb (O P X) - the buffers in pfx_glue.py's ARGS order; P the layer's parameter buffer (make_P: headers,
//       sign bits, CT, the DNO weight).  dbg (optional): the rotated fp32 activations [M][K] before the bfp16
//       conversion.  Returns 0 (or -1: no AVX-512 / bad op, -2: bad rotation block).  P's op header [12]: the
//       rotation block (0 / 1024: H_1024 a chunk; 512: two H_512 a chunk; 4096: + the 4-chunk butterfly / 2).
//   cg_set(key, val)        test knobs (0: the bfp rounding variant, 1: rsqrt variant, 2: reduce order)
#include <immintrin.h>
#include <atomic>
#include <algorithm>
#include <cmath>
#include <condition_variable>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <mutex>
#include <thread>
#include <vector>
#ifdef _WIN32
#include <intrin.h>
#include <windows.h>
#define API extern "C" __declspec(dllexport)
#else
#include <cpuid.h>
#include <dirent.h>
#include <pthread.h>
#include <sched.h>
#include <unistd.h>
#define API extern "C" __attribute__((visibility("default")))
#endif

typedef __m512 V;
typedef __m512i VI;

static int g_knob[8];

// (only the kernels below are AVX-512 code: the rest - the CPU check, the core list, the pool - runs on any x86-64,
// so a CPU without AVX-512 gets cg_info()[0] = 0 / cg_init() = -1 instead of an illegal instruction)
#if defined(__GNUC__) && !defined(_MSC_VER)
#pragma GCC push_options
#pragma GCC target("avx512f,avx512bw,avx512dq,avx512vl,fma")
#endif
// ---------------------------------------------------------------- the NPU's arithmetic (pfx_sig.h)
static inline V bc(float x) { return _mm512_set1_ps(x); }
static inline VI ibc(int x) { return _mm512_set1_epi32(x); }
// fp32 -> bf16 (round to nearest even), kept as fp32
static inline V bfr(V v) {
  VI u = _mm512_castps_si512(v);
  VI lsb = _mm512_and_si512(_mm512_srli_epi32(u, 16), ibc(1));
  u = _mm512_add_epi32(u, _mm512_add_epi32(ibc(0x7fff), lsb));
  return _mm512_castsi512_ps(_mm512_and_si512(u, ibc((int)0xffff0000u)));
}
static float bfs(float x) {   // (scalar: a bf16 constant, bbc())
  uint32_t u;
  memcpy(&u, &x, 4);
  u = (u + 0x7fff + ((u >> 16) & 1)) & 0xffff0000u;
  memcpy(&x, &u, 4);
  return x;
}
struct T2 {
  V h, l;
};
static inline T2 sp2(V v) {
  T2 r;
  r.h = bfr(v);
  r.l = bfr(_mm512_sub_ps(v, r.h));
  return r;
}
static inline V fma_(V a, V b, V c) { return _mm512_fmadd_ps(a, b, c); }   // (a * b exact for bf16 a, b)
static inline V mulf(V a, V b) {
  T2 x = sp2(a), y = sp2(b);
  V r = _mm512_mul_ps(x.h, y.h);
  r = fma_(x.h, y.l, r);
  return fma_(x.l, y.h, r);
}
static inline V addf(V a, V b, float sb) {
  T2 y = sp2(b);
  V s = bc(sb);
  V r = fma_(y.h, s, a);
  return fma_(y.l, s, r);
}
static inline V rsqrt16(V x) {
  V r;
  if (g_knob[1] == 1) {        // (variant: the AVX-512 14-bit estimate)
    r = _mm512_rsqrt14_ps(x);
  } else {
    r = _mm512_div_ps(bc(1.0f), _mm512_sqrt_ps(x));
  }
  V hx = mulf(x, bc(0.5f));
  for (int k = 0; k < 2; k++)
    r = mulf(r, addf(bc(1.5f), mulf(hx, mulf(r, r)), -1.0f));
  return r;
}
static float C_S1, C_S2, C_P4, C_R1, C_R2;
static void init_consts() {
  C_S1 = bfs(-1.4453125f);
  C_S2 = bfs(0.00262451171875f);
  C_P4 = bfs(0.00961812911f);
  C_R1 = bfs(-1.4545455f);
  C_R2 = bfs(0.32323232f);
}
// sigmoid(x) for bf16 x (pfx_sig.h sigm)
static inline V sigm(V x) {
  V ta = _mm512_mul_ps(x, bc(C_S1));
  ta = fma_(x, bc(C_S2), ta);
  V u = _mm512_castsi512_ps(_mm512_or_si512(_mm512_castps_si512(ta), ibc((int)0x80000000u)));
  u = _mm512_max_ps(u, bc(-126.0f));
  V um = _mm512_add_ps(u, bc(12582912.0f));
  V n = _mm512_sub_ps(um, bc(12582912.0f));
  V f = _mm512_sub_ps(u, n);
  VI e = _mm512_sub_epi32(_mm512_castps_si512(um), ibc(0x4B400000 - 127));
  V pw = _mm512_castsi512_ps(_mm512_slli_epi32(e, 23));
  T2 fs = sp2(f);
  V p = fma_(fs.h, bc(C_P4), bc(0.0555041087f));
  p = fma_(fs.h, bfr(p), bc(0.240226507f));
  p = fma_(fs.h, bfr(p), bc(0.693147181f));
  {
    T2 ps = sp2(p);
    V r = fma_(ps.h, fs.h, bc(1.0f));
    r = fma_(ps.h, fs.l, r);
    p = fma_(ps.l, fs.h, r);
  }
  V s2 = bfr(pw);
  T2 ps = sp2(p);
  V ea = _mm512_mul_ps(ps.h, s2);
  ea = fma_(ps.l, s2, ea);
  T2 es = sp2(ea);
  V da = _mm512_add_ps(ea, bc(1.0f));
  T2 d = sp2(da);
  V dd = bfr(_mm512_mul_ps(d.h, d.h));
  V ra = fma_(d.h, bc(C_R1), bc(2.1212121f));
  ra = fma_(dd, bc(C_R2), ra);
  {
    V r0 = ra;
    T2 rs = sp2(r0);
    V res = _mm512_fnmadd_ps(d.h, rs.h, bc(1.0f));
    res = _mm512_fnmadd_ps(d.h, rs.l, res);
    res = _mm512_fnmadd_ps(d.l, rs.h, res);
    ra = fma_(rs.h, bfr(res), r0);
  }
  V r = ra;
  T2 rs = sp2(r);
  V er = _mm512_mul_ps(es.h, rs.h);
  er = fma_(es.h, rs.l, er);
  er = fma_(es.l, rs.h, er);
  __mmask16 neg = _mm512_cmp_ps_mask(x, bc(0.0f), _CMP_LT_OQ);
  return _mm512_mask_blend_ps(neg, r, er);
}
// the accumulator's fp32 add (measured on the NPU's residual adds): the smaller operand rounded (to nearest even) to 2
// bits below the larger one's ulp, then the sum rounded to fp32 (IEEE: exact alignment). Done in fp64 (exact).
static inline __m512d addm_pd(__m512d a, __m512d b) {
  __m512d na = _mm512_abs_pd(a), nb = _mm512_abs_pd(b);
  __mmask8 ab = _mm512_cmp_pd_mask(na, nb, _CMP_GE_OQ);
  __m512d big = _mm512_mask_blend_pd(ab, b, a), sm = _mm512_mask_blend_pd(ab, a, b);
  __m512d e = _mm512_getexp_pd(big);                        // floor(log2 |big|) (-inf for 0)
  __m512d c = _mm512_scalef_pd(_mm512_set1_pd(1.5), _mm512_add_pd(e, _mm512_set1_pd(27.0)));
  sm = _mm512_sub_pd(_mm512_add_pd(sm, c), c);
  return _mm512_add_pd(big, sm);
}
static inline V addm(V a, V b) {
  if (g_knob[2] == 1 && !g_knob[3])
    return _mm512_add_ps(a, b);
  __m512d lo = addm_pd(_mm512_cvtps_pd(_mm512_castps512_ps256(a)), _mm512_cvtps_pd(_mm512_castps512_ps256(b)));
  __m512d hi = addm_pd(_mm512_cvtps_pd(_mm512_extractf32x8_ps(a, 1)), _mm512_cvtps_pd(_mm512_extractf32x8_ps(b, 1)));
  return _mm512_insertf32x8(_mm512_castps256_ps512(_mm512_cvtpd_ps(lo)), _mm512_cvtpd_ps(hi), 1);
}
// sigmoid(x) and silu(x) / 32 of every bf16 x (sigm, mulf: the NPU's results exactly; 2 x 256 KB, made once)
alignas(64) static float g_sig[65536], g_sz[65536];
static void init_tables() {
  alignas(64) int idx[16];
  for (int b = 0; b < 65536; b += 16) {
    for (int i = 0; i < 16; i++)
      idx[i] = (b + i) << 16;
    V x = _mm512_castsi512_ps(_mm512_loadu_si512(idx));
    V sb = sigm(x);
    _mm512_store_ps(g_sig + b, sb);
    _mm512_store_ps(g_sz + b, mulf(_mm512_mul_ps(x, bc(0.03125f)), sb));
  }
}
static inline VI bfidx(V x) { return _mm512_srli_epi32(_mm512_castps_si512(x), 16); }   // (x a bf16 value)
static inline V bf16x16(const uint16_t *p) {   // 16 bf16 -> fp32
  __m256i b = _mm256_loadu_si256((const __m256i *)p);
  return _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(b), 16));
}
static inline V bf16x8x2(const uint16_t *p0, const uint16_t *p1) {   // 8 + 8 bf16 -> fp32
  __m128i a = _mm_loadu_si128((const __m128i *)p0), b = _mm_loadu_si128((const __m128i *)p1);
  __m256i c = _mm256_inserti128_si256(_mm256_castsi128_si256(a), b, 1);
  return _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(c), 16));
}
static inline V ld8x2(const float *p0, const float *p1) {
  return _mm512_insertf32x8(_mm512_castps256_ps512(_mm256_loadu_ps(p0)), _mm256_loadu_ps(p1), 1);
}

// ---------------------------------------------------------------- FWHT 1024 (the NPU's stage order)
static inline V negv(V a) { return _mm512_xor_ps(a, bc(-0.0f)); }
static inline V fadd(V a, V b) { return g_knob[3] ? addm(a, b) : _mm512_add_ps(a, b); }
static inline V fsub(V a, V b) { return g_knob[3] ? addm(a, negv(b)) : _mm512_sub_ps(a, b); }
static inline void bfly(float *x, int h) {   // stage of stride h >= 16
  for (int i = 0; i < 1024; i += 2 * h)
    for (int j = i; j < i + h; j += 16) {
      V a = _mm512_load_ps(x + j), b = _mm512_load_ps(x + j + h);
      _mm512_store_ps(x + j, fadd(a, b));
      _mm512_store_ps(x + j + h, fsub(a, b));
    }
}
static VI P8, P4, P2, P1;
static void init_perm() {
  int a[16];
  for (int s : {8, 4, 2, 1}) {
    for (int i = 0; i < 16; i++)
      a[i] = i ^ s;
    VI v = _mm512_loadu_si512(a);
    (s == 8 ? P8 : s == 4 ? P4 : s == 2 ? P2 : P1) = v;
  }
}
static inline V intra(V v, VI perm, __mmask16 up) {   // stride s < 16 inside the vector: lo + hi | lo - hi
  V w = _mm512_permutexvar_ps(perm, v);
  return _mm512_mask_blend_ps(up, fadd(v, w), fsub(w, v));
}
static void fwht(float *x, bool skip512 = false) {   // (skip512: two H_512 blocks - Prism's block 512)
  bfly(x, 16);
  bfly(x, 32);
  bfly(x, 64);
  bfly(x, 128);
  bfly(x, 256);
  if (!skip512)
    bfly(x, 512);
  for (int j = 0; j < 1024; j += 16) {
    V v = _mm512_load_ps(x + j);
    v = intra(v, P8, 0xFF00);
    v = intra(v, P4, 0xF0F0);
    v = intra(v, P2, 0xCCCC);
    v = intra(v, P1, 0xAAAA);
    _mm512_store_ps(x + j, v);
  }
}

// ---------------------------------------------------------------- bfp16ebs8: 8 values -> exponent byte + 8 int8
static inline void bfp8(const float *v, uint8_t *o) {
  __m256 x = _mm256_loadu_ps(v);
  __m256 ax = _mm256_andnot_ps(_mm256_set1_ps(-0.0f), x);
  __m256 m1 = _mm256_max_ps(ax, _mm256_permute2f128_ps(ax, ax, 1));
  m1 = _mm256_max_ps(m1, _mm256_shuffle_ps(m1, m1, 0x4E));
  m1 = _mm256_max_ps(m1, _mm256_shuffle_ps(m1, m1, 0xB1));
  uint32_t mb = (uint32_t)_mm_cvtsi128_si32(_mm_castps_si128(_mm256_castps256_ps128(m1)));
  int E = (int)(mb >> 23);                      // the max's biased exponent: value = m * 2^(E - 133)
  if (E == 0) {                                 // (zero / denormal block)
    memset(o, 0, 9);
    return;
  }
  int sh = 133 - E;                             // scale 2^sh (E in 1..254: sh in -121..132)
  __m256 sc;
  {
    int e2 = sh + 127;
    if (e2 < 1 || e2 > 254) {                   // (out of fp32's range: scale in two steps)
      float s = ldexpf(1.0f, sh);
      sc = _mm256_set1_ps(s);
    } else {
      sc = _mm256_castsi256_ps(_mm256_set1_epi32(e2 << 23));
    }
  }
  __m256 y = _mm256_mul_ps(x, sc);
  __m256 r = _mm256_round_ps(y, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
  // (the NPU: a value rounding to +128 moves the block to the next exponent; -128 stays)
  if (g_knob[0] != 1 && _mm256_movemask_ps(_mm256_cmp_ps(r, _mm256_set1_ps(127.0f), _CMP_GT_OQ))) {
    E++;
    y = _mm256_mul_ps(y, _mm256_set1_ps(0.5f));
    r = _mm256_round_ps(y, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
  }
  r = _mm256_min_ps(_mm256_max_ps(r, _mm256_set1_ps(-128.0f)), _mm256_set1_ps(127.0f));
  __m256i iv = _mm256_cvtps_epi32(r);
  __m128i lo = _mm256_castsi256_si128(iv), hi = _mm256_extracti128_si256(iv, 1);
  __m128i p16 = _mm_packs_epi32(lo, hi);
  __m128i p8 = _mm_packs_epi16(p16, p16);
  o[0] = (uint8_t)E;
  _mm_storel_epi64((__m128i *)(o + 1), p8);
}

// ---------------------------------------------------------------- an op
enum { N1, R1N, R3N, R3, ACT, DNO, ROTP, ROTB, MOEG, MOEA, MOES };
static const int ROTB_T[6][4] = {{0, 1, 512, 0},    {512, 2, 256, 544}, {1312, 1, 512, 0},
                                 {2112, 1, 512, 0}, {2624, 2, 256, 544}, {3424, 1, 512, 0}};
static const int HW_ = 1536, CT_IN = 10752, CT_POST = 18944;

struct Op {
  int op, M, H, I, ND, KO, NO, ZB, KX, KS, NMT, NB, nunits, nq;
  int rb, G;                                     // the rotation block (P header [12]; 0: 1024) / the 1024-chunks it spans
  size_t S;                                      // words between Y pieces (32 M)
  float *h;
  const uint16_t *Y, *Z, *Ob;
  const float *Of;
  uint8_t *X;
  float *dbg;
  float invh;
  std::vector<float> ct, sg, dw;
  // [MoE] (cg_moe)
  const uint8_t *Xs;
  const int32_t *idx;
  const float *wt;
  int nk;                                        // (MOES: rows a token)
};

struct alignas(64) Scratch {
  float wk[8][4096];                             // 8 tokens x (up to 4) 1024-chunks of one rotation block
};

static inline uint8_t *xblk(const Op &g, int T, int ks, int i) {   // the 72 B of block T, K-tile ks, i
  int trg = T >> 3, z = T & 7, r = trg & 3, mt = trg >> 2;
  size_t w = ((size_t)(r * g.NMT + mt) * g.KS + ks) * 2304 + (z >> 1) * 576 + i * 36 + (z & 1) * 18;
  return g.X + w * 4;
}
// the 8 tokens' blocks of columns c .. c + 7 -> 72 B (as bfp8, two tokens a vector)
static inline void bfp72(const Scratch &s, int c, uint8_t *o) {
  alignas(16) uint8_t mb[16];
  for (int a = 0; a < 8; a += 2) {
    V v = ld8x2(s.wk[a] + c, s.wk[a + 1] + c);
    V ax = _mm512_abs_ps(v);
    V m = _mm512_max_ps(ax, _mm512_permute_ps(ax, 0xB1));
    m = _mm512_max_ps(m, _mm512_permute_ps(m, 0x4E));
    m = _mm512_max_ps(m, _mm512_shuffle_f32x4(m, m, 0xB1));
    VI E = _mm512_srli_epi32(_mm512_castps_si512(m), 23);
    V y = _mm512_scalef_ps(v, _mm512_cvtepi32_ps(_mm512_sub_epi32(ibc(133), E)));
    V r = _mm512_roundscale_ps(y, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
    __mmask16 over = _mm512_cmp_ps_mask(r, bc(127.0f), _CMP_GT_OQ);
    if (over) {                                 // (a value rounding to +128: the block's exponent + 1)
      __mmask16 bm = (__mmask16)(((over & 0xFF) ? 0xFF : 0) | ((over & 0xFF00) ? 0xFF00 : 0));
      E = _mm512_mask_add_epi32(E, bm, E, ibc(1));
      y = _mm512_mask_mul_ps(y, bm, y, bc(0.5f));
      r = _mm512_roundscale_ps(y, _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC);
    }
    __mmask16 nz = _mm512_cmpneq_epi32_mask(E, _mm512_setzero_si512());
    r = _mm512_maskz_mov_ps(nz, _mm512_max_ps(r, bc(-128.0f)));
    _mm_store_si128((__m128i *)mb, _mm512_cvtsepi32_epi8(_mm512_cvtps_epi32(r)));
    uint8_t *p = o + a * 9;
    p[0] = (uint8_t)_mm_cvtsi128_si32(_mm512_castsi512_si128(E));
    memcpy(p + 1, mb, 8);
    p[9] = (uint8_t)_mm_cvtsi128_si32(_mm512_extracti32x4_epi32(E, 2));
    memcpy(p + 10, mb + 8, 8);
  }
}
// the unit's 8 tokens x G chunks (q0 ..) -> FWHT per chunk (Prism's block: 512 leaves the last stage out, 4096 adds
// the 4-chunk butterfly / 2 after the chunks, as the decode kernels) -> X (and dbg)
static void emit(const Op &g, Scratch &s, int T, int q0) {
  if (!(g_knob[6] & 2))
    for (int a = 0; a < 8; a++)
      for (int j = 0; j < g.G; j++)
        fwht(s.wk[a] + j * 1024, g.rb == 512);
  if (g.G == 4)
    for (int a = 0; a < 8; a++) {
      float *x = s.wk[a];
      for (int c = 0; c < 1024; c += 16) {
        V x0 = _mm512_load_ps(x + c), x1 = _mm512_load_ps(x + 1024 + c);
        V x2 = _mm512_load_ps(x + 2048 + c), x3 = _mm512_load_ps(x + 3072 + c);
        V p = _mm512_add_ps(x0, x1), m = _mm512_sub_ps(x0, x1), u = _mm512_add_ps(x2, x3), w = _mm512_sub_ps(x2, x3);
        V h = bc(0.5f);
        _mm512_store_ps(x + c, _mm512_mul_ps(_mm512_add_ps(p, u), h));
        _mm512_store_ps(x + 1024 + c, _mm512_mul_ps(_mm512_add_ps(m, w), h));
        _mm512_store_ps(x + 2048 + c, _mm512_mul_ps(_mm512_sub_ps(p, u), h));
        _mm512_store_ps(x + 3072 + c, _mm512_mul_ps(_mm512_sub_ps(m, w), h));
      }
    }
  if (g_knob[6] & 4)
    return;
  for (int j = 0; j < g.G; j++) {
    int q = q0 + j;
    if (g.dbg)
      for (int a = 0; a < 8; a++)
        memcpy(g.dbg + (size_t)(T * 8 + a) * g.KX + q * 1024, s.wk[a] + j * 1024, 4096);
    for (int kk = 0; kk < 8; kk++)
      for (int i = 0; i < 16; i++) {
        uint8_t *o = xblk(g, T, q * 8 + kk, i);
        int c = j * 1024 + kk * 128 + i * 8;
        if (g_knob[4])
          for (int a = 0; a < 8; a++)
            bfp8(s.wk[a] + c, o + a * 9);
        else
          bfp72(s, c, o);
      }
  }
}
// 1 / rms of 16 lanes of sums of squares s (sc: the mean's scale)
static inline V rms_r(V s, float sc) { return rsqrt16(addf(mulf(s, bc(sc)), bc(1e-6f), 1.0f)); }

static void unit_norm(const Op &g, Scratch &s, int T) {
  int NPC = g.H / 64, np = g.op == R1N ? g.NO : (g.op == R3N || g.op == R3) ? g.ND : 0;
  float *hb = g.h + (size_t)T * NPC * 576;
  V acc[4] = {_mm512_setzero_ps(), _mm512_setzero_ps(), _mm512_setzero_ps(), _mm512_setzero_ps()};
  for (int p = 0; p < NPC; p++) {
    float *hp = hb + p * 576;
    for (int yp = 0; yp < np; yp++) {
      const uint16_t *y = g.Y + ((size_t)(p * np + yp) * g.S + (size_t)T * 256) * 2;
      for (int v = 0; v < 512; v += 16)
        _mm512_storeu_ps(hp + v, addm(_mm512_loadu_ps(hp + v), bf16x16(y + v)));
    }
    if (g.op == R3)
      continue;
    for (int j = 0; j < 8; j++)
      for (int k = 0; k < 4; k++) {
        T2 x = sp2(_mm512_loadu_ps(hp + j * 64 + 16 * k));
        acc[k] = fma_(x.h, x.h, acc[k]);
        acc[k] = fma_(x.h, _mm512_add_ps(x.l, x.l), acc[k]);
      }
  }
  if (g.op == R3)
    return;
  alignas(64) float sq[64], ss[16];
  for (int k = 0; k < 4; k++)
    _mm512_store_ps(sq + 16 * k, acc[k]);
  for (int i = 0; i < 16; i++)
    ss[i] = 4096.0f;
  for (int a = 0; a < 8; a++) {
    float t = 0.0f;
    for (int b = 0; b < 8; b++)
      t += sq[a * 8 + b];
    ss[a] = t;
  }
  alignas(64) float rt[16];
  _mm512_store_ps(rt, rms_r(_mm512_load_ps(ss), g.invh));
  for (int q = 0; q < g.H / 1024; q++) {
    for (int a = 0; a < 8; a++) {
      V r = bc(rt[a]);
      for (int c = 0; c < 1024; c += 16) {      // columns c .. c + 15: piece c / 64, j = (c % 64) / 8 (2 of them)
        int col = q * 1024 + c;
        const float *hp = hb + (col >> 6) * 576 + ((col & 63) >> 3) * 64 + a * 8;
        V x = ld8x2(hp, hp + 64);
        V v = mulf(mulf(x, _mm512_loadu_ps(g.ct.data() + col)), r);
        _mm512_store_ps(s.wk[a] + (q % g.G) * 1024 + c, v);
      }
    }
    if (q % g.G == g.G - 1)
      emit(g, s, T, q - g.G + 1);
  }
}
static void unit_act(const Op &g, Scratch &s, int T, int q, int off) {   // (fills chunk q at wk + off)
  size_t NB = g.NB;
  for (int u = 0; u < 16; u++) {
    int mg = q * 32 + u, mu = mg + 16;
    const uint16_t *yg = g.Y + (((size_t)(mg >> 1) * NB + T) * 512 + (mg & 1) * 256) * 2;
    const uint16_t *yu = g.Y + (((size_t)(mu >> 1) * NB + T) * 512 + (mu & 1) * 256) * 2;
    for (int a = 0; a < 8; a++)
      for (int j = 0; j < 8; j += 2) {        // columns u*64 + j*8 .. +16
        V gt = bf16x8x2(yg + j * 64 + a * 8, yg + (j + 1) * 64 + a * 8);
        V up = bf16x8x2(yu + j * 64 + a * 8, yu + (j + 1) * 64 + a * 8);
        int c = u * 64 + j * 8;
        V us = _mm512_mul_ps(up, _mm512_loadu_ps(g.sg.data() + q * 1024 + c));
        V gu = _mm512_mul_ps(gt, us);
        V sb = g_knob[7] ? sigm(gt) : _mm512_i32gather_ps(bfidx(gt), g_sig, 4);
        _mm512_store_ps(s.wk[a] + off + c, (g_knob[6] & 1) ? gu : mulf(gu, sb));
      }
  }
}
static void unit_dno(const Op &g, Scratch &s, int T, int q, int off) {
  size_t NB = g.NB;
  for (int hd = 0; hd < 8; hd++) {
    int head = q * 8 + hd;
    const float *ob = g.Of + ((size_t)head * NB + T) * 1024;
    for (int a = 0; a < 8; a++)
      for (int w = 0; w < 128; w += 16)
        _mm512_store_ps(s.wk[a] + off + hd * 128 + w, _mm512_loadu_ps(ob + (w >> 5) * 256 + a * 32 + (w & 31)));
  }
  alignas(64) float ssum[64];
  for (int i = 0; i < 64; i++) {               // (head i >> 3, token i & 7)
    const float *p = s.wk[i & 7] + off + (i >> 3) * 128;
    alignas(64) float l[64];
    for (int k = 0; k < 4; k++) {
      V acc = _mm512_setzero_ps();
      for (int h2 = 0; h2 < 2; h2++) {
        T2 x = sp2(_mm512_load_ps(p + h2 * 64 + 16 * k));
        acc = fma_(x.h, x.h, acc);
        acc = fma_(x.h, _mm512_add_ps(x.l, x.l), acc);
      }
      _mm512_store_ps(l + 16 * k, acc);
    }
    for (int n = 32; n >= 1; n >>= 1)          // (reduce: halving)
      for (int k = 0; k < n; k++)
        l[k] = l[k] + l[k + n];
    ssum[i] = l[0];
  }
  alignas(64) float rt[64];
  for (int k = 0; k < 4; k++)
    _mm512_store_ps(rt + 16 * k, rms_r(_mm512_load_ps(ssum + 16 * k), 1.0f / 128.0f));
  for (int hd = 0; hd < 8; hd++)
    for (int hz = 0; hz < 2; hz++) {
      int p = q * 16 + hd * 2 + hz;
      const uint16_t *z = g.Z + ((size_t)(g.ZB + p) * g.S + (size_t)T * 256) * 2;
      for (int a = 0; a < 8; a++) {
        V r = bc(rt[hd * 8 + a]);
        for (int j = 0; j < 8; j += 2) {
          int w = hz * 64 + j * 8;              // within the head
          V zv = bf16x8x2(z + j * 64 + a * 8, z + (j + 1) * 64 + a * 8);
          V sgv = _mm512_loadu_ps(g.sg.data() + q * 1024 + hd * 128 + w);
          V sz = g_knob[7] ? mulf(_mm512_mul_ps(zv, sgv), sigm(zv))
                           : _mm512_mul_ps(_mm512_i32gather_ps(bfidx(zv), g_sz, 4), _mm512_mul_ps(sgv, bc(32.0f)));
          V rw = mulf(_mm512_loadu_ps(g.dw.data() + w), r);
          float *pp = s.wk[a] + off + hd * 128 + w;
          _mm512_store_ps(pp, mulf(mulf(_mm512_load_ps(pp), rw), sz));
        }
      }
    }
}
static void unit_rot(const Op &g, Scratch &s, int T, int q, int off) {
  for (int a = 0; a < 8; a++) {
    int t = T * 8 + a;
    const float *sgq = g.sg.data() + q * 1024;
    if (g.op == ROTP) {
      const uint16_t *y = g.Ob + ((size_t)t * 2176 + q * 544) * 2;
      for (int i = 0; i < 1024; i += 16)
        _mm512_store_ps(s.wk[a] + off + i, _mm512_mul_ps(bf16x16(y + i), _mm512_loadu_ps(sgq + i)));
    } else {
      const int *rb = ROTB_T[q];
      int i = 0;
      for (int pc = 0; pc < rb[1]; pc++) {
        const uint16_t *y = g.Ob + ((size_t)t * 4224 + rb[0] + pc * rb[3]) * 2;
        for (int e = 0; e < 2 * rb[2]; e += 16, i += 16)
          _mm512_store_ps(s.wk[a] + off + i, _mm512_mul_ps(bf16x16(y + e), _mm512_loadu_ps(sgq + i)));
      }
    }
  }
}
// ---------------------------------------------------------------- [MoE] the prefill's expert glue (cg_moe)
// MOEG: rows of an M-token X (bfp16 tiles, K 2048: KS 16) -> segments of 256 rows (X layouts of M = 256), row r of
//   the output = token idx[r] (-1: zeros); each token's 8 values of a column group are 9 bytes of their own (one
//   exponent a token), so this is a byte gather - exactly the values of the source X. Unit: 8 output rows.
// MOEA: Y [rows][1024] bf16 (gate | up) -> silu(gate) * up, x' = H_512 (s * a) / 32 (s: g.sg, 512) -> X of K 512
//   (KS 4) per 256-row segment. Unit: 8 rows.
// MOES: h (glue layout [T][H / 64][576] fp32) += sum_k wt[t * nk + k] * Y[idx[t * nk + k]] (bf16 rows of 2048; idx
//   -1: none). Unit: 8 tokens (one h block: no two threads write a token).
static inline uint8_t *xrec(uint8_t *X, int NMT, int KS, int t, int gcol) {   // token t's 9 bytes of column group
  int T = t >> 3, a = t & 7, trg = T >> 3, z = T & 7, r = trg & 3, mt = trg >> 2, ks = gcol >> 4, i = gcol & 15;
  size_t w = ((size_t)(r * NMT + mt) * KS + ks) * 2304 + (z >> 1) * 576 + i * 36 + (z & 1) * 18;
  return X + w * 4 + a * 9;
}
static void unit_moeg(const Op &g, int u) {
  int S = g.KX, nb = S / 8, seg = u / nb, blk = u % nb;     // (S: rows a segment, 64 or 256)
  uint8_t *dst = g.X + (size_t)seg * S * 2048 * 9 / 8;
  int nmt_src = g.M / 256;
  for (int a = 0; a < 8; a++) {
    int row = blk * 8 + a, t = g.idx[seg * S + row];
    for (int gc = 0; gc < 256; gc++) {
      uint8_t *d = xrec(dst, S >= 256 ? S / 256 : 1, 16, row, gc);
      if (t < 0)
        memset(d, 0, 9);
      else
        memcpy(d, xrec((uint8_t *)g.Xs, nmt_src, 16, t, gc), 9);
    }
  }
}
static void unit_moea(const Op &g, Scratch &s, int u) {
  int S = g.KX, nb = S / 8, seg = u / nb, blk = u % nb;
  for (int a = 0; a < 8; a++) {
    const uint16_t *y = g.Y + ((size_t)seg * S + blk * 8 + a) * 1024;
    float *x = s.wk[a];
    for (int c = 0; c < 512; c += 16) {
      V gt = bf16x16(y + c), up = bf16x16(y + 512 + c);
      alignas(64) float gg[16];
      _mm512_store_ps(gg, gt);
      for (int k = 0; k < 16; k++)
        gg[k] = gg[k] / (1.0f + expf(-gg[k]));
      V v = _mm512_mul_ps(_mm512_mul_ps(_mm512_load_ps(gg), up), _mm512_loadu_ps(g.sg.data() + c));
      _mm512_store_ps(x + c, v);
    }
    for (int c = 512; c < 1024; c += 16)
      _mm512_store_ps(x + c, _mm512_setzero_ps());
    fwht(x, true);                                 // (two H_512 blocks: the second is all zeros)
  }
  uint8_t *X = g.X + (size_t)seg * S * 512 * 9 / 8;
  int nmt = S >= 256 ? S / 256 : 1;
  for (int kk = 0; kk < 4; kk++)
    for (int i = 0; i < 16; i++) {
      int T = blk;
      int trg = T >> 3, z = T & 7, r = trg & 3, mt = trg >> 2;
      size_t w = ((size_t)(r * nmt + mt) * 4 + kk) * 2304 + (z >> 1) * 576 + i * 36 + (z & 1) * 18;
      bfp72(s, kk * 128 + i * 8, X + w * 4);
    }
}
static void unit_moes(const Op &g, int u) {
  int NPC = g.H / 64;
  for (int a = 0; a < 8; a++) {
    int t = u * 8 + a;
    if (t >= g.M)
      return;
    float *hb = g.h + (size_t)u * NPC * 576;
    for (int k = 0; k < g.nk; k++) {
      int r = g.idx[t * g.nk + k];
      if (r < 0)
        continue;
      V w = bc(g.wt[t * g.nk + k]);
      const uint16_t *y = g.Y + (size_t)r * g.H;
      for (int c = 0; c < g.H; c += 16) {
        V yv = bf16x16(y + c);
        for (int hh = 0; hh < 2; hh++) {
          int cc = c + 8 * hh;
          float *hp = hb + (cc >> 6) * 576 + ((cc & 63) >> 3) * 64 + a * 8;
          __m256 y8 = hh ? _mm512_extractf32x8_ps(yv, 1) : _mm512_castps512_ps256(yv);
          _mm256_storeu_ps(hp, _mm256_fmadd_ps(y8, _mm512_castps512_ps256(w), _mm256_loadu_ps(hp)));
        }
      }
    }
  }
}

// a unit: the 2 blocks of an X z-pair (one thread writes whole 2304-B runs of X: no shared cache lines)
static void do_unit(const Op &g, Scratch &s, int u) {
  if (g.op == MOEG)
    return unit_moeg(g, u);
  if (g.op == MOEA)
    return unit_moea(g, s, u);
  if (g.op == MOES)
    return unit_moes(g, u);
  int ng = g.nq / g.G;                           // (a unit: one rotation block of the row, 2 token blocks)
  int T = 2 * (g.op <= R3 ? u : u / ng), q0 = g.op <= R3 ? 0 : (u % ng) * g.G;
  for (int t = T; t < T + 2; t++) {
    if (g.op <= R3) {
      unit_norm(g, s, t);
      continue;
    }
    for (int j = 0; j < g.G; j++)
      switch (g.op) {
      case ACT:
        unit_act(g, s, t, q0 + j, j * 1024);
        break;
      case DNO:
        unit_dno(g, s, t, q0 + j, j * 1024);
        break;
      default:
        unit_rot(g, s, t, q0 + j, j * 1024);
      }
    emit(g, s, t, q0);
  }
}

#if defined(__GNUC__) && !defined(_MSC_VER)
#pragma GCC pop_options
#endif

// ---------------------------------------------------------------- cores / the pool
static std::vector<int> g_cores;               // chosen logical processors (one a physical core)
static int g_classes = 0, g_phys = 0;
#ifdef _WIN32
static std::vector<GROUP_AFFINITY> g_aff;
static void find_cores() {
  DWORD len = 0;
  GetLogicalProcessorInformationEx(RelationProcessorCore, nullptr, &len);
  std::vector<uint8_t> buf(len);
  if (!GetLogicalProcessorInformationEx(RelationProcessorCore, (SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *)buf.data(), &len))
    return;
  int lo = 255;
  for (DWORD off = 0; off < len;) {
    auto *q = (SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *)(buf.data() + off);
    g_classes |= 1 << q->Processor.EfficiencyClass;
    lo = q->Processor.EfficiencyClass < lo ? q->Processor.EfficiencyClass : lo;
    g_phys++;
    off += q->Size;
  }
  bool two = (g_classes & (g_classes - 1)) != 0;
  for (DWORD off = 0; off < len;) {
    auto *q = (SYSTEM_LOGICAL_PROCESSOR_INFORMATION_EX *)(buf.data() + off);
    if (!two || q->Processor.EfficiencyClass == lo) {
      GROUP_AFFINITY ga = q->Processor.GroupMask[0];
      for (int b = 0; b < 64; b++)
        if (ga.Mask & (1ull << b)) {          // (one logical processor a core)
          GROUP_AFFINITY one{};
          one.Group = ga.Group;
          one.Mask = 1ull << b;
          g_aff.push_back(one);
          g_cores.push_back(ga.Group * 64 + b);
          break;
        }
    }
    off += q->Size;
  }
}
static void pin_self(int k, int mode) {
  if ((mode & 1) && !g_aff.empty())
    SetThreadGroupAffinity(GetCurrentThread(), &g_aff[k % g_aff.size()], nullptr);
  if (mode & 2) {
    THREAD_POWER_THROTTLING_STATE st{};
    st.Version = THREAD_POWER_THROTTLING_CURRENT_VERSION;
    st.ControlMask = THREAD_POWER_THROTTLING_EXECUTION_SPEED;
    st.StateMask = THREAD_POWER_THROTTLING_EXECUTION_SPEED;
    SetThreadInformation(GetCurrentThread(), ThreadPowerThrottling, &st, sizeof(st));
  }
}
static bool has_avx512() {
  int r[4];
  __cpuidex(r, 0, 0);
  if (r[0] < 7)
    return false;
  __cpuidex(r, 1, 0);
  if (!(r[2] & (1 << 27)))                     // OSXSAVE
    return false;
  unsigned long long xcr = _xgetbv(0);
  if ((xcr & 0xE6) != 0xE6)                    // the OS saves the zmm state
    return false;
  __cpuidex(r, 7, 0);
  bool f = (r[1] & (1 << 16)) && (r[1] & (1 << 17)) && (r[1] & (1 << 30)) && (r[1] & (1 << 31));   // F DQ BW VL
  return f;
}
#else
static long rd_long(const char *path) {
  FILE *f = fopen(path, "r");
  if (!f)
    return -1;
  long v = -1;
  if (fscanf(f, "%ld", &v) != 1)
    v = -1;
  fclose(f);
  return v;
}
static void find_cores() {
  int n = (int)sysconf(_SC_NPROCESSORS_CONF);
  std::vector<long> perf(n, -1), key(n, -1);
  char p[256];
  for (int c = 0; c < n; c++) {
    snprintf(p, sizeof p, "/sys/devices/system/cpu/cpu%d/topology/core_id", c);
    long core = rd_long(p);
    snprintf(p, sizeof p, "/sys/devices/system/cpu/cpu%d/topology/physical_package_id", c);
    long pkg = rd_long(p);
    if (core < 0)
      continue;
    key[c] = pkg * 100000 + core;
    snprintf(p, sizeof p, "/sys/devices/system/cpu/cpu%d/cpufreq/cpuinfo_max_freq", c);
    perf[c] = rd_long(p);
    if (perf[c] <= 0) {
      snprintf(p, sizeof p, "/sys/devices/system/cpu/cpu%d/acpi_cppc/highest_perf", c);
      perf[c] = rd_long(p);
    }
  }
  std::vector<int> first;                      // one logical cpu a physical core
  for (int c = 0; c < n; c++) {
    if (key[c] < 0)
      continue;
    bool dup = false;
    for (int f : first)
      dup |= key[f] == key[c];
    if (!dup)
      first.push_back(c);
  }
  g_phys = (int)first.size();
  long lo = -1, hi = -1;
  for (int c : first)
    if (perf[c] > 0) {
      lo = lo < 0 || perf[c] < lo ? perf[c] : lo;
      hi = perf[c] > hi ? perf[c] : hi;
    }
  bool two = lo > 0 && hi > lo * 5 / 4;        // (2 classes: the fastest cores > 1.25 x the slowest)
  g_classes = two ? 3 : 1;
  for (int c : first)
    if (!two || perf[c] <= lo * 23 / 20)
      g_cores.push_back(c);
}
static void pin_self(int k, int mode) {
  if ((mode & 1) && !g_cores.empty()) {
    cpu_set_t s;
    CPU_ZERO(&s);
    CPU_SET(g_cores[k % g_cores.size()], &s);
    pthread_setaffinity_np(pthread_self(), sizeof s, &s);
  }
}
static bool has_avx512() {
  __builtin_cpu_init();
  return __builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512bw") &&
         __builtin_cpu_supports("avx512dq") && __builtin_cpu_supports("avx512vl");
}
#endif

struct Pool {
  std::vector<std::thread> th;
  std::mutex mu;
  std::condition_variable cv, done;
  uint64_t gen = 0;
  int pending = 0;
  const Op *job = nullptr;
  std::atomic<int> next{0};
  bool quit = false;
};
static Pool *g_pool = nullptr;
static std::once_flag g_once;
static void global_init() {
  init_consts();
  find_cores();
  if (has_avx512()) {
    init_perm();
    init_tables();
  }
}

static void worker(int k, int mode) {
  pin_self(k, mode);
  Scratch *s = (Scratch *)
#ifdef _WIN32
      _aligned_malloc(sizeof(Scratch), 64);
#else
      aligned_alloc(64, sizeof(Scratch));
#endif
  uint64_t seen = 0;
  for (;;) {
    const Op *g;
    {
      std::unique_lock<std::mutex> lk(g_pool->mu);
      g_pool->cv.wait(lk, [&] { return g_pool->quit || g_pool->gen != seen; });
      if (g_pool->quit)
        break;
      seen = g_pool->gen;
      g = g_pool->job;
    }
    for (int u; (u = g_pool->next.fetch_add(1)) < g->nunits;)
      do_unit(*g, *s, u);
    {
      std::lock_guard<std::mutex> lk(g_pool->mu);
      if (--g_pool->pending == 0)
        g_pool->done.notify_all();
    }
  }
#ifdef _WIN32
  _aligned_free(s);
#else
  free(s);
#endif
}

API void cg_info(int *out) {
  std::call_once(g_once, global_init);
  out[0] = has_avx512() ? 1 : 0;
  out[1] = g_classes;
  out[2] = (int)g_cores.size();
  out[3] = g_phys;
  out[4] = g_pool ? (int)g_pool->th.size() : 0;
  for (int i = 0; i < (int)g_cores.size() && i < 3; i++)
    out[5 + i] = g_cores[i];
}
API void cg_set(int key, int val) {
  if (key >= 0 && key < 8)
    g_knob[key] = val;
}
API int cg_init(int nthreads, int mode) {
  std::call_once(g_once, global_init);
  if (!has_avx512())
    return -1;
  if (g_pool)
    return (int)g_pool->th.size();
  g_pool = new Pool();
  int n = nthreads > 0 ? nthreads : (int)(g_cores.empty() ? 4 : g_cores.size());
  for (int k = 0; k < n; k++)
    g_pool->th.emplace_back(worker, k, mode);
  return n;
}
API void cg_close() {
  if (!g_pool)
    return;
  {
    std::lock_guard<std::mutex> lk(g_pool->mu);
    g_pool->quit = true;
  }
  g_pool->cv.notify_all();
  for (auto &t : g_pool->th)
    t.join();
  delete g_pool;
  g_pool = nullptr;
}

API int cg_run(int op, int M, int H, int I, int ND, int KO, int NO, int ZB, void *a0, void *a1, void *a2, void *a3,
               float *dbg) {
  if (!g_pool && cg_init(0, 3) < 0)
    return -1;
  if (op < 0 || op > ROTB)
    return -1;
  Op g;
  g.op = op;
  g.M = M, g.H = H, g.I = I, g.ND = ND, g.KO = KO, g.NO = NO, g.ZB = ZB;
  g.S = (size_t)32 * M;
  g.NB = M / 8;
  g.NMT = M / 256;
  g.h = nullptr, g.Y = g.Z = g.Ob = nullptr, g.Of = nullptr, g.X = nullptr;
  g.dbg = dbg;
  const int32_t *P;
  void *a[4] = {a0, a1, a2, a3};
  switch (op) {
  case N1: g.h = (float *)a[0], P = (const int32_t *)a[1], g.X = (uint8_t *)a[2]; break;
  case R1N: case R3N: g.h = (float *)a[0], g.Y = (const uint16_t *)a[1], P = (const int32_t *)a[2], g.X = (uint8_t *)a[3]; break;
  case R3: g.h = (float *)a[0], g.Y = (const uint16_t *)a[1], P = (const int32_t *)a[2]; break;
  case ACT: g.Y = (const uint16_t *)a[0], P = (const int32_t *)a[1], g.X = (uint8_t *)a[2]; break;
  case DNO: g.Of = (const float *)a[0], g.Z = (const uint16_t *)a[1], P = (const int32_t *)a[2], g.X = (uint8_t *)a[3]; break;
  default: g.Ob = (const uint16_t *)a[0], P = (const int32_t *)a[1], g.X = (uint8_t *)a[2];
  }
  g.KX = op == ACT ? I : (op >= DNO ? KO : H);
  g.KS = g.KX / 128;
  int oi = op == ROTB ? ROTP : op;                // (the P slot: rotb uses rotp's)
  const int32_t *hd = P + oi * HW_, *hd2 = P + oi * HW_ + 1024;
  memcpy(&g.invh, hd + 10, 4);
  g.rb = hd[12] ? hd[12] : 1024;
  if (g.rb != 512 && g.rb != 1024 && g.rb != 4096)
    return -2;
  g.G = g.rb == 4096 ? 4 : 1;
  int nsg = op == ACT ? I : KO;
  g.sg.resize(nsg + 64);
  for (int c = 0; c < nsg; c++) {
    int w = c >> 5;
    uint32_t bits = w < 448 ? (uint32_t)hd[64 + w] : (uint32_t)hd2[128 + w - 448];
    g.sg[c] = (bits >> (c & 31)) & 1 ? -0.03125f : 0.03125f;
  }
  if (op == DNO) {
    g.dw.resize(128);
    memcpy(g.dw.data(), hd2, 512);
  }
  if (op == N1 || op == R1N || op == R3N) {
    int ctin = CT_IN, ctpost = H <= 4096 ? CT_POST : CT_IN + 2 * H;
    const float *cf = (const float *)P + (op == R1N ? ctpost : ctin);
    g.ct.resize(H);
    for (int c = 0; c < H; c++)
      g.ct[c] = cf[(c >> 8) * 512 + (c & 255)];
  }
  if (op <= R3) {
    g.nq = H / 1024;
    g.nunits = g.NB / 2;
  } else {
    g.nq = (op == ACT ? I : KO) / 1024;
    g.nunits = g.NB / 2 * (g.nq / g.G);
  }
  {
    std::lock_guard<std::mutex> lk(g_pool->mu);
    g_pool->job = &g;
    g_pool->next.store(0);
    g_pool->pending = (int)g_pool->th.size();
    g_pool->gen++;
  }
  g_pool->cv.notify_all();
  {
    std::unique_lock<std::mutex> lk(g_pool->mu);
    g_pool->done.wait(lk, [] { return g_pool->pending == 0; });
  }
  return 0;
}

// [MoE] cg_moe(op, M, H, n, a0, a1, a2, a3): MOEG (M source tokens, n segments of H rows - 64 / 256): a0 X src, a1 idx
//   [n * H], a2 X out | MOEA (n segments of H rows): a0 Y [n * H][1024] bf16, a1 the 512 signs (float +-1), a2 X out |
//   MOES (M tokens, n = rows a token; H the width): a0 h, a1 Y [rows][H] bf16, a2 idx [M * n], a3 wt [M * n]
API int cg_moe(int op, int M, int H, int n, void *a0, void *a1, void *a2, void *a3) {
  if (!g_pool && cg_init(0, 3) < 0)
    return -1;
  Op g;
  g.op = op, g.M = M, g.H = H, g.nk = n;
  g.h = nullptr, g.Y = g.Z = g.Ob = nullptr, g.Of = nullptr, g.X = nullptr, g.dbg = nullptr;
  g.Xs = nullptr, g.idx = nullptr, g.wt = nullptr;
  g.rb = 1024, g.G = 1;
  if (op == MOEG) {
    g.Xs = (const uint8_t *)a0, g.idx = (const int32_t *)a1, g.X = (uint8_t *)a2;
    g.KX = H, g.H = 2048;
    g.nunits = n * (H / 8);
  } else if (op == MOEA) {
    g.KX = H, g.H = 2048;
    g.Y = (const uint16_t *)a0, g.X = (uint8_t *)a2;
    g.sg.resize(512);
    for (int c = 0; c < 512; c++)
      g.sg[c] = ((const float *)a1)[c] * 0.03125f;
    g.nunits = n * (g.KX / 8);
  } else if (op == MOES) {
    g.h = (float *)a0, g.Y = (const uint16_t *)a1, g.idx = (const int32_t *)a2, g.wt = (const float *)a3;
    g.nunits = (M + 7) / 8;
  } else
    return -1;
  {
    std::lock_guard<std::mutex> lk(g_pool->mu);
    g_pool->job = &g;
    g_pool->next.store(0);
    g_pool->pending = (int)g_pool->th.size();
    g_pool->gen++;
  }
  g_pool->cv.notify_all();
  {
    std::unique_lock<std::mutex> lk(g_pool->mu);
    g_pool->done.wait(lk, [] { return g_pool->pending == 0; });
  }
  return 0;
}


// [0.8] the MoE router of T tokens (decode / MTP verify), on the calling thread: xn = x / rms(x) * g; the NPU router's
// 16 best candidates (approx: int4 logits, a row of astride floats a token) re-scored exactly (R's rows . xn); the
// top-k of those by score, their softmax (norm_topk); the shared expert's gate sigmoid(sg . xn)
#ifndef _MSC_VER
__attribute__((target("avx2,fma")))
#endif
static float dotf(const float *a, const float *b, int n) {   // (AVX2 + FMA: every CPU with an XDNA NPU has them)
  __m256 s0 = _mm256_setzero_ps(), s1 = s0, s2 = s0, s3 = s0;
  for (int i = 0; i < n; i += 32) {
    s0 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i), _mm256_loadu_ps(b + i), s0);
    s1 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i + 8), _mm256_loadu_ps(b + i + 8), s1);
    s2 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i + 16), _mm256_loadu_ps(b + i + 16), s2);
    s3 = _mm256_fmadd_ps(_mm256_loadu_ps(a + i + 24), _mm256_loadu_ps(b + i + 24), s3);
  }
  __m256 v = _mm256_add_ps(_mm256_add_ps(s0, s1), _mm256_add_ps(s2, s3));
  float r[8];
  _mm256_storeu_ps(r, v);
  return ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7]));
}

API int cg_route(int T, int hid, int E, int topk, const float *H, const float *approx, int astride, const float *R,
                 const float *sg, const float *g, int32_t *top, float *p, float *ws) {
  if (hid % 32 || E > 4096 || topk > 16)
    return -1;
  std::vector<float> xn(hid);
  std::vector<int> ord(E);
  for (int t = 0; t < T; t++) {
    const float *x = H + (size_t)t * hid;
    float ss = dotf(x, x, hid);
    float s = sqrtf(ss / (float)hid + 1e-6f);
    for (int i = 0; i < hid; i++)
      xn[i] = x[i] / s * g[i];
    const float *a = approx + (size_t)t * astride;
    for (int i = 0; i < E; i++)
      ord[i] = i;
    std::nth_element(ord.begin(), ord.begin() + 16, ord.end(), [a](int u, int v) { return a[u] > a[v]; });
    int cand[16];
    float lc[16];
    for (int k = 0; k < 16; k++) {
      cand[k] = ord[k];
      lc[k] = dotf(R + (size_t)cand[k] * hid, xn.data(), hid);
    }
    int o[16];
    for (int k = 0; k < 16; k++)
      o[k] = k;
    std::sort(o, o + 16, [&](int u, int v) { return lc[u] > lc[v] || (lc[u] == lc[v] && cand[u] < cand[v]); });
    float mx = lc[o[0]], sum = 0;
    for (int k = 0; k < topk; k++) {
      top[t * topk + k] = cand[o[k]];
      p[t * topk + k] = expf(lc[o[k]] - mx);
      sum += p[t * topk + k];
    }
    for (int k = 0; k < topk; k++)
      p[t * topk + k] /= sum;
    ws[t] = 1.0f / (1.0f + expf(-dotf(sg, xn.data(), hid)));
  }
  return 0;
}

// [0.8] the MTP verify's MoE slots: the union of T tokens' top-k experts (sorted) + the shared expert, n groups of 2
// slots (>= ngmin); MOEH2 group g's two header blocks get each token's weights of the group's slots at xs + 16 t
// (0: an expert the token did not choose; tokens T..VT-1: 0). Returns the groups (-1: more than ngmax).
API int cg_union(int T, int VT, int topk, const int32_t *top, const float *p, const float *ws, int shared, int ngmin,
                 int ngmax, int64_t *slots, int32_t *mh, int blk_w, int xs) {
  int u[64], n = 0;
  for (int i = 0; i < T * topk; i++)
    u[n++] = top[i];
  std::sort(u, u + n);
  n = (int)(std::unique(u, u + n) - u);
  int ng = (n + 2) / 2 > ngmin ? (n + 2) / 2 : ngmin;
  if (ng > ngmax)
    return -1;
  for (int s = 0; s < 2 * ng; s++)
    slots[s] = s < n ? u[s] : s == n ? shared : 0;
  for (int g = 0; g < ng; g++)
    for (int t = 0; t < VT; t++)
      for (int j = 0; j < 2; j++) {
        int s = 2 * g + j;
        float w = 0.0f;
        if (t < T) {
          if (s == n)
            w = ws[t];
          else if (s < n)
            for (int k = 0; k < topk; k++)
              if (top[t * topk + k] == u[s])
                w = p[t * topk + k];
        }
        int32_t bits;
        std::memcpy(&bits, &w, 4);
        mh[(g * 4 + 0) * blk_w + xs + 16 * t + j] = bits;
        mh[(g * 4 + 1) * blk_w + xs + 16 * t + j] = bits;
      }
  return ng;
}

// [0.8] a MoE program's weight fills pointed at experts: im[pos[i]] = base[i] + slots[slot_of[i]] * step (mod 2^32)
API void cg_patch(int n, uint32_t *im, const int64_t *pos, const int64_t *base, const int64_t *slot_of,
                  const int64_t *slots, int64_t step) {
  for (int i = 0; i < n; i++)
    im[pos[i]] = (uint32_t)(base[i] + slots[slot_of[i]] * step);
}

// [0.9] the experts a layer's router would pick for T rows (the expert streaming's prefetch: the next layer's router
// on this layer's h): xn = x / rms(x) * g, all E logits, the k best (any order) -> top [T * k]
API int cg_predict(int T, int hid, int E, int k, const float *H, const float *R, const float *g, int32_t *top) {
  if (hid % 32 || k > E)
    return -1;
  std::vector<float> xn(hid), lg(E);
  std::vector<int> ord(E);
  for (int t = 0; t < T; t++) {
    const float *x = H + (size_t)t * hid;
    float s = sqrtf(dotf(x, x, hid) / (float)hid + 1e-6f);
    for (int i = 0; i < hid; i++)
      xn[i] = x[i] / s * g[i];
    for (int e = 0; e < E; e++) {
      lg[e] = dotf(R + (size_t)e * hid, xn.data(), hid);
      ord[e] = e;
    }
    std::nth_element(ord.begin(), ord.begin() + k, ord.end(), [&](int u, int v) { return lg[u] > lg[v]; });
    for (int j = 0; j < k; j++)
      top[t * k + j] = ord[j];
  }
  return 0;
}