nx_f32_llm_run_v3_test.nx source
↩ module page · 63 lines · 1928 B
1// nx_f32_llm_run_v3_test.nx -- smoke for lazy runner v3.
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_lazy_weight.nx"
10import "nx_f32_llama_block_v4.nx"
11import "nx_f32_llama_stack_v4.nx"
12import "nx_f32_llama_layer_lazy_load.nx"
13import "nx_f32_llm.nx"
14import "nx_f32_llm_v4.nx"
15import "nx_f32_sampler.nx"
16import "nx_f32_llm_run_v2.nx"
17import "nx_f32_llm_run_v3.nx"
18
19func main() -> i64 {
20 let b: *NxGgufFixtureBundle = nx_gft_build_tiny_llama(42 as i64)
21 if nx_gft_is_built(b) != 1 { return 10 }
22
23 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
24 model.n_layers = b.n_layers
25 model.hidden_dim = b.hidden_dim
26 model.n_heads = b.n_heads
27 model.n_kv_heads = b.n_heads
28 model.head_dim = b.head_dim
29 model.ffn_dim = b.ffn_dim
30 model.vocab_size = b.vocab_size
31
32 let out_err: *i64 = sys_mmap(8) as *i64
33 let v_load: nx_int = nx_f32_llm_load_weights_v4_from_gguf(b.gguf_buf, b.hdr,
34 model, out_err)
35 if v_load != NX_FLV4_OK { return 20 + v_load }
36
37 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(
38 model.n_layers, model.n_kv_heads, 32, model.head_dim)
39
40 let prompt: *u8 = sys_mmap(2)
41 prompt[0] = 0x61 as u8
42 prompt[1] = 0x62 as u8
43
44 let cfg: *NxF32SamplerCfg = nx_f32_sampler_cfg_alloc()
45 let out_buf: *u8 = sys_mmap(32)
46
47 let eps: i64 = 0x322BCC77
48 let attn_scale: i64 = 0x3F3504F3
49 let rope_base: i64 = 0x4548F000
50
51 let n_emit: nx_int = nx_f32_llm_run_v3(
52 model, b.bpe, cache, prompt, 2, 5, cfg,
53 eps, attn_scale, rope_base, 0,
54 b.prng_state, 0 - 1, out_buf, 32)
55 if n_emit != 5 { return 30 }
56 var i: nx_int = 0
57 while i < 5 {
58 if out_buf[i] != (0x61 as u8) { return 40 + i }
59 i = i + 1
60 }
61
62 return 0
63}