nx_nofloat_mla.nx source
↩ module page · 129 lines · 5662 B
1// nx_nofloat_mla.nx -- MLA (Multi-head Latent Attention, DeepSeek's marquee KV-compression) in no-float
2// integer SERVE convention (2026-07-15). The 2026-frontier attention: instead of caching full K,V (2*kv_dim
3// per token), x is DOWN-projected to a small latent c (LC-dim, LC << 2*kv_dim) which is the WHOLE KV cache;
4// K,V are UP-projected from c at attention time. The KEY property (provable by construction): serving from
5// the LC-dim latent cache produces the BYTE-IDENTICAL output as serving from a full-K,V cache -- LOSSLESS
6// compression, ratio (2*kv_dim)/LC. This lib gives BOTH paths so a gate can differential them.
7// NOPE variant (no RoPE) for rung 1 -- decoupled-RoPE (RoPE on a separate dim, since RoPE doesn't commute
8// with the up-projection) is rung 2. Composes proven primitives (mm_out_in / fx_exp / qmul); Q16 throughout,
9// deterministic (exact integer sums, positional tie-breaks). license_tier: ORIGINAL (lib: no main)
10import "nx_syscalls.nx"
11import "nx_nofloat_llm.nx"
12
13// down-project x[T,D] -> latent c[T,LC] (the compressed KV cache). Q16.
14func nmla_downproj(x: *i64, Wdkv: *i64, c: *i64, T: i64, D: i64, LC: i64) -> i64 {
15 mm_out_in(x, Wdkv, c, T, D, LC, 0)
16 return 0
17}
18// up-project latent c[T,LC] -> K[T,kvd], V[T,kvd]. Q16.
19func nmla_upproj(c: *i64, Wuk: *i64, Wuv: *i64, K: *i64, V: *i64, T: i64, LC: i64, kvd: i64) -> i64 {
20 mm_out_in(c, Wuk, K, T, LC, kvd, 0)
21 mm_out_in(c, Wuv, V, T, LC, kvd, 0)
22 return 0
23}
24// shared causal MHA core: Q,K,V are [T, nh*hd]; writes concat[T, nh*hd]. Q16, deterministic.
25func nmla_mha(Q: *i64, K: *i64, V: *i64, concat: *i64, T: i64, nh: i64, hd: i64, scale: i64, sc: *i64, at: *i64) -> i64 {
26 let D: i64 = nh*hd
27 var t: i64 = 0
28 while t < T {
29 var h: i64 = 0
30 while h < nh {
31 let qb: i64 = h*hd
32 var s: i64 = 0
33 while s <= t {
34 var dot: i64 = 0
35 var i: i64 = 0
36 while i < hd { dot = dot + (Q[t*D + qb + i] * K[s*D + qb + i]); i = i + 1 }
37 sc[s] = qmul(dot >> 16, scale)
38 s = s + 1
39 }
40 var mmax: i64 = sc[0]
41 var j: i64 = 1
42 while j <= t { if sc[j] > mmax { mmax = sc[j] } j = j + 1 }
43 var sum: i64 = 0
44 j = 0
45 while j <= t { let e: i64 = fx_exp(sc[j] - mmax); at[j] = e; sum = sum + e; j = j + 1 }
46 if sum < 1 { sum = 1 }
47 j = 0
48 while j <= t { at[j] = (at[j] << 16) / sum; j = j + 1 }
49 var d2: i64 = 0
50 while d2 < hd {
51 var acc: i64 = 0
52 s = 0
53 while s <= t { acc = acc + (at[s] * V[s*D + qb + d2]); s = s + 1 }
54 concat[t*D + qb + d2] = acc >> 16
55 d2 = d2 + 1
56 }
57 h = h + 1
58 }
59 t = t + 1
60 }
61 return 0
62}
63
64// MLA forward (LATENT-cache path): x -> c (the cache) -> up-project K,V from c -> MHA -> Wo. Writes out[T,D].
65// ap: [0]=x(T*D) [1]=out(T*D) [2]=T [3]=D [4]=nh [5]=hd [6]=LC [7]=scale
66// [8]=Wdkv(D*LC) [9]=Wuk(LC*D) [10]=Wuv(LC*D) [11]=Wq(D*D) [12]=Wo(D*D) [13]=scr [14]=c_cache(T*LC)
67// scr layout: Q[T*D] K[T*D] V[T*D] sc[T] at[T] concat[T*D]
68func nmla_forward_latent(ap: *i64) -> i64 {
69 let x: *i64 = ap[0] as *i64
70 let out: *i64 = ap[1] as *i64
71 let T: i64 = ap[2]
72 let D: i64 = ap[3]
73 let nh: i64 = ap[4]
74 let hd: i64 = ap[5]
75 let LC: i64 = ap[6]
76 let scale: i64 = ap[7]
77 let Wdkv: *i64 = ap[8] as *i64
78 let Wuk: *i64 = ap[9] as *i64
79 let Wuv: *i64 = ap[10] as *i64
80 let Wq: *i64 = ap[11] as *i64
81 let Wo: *i64 = ap[12] as *i64
82 let scr: *i64 = ap[13] as *i64
83 let c: *i64 = ap[14] as *i64
84 let Q: *i64 = scr
85 let K: *i64 = ((Q as i64) + T*D*8) as *i64
86 let V: *i64 = ((K as i64) + T*D*8) as *i64
87 let sc: *i64 = ((V as i64) + T*D*8) as *i64
88 let at: *i64 = ((sc as i64) + T*8) as *i64
89 let concat: *i64 = ((at as i64) + T*8) as *i64
90 // the CACHE is c (LC-dim). Everything else is transient compute.
91 nmla_downproj(x, Wdkv, c, T, D, LC)
92 nmla_upproj(c, Wuk, Wuv, K, V, T, LC, D)
93 mm_out_in(x, Wq, Q, T, D, D, 0)
94 nmla_mha(Q, K, V, concat, T, nh, hd, scale, sc, at)
95 mm_out_in(concat, Wo, out, T, D, D, 0)
96 return 0
97}
98// MHA forward (FULL-KV baseline): identical math but the CACHE is full K,V (2*D per token). Same `out`.
99// ap same, but ap[14]=kv_cache(T*2*D). scr: Q[T*D] sc[T] at[T] concat[T*D].
100func nmla_forward_fullkv(ap: *i64) -> i64 {
101 let x: *i64 = ap[0] as *i64
102 let out: *i64 = ap[1] as *i64
103 let T: i64 = ap[2]
104 let D: i64 = ap[3]
105 let nh: i64 = ap[4]
106 let hd: i64 = ap[5]
107 let LC: i64 = ap[6]
108 let scale: i64 = ap[7]
109 let Wdkv: *i64 = ap[8] as *i64
110 let Wuk: *i64 = ap[9] as *i64
111 let Wuv: *i64 = ap[10] as *i64
112 let Wq: *i64 = ap[11] as *i64
113 let Wo: *i64 = ap[12] as *i64
114 let scr: *i64 = ap[13] as *i64
115 let kvc: *i64 = ap[14] as *i64 // full-K,V cache: K at kvc[0..T*D), V at kvc[T*D..2*T*D)
116 let K: *i64 = kvc
117 let V: *i64 = ((kvc as i64) + T*D*8) as *i64
118 let Q: *i64 = scr
119 let sc: *i64 = ((Q as i64) + T*D*8) as *i64
120 let at: *i64 = ((sc as i64) + T*8) as *i64
121 let concat: *i64 = ((at as i64) + T*8) as *i64
122 let ctmp: *i64 = ((concat as i64) + T*D*8) as *i64 // transient latent (down/up once, then store K,V)
123 nmla_downproj(x, Wdkv, ctmp, T, D, LC)
124 nmla_upproj(ctmp, Wuk, Wuv, K, V, T, LC, D) // store FULL K,V (the fat cache)
125 mm_out_in(x, Wq, Q, T, D, D, 0)
126 nmla_mha(Q, K, V, concat, T, nh, hd, scale, sc, at)
127 mm_out_in(concat, Wo, out, T, D, D, 0)
128 return 0
129}