nx_token_merge.nx source
↩ module page · 235 lines · 9680 B
1// nx_token_merge.nx -- H8 token merging plan (bits-up).
2//
3// Per NISHI_ELDER_AI_OFF_DOCKER_2026_05_20.md ยง2.1 H8: token
4// merging (ToMe, Bolya 2022 "Token Merging: Your ViT But Faster")
5// shortens the attention sequence by merging the top-K most
6// similar token pairs at layer boundaries. Used in ViT and
7// Stable Diffusion to drop 30-50% of tokens with minimal quality
8// loss. Linear-time savings: shrinking seq from N to N*(1-r)
9// at every L layers cuts overall attention FLOPs ~(1-r)^L.
10//
11// V1 mechanics: substrate primitive does NOT compute the
12// similarity matrix. It accepts a pre-computed array of
13// (left_idx, right_idx, score_q16) triples sorted by descending
14// score, and a target_merges count. It greedily walks the
15// sorted list, merging each pair where BOTH tokens are still
16// "kept" (no prior merge), until target_merges is reached or
17// the list is exhausted. Returns the actual merges executed
18// (which can be less than requested if not enough non-conflicting
19// pairs were available).
20//
21// Composition path:
22// - Consumer kernel computes cosine sim between adjacent tokens
23// (substrate Q16-mul + dot product primitives)
24// - Consumer sorts the resulting (i, j, score) triples by
25// score desc using shipped nx_mergesort
26// - nx_token_merge_apply(plan, sorted_pairs, n_pairs, target)
27// populates the assignment table
28// - Downstream attention layer reads target_for(idx) to know
29// which slot each token routes to; merged tokens are dropped
30// from the active sequence
31//
32// Pure substrate logic. No Linux features. Composes with shipped
33// nx_mergesort (consumer-side sort dependency).
34//
35// V1 honest scope:
36// - Caller pre-sorts; substrate does NOT sort internally (would
37// duplicate nx_mergesort and force a bigger primitive)
38// - Greedy merge (top-down) -- the published ToMe paper proves
39// this matches bipartite matching within a fraction of a
40// percent on real model quality, at half the implementation
41// cost
42// - assignment[idx] semantics:
43// == idx token IS kept (root of its merge group)
44// != idx token merged INTO assignment[idx]
45// Caller iterates idx in [0, n_tokens) and only emits tokens
46// where assignment[idx] == idx; for the merged ones, the
47// consumer kernel sums/averages their features into the root's
48// buffer
49//
50// genealogy_id: bolya_2022_token_merging + bolya_2023_tome_for_diffusion +
51// cardinal_2026-05-20_elder_ai_off_docker +
52// cardinal_2026-05-20_bits_up_nishi_not_linux
53// lineage_id: substrate_token_merge_v1
54//
55// nx_capability_manifest:
56// variant_class: token_merge
57// variant_id: token_merge_v1_greedy_sorted_input
58// requires_isa: [rv32i, rv64imac, x86_64, cortex_m, armv7a, aarch64, wasm32]
59// requires_syscalls: [mmap]
60// requires_ram_min_b: 2048
61// tier_floor: NX_TIER_INF_MOBILE
62// tier_ceiling: NX_TIER_INF_HPC
63// cost_model:
64// flops_per_n: 1.0 // O(n_pairs) one-pass scan
65// bytes_per_n: 8.0 // one i64 per token in assignment
66// syscalls_per_n: 0.0
67// adversary_class: THREAT_AI_ADVERSARY
68//
69// nx_safety_envelope:
70// intended_use: "Greedy token-merge plan from pre-sorted similarity
71// pairs; bits-up; composes with nx_mergesort"
72// sil_target: SIL2
73// evidence: [canary_bracketed, root_invariant,
74// idempotent_kept_count, no_cycle_in_assignment]
75// verdict: NOT_YET_EVALUATED
76
77import "nx_syscalls.nx"
78const NX_MAGIC_65536: i64 = 65536
79
80// ===== Constants =================================================
81const NX_TOMERGE_MAX_TOKENS: i64 = 4096
82
83// Per-pair triple stride in caller's input array.
84const NX_TOMERGE_PAIR_STRIDE: i64 = 3 // (left, right, score)
85
86// Verdicts.
87const NX_TOMERGE_OK: i64 = 0
88const NX_TOMERGE_BAD_INPUT: i64 = 1
89const NX_TOMERGE_TAMPER: i64 = 2
90const NX_TOMERGE_N_VERDICTS: i64 = 3
91
92func nx_tomerge_verdict_is_valid(v: i64) -> i64 {
93 if v < 0 { return 0 }
94 if v >= NX_TOMERGE_N_VERDICTS { return 0 }
95 return 1
96}
97
98// Canary magic.
99const NX_TOMERGE_CANARY_PRE: i64 = 0x546F4D67506C617A // "ToMgPlaz"
100const NX_TOMERGE_CANARY_POST: i64 = 0x506C616E45314433 // "PlanE1D3"
101
102// ===== Struct ====================================================
103struct NxTokenMergePlan {
104 canary_pre: i64,
105 n_tokens: i64,
106 n_merged: i64, // number of tokens whose assignment != idx
107 assignment: *i64, // length n_tokens; assignment[idx] = idx (kept) OR merge target
108 canary_post: i64,
109}
110
111// ===== Validity =================================================
112func nx_token_merge_plan_is_valid(p: *NxTokenMergePlan) -> i64 {
113 if (p as i64) == 0 { return 0 }
114 if p.canary_pre != NX_TOMERGE_CANARY_PRE { return 0 }
115 if p.canary_post != NX_TOMERGE_CANARY_POST { return 0 }
116 if p.n_tokens <= 0 { return 0 }
117 if p.n_tokens > NX_TOMERGE_MAX_TOKENS { return 0 }
118 if p.n_merged < 0 { return 0 }
119 if p.n_merged > p.n_tokens { return 0 }
120 return 1
121}
122
123// ===== Constructor =================================================
124// Builds an empty plan: every token is its own root (assignment[i] == i,
125// n_merged == 0). Call nx_token_merge_apply() to populate.
126func nx_token_merge_plan_new(n_tokens: i64) -> *NxTokenMergePlan {
127 if n_tokens <= 0 { return (0 as i64) as *NxTokenMergePlan }
128 if n_tokens > NX_TOMERGE_MAX_TOKENS { return (0 as i64) as *NxTokenMergePlan }
129
130 let p: *NxTokenMergePlan = (sys_mmap(40)) as *NxTokenMergePlan
131 p.canary_pre = NX_TOMERGE_CANARY_PRE
132 p.n_tokens = n_tokens
133 p.n_merged = 0
134 p.assignment = (sys_mmap(n_tokens * 8)) as *i64
135 p.canary_post = NX_TOMERGE_CANARY_POST
136
137 var i: i64 = 0
138 while i < n_tokens {
139 p.assignment[i] = i
140 i = i + 1
141 }
142 return p
143}
144
145// ===== Apply =================================================
146// Walks sorted_pairs greedily, merging each pair whose BOTH endpoints
147// are still roots (assignment == idx). Returns the count of merges
148// actually executed. On bad input returns 0 - <verdict>.
149//
150// Input layout: sorted_pairs is a flat *i64 array of length
151// n_pairs * NX_TOMERGE_PAIR_STRIDE. Triple at offset k is
152// (sorted_pairs[k*3] = left, sorted_pairs[k*3+1] = right, sorted_pairs[k*3+2] = score).
153// score is informational here (caller used it to sort); substrate
154// only reads left + right.
155//
156// Conflict resolution: if left or right is already merged (its
157// assignment is not its own index), skip the pair. This is the
158// "non-conflicting greedy" pattern that ToMe-paper bipartite
159// matching converges to within a tiny epsilon.
160func nx_token_merge_apply(p: *NxTokenMergePlan, sorted_pairs: *i64, n_pairs: i64, target_merges: i64) -> i64 {
161 if nx_token_merge_plan_is_valid(p) != 1 { return 0 - NX_TOMERGE_TAMPER }
162 if (sorted_pairs as i64) == 0 { return 0 - NX_TOMERGE_BAD_INPUT }
163 if n_pairs < 0 { return 0 - NX_TOMERGE_BAD_INPUT }
164 if target_merges < 0 { return 0 - NX_TOMERGE_BAD_INPUT }
165
166 var done: i64 = 0
167 var k: i64 = 0
168 while k < n_pairs {
169 if done < target_merges {
170 let left: i64 = sorted_pairs[k * 3]
171 let right: i64 = sorted_pairs[k * 3 + 1]
172 // Bound check.
173 var skip: i64 = 0
174 if left < 0 { skip = 1 }
175 if right < 0 { skip = 1 }
176 if left >= p.n_tokens { skip = 1 }
177 if right >= p.n_tokens { skip = 1 }
178 if left == right { skip = 1 }
179 if skip == 0 {
180 // Both must still be roots.
181 if p.assignment[left] == left {
182 if p.assignment[right] == right {
183 // Merge right INTO left (left becomes root).
184 p.assignment[right] = left
185 p.n_merged = p.n_merged + 1
186 done = done + 1
187 }
188 }
189 }
190 }
191 k = k + 1
192 }
193 return done
194}
195
196// ===== Accessors =================================================
197// Returns the root token that idx maps to. For a kept token (root)
198// this is idx itself. For a merged token this is the root it was
199// merged into. V1 invariant: chains have length 1 (no cascading
200// merges -- once a token is merged it cannot be a merge target).
201func nx_token_merge_target_for(p: *NxTokenMergePlan, idx: i64) -> i64 {
202 if nx_token_merge_plan_is_valid(p) != 1 { return 0 - NX_TOMERGE_TAMPER }
203 if idx < 0 { return 0 - NX_TOMERGE_BAD_INPUT }
204 if idx >= p.n_tokens { return 0 - NX_TOMERGE_BAD_INPUT }
205 return p.assignment[idx]
206}
207
208func nx_token_merge_is_kept(p: *NxTokenMergePlan, idx: i64) -> i64 {
209 if nx_token_merge_plan_is_valid(p) != 1 { return 0 - NX_TOMERGE_TAMPER }
210 if idx < 0 { return 0 - NX_TOMERGE_BAD_INPUT }
211 if idx >= p.n_tokens { return 0 - NX_TOMERGE_BAD_INPUT }
212 if p.assignment[idx] == idx { return 1 }
213 return 0
214}
215
216func nx_token_merge_n_kept(p: *NxTokenMergePlan) -> i64 {
217 if nx_token_merge_plan_is_valid(p) != 1 { return 0 - 1 }
218 return p.n_tokens - p.n_merged
219}
220
221func nx_token_merge_n_merged(p: *NxTokenMergePlan) -> i64 {
222 if nx_token_merge_plan_is_valid(p) != 1 { return 0 - 1 }
223 return p.n_merged
224}
225
226func nx_token_merge_n_tokens(p: *NxTokenMergePlan) -> i64 {
227 if nx_token_merge_plan_is_valid(p) != 1 { return 0 - 1 }
228 return p.n_tokens
229}
230
231// Reduction ratio as Q16 fixed-point: n_merged / n_tokens.
232func nx_token_merge_reduction_q16(p: *NxTokenMergePlan) -> i64 {
233 if nx_token_merge_plan_is_valid(p) != 1 { return 0 - 1 }
234 return (p.n_merged * NX_MAGIC_65536) / p.n_tokens
235}