code wiki / (root) / nx_gen_linear_verify.nx
`nx_gen_linear_verify.nx` source

↩ module page · 415 lines · 18940 B
1// nx_gen_linear_verify.nx -- SOVEREIGN linear projection (matmul), verified vs a reference.
2//
3//   y == x @ W^T
4//
5// This is the core compute of every diffusion and transformer architecture: qkv projection,
6// attention out-projection, every FFN matrix, every text-encoder layer. It is also where the
7// entire sovereignty gap lives -- correctness here was proven cheaply; THROUGHPUT is the wall.
8//
9// Usage:  nx_gen_linear_verify <model_id> <x_name> <w_name> <y_name> [rows] [pass_idx] [accum] [nworkers]
10//
11// Names are ABSOLUTE fixture names -- weights carry their own checkpoint name
12// ("model.diffusion_model.layers.0.attention.qkv.weight"), so nothing here is tied to one
13// architecture or one block. Shapes are read from the manifest, never assumed.
14//
15// LAYOUT (ggml): ggml_mul_mat(w, x) with w[ne0=in, ne1=out] and x[ne0=in, ne1=tokens] gives
16// y[ne0=out, ne1=tokens]. ne[0] is the contiguous dim, so
17//     y[t][o] = sum_i x[t*in + i] * w[o*in + i]
18//
19// TWO REFERENCE TIERS (measured 2026-08-06, not a stylistic choice):
20// The oracle's mul_mat on a Q8_0 weight does NOT compute dequant(W) @ x -- it quantizes the
21// ACTIVATIONS too and runs a blocked integer dot. An exact f64 reference deviates from the
22// oracle by ~0.52% median / 43% max on the Z-Image qkv projection. So grade against
23//   - `<y>.exact` (f64 reference)  -> "is our arithmetic right", tight band
24//   - `<y>` (the oracle itself)    -> only at a loose band; it is the LOWER-precision side
25// * A TOLERANCE IS A CLAIM ABOUT THE REFERENCE'S PRECISION, NOT THE IMPLEMENTATION'S.
26//
27// ACCUM MODES -- every one stays A/B-able, because a speed or accuracy claim that cannot be
28// re-measured against its predecessor is an opinion:
29//   0 naive sequential f32
30//   1 Kahan compensated
31//   2 __f32x8_dot          (8-wide, but horizontally reduces EVERY 8 elements)
32//   3 __f32x8_dot + Kahan on the chunk sum
33//   4 __f32x8_fma          (fused 8-wide multiply-add into a 32-byte accumulator, hsum ONCE)
34//   5 __f32x8_fma + multicore (forked workers over an out_dim band, shared result buffer)
35//   6 __f32x8_fma x4 independent accumulators (breaks the FMA latency chain)
36//   7 mode 6 + multicore
37//   8 Q8_0 dequant-dot (int8 weight x f32 activation via __f32_i8dot32a) -- reads the RAW
38//     quantized weight, 3.77x less memory traffic than the f32 weight
39//   9 mode 8 + multicore
40//  10 Q8_0 with per-block f16 scales pre-decoded once (software f16 decode out of the hot loop)
41//  11 mode 10 + multicore
42//  12 mode 10 with OUTPUT-OUTER loop order: each weight row loads once and serves every token
43//  13 mode 12 + multicore
44// Modes 2/4 and the fork pattern are all PRE-EXISTING estate primitives (nx_f32_linear_simd,
45// nx_fma_matmul_gate, nx_f32x8_mt_matmul). * CHECK THE ESTATE BEFORE WRITING A KERNEL.
46//
47// license_tier: ORIGINAL
48
49import "nx_syscalls.nx"
50import "nx_le.nx"
51import "nx_f32.nx"
52import "nx_f32_div.nx"
53import "nx_f32_cvt.nx"
54import "nx_f16.nx"
55import "nx_strconv.nx"
56import "nx_genfix.nx"
57import "nx_genver.nx"
58
59func zl_strlen(s: *u8) -> i64 {
  var n: i64 = 0
  while s[n] != (0 as u8) { n = n + 1 }
  return n
63}
64
65// Compute res[t*out_dim + o] for o in [o0, o1) across all `rows` tokens.
66//
67// Split out from the tally so every mode -- including the forked one, whose results have to
68// cross a process boundary -- produces its answer the same way and is scored by the same code.
69// Fusing the compare into the inner loop would have made the multicore mode a DIFFERENT
70// measurement rather than the same measurement done faster.
71func zl_band(mode: i64, x: *u8, w: *u8, res: *i64, scales: *i64,
           rows: i64, in_dim: i64, out_dim: i64, o0: i64, o1: i64) -> i64 {
  let fma_acc: *u8 = sys_mmap(256)         // 4 x 32B accumulators; forked child gets its own COW copy
  let az: *i64 = fma_acc as *i64
  // MODE 12: output-outer, token-inner.
  //
  // Every other mode here loops token-outer, so for each token it streams the ENTIRE 47MB
  // weight and nothing survives in cache to the next token: 64 tokens = 64 full passes over
  // the weight. Swapping the loops makes each 4080-byte weight row load ONCE and serve all
  // `rows` tokens from L1 -- the weight traffic drops by a factor of `rows`, and the activation
  // rows (15KB each) are what stay resident instead.
  // This is the classic blocking transform; nx_blocked_matmul already applies it, and the
  // estate's own Q8_0 notes named MEMORY ACCESS -- not the dot kernel -- as the real gap.
  if mode == 12 {
      let nb3: i64 = in_dim / 32
      var o3: i64 = o0
      while o3 < o1 {
          let qb3: i64 = (w as i64) + o3 * nb3 * 34
          let sb3: i64 = o3 * nb3
          var t3: i64 = 0
          while t3 < rows {
              let ab3: i64 = (x as i64) + t3 * in_dim * 4
              var a3: i64 = 0
              var b3: i64 = 0
              while b3 < nb3 {
                  let r3: i64 = __f32_i8dot32a((qb3 + b3 * 34 + 2) as *u8, (ab3 + b3 * 128) as *u8)
                  a3 = __f32_add(a3, __f32_mul(scales[sb3 + b3], r3))
                  b3 = b3 + 1
              }
              res[t3 * out_dim + o3] = a3
              t3 = t3 + 1
          }
          o3 = o3 + 1
      }
      return 0
  }
107
  var t: i64 = 0
  while t < rows {
      let xb: i64 = t * in_dim * 4
      var o: i64 = o0
      while o < o1 {
          let wb: i64 = o * in_dim * 4
          var acc: i64 = 0
          var comp: i64 = 0
          var i: i64 = 0
117
          if mode == 4 {
              az[0] = 0
              az[1] = 0
              az[2] = 0
              az[3] = 0
              let ib: i64 = (x as i64) + xb
              let wp: i64 = (w as i64) + wb
              let nch: i64 = in_dim / 8
              var ch: i64 = 0
              while ch < nch {
                  __f32x8_fma(fma_acc, (ib + ch * 32) as *u8, (wp + ch * 32) as *u8)
                  ch = ch + 1
              }
              acc = __f32x8_hsum(fma_acc)
              i = nch * 8
          }
          if mode == 6 {
              // FOUR independent accumulators. A single accumulator serialises on the FMA's
              // own latency: every vfmadd231ps must wait for the previous one to retire before
              // it can add into the same register, so the pipeline runs at 1/latency instead of
              // 1/throughput. Four chains let four FMAs be in flight at once. It also improves
              // accuracy for free -- 32 partial sums instead of 8.
              // Measured 1-core FMA sat at ~8% of the 8-wide FMA peak, which is the signature
              // of a latency-bound chain rather than a bandwidth-bound one.
              var q: i64 = 0
              while q < 16 { az[q] = 0; q = q + 1 }        // zero all four 32-byte accumulators
              let a0: *u8 = fma_acc
              let a1: *u8 = ((fma_acc as i64) + 32) as *u8
              let a2: *u8 = ((fma_acc as i64) + 64) as *u8
              let a3: *u8 = ((fma_acc as i64) + 96) as *u8
              let ib: i64 = (x as i64) + xb
              let wp: i64 = (w as i64) + wb
              let nch: i64 = in_dim / 8
              var ch: i64 = 0
              while ch + 4 <= nch {
                  __f32x8_fma(a0, (ib + (ch + 0) * 32) as *u8, (wp + (ch + 0) * 32) as *u8)
                  __f32x8_fma(a1, (ib + (ch + 1) * 32) as *u8, (wp + (ch + 1) * 32) as *u8)
                  __f32x8_fma(a2, (ib + (ch + 2) * 32) as *u8, (wp + (ch + 2) * 32) as *u8)
                  __f32x8_fma(a3, (ib + (ch + 3) * 32) as *u8, (wp + (ch + 3) * 32) as *u8)
                  ch = ch + 4
              }
              while ch < nch {
                  __f32x8_fma(a0, (ib + ch * 32) as *u8, (wp + ch * 32) as *u8)
                  ch = ch + 1
              }
              acc = __f32_add(__f32_add(__f32x8_hsum(a0), __f32x8_hsum(a1)),
                              __f32_add(__f32x8_hsum(a2), __f32x8_hsum(a3)))
              i = nch * 8
          }
          if mode == 10 {
              // Q8_0 with the per-block f16 scale decoded ONCE per weight, not once per
              // (token, block). nx_f16_to_f32 is SOFTWARE -- the same class of mistake as the
              // software add: it sits in the hot loop and is re-executed for every token, so at
              // rows=64 the identical 1.38M decodes were done 64 times.
              let qb2: i64 = (w as i64) + o * (in_dim / 32) * 34
              let ab2: i64 = (x as i64) + xb
              let nb2: i64 = in_dim / 32
              let sb2: i64 = o * nb2
              var b2: i64 = 0
              while b2 < nb2 {
                  let raw2: i64 = __f32_i8dot32a((qb2 + b2 * 34 + 2) as *u8, (ab2 + b2 * 128) as *u8)
                  acc = __f32_add(acc, __f32_mul(scales[sb2 + b2], raw2))
                  b2 = b2 + 1
              }
              i = nb2 * 32
          }
          if mode == 8 {
              // Q8_0 DEQUANT-DOT -- the same computation the engine's own Q8_0 path runs
              // (_lw_q8_0_dot in nx_f32_lazy_weight): per 34-byte block, f16 scale d and 32
              // int8 quants, dotted against 32 CONTIGUOUS f32 activations by __f32_i8dot32a,
              // scaled by d and accumulated.
              //
              // Note what this is NOT: ggml's q8_0 x q8_0 quantizes the ACTIVATIONS too. This
              // keeps activations in f32 and quantizes only the WEIGHT side, so it gets the
              // 3.77x weight-traffic cut WITHOUT the activation quantization error -- more
              // accurate than the engine it is replacing, at the engine's memory cost.
              let qb: i64 = (w as i64) + o * (in_dim / 32) * 34
              let ab: i64 = (x as i64) + xb
              let nblk: i64 = in_dim / 32
              var b: i64 = 0
              while b < nblk {
                  let boff: i64 = b * 34
                  let d32: i64 = nx_f16_to_f32(nx_le_read_u16(w, (o * nblk * 34) + boff))
                  let raw: i64 = __f32_i8dot32a((qb + boff + 2) as *u8, (ab + b * 32 * 4) as *u8)
                  // __f32_add/__f32_mul are the HARDWARE SSE ops; nx_f32_add/nx_f32_mul are
                  // full software IEEE-754 (dozens of integer ops each). Using the software
                  // pair here -- once per 32 values -- measured 8.4x slower and made the whole
                  // Q8_0 lever look like a dead end. The known-good _lw_q8_0_dot uses the
                  // hardware pair; deviating from it in a detail this small is what cost the
                  // measurement. * MATCH THE KNOWN GOOD EXACTLY BEFORE CONCLUDING IT LOSES.
                  acc = __f32_add(acc, __f32_mul(d32, raw))
                  b = b + 1
              }
              i = nblk * 32
          }
          if mode == 2 || mode == 3 {
              let ib: i64 = (x as i64) + xb
              let wp: i64 = (w as i64) + wb
              let nch: i64 = in_dim / 8
              var ch: i64 = 0
              while ch < nch {
                  let part: i64 = __f32x8_dot((ib + ch * 32) as *i64, (wp + ch * 32) as *i64)
                  // HARDWARE add: nx_f32_add is software IEEE-754 and runs 480x per dot here,
                  // which made mode 2 look ~9x worse than its kernel actually is. The A/B ladder
                  // has to differ ONLY in the thing being compared.
                  if mode == 2 { acc = __f32_add(acc, part) }
                  if mode == 3 {
                      let yv: i64 = nx_f32_sub(part, comp)
                      let tv: i64 = nx_f32_add(acc, yv)
                      comp = nx_f32_sub(nx_f32_sub(tv, acc), yv)
                      acc = tv
                  }
                  ch = ch + 1
              }
              i = nch * 8
          }
          // Scalar path, and the tail for every vector mode when in_dim is not a multiple of 8
          // -- a silently truncated tail would read as a clean result.
          while i < in_dim {
              let prod: i64 = nx_f32_mul(nx_le_read_u32(x, xb + i * 4),
                                         nx_le_read_u32(w, wb + i * 4))
              if mode == 1 {
                  let yv: i64 = nx_f32_sub(prod, comp)
                  let tv: i64 = nx_f32_add(acc, yv)
                  comp = nx_f32_sub(nx_f32_sub(tv, acc), yv)
                  acc = tv
              }
              if mode != 1 { acc = nx_f32_add(acc, prod) }
              i = i + 1
          }
          res[t * out_dim + o] = acc
          o = o + 1
      }
      t = t + 1
  }
  return 0
254}
255
256func main(argc: i64, argv: *i64) -> i64 {
  if argc < 5 {
      nx_genver_emit("usage_model_x_w_y_rows_passidx_accum_workers" as *u8, argc)
      return 2
  }
  let model: *u8 = argv[1] as *u8
  let xn: *u8 = argv[2] as *u8
  let wn: *u8 = argv[3] as *u8
  let yn: *u8 = argv[4] as *u8
  let lx: i64 = zl_strlen(xn)
  let lw: i64 = zl_strlen(wn)
  let ly: i64 = zl_strlen(yn)
268
  let errp: *i64 = sys_mmap(32) as *i64
  var rows: i64 = 1
  if argc >= 6 {
      errp[0] = 0
      rows = nx_strconv_parse_i64(argv[5] as *u8, errp)
      if errp[0] != 0 { nx_genver_emit("bad_rows" as *u8, 1); return 3 }
      if rows <= 0 { nx_genver_emit("bad_rows" as *u8, rows); return 4 }
  }
  var pass_idx: i64 = 0
  if argc >= 7 {
      errp[0] = 0
      pass_idx = nx_strconv_parse_i64(argv[6] as *u8, errp)
      if errp[0] != 0 { nx_genver_emit("bad_pass_idx" as *u8, 1); return 5 }
  }
  // Default 4 (__f32x8_fma): measured fastest AND 0 fails at 1e-4 on both projections.
  var accum: i64 = 6
  if argc >= 8 {
      errp[0] = 0
      accum = nx_strconv_parse_i64(argv[7] as *u8, errp)
      if errp[0] != 0 { nx_genver_emit("bad_accum_mode" as *u8, 1); return 6 }
  }
  var nworkers: i64 = 8
  if argc >= 9 {
      errp[0] = 0
      nworkers = nx_strconv_parse_i64(argv[8] as *u8, errp)
      if errp[0] != 0 { nx_genver_emit("bad_nworkers" as *u8, 1); return 7 }
      if nworkers <= 0 { nworkers = 1 }
  }
297
  let ne_x: *i64 = sys_mmap(64) as *i64
  let ne_w: *i64 = sys_mmap(64) as *i64
  let ne_y: *i64 = sys_mmap(64) as *i64
  let c_x: i64 = nx_genfix_dims(model, xn, lx, ne_x)
  let c_w: i64 = nx_genfix_dims(model, wn, lw, ne_w)
  let c_y: i64 = nx_genfix_dims(model, yn, ly, ne_y)
  if c_x < 0 { nx_genver_emit("missing_x" as *u8, 1); return 30 }
  if c_w < 0 { nx_genver_emit("missing_w" as *u8, 1); return 31 }
  if c_y < 0 { nx_genver_emit("missing_y" as *u8, 1); return 32 }
307
  let in_dim:  i64 = ne_w[0]
  let out_dim: i64 = ne_w[1]
  let n_tok:   i64 = ne_x[1]
311
  // Refuse a shape story that does not hold together, rather than produce a number from it.
  if ne_x[0] != in_dim  { nx_genver_emit("x_indim_mismatch" as *u8, ne_x[0]); return 33 }
  if ne_y[0] != out_dim { nx_genver_emit("y_outdim_mismatch" as *u8, ne_y[0]); return 34 }
  if ne_y[1] != n_tok   { nx_genver_emit("y_tokens_mismatch" as *u8, ne_y[1]); return 35 }
316
  let x: *u8 = nx_genfix_load(model, xn, lx, c_x)
  if (x as i64) == 0 { nx_genver_emit("load_failed_x" as *u8, 1); return 40 }
  // Modes 8/9 read the native quantized blocks; loading the 177MB f32 shadow too would
  // triple the footprint and measure a kernel that is not the one under test.
  var w: *u8 = 0 as *u8
  if accum >= 8 {
      let q_bytes: i64 = (in_dim / 32) * 34 * out_dim
      w = nx_genfix_load_raw(model, wn, lw, q_bytes)
      if (w as i64) == 0 { nx_genver_emit("load_failed_w_raw" as *u8, 1); return 43 }
      nx_genver_emit("q8_weight_bytes" as *u8, q_bytes)
  }
  if accum < 8 {
      w = nx_genfix_load(model, wn, lw, c_w)
      if (w as i64) == 0 { nx_genver_emit("load_failed_w" as *u8, 1); return 41 }
  }
  let y: *u8 = nx_genfix_load(model, yn, ly, c_y)
  if (y as i64) == 0 { nx_genver_emit("load_failed_y" as *u8, 1); return 42 }
334
  if n_tok < rows { rows = n_tok }
336
  nx_genver_emit("in_dim" as *u8, in_dim)
  nx_genver_emit("out_dim" as *u8, out_dim)
  nx_genver_emit("tokens_total" as *u8, n_tok)
  nx_genver_emit("tokens_checked" as *u8, rows)
  nx_genver_emit("multiplies" as *u8, rows * out_dim * in_dim)
  nx_genver_emit("accum_mode" as *u8, accum)
343
  // SHARED so forked workers' results are visible to the parent that scores them.
  let res: *i64 = sys_mmap_shared(rows * out_dim * 8) as *i64
346
  // Pre-decode every block scale once. SHARED so forked workers read it instead of each
  // redoing the software f16 decode.
  let nblk_row: i64 = in_dim / 32
  var scales: *i64 = 0 as *i64
  if accum >= 10 {
      scales = sys_mmap_shared(out_dim * nblk_row * 8) as *i64
      var so: i64 = 0
      while so < out_dim {
          var sb: i64 = 0
          while sb < nblk_row {
              scales[so * nblk_row + sb] = nx_f16_to_f32(nx_le_read_u16(w, (so * nblk_row + sb) * 34))
              sb = sb + 1
          }
          so = so + 1
      }
  }
363
  let t_start: i64 = sys_now_us()
  if accum == 5 || accum == 7 || accum == 9 || accum == 11 || accum == 13 {
      nx_genver_emit("workers" as *u8, nworkers)
      let pids: *i64 = sys_mmap(nworkers * 8 + 64) as *i64
      var k: i64 = 0
      while k < nworkers {
          let o0: i64 = k * out_dim / nworkers
          let o1: i64 = (k + 1) * out_dim / nworkers
          let pid: i64 = sys_fork()
          if pid == 0 {
              var inner: i64 = 4
              if accum == 7 { inner = 6 }
              if accum == 9 { inner = 8 }
              if accum == 11 { inner = 10 }
              if accum == 13 { inner = 12 }
              zl_band(inner, x, w, res, scales, rows, in_dim, out_dim, o0, o1)
              sys_exit(0)
          }
          pids[k] = pid
          k = k + 1
      }
      let st: *i64 = sys_mmap(64) as *i64
      k = 0
      while k < nworkers {
          sys_wait4(pids[k], st, 0)
          k = k + 1
      }
  }
  if accum != 5 && accum != 7 && accum != 9 && accum != 11 && accum != 13 {
      zl_band(accum, x, w, res, scales, rows, in_dim, out_dim, 0, out_dim)
  }
  let t_end: i64 = sys_now_us()
  nx_genver_emit("compute_us" as *u8, t_end - t_start)
397
  let tol: *i64 = sys_mmap(64) as *i64
  nx_genver_tols(tol)
  let c: *i64 = sys_mmap(128) as *i64
  nx_genver_init(c, pass_idx)
402
  var t2: i64 = 0
  while t2 < rows {
      var o2: i64 = 0
      while o2 < out_dim {
          let flat: i64 = t2 * out_dim + o2
          nx_genver_tally(c, tol, res[flat], nx_le_read_u32(y, flat * 4), flat)
          o2 = o2 + 1
      }
      t2 = t2 + 1
  }
413
  return nx_genver_report(c)
415}
nx_gen_linear_verify.nx source

`nx_gen_linear_verify.nx` source
