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.
166 lines
7.8 KiB
Python
166 lines
7.8 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate public synthetic K replay fixtures and validate native outputs.
|
|
|
|
Uses the upstream replay format and NumPy model for all eleven K layers.
|
|
No NVIDIA weights, PTX or DLLs are used by this fixture generator.
|
|
"""
|
|
import argparse
|
|
import math
|
|
import pathlib
|
|
import struct
|
|
import sys
|
|
|
|
import numpy as np
|
|
|
|
ROOT = pathlib.Path(__file__).resolve().parents[2]
|
|
REFERENCE = pathlib.Path(__file__).resolve().parent / 'reference'
|
|
if not (REFERENCE / 'pwin_check.py').exists():
|
|
REFERENCE = ROOT / 'kernels/tools'
|
|
sys.path.insert(0, str(REFERENCE))
|
|
import pwin_model as pm
|
|
pm.NORM_FRAGMENT_ORDER = True
|
|
from pwin_check import LAYERS, model_window, Outputs, stored, windows
|
|
|
|
|
|
def generate(layer, directory):
|
|
kind, heads, channels, other, extra = LAYERS[layer]
|
|
cin = 16 if kind == 'enc0' else channels
|
|
low = kind in ('dec', 'dec0')
|
|
cout = other if kind in ('enc', 'enc0') else channels
|
|
layout = pm.Layout(heads, channels, cout)
|
|
core_base = 2 * channels + 2 * 16 * channels if kind == 'enc0' else 0
|
|
if low:
|
|
core_base = 2 * other * 4 * channels + 8 * channels
|
|
weight_bytes = core_base + layout.PMB + 2 * layout.NPA * heads
|
|
weights = np.zeros(weight_bytes, np.uint8)
|
|
|
|
def half(offset, value):
|
|
weights[offset:offset + 2] = np.array([value], np.float16).view(np.uint8)
|
|
|
|
def matrix(base, k, n, stride, value=1.0):
|
|
half(base + 512 * (k // 16) + stride * (n // 16) + pm.frag_offset(k % 16, n % 16), value)
|
|
|
|
if kind == 'enc0':
|
|
for k in range(16):
|
|
matrix(2 * channels, k, k, 512)
|
|
matrix(2 * channels, k, k + 16, 512, .5)
|
|
if low:
|
|
for q in range(4):
|
|
for k in range(min(other, channels)):
|
|
matrix(0, k, q * channels + k, 512 * (other // 16), .25)
|
|
for c in range(channels):
|
|
half(core_base + layout.G1 + 2 * c, 1)
|
|
half(core_base + layout.G2 + 2 * c, 1)
|
|
half(core_base + layout.BO + 2 * c, 1 / 64)
|
|
half(core_base + layout.B2 + 2 * c, 1 / 64)
|
|
position_only = kind in ('enc0', 'dec0') or bool(extra)
|
|
for head in range(heads):
|
|
source_group = head % (channels // 32)
|
|
for j in range(3):
|
|
for k in range(32):
|
|
matrix(core_base + layout.QKV + layout.HEAD * head + 2048 * (3 * source_group + j),
|
|
k, k, 1024, .5 if j == 2 else .25)
|
|
for k in range(32):
|
|
matrix(core_base + layout.WO + 64 * channels * head, k,
|
|
32 * source_group + k, 1024, .5)
|
|
if position_only:
|
|
for row in range(64):
|
|
for col in range(64):
|
|
r, c = row % 16, col % 16
|
|
offset = core_base + layout.BIAS + head * 8192 + 512 * (4 * (col // 16) + row // 16)
|
|
offset += 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3)
|
|
half(offset, 1 / 64)
|
|
for group in range(channels // 32):
|
|
for k in range(32):
|
|
matrix(core_base + layout.W1 + 64 * channels * group, 32 * group + k, k, 32 * channels, .5)
|
|
matrix(core_base + layout.W2 + 64 * channels * group, k, 32 * group + k, 1024, .5)
|
|
if kind in ('enc', 'enc0'):
|
|
for n in range(cout):
|
|
matrix(core_base + layout.PM, n % channels + channels * ((n // channels) % 4),
|
|
n, 128 * channels, .5)
|
|
if kind == 'dec0':
|
|
for n in range(40):
|
|
matrix(core_base + layout.PM + 2 * 48, n % channels, n, 1024, .5)
|
|
half(core_base + layout.PM + 2 * n, 1 / 64)
|
|
|
|
input_elements = 16 * other if low else 64 * cin
|
|
x = ((np.arange(input_elements) % 9 - 4) / 4).astype(np.float16)
|
|
skip = ((np.arange(64 * channels) % 7 - 3) / 8).astype(np.float16)
|
|
output_channels = 40 if kind == 'dec0' else channels
|
|
allocations = {0: x.view(np.uint8), 1: skip.view(np.uint8), 2: weights}
|
|
if kind in ('enc0', 'enc'):
|
|
allocations[3] = np.full(16 * cout, np.nan, np.float16).view(np.uint8)
|
|
allocations[4] = np.full(64 * channels, np.nan, np.float16).view(np.uint8)
|
|
else:
|
|
allocations[3] = np.full(64 * output_channels, np.nan, np.float16).view(np.uint8)
|
|
parameters = bytearray(176)
|
|
struct.pack_into('<ii', parameters, 0, 8, 8)
|
|
pointers = {8: 0, 64: 2, 24: 3}
|
|
if low:
|
|
pointers[16] = 1
|
|
if kind in ('enc0', 'enc'):
|
|
pointers[32] = 4
|
|
bases = {i: 0x100000000 + i * 0x1000000 for i in allocations}
|
|
for offset, identifier in pointers.items():
|
|
struct.pack_into('<Q', parameters, offset, bases[identifier])
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
for identifier, data in allocations.items():
|
|
data.tofile(directory / f'alloc-{identifier}.bin')
|
|
(directory / 'args.bin').write_bytes(parameters)
|
|
lines = [f'kernel dltss_pwin_{layer}_layer', 'launch 1 1 1 32 1 4 0']
|
|
lines += [f'alloc {i} {bases[i]:016x} {data.size}' for i, data in allocations.items()]
|
|
lines += [f'pointer {offset} {i} 0' for offset, i in pointers.items()]
|
|
(directory / 'manifest.txt').write_text('\n'.join(lines) + '\n', encoding='utf-8')
|
|
print(f'PASS K_FIXTURE kernel={layer} weights={weight_bytes} public_synthetic=1')
|
|
|
|
|
|
def validate(layer, directory, output, window_count=1, real_capture=False):
|
|
dump = pm.Dump(str(directory))
|
|
outputs = Outputs(dump, str(output))
|
|
pairs = []
|
|
for bx, by in windows(dump.grid, window_count):
|
|
model, extra = model_window(dump, layer, bx, by)
|
|
pairs.append(stored(dump, outputs, layer, bx, by, model, extra))
|
|
expected, actual = (np.concatenate(values) for values in zip(*pairs))
|
|
if not np.isfinite(expected).all() or not np.isfinite(actual).all():
|
|
raise RuntimeError('K replay contains NaN/Inf or unwritten output')
|
|
difference = np.abs(actual.astype(np.float64) - expected.astype(np.float64))
|
|
maximum = float(difference.max())
|
|
relative = float((difference / np.maximum(np.abs(expected.astype(np.float64)), 1e-3)).max())
|
|
rms = float(np.sqrt(np.mean(difference ** 2)))
|
|
peak = max(1., float(np.abs(expected.astype(np.float64)).max()))
|
|
psnr = math.inf if not rms else 20 * math.log10(peak / rms)
|
|
print(f'K_REPLAY_REFERENCE kernel={layer} elements={actual.size} max_abs={maximum:.9g} '
|
|
f'max_rel={relative:.9g} rms={rms:.9g} psnr_db={psnr:.6g}')
|
|
# Bound ISA approximate reciprocal/rsqrt and matrix accumulation differences;
|
|
# require the independent PTX reference gate for every real-weight layer.
|
|
absolute_limit = max(.001, .005 * peak) if real_capture else .001
|
|
print(f'K_REPLAY_LIMIT windows={len(pairs)} max_abs_limit={absolute_limit:.9g} min_psnr_db=60 real_capture={int(real_capture)}')
|
|
if maximum > absolute_limit or psnr < 60:
|
|
worst = int(difference.argmax())
|
|
raise RuntimeError(f'K mismatch at {worst}: reference={expected[worst]}, gpu={actual[worst]}')
|
|
print(f'PASS K_REPLAY_REFERENCE kernel={layer} reference=logical nan_inf=0')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('mode', choices=['generate', 'validate'])
|
|
parser.add_argument('--layer', choices=list(LAYERS), required=True)
|
|
parser.add_argument('--fixture-dir', type=pathlib.Path, required=True)
|
|
parser.add_argument('--output-dir', type=pathlib.Path)
|
|
parser.add_argument('--windows', type=int, default=1)
|
|
parser.add_argument('--real-capture', action='store_true')
|
|
args = parser.parse_args()
|
|
try:
|
|
if args.mode == 'generate':
|
|
generate(args.layer, args.fixture_dir)
|
|
else:
|
|
if not args.output_dir:
|
|
parser.error('validate requires --output-dir')
|
|
if args.windows < 1:
|
|
parser.error('--windows must be positive')
|
|
validate(args.layer, args.fixture_dir, args.output_dir, args.windows, args.real_capture)
|
|
except Exception as error:
|
|
print(f'FAIL K_REPLAY {error}', file=sys.stderr)
|
|
raise SystemExit(5)
|