nx_f32_attn_cached.nx source
↩ module page · 151 lines · 5729 B
1// nx_f32_attn_cached.nx -- bits-up f32 multi-head attention with KV cache.
2//
3// L7 / L8 integration brick. Closes the substrate-side gap between
4// the math primitives (RMSNorm + matmul + softmax + ...) and the
5// real LLM inference pattern (prefill once, then decode token-by-
6// token attending to the growing cache).
7//
8// Caller responsibility:
9// 1. Project x -> Q, K_new, V_new via matmul(x, W_q/k/v).
10// 2. Optionally apply RoPE to each head slice (per token position
11// cache_seq_before + i) via nx_f32_rope_apply_vector.
12// 3. Call this kernel: appends K_new/V_new to cache, attends over
13// full history.
14// 4. Project the concatenated attention output via matmul(., W_o)
15// to get attn_proj.
16// 5. Residual: x = x + attn_proj.
17//
18// Algorithm:
19// 1. Append K_new, V_new to cache at layer_idx. (Caller has NOT
20// yet advanced cache.seq_len.)
21// 2. Total key/value count visible to this layer: total_k =
22// cache.seq_len + n_tokens (the new rows are now in cache at
23// positions seq_len..seq_len+n_tokens-1 ready to be read).
24// 3. For each query head h in 0..n_heads:
25// kv_head = h / group_size
26// Slice Q_h from Q at columns [h*head_dim, (h+1)*head_dim]
27// Read K_all = cache.K[layer_idx] (first total_k rows valid)
28// Slice K_h from K_all (per-token slicing over kv_head's stripe)
29// Same for V_h.
30// attn_h = nx_f32_attn_multi(Q_h, K_h, V_h, n_tokens, total_k,
31// head_dim, causal, attn_scale)
32// Scatter attn_h to attn_concat[:, h*head_dim:(h+1)*head_dim].
33// 4. The caller advances cache.seq_len AFTER all layers done for
34// this forward pass via nx_f32_kv_cache_advance.
35//
36// Returns: NX_F32_AC_OK on success, error verdicts otherwise.
37//
38// genealogy_id: standard_kv_cache_decode_pattern + mha + gqa
39// lineage_id: substrate_f32_attn_cached_v1
40
41import "nx_syscalls.nx"
42import "nx_tier.nx"
43import "nx_f32.nx"
44import "nx_f32_attn_multi.nx"
45import "nx_f32_kv_cache.nx"
46
47const NX_F32_AC_OK: nx_int = 0
48const NX_F32_AC_ERR_BAD_DIM: nx_int = 1
49const NX_F32_AC_ERR_BAD_GQA: nx_int = 2
50const NX_F32_AC_ERR_CACHE: nx_int = 3
51const NX_F32_AC_N_VERDICTS: nx_int = 4
52
53func nx_f32_ac_verdict_is_valid(v: nx_int) -> nx_int {
54 if v < 0 { return 0 }
55 if v >= NX_F32_AC_N_VERDICTS { return 0 }
56 return 1
57}
58
59// 13 args -- under the 16-arg limit.
60
61func nx_f32_attn_with_cache(
62 Q: *i64, K_new: *i64, V_new: *i64,
63 n_tokens: nx_int,
64 n_heads: nx_int,
65 n_kv_heads: nx_int,
66 head_dim: nx_int,
67 cache: *NxF32KVCache,
68 layer_idx: nx_int,
69 causal: nx_int,
70 attn_scale: i64,
71 attn_concat: *i64) -> nx_int {
72
73 if n_tokens <= 0 { return NX_F32_AC_ERR_BAD_DIM }
74 if n_heads <= 0 { return NX_F32_AC_ERR_BAD_DIM }
75 if n_kv_heads <= 0 { return NX_F32_AC_ERR_BAD_DIM }
76 if head_dim <= 0 { return NX_F32_AC_ERR_BAD_DIM }
77 if n_heads - (n_heads / n_kv_heads) * n_kv_heads != 0 {
78 return NX_F32_AC_ERR_BAD_GQA
79 }
80 if cache == (0 as *NxF32KVCache) { return NX_F32_AC_ERR_CACHE }
81
82 let q_dim: nx_int = n_heads * head_dim
83 let kv_dim: nx_int = n_kv_heads * head_dim
84 let group_size: nx_int = n_heads / n_kv_heads
85
86 // 1) Append K_new, V_new to cache.
87 let v_app: nx_int = nx_f32_kv_cache_append_layer(cache, layer_idx,
88 K_new, V_new, n_tokens)
89 if v_app != NX_F32_KVC_OK { return NX_F32_AC_ERR_CACHE }
90
91 // 2) Total key/value count for this attention call.
92 let cache_before: nx_int = nx_f32_kv_cache_get_seq_len(cache)
93 let total_k: nx_int = cache_before + n_tokens
94
95 let K_layer: *i64 = nx_f32_kv_cache_get_K_layer(cache, layer_idx)
96 let V_layer: *i64 = nx_f32_kv_cache_get_V_layer(cache, layer_idx)
97
98 // 3) Per-head attention.
99 let Q_h: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
100 let K_h: *i64 = sys_mmap(total_k * head_dim * 8) as *i64
101 let V_h: *i64 = sys_mmap(total_k * head_dim * 8) as *i64
102 let attn_h: *i64 = sys_mmap(n_tokens * head_dim * 8) as *i64
103
104 var h: nx_int = 0
105 while h < n_heads {
106 let kv_head: nx_int = h / group_size
107
108 // Gather Q_h: query columns [h*head_dim, (h+1)*head_dim] over n_tokens rows.
109 var t: nx_int = 0
110 while t < n_tokens {
111 var d: nx_int = 0
112 while d < head_dim {
113 Q_h[t * head_dim + d] = Q[t * q_dim + h * head_dim + d]
114 d = d + 1
115 }
116 t = t + 1
117 }
118
119 // Gather K_h, V_h from the cache layer: kv_head's stripe over total_k rows.
120 var t2: nx_int = 0
121 while t2 < total_k {
122 var d2: nx_int = 0
123 while d2 < head_dim {
124 K_h[t2 * head_dim + d2] = K_layer[t2 * kv_dim + kv_head * head_dim + d2]
125 V_h[t2 * head_dim + d2] = V_layer[t2 * kv_dim + kv_head * head_dim + d2]
126 d2 = d2 + 1
127 }
128 t2 = t2 + 1
129 }
130
131 // Multi-token attention. Note n_tokens_q = n_tokens, n_tokens_k = total_k
132 // (q_pos for last new token = total_k - 1; causal mask handles it).
133 nx_f32_attn_multi(Q_h, K_h, V_h, n_tokens, total_k, head_dim,
134 causal, attn_scale, attn_h)
135
136 // Scatter into attn_concat at columns [h*head_dim, (h+1)*head_dim].
137 var t3: nx_int = 0
138 while t3 < n_tokens {
139 var d3: nx_int = 0
140 while d3 < head_dim {
141 attn_concat[t3 * q_dim + h * head_dim + d3] = attn_h[t3 * head_dim + d3]
142 d3 = d3 + 1
143 }
144 t3 = t3 + 1
145 }
146
147 h = h + 1
148 }
149
150 return NX_F32_AC_OK
151}