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

↩ module page · 323 lines · 14871 B
1// nx_f32_llama_layer_lazy_load.nx -- bind one layer's 9 tensors into
2// NxF32LlamaLayerLazy.
3//
4// For each tensor:
5//   * F32 source -> dequant once into f32 storage, wrap as
6//     NxF32LazyWeight tagged F32.
7//   * F16 source -> widen to f32, store, tag F32.
8//   * Q4_K source -> DO NOT DEQUANT.  Record the GGUF buffer + byte
9//     offset directly, tag Q4_K.  The forward path (v4 block) will
10//     stream Q4_K super-blocks during matmul.
11//
12// gamma_attn / gamma_ffn (RMSNorm scales) are always loaded as f32
13// raw bits (they're 1D vectors, small).
14//
15// genealogy_id: ggml_format_taxonomy + standard_layer_binding
16// lineage_id:   substrate_f32_llama_layer_lazy_load_v1
17
18import "nx_syscalls.nx"
19import "nx_tier.nx"
20import "nx_dec_emit.nx"
21import "nx_le.nx"
22import "nx_gguf.nx"
23import "nx_gguf_load.nx"
24import "nx_gguf_load_f32.nx"
25import "nx_f32_cvt.nx"
26import "nx_f32_lazy_weight.nx"
27import "nx_f32_llama_block_v4.nx"
28import "nx_q5_k_to_f32.nx"
29import "nx_q6_k_to_f32.nx"
30import "nx_q4k_to_f32.nx"
31import "nx_q8_0_to_f32.nx"
32import "nx_q8_0_from_q5_0.nx"
33import "nx_q8_0_from_f32.nx"
34
35const NX_FLLL_OK:              nx_int = 0
36const NX_FLLL_ERR_BAD_LAYER:   nx_int = 1
37const NX_FLLL_ERR_NOT_FOUND:   nx_int = 2
38const NX_FLLL_ERR_NULL:        nx_int = 3
39const NX_FLLL_ERR_BAD_TYPE:    nx_int = 4
40const NX_FLLL_N_VERDICTS:      nx_int = 5
41
42func nx_flll_verdict_is_valid(v: nx_int) -> nx_int {
  if v < 0 { return 0 }
  if v >= NX_FLLL_N_VERDICTS { return 0 }
  return 1
46}
47
48func _flll_fmt_name(layer_idx: nx_int, suffix: *u8, suffix_len: nx_int,
                  name_out: *u8) -> nx_int {
  name_out[0] = 0x62 as u8
  name_out[1] = 0x6c as u8
  name_out[2] = 0x6b as u8
  name_out[3] = 0x2e as u8
  let li_len: i64 = nx_dec_emit_u63(name_out, 4, layer_idx as i64)
  let after_idx: nx_int = 4 + (li_len as nx_int)
  name_out[after_idx] = 0x2e as u8
  var i: nx_int = 0
  while i < suffix_len {
      name_out[after_idx + 1 + i] = suffix[i]
      i = i + 1
  }
  return after_idx + 1 + suffix_len
63}
64
65// Build a NxF32LazyWeight for a named tensor.  F32/F16 -> eager f32;
66// Q4_K -> record offset only.  Returns 0 on not-found.
67
68func _flll_build_lazy(buf: *u8, hdr: *NxGgufHeader,
                    name: *u8, name_len: nx_int) -> *NxF32LazyWeight {
  let idx: nx_int = nx_gguf_find_tensor(hdr, name, name_len)
  if idx < 0 { return 0 as *NxF32LazyWeight }
  let ti: *NxGgufTensorInfo = nx_gguf_tensor_at(hdr, idx)
73
  // Dims: tensor info stores up to 4 dims; we want rows/cols.
  let rows: nx_int = ti.dim_0 as nx_int
  var cols: nx_int = 1
  if ti.n_dims >= 2 { cols = ti.dim_1 as nx_int }
78
  let data_off: i64 = hdr.data_off + ti.offset
80
  if ti.ggml_type == NX_GGML_TYPE_Q4_K {
      // LAZY: keep the Q4_K bytes + offset; the fixed streaming nx_f32_q4k_matmul (fq4m_rows) dequants one
      // output row at a time so the weight stays quantized in memory (a 7B = ~4.5GB, not ~28GB f32). The
      // earlier eager-dequant here was a memory-blowup shortcut; the streaming kernel is now correct.
      return nx_f32_lazy_weight_new_q4k(buf, data_off, rows, cols)
  }
  if ti.ggml_type == NX_GGML_TYPE_Q8_0 {
      // LAZY (2026-07-08): keep Q8_0 quantized -> the dispatcher's SIMD
      // __f32_i8dot32 dequant-dot (PROVEN 10.2x).  This model has NO Q8_0
      // BLOCK weights (they're Q5_0; Q8_0 is only embed/lm_head), so this
      // is a NO-OP here -- but a Q8_0-quantized model runs EVERY block
      // matmul through the SIMD path with zero further work. k is 32-aligned
      // for all LLM reduction dims.
      return nx_f32_lazy_weight_new_q8_0(buf, data_off, rows, cols)
  }
  if ti.ggml_type == NX_GGML_TYPE_Q5_0 {
      // ★ RE-QUANTIZE Q5_0 -> Q8_0 ONCE AT LOAD (2026-07-08).  Q5_0->Q8_0 is
      // EXACT (5-bit vals -16..15 fit int8, same per-block f16 d).  This
      // AMORTIZES the nibble+qh unpack to load-time (the very cost that made
      // a per-token Q5_0 SIMD dequant-dot 2x SLOWER) so the forward then runs
      // the TRIVIAL Q8_0 SIMD dequant-dot (__f32_i8dot32, PROVEN 10.2x) on
      // ALL 79%-Q5_0 weights.  RAM 1.5x the Q5_0 bytes (still << F32 eager).
      let nvq: i64 = nx_gguf_tensor_n_values(ti)
      let nblkq: i64 = (nvq + NX_Q8_0_VPB - 1) / NX_Q8_0_VPB
      let q8buf: *u8 = sys_mmap(nblkq * NX_Q8_0_BPB)
      nx_q8_0_from_q5_0(buf, data_off, nvq, q8buf)
      return nx_f32_lazy_weight_new_q8_0(q8buf, 0, rows, cols)
  }
  if ti.ggml_type == NX_GGML_TYPE_F32 {
      let nv: i64 = nx_gguf_tensor_n_values(ti)
      let storage: *i64 = sys_mmap(nv * 8) as *i64
      var i: i64 = 0
      while i < nv {
          storage[i] = nx_le_read_u32(buf, data_off + i * 4)
          i = i + 1
      }
      return nx_f32_lazy_weight_new_f32(storage, rows, cols)
  }
  if ti.ggml_type == NX_GGML_TYPE_F16 {
      let nv: i64 = nx_gguf_tensor_n_values(ti)
      let storage: *i64 = sys_mmap(nv * 8) as *i64
      var i2: i64 = 0
      while i2 < nv {
          let raw_f16: i64 = nx_le_read_u16(buf, data_off + i2 * 2)
          storage[i2] = nx_f16_to_f32(raw_f16)
          i2 = i2 + 1
      }
      return nx_f32_lazy_weight_new_f32(storage, rows, cols)
  }
  if ti.ggml_type == NX_GGML_TYPE_Q5_K {
      let nv5: i64 = nx_gguf_tensor_n_values(ti)
      let st5: *i64 = sys_mmap(nv5 * 8) as *i64
      nx_q5_k_to_f32(buf, data_off, nv5, st5)
      return nx_f32_lazy_weight_new_f32(st5, rows, cols)
  }
  if ti.ggml_type == NX_GGML_TYPE_Q6_K {
      // dequant Q6_K -> F32, then re-quantize -> Q8_0 (near-lossless: 8-bit >=
      // 6-bit) so this 8% runs the SIMD Q8_0 dequant-dot too -> FULLY-Q8_0
      // block matmul. Q8_0 temp replaces the F32 temp (freed).
      let nv6: i64 = nx_gguf_tensor_n_values(ti)
      let st6: *i64 = sys_mmap(nv6 * 8) as *i64
      nx_q6_k_to_f32(buf, data_off, nv6, st6)
      let nb6: i64 = (nv6 + NX_Q8_0_VPB - 1) / NX_Q8_0_VPB
      let q86: *u8 = sys_mmap(nb6 * NX_Q8_0_BPB)
      nx_q8_0_from_f32(st6, nv6, q86)
      sys_munmap(st6, nv6 * 8)
      return nx_f32_lazy_weight_new_q8_0(q86, 0, rows, cols)
  }
  // Q5_0 / Q8_0: eager dequant via the unified loader (which
  // handles all 6 supported types).  Wraps result as F32 lazy weight.
  let nv_out: *i64 = sys_mmap(8) as *i64
  let inner_err: *i64 = sys_mmap(8) as *i64
  let name_buf: *u8 = sys_mmap(64)
  var n_len: nx_int = 0
  var k: nx_int = 0
  while k < ti.name_len {
      name_buf[k] = ti.name[k]
      n_len = n_len + 1
      k = k + 1
  }
  let storage: *i64 = nx_gguf_load_tensor_to_f32(buf, hdr, name_buf, n_len,
                                                 nv_out, inner_err)
  if inner_err[0] != NX_GLF_OK { return 0 as *NxF32LazyWeight }
  return nx_f32_lazy_weight_new_f32(storage, rows, cols)
165}
166
167// Load a 1-D f32 gamma vector (RMSNorm scale) into a raw *i64.
168
169func _flll_load_gamma(buf: *u8, hdr: *NxGgufHeader,
                    name: *u8, name_len: nx_int) -> *i64 {
  let n_out: *i64 = sys_mmap(8) as *i64
  let err:   *i64 = sys_mmap(8) as *i64
  let storage: *i64 = nx_gguf_load_tensor_to_f32(buf, hdr, name, name_len,
                                                 n_out, err)
  if err[0] != NX_GLF_OK { return 0 as *i64 }
  return storage
177}
178
179// Public: populate NxF32LlamaLayerLazy by loading all 9 tensors for layer_idx.
180
181func nx_f32_llama_layer_lazy_load_from_gguf(buf: *u8, hdr: *NxGgufHeader,
                                          layer_idx: nx_int,
                                          layer_out: *NxF32LlamaLayerLazy,
                                          out_err: *i64) -> nx_int {
  if layer_idx < 0 {
      out_err[0] = NX_FLLL_ERR_BAD_LAYER
      return NX_FLLL_ERR_BAD_LAYER
  }
  if layer_out == (0 as *NxF32LlamaLayerLazy) {
      out_err[0] = NX_FLLL_ERR_NULL
      return NX_FLLL_ERR_NULL
  }
193
  // Suffix byte arrays (mirror nx_f32_llama_layer_load).
  let s_an: *u8 = sys_mmap(16)
  s_an[0]=0x61 as u8; s_an[1]=0x74 as u8; s_an[2]=0x74 as u8; s_an[3]=0x6e as u8
  s_an[4]=0x5f as u8; s_an[5]=0x6e as u8; s_an[6]=0x6f as u8; s_an[7]=0x72 as u8
  s_an[8]=0x6d as u8; s_an[9]=0x2e as u8; s_an[10]=0x77 as u8; s_an[11]=0x65 as u8
  s_an[12]=0x69 as u8; s_an[13]=0x67 as u8; s_an[14]=0x68 as u8; s_an[15]=0x74 as u8
200
  let s_q: *u8 = sys_mmap(13)
  s_q[0]=0x61 as u8; s_q[1]=0x74 as u8; s_q[2]=0x74 as u8; s_q[3]=0x6e as u8
  s_q[4]=0x5f as u8; s_q[5]=0x71 as u8; s_q[6]=0x2e as u8; s_q[7]=0x77 as u8
  s_q[8]=0x65 as u8; s_q[9]=0x69 as u8; s_q[10]=0x67 as u8; s_q[11]=0x68 as u8
  s_q[12]=0x74 as u8
206
  let s_k: *u8 = sys_mmap(13)
  s_k[0]=0x61 as u8; s_k[1]=0x74 as u8; s_k[2]=0x74 as u8; s_k[3]=0x6e as u8
  s_k[4]=0x5f as u8; s_k[5]=0x6b as u8; s_k[6]=0x2e as u8; s_k[7]=0x77 as u8
  s_k[8]=0x65 as u8; s_k[9]=0x69 as u8; s_k[10]=0x67 as u8; s_k[11]=0x68 as u8
  s_k[12]=0x74 as u8
212
  let s_v: *u8 = sys_mmap(13)
  s_v[0]=0x61 as u8; s_v[1]=0x74 as u8; s_v[2]=0x74 as u8; s_v[3]=0x6e as u8
  s_v[4]=0x5f as u8; s_v[5]=0x76 as u8; s_v[6]=0x2e as u8; s_v[7]=0x77 as u8
  s_v[8]=0x65 as u8; s_v[9]=0x69 as u8; s_v[10]=0x67 as u8; s_v[11]=0x68 as u8
  s_v[12]=0x74 as u8
218
  let s_o: *u8 = sys_mmap(18)
  s_o[0]=0x61 as u8; s_o[1]=0x74 as u8; s_o[2]=0x74 as u8; s_o[3]=0x6e as u8
  s_o[4]=0x5f as u8; s_o[5]=0x6f as u8; s_o[6]=0x75 as u8; s_o[7]=0x74 as u8
  s_o[8]=0x70 as u8; s_o[9]=0x75 as u8; s_o[10]=0x74 as u8; s_o[11]=0x2e as u8
  s_o[12]=0x77 as u8; s_o[13]=0x65 as u8; s_o[14]=0x69 as u8; s_o[15]=0x67 as u8
  s_o[16]=0x68 as u8; s_o[17]=0x74 as u8
225
  let s_fn: *u8 = sys_mmap(15)
  s_fn[0]=0x66 as u8; s_fn[1]=0x66 as u8; s_fn[2]=0x6e as u8; s_fn[3]=0x5f as u8
  s_fn[4]=0x6e as u8; s_fn[5]=0x6f as u8; s_fn[6]=0x72 as u8; s_fn[7]=0x6d as u8
  s_fn[8]=0x2e as u8; s_fn[9]=0x77 as u8; s_fn[10]=0x65 as u8; s_fn[11]=0x69 as u8
  s_fn[12]=0x67 as u8; s_fn[13]=0x68 as u8; s_fn[14]=0x74 as u8
231
  let s_fg: *u8 = sys_mmap(15)
  s_fg[0]=0x66 as u8; s_fg[1]=0x66 as u8; s_fg[2]=0x6e as u8; s_fg[3]=0x5f as u8
  s_fg[4]=0x67 as u8; s_fg[5]=0x61 as u8; s_fg[6]=0x74 as u8; s_fg[7]=0x65 as u8
  s_fg[8]=0x2e as u8; s_fg[9]=0x77 as u8; s_fg[10]=0x65 as u8; s_fg[11]=0x69 as u8
  s_fg[12]=0x67 as u8; s_fg[13]=0x68 as u8; s_fg[14]=0x74 as u8
237
  let s_fu: *u8 = sys_mmap(13)
  s_fu[0]=0x66 as u8; s_fu[1]=0x66 as u8; s_fu[2]=0x6e as u8; s_fu[3]=0x5f as u8
  s_fu[4]=0x75 as u8; s_fu[5]=0x70 as u8; s_fu[6]=0x2e as u8; s_fu[7]=0x77 as u8
  s_fu[8]=0x65 as u8; s_fu[9]=0x69 as u8; s_fu[10]=0x67 as u8; s_fu[11]=0x68 as u8
  s_fu[12]=0x74 as u8
243
  let s_fd: *u8 = sys_mmap(15)
  s_fd[0]=0x66 as u8; s_fd[1]=0x66 as u8; s_fd[2]=0x6e as u8; s_fd[3]=0x5f as u8
  s_fd[4]=0x64 as u8; s_fd[5]=0x6f as u8; s_fd[6]=0x77 as u8; s_fd[7]=0x6e as u8
  s_fd[8]=0x2e as u8; s_fd[9]=0x77 as u8; s_fd[10]=0x65 as u8; s_fd[11]=0x69 as u8
  s_fd[12]=0x67 as u8; s_fd[13]=0x68 as u8; s_fd[14]=0x74 as u8
249
  // Qwen2 attention biases: "attn_q.bias" / "attn_k.bias" / "attn_v.bias" (11 bytes each).
  let s_bq: *u8 = sys_mmap(11)
  s_bq[0]=0x61 as u8; s_bq[1]=0x74 as u8; s_bq[2]=0x74 as u8; s_bq[3]=0x6e as u8
  s_bq[4]=0x5f as u8; s_bq[5]=0x71 as u8; s_bq[6]=0x2e as u8; s_bq[7]=0x62 as u8
  s_bq[8]=0x69 as u8; s_bq[9]=0x61 as u8; s_bq[10]=0x73 as u8
255
  let s_bk: *u8 = sys_mmap(11)
  s_bk[0]=0x61 as u8; s_bk[1]=0x74 as u8; s_bk[2]=0x74 as u8; s_bk[3]=0x6e as u8
  s_bk[4]=0x5f as u8; s_bk[5]=0x6b as u8; s_bk[6]=0x2e as u8; s_bk[7]=0x62 as u8
  s_bk[8]=0x69 as u8; s_bk[9]=0x61 as u8; s_bk[10]=0x73 as u8
260
  let s_bv: *u8 = sys_mmap(11)
  s_bv[0]=0x61 as u8; s_bv[1]=0x74 as u8; s_bv[2]=0x74 as u8; s_bv[3]=0x6e as u8
  s_bv[4]=0x5f as u8; s_bv[5]=0x76 as u8; s_bv[6]=0x2e as u8; s_bv[7]=0x62 as u8
  s_bv[8]=0x69 as u8; s_bv[9]=0x61 as u8; s_bv[10]=0x73 as u8
265
  let name_buf: *u8 = sys_mmap(64)
267
  // gamma_attn (1D)
  let n_an: nx_int = _flll_fmt_name(layer_idx, s_an, 16, name_buf)
  layer_out.gamma_attn = _flll_load_gamma(buf, hdr, name_buf, n_an)
  if (layer_out.gamma_attn as i64) == 0 {
      out_err[0] = NX_FLLL_ERR_NOT_FOUND
      return NX_FLLL_ERR_NOT_FOUND
  }
275
  // W_q (2D, possibly Q4_K)
  let n_q: nx_int = _flll_fmt_name(layer_idx, s_q, 13, name_buf)
  layer_out.W_q = _flll_build_lazy(buf, hdr, name_buf, n_q)
  if (layer_out.W_q as i64) == 0 {
      out_err[0] = NX_FLLL_ERR_NOT_FOUND
      return NX_FLLL_ERR_NOT_FOUND
  }
283
  let n_k: nx_int = _flll_fmt_name(layer_idx, s_k, 13, name_buf)
  layer_out.W_k = _flll_build_lazy(buf, hdr, name_buf, n_k)
  if (layer_out.W_k as i64) == 0 { out_err[0] = NX_FLLL_ERR_NOT_FOUND; return NX_FLLL_ERR_NOT_FOUND }
287
  let n_v: nx_int = _flll_fmt_name(layer_idx, s_v, 13, name_buf)
  layer_out.W_v = _flll_build_lazy(buf, hdr, name_buf, n_v)
  if (layer_out.W_v as i64) == 0 { out_err[0] = NX_FLLL_ERR_NOT_FOUND; return NX_FLLL_ERR_NOT_FOUND }
291
  let n_o: nx_int = _flll_fmt_name(layer_idx, s_o, 18, name_buf)
  layer_out.W_o = _flll_build_lazy(buf, hdr, name_buf, n_o)
  if (layer_out.W_o as i64) == 0 { out_err[0] = NX_FLLL_ERR_NOT_FOUND; return NX_FLLL_ERR_NOT_FOUND }
295
  let n_fn: nx_int = _flll_fmt_name(layer_idx, s_fn, 15, name_buf)
  layer_out.gamma_ffn = _flll_load_gamma(buf, hdr, name_buf, n_fn)
  if (layer_out.gamma_ffn as i64) == 0 { out_err[0] = NX_FLLL_ERR_NOT_FOUND; return NX_FLLL_ERR_NOT_FOUND }
299
  let n_fg: nx_int = _flll_fmt_name(layer_idx, s_fg, 15, name_buf)
  layer_out.W_gate = _flll_build_lazy(buf, hdr, name_buf, n_fg)
  if (layer_out.W_gate as i64) == 0 { out_err[0] = NX_FLLL_ERR_NOT_FOUND; return NX_FLLL_ERR_NOT_FOUND }
303
  let n_fu: nx_int = _flll_fmt_name(layer_idx, s_fu, 13, name_buf)
  layer_out.W_up = _flll_build_lazy(buf, hdr, name_buf, n_fu)
  if (layer_out.W_up as i64) == 0 { out_err[0] = NX_FLLL_ERR_NOT_FOUND; return NX_FLLL_ERR_NOT_FOUND }
307
  let n_fd: nx_int = _flll_fmt_name(layer_idx, s_fd, 15, name_buf)
  layer_out.W_down = _flll_build_lazy(buf, hdr, name_buf, n_fd)
  if (layer_out.W_down as i64) == 0 { out_err[0] = NX_FLLL_ERR_NOT_FOUND; return NX_FLLL_ERR_NOT_FOUND }
311
  // Qwen2 attention biases (1D f32). NOT an error if absent -- Llama-style models have none;
  // block_v4 guards on (bias_q != 0). This was the DROPPED TERM behind the garbage forward.
  let n_bq: nx_int = _flll_fmt_name(layer_idx, s_bq, 11, name_buf)
  layer_out.bias_q = _flll_load_gamma(buf, hdr, name_buf, n_bq)
  let n_bk: nx_int = _flll_fmt_name(layer_idx, s_bk, 11, name_buf)
  layer_out.bias_k = _flll_load_gamma(buf, hdr, name_buf, n_bk)
  let n_bv: nx_int = _flll_fmt_name(layer_idx, s_bv, 11, name_buf)
  layer_out.bias_v = _flll_load_gamma(buf, hdr, name_buf, n_bv)
320
  out_err[0] = NX_FLLL_OK
  return NX_FLLL_OK
323}
nx_f32_llama_layer_lazy_load.nx source

`nx_f32_llama_layer_lazy_load.nx` source
