ayo
7c5d3c9157
Add native K/M/RR inference engines, HIP kernels, replay and validation harnesses, and Windows DLSS setup and runtime support. Include GPU profiling tools, build and packaging scripts, platform integration patches, and implementation documentation.
171 lines
7.9 KiB
Python
171 lines
7.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Parametrised numpy model of the rrlite Swin blocks (enc3 spec generalised); checked stage by stage
|
|
against ptxsim MMA captures."""
|
|
import numpy as np
|
|
from model_enc3 import (q8, f16, Params, E4M3, F16, exp_trick, row_sum16, C_SCALE, CLAMP, P1, P0,
|
|
gelu_poly, mma_chain)
|
|
|
|
def gp(c):
|
|
c = np.asarray(c)
|
|
return 32 * (c // 32) + 2 * ((c % 32) // 4) + (c % 2) + 16 * ((c % 4) // 2)
|
|
|
|
def pair_to_nat(M):
|
|
return M[:, gp(np.arange(M.shape[1]))]
|
|
|
|
def nat_to_pair(M):
|
|
out = np.empty_like(M); out[:, gp(np.arange(M.shape[1]))] = M; return out
|
|
|
|
def woff_table(base, Ks, Ns, K, N):
|
|
k = np.arange(K)[:, None]; n = np.arange(N)[None, :]; j = n // 8
|
|
return base + Ks * (k // 32) + Ns * (j // 2) + 64 * (n % 8) + 16 * ((k % 16) // 4) + 8 * (j % 2) + 4 * ((k % 32) // 16) + k % 4
|
|
|
|
def load_tokens(P, bx, by, planes):
|
|
X, Y = P.token_xy(bx, by); X, Y = P.mirror(X, P.W), P.mirror(Y, P.H)
|
|
out = np.zeros((64, 32 * planes))
|
|
for p in range(planes):
|
|
off = P.inp - P.abase + ((p * P.H + Y) * P.W + X) * 32
|
|
out[:, 32 * p:32 * p + 32] = E4M3[P.arena[off[:, None] + np.arange(32)]]
|
|
return out
|
|
|
|
def sumsq_tokens(x):
|
|
"""Kernel-exact f16 sum of squares of natural-order token rows (planes of 32; lane t holds words t, t+4)."""
|
|
sq = f16(x * x); planes = x.shape[1] // 32
|
|
def add(a, b): return f16(a[0] + b[0]), f16(a[1] + b[1])
|
|
lanes = []
|
|
for t in range(4):
|
|
def pair(pp, word, k):
|
|
c = 32 * pp + 4 * word + 2 * k
|
|
return sq[:, c], sq[:, c + 1]
|
|
def plane_sum(word):
|
|
p = [add(pair(pp, word, 0), pair(pp, word, 1)) for pp in range(planes)]
|
|
if planes == 4: return add(add(p[0], p[1]), add(p[2], p[3]))
|
|
if planes == 3: return add(add(p[0], p[1]), p[2])
|
|
if planes == 2: return add(p[0], p[1])
|
|
raise ValueError(planes)
|
|
lanes.append(add(plane_sum(t), plane_sum(t + 4)))
|
|
a = [(f16(lanes[t][0] + lanes[t ^ 1][0]), f16(lanes[t][1] + lanes[t ^ 1][1])) for t in range(4)]
|
|
b = (f16(a[0][0] + a[2][0]), f16(a[0][1] + a[2][1]))
|
|
return f16(b[0] + b[1])[:, None]
|
|
|
|
def rms_from_sum(x, s, gamma):
|
|
r = f16(1.0 / np.sqrt(f16(s + f16(2.0 ** -13))))
|
|
return q8(f16(x * f16(r * gamma[None, :])))
|
|
|
|
def rel_bias(P, off):
|
|
raw = P.f16vec(off, 256)
|
|
i, j = np.meshgrid(np.arange(16), np.arange(16), indexing='ij')
|
|
lane = 4 * (i % 8) + (j % 8) // 2; word = i // 8 + 2 * (j // 8)
|
|
return raw[8 * lane + 2 * word + (j % 2)]
|
|
|
|
def attention(Qq, V, bias):
|
|
S = f16(Qq @ Qq.T + bias)
|
|
t = np.clip(f16(S * C_SCALE), -CLAMP, CLAMP)
|
|
e1 = f16(t * (-t) + P1); poly = f16(t * e1 + P0)
|
|
w = exp_trick(poly); s = row_sum16(w)
|
|
r = f16(np.float32(1.0) / s.astype(np.float32)); Pm = f16(w * r)
|
|
return f16(Pm @ V)
|
|
|
|
def tree_norm2(x):
|
|
"""Kernel-exact f16 sum of squares of pair-ordered C-fragment rows (NT8 = C/8 n8 tiles)."""
|
|
sq = f16(x * x); nt8 = x.shape[1] // 8
|
|
def add(a, b): return f16(a[0] + b[0]), f16(a[1] + b[1])
|
|
lanes = []
|
|
for t in range(4):
|
|
T = [(sq[:, 8 * j + 2 * t], sq[:, 8 * j + 2 * t + 1]) for j in range(nt8)]
|
|
def quad(js): return add(add(T[js[0]], T[js[1]]), add(T[js[2]], T[js[3]]))
|
|
if nt8 == 16:
|
|
even = add(quad([0, 2, 4, 6]), quad([8, 10, 12, 14])); odd = add(quad([1, 3, 5, 7]), quad([9, 11, 13, 15]))
|
|
elif nt8 == 12:
|
|
even = add(quad([0, 2, 4, 6]), add(T[8], T[10])); odd = add(quad([1, 3, 5, 7]), add(T[9], T[11]))
|
|
elif nt8 == 8:
|
|
even = quad([0, 2, 4, 6]); odd = quad([1, 3, 5, 7])
|
|
else:
|
|
raise ValueError(nt8)
|
|
lanes.append(add(even, odd))
|
|
a = [(f16(lanes[t][0] + lanes[t ^ 1][0]), f16(lanes[t][1] + lanes[t ^ 1][1])) for t in range(4)]
|
|
b = (f16(a[0][0] + a[2][0]), f16(a[0][1] + a[2][1]))
|
|
return f16(b[0] + b[1])[:, None]
|
|
|
|
def layout(C, heads, pre=0):
|
|
"""Weight offsets of a Swin block with C channels starting at byte `pre`."""
|
|
HS = 96 * C + 512
|
|
L = dict(g1=pre, head=[pre + 2 * C + HS * h for h in range(heads)])
|
|
E = pre + 2 * C + heads * HS
|
|
L.update(bo=E, g2=E + 2 * C)
|
|
M0 = E + 4 * C; NC = C // 8
|
|
L['b1'] = [M0 if c == 0 else M0 + 66 * C + 64 + (64 * C + 64) * (c - 1) for c in range(NC)]
|
|
L['b2'] = M0 + 64 * C + 64
|
|
L['end'] = M0 + 66 * C + 64 + (64 * C + 64) * (NC - 1)
|
|
return L
|
|
|
|
def swin_block(P, x0, C, heads, pre=0, stages=None):
|
|
"""x0: 64 x C natural (f16 values). Returns (x2 pair order, out natural fp8 values)."""
|
|
fp8 = lambda T: E4M3[P.wbuf[T]]
|
|
L = layout(C, heads, pre)
|
|
h1 = rms_from_sum(x0, sumsq_tokens(x0), P.f16vec(L['g1'], C)[gp(np.arange(C))])
|
|
O = np.zeros((64, 32 * heads))
|
|
for h in range(heads):
|
|
hb = L['head'][h]
|
|
Q = mma_chain(h1, fp8(woff_table(hb, 1024, 512, C, 32)))
|
|
V = mma_chain(h1, fp8(woff_table(hb + 32 * C, 1024, 512, C, 32)))
|
|
B = rel_bias(P, hb + 64 * C)
|
|
for m in range(4):
|
|
O[16 * m:16 * m + 16, 32 * h:32 * h + 32] = attention(q8(Q[16 * m:16 * m + 16]), V[16 * m:16 * m + 16], B)
|
|
Wo = np.zeros((32 * heads, C), np.int64)
|
|
for s in range(heads):
|
|
Wo[32 * s:32 * s + 32] = woff_table(L['head'][s] + 64 * C + 512, 0, 512, 32, C)
|
|
x1 = mma_chain(pair_to_nat(q8(O)), fp8(Wo), C0=f16(nat_to_pair(x0) + P.f16vec(L['bo'], C)[None, :]))
|
|
h2 = pair_to_nat(rms_from_sum(x1, tree_norm2(x1), P.f16vec(L['g2'], C)))
|
|
if stages is not None:
|
|
stages.update(h1=h1, x1=pair_to_nat(x1), h2=h2, O=pair_to_nat(q8(O)))
|
|
x2 = f16(x1 + P.f16vec(L['b2'], C)[None, :])
|
|
idx = gp(np.arange(32))
|
|
for c, B1 in enumerate(L['b1']):
|
|
z = mma_chain(h2, fp8(woff_table(B1, 512, 16 * C, C, 32)), C0=np.broadcast_to(P.f16vec(B1 + 32 * C, 32), (64, 32)).copy())
|
|
g = q8(gelu_poly(z))[:, idx]
|
|
x2 = f16(g @ fp8(woff_table(B1 + 32 * C + 64, 0, 512, 32, C)) + x2)
|
|
return x2, q8(pair_to_nat(x2)), L
|
|
|
|
TOK_X = (np.arange(64) & 3) + 4 * ((np.arange(64) >> 4) & 1)
|
|
TOK_Y = ((np.arange(64) >> 2) & 3) + 4 * (np.arange(64) >> 5)
|
|
|
|
def patch_merge_input(out, C):
|
|
"""16 merged rows (r -> (r%4, r//4)) x 4C: column C*q + ch = out[token(2mx+q%2, 2my+q//2), ch]."""
|
|
A = np.zeros((16, 4 * C))
|
|
for r in range(16):
|
|
mx, my = r % 4, r // 4
|
|
for q in range(4):
|
|
t = np.nonzero((TOK_X == 2 * mx + (q & 1)) & (TOK_Y == 2 * my + (q >> 1)))[0][0]
|
|
A[r, C * q:C * q + C] = out[t]
|
|
return A
|
|
|
|
def load_at(P, ptr, bx, by, planes, W=None, H=None, low=False):
|
|
W = W or P.W; H = H or P.H
|
|
X, Y = P.token_xy(bx, by); X, Y = P.mirror(X, W), P.mirror(Y, H)
|
|
out = np.zeros((64, 32 * planes))
|
|
for p in range(planes):
|
|
off = ptr - P.abase + ((p * H + Y) * W + X) * 32
|
|
out[:, 32 * p:32 * p + 32] = E4M3[P.arena[off[:, None] + np.arange(32)]]
|
|
return out
|
|
|
|
def decoder_input(P, bx, by, C, CIN, NW, skip_ptr):
|
|
"""x0 = f16(q8(patch_expand(low-res CIN)) + skip), expand output col G = (4C/NW) w + n -> q = G // C, ch = ginv(G % C)."""
|
|
fp8 = lambda T: E4M3[P.wbuf[T]]
|
|
W2, H2 = P.W // 2, P.H // 2
|
|
A = np.zeros((16, CIN))
|
|
for r in range(16):
|
|
MX = P.mirror(np.array((8 * bx - P.sx) // 2 + r % 4), W2); MY = P.mirror(np.array((8 * by - P.sy) // 2 + r // 4), H2)
|
|
for p in range(CIN // 32):
|
|
off = P.inp - P.abase + ((p * H2 + MY) * W2 + MX) * 32
|
|
A[r, 32 * p:32 * p + 32] = E4M3[P.arena[off:off + 32]]
|
|
cols = 4 * C // NW
|
|
E_ = np.concatenate([mma_chain(A, fp8(woff_table(w * CIN * cols, 512, 16 * CIN, CIN, cols)),
|
|
C0=np.broadcast_to(P.f16vec(4 * C * CIN + 2 * cols * w, cols), (16, cols)).copy()) for w in range(NW)], 1)
|
|
skip = load_at(P, skip_ptr, bx, by, C // 32)
|
|
x0 = np.zeros((64, C))
|
|
for T in range(64):
|
|
tx, ty = TOK_X[T], TOK_Y[T]; r = tx // 2 + 4 * (ty // 2); q = (tx & 1) + 2 * (ty & 1)
|
|
G = np.arange(C) + q * C # global expand columns of sub-position q (pair order within C)
|
|
x0[T] = q8(E_[r, G])[gp(np.arange(C))]
|
|
return f16(x0 + skip)
|