nx_f32_attn_multi.nx source
↩ module page · 137 lines · 4409 B
1// nx_f32_attn_multi.nx -- bits-up f32 multi-token attention kernel.
2//
3// L7 / L8 composition brick. The prerequisite for both prefill
4// (processing a whole prompt) and KV cache decode (attending to
5// growing cache).
6//
7// Composes nx_f32_dot + nx_f32_mul + nx_f32_softmax + nx_f32_add.
8// No libm.
9//
10// Algorithm:
11//
12// Inputs:
13// Q [n_tokens_q, head_dim] query tensor
14// K [n_tokens_k, head_dim] key tensor (may be larger when
15// attending to a KV cache)
16// V [n_tokens_k, head_dim] value tensor
17//
18// For each query token i in 0..n_tokens_q:
19// For each key token j in 0..n_tokens_k:
20// scores[j] = dot(Q[i,:], K[j,:]) * attn_scale
21// (if causal: positions j > q_pos_i are masked with -inf)
22// probs = softmax(scores)
23// For each output dim d:
24// out[i, d] = sum_j probs[j] * V[j, d]
25//
26// For causal LM during prefill: q_pos_i = i (query i can attend
27// to keys 0..=i).
28// For causal LM during KV-cache decode: q_pos_i = n_tokens_k -
29// n_tokens_q + i (query at end of cache, attends to all prior).
30//
31// K^T transpose is handled by index arithmetic (we access K row-by-
32// row in its natural [n_tokens_k, head_dim] layout); no physical
33// transpose needed.
34//
35// genealogy_id: vaswani_2017_attention + causal_mask_standard
36// lineage_id: substrate_f32_attn_multi_v1
37
38import "nx_syscalls.nx"
39import "nx_tier.nx"
40import "nx_f32.nx"
41import "nx_f32_matmul.nx"
42import "nx_f32_softmax.nx"
43
44const NX_F32_AM_OK: nx_int = 0
45const NX_F32_AM_ERR_BAD_DIM: nx_int = 1
46const NX_F32_AM_N_VERDICTS: nx_int = 2
47
48func nx_f32_am_verdict_is_valid(v: nx_int) -> nx_int {
49 if v < 0 { return 0 }
50 if v >= NX_F32_AM_N_VERDICTS { return 0 }
51 return 1
52}
53
54// Multi-token attention kernel.
55//
56// causal: 1 = apply causal mask (q_pos_i = n_tokens_k - n_tokens_q + i)
57// 0 = full attention (every query attends to every key)
58// attn_scale: 1/sqrt(head_dim) f32 raw bits
59//
60// out shape matches Q shape: [n_tokens_q, head_dim].
61
62func nx_f32_attn_multi(
63 Q: *i64, K: *i64, V: *i64,
64 n_tokens_q: nx_int,
65 n_tokens_k: nx_int,
66 head_dim: nx_int,
67 causal: nx_int,
68 attn_scale: i64,
69 out: *i64) -> nx_int {
70
71 if n_tokens_q <= 0 { return NX_F32_AM_ERR_BAD_DIM }
72 if n_tokens_k <= 0 { return NX_F32_AM_ERR_BAD_DIM }
73 if head_dim <= 0 { return NX_F32_AM_ERR_BAD_DIM }
74 if n_tokens_q > n_tokens_k {
75 return NX_F32_AM_ERR_BAD_DIM
76 }
77
78 let scores: *i64 = sys_mmap(n_tokens_k * 8) as *i64
79 let probs: *i64 = sys_mmap(n_tokens_k * 8) as *i64
80
81 // -inf as f32 raw bits for masking.
82 let NEG_INF: i64 = 0xFF800000
83
84 var i: nx_int = 0
85 while i < n_tokens_q {
86 // Cache position of this query (for causal mask).
87 let q_pos: nx_int = n_tokens_k - n_tokens_q + i
88
89 let Q_i_base: i64 = (Q as i64) + i * head_dim * 8
90 let Q_i: *i64 = Q_i_base as *i64
91
92 // Compute scores[0..n_tokens_k]
93 var j: nx_int = 0
94 while j < n_tokens_k {
95 let K_j_base: i64 = (K as i64) + j * head_dim * 8
96 let K_j: *i64 = K_j_base as *i64
97
98 var masked: nx_int = 0
99 if causal == 1 {
100 if j > q_pos { masked = 1 }
101 }
102 if masked == 1 {
103 scores[j] = NEG_INF
104 } else {
105 let raw: i64 = nx_f32_dot(Q_i, K_j, head_dim)
106 scores[j] = __f32_mul(raw, attn_scale)
107 }
108 j = j + 1
109 }
110
111 // Softmax over scores[0..n_tokens_k]
112 nx_f32_softmax(scores, n_tokens_k, probs)
113
114 // out[i, d] = sum_j probs[j] * V[j, d]
115 let out_i_base: i64 = (out as i64) + i * head_dim * 8
116 let out_i: *i64 = out_i_base as *i64
117
118 var d: nx_int = 0
119 while d < head_dim {
120 var acc: i64 = 0
121 var k: nx_int = 0
122 while k < n_tokens_k {
123 let V_k_base: i64 = (V as i64) + k * head_dim * 8
124 let V_k: *i64 = V_k_base as *i64
125 let p: i64 = probs[k]
126 let v: i64 = V_k[d]
127 acc = __f32_add(acc, __f32_mul(p, v)) // hw SSE, bit-identical (2026-07-10)
128 k = k + 1
129 }
130 out_i[d] = acc
131 d = d + 1
132 }
133
134 i = i + 1
135 }
136 return NX_F32_AM_OK
137}