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

↩ module page · 275 lines · 11119 B
1// nx_q4k_dot_simd.nx -- SIMD-accelerated Q4_K fused dot (vpmaddwd), bit-exact vs nx_q4k_dot_row_col.
2//
3// Perf-exceed applied to the Qwen GEMM. Reformulation (exact integer algebra): per sub-block,
4//   dot += Σ (d1*q - m0)*col  =  d1*(Σ q*col) - m0*(Σ col)  =  d1*sq - m0*sc
5// sq = Σ q*col via one vpmaddwd accumulate (q i16 0..15, col i16 Q10). Optimizations vs v1: read qs bytes
6// ONCE (extract lo+hi together), and take sc = Σ col from a PRECOMPUTED sub-block table (activation-only,
7// constant across the dot -- free in a GEMM). col: both Q10 i64 (unused now except pack) + i16-packed.
8// license_tier: ORIGINAL
9import "nx_syscalls.nx"
10import "nx_tier.nx"
11import "nx_le.nx"
12import "nx_strconv.nx"
13import "nx_tensor.nx"
14import "nx_gguf.nx"
15import "nx_gguf_load.nx"
16import "nx_gguf_meta.nx"
17import "nx_placement.nx"
18import "nx_gguf_load_lazy.nx"
19import "nx_q4k_matmul.nx"
20import "nx_dequant_iter.nx"
21import "nx_clock.nx"
22
23func ds_pack4(a: i64, b: i64, c: i64, d: i64) -> i64 {
  return (a & 0xFFFF) | ((b & 0xFFFF) << 16) | ((c & 0xFFFF) << 32) | ((d & 0xFFFF) << 48)
25}
26
27func ds_hsum(acc: *i64) -> i64 {
  var sum: i64 = 0
  var i: i64 = 0
  while i < 4 {
      let v: i64 = acc[i]
      var lo: i64 = v & 0xFFFFFFFF
      if lo >= 0x80000000 { lo = lo - 0x100000000 }
      var hi: i64 = (v >> 32) & 0xFFFFFFFF
      if hi >= 0x80000000 { hi = hi - 0x100000000 }
      sum = sum + lo + hi
      i = i + 1
  }
  return sum
40}
41
42// sc_pre[blk*8 + is] = Σ_{l<32} col[blk*256 + is*32 + l]  (activation-only, precompute once)
43func nx_q4k_sc_precompute(col_q10: *i64, n_blocks: i64, sc_pre: *i64) -> i64 {
  var blk: i64 = 0
  while blk < n_blocks {
      var is_: i64 = 0
      while is_ < 8 {
          var s: i64 = 0
          var l: i64 = 0
          let base: i64 = blk * 256 + is_ * 32
          while l < 32 { s = s + col_q10[base + l]; l = l + 1 }
          sc_pre[blk * 8 + is_] = s
          is_ = is_ + 1
      }
      blk = blk + 1
  }
  return 0
58}
59
60func nx_q4k_dot_simd(buf: *u8, base_off: i64, n_blocks: i64, col_i16: *i64,
                   qpk: *i64, qhi: *i64, acc: *i64, sc_pre: *i64) -> i64 {
  var dot: i64 = 0
  var blk: i64 = 0
  while blk < n_blocks {
      let sb: i64 = base_off + blk * 144
      let d_q24:    i64 = _gguf_f16_to_q24(nx_le_read_u16(buf, sb))
      let dmin_q24: i64 = _gguf_f16_to_q24(nx_le_read_u16(buf, sb + 2))
      let scales_off: i64 = sb + 4
      let qs_off:     i64 = sb + 16
      let col_base:   i64 = blk * 256
      var g: i64 = 0
      while g < 4 {
          let is0: i64 = g + g
          let is1: i64 = is0 + 1
          var sc0: i64 = 0
          var m0: i64 = 0
          var sc1: i64 = 0
          var m1s: i64 = 0
          if is0 < 4 {
              sc0 = nx_le_read_u8(buf, scales_off + is0) & 0x3F
              m0  = nx_le_read_u8(buf, scales_off + is0 + 4) & 0x3F
          } else {
              let k0: i64 = is0 - 4
              sc0 = ((nx_le_read_u8(buf, scales_off + k0) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k0) & 0x0F)
              m0  = ((nx_le_read_u8(buf, scales_off + 4 + k0) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k0) >> 4)
          }
          if is1 < 4 {
              sc1 = nx_le_read_u8(buf, scales_off + is1) & 0x3F
              m1s = nx_le_read_u8(buf, scales_off + is1 + 4) & 0x3F
          } else {
              let k1: i64 = is1 - 4
              sc1 = ((nx_le_read_u8(buf, scales_off + k1) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k1) & 0x0F)
              m1s = ((nx_le_read_u8(buf, scales_off + 4 + k1) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k1) >> 4)
          }
          let d1: i64 = d_q24 * sc0
          let m0v: i64 = dmin_q24 * m0
          let d2: i64 = d_q24 * sc1
          let m1v: i64 = dmin_q24 * m1s
          let grp: i64 = qs_off + g * 32
100
          // read 32 qs bytes ONCE; fill qpk (low nibbles) + qhi (high nibbles)
          var j: i64 = 0
          while j < 8 {
              let b0: i64 = nx_le_read_u8(buf, grp + j * 4 + 0)
              let b1: i64 = nx_le_read_u8(buf, grp + j * 4 + 1)
              let b2: i64 = nx_le_read_u8(buf, grp + j * 4 + 2)
              let b3: i64 = nx_le_read_u8(buf, grp + j * 4 + 3)
              qpk[j] = ds_pack4(b0 & 0x0F, b1 & 0x0F, b2 & 0x0F, b3 & 0x0F)
              qhi[j] = ds_pack4(b0 >> 4, b1 >> 4, b2 >> 4, b3 >> 4)
              j = j + 1
          }
112
          acc[0] = 0; acc[1] = 0; acc[2] = 0; acc[3] = 0
          let cl0: i64 = (col_i16 as i64) + ((col_base + is0 * 32) / 4) * 8
          __i16x16_madd(acc as *i64, qpk as *i64, cl0 as *i64)
          __i16x16_madd(acc as *i64, ((qpk as i64) + 32) as *i64, (cl0 + 32) as *i64)
          let sq_lo: i64 = ds_hsum(acc)
118
          acc[0] = 0; acc[1] = 0; acc[2] = 0; acc[3] = 0
          let cl1: i64 = (col_i16 as i64) + ((col_base + is1 * 32) / 4) * 8
          __i16x16_madd(acc as *i64, qhi as *i64, cl1 as *i64)
          __i16x16_madd(acc as *i64, ((qhi as i64) + 32) as *i64, (cl1 + 32) as *i64)
          let sq_hi: i64 = ds_hsum(acc)
124
          let sc_lo: i64 = sc_pre[blk * 8 + is0]
          let sc_hi: i64 = sc_pre[blk * 8 + is1]
          dot = dot + d1 * sq_lo - m0v * sc_lo + d2 * sq_hi - m1v * sc_hi
          g = g + 1
      }
      blk = blk + 1
  }
  return dot
133}
134
135func ds_emit(fd: i64, key: *u8, kl: i64, v: i64) -> i64 {
  let line: *u8 = sys_mmap(64)
  var lo: i64 = 0
  var i: i64 = 0
  while i < kl { line[lo] = key[i]; lo = lo + 1; i = i + 1 }
  line[lo] = 0x3D; lo = lo + 1
  let dec: *u8 = sys_mmap(32)
  let nd: i64 = nx_strconv_format_i64(v, dec)
  var k: i64 = 0
  while k < nd { line[lo] = dec[k]; lo = lo + 1; k = k + 1 }
  line[lo] = 0x0A; lo = lo + 1
  return sys_write(fd, line, lo)
147}
148
149// Synthesise `nb` Q4_K blocks in place. A Q4_K block is 144 bytes: 2B d (fp16) + 2B dmin (fp16) +
150// 12B packed 6-bit scales + 128B of packed 4-bit quants. The PERF question -- is the SIMD dot faster
151// than the scalar one -- does not need real weights at all: both paths read the SAME bytes, so any
152// deterministic fill answers it exactly. This removes a 1.1GB model dependency from a timing bench.
153func ds_synth(buf: *u8, nb: i64) -> i64 {
  var b: i64 = 0
  while b < nb {
      let sb: i64 = b * 144
      buf[sb + 0] = 0 as u8
      buf[sb + 1] = 60 as u8
      buf[sb + 2] = 0 as u8
      buf[sb + 3] = 52 as u8
      var i: i64 = 0
      while i < 12 { buf[sb + 4 + i] = ((b * 7 + i * 13) % 64) as u8; i = i + 1 }
      i = 0
      while i < 128 { buf[sb + 16 + i] = ((b * 31 + i * 17) % 256) as u8; i = i + 1 }
      b = b + 1
  }
  return 0
168}
169
170// argv[1] = OPTIONAL gguf path (rule 17: never bake one). Absent or unopenable -> SYNTHETIC blocks, and
171// the run reports synthetic=1 so no reader can mistake a fixture number for a real-weights number.
172// The old build baked a /mnt/c laptop path and returned 30 on the NAS, writing zero bytes: a silent
173// failure that read as a hang. Results now also go to STDOUT, because nx_job_run captures stdout and a
174// file-only emitter looks empty even when it succeeded.
175func main(argc: i64, argv: *i64) -> i64 {
  var buf: *u8 = 0 as *u8
  var total: i64 = 0
  var IN: i64 = 0
  var w_off: i64 = 0
  var synth: i64 = 1
  if argc >= 2 {
      let path: *u8 = argv[1] as *u8
      let fd: i64 = sys_openat_rd(path)
      if fd >= 0 {
          let CAP: i64 = 1153433600
          buf = sys_mmap(CAP)
          var go: i64 = 1
          while go == 1 {
              let r: i64 = sys_read(fd, ((buf as i64) + total) as *u8, CAP - total)
              if r <= 0 { go = 0 } else { total = total + r; if total >= CAP { go = 0 } }
          }
          sys_close(fd)
          if total >= 100000000 {
              let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
              if nx_gguf_parse(buf, total, hdr) == NX_GGUF_OK {
                  let qi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_q.weight" as *u8, 19)
                  if qi >= 0 {
                      let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, qi)
                      IN = ti.dim_0
                      w_off = hdr.data_off + ti.offset
                      synth = 0
                  }
              }
          }
      }
  }
  var n_blocks: i64 = 0
  if synth == 1 {
      n_blocks = 14
      IN = n_blocks * 256
      total = n_blocks * 144
      buf = sys_mmap(total)
      ds_synth(buf, n_blocks)
      w_off = 0
  } else {
      n_blocks = IN / 256
      if w_off + n_blocks * 144 > total { return 63 }
  }
