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.
107 lines
4.5 KiB
Python
107 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""Pack encoder 1's extracted weights (engine/rr/enc1_ref.py --extract) for engine/rr/enc1_m.comp.
|
|
|
|
usage: enc1_pack.py WEIGHTS.npz OUT.bin OUT.h
|
|
|
|
The shader runs every product transposed (D^T = W^T X^T: weights are the A operand, activations the B operand), so a
|
|
result's rows are output channels and its columns tokens, and the result is the next product's B operand in place.
|
|
The rows come out in the order of W's columns, so every weight's columns are permuted here into the k order NVIDIA's
|
|
next product contracts them in (engine/rr/enc1_ref.py: q8(x)[:, inter(n)]), and the f16 vectors that act on those
|
|
rows likewise. Keys and queries are token slots: slot s is token inter(96)[s] (the prologue's key list), so the
|
|
attention bias is stored transposed in slot order.
|
|
|
|
A tiles W (K x N) -> W^T 16x16 tiles [mt][kt], row-major (m, k): 256 * (mt * K/16 + kt) bytes
|
|
B tiles the V projection runs X Wv (A = tokens): tiles [kt][nt], row-major (k, n)
|
|
f16 vectors as is; the attention bias as 96 x 64 [key slot][query slot]
|
|
"""
|
|
import os
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
from enc1_ref import inter # noqa: E402
|
|
|
|
I32, I64, I96 = inter(32), inter(64), inter(96)
|
|
|
|
|
|
def a_tiles(w):
|
|
k, n = w.shape
|
|
return np.ascontiguousarray(w.T).reshape(n // 16, 16, k // 16, 16).transpose(0, 2, 1, 3).reshape(-1)
|
|
|
|
|
|
def b_tiles(w):
|
|
k, n = w.shape
|
|
return w.reshape(k // 16, 16, n // 16, 16).transpose(0, 2, 1, 3).reshape(-1)
|
|
|
|
|
|
CUBIC_MAX = 0x3873 # binary16 bits of 0.55615234375, the clamp of the scaled score
|
|
|
|
|
|
def act_table():
|
|
"""the MLP's hidden activation and its FP8 conversion, q8(act(h)), for every binary16 h, indexed by its bits: one
|
|
byte load in place of the six binary16 operations and the conversion (the shaders' ACTQ)"""
|
|
from enc1_ref import act
|
|
from rrmath import q8
|
|
h = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16)
|
|
with np.errstate(all='ignore'):
|
|
return np.asarray(q8(act(h).astype(np.float64)), np.uint8)
|
|
|
|
|
|
def cubic_table():
|
|
"""the cubic softmax numerator's bits (enc1_ref.softmax_cubic) for every binary16 score, indexed by its bits (the
|
|
scaling and the clamp included; NaN scores take NVIDIA's min / max, which return the other operand)"""
|
|
x = np.arange(65536, dtype=np.uint32).astype(np.uint16).view(np.float16)
|
|
with np.errstate(invalid='ignore', over='ignore'):
|
|
m = (x * np.float16(0.017181396484375)).astype(np.float16)
|
|
lim = np.float16(0.55615234375)
|
|
m = np.where(np.isnan(m), lim, np.clip(m, -lim, lim)).astype(np.float16)
|
|
|
|
def fma16(a, b, c):
|
|
return (a.astype(np.float64) * b.astype(np.float64) + c.astype(np.float64)).astype(np.float16)
|
|
t = fma16(m, -m, np.full_like(m, np.float16(0.92724609375)))
|
|
u = fma16(m, t, np.full_like(m, np.float16(1.375)))
|
|
return ((u.view(np.uint16).astype(np.uint32) << 5) & 0x7FE0).astype(np.uint16)
|
|
|
|
|
|
def main():
|
|
src, out_bin, out_h = sys.argv[1:4]
|
|
w = dict(np.load(src))
|
|
e = {}
|
|
for h in range(2):
|
|
e[f'wk{h}'] = ('A', np.vstack([w[f'wk{h}a'], w[f'wk{h}b']])[:, I32])
|
|
e[f'wv{h}'] = ('B', np.vstack([w[f'wv{h}a'], w[f'wv{h}b']])[:, I32])
|
|
e[f'wo{h}'] = ('A', w[f'wo{h}'][:, I64])
|
|
e[f'bias{h}'] = ('H', w[f'bias{h}'][np.ix_(I64, I96)].T)
|
|
for b in range(8):
|
|
e[f'w1_{b}'] = ('A', w[f'w1_{b}'][:, I32])
|
|
e[f'b1_{b}'] = ('H', w[f'b1_{b}'][I32])
|
|
e[f'w2_{b}'] = ('A', w[f'w2_{b}'][:, I64])
|
|
e['wm'], e['bm'] = ('A', w['wm'][:, I96]), ('H', w['bm'][I96])
|
|
e['wh'], e['bh'] = ('A', w['wh']), ('H', w['bh'])
|
|
e['gain0'] = ('H', w['gain0'])
|
|
for k in ('table7', 'gain14', 'table16'):
|
|
e[k] = ('H', w[k][I64])
|
|
e['cubic'] = ('U', cubic_table())
|
|
blob, defines = bytearray(), []
|
|
for name, (kind, a) in e.items():
|
|
while len(blob) % 256:
|
|
blob.append(0)
|
|
off = len(blob)
|
|
if kind == 'A':
|
|
blob += a_tiles(a).tobytes()
|
|
elif kind == 'B':
|
|
blob += b_tiles(a).tobytes()
|
|
elif kind == 'U':
|
|
blob += np.ascontiguousarray(a, np.uint16).tobytes()
|
|
else:
|
|
blob += np.ascontiguousarray(a, np.float16).tobytes()
|
|
defines.append(f'#define W_{name.upper()} {off}u // {kind} {a.shape}')
|
|
open(out_bin, 'wb').write(bytes(blob))
|
|
open(out_h, 'w').write('// generated by engine/rr/enc1_pack.py\n' + '\n'.join(defines) + '\n')
|
|
print(f'{len(blob)} bytes, {len(e)} entries')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|