nx_gqa_head_map.nx source
↩ module page · 228 lines · 9244 B
1// nx_gqa_head_map.nx -- H6 grouped/multi-query attention bits-up.
2//
3// Per NISHI_ELDER_AI_OFF_DOCKER_2026_05_20.md ยง2.1 H6: grouped-query
4// attention (Ainslie 2023) + multi-query attention (Shazeer 2019)
5// collapse the K and V projection dimensions so multiple query
6// heads share one KV head. Memory + bandwidth wins on the KV
7// cache scale with (n_q_heads / n_kv_heads): a 32-Q-head model with
8// 8 KV heads cuts KV cache 4x; with 1 KV head (MQA) cuts it 32x.
9// Llama 3, Mistral, Qwen, every modern open-weight model ships
10// some flavour of this.
11//
12// V1 mechanics: a static mapping table head_to_kv_group[q_head_idx]
13// = kv_head_idx, computed once at model-load time from
14// (n_q_heads, n_kv_heads). The mapping is contiguous-group:
15// q heads [0..group_size) -> kv head 0, [group_size..2*group_size)
16// -> kv head 1, etc. This matches the layout every published
17// implementation uses (Llama / Mistral / Qwen). Three named
18// attention kinds fall out as sealed-enum special cases:
19// MHA: n_kv_heads == n_q_heads (group_size = 1)
20// GQA: 1 < n_kv_heads < n_q_heads (group_size > 1, divides n_q_heads)
21// MQA: n_kv_heads == 1 (group_size = n_q_heads)
22//
23// Composition path: a downstream attention kernel (FlashAttention
24// or naive softmax) loops over q heads, calls
25// nx_gqa_head_map_kv_head_for_q_head(map, q_h) to find which KV
26// head to read from, and computes attention against that smaller
27// shared K/V. H5 nx_flash_attention can compose with this --
28// the tile plan reads from the shared KV head.
29//
30// Pure substrate logic. No Linux features. Composes with shipped
31// H1 nx_kv_arena (KV pages are now sized by n_kv_heads not
32// n_q_heads -- the arena allocator's hidden_dim parameter should
33// reflect the smaller GQA/MQA layout), H5 nx_flash_attention
34// (tile plan reads from shared KV head).
35//
36// V1 honest scope:
37// - Contiguous-group mapping only (V2 could add strided or
38// learned mapping for research variants; not in any production
39// open-weight model as of 2026-05-20)
40// - Sealed attention kind derived from (n_q, n_kv) -- not stored
41// as a separate field, computed via classify()
42// - n_q_heads must be divisible by n_kv_heads (no jagged groups)
43//
44// genealogy_id: shazeer_2019_multi_query_attention +
45// ainslie_2023_gqa + llama2_2023_paper +
46// cardinal_2026-05-20_elder_ai_off_docker +
47// cardinal_2026-05-20_bits_up_nishi_not_linux
48// lineage_id: substrate_gqa_head_map_v1
49//
50// nx_capability_manifest:
51// variant_class: gqa_head_map
52// variant_id: gqa_head_map_v1_contiguous
53// requires_isa: [rv32i, rv64imac, x86_64, cortex_m, armv7a, aarch64, wasm32]
54// requires_syscalls: [mmap]
55// requires_ram_min_b: 1024
56// tier_floor: NX_TIER_INF_EDGE
57// tier_ceiling: NX_TIER_INF_HPC
58// cost_model:
59// flops_per_n: 0.0 // O(1) lookup, no flops
60// bytes_per_n: 8.0 // one i64 entry per q head
61// syscalls_per_n: 0.0
62// adversary_class: THREAT_AI_ADVERSARY
63//
64// nx_safety_envelope:
65// intended_use: "Q-head to KV-head mapping for grouped/multi-query
66// attention; bits-up; composes with nx_kv_arena
67// and nx_flash_attention"
68// sil_target: SIL2
69// evidence: [canary_bracketed, kind_sealed_enum,
70// divisibility_invariant, contiguous_group]
71// verdict: NOT_YET_EVALUATED
72
73import "nx_syscalls.nx"
74const NX_MAGIC_65536: i64 = 65536
75
76// ===== Constants =================================================
77const NX_GQA_MAX_Q_HEADS: i64 = 256
78const NX_GQA_MAX_KV_HEADS: i64 = 256
79
80// Attention kinds (sealed).
81const NX_ATTN_MHA: i64 = 0 // n_kv == n_q
82const NX_ATTN_GQA: i64 = 1 // 1 < n_kv < n_q
83const NX_ATTN_MQA: i64 = 2 // n_kv == 1
84const NX_ATTN_N_KINDS: i64 = 3
85
86func nx_attn_kind_is_valid(k: i64) -> i64 {
87 if k < 0 { return 0 }
88 if k >= NX_ATTN_N_KINDS { return 0 }
89 return 1
90}
91
92// Verdicts.
93const NX_GQA_OK: i64 = 0
94const NX_GQA_BAD_INPUT: i64 = 1
95const NX_GQA_BAD_DIVISIBILITY: i64 = 2
96const NX_GQA_OUT_OF_RANGE: i64 = 3
97const NX_GQA_TAMPER: i64 = 4
98const NX_GQA_N_VERDICTS: i64 = 5
99
100func nx_gqa_verdict_is_valid(v: i64) -> i64 {
101 if v < 0 { return 0 }
102 if v >= NX_GQA_N_VERDICTS { return 0 }
103 return 1
104}
105
106// Canary magic.
107const NX_GQA_MAP_CANARY_PRE: i64 = 0x4751414D6170503A // "GQAMapP:"
108const NX_GQA_MAP_CANARY_POST: i64 = 0x4D617045314532A0 // "MapE1E2 "
109
110// ===== Struct ====================================================
111struct NxGqaHeadMap {
112 canary_pre: i64,
113 n_q_heads: i64,
114 n_kv_heads: i64,
115 group_size: i64,
116 kind: i64, // NX_ATTN_MHA / GQA / MQA
117 head_to_kv_group: *i64, // length n_q_heads
118 canary_post: i64,
119}
120
121// ===== Validity =================================================
122func nx_gqa_head_map_is_valid(m: *NxGqaHeadMap) -> i64 {
123 if (m as i64) == 0 { return 0 }
124 if m.canary_pre != NX_GQA_MAP_CANARY_PRE { return 0 }
125 if m.canary_post != NX_GQA_MAP_CANARY_POST { return 0 }
126 if m.n_q_heads <= 0 { return 0 }
127 if m.n_q_heads > NX_GQA_MAX_Q_HEADS { return 0 }
128 if m.n_kv_heads <= 0 { return 0 }
129 if m.n_kv_heads > m.n_q_heads { return 0 }
130 if m.group_size <= 0 { return 0 }
131 if m.group_size * m.n_kv_heads != m.n_q_heads { return 0 }
132 if nx_attn_kind_is_valid(m.kind) != 1 { return 0 }
133 return 1
134}
135
136// Classify by counts (pure, doesn't allocate).
137func nx_attn_classify(n_q_heads: i64, n_kv_heads: i64) -> i64 {
138 if n_kv_heads == n_q_heads { return NX_ATTN_MHA }
139 if n_kv_heads == 1 { return NX_ATTN_MQA }
140 return NX_ATTN_GQA
141}
142
143// ===== Constructor =================================================
144// Returns null on bad input. Caller checks with nx_gqa_head_map_is_valid.
145func nx_gqa_head_map_new(n_q_heads: i64, n_kv_heads: i64) -> *NxGqaHeadMap {
146 if n_q_heads <= 0 { return (0 as i64) as *NxGqaHeadMap }
147 if n_q_heads > NX_GQA_MAX_Q_HEADS { return (0 as i64) as *NxGqaHeadMap }
148 if n_kv_heads <= 0 { return (0 as i64) as *NxGqaHeadMap }
149 if n_kv_heads > n_q_heads { return (0 as i64) as *NxGqaHeadMap }
150 // Must divide evenly (contiguous group invariant).
151 let rem: i64 = n_q_heads - (n_q_heads / n_kv_heads) * n_kv_heads
152 if rem != 0 { return (0 as i64) as *NxGqaHeadMap }
153
154 let m: *NxGqaHeadMap = (sys_mmap(64)) as *NxGqaHeadMap
155 m.canary_pre = NX_GQA_MAP_CANARY_PRE
156 m.n_q_heads = n_q_heads
157 m.n_kv_heads = n_kv_heads
158 m.group_size = n_q_heads / n_kv_heads
159 m.kind = nx_attn_classify(n_q_heads, n_kv_heads)
160 m.head_to_kv_group = (sys_mmap(n_q_heads * 8)) as *i64
161 m.canary_post = NX_GQA_MAP_CANARY_POST
162
163 var q: i64 = 0
164 while q < n_q_heads {
165 m.head_to_kv_group[q] = q / m.group_size
166 q = q + 1
167 }
168 return m
169}
170
171// ===== Lookup =================================================
172// Returns the KV head index that serves a given query head. On
173// invalid input returns 0 - <verdict>.
174func nx_gqa_head_map_kv_head_for_q_head(m: *NxGqaHeadMap, q_head: i64) -> i64 {
175 if nx_gqa_head_map_is_valid(m) != 1 { return 0 - NX_GQA_TAMPER }
176 if q_head < 0 { return 0 - NX_GQA_OUT_OF_RANGE }
177 if q_head >= m.n_q_heads { return 0 - NX_GQA_OUT_OF_RANGE }
178 return m.head_to_kv_group[q_head]
179}
180
181// Count of Q heads that share a given KV head. For contiguous-group
182// mapping this is always m.group_size, but exposing the count by
183// scan makes the invariant auditable from downstream.
184func nx_gqa_head_map_count_q_heads_in_group(m: *NxGqaHeadMap, kv_head: i64) -> i64 {
185 if nx_gqa_head_map_is_valid(m) != 1 { return 0 - NX_GQA_TAMPER }
186 if kv_head < 0 { return 0 - NX_GQA_OUT_OF_RANGE }
187 if kv_head >= m.n_kv_heads { return 0 - NX_GQA_OUT_OF_RANGE }
188 var count: i64 = 0
189 var q: i64 = 0
190 while q < m.n_q_heads {
191 if m.head_to_kv_group[q] == kv_head { count = count + 1 }
192 q = q + 1
193 }
194 return count
195}
196
197// ===== Accessors =================================================
198func nx_gqa_head_map_n_q_heads(m: *NxGqaHeadMap) -> i64 {
199 if nx_gqa_head_map_is_valid(m) != 1 { return 0 - 1 }
200 return m.n_q_heads
201}
202
203func nx_gqa_head_map_n_kv_heads(m: *NxGqaHeadMap) -> i64 {
204 if nx_gqa_head_map_is_valid(m) != 1 { return 0 - 1 }
205 return m.n_kv_heads
206}
207
208func nx_gqa_head_map_group_size(m: *NxGqaHeadMap) -> i64 {
209 if nx_gqa_head_map_is_valid(m) != 1 { return 0 - 1 }
210 return m.group_size
211}
212
213func nx_gqa_head_map_kind(m: *NxGqaHeadMap) -> i64 {
214 if nx_gqa_head_map_is_valid(m) != 1 { return 0 - 1 }
215 return m.kind
216}
217
218// KV cache savings ratio as Q16 fixed-point. Returns
219// (1 - n_kv_heads/n_q_heads) * 65536 == bytes_saved / bytes_naive.
220// For MHA this is 0; for n_q=32,n_kv=8 (4x compress) this is 49152
221// (= 0.75 * 65536); for MQA n_q=32,n_kv=1 this is 63488 (~0.969).
222func nx_gqa_head_map_kv_savings_q16(m: *NxGqaHeadMap) -> i64 {
223 if nx_gqa_head_map_is_valid(m) != 1 { return 0 - 1 }
224 let total: i64 = m.n_q_heads * NX_MAGIC_65536
225 let kept: i64 = m.n_kv_heads * NX_MAGIC_65536
226 let saved: i64 = total - kept
227 return saved / m.n_q_heads
228}