nx_f32_llm.nx source
↩ module page · 158 lines · 5686 B
1// nx_f32_llm.nx -- full Llama-class model forward pass.
2//
3// The API a real text-in -> logits-out inference loop calls.
4//
5// Composes:
6// embed lookup (token_id -> embedding row copy)
7// nx_f32_llama_stack (N-layer transformer forward)
8// nx_f32_rmsnorm (final pre-LM-head normalization)
9// nx_f32_matmul (LM head: hidden -> logits)
10//
11// Algorithm:
12// 1. Embedding lookup: for each token t in 0..n_tokens:
13// x[t, :] = embed_weights[token_ids[t], :]
14// 2. Stack forward through N layers (advances cache by n_tokens).
15// 3. Final RMSNorm per token with gamma_out.
16// 4. Logits = matmul(out_norm, lm_head) -> [n_tokens, vocab_size]
17//
18// The caller reads logits[(n_tokens-1)*vocab_size:n_tokens*vocab_size]
19// to get next-token logits for autoregressive sampling. Earlier
20// rows are valid for any per-position decode/classification.
21//
22// Tied-embed optimization: real Qwen/Llama variants often use the
23// same matrix for embedding and LM head. Caller can pass the
24// transposed embed_weights as lm_head -- this kernel does not care.
25//
26// genealogy_id: standard_llama_forward_pass
27// lineage_id: substrate_f32_llm_v1
28
29import "nx_syscalls.nx"
30import "nx_tier.nx"
31import "nx_f32.nx"
32import "nx_f32_rmsnorm.nx"
33import "nx_f32_matmul.nx"
34import "nx_f32_kv_cache.nx"
35import "nx_f32_llama_block.nx"
36import "nx_f32_llama_stack.nx"
37
38const NX_F32_LLM_OK: nx_int = 0
39const NX_F32_LLM_ERR_BAD_DIM: nx_int = 1
40const NX_F32_LLM_ERR_NULL: nx_int = 2
41const NX_F32_LLM_ERR_TOKEN: nx_int = 3
42const NX_F32_LLM_ERR_STACK: nx_int = 4
43const NX_F32_LLM_N_VERDICTS: nx_int = 5
44
45func nx_f32_llm_verdict_is_valid(v: nx_int) -> nx_int {
46 if v < 0 { return 0 }
47 if v >= NX_F32_LLM_N_VERDICTS { return 0 }
48 return 1
49}
50
51struct NxF32LlamaModel {
52 n_layers: nx_int,
53 hidden_dim: nx_int,
54 n_heads: nx_int,
55 n_kv_heads: nx_int,
56 head_dim: nx_int,
57 ffn_dim: nx_int,
58 vocab_size: nx_int,
59 embed_weights: *i64, // [vocab_size, hidden_dim]
60 layers: *i64, // array of n_layers *NxF32LlamaLayer (as i64)
61 gamma_out: *i64, // [hidden_dim]
62 lm_head: *i64, // [hidden_dim, vocab_size] (eager F32; the v3 path + v4 fallback)
63 lm_head_q8: i64 // v4: *NxF32LazyWeight (Q8_0) as i64, or 0 -> SIMD dequant-dot for lm_head
64}
65
66const NX_F32_LLAMA_MODEL_BYTES: nx_int = 96 // 12 fields * 8
67
68func nx_f32_llama_model_alloc() -> *NxF32LlamaModel {
69 return sys_mmap(NX_F32_LLAMA_MODEL_BYTES) as *NxF32LlamaModel
70}
71
72// Full model forward. 9 args.
73//
74// token_ids: i64 array (each value is a token id in [0, vocab_size)).
75// logits: [n_tokens, vocab_size] output.
76
77func nx_f32_llm_forward(
78 model: *NxF32LlamaModel,
79 token_ids: *i64,
80 n_tokens: nx_int,
81 cache: *NxF32KVCache,
82 eps: i64,
83 attn_scale: i64,
84 rope_log_base: i64,
85 apply_rope: nx_int,
86 logits: *i64) -> nx_int {
87
88 if model == (0 as *NxF32LlamaModel) { return NX_F32_LLM_ERR_NULL }
89 if token_ids == (0 as *i64) { return NX_F32_LLM_ERR_NULL }
90 if cache == (0 as *NxF32KVCache) { return NX_F32_LLM_ERR_NULL }
91 if logits == (0 as *i64) { return NX_F32_LLM_ERR_NULL }
92 if n_tokens <= 0 { return NX_F32_LLM_ERR_BAD_DIM }
93
94 let hidden_dim: nx_int = model.hidden_dim
95 let vocab_size: nx_int = model.vocab_size
96
97 // Scratch.
98 let x_embed: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
99 let x_stacked: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
100 let x_normed: *i64 = sys_mmap(n_tokens * hidden_dim * 8) as *i64
101
102 // 1) Embed lookup.
103 var t: nx_int = 0
104 while t < n_tokens {
105 let tok: nx_int = token_ids[t] as nx_int
106 if tok < 0 { return NX_F32_LLM_ERR_TOKEN }
107 if tok >= vocab_size { return NX_F32_LLM_ERR_TOKEN }
108 var d: nx_int = 0
109 while d < hidden_dim {
110 x_embed[t * hidden_dim + d] = model.embed_weights[tok * hidden_dim + d]
111 d = d + 1
112 }
113 t = t + 1
114 }
115
116 // 2) Stack forward.
117 let v_stk: nx_int = nx_f32_llama_stack_forward(
118 x_embed, x_stacked, n_tokens, model.n_layers,
119 hidden_dim, model.n_heads, model.n_kv_heads, model.head_dim, model.ffn_dim,
120 model.layers, cache, eps, attn_scale, rope_log_base, apply_rope)
121 if v_stk != NX_F32_STK_OK { return NX_F32_LLM_ERR_STACK }
122
123 // 3) Final RMSNorm per token.
124 var t2: nx_int = 0
125 while t2 < n_tokens {
126 let src: *i64 = (((x_stacked as i64) + t2 * hidden_dim * 8)) as *i64
127 let dst: *i64 = (((x_normed as i64) + t2 * hidden_dim * 8)) as *i64
128 nx_f32_rmsnorm(src, model.gamma_out, hidden_dim, eps, dst)
129 t2 = t2 + 1
130 }
131
132 // 4) LM head.
133 nx_f32_matmul(x_normed, model.lm_head, logits, n_tokens, hidden_dim, vocab_size)
134
135 return NX_F32_LLM_OK
136}
137
138// Convenience: return the i64 value (= sign-extended logit f32 bits) of the
139// argmax token id for the last token's logits. Greedy sampler.
140
141func nx_f32_llm_argmax_last(logits: *i64, n_tokens: nx_int,
142 vocab_size: nx_int) -> nx_int {
143 let base: nx_int = (n_tokens - 1) * vocab_size
144 var best_id: nx_int = 0
145 var best_val: i64 = logits[base + 0]
146 var i: nx_int = 1
147 while i < vocab_size {
148 let v: i64 = logits[base + i]
149 // f32 lexicographic comparison: for non-negative floats, raw int
150 // compare works. For negative, flip. Use nx_f32_lt for safety.
151 if nx_f32_lt(best_val, v) != 0 {
152 best_val = v
153 best_id = i
154 }
155 i = i + 1
156 }
157 return best_id
158}