219
  let col: *i64 = sys_mmap(IN * 8) as *i64
  let col_i16: *i64 = sys_mmap(IN * 2) as *i64
  var i: i64 = 0
  while i < IN { col[i] = 1024 + (i - (i / 5) * 5) * 256; i = i + 1 }
  var jj: i64 = 0
  while jj < IN / 4 {
      col_i16[jj] = ds_pack4(col[jj * 4], col[jj * 4 + 1], col[jj * 4 + 2], col[jj * 4 + 3])
      jj = jj + 1
  }
  let sc_pre: *i64 = sys_mmap(n_blocks * 8 * 8) as *i64
  nx_q4k_sc_precompute(col, n_blocks, sc_pre)
231
  let qpk: *i64 = sys_mmap(64) as *i64
  let qhi: *i64 = sys_mmap(64) as *i64
  let acc: *i64 = sys_mmap(32) as *i64
  let it: *NxQ4KBlockIter = nx_q4k_iter_alloc()
236
  let scalar_dot: i64 = nx_q4k_dot_row_col(buf, w_off, n_blocks, col, it)
  let simd_dot: i64 = nx_q4k_dot_simd(buf, w_off, n_blocks, col_i16, qpk, qhi, acc, sc_pre)
239
  let IT: i64 = 20000
  let t0: i64 = nx_clock_monotonic_ns()
  var s1: i64 = 0
  var k: i64 = 0
  while k < IT { s1 = s1 + nx_q4k_dot_row_col(buf, w_off, n_blocks, col, it); k = k + 1 }
  let t1: i64 = nx_clock_monotonic_ns()
  let scalar_ns: i64 = t1 - t0
