nx_attribution.nx source
↩ module page · 309 lines · 11857 B
1// nx_attribution.nx -- per-output bottom-up contribution tracking.
2//
3// User correction 2026-05-16: "your feedback is top down not bottom
4// up as well i dont know wtf the bits are doing and no system
5// exposes that thats the point of all the nishi lang no more black
6// boxes or secrets".
7//
8// THIS primitive is the bits-up answer. For linear/matmul-class
9// operations, the substrate EMITS what every contribution was --
10// signed, ranked, named-input -- so callers can ASK "which input
11// contributed most to this output?" and get an exact answer at
12// full bit precision.
13//
14// No black boxes. No hidden activations. No post-hoc inference
15// from outputs. The substrate exposes the bits.
16//
17// ===== Why this matters ===========================================
18//
19// Every existing ML stack stores tensors as opaque blocks. When
20// output[i] = 0.7, no one knows which input element was responsible.
21// Mechanistic-interpretability research (Olah, Conmy, Nanda 2022+;
22// Anthropic's circuits) reverse-engineers this MANUALLY for tiny
23// models. Production stacks ship without it.
24//
25// NishiLang substrate makes attribution a first-class primitive:
26// any matmul / linear / attention op can be called with attribution
27// recording on, and the substrate emits the (input_idx, contribution)
28// pairs that produced each output.
29//
30// ===== Math =======================================================
31//
32// For y = W @ x where y is [M], x is [N], W is [M, N]:
33//
34// y[i] = sum_{j=0..N} W[i, j] * x[j]
35//
36// Each (j, W[i, j] * x[j]) pair is one contribution. We store the
37// top-K largest-by-absolute-magnitude per output element. Signed:
38// contribution[i, k] = W[i, j*] * x[j*] where j* is the k-th rank.
39//
40// For attention output[i, d] = sum_j attn_weights[i, j] * V[j, d]:
41// Top-K attention tokens per output element.
42//
43// For activation-steered hidden h' = h + alpha * s:
44// Contribution[i, d] of steering = alpha * s[d] (the magnitude
45// added). Caller compares h vs h' to see which dims moved most.
46//
47// ===== Algorithm ==================================================
48//
49// Per output element y[i]:
50// 1. Initialise top-K array with -inf placeholder magnitudes.
51// 2. For each j in 0..N:
52// contrib = W[i, j] * x[j]
53// abs_contrib = |contrib|
54// If abs_contrib > top-K min: insertion-sort it in.
55// 3. Emit (j*, contrib) pairs in descending |contrib| order.
56//
57// O(M * N * K) compute; for typical M=N=4096, K=8 = ~135M ops per
58// matmul. Affordable for substrate-side post-hoc attribution
59// queries; for production inference the caller skips the
60// attribution path.
61//
62// ===== Per the cardinals ==========================================
63//
64// Bounded-loop + bits-up + sealed-verdict. Tracking is OPT-IN via
65// a separate entrypoint so production inference doesn't pay the
66// cost.
67//
68// genealogy_id: olah_2018_circuits + conmy_2023_acdc +
69// nanda_2023_progress_measures +
70// sundararajan_2017_integrated_gradients
71// lineage_id: substrate_attribution_v1
72
73// nx_safety_envelope:
74// intended_use: AUTO_APPLIED -- primitive-specific tuning queued
75// sil_target: SIL1
76// evidence: [bulk_applied_2026-05-16, see-file-comment-for-detail]
77// verdict: NOT_YET_EVALUATED
78
79import "nx_syscalls.nx"
80import "nx_tier.nx"
81import "nx_loop.nx"
82import "nx_tensor.nx"
83const NX_MAGIC_1024: i64 = 1024
84
85const NX_ATTR_NEG_INF: i64 = -9223372036854775807
86
87// ===== Sealed-enum: AttributionVerdict ============================
88
89const NX_ATTR_OK: nx_int = 0
90const NX_ATTR_ERR_BAD_DTYPE: nx_int = 1
91const NX_ATTR_ERR_BAD_NDIM: nx_int = 2
92const NX_ATTR_ERR_SHAPE_MISMATCH: nx_int = 3
93const NX_ATTR_ERR_BAD_K: nx_int = 4
94const NX_ATTR_N_VERDICTS: nx_int = 5
95
96func nx_attr_verdict_is_valid(v: nx_int) -> nx_int {
97 if v < 0 { return 0 }
98 if v >= NX_ATTR_N_VERDICTS { return 0 }
99 return 1
100}
101
102// ===== Helper: |x| via two's complement ==========================
103
104func _attr_abs(x: i64) -> i64 {
105 if x < 0 { return 0 - x }
106 return x
107}
108
109// ===== Matmul with top-K attribution ============================
110//
111// W: *NxTensor [M, N] weight matrix (Q10)
112// x: *NxTensor [N] input vector (Q10)
113// y: *NxTensor [M] output vector (Q10) -- written
114// top_k: K contributions to track per output
115// attr_idx: *i64 [M, K] top-K input indices per output -- written
116// attr_val: *i64 [M, K] their signed contributions (Q20) -- written
117//
118// Result: y[i] = sum_j W[i, j] * x[j] / Q10, and for each i the
119// top-K (j, W[i, j] * x[j]) pairs by absolute magnitude.
120
121func nx_attr_matmul_topk(W: *NxTensor, x: *NxTensor, y: *NxTensor,
122 top_k: nx_int, attr_idx: *i64, attr_val: *i64) -> nx_int {
123 if W.dtype != NX_DT_I64 { return NX_ATTR_ERR_BAD_DTYPE }
124 if x.dtype != NX_DT_I64 { return NX_ATTR_ERR_BAD_DTYPE }
125 if y.dtype != NX_DT_I64 { return NX_ATTR_ERR_BAD_DTYPE }
126 if W.ndim != 2 { return NX_ATTR_ERR_BAD_NDIM }
127 if x.ndim != 1 { return NX_ATTR_ERR_BAD_NDIM }
128 if y.ndim != 1 { return NX_ATTR_ERR_BAD_NDIM }
129 let M: nx_int = W.shape[0]
130 let N: nx_int = W.shape[1]
131 if x.shape[0] != N { return NX_ATTR_ERR_SHAPE_MISMATCH }
132 if y.shape[0] != M { return NX_ATTR_ERR_SHAPE_MISMATCH }
133 if nx_t_is_contiguous(W) == 0 { return NX_ATTR_ERR_SHAPE_MISMATCH }
134 if nx_t_is_contiguous(x) == 0 { return NX_ATTR_ERR_SHAPE_MISMATCH }
135 if nx_t_is_contiguous(y) == 0 { return NX_ATTR_ERR_SHAPE_MISMATCH }
136 if top_k <= 0 { return NX_ATTR_ERR_BAD_K }
137 if top_k > N { return NX_ATTR_ERR_BAD_K }
138
139 let pw: *i64 = W.storage as *i64
140 let px: *i64 = x.storage as *i64
141 let py: *i64 = y.storage as *i64
142
143 var i: nx_int = 0
144 var i_iter: nx_int = 0
145 var i_verdict: nx_int = NX_LOOP_RUNNING
146 let I_BUDGET: nx_int = M
147 while i_verdict == NX_LOOP_RUNNING && i_iter < I_BUDGET {
148 // Initialise this row's top-K to sentinel-min.
149 var k: nx_int = 0
150 while k < top_k {
151 attr_idx[i * top_k + k] = 0 - 1
152 attr_val[i * top_k + k] = 0
153 k = k + 1
154 }
155 // Track top-K by absolute magnitude. Store min-abs threshold
156 // for early-skip.
157 var min_abs_in_topk: i64 = 0 // anything beats the 0-sentinel initially
158
159 var acc: i64 = 0
160 var j: nx_int = 0
161 var j_iter: nx_int = 0
162 var j_verdict: nx_int = NX_LOOP_RUNNING
163 let J_BUDGET: nx_int = N
164 while j_verdict == NX_LOOP_RUNNING && j_iter < J_BUDGET {
165 let w_ij: i64 = pw[i * N + j]
166 let x_j: i64 = px[j]
167 let contrib: i64 = w_ij * x_j
168 acc = acc + contrib
169
170 // Top-K insertion if abs(contrib) > current min.
171 let abs_c: i64 = _attr_abs(contrib)
172 if abs_c > min_abs_in_topk {
173 // Find insertion position from the end.
174 var pos: nx_int = top_k - 1
175 var shift_iter: nx_int = 0
176 var shift_verdict: nx_int = NX_LOOP_RUNNING
177 while shift_verdict == NX_LOOP_RUNNING && shift_iter < top_k {
178 if pos > 0 {
179 let prev_abs: i64 = _attr_abs(attr_val[i * top_k + pos - 1])
180 if prev_abs < abs_c {
181 attr_idx[i * top_k + pos] = attr_idx[i * top_k + pos - 1]
182 attr_val[i * top_k + pos] = attr_val[i * top_k + pos - 1]
183 pos = pos - 1
184 } else {
185 shift_verdict = NX_LOOP_DONE_EXIT
186 }
187 } else {
188 shift_verdict = NX_LOOP_DONE_EXIT
189 }
190 shift_iter = shift_iter + 1
191 }
192 attr_idx[i * top_k + pos] = j
193 attr_val[i * top_k + pos] = contrib
194
195 // Update min-abs threshold from the new last slot.
196 let last_idx: nx_int = i * top_k + top_k - 1
197 if attr_idx[last_idx] >= 0 {
198 min_abs_in_topk = _attr_abs(attr_val[last_idx])
199 }
200 }
201 j = j + 1
202 j_iter = j_iter + 1
203 }
204 py[i] = acc / NX_MAGIC_1024 // Q10
205 i = i + 1
206 i_iter = i_iter + 1
207 }
208 return NX_ATTR_OK
209}
210
211// ===== Convenience: explain one output ===========================
212//
213// For output index `i`, fills out_indices + out_signed_contribs
214// with the top-K contributions in descending |contrib| order.
215// Caller has already run attr_matmul_topk; this is just a read
216// helper.
217
218func nx_attr_explain(attr_idx: *i64, attr_val: *i64,
219 output_i: nx_int, top_k: nx_int,
220 out_indices: *i64, out_contribs: *i64) -> nx_int {
221 var k: nx_int = 0
222 while k < top_k {
223 out_indices[k] = attr_idx[output_i * top_k + k]
224 out_contribs[k] = attr_val[output_i * top_k + k]
225 k = k + 1
226 }
227 return NX_ATTR_OK
228}
229
230// ===== Self-test ==================================================
231//
232// Tiny matmul: W [3, 4], x [4], y [3].
233//
234// W = [[10, 1, 100, 2],
235// [ 5, 50, 1, 3],
236// [ 1, 1, 1, 77]]
237// x = [1, 1, 1, 1] (Q10 = 1024 each for cleanness; we use raw
238// integers here for easier verification)
239//
240// y[0] = 10 + 1 + 100 + 2 = 113. Top contributor: W[0,2]*x[2] = 100.
241// y[1] = 5 + 50 + 1 + 3 = 59. Top contributor: W[1,1]*x[1] = 50.
242// y[2] = 1 + 1 + 1 + 77 = 80. Top contributor: W[2,3]*x[3] = 77.
243//
244// Top-K with K=2 should rank these correctly per row.
245//
246// Closed-form invariants:
247// (a) Attr-idx[0, 0] == 2 (column with the 100 contribution)
248// (b) Attr-val[0, 0] == 100
249// (c) Attr-idx[0, 1] == 0 (next biggest: 10)
250// (d) Attr-idx[1, 0] == 1 (column with 50)
251// (e) Attr-idx[2, 0] == 3 (column with 77)
252// (f) Verdict gate
253
254func main() -> i64 {
255 let M: nx_int = 3
256 let N: nx_int = 4
257
258 let w_sh: *nx_int = sys_mmap(2 * 8) as *nx_int
259 w_sh[0] = M; w_sh[1] = N
260 let x_sh: *nx_int = sys_mmap(1 * 8) as *nx_int
261 x_sh[0] = N
262 let y_sh: *nx_int = sys_mmap(1 * 8) as *nx_int
263 y_sh[0] = M
264
265 let err: *nx_int = sys_mmap(8) as *nx_int
266 err[0] = 0
267 let W: *NxTensor = nx_t_alloc(NX_DT_I64, w_sh, 2, err)
268 let x: *NxTensor = nx_t_alloc(NX_DT_I64, x_sh, 1, err)
269 let y: *NxTensor = nx_t_alloc(NX_DT_I64, y_sh, 1, err)
270 if err[0] != 0 { return 5 }
271
272 let pw: *i64 = W.storage as *i64
273 let px: *i64 = x.storage as *i64
274 pw[0] = 10; pw[1] = 1; pw[2] = 100; pw[3] = 2
275 pw[4] = 5; pw[5] = 50; pw[6] = 1; pw[7] = 3
276 pw[8] = 1; pw[9] = 1; pw[10] = 1; pw[11] = 77
277 px[0] = 1; px[1] = 1; px[2] = 1; px[3] = 1
278
279 let TOP_K: nx_int = 2
280 let attr_idx: *i64 = sys_mmap(M * TOP_K * 8) as *i64
281 let attr_val: *i64 = sys_mmap(M * TOP_K * 8) as *i64
282
283 let v: nx_int = nx_attr_matmul_topk(W, x, y, TOP_K, attr_idx, attr_val)
284 if v != NX_ATTR_OK { return 10 + v }
285
286 // Row 0 top contributor: column 2, contribution 100.
287 if attr_idx[0 * TOP_K + 0] != 2 { return 20 }
288 if attr_val[0 * TOP_K + 0] != 100 { return 21 }
289 // Row 0 second: column 0, contribution 10.
290 if attr_idx[0 * TOP_K + 1] != 0 { return 22 }
291 if attr_val[0 * TOP_K + 1] != 10 { return 23 }
292
293 // Row 1 top: column 1, contribution 50.
294 if attr_idx[1 * TOP_K + 0] != 1 { return 30 }
295 if attr_val[1 * TOP_K + 0] != 50 { return 31 }
296
297 // Row 2 top: column 3, contribution 77.
298 if attr_idx[2 * TOP_K + 0] != 3 { return 40 }
299 if attr_val[2 * TOP_K + 0] != 77 { return 41 }
300
301 // Verdict gate.
302 var vi: nx_int = 0
303 while vi < NX_ATTR_N_VERDICTS {
304 if nx_attr_verdict_is_valid(vi) != 1 { return 50 + vi }
305 vi = vi + 1
306 }
307
308 return 0
309}