nx_f32_lazy_weight_test.nx source

1// nx_f32_lazy_weight_test.nx -- smoke for nx_f32_lazy_weight.nx. 2// 3// Two dispatch paths: 4// F32 path: build a 1x4 f32 weight as raw bits, do lazy_matmul, 5// verify result equals direct nx_f32_matmul output. 6// Q4_K path: build a 1x256 Q4_K super-block with known values, 7// do lazy_matmul, verify result equals direct 8// nx_f32_q4k_matmul output. 9 10import "nx_syscalls.nx" 11import "nx_tier.nx" 12import "nx_le.nx" 13import "nx_gguf_load.nx" 14import "nx_f32.nx" 15import "nx_f32_matmul.nx" 16import "nx_q4k_to_f32.nx" 17import "nx_f32_q4k_matmul.nx" 18import "nx_f32_lazy_weight.nx" 19 20func _write_block(buf: *u8, off: i64) -> i64 { 21 nx_le_write_u16(buf, off + 0, 0x3C00) // d = 1.0 22 nx_le_write_u16(buf, off + 2, 0x0000) 23 buf[off + 4] = 0x01 as u8 24 buf[off + 5] = 0x01 as u8 25 buf[off + 6] = 0x01 as u8 26 buf[off + 7] = 0x01 as u8 27 buf[off + 8] = 0 as u8 28 buf[off + 9] = 0 as u8 29 buf[off + 10] = 0 as u8 30 buf[off + 11] = 0 as u8 31 buf[off + 12] = 0x02 as u8 32 buf[off + 13] = 0 as u8 33 buf[off + 14] = 0 as u8 34 buf[off + 15] = 0 as u8 35 var z: i64 = 0 36 while z < 128 { 37 buf[off + 16 + z] = 0 as u8 38 z = z + 1 39 } 40 buf[off + 16 + 0] = 0x32 as u8 41 buf[off + 16 + 64] = 0x32 as u8 42 return off + 144 43} 44 45func main() -> i64 { 46 var vi: nx_int = 0 47 while vi < NX_LW_N_VERDICTS { 48 if nx_lw_verdict_is_valid(vi) != 1 { return 5 + vi } 49 vi = vi + 1 50 } 51 52 // ===== F32 path ===== 53 let W_f32_storage: *i64 = sys_mmap(4 * 8) as *i64 54 W_f32_storage[0] = 0x3F800000 // 1.0 55 W_f32_storage[1] = 0x40000000 // 2.0 56 W_f32_storage[2] = 0x40400000 // 3.0 57 W_f32_storage[3] = 0x40800000 // 4.0 58 59 let A1: *i64 = sys_mmap(8) as *i64 60 A1[0] = 0x40000000 // 2.0 61 62 // Direct call: A [1,1] @ B [1,4] -> C [1,4]. With A=2.0, expected: [2, 4, 6, 8]. 63 let C_direct: *i64 = sys_mmap(4 * 8) as *i64 64 nx_f32_matmul(A1, W_f32_storage, C_direct, 1, 1, 4) 65 66 // Lazy call: same inputs via dispatcher. 67 let W_lazy_f32: *NxF32LazyWeight = 68 nx_f32_lazy_weight_new_f32(W_f32_storage, 1, 4) 69 let C_lazy: *i64 = sys_mmap(4 * 8) as *i64 70 let v_f32: nx_int = nx_f32_lazy_matmul(A1, W_lazy_f32, C_lazy, 1, 1, 4) 71 if v_f32 != NX_LW_OK { return 20 + v_f32 } 72 73 // Both results must be bit-exact equal. 74 var i: nx_int = 0 75 while i < 4 { 76 if C_lazy[i] != C_direct[i] { return 30 + i } 77 i = i + 1 78 } 79 if C_lazy[0] != 0x40000000 { return 40 } // 2.0 80 if C_lazy[1] != 0x40800000 { return 41 } // 4.0 81 if C_lazy[2] != 0x40C00000 { return 42 } // 6.0 82 if C_lazy[3] != 0x41000000 { return 43 } // 8.0 83 84 // ===== Q4_K path ===== 85 // UPDATED 2026-07-07 to the organ's CURRENT axis semantics (B = n 86 // weight rows x k quantized values, k % 256 == 0). The previous 87 // k=1/n=256 shape encoded the OLD transposed layout: after the 88 // k-axis fix it tripped ERR_ALIGN, which the dispatcher then 89 // SWALLOWED (returned OK over a zero C) -- this test had been 90 // failing exit 70 unnoticed. The dispatcher now surfaces inner 91 // verdicts as NX_LW_ERR_INNER (asserted below) and routes Q4_K 92 // through the shared thread pool (nx_f32_q4k_matmul_pool). 93 let buf_q4k: *u8 = sys_mmap(256) 94 let _end: i64 = _write_block(buf_q4k, 0) 95 96 // A row: 1.0 at the block's nonzero weight positions [0,32,128]. 97 let A2: *i64 = sys_mmap(256 * 8) as *i64 98 var p: nx_int = 0 99 while p < 256 { 100 A2[p] = 0 101 p = p + 1 102 } 103 A2[0] = 0x3F800000 // 1.0 104 A2[32] = 0x3F800000 105 A2[128] = 0x3F800000 106 107 // Direct serial call: C[0] = 1*2 + 1*3 + 1*4 = 9.0. 108 let C_direct2: *i64 = sys_mmap(8) as *i64 109 let v_dir: nx_int = nx_f32_q4k_matmul(A2, buf_q4k, 0, C_direct2, 1, 256, 1) 110 if v_dir != NX_FQ4M_OK { return 45 + v_dir } 111 112 // Lazy dispatcher call (pooled path) must match bit-exactly. 113 let W_lazy_q4k: *NxF32LazyWeight = 114 nx_f32_lazy_weight_new_q4k(buf_q4k, 0, 256, 1) 115 let C_lazy2: *i64 = sys_mmap(8) as *i64 116 C_lazy2[0] = 0 - 777777 // poison 117 let v_q4k: nx_int = nx_f32_lazy_matmul(A2, W_lazy_q4k, C_lazy2, 1, 256, 1) 118 if v_q4k != NX_LW_OK { return 50 + v_q4k } 119 if C_lazy2[0] != C_direct2[0] { return 60 } 120 if C_lazy2[0] != 0x41100000 { return 70 } // 9.0 121 122 // The stale pre-axis-fix shape (k=1) must now SURFACE as 123 // NX_LW_ERR_INNER (inner ERR_ALIGN), not silently return OK. 124 let C_bad: *i64 = sys_mmap(8) as *i64 125 let v_bad: nx_int = nx_f32_lazy_matmul(A2, W_lazy_q4k, C_bad, 1, 1, 256) 126 if v_bad != NX_LW_ERR_INNER { return 75 } 127 128 return 0 129}