nx_f32_llama_stack.nx source
↩ module page · 121 lines · 4322 B
1// nx_f32_llama_stack.nx -- multi-block transformer stack forward.
2//
3// Calls nx_f32_llama_block_forward_v3 N times across the layers
4// array, then advances the KV cache ONCE per forward pass.
5//
6// Double-buffered execution: alternates between buf_a and buf_b
7// to avoid allocating per-layer scratch. The final result lands
8// in x_out (which may NOT alias x_in).
9//
10// Layers array: layers[i] is the i64 value of a *NxF32LlamaLayer.
11// Caller is responsible for allocating + populating each layer.
12//
13// genealogy_id: standard_transformer_stack
14// lineage_id: substrate_f32_llama_stack_v1
15
16import "nx_syscalls.nx"
17import "nx_tier.nx"
18import "nx_f32.nx"
19import "nx_f32_kv_cache.nx"
20import "nx_f32_llama_block.nx"
21
22const NX_F32_STK_OK: nx_int = 0
23const NX_F32_STK_ERR_BAD_DIM: nx_int = 1
24const NX_F32_STK_ERR_NULL: nx_int = 2
25const NX_F32_STK_ERR_LAYER: nx_int = 3
26const NX_F32_STK_N_VERDICTS: nx_int = 4
27
28func nx_f32_stk_verdict_is_valid(v: nx_int) -> nx_int {
29 if v < 0 { return 0 }
30 if v >= NX_F32_STK_N_VERDICTS { return 0 }
31 return 1
32}
33
34// 15 args -- at the 16-arg limit.
35
36func nx_f32_llama_stack_forward(
37 x_in: *i64,
38 x_out: *i64,
39 n_tokens: nx_int,
40 n_layers: nx_int,
41 hidden_dim: nx_int,
42 n_heads: nx_int,
43 n_kv_heads: nx_int,
44 head_dim: nx_int,
45 ffn_dim: nx_int,
46 layers: *i64,
47 cache: *NxF32KVCache,
48 eps: i64,
49 attn_scale: i64,
50 rope_log_base: i64,
51 apply_rope: nx_int) -> nx_int {
52
53 if n_tokens <= 0 { return NX_F32_STK_ERR_BAD_DIM }
54 if n_layers <= 0 { return NX_F32_STK_ERR_BAD_DIM }
55 if hidden_dim <= 0 { return NX_F32_STK_ERR_BAD_DIM }
56 if x_in == (0 as *i64) { return NX_F32_STK_ERR_NULL }
57 if x_out == (0 as *i64) { return NX_F32_STK_ERR_NULL }
58 if layers == (0 as *i64) { return NX_F32_STK_ERR_NULL }
59 if cache == (0 as *NxF32KVCache) { return NX_F32_STK_ERR_NULL }
60
61 let n_elem: nx_int = n_tokens * hidden_dim
62 let buf_a: *i64 = sys_mmap(n_elem * 8) as *i64
63 let buf_b: *i64 = sys_mmap(n_elem * 8) as *i64
64
65 // Copy x_in -> buf_a.
66 var i: nx_int = 0
67 while i < n_elem {
68 buf_a[i] = x_in[i]
69 i = i + 1
70 }
71
72 // For each layer: forward through block. Source = buf_a, dest = buf_b,
73 // then swap. After n_layers, the final result is in buf_a if n_layers
74 // is even, buf_b if odd.
75 var L: nx_int = 0
76 var use_a_as_src: nx_int = 1 // 1: src=a, dst=b 0: src=b, dst=a
77 while L < n_layers {
78 let layer_ptr_i64: i64 = layers[L]
79 let layer: *NxF32LlamaLayer = layer_ptr_i64 as *NxF32LlamaLayer
80 if layer == (0 as *NxF32LlamaLayer) { return NX_F32_STK_ERR_LAYER }
81
82 if use_a_as_src != 0 {
83 let v: nx_int = nx_f32_llama_block_forward_v3(
84 buf_a, n_tokens, hidden_dim, n_heads, n_kv_heads, head_dim, ffn_dim,
85 layer, cache, L, eps, attn_scale, rope_log_base, apply_rope, buf_b)
86 if v != NX_F32_BLK_OK { return NX_F32_STK_ERR_LAYER }
87 use_a_as_src = 0
88 } else {
89 let v: nx_int = nx_f32_llama_block_forward_v3(
90 buf_b, n_tokens, hidden_dim, n_heads, n_kv_heads, head_dim, ffn_dim,
91 layer, cache, L, eps, attn_scale, rope_log_base, apply_rope, buf_a)
92 if v != NX_F32_BLK_OK { return NX_F32_STK_ERR_LAYER }
93 use_a_as_src = 1
94 }
95 L = L + 1
96 }
97
98 // Advance cache.seq_len by n_tokens (caller-relative -- once per stack call).
99 nx_f32_kv_cache_advance(cache, n_tokens)
100
101 // Copy final result to x_out. After the last iter, the SOURCE buffer is
102 // whichever was written last (the buffer pointed to by NOT use_a_as_src).
103 // After L=0: src=a, dst=b, output in b, use_a_as_src=0
104 // After L=1: src=b, dst=a, output in a, use_a_as_src=1
105 // So the final output is in buf_a if use_a_as_src==1, else buf_b.
106 if use_a_as_src != 0 {
107 var k: nx_int = 0
108 while k < n_elem {
109 x_out[k] = buf_a[k]
110 k = k + 1
111 }
112 } else {
113 var k2: nx_int = 0
114 while k2 < n_elem {
115 x_out[k2] = buf_b[k2]
116 k2 = k2 + 1
117 }
118 }
119
120 return NX_F32_STK_OK
121}