nx_f32_llm_run_test.nx source
↩ module page · 111 lines · 3752 B
1// nx_f32_llm_run_test.nx -- smoke for nx_f32_llm_run.nx.
2//
3// Composes tiny-Llama fixture + run loop. Zero-weight fixture
4// means argmax always returns 0 ('a' = byte 0x61). Verifies the
5// runner generates max_new_tokens bytes of 'a's.
6
7import "nx_syscalls.nx"
8import "nx_tier.nx"
9import "nx_bpe.nx"
10import "nx_gguf.nx"
11import "nx_gguf_fixture_tiny.nx"
12import "nx_f32_kv_cache.nx"
13import "nx_f32_llama_block.nx"
14import "nx_f32_llama_stack.nx"
15import "nx_f32_llama_layer_load.nx"
16import "nx_f32_llm.nx"
17import "nx_f32_llm_load.nx"
18import "nx_f32_sampler.nx"
19import "nx_f32_llm_run.nx"
20
21func main() -> i64 {
22 let b: *NxGgufFixtureBundle = nx_gft_build_tiny_llama(42 as i64)
23 if nx_gft_is_built(b) != 1 { return 10 }
24
25 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
26 model.n_layers = b.n_layers
27 model.hidden_dim = b.hidden_dim
28 model.n_heads = b.n_heads
29 model.n_kv_heads = b.n_heads
30 model.head_dim = b.head_dim
31 model.ffn_dim = b.ffn_dim
32 model.vocab_size = b.vocab_size
33
34 let out_err: *i64 = sys_mmap(8) as *i64
35 let v_load: nx_int = nx_f32_llm_load_weights_from_gguf(b.gguf_buf, b.hdr,
36 model, out_err)
37 if v_load != NX_FLM_OK { return 20 + v_load }
38
39 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(
40 model.n_layers, model.n_kv_heads, 32, model.head_dim)
41
42 let prompt: *u8 = sys_mmap(2)
43 prompt[0] = 0x61 as u8 // 'a'
44 prompt[1] = 0x62 as u8 // 'b'
45
46 let out_buf: *u8 = sys_mmap(32)
47
48 let eps: i64 = 0x322BCC77
49 let attn_scale: i64 = 0x3F3504F3
50 let rope_base: i64 = 0x4548F000
51
52 // top_k = 0 (greedy), max 5 tokens, eos = -1 (disabled).
53 let n_emit: nx_int = nx_f32_llm_run(
54 model, b.bpe, cache, prompt, 2,
55 5, // max_new_tokens
56 0x3F800000 as i64, // inv_temp (irrelevant for greedy)
57 0, // top_k = 0 -> greedy
58 eps, attn_scale, rope_base, 0,
59 b.prng_state,
60 0 - 1, // eos disabled
61 out_buf, 32)
62
63 if n_emit < 0 { return 30 }
64 if n_emit != 5 { return 31 }
65
66 // All zero-weight => argmax = 0 => 'a'.
67 var i: nx_int = 0
68 while i < 5 {
69 if out_buf[i] != (0x61 as u8) { return 40 + i }
70 i = i + 1
71 }
72
73 // Cache.seq_len = 2 prompt + 5 decoded = 7.
74 if nx_f32_kv_cache_get_seq_len(cache) != 7 { return 50 }
75
76 // ===== Run with EOS = 0 (the only token argmax can emit) =====
77 // Should emit ONE byte then stop (the very first sampled id is 0
78 // which matches eos, so the loop breaks before emitting that token).
79 let cache2: *NxF32KVCache = nx_f32_kv_cache_alloc(
80 model.n_layers, model.n_kv_heads, 32, model.head_dim)
81 let out2: *u8 = sys_mmap(32)
82 let n_emit2: nx_int = nx_f32_llm_run(
83 model, b.bpe, cache2, prompt, 2,
84 5,
85 0x3F800000 as i64,
86 0,
87 eps, attn_scale, rope_base, 0,
88 b.prng_state,
89 0, // eos = 0 (zero-weight argmax always emits 0)
90 out2, 32)
91 if n_emit2 != 0 { return 60 } // EOS hit immediately, nothing emitted
92
93 // ===== Run with cap = 2 =====
94 let cache3: *NxF32KVCache = nx_f32_kv_cache_alloc(
95 model.n_layers, model.n_kv_heads, 32, model.head_dim)
96 let out3: *u8 = sys_mmap(8)
97 let n_emit3: nx_int = nx_f32_llm_run(
98 model, b.bpe, cache3, prompt, 2,
99 10,
100 0x3F800000 as i64,
101 0,
102 eps, attn_scale, rope_base, 0,
103 b.prng_state,
104 0 - 1,
105 out3, 2)
106 if n_emit3 != 2 { return 70 }
107 if out3[0] != (0x61 as u8) { return 71 }
108 if out3[1] != (0x61 as u8) { return 72 }
109
110 return 0
111}