Paper deep dive
MoA-Structured Decode Attention DNF Derivation, KV-Cache Accumulation, GQA/MQA, and OpenACC Kernel
Lenore Mulin, Gaetan Hains
Intelligence
Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 93%
Last extracted: 7/23/2026, 2:19:30 AM
Summary
The paper introduces a Memory-optimal inference artifacts for transformer attention using the Mathematics of Arrays (MoA). It derives a single-query decode Denotational Normal Form (DNF) that eliminates the K^T buffer, a C/OpenACC GPU kernel with Operational Normal Form (ONF) stride arithmetic, a multi-step KV-cache accumulation method, and Grouped-Query Attention (GQA)/Multi-Query Attention (MQA) derivations. The approach proves memory minimality via the Storage Theorem, reducing DRAM traffic compared to conventional methods.
Entities (9)
Relation Signals (7)
Mathematics of Arrays → uses → Denotational Normal Form
confidence 95% · We derive four memory-optimal inference artifacts for transformer attention using the Mathematics of Arrays (MoA), each following directly from the forward-pass Denotational Normal Form (DNF)
Mathematics of Arrays → uses → Operational Normal Form
confidence 95% · a C/OpenACC Graphics Processing Unit (GPU) kernel with Operational Normal Form (ONF) stride arithmetic
Mathematics of Arrays → enables → KV cache
confidence 94% · a multi-step KV-cache with O(d_k+d_v) per-step append via MoA concatenation
Mathematics of Arrays → derives → Grouped Query Attention
confidence 93% · Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) derived via ψ-selection
Mathematics of Arrays → derives → Multi-Query Attention
confidence 93% · Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) derived via ψ-selection
Storage Theorem → proves → Memory Minimality
confidence 92% · Storage Theorem 2.7 of [1]: every array read once; traffic proven minimal before code is written
OpenACC → implements → Operational Normal Form
confidence 90% · a C/OpenACC Graphics Processing Unit (GPU) kernel with Operational Normal Form (ONF) stride arithmetic
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:We derive four memory-optimal inference artifacts for transformer attention using the Mathematics of Arrays (MoA), each following directly from the forward-pass Denotational Normal Form (DNF) of with the query-row index fixed to the current decode step. The artifacts are: (1)~a single-query decode DNF in which the $\psi$-reduction eliminates the $K^\top$ buffer algebraically, achieving $(d_k + nd_k+ nd_v+ d_v)\times4\,{B}$ Dynamic Random Access Memory (DRAM) traffic result numerically verified to $\|{err}\|_\leq2\times10^{-7}$; (2)~a C/OpenACC Graphics Processing Unit (GPU) kernel with Operational Normal Form (ONF) stride arithmetic and hardware-coalesced memory access, verified to $\|\mathrm{err}\|_\infty=0$ (exact IEEE-754 floating-point arithmetic); (3)~a multi-step KV-cache with $O(d_k+d_v)$ per-step append via MoA concatenation $\#$; and (4)~Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) derived via $\psi$-selection, achieving a proven $\frac {h_q} { h_{kv} }$ reduction in KV traffic. All programs are verified against PyTorch scaled_dot_product_attention.
Tags
Links
- Source: https://arxiv.org/abs/2607.19456v1
- Canonical: https://arxiv.org/abs/2607.19456v1
Trouble viewing inline? Open PDF directly →
Full Text
54,116 characters extracted from source content.
Expand or collapse full text
MoA-Structured Decode Attention DNF Derivation, KV-Cache Accumulation, GQA/MQA, and OpenACC Kernel Lenore Mullin1,∗ 1University at Albany, SUNY lmullin@albany.edu ∗Corresponding author Gaétan Hains2 2LACL, Université Paris-Est Créteil gaetan.hains@u-pec.fr (June 2026) Abstract We derive four memory-optimal inference artifacts for transformer attention using the Mathematics of Arrays (MoA), each following directly from the forward-pass Denotational Normal Form (DNF) of [1] with the query-row index fixed to the current decode step. The artifacts are: (1) a single-query decode DNF in which the ψ-reduction eliminates the K⊤K buffer algebraically, achieving (dk+ndk+ndv+dv)×4B(d_k+nd_k+nd_v+d_v)× 4\,B Dynamic Random Access Memory (DRAM) traffic (Storage Theorem 2.7 of [1]), result numerically verified to ‖err‖∞≤2×10−7\|err\|_∞≤ 2× 10^-7; (2) a C/OpenACC Graphics Processing Unit (GPU) kernel with Operational Normal Form (ONF) stride arithmetic and hardware-coalesced memory access, verified to ‖err‖∞=0\|err\|_∞=0 (exact IEEE-754 floating-point arithmetic); (3) a multi-step KV-cache with O(dk+dv)O(d_k+d_v) per-step append via MoA concatenation #\#; and (4) Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) [7, 8] derived via ψ-selection, achieving a proven hq/hkvh_q/h_kv reduction in KV traffic. All programs are verified against PyTorch scaled_dot_product_attention. Phase L=96 layers Conventional allocates to DRAM per layer MoA ψ-reduction eliminates per layer Forward training arXiv ∼ 100ms/batch Per layer, per batch: K⊤K ρ=⟨dk,n⟩ρ= d_k,n (buffer) S=QK⊤S=QK ρ=⟨n,n⟩ρ= n,n A=softmax(S)A=softmax(S) ρ=⟨n,n⟩ρ= n,n traffic: 96×O(n2dk+n2dv)96\!×\!O(n^2d_k\!+\!n^2d_v) Per layer, per batch: K⊤K — never materialised S — on-stack per row, discarded A — on-stack per row, discarded traffic: 96×O(ndk+ndv)96\!×\!O(nd_k\!+\!nd_v) ✓ ×102× 10^2 epochs×104× 10^4 batches×96× 96 layers O(n2)→O(n)O(n^2)→ O(n) per layer per batch ×106× 10^6 batches ≈ 40 kWh saved per training run Fused fwd+bwd HAL ∼ 180ms/batch Per layer, per batch: A saved fwd→ ρ=⟨n,n⟩ρ= n,n Jacobian ρ=⟨n,n⟩ρ= n,n dSdS ρ=⟨n,n⟩ρ= n,n traffic: 96×(8ndk+8ndv)96\!×\!(8nd_k\!+\!8nd_v) Per layer, per batch: A recomputed on-fly, not stored Jacobian — outer-product compat. dSdS — substitution traffic: 96×(4ndk+4ndv)96\!×\!(4nd_k\!+\!4nd_v) ✓ ×102× 10^2 epochs×104× 10^4 batches×96× 96 layers 2×2× vs separate passes ×96× 96 layers ×106× 10^6 batches ≈ 80 kWh saved per training run Inference decode This paper 1–10 μ /step per layer Per token, per layer: K⊤K view (q @ K.T) score ρ=⟨1,n⟩ρ= 1,n (n grows) weight ρ=⟨1,n⟩ρ= 1,n n=4096n\!=\!4096, L=96L\!=\!96: ≈ 200 MB/step, no proof Per token, per layer: K⊤K — fix i′=1i \!=\!1, ψ reads rows s→,a→ s, a — rank-1, on-stack KV cache: +1 row at γ offset n=4096n\!=\!4096, L=96L\!=\!96: ≈ 200 MB/step ✓proven ×109× 10^9 req/day×512× 512 tokens×96× 96 layers KV cache n=1→512n\!=\!1\!→\!512 per request; total: 96×∑t=1512(dk+tdk+tdv+dv)×4B≈96×1.6GB96\!×\! _t=1^512(d_k+td_k+td_v+d_v)\!×\!4\,B\!≈\!96\!×\!1.6\,GB Energy: 300 W × 10 μ × 96 × 512 × 10910^9/day ≈ 14,700 kWh/day — MoA minimises every byte transferred FlashAttention: similar traffic via tiling, no algebraic proof. MoA: proven before code. Same DNF → CPU/GPU/FPGA. Storage Thm. 2.7 [1]: every array read once; traffic proven minimal before code is written Figure 1: Realistic nested loops: conventional vs. MoA across training and inference. Solid loops = outer; dashed = inner. Training runs ∼106 \!10^6 total batches once; inference runs 10910^9 requests/day indefinitely, each generating 512 tokens across 96 layers. At n=4096n\!=\!4096, L=96L\!=\!96: both conventional and MoA use ≈200≈\!200 MB per decode step — but only MoA’s is proven minimal (Storage Thm. 2.7). FlashAttention achieves similar traffic via tiling without algebraic proof. Energy: ≈14,700≈\!14,700 kWh/day per deployed model; MoA minimises every byte transferred. 1 Introduction 1.1 Background The Mathematics of Arrays (MoA) [3] is a formal algebra of multi-dimensional arrays whose primitive operations — ψ-selection, ρ (shape), ι (index generation), γ (linear offset), and the Omega (Ω ) family of inner products — compose provably minimal computations. The Denotational Normal Form (DNF) specifies what to compute; the Operational Normal Form (ONF) specifies how to lay it out in memory. The companion paper [1] applies MoA to the forward training pass of transformer attention [6], deriving the DNF and proving via the Storage Theorem (Theorem 2.7 of [1]) that the resulting ONF achieves the information-theoretic minimum DRAM traffic. The HAL companion [2] extends this to the backward pass and fused kernel. 1.2 Contributions This paper shows that the same DNF, with the query-row index i′i fixed to the current decode step, yields four inference artifacts. 1. Decode DNF (§2): the ψ-reduction eliminates the K⊤K buffer implicit in the standard formula A=softmax(QK⊤/dk)A=softmax(QK / d_k), reducing the score computation to s→=(q→(+.×Ω⟨1,2⟩)K) s=scale\,( q\;(+.× _ 1,2 )\;K) with K accessed row-by-row in its natural layout. Memory traffic: MMoAdec=(dk+ndk+ndv+dv)×4BM^dec_MoA=(d_k+nd_k+nd_v+d_v)× 4\,B, is then proven minimal by Storage Theorem 2.7 of [1]. 2. C/OpenACC GPU kernel (§4): ONF γ stride arithmetic gives coalesced memory access; the coalescing criterion is stated as a γ-difference condition. 3. KV-cache accumulation (§5): MoA concatenation #\# appends one row per step at the pre-computed γ offset, with proven O(dk+dv)O(d_k+d_v) per-step traffic. 4. GQA/MQA (§6): ψ-selection ⟨gkv⟩ψK g_kv ψ\,K maps each query head to its KV group without materialising the broadcast, reducing traffic by hq/hkvh_q/h_kv. 1.3 Notation and Terminology Figure 1 summarises the relationship between the three companion papers and the importance of memory minimality at every stage. A transformer [6] is a neural network in which each output element is a learned weighted combination of all input elements (attention). A token is the basic unit of text processed by the model (typically a word or sub-word). Autoregressive inference generates output one token at a time, feeding each generated token back as input; at each step a single query vector q→=WQx→ q=W_Q x is produced from the current token embedding x→ x via learned projection WQW_Q. All past key vectors k→l=WKx→l k_l=W_K x_l and value vectors v→l=WVx→l v_l=W_V x_l are accumulated in the KV cache: K (shape, ρ(K)=⟨n,dk⟩ρ(K)= n,d_k ) and V (shape, ρ(V)=⟨n,dv⟩ρ(V)= n,d_v ), where n grows by one per step. The dominant cost at each step is reading the full KV cache from DRAM: O(ndk+ndv)O(nd_k+nd_v) bytes, unavoidable by any correct implementation (Storage Theorem 2.7 of [1]). MoA proves this bound algebraically before code is written. MoA notation used throughout: ρ(A)ρ(A) denotes the shape of array A; ⟨i0,…,ir−1⟩ψA i_0,…,i_r-1 ψ\,A selects element (i0,…,ir−1)(i_0,…,i_r-1); γ(⟨i0,…⟩,ρ(A))γ( i_0,… ,ρ(A)) is the linear memory offset; A(+.×Ω⟨σl,σr⟩)BA\;(+.× _ _l, _r )\;B is the Omega inner product with left rank σl _l, right rank σr _r; A#BA\;\#\;B concatenates along the leading axis. All multiplication uses ×. 2 Decode DNF Formal Inputs At decode step t with n cached tokens: • q→ q: ρ(q→)=⟨dk⟩ρ( q)= d_k — query vector for the current token. • K: ρ(K)=⟨n,dk⟩ρ(K)= n,d_k — n rows, each row ⟨l⟩ψK l ψ\,K (shape ⟨dk⟩ d_k ) is the key vector for cached token l. K⊤K would have shape ⟨dk,n⟩ d_k,n ; the ψ-reduction accesses rows of K directly so K⊤K is never computed thus saving transfers and storage. • V: ρ(V)=⟨n,dv⟩ρ(V)= n,d_v — n rows, each row ⟨l⟩ψV l ψ\,V (shape ⟨dv⟩ d_v ) is the value vector for cached token l. Output: → out, ρ(→)=⟨dv⟩ρ( out)= d_v ; its d-th element is ⟨d⟩ψ→ d ψ\, out. An attention head is an independent copy of the attention computation with its own learned key, query, and value projections. Using h heads in parallel (multi-head attention) allows the model to extract different types of relationships from the input simultaneously [6]. In the multi-head case K becomes rank-3 (ρ(K)=⟨h,n,dk⟩ρ(K)= h,n,d_k ) or rank-4 (ρ(K)=⟨B,h,n,dk⟩ρ(K)= B,h,n,d_k ), where batch B is the number of independent sequences processed simultaneously. The Omega operator Ω⟨1,2⟩ _ 1,2 generalizes automatically: σl=1 _l=1 uses a rank-1 vector ; σr=2 _r=2 uses a rank-2 matrix as l/r arguments. Leading batch or head dimensions are outer indices iterated over by Omega — no re-derivation is needed. How K⊤K is eliminated by ψ-reduction. The standard attention formula contains K⊤K : A A =softmax(QK⊤/dk), =softmax(QK / d_k), Q,K:ρ=⟨n,dk⟩,QK⊤:ρ=⟨n,n⟩. Q,K:ρ= n,d_k , QK :ρ= n,n . The arXiv paper [1] starts from this formula and applies the ψ-reduction. The key step: instead of forming QK⊤QK (which requires materialising K⊤K in memory), the ψ-reduction shows that each element can be computed by accessing K directly: ⟨i′,l⟩ψ(QK⊤)=⟨i′⟩ψQ(+.×Ω⟨1,1⟩)⟨l⟩ψK. i ,l ψ\,(QK )= i ψ\,Q\;(+.× _ 1,1 )\; l ψ\,K. where this specific Ω expression implements the product. The columns of K⊤K are exactly the rows ⟨l⟩ψK l ψ\,K of K — so K⊤K is algebraically present in the starting formula but the ψ-reduction proves it never needs to be materialised. The Storage Theorem (Theorem 2.7 of [1]) then guarantees this is the minimum-traffic access pattern. For the decode case (this paper), Q reduces to a single query vector q→ q, so the ψ-reduction over i′i disappears (there is only one query row), giving directly: s→[l]=q→(+.×Ω⟨1,1⟩)⟨l⟩ψK. s[l]=scale\; q\;(+.× _ 1,1 )\; l ψ\,K. K⊤K was in the original formulation; the ψ-reduction removed the need to materialise it before any code is written. This is the central result of [1] applied to the decode setting. This is equations (4)–(6) of [1] with the query-row index i′i fixed to the current token: eq. (4) is the score ψ-reduction, eq. (5) the softmax, eq. (6) the weighted sum. No new Omega derivation is required. Independence of q→ q and K. q→ q and K are produced independently: • q→=WQx→current q=W_Q x_current is the projection of the current token embedding through the learned query matrix WQW_Q, ρ(WQ)=⟨dk,d⟩ρ(W_Q)= d_k,d . • ⟨l⟩ψK=WKx→l l ψ\,K=W_K x_l is the projection of the l-th past token embedding through the learned key matrix WKW_K, stored as row l of the KV cache. Neither q→ q nor K is derived from the other. Step I asks: how similar is my current query q→ q to each of the n stored key vectors ⟨l⟩ψK l ψ\,K? Step I – Scores. The ψ-reduction selects row l of K as ⟨l⟩ψK l ψ\,K (shape ⟨dk⟩ d_k ) and contracts it with q→ q via the inner product: Remark 1. The standard formula QK⊤QK is algebraically present as the starting point (see §2 preamble); the ψ-reduction proves K⊤K need never be materialised. In MoA there are no column vectors: q→ q and K are both rank-≥1≥ 1 arrays and contraction is expressed directly via Ω⟨1,2⟩ _ 1,2 . s→[l] s[l] =(q→(+.×Ω⟨1,1⟩)⟨l⟩ψK) =scale\; ( q\;(+.× _ 1,1 )\; l ψ\,K ) =∑j=0dk−1⟨j⟩ψq→×⟨l,j⟩ψK. =scale _j=0^d_k-1 j ψ\, q× l,j ψ\,K. Over all n rows simultaneously, with σl=1 _l=1 exposing q→ q as a whole rank-1 vector and σr=2 _r=2 exposing K (shape ⟨n,dk⟩ n,d_k ) as a whole rank-2 array, the Omega expression is: s→ s =(q→(+.×Ω⟨1,2⟩)K),ρ(s→)=⟨n⟩. =scale\; ( q\;(+.× _ 1,2 )\;K ), ρ( s)= n . (1) This equals Kq→scale\,K q (matrix-vector product with K in its natural row-major layout). No transpose is formed: the Cartesian coordinate interchange is implicit in the ψ-reduction, which reads K row by row as stored. The Python equivalent is scale*(K@q), not scale*(q@K.T). Ω⟨1,2⟩ _ 1,2 is fixed: it always pairs one rank-1 vector from the left with one rank-2 matrix from the right, applying +.×+.× to produce a scalar for each pair. As the rank of the arguments grows, the same partitioning lifts naturally — superscripts denote rank, Ω⟨1,2⟩ _ 1,2 never changes, only the result rank grows: • q1→(+.×Ω⟨1,2⟩)K2→s1→ q^1\;(+.× _ 1,2 )\;K^2→ s^1 (vector-matrix product) • A2(+.×Ω⟨1,2⟩)K3→A^2\;(+.× _ 1,2 )\;K^3→ rank-2 (each row of A with each matrix of K) • A3(+.×Ω⟨1,2⟩)K4→A^3\;(+.× _ 1,2 )\;K^4→ rank-3 (same, one rank higher) This is how MoA handles batched and multi-head attention without re-derivation. No K⊤K buffer is formed; K is read once and fully consumed in Step I. Step I – Numerically stable softmax. K does not appear in Step I at all. By the time Step I runs, K has already been fully processed and discarded. Step I operates only on the rank-1 score vector s→ s (shape ⟨n⟩ n ) produced by Step I — no matrix, no K⊤K , no K. In the standard training formula, softmax would be applied to the full n×n× n score matrix QK⊤/dkQK / d_k and its Jacobian would be n×n× n. Here, softmax operates on a single rank-1 vector s→ s and the MoA ψ-reduction eliminates all intermediates: s→ s, e→ e, a→ a and scalars m, Z are O(n)O(n) on-stack and never written to DRAM. m m =maxls→[l], = _l s[l], e→[l] e[l] =exp(s→[l]−m), = ( s[l]-m), Z Z =∑l′=0n−1e→[l′], = _l =0^n-1 e[l ], a→[l] a[l] =e→[l]/Z,0≤l<n. = e[l]/Z, 0≤ l<n. (2) No n×n× n score matrix or Jacobian is allocated. Numerical stability: subtracting m=maxls→[l]m= _l s[l] before exponentiating does not change the softmax value (the m terms cancel) but prevents overflow. =1/dkscale=1/ d_k normalises variance to 1 for random unit-norm inputs. Step I – Weighted sum: ⟨d⟩ψ→ d ψ\, out =∑l=0n−1⟨l⟩ψa→×⟨l,d⟩ψV, = _l=0^n-1 l ψ\, a× l,d ψ\,V, 0≤d<dv, 0≤l<n. 40.00006pt0≤ d<d_v,\;0≤ l<n. (3) V is read once in row-major order; no intermediate array exists. Concrete example (n=2n=2, dk=2d_k=2, dv=2d_v=2, =1/2scale=1/ 2). Let q→=⟨1,0⟩,K=(1001),V=(2468). q= 1,0 , K= pmatrix1&0\\ 0&1 pmatrix, V= pmatrix2&4\\ 6&8 pmatrix. • Step I: let qj=⟨j⟩ψq→q_j= j ψ\, q, Klj=⟨l,j⟩ψK_lj= l,j ψ\,K. s→[0] s[0] =(q0K00+q1K01)=12(1+0)=12, =scale(q_0K_00+q_1K_01)= 1 2(1+0)= 1 2, s→[1] s[1] =(q0K10+q1K11)=12(0+0)=0. =scale(q_0K_10+q_1K_11)= 1 2(0+0)=0. • Step I (softmax, m=1/2m=1/ 2): e→[0] e[0] =exp(0)=1, = (0)=1, e→[1] e[1] =exp−1/2≈0.493, = -1/ 2≈ 0.493, Z Z ≈1.493, ≈ 1.493, a→[0] a[0] ≈0.670,a→[1]≈0.330. ≈ 0.670,\; a[1]≈ 0.330. • Step I (weighted sum): ⟨0⟩ψ→ 0 ψ\, out =a→[0]×⟨0,0⟩ψV+a→[1]×⟨1,0⟩ψV = a[0]× 0,0 ψ\,V+ a[1]× 1,0 ψ\,V =0.670×2+0.330×6≈3.32, =0.670× 2+0.330× 6≈ 3.32, ⟨1⟩ψ→ 1 ψ\, out =a→[0]×⟨0,1⟩ψV+a→[1]×⟨1,1⟩ψV = a[0]× 0,1 ψ\,V+ a[1]× 1,1 ψ\,V =0.670×4+0.330×8≈5.32. =0.670× 4+0.330× 8≈ 5.32. Proposition 1 (Decode memory minimality). The decode DNF achieves MMoAdec=(dk+ndk+ndv+dv)×4B=O(ndk+ndv)M^dec_MoA=(d_k+nd_k+nd_v+d_v)× 4\,B=O(nd_k+nd_v) memory usage.. Algorithm 1 MoA decode (single query step). Scratch: s→[n] s[n], e→[n] e[n] – O(n)O(n); no n2n^2 array. 1:q→:⟨dk⟩ q\!:\! d_k , K:⟨n,dk⟩K\!:\! n,d_k , V:⟨n,dv⟩V\!:\! n,d_v , =1/dkscale\!=\!1/\! d_k 2:→:⟨dv⟩ out\!:\! d_v 3:for l=0…n−1l=0… n-1 do 4: s→[l]←∑j⟨j⟩ψq→×⟨l,j⟩ψK s[l] _j j ψ\, q× l,j ψ\,K ⊳ eq. (1) 5:end for 6:m←maxls→[l]m← _l s[l] 7:for l=0…n−1l=0… n-1 do 8: e→[l]←exp(s→[l]−m) e[l]← ( s[l]-m) 9:end for 10:Z←∑le→[l]Z← _l e[l]; a→[l]←e→[l]/Z\; a[l]← e[l]/Z for all l 11:for d=0…dv−1d=0… d_v-1 do 12: ⟨d⟩ψ→←∑la→[l]×⟨l,d⟩ψV d ψ\, out← _l a[l]× l,d ψ\,V ⊳ eq. (3) 13:end for Python start here q @ K.T K⊤K buffer softmax ⟨1,n⟩ 1,n row a @ V no proof MoA algebra ρ(Q)=⟨n,dk⟩ρ(Q)= n,d_k ρ(K)=⟨n,dk⟩ρ(K)= n,d_k ρ(V)=⟨n,dv⟩ρ(V)= n,d_v Q(+.×Ω⟨2,2⟩)KQ(+.\!×\! _ 2,2 )K softmax A(+.×Ω⟨2,2⟩)VA(+.\!×\! _ 2,2 )V K⊤K still present DNF ψ-reduction Step I s→=q→(+.×Ω⟨1,2⟩)K s\!=\!scale\, q(+.\!×\! _ 1,2 )K Step I a→=softmax(s→) a\!=\!softmax( s) Step I →=a→(+.×Ω⟨1,2⟩)V out\!=\! a(+.\!×\! _ 1,2 )V K⊤K eliminated ONF γ strides K at offset: l×dk+jl\!×\!d_k+j V at offset: l×dv+dl\!×\!d_v+d d-inner loop: stride(d)=1(d)\!=\!1 coalesced C / Fortran OpenACC C: K[l*dk+j] F90: K(j,l) OpenACC: gang/vector GPU coalesced ✓ proven expressψ-red.γ mapcompile K⊤K , n2n^2 score K⊤K in formula K⊤K gone Jacobian gone layout fixed coalescing minimal by Thm 2.7 Storage Theorem 2.7 [1]: DNF proves K⊤K never forms; ONF proves every element read once; code correct by construction. Figure 2: The MoA pipeline: Python → MoA → DNF → ONF → C/Fortran/OpenACC. Starting from conventional Python (q @ K.T), the MoA formulation expresses attention in array algebra. ψ-reduction (DNF) eliminates K⊤K and the softmax Jacobian algebraically. The γ stride map (ONF) fixes the row-major layout and proves coalescing. The resulting C/Fortran 90/OpenACC code is correct by construction (Storage Theorem 2.7). 3 Python Implementation Standard PyTorch Reference Two PyTorch shape operations appear throughout: .unsqueeze(d) inserts a new size-1 dimension at position d (e.g. ⟨dk⟩→⟨1,dk⟩ d_k → 1,d_k ); .squeeze() removes all size-1 dimensions (e.g. ⟨1,1,1,dv⟩→⟨dv⟩ 1,1,1,d_v → d_v ), recovering a MoA rank-1 vector. The conventional single-query decode in PyTorch unsqueezes q→ q to shape ⟨1,dk⟩ 1,d_k , forms the score row via q_ @ K.T (note: .T creates a transposed view; MoA avoids this entirely), applies softmax, and squeezes the result: ⬇ 1def std_decode(q, K, V, scale=None): 2 if scale is None: 3 scale = 1.0 / math.sqrt(K.shape[-1]) 4 q_ = q.unsqueeze(0) # (1, dk) 5 # .T creates a transposed view of K 6 scores = scale * (q_ @ K.T) # (1, n) 7 weights = torch.softmax(scores, dim=-1) 8 return (weights @ V).squeeze(0) # (dv,) Listing 1: Standard PyTorch single-query decode (pytorch_reference.py: std_decode). Compare with Listing 2: same result, no algebraic minimality proof. The MoA version (below) replaces q_ @ K.T with q→(+.×Ω⟨1,2⟩)K q\;(+.× _ 1,2 )\;K: the ψ-reduction reads K row by row in its natural layout — no transposed view is created — and the Storage Theorem proves the result is memory-minimal by construction. MoA - Based Implementation moa_decode.py mirrors the three DNF steps directly. The MoA formulation: • Step I (scores, eq. (1)): s→=(q→(+.×Ω⟨1,2⟩)K),ρ(s→)=⟨n⟩. s=scale\; ( q\;(+.× _ 1,2 )\;K ), ρ( s)= n . • Step I (softmax, eq. (2)): a→[l]=e→[l]/∑l′=0n−1e→[l′],ρ(a→)=⟨n⟩. a[l]= e[l] / _l =0^n-1 e[l ], ρ( a)= n . • Step I (weighted sum, eq. (3)): →=a→(+.×Ω⟨1,2⟩)V,ρ(→)=⟨dv⟩. out= a\;(+.× _ 1,2 )\;V, ρ( out)= d_v . Each listing block below corresponds to one step. ⬇ 1from typing import Optional 2import torch, math 3def moa_decode(q, K, V, scale=None): 4n, dk = K.shape 5if scale is None: 6scale = 1.0 / math.sqrt(dk) 7# Step I: scores -- no K^T buffer (eq. 1) 8s = scale * (K @ q) # (n,) 9# Step I: numerically stable softmax 10s = s - s.max() 11e = s.exp() 12a = e / e.sum() # (n,) O(n) scratch 13# Step I: weighted sum (eq. 3) 14return a @ V # (dv,) Listing 2: moa_decode.py: core decode function. Steps I–I map to eqs. (1)–(3). Reference Implementation PyTorch — a widely-used open-source numerical computing library for neural networks (A. Paszke et al., NeurIPS 2019) — provides scaled_dot_product_attention (SDPA) as a numerically verified fused reference implementation of the attention operation. In the full training setting SDPA computes softmax(Q(+.×Ω⟨2,2⟩)K)(+.×Ω⟨2,2⟩)Vsoftmax(scale\;Q\;(+.× _ 2,2 )\;K)\;(+.× _ 2,2 )\;V where ρ(Q)=⟨n,dk⟩ρ(Q)= n,d_k and the output has shape ⟨n,dv⟩ n,d_v . Here Ω⟨2,2⟩ _ 2,2 has σl=σr=2 _l= _r=2 (both arguments rank-2); for the single decode step, Q reduces to rank-1 q→ q so Ω⟨1,2⟩ _ 1,2 is used throughout this paper. For a single decode step there is only one query vector q→∈ℝdk q ^d_k, ρ(q→)=⟨dk⟩ρ( q)= d_k . In MoA there are no row or column vectors: the score vector s→ s, weight vector a→ a, and output → out are all rank-1 arrays. The decode computation proceeds in three steps, each expressed as a ψ-reduction that eliminates intermediates: • Step I (scores, ψ-reduction over K): q→(+.×Ω⟨1,2⟩)K∈ℝn q\;(+.× _ 1,2 )\;K ^n — no K⊤K buffer. • Step I (softmax, ψ-reduction over s→ s): softmax(s→/dk)∈ℝnsoftmax( s/ d_k) ^n — no n×n× n Jacobian. • Step I (weighted sum, ψ-reduction over V): a→(+.×Ω⟨1,2⟩)V∈ℝdv a\;(+.× _ 1,2 )\;V ^d_v (=→= out). Every intermediate and output is a rank-1 vector in MoA: s→,a→∈ℝn s, a ^n and →∈ℝdv out ^d_v. This is exactly what the decode DNF computes in equations (1)–(3): s→ s is the score vector, a→ a is the weight vector, and → out is the rank-1 output vector. SDPA is PyTorch’s own numerically verified fused kernel, making it a valid ground truth. Note: PyTorch’s SDPA requires (B, H, S, D) input shape (batch, heads, sequence, dimension), so dummy size-1 dimensions must be added before the call and removed with .squeeze() afterwards (defined above). To call it as a reference, q→ q, K, V are wrapped with dummy batch B=1B=1 and head H=1H=1 dimensions to match the required (B, H, S, D) input shape: ⬇ 1## Single decode step: vq:(dk,) K:(n,dk) V:(n,dv) 2q_ = vq.unsqueeze(0).unsqueeze(0).unsqueeze(0) 3## -> (1, 1, 1, dk) 4k_ = K.unsqueeze(0).unsqueeze(0) # -> (1, 1, n, dk) 5v_ = V.unsqueeze(0).unsqueeze(0) # -> (1, 1, n, dv) 6## SDPA output: (1, 1, 1, dv) -> squeeze -> (dv,) 7out_ref = F.scaled_dot_product_attention( 8q_, k_, v_, scale=scale).squeeze() The .squeeze() removes all size-1 dimensions, giving →ref∈ℝdv out_ref ^d_v for direct comparison with moa_decode output. Table 1 reports the verification. Errors are float32 round-off only, an order of magnitude smaller than the training kernel (9×10−59× 10^-5) because the decode path has no backward accumulation. Table 1: Verification: moa_decode vs. PyTorch SDPA. dk=dv=64d_k=d_v=64, float32 (32-bit floating point, 4 bytes per element), seed(0). n ‖err‖∞\|err\|_∞ 4 1.79×10−71.79× 10^-7 ✓ 64 5.96×10−85.96× 10^-8 ✓ 256 4.47×10−84.47× 10^-8 ✓ 1024 5.96×10−85.96× 10^-8 ✓ 4096 2.98×10−82.98× 10^-8 ✓ 4 C/OpenACC Kernel 4.1 ONF Stride Arithmetic Applying the γ primitive (Definition 2.1 of [1]) to row-major storage: γ(⟨l,j⟩,⟨n,dk⟩)=l×dk+j(K offset),γ( l,j , n,d_k )=l× d_k+j (K offset), γ(⟨l,d⟩,⟨n,dv⟩)=l×dv+d(V offset).γ( l,d , n,d_v )=l× d_v+d (V offset). These stride expressions appear verbatim as K[l*dk+j] and V[l*dv+d] in Listing 3. Concrete example (n=3n=3, dk=2d_k=2, dv=2d_v=2). The six elements of K (shape ⟨3,2⟩ 3,2 , row-major) lie at flat offsets: ⟨0,0⟩:γ(⟨0,0⟩,⟨3,2⟩)=0×2+0=0⟨0,1⟩:0×2+1=1⟨1,0⟩:1×2+0=2⟨1,1⟩:1×2+1=3⟨2,0⟩:2×2+0=4⟨2,1⟩:2×2+1=5 array[]l 0,0 :&γ( 0,0 , 3,2 )=0× 2+0=0\\ 0,1 :&0× 2+1=1\\ 1,0 :&1× 2+0=2\\ 1,1 :&1× 2+1=3\\ 2,0 :&2× 2+0=4\\ 2,1 :&2× 2+1=5 array So K[l*2+j] in C accesses ⟨l,j⟩ψK l,j ψ\,K directly – no transpose buffer is needed. 4.2 Dimension Lifting and Hardware-Coalescing Proof Dimension lifting maps a flat index space ⟨N⟩ N to a structured processing array ⟨Π,⌈N/Π⌉⟩ , N/ , partitioning work across hierarchical execution blocks. In modern streaming multiprocessors (GPUs), memory throughput is governed strictly by hardware coalescing: global memory requests across a warp — a group of 32 threads that execute in lock-step on a GPU streaming multiprocessor — are merged if those threads access contiguous memory locations. Let the core evaluation of Step I map an abstract 2-D sequence-to-feature reduction domain D=⟨n⟩×⟨dv⟩D= n × d_v . In parallel architectures, an MoA compiler solves layout compatibility by formalizing a hardware mapping tree. We define multi-level lifting primitives: ⟨n⟩ n ↦⟨Πg,⌈n/Πg⌉⟩(Sequence Rows), \; \; _g,\, n/ _g (Sequence Rows), (4) ⟨dv⟩ d_v ↦⟨Πv,⌈dv/Πv⌉⟩(Feature Columns). \; \; _v,\, d_v/ _v (Feature Columns). (5) The combined domain is reshaped via ρ to mirror the GPU topology: S=⟨Πg,⌈n/Πg⌉,Πv,⌈dv/Πv⌉⟩,S= _g, n/ _g , _v, d_v/ _v , A coordinate index vector in this space is I=⟨g,i,v,j⟩I= g,i,v,j , where g is the gang, i is sequential work inside the block, v is the vector thread (warp channel), and j is the register stride. To detect memory bottlenecks, the compiler evaluates the γ stride between adjacent vector threads v and v+1v+1 for a fixed (g,i,j)(g,i,j). Under row-major layout with the hardware-to-logical mapping l=g⌈n/Πg⌉+il=g n/ _g +i, d=v⌈dv/Πv⌉+jd=v d_v/ _v +j: γ(⟨g,i,v,j⟩,S)=(g⌈n/Πg⌉+i)×dv+v⌈dv/Πv⌉+j.γ( g,i,v,j ,S)= (g n/ _g +i )× d_v+v d_v/ _v +j. The adjacent-thread γ difference is: γ(⟨g,i,v+1,j⟩,S)−γ(⟨g,i,v,j⟩,S)=⌈dv/Πv⌉.γ( g,i,v+1,j ,S)-γ( g,i,v,j ,S)= d_v/ _v . Coalescing Criterion (in MoA γ terms): access is perfectly coalesced if and only if γ(⟨g,i,v+1,j⟩,S)−γ(⟨g,i,v,j⟩,S)=1,γ( g,i,v+1,j ,S)-γ( g,i,v,j ,S)=1, i.e., adjacent vector threads access physically adjacent memory locations (Definition 2.1 of [1]). In the naive DNF, d maps to gangs and l to vectors. The adjacent-thread γ difference becomes ⌈dv/Πv⌉=dv=64 d_v/ _v =d_v=64, meaning threads hit data 256 bytes apart — shattering coalescing. By transposing the execution array (hoisting the d loop as the inner vectorized axis), the MoA compiler sets Πv=dv _v=d_v, so ⌈dv/Πv⌉=1 d_v/ _v =1, and the γ difference equals 1 — aligning threads with contiguous 128-bit cache lines. The transformed, fully coalesced algorithm is implemented in Listing 3; this coalescing strategy follows the GPU energy-efficiency approach of [5]. The column space loop over d is hoisted as the inner vectorized axis, ensuring maximal memory saturation. 4.3 OpenACC Annotations Three OpenACC pragma types appear in Listing 3: • #pragma acc parallel loop gang distributes the outer loop across GPU gangs (groups of warps). • #pragma acc loop vector distributes the inner loop across threads within a gang (the vector lane). • #pragma acc atomic ensures that when multiple gangs update the same out[d] element simultaneously (Pass 3, the weighted sum), the read-modify-write is indivisible — preventing a race condition (a data corruption that occurs when two threads read, then both write, the same location without coordination). The reduction clauses in Passes 1 and 2 use a more efficient tree-reduction pattern rather than serialising through a single atomic; Pass 3 requires atomic because each out[d] is accumulated from all n gang iterations. ⬇ 1void moa_decode_acc( 2const float *restrict q, /* rho(q) = <dk> */ 3const float *restrict K, /* rho(K) = <n,dk> */ 4const float *restrict V, /* rho(V) = <n,dv> */ 5float *restrict out, /* rho(out)= <dv> */ 6int n, int dk, int dv) 7 8const float sc = 1.0f / sqrtf((float)dk); 9float *s = malloc(n * sizeof(float)); /* rho(s)=<n> */ 10float *e = malloc(n * sizeof(float)); /* rho(e)=<n> */ 11/* Pass 1: scores; l-loop->gangs, j-loop->vector 12 ONF: gamma(<l,j>,<n,dk>) = l*dk+j */ 13float m = -1e38f; 14#pragma acc data copyin(q[0:dk],K[0:n*dk],V[0:n*dv]) \ 15 copyout(out[0:dv]) create(s[0:n],e[0:n]) 16 17#pragma acc parallel loop gang reduction(max:m) 18for (int l = 0; l < n; l++) 19float dot = 0.0f; 20#pragma acc loop vector reduction(+:dot) 21for (int j = 0; j < dk; j++) 22dot += q[j] * K[l*dk+j]; /* ONF: l*dk+j */ 23s[l] = sc * dot; 24if (s[l] > m) m = s[l]; 25 26 27/* Pass 2: exponentiate and sum denominator Z */ 28float Z = 0.0f; 29#pragma acc parallel loop gang reduction(+:Z) 30for (int l = 0; l < n; l++) e[l] = expf(s[l]-m); Z += e[l]; 31/* Pass 3: weighted sum; d-loop->vector (coalesced) 32 ONF: gamma(<l,d>,<n,dv>) = l*dv+d 33 gamma-diff(v,v+1) = 1: coalesced */ 34float inv_Z = 1.0f / Z; 35 36#pragma acc parallel loop gang 37for (int d = 0; d < dv; d++) 38out[d] = 0.0f; 39 40#pragma acc parallel loop gang 41for (int l = 0; l < n; l++) 42float scaled_weight = e[l] * inv_Z; 43#pragma acc loop vector 44for (int d = 0; d < dv; d++) 45#pragma acc atomic 46out[d] += scaled_weight * V[l*dv+d]; /* ONF: l*dv+d coalesced */ 47 48 49 /* end acc data region */ 50free(s); free(e); 51 Listing 3: moa_decode_openacc.c: core GPU kernel. ONF stride arithmetic and hardware-coalesced loops. Compiles with clang -O3 -o moa_decode_cpu moa_decode_openacc.c -lm on macOS/Linux (no OpenMP needed; #pragma acc ignored by plain C compilers). 4.4 Verification Compiled with clang -O3 (macOS) or gcc -O3 (Linux) — no flags needed beyond -lm for CPU verification. The #pragma acc directives are silently ignored by plain C compilers; the moa_decode_cpu sequential reference runs in the same binary. Table 2: C/OpenACC kernel: CPU verification, float32. ‖err‖∞=0\|err\|_∞=0 (exact match; same IEEE-754 operation order). n MoA (MB) Std n2n^2 (MB) k 1,024 0.52 4.19 8×8× 4,096 2.10 67.11 32×32× The CPU kernel matches the sequential reference to ‖err‖∞=0\|err\|_∞=0 (exact IEEE-754 equality for the same operation order). For GPU execution: replace clang with nvc -acc=gpu -gpu=c120 (NVIDIA B200), clang -fopenacc --offload-arch=gfx942 (AMD MI300X), or icx -fopenacc (Intel Aurora). 5 Multi-Step KV-Cache Accumulation 5.1 Cache Append Expression At each decode step t, one new key vector k→t∈ℝdk k_t ^d_k, ρ(k→t)=⟨dk⟩ρ( k_t)= d_k , and value vector v→t∈ℝdv v_t ^d_v, ρ(v→t)=⟨dv⟩ρ( v_t)= d_v , are appended to the caches using MoA concatenation #\# along axis 0, defined as: A#b→A\;\#\; b appends rank-1 vector b→ b (shape ⟨dk⟩ d_k ) as a new final row of rank-2 array A (shape ⟨t,dk⟩ t,d_k ), giving ρ(A#b→)=⟨t+1,dk⟩ρ(A\;\#\; b)= t+1,d_k without copying any existing data: Kt+1 K_t+1 =Kt#k→t, =K_t\;\#\; k_t, (6) Vt+1 V_t+1 =Vt#v→t. =V_t\;\#\; v_t. (7) In MoA there are no row or column vectors, so no reshape is needed: #\# appends k→t k_t (shape ⟨dk⟩ d_k ) as a new final row of KtK_t (shape ⟨t,dk⟩ t,d_k ), giving ρ(Kt+1)=⟨t+1,dk⟩ρ(K_t+1)= t+1,d_k . The trailing dimension dkd_k matches, so #\# is well-defined without Omega lifting. Omega would be needed only if #\# were being lifted over a batch of such pairs. Concrete example (t=2t=2, dk=3d_k=3). Let kij≡⟨i,j⟩ψK2k_ij≡ i,j ψ\,K_2. K2K_2 has shape ⟨2,3⟩ 2,3 ; k→2=⟨a,b,c⟩ k_2= a,b,c has shape ⟨3⟩ 3 : K2=(k00k01k02k10k11k12),ρ(K2)=⟨2,3⟩.K_2= pmatrixk_00&k_01&k_02\\ k_10&k_11&k_12 pmatrix, ρ(K_2)= 2,3 . K3≡K2#k→2K_3≡ K_2\;\#\; k_2 is defined as the result of appending k→2 k_2 as the new final row of K2K_2: K3=K2#k→2=(k00k01k02k10k11k12abc).K_3=K_2\;\#\; k_2= pmatrixk_00&k_01&k_02\\ k_10&k_11&k_12\\ a&b&c pmatrix. ρ(K3)=⟨3,3⟩ρ(K_3)= 3,3 . In the ONF, the new row is written at the gamma offset γ(⟨t,j⟩,⟨T,dk⟩)=t×dk+jγ( t,j , T,d_k )=t× d_k+j, so elements a, b, c land at offsets 2×3+0=62× 3+0=6, 77, 88 in the flat array. No existing data is moved. Proposition 2 (Cache-append traffic). Equations (6)–(7) append k→t k_t and v→t v_t directly, moving exactly (dk+dv)×4B(d_k+d_v)× 4\,\,B per step, independent of t. This is the information-theoretic minimum: one key row and one value row must be written regardless of algorithm. 5.2 Combined Per-Step Cost The total traffic at step t is: Mstep(t) M_step(t) =(dk+dv)×4⏟append O(dk+dv)+(tdk+tdv)×4⏟decode O(tdk+tdv). = (d_k+d_v)× 4_append O(d_k+d_v)+ (td_k+td_v)× 4_decode O(td_k+td_v). (8) Summed over T output tokens: Mtotal=∑t=1TMstep(t)=O(T2(dk+dv)).M_total= _t=1^TM_step(t)=O\! (T^2(d_k+d_v) ). This quadratic growth in T is unavoidable (every token must attend over all previous tokens), but MoA makes the constant factor minimal and the bound a theorem. 5.3 Python Implementation Standard PyTorch KV-cache uses torch.cat to grow the cache at each step, materialising K⊤K at each attention call: ⬇ 1K_cache = torch.zeros(0, dk) 2V_cache = torch.zeros(0, dv) 3for t in range(T): 4 # Append -- torch.cat copies the whole cache 5 K_cache = torch.cat([K_cache, keys[t:t+1]], dim=0) 6 V_cache = torch.cat([V_cache, values[t:t+1]], dim=0) 7 q_ = queries[t].unsqueeze(0) # (1,dk) 8 scores = scale * (q_ @ K_cache.T) # K^T buffer 9 weights = torch.softmax(scores, -1) 10 outputs[t] = (weights @ V_cache).squeeze(0) Listing 4: Standard autoregressive decode (pytorch_reference.py: std_kvcache_decode). Uses torch.cat to grow cache; no algebraic minimality proof. The MoA version replaces torch.cat with plain #\# concatenation (writing one row at the pre-computed γ offset, no copy) and proves append cost is O(dk+dv)O(d_k+d_v) per step. The MoA cache-append expressions (eqs. (6)–(7)) and per-step traffic bound (Proposition 2) are: Kt+1 K_t+1 =Kt#k→t,ρ(Kt+1)=⟨t+1,dk⟩, =K_t\;\#\; k_t, ρ(K_t+1)= t+1,d_k , Vt+1 V_t+1 =Vt#v→t,ρ(Vt+1)=⟨t+1,dv⟩, =V_t\;\#\; v_t, ρ(V_t+1)= t+1,d_v , Mappend M_append =(dk+dv)×4B(constant in t). =(d_k+d_v)× 4\,B (constant in $t$). The implementation pre-allocates to max_len and writes each new row at the ONF gamma offset t×dkt× d_k: ⬇ 1class MoAKVCache: 2def __init__(self, dk, dv, max_len=32768): 3self._K = torch.zeros(max_len, dk) 4self._V = torch.zeros(max_len, dv) 5self._t = 0 6self.dk, self.dv = dk, dv 7@property 8def K(self): return self._K[:self._t] 9@property 10def V(self): return self._V[:self._t] 11def append(self, k, v): 12# Write at gamma(t; max_len, dk) = tdk 13# O(dk + dv) traffic, independent of t 14self._K[self._t] = k # ONF: gamma(t,j;<T,dk>)=t*dk+j 15self._V[self._t] = v # ONF: gamma(t,d;<T,dv>)=t*dv+d 16self._t += 1 Listing 5: moa_kvcache.py: pre-allocated cache with O(dk+dv)O(d_k+d_v) append. MoAKVCache.append writes at the gamma offset t*dk without copying the existing cache. 5.4 Verification Table 3: KV-cache multi-step decode: moa_autoregressive_decode vs. causal SDPA (Scaled Dot-Product Attention). dk=dv=64d_k=d_v=64, float32. T dkd_k dvd_v ‖err‖∞\|err\|_∞ 4 8 8 1.19×10−71.19× 10^-7 ✓ 16 64 64 4.77×10−74.77× 10^-7 ✓ 64 64 64 5.96×10−75.96× 10^-7 ✓ Table 4: KV-cache traffic accounting: append cost is constant; decode cost grows as O(tdk)O(td_k) per step. MoA total beats the n2n^2 score-matrix cost for T≳256T 256; at small T the n2n^2 term is also small. T App. (MB) Dec. (MB) n2n^2 (MB) 16 0.01 0.08 0.01 64 0.03 1.10 0.36 256 0.13 16.97 22.50 1024 0.52 269.22 1433.75 At T=1024T=1024 the MoA total (≈ 270 MB) is 5.3×5.3× less than the score-matrix cost alone (1434 MB), with the gap widening quadratically in T. hq=8h_q=8 query headsgroup 0 (g=0,1,2,3→gkv=0g\!=\!0,1,2,3→ g_kv\!=\!0) g=0 gkv=0g_kv=0 g=1 gkv=0g_kv=0 g=2 gkv=0g_kv=0 g=3 gkv=0g_kv=0 group 1 (g=4,5,6,7→gkv=1g\!=\!4,5,6,7→ g_kv\!=\!1) g=4 gkv=1g_kv=1 g=5 gkv=1g_kv=1 g=6 gkv=1g_kv=1 g=7 gkv=1g_kv=1 gkv=⌊g×hkv/hq⌋g_kv= g\!×\!h_kv/h_q four read ⟨0⟩ψK 0 ψ\,K and ⟨0⟩ψV 0 ψ\,Vno K⊤K , no broadcastall four read ⟨1⟩ψK 1 ψ\,K and ⟨1⟩ψV 1 ψ\,Vno K⊤K , no broadcasthkv=2h_kv=2 KV slabs ⟨0⟩ψK 0 ψ\,K, ⟨0⟩ψV 0 ψ\,V ρ=⟨n,dk⟩ρ= n,d_k , ⟨n,dv⟩ n,d_v shared by g=0,1,2,3g\!=\!0,1,2,3 read once, not 4 times ⟨1⟩ψK 1 ψ\,K, ⟨1⟩ψV 1 ψ\,V ρ=⟨n,dk⟩ρ= n,d_k , ⟨n,dv⟩ n,d_v shared by g=4,5,6,7g\!=\!4,5,6,7 read once, not 4 times ⟨0⟩ψK 0 ψ K⟨1⟩ψK 1 ψ K Conventional (repeat_interleave) K: ⟨2,n,dk⟩→⟨8,n,dk⟩ 2,n,d_k \!→\! 8,n,d_k (4× copy) traffic: hq×n×dk×4h_q\!×\!n\!×\!d_k\!×\!4 B, no proof MoA (ψ-selection) K stays ⟨2,n,dk⟩ 2,n,d_k (no copy) traffic: hkv×n×dk×4h_kv\!×\!n\!×\!d_k\!×\!4 B ✓proven hq/hkv=4×h_q/h_kv=4×less traffic MQA (hkv=1h_kv\!=\!1): hq×h_q× less traffic. GQA general saving: hq/hkvh_q/h_kv. Same output. Storage Thm 2.7 [1]. Figure 3: GQA/MQA via MoA ψ-selection. hq=8h_q\!=\!8 query heads are grouped into hkv=2h_kv\!=\!2 KV groups. ψ-selection ⟨gkv⟩ψK g_kv ψ\,K maps each query head to its KV slab directly, without materialising the ⟨hq,n,dk⟩ h_q,n,d_k broadcast that repeat_interleave would create. Traffic reduction: hq/hkv=4×h_q/h_kv\!=\!4× here; hq×h_q× for MQA (hkv=1h_kv\!=\!1). Proven minimal by Storage Theorem 2.7. 6 Grouped-Query and Multi-Query Attention 6.1 Psi-Selection Expression In GQA, hqh_q query heads share hkv<hqh_kv<h_q KV heads (hkvh_kv divides hqh_q), with K of shape ⟨hkv,n,dk⟩ h_kv,n,d_k . The scalar elements of K are defined by ψ-selection with a full index vector: ⟨g,i,j⟩ψK0≤g<hkv(KV head)0≤i<n(sequence position)0≤j<dk(feature) g,i,j \,ψ\,K cases0≤ g<h_kv&(KV head)\\ 0≤ i<n&(sequence position)\\ 0≤ j<d_k&(feature) cases (9) A partial index vector ⟨gkv⟩ g_kv selects the entire gkvg_kv-th KV head slab: Kg=⟨gkv⟩ψK,ρ(Kg)=⟨n,dk⟩,K_g= g_kv \,ψ\,K, ρ(K_g)= n,d_k , (10) whose elements satisfy ⟨i,j⟩ψ(⟨gkv⟩ψK)=⟨gkv,i,j⟩ψK i,j \,ψ\,( g_kv \,ψ\,K)= g_kv,i,j \,ψ\,K. The group assignment for query head g is gkv=⌊g×hkv/hq⌋g_kv= g× h_kv/h_q . This drops the leading head dimension, leaving a 2-D array of shape ⟨n,dk⟩ n,d_k passed directly to the decode DNF (equation (1)). No Omega lifting is needed and no expanded ⟨hq,n,dk⟩ h_q,n,d_k tensor is materialized. Concrete example (hq=4h_q=4, hkv=2h_kv=2, n=3n=3, dk=2d_k=2, _=2group\_size=2). K has shape ⟨2,3,2⟩ 2,3,2 with elements ⟨g,i,j⟩ψK g,i,j \,ψ\,K. The two KV slabs (each shape ⟨3,2⟩ 3,2 ) are: ⟨0⟩ψK=(⟨0,0,0⟩ψK⟨0,0,1⟩ψK⟨0,1,0⟩ψK⟨0,1,1⟩ψK⟨0,2,0⟩ψK⟨0,2,1⟩ψK) 0 \,ψ\,K= pmatrix 0,0,0 \,ψ\,K& 0,0,1 \,ψ\,K\\ 0,1,0 \,ψ\,K& 0,1,1 \,ψ\,K\\ 0,2,0 \,ψ\,K& 0,2,1 \,ψ\,K pmatrix ⟨1⟩ψK=(⟨1,0,0⟩ψK⟨1,0,1⟩ψK⟨1,1,0⟩ψK⟨1,1,1⟩ψK⟨1,2,0⟩ψK⟨1,2,1⟩ψK) 1 \,ψ\,K= pmatrix 1,0,0 \,ψ\,K& 1,0,1 \,ψ\,K\\ 1,1,0 \,ψ\,K& 1,1,1 \,ψ\,K\\ 1,2,0 \,ψ\,K& 1,2,1 \,ψ\,K pmatrix The four query heads select slabs as follows: g=0: g=0: gkv=0,⟨0⟩ψK,ρ=⟨3,2⟩ \;g_kv=0, 0 \,ψ\,K,\;ρ= 3,2 g=1: g=1: gkv=0,⟨0⟩ψK(same slab as g=0) \;g_kv=0, 0 \,ψ\,K (same slab as $g=0$) g=2: g=2: gkv=1,⟨1⟩ψK,ρ=⟨3,2⟩ \;g_kv=1, 1 \,ψ\,K,\;ρ= 3,2 g=3: g=3: gkv=1,⟨1⟩ψK(same slab as g=2) \;g_kv=1, 1 \,ψ\,K (same slab as $g=2$) Query heads 0 and 1 both evaluate ⟨0⟩ψK 0 \,ψ\,K; heads 2 and 3 both evaluate ⟨1⟩ψK 1 \,ψ\,K. By the Storage Theorem, each slab need only be read from DRAM once, giving the hq/hkv=2×h_q/h_kv=2× KV traffic reduction algebraically. 6.2 Dimension Lifting Dimension lifting is MoA’s way of mapping to the abstract shape of the target architecture. Starting with the ONF shape, each component of the shape is partitioned one or more times to conform to each architectural target, e.g. processors, memory, etc. For example, The q-head axis and kvkv-head axis are lifted independently using the dimension-lifting map N↦⟨Π,⌈N/Π⌉⟩N , N/ of Section 2.5 of [1]: ⟨hq⟩ h_q ↦⟨Πq,⌈hq/Πq⌉⟩(q-head gang axis), _q, h_q/ _q ($q$-head gang axis), (11) ⟨hkv⟩ h_kv ↦⟨Πkv,⌈hkv/Πkv⌉⟩(KV read axis). _kv, h_kv/ _kv (KV read axis). (12) With shared-memory caching (⟨gkv⟩ψK g_kv \,ψ\,K held in L1 for the hq/hkvh_q/h_kv gangs sharing it), the KV traffic reduces from hq×n×dkh_q× n× d_k to hkv×n×dkh_kv× n× d_k bytes – a factor of hq/hkvh_q/h_kv. MQA is the special case hkv=1h_kv=1. 6.3 Python Implementation Standard PyTorch GQA expands the KV heads to hqh_q by repeat_interleave, materialising a full ⟨hq,n,dk⟩ h_q,n,d_k broadcast tensor: ⬇ 1def std_gqa_decode(Q, K, V, scale=None): 2 group_size = Q.shape[0] // K.shape[0] 3 # Expand: (hkv,n,dk)->(hq,n,dk) -- materialised 4 K_exp = K.repeat_interleave(group_size, dim=0) 5 V_exp = V.repeat_interleave(group_size, dim=0) 6 for g in range(Q.shape[0]): 7 scores = scale*(Q[g:g+1] @ K_exp[g].T) 8 weights = torch.softmax(scores, dim=-1) 9 outputs[g] = (weights @ V_exp[g]).squeeze(0) Listing 6: Standard GQA decode (pytorch_reference.py: std_gqa_decode). Materialises the KV broadcast; MoA avoids this via ψ-selection. The MoA version replaces repeat_interleave with ψ-selection ⟨gkv⟩ψK g_kv ψ\,K, reading each KV slab once per group rather than hq/hkvh_q/h_kv times, giving a proven hq/hkvh_q/h_kv traffic reduction. The MoA formulation for GQA uses ψ-selection (eq. (10)) to map each query head g to its KV group, then applies the decode DNF: gkv g_kv =⌊g×hkv/hq⌋,0≤g<hq, = g× h_kv/h_q , 0≤ g<h_q, Kg K_g =⟨gkv⟩ψK,ρ(Kg)=⟨n,dk⟩, = g_kv \,ψ\,K, ρ(K_g)= n,d_k , →g out_g =a→g(+.×Ω⟨1,2⟩)Vg,ρ(→g)=⟨dv⟩. = a_g\;(+.× _ 1,2 )\;V_g, ρ( out_g)= d_v . No ⟨hq,n,dk⟩ h_q,n,d_k tensor is materialized; each query head reads its assigned slab directly: ⬇ 1def moa_gqa_decode(Q, K, V, scale=None): 2# Q : (h_q, dk) 3# K : (h_kv, n, dk) -- one slab per KV head 4# V : (h_kv, n, dv) 5h_q, dk = Q.shape 6h_kv = K.shape[0] 7if scale is None: 8scale = 1.0 / math.sqrt(dk) 9group_size = h_q // h_kv 10out = torch.zeros(h_q, V.shape[2]) 11for g in range(h_q): 12g_kv = g // group_size # psi-selection: <g_kv> psi K 13# <g_kv> psi K shared across group_size query heads 14out[g] = _decode_one_head( 15Q[g], K[g_kv], V[g_kv], scale) 16return out Listing 7: moa_gqa.py: GQA decode. The ψ-selection eq. (10) is g_kv = g // group_size; no repeat_interleave materialization. 6.4 Verification and Traffic Table 5: GQA/MQA verification: moa_gqa_decode vs. per-head SDPA reference. n=64n=64, dk=dv=32d_k=d_v=32, float32. Configuration ‖err‖∞\|err\|_∞ MHA (hq=8,hkv=8h_q=8,h_kv=8) 1.19×10−71.19× 10^-7 ✓ GQA (hq=8,hkv=4h_q=8,h_kv=4) 1.19×10−71.19× 10^-7 ✓ GQA (hq=8,hkv=2h_q=8,h_kv=2) 1.49×10−71.49× 10^-7 ✓ MQA (hq=8,hkv=1h_q=8,h_kv=1) 1.79×10−71.79× 10^-7 ✓ Table 6: GQA/MQA KV traffic per decode step. n=1024n=1024, dk=dv=64d_k=d_v=64, float32, hq=32h_q=32. KV amortized over hkvh_kv reads (shared memory). Config KV (MB) Total (MB) MHA hkv=32h_kv\!=\!32 16.78 16.79 GQA hkv=8h_kv\!=\!8 4.19 4.21 GQA hkv=4h_kv\!=\!4 2.10 2.11 MQA hkv=1h_kv\!=\!1 0.52 0.54 The hq/hkvh_q/h_kv traffic reduction is exact and follows algebraically from the ψ-selection expression (10): at hkv=1h_kv=1 (MQA), KV traffic falls from 16.78 MB to 0.52 MB (32×32×) on top of the baseline MoA decode savings. 7 Summary and Relationship to Training Kernel Training – Algorithm 3 of [1] Inputs: Q,K,V,OQ,K,V,O from DRAM Ω I: GV→GV\!→\!DRAM; Ω I–I: GAGA on-stack, gsg_s in register Ω IV/V: GQ,GK→GQ,GK\!→\!DRAM Traffic: 4ndk+4ndv4nd_k+4nd_v B Scratch: O(n)O(n) Decode – Algorithm 1 (this report) Inputs: q→,K,V q,K,V from DRAM (++ GQA: hkvh_kv heads) Step I: s→ s, ρ(s→)=⟨n⟩ρ( s)= n , on stack, Step I: a→ a, ρ(a→)=⟨n⟩ρ( a)= n , on stack Step I: → out\!→\!DRAM Traffic: O(ndk+ndv)O(nd_k+nd_v) B GQA bonus: ÷(hq/hkv) \,(h_q/h_kv) KV-cache append (Section 5) Write k→t k_t at t×dkt\!×\!d_k, v→t v_t at t×dvt\!×\!d_v Traffic: (dk+dv)×4(d_k+d_v)× 4 B per step – constant in t Figure 4: MoA inference pipeline. All three components are specializations or extensions of the training DNF; no new Omega derivations are needed. Table 7 collects the verified artifacts. Table 7: Verified inference artifacts. All errors float32 round-off only. Artifact File ‖err‖∞\|err\|_∞ Python decode moa_decode.py ≤2×10−7≤ 2× 10^-7 C/OpenACC moa_decode_openacc.c 0 (exact) KV-cache moa_kvcache.py 6×10−76× 10^-7 GQA/MQA moa_gqa.py ≤2×10−7≤ 2× 10^-7 8 Conclusion Four MoA inference artifacts have been derived, implemented, and numerically verified, each following directly from the forward-pass DNF of [1]. The decode DNF is eqs. (4)–(6) of that paper with i′i fixed to the current token: eq. (4) score ψ-reduction, eq. (5) softmax, eq. (6) weighted sum. The C/OpenACC kernel applies the γ ONF stride arithmetic of [1] §7.2; KV-cache append uses plain #\# concatenation; GQA/MQA uses ψ-selection and the dimension-lifting Section 2.5 map of [1]. In every case the memory bound is a theorem (Storage Theorem 2.7 of [1]) before code is written, and the DNF→ → pipeline ports the result to any architecture without re-derivation. References [1] L. Mullin and G. Hains, “Attention at the Theoretical Minimum: A Mathematics of Arrays Framework for Memory-Optimal Transformer Kernels,” arXiv:2606.07713v1 [cs.LG], June 2026. [2] L. Mullin and G. Hains, “Memory-Optimal Transformer Training: Backward Pass and Fused Kernel via the Mathematics of Arrays,” HAL-05659212, June 2026. [3] L. M. Mullin, A Mathematics of Arrays, Ph.D. dissertation, Syracuse University, 1988. [4] S. Thomas, L. Mullin, and K. Świrydowicz, “Improving the Performance of DGEMM with MoA and Cache-Blocking,” in Proc. ACM ARRAY Workshop (PLDI), 2021. https://doi.org/10.1145/1122445.1122456 [5] L. Mullin, “From Array Algebra to Energy Efficiency on GPUs,” arXiv:2306.11148, 2023. [6] A. Vaswani et al., “Attention is all you need,” in Advances in Neural Information Processing Systems (NeurIPS), 2017. [7] N. Shazeer, “Fast Transformer Decoding: One Write-Head is All You Need,” arXiv:1911.02150, 2019. [8] J. Ainslie et al., “GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,” in Proc. EMNLP, p. 4895–4901, Singapore, 2023. [9] T. Dao, D. Y. Fu, S. Ermon, A. Rudra, and C. Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness,” in Advances in Neural Information Processing Systems (NeurIPS), 2022. [10] T. Dao, “FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning,” in Proc. ICLR, 2024. arXiv:2307.08691. [11] T. Dao and D. Y. Fu, “Flash-Decoding for Long-Context Inference,” https://crfm.stanford.edu/2023/10/12/flash-decoding.html, 2023.