nx_f32_llm_load_test.nx source
↩ module page · 92 lines · 3307 B
1// nx_f32_llm_load_test.nx -- smoke for nx_f32_llm_load.nx.
2//
3// Uses the tiny-Llama fixture (1 layer, F32 zero weights) to verify
4// the model-level binder populates embed + gamma_out + lm_head + 1
5// layer correctly. Then composes with nx_f32_llm_forward to verify
6// the loaded model can run a forward pass.
7
8import "nx_syscalls.nx"
9import "nx_tier.nx"
10import "nx_gguf.nx"
11import "nx_gguf_fixture_tiny.nx"
12import "nx_gguf_load_f32.nx"
13import "nx_f32_kv_cache.nx"
14import "nx_f32_llama_block.nx"
15import "nx_f32_llama_stack.nx"
16import "nx_f32_llama_layer_load.nx"
17import "nx_f32_llm.nx"
18import "nx_f32_llm_load.nx"
19
20func main() -> i64 {
21 var vi: nx_int = 0
22 while vi < NX_FLM_N_VERDICTS {
23 if nx_flm_verdict_is_valid(vi) != 1 { return 5 + vi }
24 vi = vi + 1
25 }
26
27 // Build the tiny fixture.
28 let b: *NxGgufFixtureBundle = nx_gft_build_tiny_llama(42 as i64)
29 if nx_gft_is_built(b) != 1 { return 10 }
30
31 // Allocate model + pre-fill dims (v1 -- metadata walker is a future brick).
32 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
33 model.n_layers = b.n_layers
34 model.hidden_dim = b.hidden_dim
35 model.n_heads = b.n_heads
36 model.n_kv_heads = b.n_heads // tiny fixture has no GQA
37 model.head_dim = b.head_dim
38 model.ffn_dim = b.ffn_dim
39 model.vocab_size = b.vocab_size
40
41 let out_err: *i64 = sys_mmap(8) as *i64
42 let v: nx_int = nx_f32_llm_load_weights_from_gguf(b.gguf_buf, b.hdr,
43 model, out_err)
44 if v != NX_FLM_OK { return 20 + v }
45 if out_err[0] != NX_FLM_OK { return 30 + (out_err[0] as nx_int) }
46
47 // All four top-level pointers non-null.
48 if (model.embed_weights as i64) == 0 { return 40 }
49 if (model.gamma_out as i64) == 0 { return 41 }
50 if (model.lm_head as i64) == 0 { return 42 }
51 if (model.layers as i64) == 0 { return 43 }
52
53 // Fixture has zero F32 weights -> first values are 0.
54 if model.embed_weights[0] != 0 { return 50 }
55 if model.gamma_out[0] != 0 { return 51 }
56 if model.lm_head[0] != 0 { return 52 }
57
58 // Layer 0 was loaded.
59 let layer0: *NxF32LlamaLayer = (model.layers[0]) as *NxF32LlamaLayer
60 if (layer0 as i64) == 0 { return 60 }
61 if (layer0.W_q as i64) == 0 { return 61 }
62 if layer0.W_q[0] != 0 { return 62 }
63
64 // ===== End-to-end forward pass with the loaded model =====
65 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(
66 model.n_layers, model.n_kv_heads, 8, model.head_dim)
67 if cache == (0 as *NxF32KVCache) { return 70 }
68
69 let token_ids: *i64 = sys_mmap(2 * 8) as *i64
70 token_ids[0] = 0
71 token_ids[1] = 1
72
73 let logits: *i64 = sys_mmap(2 * model.vocab_size * 8) as *i64
74
75 let eps: i64 = 0x322BCC77
76 let attn_scale: i64 = 0x3F3504F3
77 let rope_base: i64 = 0x4548F000
78
79 let vF: nx_int = nx_f32_llm_forward(model, token_ids, 2, cache,
80 eps, attn_scale, rope_base, 0, logits)
81 if vF != NX_F32_LLM_OK { return 80 + vF }
82
83 // Zero-weight fixture -> all logits zero.
84 var li: nx_int = 0
85 while li < 2 * model.vocab_size {
86 if logits[li] != 0 { return 100 + li }
87 li = li + 1
88 }
89 if nx_f32_kv_cache_get_seq_len(cache) != 2 { return 130 }
90
91 return 0
92}