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

↩ module page · 254 lines · 9977 B
1// nx_q5k_dot_simd.nx -- SIMD-accelerated Q5_K fused dot (vpmaddwd), bit-exact vs a scalar Q5_K dot.
2//
3// Extends the perf-exceed to the Q5_K layer matmuls (attn_v/output/ffn_down). Q5_K = Q4_K + a 5th bit from
4// qh: q5 = q4 + (qh_bit ? 16 : 0). Reformulation is the same: dot += d1*sq - m0*sc, sq = Σ q5*col via
5// vpmaddwd (q5 i16 0..31, col i16 Q10). Read-once + sc-precompute like nx_q4k_dot_simd. Includes an inline
6// scalar Q5_K dot as the bit-exact oracle. Verified on real blk.0.attn_v + measured.
7// license_tier: ORIGINAL
8import "nx_syscalls.nx"
9import "nx_tier.nx"
10import "nx_le.nx"
11import "nx_strconv.nx"
12import "nx_tensor.nx"
13import "nx_gguf.nx"
14import "nx_gguf_load.nx"
15import "nx_gguf_meta.nx"
16import "nx_placement.nx"
17import "nx_gguf_load_lazy.nx"
18import "nx_clock.nx"
19
20func q5s_pack4(a: i64, b: i64, c: i64, d: i64) -> i64 {
  return (a & 0xFFFF) | ((b & 0xFFFF) << 16) | ((c & 0xFFFF) << 32) | ((d & 0xFFFF) << 48)
22}
23
24func q5s_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
37}
38
39// unpack (sc,m) for sub-block j from the 12-byte packed scales at scales_off
40func q5s_scmin(buf: *u8, scales_off: i64, j: i64, out: *i64) -> i64 {
  if j < 4 {
      out[0] = nx_le_read_u8(buf, scales_off + j) & 0x3F
      out[1] = nx_le_read_u8(buf, scales_off + j + 4) & 0x3F
  } else {
      let k: i64 = j - 4
      out[0] = ((nx_le_read_u8(buf, scales_off + k) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k) & 0x0F)
      out[1] = ((nx_le_read_u8(buf, scales_off + 4 + k) >> 6) << 4) | (nx_le_read_u8(buf, scales_off + 8 + k) >> 4)
  }
  return 0
50}
51
52// scalar Q5_K dot oracle: Σ (d*sc*q5 - dmin*m) * col, in Q34.
53func q5k_dot_scalar(buf: *u8, base_off: i64, n_blocks: i64, col_q10: *i64, sm: *i64) -> i64 {
  var dot: i64 = 0
  var blk: i64 = 0
  while blk < n_blocks {
      let sb: i64 = base_off + blk * 176
      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 qh_off: i64 = sb + 16
      let qs_off: i64 = sb + 48
      let col_base: i64 = blk * 256
      var g: i64 = 0
      while g < 4 {
          let is0: i64 = g + g
          let is1: i64 = is0 + 1
          q5s_scmin(buf, scales_off, is0, sm)
          let d1: i64 = d_q24 * sm[0]
          let m0v: i64 = dmin_q24 * sm[1]
          q5s_scmin(buf, scales_off, is1, sm)
          let d2: i64 = d_q24 * sm[0]
          let m1v: i64 = dmin_q24 * sm[1]
          let grp: i64 = qs_off + g * 32
          let u1: i64 = 1 << (g + g)
          let u2: i64 = 1 << (g + g + 1)
          var l: i64 = 0
          while l < 32 {
              let byte_v: i64 = nx_le_read_u8(buf, grp + l)
              let qh_l: i64 = nx_le_read_u8(buf, qh_off + l)
              var q_lo: i64 = byte_v & 0x0F
              var q_hi: i64 = byte_v >> 4
              if (qh_l & u1) != 0 { q_lo = q_lo + 16 }
              if (qh_l & u2) != 0 { q_hi = q_hi + 16 }
              dot = dot + (d1 * q_lo - m0v) * col_q10[col_base + is0 * 32 + l]
              dot = dot + (d2 * q_hi - m1v) * col_q10[col_base + is1 * 32 + l]
              l = l + 1
          }
          g = g + 1
      }
      blk = blk + 1
  }
  return dot
94}
95
96// SIMD Q5_K dot: sq = Σ q5*col via vpmaddwd; sc from precompute.
97func q5k_dot_simd(buf: *u8, base_off: i64, n_blocks: i64, col_i16: *i64,
                qpk: *i64, qhi: *i64, acc: *i64, sc_pre: *i64, sm: *i64) -> i64 {
  var dot: i64 = 0
  var blk: i64 = 0
  while blk < n_blocks {
      let sb: i64 = base_off + blk * 176
      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 qh_off: i64 = sb + 16
      let qs_off: i64 = sb + 48
      let col_base: i64 = blk * 256
      var g: i64 = 0
      while g < 4 {
          let is0: i64 = g + g
          let is1: i64 = is0 + 1
          q5s_scmin(buf, scales_off, is0, sm)
          let d1: i64 = d_q24 * sm[0]
          let m0v: i64 = dmin_q24 * sm[1]
          q5s_scmin(buf, scales_off, is1, sm)
          let d2: i64 = d_q24 * sm[0]
          let m1v: i64 = dmin_q24 * sm[1]
          let grp: i64 = qs_off + g * 32
          let u1: i64 = 1 << (g + g)
          let u2: i64 = 1 << (g + g + 1)
122
          var jl: i64 = 0
          while jl < 32 {
              let byte_v: i64 = nx_le_read_u8(buf, grp + jl)
              let qh_l: i64 = nx_le_read_u8(buf, qh_off + jl)
              var q_lo: i64 = byte_v & 0x0F
              var q_hi: i64 = byte_v >> 4
              if (qh_l & u1) != 0 { q_lo = q_lo + 16 }
              if (qh_l & u2) != 0 { q_hi = q_hi + 16 }
              let w: i64 = jl / 4
              let sh: i64 = (jl - w * 4) * 16
              if sh == 0 { qpk[w] = q_lo; qhi[w] = q_hi }
              else { qpk[w] = qpk[w] | (q_lo << sh); qhi[w] = qhi[w] | (q_hi << sh) }
              jl = jl + 1
          }
137
          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 = q5s_hsum(acc)
          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 = q5s_hsum(acc)
148
          dot = dot + d1 * sq_lo - m0v * sc_pre[blk * 8 + is0] + d2 * sq_hi - m1v * sc_pre[blk * 8 + is1]
          g = g + 1
      }
      blk = blk + 1
  }
  return dot
155}
156
157func q5s_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)
169}
170
171func main() -> i64 {
  let path: *u8 = "/mnt/c/Users/elder/elder-ai-platform/models/unified/text_encoder/Huihui-Qwen3-4B-Instruct-2507-abliterated-Q4_K_M.gguf" as *u8
  let fd: i64 = sys_openat_rd(path)
  if fd < 0 { return 30 }
  let CAP: i64 = 1153433600
  let buf: *u8 = sys_mmap(CAP)
  var total: i64 = 0
  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 { return 31 }
  let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
  if nx_gguf_parse(buf, total, hdr) != NX_GGUF_OK { return 40 }
  let vi: nx_int = nx_gguf_find_tensor(hdr, "blk.0.attn_v.weight" as *u8, 19)
  if vi < 0 { return 60 }
  let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, vi)
  if ti.ggml_type != 13 { return 61 }
  let IN: i64 = ti.dim_0
  let w_off: i64 = hdr.data_off + ti.offset
  let n_blocks: i64 = IN / 256
  if w_off + n_blocks * 176 > total { return 63 }
