feat: add d4r inference engines and Windows parity tooling

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.
This commit is contained in:
ayo committed 2026-10-16 08:54:47 -04:00
1 parent 25493b3198
commit 7c5d3c9157
430 files changed
+76191

No files matched your search

+106
View File
@@ -0,0 +1,106 @@
#!/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()