247
  let t2: i64 = nx_clock_monotonic_ns()
  var s2: i64 = 0
  k = 0
  while k < IT { s2 = s2 + nx_q4k_dot_simd(buf, w_off, n_blocks, col_i16, qpk, qhi, acc, sc_pre); k = k + 1 }
  let t3: i64 = nx_clock_monotonic_ns()
  let simd_ns: i64 = t3 - t2
254
  let ofd: i64 = sys_openat_wr("/tmp/q4k_simd.txt" as *u8, 0x1a4)
  if ofd >= 0 {
      ds_emit(ofd, "scalar_dot" as *u8, 10, scalar_dot)
      ds_emit(ofd, "simd_dot" as *u8, 8, simd_dot)
      ds_emit(ofd, "scalar_ns" as *u8, 9, scalar_ns)
      ds_emit(ofd, "simd_ns" as *u8, 7, simd_ns)
      if simd_ns > 0 { ds_emit(ofd, "speedup_x100" as *u8, 12, scalar_ns * 100 / simd_ns) }
      sys_close(ofd)
  }
  ds_emit(1, "synthetic" as *u8, 9, synth)
  ds_emit(1, "n_blocks" as *u8, 8, n_blocks)
  ds_emit(1, "scalar_dot" as *u8, 10, scalar_dot)
  ds_emit(1, "simd_dot" as *u8, 8, simd_dot)
  ds_emit(1, "bit_exact" as *u8, 9, 1 - ((scalar_dot - simd_dot) * (scalar_dot - simd_dot)))
  ds_emit(1, "scalar_ns" as *u8, 9, scalar_ns)
  ds_emit(1, "simd_ns" as *u8, 7, simd_ns)
  if simd_ns > 0 { ds_emit(1, "speedup_x100" as *u8, 12, scalar_ns * 100 / simd_ns) }
  if simd_dot != scalar_dot { return 80 }
  if simd_ns >= scalar_ns { return 81 }
  return 0
275}
nx_q4k_dot_simd.nx source

`nx_q4k_dot_simd.nx` source
