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.
129 lines
6.4 KiB
Python
129 lines
6.4 KiB
Python
#!/usr/bin/env python3
|
|
"""Extract the preset K network from one captured frame into d4r's own model file.
|
|
|
|
usage: extract_k.py REPLAY_DIR OUT_PREFIX
|
|
REPLAY_DIR D4R_CUDA_REPLAY_DUMP_DIR of a run with D4R_CUDA_REPLAY_DUMP_FILTER=dltss_pwin (the 11 layer launches)
|
|
Writes OUT_PREFIX.bin (f16) and OUT_PREFIX.json (per layer: shape parameters and the offset of every tensor, in f16
|
|
elements). Tensors are plain row-major matrices W[K][N] (y = x W), unlike the pre-swizzled fragments NGX keeps on
|
|
the GPU; per-channel vectors are stored 16 times (one 16-row tile) so a kernel can load them as a matrix operand.
|
|
No NVIDIA code is involved: only the weight values, read from the buffers NGX uploaded.
|
|
"""
|
|
import argparse, glob, json, os, sys
|
|
import numpy as np
|
|
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'kernels', 'tools'))
|
|
import pwin_model as pm
|
|
|
|
F16 = np.float16
|
|
# layer -> (kind, heads, C, COUT or CL, position-only attention)
|
|
LAYERS = {'enc0': ('enc0', 2, 32, 64, True), 'enc1': ('enc', 2, 64, 64, True), 'enc2': ('enc', 2, 64, 96, False),
|
|
'enc3': ('enc', 4, 96, 128, False), 'enc4': ('enc', 4, 128, 160, False), 'dec5': ('plain', 8, 160, 0, False),
|
|
'dec4': ('dec', 4, 128, 160, False), 'dec3': ('dec', 4, 96, 128, False), 'dec2': ('dec', 2, 64, 96, False),
|
|
'dec1': ('dec', 2, 64, 64, False), 'dec0': ('dec0', 2, 32, 64, True)}
|
|
|
|
|
|
class Blob:
|
|
def __init__(self):
|
|
self.parts, self.n = [], 0
|
|
|
|
def add(self, a):
|
|
a = np.ascontiguousarray(a, F16).ravel()
|
|
pad = (-self.n) % 16 # keep every tensor 32-byte aligned
|
|
if pad:
|
|
self.parts.append(np.zeros(pad, F16)); self.n += pad
|
|
off = self.n
|
|
self.parts.append(a); self.n += a.size
|
|
return off
|
|
|
|
|
|
def tile16(v):
|
|
return np.tile(np.asarray(v, F16)[None, :], (16, 1))
|
|
|
|
|
|
def pad_cols(W, n):
|
|
out = np.zeros((W.shape[0], n), F16); out[:, :W.shape[1]] = W
|
|
return out
|
|
|
|
|
|
def core(P, L, blob, posattn, fold_gains=True):
|
|
C, H = L.C, L.H
|
|
# The two per-channel scales are folded into the rows of the weights they feed (x * g) W = x (diag(g) W):
|
|
# norm 1's gain into Q, K, V and the MLP input gain into fc1. Products in f64, rounded once to f16.
|
|
g1 = pm.vec(P, L.G1, C).astype(np.float64)[:, None]
|
|
g2 = pm.vec(P, L.G2, C).astype(np.float64)[:, None]
|
|
t = {'g1': blob.add(tile16(pm.vec(P, L.G1, C)))}
|
|
t['qkv'], t['table'], t['wo'] = [], [], []
|
|
for h in range(H):
|
|
hb = L.QKV + L.HEAD * h
|
|
for j in range(3):
|
|
W = np.concatenate([pm.weights(P, hb + 2048 * (3 * ch + j), 32, 32, 1024) for ch in range(C // 32)], axis=0) # [C][32]
|
|
t['qkv'].append(blob.add((W.astype(np.float64) * g1).astype(F16) if fold_gains else W))
|
|
t['table'].append(blob.add(pm.bias_table(P, L.BIAS + 8192 * h))) # [64][64]
|
|
t['wo'].append(blob.add(pm.weights(P, L.WO + 64 * C * h, 32, C, 1024))) # [32][C]
|
|
t['bo'] = blob.add(tile16(pm.vec(P, L.BO, C)))
|
|
t['g2'] = blob.add(tile16(pm.vec(P, L.G2, C)))
|
|
t['b2'] = blob.add(tile16(pm.vec(P, L.B2, C)))
|
|
t['w1'] = []
|
|
for c in range(C // 8):
|
|
W = pm.weights(P, L.W1 + 64 * C * c, C, 32, 32 * C)
|
|
t['w1'].append(blob.add((W.astype(np.float64) * g2).astype(F16) if fold_gains else W))
|
|
t['b1'] = [blob.add(tile16(pm.vec(P, L.B1 + 64 * c, 32))) for c in range(C // 8)]
|
|
t['w2'] = [blob.add(pm.weights(P, L.W2 + 64 * C * c, 32, C, 1024)) for c in range(C // 8)] # [32][C]
|
|
return t
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument('replay_dir')
|
|
parser.add_argument('out_prefix')
|
|
parser.add_argument('--unfold-gains', action='store_true', help='preserve explicit f16 gain and residual rounding instead of folding gains into weights')
|
|
args = parser.parse_args()
|
|
rdir, prefix = args.replay_dir, args.out_prefix
|
|
blob, index = Blob(), {}
|
|
for name, (kind, H, C, X, posattn) in LAYERS.items():
|
|
matches = sorted(glob.glob(f'{rdir}/replay-*-dltss_pwin_{name}_layer'))
|
|
if not matches:
|
|
raise ValueError(f'missing preset K layer {name} in {rdir}')
|
|
d = matches[0]
|
|
P = pm.Dump(d)
|
|
e = {'kind': kind, 'H': H, 'C': C, 'X': X, 'posattn': posattn, 'fold_gains': not args.unfold_gains}
|
|
if kind in ('enc0', 'enc'):
|
|
L = pm.Layout(H, C, X)
|
|
Pw = P
|
|
if kind == 'enc0':
|
|
CIN = 16
|
|
e['embed_w'] = blob.add(pm.weights(P, 2 * C, CIN, C, 512 * (CIN // 16))) # [16][C]
|
|
e['embed_b'] = blob.add(tile16(pm.vec(P, 0, C)))
|
|
Pw = pm.Shifted(P, 2 * C + 2 * CIN * C)
|
|
e.update(core(Pw, L, blob, posattn, not args.unfold_gains))
|
|
Wp = pm.weights(Pw, L.PM, 4 * C, L.NPA * H, 128 * C) # [4C][NPA*H]
|
|
bias = pm.vec(Pw, L.PMB, L.NPA * H) if kind == 'enc0' or name != 'enc1x' else None
|
|
cols = np.concatenate([Wp[:, L.NPA * w:L.NPA * w + L.NPW] for w in range(H)], axis=1) # [4C][COUT]
|
|
b = np.concatenate([bias[L.NPA * w:L.NPA * w + L.NPW] for w in range(H)])
|
|
n16 = -(-X // 16) * 16
|
|
e['merge_w'] = blob.add(pad_cols(cols, n16)); e['merge_b'] = blob.add(tile16(np.concatenate([b, np.zeros(n16 - X, F16)])))
|
|
e['merge_n'] = n16
|
|
elif kind == 'plain':
|
|
e.update(core(P, pm.Layout(H, C, C), blob, posattn, not args.unfold_gains))
|
|
else:
|
|
CL = X
|
|
EXP = 2 * CL * 4 * C
|
|
e['expand_w'] = blob.add(pm.weights(P, 0, CL, 4 * C, 512 * (CL // 16))) # [CL][4C]
|
|
e['expand_b'] = blob.add(tile16(pm.vec(P, EXP, 4 * C)))
|
|
L = pm.Layout(H, C, C)
|
|
Pw = pm.Shifted(P, EXP + 8 * C)
|
|
e.update(core(Pw, L, blob, posattn, not args.unfold_gains))
|
|
if kind == 'dec0':
|
|
NOUTA = 48
|
|
e['head_w'] = blob.add(pm.weights(Pw, L.PM + 2 * NOUTA, C, NOUTA, 1024)) # [C][48]
|
|
e['head_b'] = blob.add(tile16(pm.vec(Pw, L.PM, NOUTA)))
|
|
index[name] = e
|
|
data = np.concatenate(blob.parts)
|
|
data.tofile(prefix + '.bin')
|
|
with open(prefix + '.json', 'w') as f:
|
|
json.dump(index, f, indent=1)
|
|
print(f'{prefix}.bin: {data.size * 2} bytes, {len(index)} layers')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|