nx_f32_llm_e2e_test.nx source
↩ module page · 105 lines · 4016 B
1// nx_f32_llm_e2e_test.nx -- end-to-end bits-up text-to-text smoke.
2//
3// Composes the entire f32 LLM pipeline on the tiny-Llama fixture:
4// 1. Build fixture (12-tensor GGUF + BPE vocab + spec)
5// 2. Allocate NxF32LlamaModel + populate dims from fixture
6// 3. Load weights via nx_f32_llm_load_weights_from_gguf
7// 4. Tokenize 'ab' via nx_bpe_encode -> [0, 1]
8// 5. Allocate KV cache
9// 6. Run nx_f32_llm_forward over the 2 prompt tokens
10// 7. Argmax last token's logits -> next token id
11// 8. Detokenize -> byte
12//
13// Zero-weight fixture means logits are all zero -> argmax returns 0
14// (the first token id 'a' = byte 0x61). This proves the entire
15// pipeline executes without segfault and produces a deterministic
16// known-good result.
17
18import "nx_syscalls.nx"
19import "nx_tier.nx"
20import "nx_bpe.nx"
21import "nx_gguf.nx"
22import "nx_gguf_fixture_tiny.nx"
23import "nx_f32_kv_cache.nx"
24import "nx_f32_llama_block.nx"
25import "nx_f32_llama_stack.nx"
26import "nx_f32_llama_layer_load.nx"
27import "nx_f32_llm.nx"
28import "nx_f32_llm_load.nx"
29import "nx_f32_sampler.nx"
30
31func main() -> i64 {
32 let b: *NxGgufFixtureBundle = nx_gft_build_tiny_llama(42 as i64)
33 if nx_gft_is_built(b) != 1 { return 10 }
34
35 // ===== Phase 1: build + load model =====
36 let model: *NxF32LlamaModel = nx_f32_llama_model_alloc()
37 model.n_layers = b.n_layers
38 model.hidden_dim = b.hidden_dim
39 model.n_heads = b.n_heads
40 model.n_kv_heads = b.n_heads
41 model.head_dim = b.head_dim
42 model.ffn_dim = b.ffn_dim
43 model.vocab_size = b.vocab_size
44
45 let out_err: *i64 = sys_mmap(8) as *i64
46 let v_load: nx_int = nx_f32_llm_load_weights_from_gguf(
47 b.gguf_buf, b.hdr, model, out_err)
48 if v_load != NX_FLM_OK { return 20 + v_load }
49
50 // ===== Phase 2: tokenize prompt =====
51 // Fixture's BPE vocab has 'a'=0, 'b'=1, 'c'=2, 'd'=3.
52 let prompt: *u8 = sys_mmap(2)
53 prompt[0] = 0x61 as u8 // 'a'
54 prompt[1] = 0x62 as u8 // 'b'
55
56 let tokens: *i64 = sys_mmap(8 * 8) as *i64
57 let n_tok: nx_int = nx_bpe_encode(b.bpe, prompt, 2, tokens)
58 if n_tok != 2 { return 30 }
59 if tokens[0] != 0 { return 31 } // 'a' -> 0
60 if tokens[1] != 1 { return 32 } // 'b' -> 1
61
62 // ===== Phase 3: forward pass =====
63 let cache: *NxF32KVCache = nx_f32_kv_cache_alloc(
64 model.n_layers, model.n_kv_heads, 16, model.head_dim)
65 if cache == (0 as *NxF32KVCache) { return 40 }
66
67 let logits: *i64 = sys_mmap(2 * model.vocab_size * 8) as *i64
68
69 let eps: i64 = 0x322BCC77
70 let attn_scale: i64 = 0x3F3504F3
71 let rope_base: i64 = 0x4548F000
72
73 let v_fwd: nx_int = nx_f32_llm_forward(model, tokens, 2, cache,
74 eps, attn_scale, rope_base, 0, logits)
75 if v_fwd != NX_F32_LLM_OK { return 50 + v_fwd }
76
77 if nx_f32_kv_cache_get_seq_len(cache) != 2 { return 60 }
78
79 // ===== Phase 4: sample next token =====
80 // Zero weights -> all logits zero -> argmax returns 0 ('a').
81 let next_id: nx_int = nx_f32_sampler_argmax(
82 ((logits as i64) + 1 * model.vocab_size * 8) as *i64, model.vocab_size)
83 if next_id != 0 { return 70 }
84
85 // ===== Phase 5: detokenize =====
86 let next_buf: *i64 = sys_mmap(8) as *i64
87 next_buf[0] = next_id as i64
88 let out_text: *u8 = sys_mmap(16)
89 let n_bytes: nx_int = nx_bpe_decode(b.bpe, next_buf, 1, out_text)
90 if n_bytes != 1 { return 80 }
91 if out_text[0] != (0x61 as u8) { return 81 } // 'a'
92
93 // ===== Phase 6: continue autoregressive decode 1 more step =====
94 // Append next_id and run a 1-token decode forward.
95 let v_dec: nx_int = nx_f32_llm_forward(model, next_buf, 1, cache,
96 eps, attn_scale, rope_base, 0, logits)
97 if v_dec != NX_F32_LLM_OK { return 90 + v_dec }
98 if nx_f32_kv_cache_get_seq_len(cache) != 3 { return 100 }
99
100 // Still all zero -> argmax 0.
101 let next_id2: nx_int = nx_f32_sampler_argmax(logits, model.vocab_size)
102 if next_id2 != 0 { return 110 }
103
104 return 0
105}