nx_gguf_load_lazy_test.nx source
↩ module page · 158 lines · 6305 B
1// nx_gguf_load_lazy_test.nx -- smoke for nx_gguf_load_lazy.nx +
2// QUANTITATIVE proof of lazy-vs-eager RAM footprint.
3//
4// Builds a GGUF with 2 tensors:
5// "score" F32 [2,2] (4 values, 16 bytes raw)
6// "weight" Q8_0 [32] (32 values, 34 bytes raw block)
7//
8// Test classes:
9// A) Lazy load returns NxPlacedTensor with placement=MMAP and
10// storage_bytes=0.
11// B) Materialize on demand via nx_placed_tensor_get reproduces
12// the same values nx_gguf_load_tensor would have produced.
13// C) Demote returns to zero-storage state; re-materialize works.
14// D) Footprint accounting: handle_bytes scales linearly with
15// tensor count; eager would scale with tensor-DATA bytes.
16// Verify lazy_handle_bytes(n) = n * 88.
17// E) Not-found verdict propagated.
18
19import "nx_syscalls.nx"
20import "nx_tier.nx"
21import "nx_le.nx"
22import "nx_tensor.nx"
23import "nx_gguf.nx"
24import "nx_gguf_load.nx"
25import "nx_placement.nx"
26import "nx_gguf_load_lazy.nx"
27
28func main() -> i64 {
29 // ----- Build GGUF (same layout as nx_gguf_load_test.nx) -----
30 let buf: *u8 = sys_mmap(1024)
31
32 buf[0]=0x47; buf[1]=0x47; buf[2]=0x55; buf[3]=0x46
33 buf[4]=3
34 nx_le_write_u64(buf, 8, 2)
35 nx_le_write_u64(buf, 16, 0)
36
37 nx_le_write_u64(buf, 24, 5)
38 buf[32]=0x73; buf[33]=0x63; buf[34]=0x6F; buf[35]=0x72; buf[36]=0x65
39 nx_le_write_u32(buf, 37, 2)
40 nx_le_write_u64(buf, 41, 2)
41 nx_le_write_u64(buf, 49, 2)
42 nx_le_write_u32(buf, 57, NX_GGML_TYPE_F32)
43 nx_le_write_u64(buf, 61, 0)
44
45 let t1_off: i64 = 69
46 nx_le_write_u64(buf, t1_off + 0, 6)
47 buf[t1_off + 8 + 0]=0x77; buf[t1_off + 8 + 1]=0x65
48 buf[t1_off + 8 + 2]=0x69; buf[t1_off + 8 + 3]=0x67
49 buf[t1_off + 8 + 4]=0x68; buf[t1_off + 8 + 5]=0x74
50 nx_le_write_u32(buf, t1_off + 14, 1)
51 nx_le_write_u64(buf, t1_off + 18, 32)
52 nx_le_write_u32(buf, t1_off + 26, NX_GGML_TYPE_Q8_0)
53 nx_le_write_u64(buf, t1_off + 30, 16)
54
55 let data_off: i64 = 128
56 nx_le_write_u32(buf, data_off + 0, 0x3F800000)
57 nx_le_write_u32(buf, data_off + 4, 0x40000000)
58 nx_le_write_u32(buf, data_off + 8, 0x3F000000)
59 nx_le_write_u32(buf, data_off + 12, 0xBF800000)
60 let q8_off: i64 = data_off + 16
61 nx_le_write_u16(buf, q8_off + 0, 0x3800)
62 var qi: nx_int = 0
63 while qi < 32 {
64 buf[q8_off + 2 + qi] = qi
65 qi = qi + 1
66 }
67
68 let hdr: *NxGgufHeader = sys_mmap(NX_GGUF_HDR_BYTES) as *NxGgufHeader
69 let v_p: nx_int = nx_gguf_parse(buf, 1024, hdr)
70 if v_p != NX_GGUF_OK { return 10 + v_p }
71 if hdr.n_tensors != 2 { return 19 }
72
73 // ----- A) Lazy load -----
74 let score_name: *u8 = sys_mmap(5)
75 score_name[0]=0x73; score_name[1]=0x63; score_name[2]=0x6F
76 score_name[3]=0x72; score_name[4]=0x65
77 let err: *i64 = sys_mmap(8) as *i64
78 err[0] = 0
79 let lazy_score: *NxPlacedTensor = nx_gguf_load_tensor_lazy(
80 buf, hdr, score_name, 5, err)
81 if err[0] != NX_GL_OK { return 20 }
82 if lazy_score.placement != NX_PLACE_MMAP { return 21 }
83 if nx_placed_tensor_is_live(lazy_score) != 0 { return 22 }
84 if nx_placed_tensor_storage_bytes(lazy_score) != 0 { return 23 }
85 if lazy_score.src_n_values != 4 { return 24 }
86 if lazy_score.src_ggml_type != NX_GGML_TYPE_F32 { return 25 }
87 if lazy_score.shape_0 != 2 { return 26 }
88 if lazy_score.shape_1 != 2 { return 27 }
89 if lazy_score.n_dims != 2 { return 28 }
90
91 // ----- B) Materialize on demand -----
92 let t_score: *NxTensor = nx_placed_tensor_get(lazy_score)
93 if t_score == (0 as *NxTensor) { return 30 }
94 if nx_placed_tensor_is_live(lazy_score) != 1 { return 31 }
95 if nx_placed_tensor_storage_bytes(lazy_score) != 4 * 8 { return 32 }
96 let sp: *i64 = t_score.storage as *i64
97 if sp[0] != 1024 { return 40 } // 1.0 -> Q10 1024
98 if sp[1] != 2048 { return 41 }
99 if sp[2] != 512 { return 42 }
100 if sp[3] != (0 - 1024) { return 43 }
101
102 // ----- C) Lazy-load Q8_0 + materialize -----
103 let weight_name: *u8 = sys_mmap(6)
104 weight_name[0]=0x77; weight_name[1]=0x65; weight_name[2]=0x69
105 weight_name[3]=0x67; weight_name[4]=0x68; weight_name[5]=0x74
106 err[0] = 0
107 let lazy_weight: *NxPlacedTensor = nx_gguf_load_tensor_lazy(
108 buf, hdr, weight_name, 6, err)
109 if err[0] != NX_GL_OK { return 50 }
110 if nx_placed_tensor_storage_bytes(lazy_weight) != 0 { return 51 }
111
112 let t_weight: *NxTensor = nx_placed_tensor_get(lazy_weight)
113 if t_weight == (0 as *NxTensor) { return 60 }
114 if nx_placed_tensor_storage_bytes(lazy_weight) != 32 * 8 { return 61 }
115 let wp: *i64 = t_weight.storage as *i64
116 // scale_q10 = 512; qs[i] = i; expected[i] = 512 * i
117 var ki: nx_int = 0
118 while ki < 32 {
119 if wp[ki] != 512 * ki { return 70 + ki }
120 ki = ki + 1
121 }
122
123 // ----- D) Demote + re-materialize -----
124 let v_d: nx_int = nx_placed_tensor_demote(lazy_score)
125 if v_d != NX_PT_OK { return 120 }
126 if nx_placed_tensor_is_live(lazy_score) != 0 { return 121 }
127 if nx_placed_tensor_storage_bytes(lazy_score) != 0 { return 122 }
128
129 let t_re: *NxTensor = nx_placed_tensor_get(lazy_score)
130 if t_re == (0 as *NxTensor) { return 130 }
131 let rp: *i64 = t_re.storage as *i64
132 if rp[0] != 1024 { return 140 }
133 if rp[3] != (0 - 1024) { return 141 }
134
135 // ----- E) Handle-size accounting (STRUCTURAL fact, not a perf claim) -----
136 //
137 // Each NxPlacedTensor handle is fixed at 88 bytes regardless of
138 // the underlying tensor size. This is a property of the struct
139 // layout, NOT a claim of "Nx less RAM than CUDA/llama.cpp" -- the
140 // actual runtime working-set depends on which tensors get
141 // materialized when, and llama.cpp's mmap+per-block-dequant path
142 // achieves the same architectural property by a different name.
143 // Real comparisons require running both stacks on the same model
144 // + workload + hardware.
145 if nx_gguf_lazy_handle_bytes(1) != 88 { return 150 }
146 if nx_gguf_lazy_handle_bytes(290) != 25520 { return 151 }
147 if nx_gguf_lazy_handle_bytes(1024) != 90112 { return 152 }
148
149 // ----- F) Not-found propagates -----
150 let bad: *u8 = sys_mmap(3)
151 bad[0]=0x78; bad[1]=0x78; bad[2]=0x78
152 err[0] = 0
153 let _ng: *NxPlacedTensor = nx_gguf_load_tensor_lazy(
154 buf, hdr, bad, 3, err)
155 if err[0] != NX_GL_ERR_NOT_FOUND { return 160 }
156
157 return 0
158}