195
  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] = q5s_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
  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[base + l]; l = l + 1 }
          sc_pre[blk * 8 + is_] = s
          is_ = is_ + 1
      }
      blk = blk + 1
  }
219
  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 sm: *i64 = sys_mmap(16) as *i64
224
  let scalar_dot: i64 = q5k_dot_scalar(buf, w_off, n_blocks, col, sm)
  let simd_dot: i64 = q5k_dot_simd(buf, w_off, n_blocks, col_i16, qpk, qhi, acc, sc_pre, sm)
227
  let IT: i64 = 20000
  let t0: i64 = nx_clock_monotonic_ns()
  var s1: i64 = 0
  var k: i64 = 0
  while k < IT { s1 = s1 + q5k_dot_scalar(buf, w_off, n_blocks, col, sm); k = k + 1 }
  let t1: i64 = nx_clock_monotonic_ns()
  let scalar_ns: i64 = t1 - t0
  let t2: i64 = nx_clock_monotonic_ns()
  var s2: i64 = 0
  k = 0
  while k < IT { s2 = s2 + q5k_dot_simd(buf, w_off, n_blocks, col_i16, qpk, qhi, acc, sc_pre, sm); k = k + 1 }
  let t3: i64 = nx_clock_monotonic_ns()
  let simd_ns: i64 = t3 - t2
241
  let ofd: i64 = sys_openat_wr("/tmp/q5k_simd.txt" as *u8, 0x1a4)
  if ofd >= 0 {
      q5s_emit(ofd, "scalar_dot" as *u8, 10, scalar_dot)
      q5s_emit(ofd, "simd_dot" as *u8, 8, simd_dot)
      q5s_emit(ofd, "scalar_ns" as *u8, 9, scalar_ns)
      q5s_emit(ofd, "simd_ns" as *u8, 7, simd_ns)
      if simd_ns > 0 { q5s_emit(ofd, "speedup_x100" as *u8, 12, scalar_ns * 100 / simd_ns) }
      sys_close(ofd)
  }
  if simd_dot != scalar_dot { return 80 }
  if simd_ns >= scalar_ns { return 81 }
  return 0
254}
nx_q5k_dot_simd.nx source

`nx_q5k_dot_simd.nx` source
