nx_llm_run_v2.nx source

1// nx_llm_run_v2.nx -- end-to-end Llama-class inference composer (real 2// weights edition). 3// 4// L6 brick. The CLOSER for the inference path: composes every loader 5// + per-layer compute + sampler primitive shipped this session into a 6// single callable that maps: 7// 8// text prompt + GGUF + BPE vocab + sampling knobs -> next-token-id 9// 10// Compared to nx_llm_run.nx (v1, scaffold-only): 11// * v1 omitted the per-layer transformer compute (no weight loader) 12// * v1 simulated output projection as identity copy 13// * v2 sources ALL weights from a parsed GGUF and runs the full 14// forward stack 15// 16// Per the bits-up + no-skipping cardinals: no new math, no reinvented 17// kernels. Pure composition of: 18// 19// nx_bpe.nx_bpe_encode -- text -> token ids 20// nx_gguf_load_model.nx_..._weights -- top-level tensor bundle 21// nx_tensor.nx_t_alloc -- working buffers 22// nx_embedding.nx_embedding_lookup -- token-id -> hidden vector 23// nx_transformer_stack.nx_..._fwd -- N-layer forward 24// nx_blas_i64.nx_blas_matmul -- output projection to vocab 25// nx_token_sample.* + nx_attention.nx_attn_softmax_row_q10 26// -- temp + top_k + softmax + sample 27// 28// genealogy_id: standard_decoder_only_llm_inference + 29// touvron_2023_llama + radford_2019_gpt2 30// lineage_id: substrate_llm_run_v2_real_weights 31 32// nx_safety_envelope: 33// intended_use: "End-to-end Llama-class next-token generation 34// from a parsed GGUF; single-prompt scope; no 35// KV cache reuse across calls (v3 work); no 36// multi-head (single-head v1)" 37// sil_target: SIL2 38// asil_target: QM 39// dal_target: DAL C 40// evidence: [composes_only_shipped_canonical_bricks, 41// each_loader_smoke_green, stack_smoke_green] 42// hazard_register: [bug-tape-tied-embed-no-transpose, 43// bug-tape-vocab-mismatch-vs-bpe, 44// bug-tape-attn-scale-default, 45// bug-tape-overflow-output-projection] 46// residual_risk: "Tied-embedding GGUFs (Llama-2/3) are 47// REJECTED in v2 -- v3 adds matmul-with- 48// transpose-view to handle them" 49// verdict: NOT_YET_EVALUATED 50 51import "nx_syscalls.nx" 52import "nx_tier.nx" 53import "nx_loop.nx" 54import "nx_tensor.nx" 55import "nx_bpe.nx" 56import "nx_embedding.nx" 57import "nx_blas_i64.nx" 58import "nx_attention.nx" 59import "nx_token_sample.nx" 60import "nx_model_spec.nx" 61import "nx_prng.nx" 62import "nx_gguf.nx" 63import "nx_gguf_load.nx" 64import "nx_gguf_load_model.nx" 65import "nx_transformer_stack.nx" 66 67// ===== Sealed-enum: LlmRunV2Verdict ================================ 68 69const NX_LR2_OK: nx_int = 0 70const NX_LR2_ERR_BAD_SPEC: nx_int = 1 71const NX_LR2_ERR_TOKENIZE: nx_int = 2 72const NX_LR2_ERR_MODEL_LOAD: nx_int = 3 73const NX_LR2_ERR_EMBED: nx_int = 4 74const NX_LR2_ERR_STACK_FWD: nx_int = 5 75const NX_LR2_ERR_TIED_NOT_SUPPORTED: nx_int = 6 76const NX_LR2_ERR_OOM: nx_int = 7 77const NX_LR2_ERR_BAD_ROPE: nx_int = 8 78const NX_LR2_N_VERDICTS: nx_int = 9 79 80func nx_lr2_verdict_is_valid(v: nx_int) -> nx_int { 81 if v < 0 { return 0 } 82 if v >= NX_LR2_N_VERDICTS { return 0 } 83 return 1 84} 85 86// Compute 1/sqrt(d) in Q10 with d small (head_dim typically 64-128). 87// For our v1 single-head usage with tiny d (smoke uses d=2), we use 88// the substrate's existing integer-sqrt path. Output is Q10. 89// 90// 1/sqrt(d) Q10 = Q10 * Q10 / sqrt(d * Q20) approximately 91// 92// Simpler exact-when-d-small recipe used here: 93// sqrt_d_q10 = nx_isqrt(d * Q20) (NOT shipped under that name) -- 94// instead: for d=2, 1/sqrt(2) = 0.7071 -> Q10 = 724 95// for d=128, 1/sqrt(128) = 0.0884 -> Q10 = 91 96// 97// The substrate already has nx_isqrt; we use it iff available. 98// For this v2 brick, caller passes attn_scale_q10 directly so this 99// helper is not needed. Comment retained for the v3 derivation. 100 101// ===== Public: generate one next token ============================ 102// 103// Args (12 -- under 16-arg limit): 104// spec model architecture envelope 105// gguf_buf, hdr parsed GGUF (caller pre-parses) 106// bpe BPE vocab (caller pre-loads; v3 reads from GGUF metadata) 107// prompt_text, prompt_len UTF-8 prompt bytes 108// temperature_q10 sampling temperature (Q10; 1024 = 1.0) 109// top_k top-K mask (e.g. 40) 110// prng_state *i64 mutable PRNG state 111// rope_base RoPE theta base (10000 / 500000) 112// attn_scale_q10 attention scale 1/sqrt(head_dim) in Q10 113// 114// Returns: next token ID (>= 0) on success, or 0 - verdict_id on error. 115 116func nx_llm_generate_one_v2( 117 spec: *NxModelSpec, 118 gguf_buf: *u8, hdr: *NxGgufHeader, 119 bpe: *NxBpeVocab, 120 prompt_text: *u8, prompt_len: nx_int, 121 temperature_q10: nx_int, top_k: nx_int, 122 prng_state: *i64, 123 rope_base: nx_int, attn_scale_q10: nx_int) -> nx_int { 124 125 // Boundary null guards -- public-API runner must reject null inputs 126 // cleanly instead of SEGVing inside the encode/load steps. Composes 127 // with the four-pillar null-deref preventative scan. 128 if (spec as i64) == 0 { return 0 - NX_LR2_ERR_BAD_SPEC } 129 if (gguf_buf as i64) == 0 { return 0 - NX_LR2_ERR_MODEL_LOAD } 130 if (hdr as i64) == 0 { return 0 - NX_LR2_ERR_MODEL_LOAD } 131 if (bpe as i64) == 0 { return 0 - NX_LR2_ERR_TOKENIZE } 132 if (prompt_text as i64) == 0 { return 0 - NX_LR2_ERR_TOKENIZE } 133 if (prng_state as i64) == 0 { return 0 - NX_LR2_ERR_BAD_SPEC } 134 if prompt_len <= 0 { return 0 - NX_LR2_ERR_TOKENIZE } 135 if temperature_q10 <= 0 { return 0 - NX_LR2_ERR_BAD_SPEC } 136 if top_k <= 0 { return 0 - NX_LR2_ERR_BAD_SPEC } 137 if attn_scale_q10 <= 0 { return 0 - NX_LR2_ERR_BAD_SPEC } 138 139 let v_spec: nx_int = nx_model_spec_validate(spec) 140 if v_spec != NX_MS_OK { return 0 - NX_LR2_ERR_BAD_SPEC } 141 if rope_base <= 1 { return 0 - NX_LR2_ERR_BAD_ROPE } 142 143 // ----- Step 1: tokenize prompt ----- 144 let token_buf: *i64 = sys_mmap((prompt_len + 16) * 8) as *i64 145 let n_tokens: nx_int = nx_bpe_encode(bpe, prompt_text, prompt_len, token_buf) 146 if n_tokens <= 0 { return 0 - NX_LR2_ERR_TOKENIZE } 147 if n_tokens > spec.max_seq_len { return 0 - NX_LR2_ERR_BAD_SPEC } 148 149 // ----- Step 2: load model-level weights ----- 150 let mw: *NxGgufModelWeights = sys_mmap(NX_GML_BUNDLE_BYTES) as *NxGgufModelWeights 151 let err: *i64 = sys_mmap(8) as *i64 152 err[0] = 0 153 let v_ml: nx_int = nx_gguf_load_model_weights(gguf_buf, hdr, mw, err) 154 if v_ml != NX_GML_OK { return 0 - NX_LR2_ERR_MODEL_LOAD } 155 if mw.is_output_tied == 1 { 156 // v3 path adds transposed-view matmul. v2 rejects. 157 return 0 - NX_LR2_ERR_TIED_NOT_SUPPORTED 158 } 159 160 // ----- Step 3: embed ----- 161 let sh_x: *i64 = sys_mmap(2 * 8) as *i64 162 sh_x[0] = n_tokens; sh_x[1] = spec.hidden_dim 163 let x: *NxTensor = nx_t_alloc(NX_DT_I64, sh_x, 2, err) 164 if err[0] != 0 { return 0 - NX_LR2_ERR_OOM } 165 let v_emb: nx_int = nx_embedding_lookup(mw.token_embd, token_buf, n_tokens, x) 166 if v_emb != NX_EMB_OK { return 0 - NX_LR2_ERR_EMBED } 167 168 // ----- Step 4: positions [n_tokens] ----- 169 let positions: *i64 = sys_mmap(n_tokens * 8) as *i64 170 var p_i: nx_int = 0 171 while p_i < n_tokens { 172 positions[p_i] = p_i 173 p_i = p_i + 1 174 } 175 176 // ----- Step 5: transformer stack forward ----- 177 let v_st: nx_int = nx_transformer_stack_forward( 178 x, positions, gguf_buf, hdr, 179 spec.n_layers, spec.hidden_dim, spec.head_dim, spec.ffn_dim, 180 rope_base, attn_scale_q10, 181 mw.output_norm_gamma) 182 if v_st != NX_TS_OK { return 0 - NX_LR2_ERR_STACK_FWD } 183 184 // ----- Step 6: extract last-token hidden vector as a [1, hidden] tensor ----- 185 let sh_h: *i64 = sys_mmap(2 * 8) as *i64 186 sh_h[0] = 1; sh_h[1] = spec.hidden_dim 187 let hidden_last: *NxTensor = nx_t_alloc(NX_DT_I64, sh_h, 2, err) 188 if err[0] != 0 { return 0 - NX_LR2_ERR_OOM } 189 let phl: *i64 = hidden_last.storage as *i64 190 let px: *i64 = x.storage as *i64 191 let base: nx_int = (n_tokens - 1) * spec.hidden_dim 192 var c: nx_int = 0 193 while c < spec.hidden_dim { 194 phl[c] = px[base + c] 195 c = c + 1 196 } 197 198 // ----- Step 7: output projection -- matmul([1, hidden], [hidden, vocab]) ----- 199 let sh_l: *i64 = sys_mmap(2 * 8) as *i64 200 sh_l[0] = 1; sh_l[1] = spec.vocab_size 201 let logits_t: *NxTensor = nx_t_alloc(NX_DT_I64, sh_l, 2, err) 202 if err[0] != 0 { return 0 - NX_LR2_ERR_OOM } 203 nx_blas_matmul(hidden_last, mw.output_weight, logits_t) 204 205 let logits: *i64 = logits_t.storage as *i64 206 207 // ----- Step 8: temperature + top-K mask ----- 208 nx_logit_apply_temperature(logits, spec.vocab_size, temperature_q10) 209 nx_logit_top_k_mask(logits, spec.vocab_size, top_k) 210 211 // ----- Step 9: softmax (treat logits as [1, vocab] tensor) ----- 212 nx_attn_softmax_row_q10(logits_t) 213 214 // ----- Step 10: sample ----- 215 let next_token: nx_int = nx_sample_categorical(logits, spec.vocab_size, prng_state) 216 return next_token 217}