nx_f32_llm_run_v2_test.nx source
↩ module page · 103 lines · 3803 B
1// nx_f32_llm_run_v2_test.nx -- smoke for nx_f32_llm_run_v2.nx.
2
3import "nx_syscalls.nx"
4import "nx_tier.nx"
5import "nx_bpe.nx"
6import "nx_gguf.nx"
7import "nx_gguf_fixture_tiny.nx"
8import "nx_f32_kv_cache.nx"
9import "nx_f32_llama_block.nx"
10import "nx_f32_llama_stack.nx"
11import "nx_f32_llama_layer_load.nx"
12import "nx_f32_llm.nx"
13import "nx_f32_llm_load.nx"
14import "nx_f32_sampler.nx"
15import "nx_f32_llm_run_v2.nx"
16
17func main() -> i64 {
18 let b: *NxGgufFixtureBundle = nx_gft_build_tiny_llama(42 as i64)
19 if nx_gft_is_built(b) != 1 { return 10 }
20
21 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
22 model.n_layers = b.n_layers
23 model.hidden_dim = b.hidden_dim
24 model.n_heads = b.n_heads
25 model.n_kv_heads = b.n_heads
26 model.head_dim = b.head_dim
27 model.ffn_dim = b.ffn_dim
28 model.vocab_size = b.vocab_size
29
30 let out_err: *i64 = sys_mmap(8) as *i64
31 let v_load: nx_int = nx_f32_llm_load_weights_from_gguf(b.gguf_buf, b.hdr,
32 model, out_err)
33 if v_load != NX_FLM_OK { return 20 + v_load }
34
35 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(
36 model.n_layers, model.n_kv_heads, 32, model.head_dim)
37
38 let prompt: *u8 = sys_mmap(2)
39 prompt[0] = 0x61 as u8 // 'a'
40 prompt[1] = 0x62 as u8 // 'b'
41
42 let out_buf: *u8 = sys_mmap(32)
43
44 let eps: i64 = 0x322BCC77
45 let attn_scale: i64 = 0x3F3504F3
46 let rope_base: i64 = 0x4548F000
47
48 // ===== Test 1: Default config (greedy, no penalty) =====
49 // Zero-weight => argmax of zero logits = 0 => 'a' emitted repeatedly.
50 let cfg1: *NxF32SamplerCfg = nx_f32_sampler_cfg_alloc()
51 // cfg1 has defaults: greedy (top_k=0, top_p=0 -> argmax), penalty=1.0
52
53 let n_emit: nx_int = nx_f32_llm_run_v2(
54 model, b.bpe, cache, prompt, 2, 5, cfg1,
55 eps, attn_scale, rope_base, 0,
56 b.prng_state, 0 - 1, out_buf, 32)
57 if n_emit != 5 { return 30 }
58 var i: nx_int = 0
59 while i < 5 {
60 if out_buf[i] != (0x61 as u8) { return 40 + i }
61 i = i + 1
62 }
63
64 // ===== Test 2: Config with recent_cap=4 + penalty=2.0 =====
65 // With penalty applied to recent, the argmax would shift if any logit
66 // had been non-zero -- but for zero weights, all logits are 0, so
67 // penalty has no effect (mul/div on 0 is still 0). Argmax still 0.
68 let cache2: *NxF32KVCache = nx_f32_kv_cache_alloc(
69 model.n_layers, model.n_kv_heads, 32, model.head_dim)
70 let cfg2: *NxF32SamplerCfg = nx_f32_sampler_cfg_alloc()
71 cfg2.penalty_f32 = 0x40000000 // 2.0
72 cfg2.recent_cap = 4
73 let out2: *u8 = sys_mmap(32)
74 let n_emit2: nx_int = nx_f32_llm_run_v2(
75 model, b.bpe, cache2, prompt, 2, 5, cfg2,
76 eps, attn_scale, rope_base, 0,
77 b.prng_state, 0 - 1, out2, 32)
78 if n_emit2 != 5 { return 50 }
79
80 // ===== Test 3: EOS termination at id=0 =====
81 let cache3: *NxF32KVCache = nx_f32_kv_cache_alloc(
82 model.n_layers, model.n_kv_heads, 32, model.head_dim)
83 let cfg3: *NxF32SamplerCfg = nx_f32_sampler_cfg_alloc()
84 let out3: *u8 = sys_mmap(32)
85 let n_emit3: nx_int = nx_f32_llm_run_v2(
86 model, b.bpe, cache3, prompt, 2, 5, cfg3,
87 eps, attn_scale, rope_base, 0,
88 b.prng_state, 0, out3, 32)
89 if n_emit3 != 0 { return 60 } // EOS immediately
90
91 // ===== Test 4: out_cap=2 =====
92 let cache4: *NxF32KVCache = nx_f32_kv_cache_alloc(
93 model.n_layers, model.n_kv_heads, 32, model.head_dim)
94 let cfg4: *NxF32SamplerCfg = nx_f32_sampler_cfg_alloc()
95 let out4: *u8 = sys_mmap(8)
96 let n_emit4: nx_int = nx_f32_llm_run_v2(
97 model, b.bpe, cache4, prompt, 2, 10, cfg4,
98 eps, attn_scale, rope_base, 0,
99 b.prng_state, 0 - 1, out4, 2)
100 if n_emit4 != 2 { return 70 }
101
102 return 0
103}