wouter/d4r
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

engine: this GPU's own arithmetic and shared weight tiles as the builders' defaults

Browse Source 
- layer_m.comp builds relaxed by default (M_RELAXED: f32 accumulators, paired FP8
  conversions, no activation table) with two windows per wave sharing weight tiles
  (WPW=2); M_EXACT=1 keeps NVIDIA's roundings, bit-identical to the CUDA path
- layer.comp reuses weight tiles across its waves (TPW=2 in the 32-channel layers)
- engine/rr/rr2_w.comp: one wave per window for preset F's dec0_0 and enc0, used by
  compile_rr2.py at the new default D4R_RR_RELAXED=3 (build_native_model.sh too);
  D4R_RR_FUSE=1 matches the reference's fused multiply-adds, ptx2rr.py --snorm16
- [Engine] Network defaults to vulkan; the HIP network is opt-in
This commit is contained in:
woutermans committed 2026-10-14 00:52:41 +02:00
1 parent 8eb7a94533
commit f0b55a850a
24 files changed
+1108 -254
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

config/d4r.ini.default
+4 -5
View File
@@ -126,11 +126,10 @@ Enabled = false
; puts D4R_ENGINE_MODEL. The other folders: D4R_ENGINE_MODEL_LDR_DIR, D4R_ENGINE_MODEL_M_DIR,
; D4R_ENGINE_MODEL_M_LDR_DIR, D4R_ENGINE_MODEL_L_DIR and D4R_ENGINE_MODEL_L_LDR_DIR. (D4R_ENGINE_MODEL_DIR)
ModelDir = auto
; Where the network between the engine's Vulkan stages runs (K and M). hip = d4r's native HIP layer kernels
; (M: about 1.7x faster than the shader, 3.8 instead of 5.5 ms per frame at 1080p -> 4K on an RX 9070 XT);
; the game's command list is split around the network like a split frame, so it needs VkD3DDir, the CUDA
; bridge and a model built with --hip. vulkan = the engine's cooperative-matrix shader, everything in one
; command list. auto and hip = hip when available, else vulkan (the log says why). (D4R_ENGINE_NETWORK)
; Where the network between the engine's Vulkan stages runs (K and M). vulkan = the engine's cooperative-matrix
; shader, everything in one command list. hip = d4r's native HIP layer kernels: the game's command list is split
; around the network like a split frame, so it needs VkD3DDir, the CUDA bridge and a model built with --hip; when
; they are missing the Vulkan network runs and the log says why. auto = vulkan. (D4R_ENGINE_NETWORK)
Network = auto
[Paths]
engine/README.md
+61 -2
View File
@@ -290,8 +290,8 @@ The shim loads each model on first use from a folder next to it, or from a varia
`scripts/install_d4r_runtime.sh` stages them from `D4R_ENGINE_MODEL`, `D4R_ENGINE_MODEL_LDR`, `D4R_ENGINE_MODEL_M`,
`D4R_ENGINE_MODEL_M_LDR`, `D4R_ENGINE_MODEL_L` and `D4R_ENGINE_MODEL_L_LDR`, copying each folder whole (with `hipnet.bin` and `hip/` from `--hip`). A model whose
`hipnet.bin` belongs to the other preset keeps the Vulkan network. `D4R_ENGINE_NETWORK=vulkan` (`[Engine] Network`)
forces the Vulkan network for both presets; `D4R_SHIM_ENGINE_GPU_WAIT=0` makes the HIP network wait for its inputs on
`hipnet.bin` belongs to the other preset keeps the Vulkan network. The Vulkan network is the default since 2026-10-14;
`D4R_ENGINE_NETWORK=hip` (`[Engine] Network = hip`) selects the HIP network for both presets when the model has it; `D4R_SHIM_ENGINE_GPU_WAIT=0` makes the HIP network wait for its inputs on
the CPU-visible marker instead of a GPU-side wait, for comparisons.
With the engine enabled, an evaluation of preset K is recorded entirely into the game's own command list through vkd3d-proton's
@@ -1010,6 +1010,65 @@ flag; they are not current accuracy results. Feature release and re-creation at
clean. The current game timing and longer accuracy comparison are reported above. The kernels are built per
GPU architecture (gfx1201 with FP8 here); another GPU falls back to Vulkan when they do not load.
### This GPU's own arithmetic for the network: `M_RELAXED` (2026-10-14)
`M_DEFS=-DM_RELAXED python3 engine/compile_m.py CAPTURE OUT` builds `layer_m.comp` (presets M and L) without NVIDIA's
roundings: the accumulator stays f32 across the 32-deep steps (no scratch round trip), values go from f32 into FP8 two
per saturating conversion, and the activation is computed in f32. The default build is unchanged (first frame
byte-identical to the CUDA path through ZLUDA). Against NVIDIA's kernels on an RTX PRO 4500 Blackwell (captures
capHM / capHL: harness create flags 0x4B, 640x360 -> 1280x720, three frames; `build/rr/f/nv/srm.sh`, `srt4k.sh`):
| | M, dB | M at 1280x720 -> 4K | L, dB | L at 1280x720 -> 4K |
|---|---|---|---|---|
| default | 38.4 / 38.2 / 37.9 | 3.19 ms (network 1.42) | 40.8 / 40.9 / 40.9 | 4.47 ms (network 1.46, input 0.67, expansion 0.54) |
| `M_RELAXED` | 37.6 / 37.2 / 37.2 | 2.90 ms (network 1.13) | 40.1 / 39.8 / 39.8 | 3.83 ms (network 1.13, input 0.50, expansion 0.41) |
`M_WPW=2` (2026-10-14, `layer_m.comp` `WPW`) adds two 4x4-token windows per wave in the five network layers: a
workgroup is still one 8x8-token block, but its two waves each walk two windows and load every weight tile once for
both (a window reads the whole layer's weights; with every weight tile replaced by one tile loaded once the network
took 0.79 ms instead of 1.13). The default build with `M_WPW=2` is still byte-identical to the capture, on three
reset frames and on a ten-frame temporal capture (capTM), so the restructuring itself changes nothing; it only pays
together with `M_RELAXED` (the exact build's scratch round trips hide it: 1.41 ms network against 1.38). Four windows
per wave, or two in the 128-channel layers without spilling, would need fewer registers than the residual stream
takes (two windows: 256 vector registers, 119 spilled values in the tube layers, none in the others).
| 1280x720 -> 4K | M frame | M network | L frame | L network |
|---|---|---|---|---|
| default | 3.16 ms | 1.38 | 4.44 ms | 1.42 |
| `M_RELAXED` | 2.86 ms | 1.10 | 3.81 ms | 1.11 |
| `M_RELAXED`, `M_WPW=2` | 2.68 ms | 0.93 | 3.67 ms | 0.98 |
| the native HIP network (what `[Engine] Network = auto` picks when the model has `hipnet.bin`) | 2.90 ms | 1.08 | | |
`M_RELAXED M_WPW=2` against the NVIDIA card: M 37.6 / 37.3 / 37.2 dB, L 40.1 / 39.9 / 39.8 dB (default: 38.4 / 38.2 /
37.9 and 40.8 / 40.9 / 40.9). On the temporal capture the relaxed build settles at 34-35 dB from the ZLUDA
reference after four frames. Replay notes: the D3D12 harness resets every frame unless `D4R_HARNESS_TEMPORAL=1`, so
its captures are replayed with `D4R_M_RESET_ALL=1 d4r-m ...` (an earlier comparison without it looked like a
temporal drift of the engine; there is none). `M_DEFS_ENC1`, `M_DEFS_ENC3_TUBE`, ... pass defines to one layer.
Preset K (`layer.comp`) already accumulates in f32 and is 59.9 dB from the same card (the ZLUDA capture: 49.5 dB),
so there is no rounding to relax. Its four waves per window each load the layer's weights, and with every weight
tile replaced by one tile loaded once the eleven layers take 1.16 ms instead of 1.51 at 1280x720 -> 4K, so
`layer.comp` got the same restructuring (`TPW`: 16-token tiles per wave; `K_DEFS`, `K_DEFS_<LAYER>`, `K_TPW`). The
output is byte-identical at one, two and four tiles per wave (15 files of a three-frame capture), but K's operands
are binary16 (eight registers a tile where FP8 takes two): two tiles per wave only fit in the 32-channel layers
(enc0 277 -> 272 us, dec0 unchanged), four tiles spill there too (enc0 352 us), and in the wider layers two tiles
spill 15 to 568 values and the frame takes 3.16 ms. `compile_k.py` therefore builds two tiles per wave in the
32-channel layers only: 2.86 -> 2.81 ms. The native HIP network measures 2.83 - 2.85 ms on the same frames.
### Defaults since 2026-10-14
Every model builder now makes the fastest build measured above unless told otherwise, and the shim runs the network
on the Vulkan shaders unless `[Engine] Network = hip`:
| model | default build | the exact build (NVIDIA's roundings, bit-identical to the CUDA path) |
|---|---|---|
| K (`compile_k.py`) | two tiles per wave in the 32-channel layers (same output) | `K_TPW=1` for the previous layout |
| M, L (`compile_m.py`) | `M_RELAXED`, two windows per wave | `M_EXACT=1` |
| RR D, E (`rr/build_native_model.sh`) | `D4R_RR_RELAXED=3` | `D4R_RR_RELAXED=0` |
| RR F (`compile_rr2.py`) | `D4R_RR_RELAXED=3` | `D4R_RR_RELAXED=0` |
A model built with no options was compared shader by shader with the explicitly optioned build of each row, and the
exact switches still reproduce the captures (`build/rr/f/nv/defaults.sh`).
## Preset L (runtime and shim path exist; experimental)
Preset L (render preset 12, NVIDIA's model for Ultra Performance) is M's pipeline with its own input, expansion and
engine/compile_k.py
+1 -1
View File
@@ -186,7 +186,7 @@ def main():
offsets = b''.join(params(ZeroShape(), e)[0][16:] for e in index.values())
(a.output / 'offsets.bin').write_bytes(MAGIC + struct.pack('<II', 11, weights.stat().st_size) + offsets)
for name, e in index.items():
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines(e),
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defines(e, os.environ.get('K_DEFS_' + name.upper(), '').split()), # per layer, e.g. K_DEFS_DEC0=-DTPW=2
str(HERE / 'layer.comp'), '-o', str(a.output / (name + '.spv'))], check=True)
for name in ('input_k', 'output_k'):
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *(['-DLDR=1'] if a.ldr else []),
engine/compile_m.py
+10 -4
View File
@@ -26,6 +26,12 @@ import model_enc3 as me
import swin_model as sm
from test_m import pack_layer, defines
# The network is built with this GPU's own arithmetic and two windows per wave by default (layer_m.comp M_RELAXED, WPW;
# engine/README.md): the fastest build, about 1 dB further from an NVIDIA card than NVIDIA's roundings reproduced
# exactly. M_EXACT=1 builds those roundings instead (bit-identical to the CUDA path); M_WPW / M_DEFS still apply.
EXACT = os.environ.get('M_EXACT') == '1'
FAST = [] if EXACT or '-DM_RELAXED' in os.environ.get('M_DEFS', '').split() else ['-DM_RELAXED']
FAST_WPW = [] if EXACT or os.environ.get('M_WPW') or 'WPW=' in os.environ.get('M_DEFS', '') else ['-DWPW=2']
SHADERS = {'enc1': 'm_enc1', 'enc2': 'm_enc2', 'enc3_tube': 'm_tube', 'dec2': 'm_dec2', 'dec1': 'm_dec1'}
@@ -95,7 +101,7 @@ def main():
# the network in the arithmetic of ZLUDA's lowering of NVIDIA's PTX (TEXL): byte-identical to the CUDA backend
# with NVIDIA's translated kernels in the harness (2026-10-11), where the native accuracy kernels' K32_MIX order
# drifted to 36-41 dB through the FP8 tube
for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'])
for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'] + FAST + FAST_WPW)
for s in ('enc0_m', 'dec0_m', 'post_m', 'down_m'): glsl(s + '.comp', s + '.spv', ['-DLDR=1'] if a.ldr and s != 'dec0_m' else [])
glsl('down_m.comp', 'down_m_rgb10.spv', ['-DOUTPUT_RGB10=1'] + (['-DLDR=1'] if a.ldr else []))
glsl('post_m.comp', 'post_m_q.spv', ['-DTILE_W=15', '-DTILE_H=15'] + (['-DLDR=1'] if a.ldr else []))
@@ -146,10 +152,10 @@ def compile_l(a, one, R):
subprocess.run(['glslangValidator', '-V', '--target-env', 'vulkan1.3', *flags, str(HERE / src), '-o', str(a.output / out)], check=True, stdout=subprocess.DEVNULL)
ldr = ['-DLDR=1'] if a.ldr else []
# the network in the arithmetic of ZLUDA's lowering of NVIDIA's PTX (TEXL), as for preset M (L_NOTES.md)
for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'])
for name, out in SHADERS.items(): glsl('layer_m.comp', out + '.spv', defines(name) + ['-DTEXL=1'] + FAST + FAST_WPW)
for planes, suffix in ((0, ''), (1, '_planes')):
glsl('layer_m.comp', f'enc0_l{suffix}.spv', defines('enc0l') + ldr + [f'-DPLANES={planes}'])
glsl('layer_m.comp', f'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'])
glsl('layer_m.comp', f'enc0_l{suffix}.spv', defines('enc0l') + ldr + [f'-DPLANES={planes}'] + FAST)
glsl('layer_m.comp', f'dec0_l{suffix}.spv', defines('dec0l') + [f'-DPLANES={planes}'] + FAST)
# the reconstruction one render texel wider per tier than M's (4x4 taps): 11 (reconstruction/render >= 3),
# 16 (the whole post_3_2 range), 12 (the compact tier: render/reconstruction <= 0.45, jitter within half a texel)
for out, tile in (('post_l', 11), ('post_l_q', 16), ('post_l_c', 12)):
engine/compile_rr2.py
+24 -5
View File
@@ -30,6 +30,21 @@ sys.path.insert(0, str(HERE / 'rr'))
from rr_capture import Capture # noqa: E402
from test_layer import stage_params # noqa: E402
# D4R_RR_FUSE=1: the fused multiply-adds of an NVIDIA card (rr2_m.comp RR_FUSE, enc1_ref.FUSE); give --ptx-dir the
# modules rewritten the same way (build/rr/f/nv/mkptx.sh) so the translated parts match
FUSE_DEFS = ['-DRR_FUSE'] if os.environ.get('D4R_RR_FUSE') == '1' else []
# D4R_RR_RELAXED=1|2|3: this GPU's own arithmetic in the native stages instead of the reference's roundings
# (engine/rr/rr_exact.glsl; 3 = also the residual steps in f32, rr2_m.comp RR_NATIVE): faster, ~40 dB from the
# reference, which is also how far two NVIDIA cards are from each other (engine/rr/README.md)
# level 3 (this GPU's own arithmetic, one wave per window where the shader has it) is the default: the fastest build,
# as far from an NVIDIA card as two NVIDIA setups are from each other. D4R_RR_RELAXED=0 builds NVIDIA's roundings
# exactly (bit-identical to the CUDA path through ZLUDA)
RELAXED = int(os.environ.get('D4R_RR_RELAXED', '3') or 0)
# the stages that have a one-wave shader for level 3 (engine/rr/rr2_w.comp: byte-identical to rr2_m.comp there, faster)
ONE_WAVE = {'dec0_0', 'enc0'}
RELAX_DEFS = ([f'-DRR_RELAXED={min(RELAXED, 2)}'] if RELAXED else []) + (['-DRR_NATIVE'] if RELAXED >= 3 else [])
# ... and the normals' filtered fetch of encoder 0 (argument 528) rounded to 16-bit steps, as an NVIDIA card returns it
SNORM16 = {'enc0': [528]} if FUSE_DEFS else {}
NETWORK = ['enc0', 'enc1', 'enc2', 'enc3', 'enc4', 'enc5'] + [f'dec{n}_{k}' for n in (4, 3, 2, 1, 0) for k in (0, 1)]
STAGES = NETWORK + ['post', 'down', 'pooled']
KERNEL = {s: f'rr2_{s}_kernel' for s in NETWORK[:6]}
@@ -101,12 +116,13 @@ def main():
weights = f'weights-{stage}m.bin'
subprocess.run([sys.executable, str(HERE / 'rr' / 'rr2_pack.py'), stage, a.native, str(out / weights),
str(work / 'rr2_stage.h')], check=True, stdout=subprocess.DEVNULL)
defs = ['-DRR_EXACT_CONVERT'] + (['-DRR_MIX'] if os.environ.get('D4R_RR_MIX') == '1' else [])
defs = ['-DRR_EXACT_CONVERT'] + (['-DRR_MIX'] if os.environ.get('D4R_RR_MIX') == '1' else []) + FUSE_DEFS + RELAX_DEFS
subprocess.run(capped + ['glslangValidator', '-V', '--target-env', 'vulkan1.3', *defs, f'-I{work}',
f'-I{HERE / "rr"}', str(HERE / 'rr' / 'rr2_m.comp'), '-o', str(out / f'{stage}_m.spv')],
check=True, stdout=subprocess.DEVNULL)
subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / f'{stage}_m.spv')],
f'-I{HERE / "rr"}', str(HERE / 'rr' / ('rr2_w.comp' if RELAXED >= 3 and stage in ONE_WAVE else 'rr2_m.comp')), '-o', str(out / f'{stage}_m.spv')],
check=True, stdout=subprocess.DEVNULL)
if RELAXED < 2: # (the relaxed stages leave their conversions to the driver)
subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / f'{stage}_m.spv')],
check=True, stdout=subprocess.DEVNULL)
(out / f'{stage}.args').write_bytes(stage_params(launch, 24)[:len(launch.args)])
arena = launch.arena_alloc
unit = launch.allocs[arena][1] / 3152
@@ -118,6 +134,7 @@ def main():
cmd = [sys.executable, str(HERE / 'rr' / 'ptx2rr.py'), str(find_module(a.ptx_dir, kernel)), kernel, str(generated),
'--block', BLOCK[stage], '--weights-arg', '24' if network else '9999']
cmd += [f'--buffer={off}' for off in BUFFERS.get(stage, [])]
cmd += [f'--snorm16={off}' for off in SNORM16.get(stage, [])]
cmd += [f'--surface={off}={fmt}' for off, fmt in SURFACES.get(stage, {}).items()]
subprocess.run(cmd, check=True, stderr=subprocess.DEVNULL)
info = json.loads(generated.with_suffix('.json').read_text())
@@ -129,7 +146,7 @@ def main():
weights = f'weights-{stage}m.bin'
subprocess.run([sys.executable, str(HERE / 'rr' / 'rr2_pack.py'), stage, a.native, str(out / weights),
str(work / 'rr2_stage.h')], check=True, stdout=subprocess.DEVNULL)
defs = ['-DRR_EXACT_CONVERT'] + (['-DRR_MIX'] if os.environ.get('D4R_RR_MIX') == '1' else [])
defs = ['-DRR_EXACT_CONVERT'] + (['-DRR_MIX'] if os.environ.get('D4R_RR_MIX') == '1' else []) + FUSE_DEFS + RELAX_DEFS
# encoder 0 in two passes (gen_rr2_enc0.py --split): its gather as the stage enc0g, whose record's address
# is a word appended to the argument block
split = stage == 'enc0' and SPLIT and GROUP == 1
@@ -144,6 +161,8 @@ def main():
subprocess.run([sys.executable, str(HERE / 'rr' / 'spv_exact_convert.py'), str(out / 'enc0g.spv')],
check=True, stdout=subprocess.DEVNULL)
record = int((work / 'enc0_g.record').read_text())
if RELAXED >= 3 and stage in ONE_WAVE and split:
how = how + ['--one-wave']
subprocess.run([sys.executable, str(HERE / 'rr' / 'gen_rr2_enc0.py'), str(find_module(a.ptx_dir, kernel)),
str(generated), str(work / f'{stage}_m.comp'), *HYBRID[stage], *how],
check=True, stdout=subprocess.DEVNULL)
engine/layer.comp
+171 -116
View File
@@ -19,6 +19,7 @@
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#ifndef PADX
#define PADX 8
#endif
@@ -37,7 +38,13 @@
#ifndef ABL_CHUNKS
#define ABL_CHUNKS (UC / 8u)
#endif
layout(local_size_x = 128) in;
// TPW: 16-token tiles of the 8x8 window per wave (1, 2 or 4). With TPW > 1 a wave walks TPW tiles and loads every
// weight tile once for all of them (four waves per window each read the whole layer's weights otherwise, about a
// fifth of the network's time). The arithmetic and its order per tile do not change.
#ifndef TPW
#define TPW 1
#endif
layout(local_size_x = 128 / TPW) in;
layout(set = 0, binding = 0, std430) readonly buffer Pb { int W; int Hh; int sx; int sy; uint o[108]; } p;
layout(set = 0, binding = 1, std430) readonly buffer Wb { f16vec4 w[]; } wb;
layout(set = 0, binding = 2, std430) readonly buffer Ib { f16vec4 a[]; } inb;
@@ -144,65 +151,85 @@ f16vec4 gelu(f16vec4 x)
}
void main()
{
const uint mt = gl_LocalInvocationID.x >> 5, lane = gl_LocalInvocationID.x & 31u; // wave = token tile
const uint row = lane & 15u, hf = lane >> 4, tok = 16u * mt + row; // two lanes per token
const uint mt0 = (gl_LocalInvocationID.x >> 5) * uint(TPW), lane = gl_LocalInvocationID.x & 31u; // wave = TPW token tiles
const uint row = lane & 15u, hf = lane >> 4; // two lanes per token
const int X0 = 8 * int(gl_WorkGroupID.x) - p.sx, Y0 = 8 * int(gl_WorkGroupID.y) - p.sy;
const int ty = Y0 + int(tok >> 3), tx = X0 + int(tok & 7u);
const uint ch0 = hf * HC; // this lane's first channel
ACC acc[C / 16]; // the residual stream of this tile
HACC residual[C / 16]; // x0, added back at the MLP boundary as one f16 add
// per tile w of this wave (FORW): its index mt in the window and this lane's token (WINV)
#define FORW [[unroll]] for (uint w = 0u; w < uint(TPW); ++w)
#define WINV const uint mt = mt0 + w, tok = 16u * mt + row; const int ty = Y0 + int(tok >> 3), tx = X0 + int(tok & 7u);
ACC ACC_[TPW][C / 16]; // the residual stream of each tile
HACC RES_[TPW][C / 16]; // x0, added back at the MLP boundary as one f16 add
#define acc ACC_[w]
#define residual RES_[w]
// ------------------------------------------------------------ input: x0 into acc, x0 * rsqrt(sum x0^2) into BF
#if KIND == 2
{
// low-resolution tokens (one per row index), channels split over the 8 lane groups of the window
const int W2 = p.W / 2, H2 = p.Hh / 2;
const uint part = uint(XC) / 8u, c0 = (mt * 2u + hf) * part;
const uint part = uint(XC) / 8u;
const uint src = uint(mirror(tdiv2(Y0) + int(row >> 2), H2) * W2 + mirror(tdiv2(X0) + int(row & 3u), W2)) * uint(XC);
for (uint c = c0; c < c0 + part; c += 4u) writeTRaw(row*SLO+c,inb.a[(src+c)/4u]);
FORW {
WINV
const uint c0 = (mt * 2u + hf) * part;
for (uint c = c0; c < c0 + part; c += 4u) writeTRaw(row*SLO+c,inb.a[(src+c)/4u]);
}
SYNC
// expand: this wave computes quarter mt of the 4C columns
for (uint nt = 0u; nt < NT; ++nt) acc[nt] = bias_tile(p.o[105] + mt * UC + nt * 16u, 4u * UC);
for (uint kt = 0u; kt < uint(XC) / 16u; ++kt)
{
MA a; LOAD_T_RAW(a,kt*16u,SLO);
MB b[C / 16];
for (uint nt = 0u; nt < NT; ++nt) LD(b[nt], wb.w, p.o[104] + kt * 16u * 4u * UC + mt * UC + nt * 16u, 4u * UC);
for (uint nt = 0u; nt < NT; ++nt) acc[nt] = coopMatMulAdd(a, b[nt], acc[nt]);
// expand: tile mt computes quarter mt of the 4C columns
FORW {
WINV
for (uint nt = 0u; nt < NT; ++nt) acc[nt] = bias_tile(p.o[105] + mt * UC + nt * 16u, 4u * UC);
for (uint kt = 0u; kt < uint(XC) / 16u; ++kt)
{
MA a; LOAD_T_RAW(a,kt*16u,SLO);
MB b[C / 16];
for (uint nt = 0u; nt < NT; ++nt) LD(b[nt], wb.w, p.o[104] + kt * 16u * 4u * UC + mt * UC + nt * 16u, 4u * UC);
for (uint nt = 0u; nt < NT; ++nt) acc[nt] = coopMatMulAdd(a, b[nt], acc[nt]);
}
}
SYNC // everyone has read the low-resolution tokens
for (uint nt = 0u; nt < NT; ++nt) STORE_B_PATCH(acc[nt],0u,mt*UC+nt*16u,S4);
FORW { WINV for (uint nt = 0u; nt < NT; ++nt) STORE_B_PATCH(acc[nt],0u,mt*UC+nt*16u,S4); }
SYNC
const uint ly = (tok >> 3) >> 1, lx = (tok & 7u) >> 1, q = ((tok >> 3) & 1u) * 2u + (tok & 1u);
const uint sk = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * UC + ch0;
for (uint c = 0u; c < HC; c += 4u) writeTToken(tok,ch0+c,SX,readBPatch(ly*4u+lx,q*UC+ch0+c,S4)+skb.a[(sk+c)/4u]);
FORW {
WINV
const uint ly = (tok >> 3) >> 1, lx = (tok & 7u) >> 1, q = ((tok >> 3) & 1u) * 2u + (tok & 1u);
const uint sk = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * UC + ch0;
for (uint c = 0u; c < HC; c += 4u) writeTToken(tok,ch0+c,SX,readBPatch(ly*4u+lx,q*UC+ch0+c,S4)+skb.a[(sk+c)/4u]);
}
SYNC // BF (the expanded tokens) is free
}
#define LOAD_X0(m,c) LOAD_T_TOKEN(m,mt*16u,c,SX)
#define READ_X0(r,c) readTToken(r,c,SX)
#elif EMBED
{
const uint src = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * 16u + hf * 8u;
for (uint c = 0u; c < 8u; c += 4u) writeTRaw(tok*SEM+hf*8u+c,inb.a[(src+c)/4u]);
MA a; LOAD_T_RAW(a,mt*16u*SEM,SEM);
MA EA_[TPW];
FORW {
WINV
const uint src = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * 16u + hf * 8u;
for (uint c = 0u; c < 8u; c += 4u) writeTRaw(tok*SEM+hf*8u+c,inb.a[(src+c)/4u]);
LOAD_T_RAW(EA_[w],mt*16u*SEM,SEM);
}
for (uint nt = 0u; nt < NT; ++nt)
{
MB b; LD(b, wb.w, p.o[106] + nt * 16u, UC);
STORE_B_TOKEN(coopMatMulAdd(a,b,bias_tile(p.o[107]+nt*16u,UC)),mt*16u,nt*16u);
FORW { WINV STORE_B_TOKEN(coopMatMulAdd(EA_[w],b,bias_tile(p.o[107]+nt*16u,UC)),mt*16u,nt*16u); }
}
for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,max(readBToken(tok,ch0+c),f16vec4(0.0)));
FORW { WINV for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,max(readBToken(tok,ch0+c),f16vec4(0.0))); }
SYNC // TT is laid out differently from here on
}
#define LOAD_X0(m,c) LOAD_B_TOKEN(m,mt*16u,c)
#define READ_X0(r,c) readBToken(r,c)
#else
{
FORW {
WINV
const uint src = uint(mirror(ty, p.Hh) * p.W + mirror(tx, p.W)) * UC + ch0;
for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,inb.a[(src+c)/4u]);
}
#define LOAD_X0(m,c) LOAD_B_TOKEN(m,mt*16u,c)
#define READ_X0(r,c) readBToken(r,c)
#endif
{
FORW {
WINV
// the residual stream starts as x0 + the output projection's bias
for (uint nt = 0u; nt < NT; ++nt) {
HACC x; LOAD_X0(x,nt*16u);
@@ -229,51 +256,42 @@ void main()
writeBToken(tok,ch0+c,x[c/4u]*(r*wb.w[(p.o[0]+ch0+c)/4u]));
#endif
}
#if KIND == 2
SYNC // x0 lived in TT; the heads overwrite it
#endif
}
MA xn[C / 16]; // the normalised input of this tile, for every head's Q, K, V
for (uint kt = 0u; kt < NT; ++kt) LOAD_B_TOKEN(xn[kt],mt*16u,kt*16u);
#if KIND == 2
SYNC // x0 lived in TT; the heads overwrite it
#endif
MA XN_[TPW][C / 16]; // the normalised input of each tile, for every head's Q, K, V
#define xn XN_[w]
FORW { WINV for (uint kt = 0u; kt < NT; ++kt) LOAD_B_TOKEN(xn[kt],mt*16u,kt*16u); }
// ------------------------------------------------------------ attention, one head at a time
for (uint h = 0u; h < ABL_HEADS; ++h)
{
const uint otab = p.o[7u + 5u * h], owo = p.o[8u + 5u * h];
{
// Q, K, V of this tile, two K steps at a time: their weight tiles are requested before the first product
ACC q0 = ACC(0.0), q1 = ACC(0.0), k0 = ACC(0.0), k1 = ACC(0.0), v0 = ACC(0.0), v1 = ACC(0.0);
for (uint kt = 0u; kt < NT; kt += 2u)
{
MB bv[4];
for (uint i = 0u; i < 2u; ++i)
{
LD(bv[2u * i], wb.w, p.o[6u + 5u * h] + (kt + i) * 512u, 32u);
LD(bv[2u * i + 1u], wb.w, p.o[6u + 5u * h] + (kt + i) * 512u + 16u, 32u);
}
// Q, K, V of every tile, two K steps at a time: their weight tiles are requested before the first product,
// and each is loaded once for the wave's tiles (one projection at a time: fewer live accumulators)
#define PROJ(WOFF, STORE0, STORE1) { \
ACC P0_[TPW], P1_[TPW]; \
FORW { P0_[w] = ACC(0.0); P1_[w] = ACC(0.0); } \
for (uint kt = 0u; kt < NT; kt += 2u) \
{ \
MB bw[4]; \
for (uint i = 0u; i < 2u; ++i) \
{ \
LD(bw[2u * i], wb.w, (WOFF) + (kt + i) * 512u, 32u); \
LD(bw[2u * i + 1u], wb.w, (WOFF) + (kt + i) * 512u + 16u, 32u); \
} \
FORW { for (uint i = 0u; i < 2u; ++i) { P0_[w] = coopMatMulAdd(xn[kt + i], bw[2u * i], P0_[w]); P1_[w] = coopMatMulAdd(xn[kt + i], bw[2u * i + 1u], P1_[w]); } } \
} \
FORW { WINV STORE0; STORE1; } }
#if !POSATTN
MB bq[4], bk[4];
for (uint i = 0u; i < 2u; ++i)
{
LD(bq[2u * i], wb.w, p.o[4u + 5u * h] + (kt + i) * 512u, 32u);
LD(bq[2u * i + 1u], wb.w, p.o[4u + 5u * h] + (kt + i) * 512u + 16u, 32u);
LD(bk[2u * i], wb.w, p.o[5u + 5u * h] + (kt + i) * 512u, 32u);
LD(bk[2u * i + 1u], wb.w, p.o[5u + 5u * h] + (kt + i) * 512u + 16u, 32u);
}
for (uint i = 0u; i < 2u; ++i)
{
q0 = coopMatMulAdd(xn[kt + i], bq[2u * i], q0); q1 = coopMatMulAdd(xn[kt + i], bq[2u * i + 1u], q1);
k0 = coopMatMulAdd(xn[kt + i], bk[2u * i], k0); k1 = coopMatMulAdd(xn[kt + i], bk[2u * i + 1u], k1);
}
PROJ(p.o[4u + 5u * h], STORE_B_TOKEN(P0_[w],mt*16u,0u), STORE_B_TOKEN(P1_[w],mt*16u,16u)) // Q: this tile's own rows
PROJ(p.o[5u + 5u * h], STORE_T_TOKEN(P0_[w],mt*16u,0u,SQ), STORE_T_TOKEN(P1_[w],mt*16u,16u,SQ))
#endif
for (uint i = 0u; i < 2u; ++i) { v0 = coopMatMulAdd(xn[kt + i], bv[2u * i], v0); v1 = coopMatMulAdd(xn[kt + i], bv[2u * i + 1u], v1); }
}
#if !POSATTN
STORE_B_TOKEN(q0,mt*16u,0u); STORE_B_TOKEN(q1,mt*16u,16u); // Q: this wave's own rows
STORE_T_TOKEN(k0,mt*16u,0u,SQ); STORE_T_TOKEN(k1,mt*16u,16u,SQ);
#endif
STORE_T_TOKEN(v0,VOFF/SQ+mt*16u,0u,SQ); STORE_T_TOKEN(v1,VOFF/SQ+mt*16u,16u,SQ);
}
PROJ(p.o[6u + 5u * h], STORE_T_TOKEN(P0_[w],VOFF/SQ+mt*16u,0u,SQ), STORE_T_TOKEN(P1_[w],VOFF/SQ+mt*16u,16u,SQ))
SYNC
MA OA0_[TPW], OA1_[TPW];
FORW {
WINV
#if !POSATTN
{
ACC s[4];
@@ -316,111 +334,148 @@ void main()
}
// the two K tiles of this head for the output projection: through memory, to become A operands
STORE_B_TOKEN(o0,mt*16u,0u); STORE_B_TOKEN(o1,mt*16u,16u);
MA a0, a1; LOAD_B_TOKEN(a0,mt*16u,0u); LOAD_B_TOKEN(a1,mt*16u,16u);
for (uint n0 = 0u; n0 < NT; n0 += 2u)
{
MB b0[2], b1[2];
for (uint i = 0u; i < 2u; ++i) { LD(b0[i], wb.w, owo + (n0 + i) * 16u, UC); LD(b1[i], wb.w, owo + 16u * UC + (n0 + i) * 16u, UC); }
for (uint i = 0u; i < 2u; ++i) {
acc[n0 + i] = coopMatMulAdd(a1, b1[i], coopMatMulAdd(a0, b0[i], acc[n0 + i]));
}
}
LOAD_B_TOKEN(OA0_[w],mt*16u,0u); LOAD_B_TOKEN(OA1_[w],mt*16u,16u);
}
}
for (uint n0 = 0u; n0 < NT; n0 += 2u)
{
MB b0[2], b1[2];
for (uint i = 0u; i < 2u; ++i) { LD(b0[i], wb.w, owo + (n0 + i) * 16u, UC); LD(b1[i], wb.w, owo + 16u * UC + (n0 + i) * 16u, UC); }
FORW { for (uint i = 0u; i < 2u; ++i) {
acc[n0 + i] = coopMatMulAdd(OA1_[w], b1[i], coopMatMulAdd(OA0_[w], b0[i], acc[n0 + i]));
} }
}
SYNC // every wave is done with K and V before the next head overwrites them
}
// ------------------------------------------------------------ MLP: its input is the stream itself (gain folded into fc1)
{
MA M_[TPW][C / 16];
#define m M_[w]
FORW {
WINV
for (uint nt = 0u; nt < NT; ++nt) acc[nt] = ACC(HACC(acc[nt]) + residual[nt]);
for (uint nt = 0u; nt < NT; ++nt) STORE_B_TOKEN(acc[nt],mt*16u,nt*16u);
#if !FOLD_GAINS
for (uint c = 0u; c < HC; c += 4u) writeBToken(tok,ch0+c,readBToken(tok,ch0+c)*wb.w[(p.o[2]+ch0+c)/4u]);
#endif
MA m[C / 16];
for (uint kt = 0u; kt < NT; ++kt) LOAD_B_TOKEN(m[kt],mt*16u,kt*16u);
for (uint nt = 0u; nt < NT; ++nt) {
acc[nt] = ACC(HACC(acc[nt]) + HACC(bias_tile(p.o[3] + nt * 16u, UC)));
}
}
for (uint ch = 0u; ch < ABL_CHUNKS; ++ch)
{
const uint ow1 = p.o[44u + 3u * ch], ob1 = p.o[45u + 3u * ch], ow2 = p.o[46u + 3u * ch];
ACC h0 = bias_tile(ob1, 32u), h1 = bias_tile(ob1 + 16u, 32u);
ACC H0_[TPW], H1_[TPW];
FORW { H0_[w] = bias_tile(ob1, 32u); H1_[w] = bias_tile(ob1 + 16u, 32u); }
for (uint kt = 0u; kt < NT; kt += 2u)
{
MB b[4];
for (uint i = 0u; i < 2u; ++i) { LD(b[2u * i], wb.w, ow1 + (kt + i) * 512u, 32u); LD(b[2u * i + 1u], wb.w, ow1 + (kt + i) * 512u + 16u, 32u); }
for (uint i = 0u; i < 2u; ++i) { h0 = coopMatMulAdd(m[kt + i], b[2u * i], h0); h1 = coopMatMulAdd(m[kt + i], b[2u * i + 1u], h1); }
FORW { for (uint i = 0u; i < 2u; ++i) { H0_[w] = coopMatMulAdd(m[kt + i], b[2u * i], H0_[w]); H1_[w] = coopMatMulAdd(m[kt + i], b[2u * i + 1u], H1_[w]); } }
}
STORE_T_TOKEN(h0,mt*16u,0u,SQ); STORE_T_TOKEN(h1,mt*16u,16u,SQ);
FORW { WINV STORE_T_TOKEN(H0_[w],mt*16u,0u,SQ); STORE_T_TOKEN(H1_[w],mt*16u,16u,SQ); }
// fc2's first weight tiles are requested before the activation pass, so they arrive while it runs
MB w2a[2], w2b[2];
for (uint i = 0u; i < 2u; ++i) { LD(w2a[i], wb.w, ow2 + i * 16u, UC); LD(w2b[i], wb.w, ow2 + 16u * UC + i * 16u, UC); }
for (uint j = 0u; j < 4u; ++j) writeTToken(tok,hf*16u+j*4u,SQ,gelu(readTToken(tok,hf*16u+j*4u,SQ)));
MA a0, a1; LOAD_T_TOKEN(a0,mt*16u,0u,SQ); LOAD_T_TOKEN(a1,mt*16u,16u,SQ);
for (uint i = 0u; i < 2u; ++i) acc[i] = coopMatMulAdd(a1, w2b[i], coopMatMulAdd(a0, w2a[i], acc[i]));
MA GA0_[TPW], GA1_[TPW];
FORW {
WINV
for (uint j = 0u; j < 4u; ++j) writeTToken(tok,hf*16u+j*4u,SQ,gelu(readTToken(tok,hf*16u+j*4u,SQ)));
LOAD_T_TOKEN(GA0_[w],mt*16u,0u,SQ); LOAD_T_TOKEN(GA1_[w],mt*16u,16u,SQ);
}
FORW { for (uint i = 0u; i < 2u; ++i) acc[i] = coopMatMulAdd(GA1_[w], w2b[i], coopMatMulAdd(GA0_[w], w2a[i], acc[i])); }
for (uint n0 = 2u; n0 < NT; n0 += 2u)
{
for (uint i = 0u; i < 2u; ++i) { LD(w2a[i], wb.w, ow2 + (n0 + i) * 16u, UC); LD(w2b[i], wb.w, ow2 + 16u * UC + (n0 + i) * 16u, UC); }
for (uint i = 0u; i < 2u; ++i) acc[n0 + i] = coopMatMulAdd(a1, w2b[i], coopMatMulAdd(a0, w2a[i], acc[n0 + i]));
FORW { for (uint i = 0u; i < 2u; ++i) acc[n0 + i] = coopMatMulAdd(GA1_[w], w2b[i], coopMatMulAdd(GA0_[w], w2a[i], acc[n0 + i])); }
}
}
for (uint nt = 0u; nt < NT; ++nt) STORE_B_TOKEN(acc[nt],mt*16u,nt*16u);
FORW { WINV for (uint nt = 0u; nt < NT; ++nt) STORE_B_TOKEN(acc[nt],mt*16u,nt*16u); }
#undef m
}
// ------------------------------------------------------------ outputs (the block output is in BF, token layout)
const bool inside = ty >= 0 && ty < p.Hh && tx >= 0 && tx < p.W;
const uint dst = uint(ty * p.W + tx);
#define OUTV const bool inside = ty >= 0 && ty < p.Hh && tx >= 0 && tx < p.W; const uint dst = uint(ty * p.W + tx);
#if HEAD
SYNC // the head's rows of TT overlap other tiles' MLP rows
{
ACC ho[3];
for (uint nt = 0u; nt < 3u; ++nt) ho[nt] = bias_tile(p.o[107] + nt * 16u, 48u);
ACC HO_[TPW][3];
FORW { for (uint nt = 0u; nt < 3u; ++nt) HO_[w][nt] = bias_tile(p.o[107] + nt * 16u, 48u); }
for (uint kt = 0u; kt < NT; ++kt)
{
MA a; LOAD_B_TOKEN(a,mt*16u,kt*16u);
MB b[3];
for (uint nt = 0u; nt < 3u; ++nt) LD(b[nt], wb.w, p.o[106] + kt * 16u * 48u + nt * 16u, 48u);
for (uint nt = 0u; nt < 3u; ++nt) ho[nt] = coopMatMulAdd(a, b[nt], ho[nt]);
FORW {
WINV
MA a; LOAD_B_TOKEN(a,mt*16u,kt*16u);
for (uint nt = 0u; nt < 3u; ++nt) HO_[w][nt] = coopMatMulAdd(a, b[nt], HO_[w][nt]);
}
}
FORW {
WINV OUTV
for (uint nt = 0u; nt < 3u; ++nt) STORE_T_TOKEN(HO_[w][nt],mt*16u,nt*16u,SHD);
if (inside)
for (uint c = hf * 20u; c < hf * 20u + 20u; c += 4u) out0.a[(dst * 40u + c) / 4u] = readTToken(tok,c,SHD);
}
for (uint nt = 0u; nt < 3u; ++nt) STORE_T_TOKEN(ho[nt],mt*16u,nt*16u,SHD);
if (inside)
for (uint c = hf * 20u; c < hf * 20u + 20u; c += 4u) out0.a[(dst * 40u + c) / 4u] = readTToken(tok,c,SHD);
}
#elif KIND == 1
if (inside)
for (uint c = 0u; c < HC; c += 4u) out1.a[(dst * UC + ch0 + c) / 4u] = readBToken(tok,ch0+c);
FORW {
WINV OUTV
if (inside)
for (uint c = 0u; c < HC; c += 4u) out1.a[(dst * UC + ch0 + c) / 4u] = readBToken(tok,ch0+c);
}
#else
if (inside)
for (uint c = 0u; c < HC; c += 4u) out0.a[(dst * UC + ch0 + c) / 4u] = readBToken(tok,ch0+c);
FORW {
WINV OUTV
if (inside)
for (uint c = 0u; c < HC; c += 4u) out0.a[(dst * UC + ch0 + c) / 4u] = readBToken(tok,ch0+c);
}
#endif
#if KIND == 1
SYNC
{
// patch merge: wave mt gathers sub-token mt of the 16 merged tokens into TT, then computes output tiles mt, mt + 4, ...
const uint sourceToken = 8u * (2u * (row >> 2) + (mt >> 1)) + 2u * (row & 3u) + (mt & 1u);
for (uint c = 0u; c < HC; c += 4u)
writeTRaw(row * S4 + mt * UC + ch0 + c, readBToken(sourceToken, ch0 + c));
// patch merge: tile mt gathers sub-token mt of the 16 merged tokens into TT, then computes output tiles mt, mt + 4, ...
FORW {
WINV
const uint sourceToken = 8u * (2u * (row >> 2) + (mt >> 1)) + 2u * (row & 3u) + (mt & 1u);
for (uint c = 0u; c < HC; c += 4u)
writeTRaw(row * S4 + mt * UC + ch0 + c, readBToken(sourceToken, ch0 + c));
}
SYNC
const uint MN = uint(MERGE_N), NI = (MN / 16u + 3u) / 4u, ntn = (MN / 16u + 3u - mt) / 4u;
ACC mg[(MERGE_N / 16 + 3) / 4];
for (uint i = 0u; i < NI; ++i) mg[i] = bias_tile(p.o[105] + min(mt + 4u * i, MN / 16u - 1u) * 16u, MN);
for (uint kt = 0u; kt < 4u * UC / 16u; kt += 2u)
{
MA a[2]; MB b[2 * ((MERGE_N / 16 + 3) / 4)];
for (uint j = 0u; j < 2u; ++j)
const uint MN = uint(MERGE_N), NI = (MN / 16u + 3u) / 4u;
ACC MG_[TPW][(MERGE_N / 16 + 3) / 4];
FORW {
WINV
for (uint i = 0u; i < NI; ++i) MG_[w][i] = bias_tile(p.o[105] + min(mt + 4u * i, MN / 16u - 1u) * 16u, MN);
for (uint kt = 0u; kt < 4u * UC / 16u; kt += 2u)
{
LOAD_T_RAW(a[j],(kt+j)*16u,S4);
for (uint i = 0u; i < NI; ++i) LD(b[j * NI + i], wb.w, p.o[104] + (kt + j) * 16u * MN + min(mt + 4u * i, MN / 16u - 1u) * 16u, MN);
MA a[2]; MB b[2 * ((MERGE_N / 16 + 3) / 4)];
for (uint j = 0u; j < 2u; ++j)
{
LOAD_T_RAW(a[j],(kt+j)*16u,S4);
for (uint i = 0u; i < NI; ++i) LD(b[j * NI + i], wb.w, p.o[104] + (kt + j) * 16u * MN + min(mt + 4u * i, MN / 16u - 1u) * 16u, MN);
}
for (uint j = 0u; j < 2u; ++j)
for (uint i = 0u; i < NI; ++i) MG_[w][i] = coopMatMulAdd(a[j], b[j * NI + i], MG_[w][i]);
}
for (uint j = 0u; j < 2u; ++j)
for (uint i = 0u; i < NI; ++i) mg[i] = coopMatMulAdd(a[j], b[j * NI + i], mg[i]);
}
// BF is free (every wave gathered before the barrier above, and every tile's products are done): the merged
// tokens go there
FORW {
WINV
const uint ntn = (MN / 16u + 3u - mt) / 4u;
for (uint i = 0u; i < NI; ++i)
if (i < ntn) STORE_B_PATCH(MG_[w][i],0u,(mt+4u*i)*16u,SMG);
}
// BF is free (every wave gathered before the barrier above): the merged tokens go there
for (uint i = 0u; i < NI; ++i)
if (i < ntn) STORE_B_PATCH(mg[i],0u,(mt+4u*i)*16u,SMG);
SYNC
const int ym = tdiv2(Y0) + int(row >> 2), xm = tdiv2(X0) + int(row & 3u);
const uint part = uint(XC) / 8u, c0 = (mt * 2u + hf) * part;
if (ym >= 0 && ym < p.Hh / 2 && xm >= 0 && xm < p.W / 2)
for (uint c = c0; c < c0 + part; c += 4u) out0.a[(uint(ym * (p.W / 2) + xm) * uint(XC) + c) / 4u] = readBPatch(row,c,SMG);
const uint part = uint(XC) / 8u;
FORW {
WINV
const uint c0 = (mt * 2u + hf) * part;
if (ym >= 0 && ym < p.Hh / 2 && xm >= 0 && xm < p.W / 2)
for (uint c = c0; c < c0 + part; c += 4u) out0.a[(uint(ym * (p.W / 2) + xm) * uint(XC) + c) / 4u] = readBPatch(row,c,SMG);
}
}
#endif
}
engine/layer_m.comp
+163 -68
View File
@@ -57,7 +57,17 @@
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_GOOGLE_include_directive : require
layout(local_size_x = 128) in;
#extension GL_EXT_control_flow_attributes : require
// WPW: 4x4-token windows per wave (1, 2 or 4). A workgroup is always one 8x8-token block of four windows; with
// WPW > 1 a wave walks WPW of them and every weight tile is loaded once for all of them (a window reads the whole
// layer's weights, which is a third of the network's time with one window per wave). Values do not change.
#ifndef WPW
#define WPW 1
#endif
#if WPW > 1 && KIND >= 3
#error layer_m.comp: L's input and expansion stages stage their pixels per thread (WPW = 1)
#endif
layout(local_size_x = 128 / WPW) in;
// o: 0 g1, 1 bo, 2 g2, 3 b2 (f16); per head h at 4 + 4h: q, v (FP8), bias (f16), wo (FP8);
// per MLP chunk c at 20 + 3c: w1 (FP8), b1 (f16), w2 (FP8); 68, 69: merge or expand weights (FP8) and bias (f16)
layout(set = 0, binding = 0, std430) readonly buffer Pb { int W; int Hh; int sx; int sy; uint o[72]; } p;
@@ -144,6 +154,17 @@ const float16_t MCL = float16_t(6.2831855);
// RADV folds a register-only f32->f16->f32 conversion away (measured: half of all values and every tie
// midpoint unchanged), so the rounding has to pass through binary16 memory. Each invocation owns 8 halves of
// its own scratch: no cross-lane traffic, so a subgroup fence is enough (no workgroup barrier).
// M_RELAXED (M_DEFS=-DM_RELAXED for compile_m.py): this GPU's own arithmetic instead. The accumulator stays f32
// across the 32-deep steps and is rounded once where it is used, with no scratch round trip; a layer is then no
// longer bit-identical to NVIDIA's kernel (FP8 codes flip by one step here and there), which two NVIDIA setups
// are not either (engine/rr/README.md).
#ifdef M_RELAXED
#define R16(m) { }
#define K32H(dst, a0, b0, a1, b1) { \
ACC P_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(dst))); \
(dst)[0] = float16_t(P_[0]); (dst)[1] = float16_t(P_[1]); (dst)[2] = float16_t(P_[2]); (dst)[3] = float16_t(P_[3]); \
(dst)[4] = float16_t(P_[4]); (dst)[5] = float16_t(P_[5]); (dst)[6] = float16_t(P_[6]); (dst)[7] = float16_t(P_[7]); }
#else
shared f16vec4 RS[128 * 2];
#define R16(m) { \
RS[gl_LocalInvocationID.x * 2u + 0u] = f16vec4(float16_t((m)[0]), float16_t((m)[1]), float16_t((m)[2]), float16_t((m)[3])); \
@@ -152,6 +173,7 @@ shared f16vec4 RS[128 * 2];
const f16vec4 r0_ = RS[gl_LocalInvocationID.x * 2u + 0u], r1_ = RS[gl_LocalInvocationID.x * 2u + 1u]; \
(m)[0] = float(r0_.x); (m)[1] = float(r0_.y); (m)[2] = float(r0_.z); (m)[3] = float(r0_.w); \
(m)[4] = float(r1_.x); (m)[5] = float(r1_.y); (m)[6] = float(r1_.z); (m)[7] = float(r1_.w); }
#endif
// one logical k32, exactly as the native K32_MIX: P = A0*B0 + A1*B1 from a zero seed, then one binary16
// rounding of P + the previous accumulator (the native v_fma_mix), through the scratch so it is not folded.
// Packed fixed-16 weight tiles (test_m.pack8): matrix W[K][N] as [K/16][N/16][N16][K16], so one 16x16 A tile
@@ -159,7 +181,8 @@ shared f16vec4 RS[128 * 2];
// operand values as the old plain W[K][N] ColumnMajor stride-N load, but a plain (non-transposed) load.
#define LDA8(dst, base, ks, stride, moff) coopMatLoad((dst), w8.w, (base) + 256u * (uint(ks) * (uint(stride) / 16u) + (uint(moff) / 16u)), 16u, RM)
#if TEXL
#if defined(M_RELAXED)
#elif TEXL
// ZLUDA's lowering of m16n8k32 (K32_NO_MIX): the accumulator enters the first product as C, f32 throughout the step,
// one binary16 rounding at its end
#define K32H(dst, a0, b0, a1, b1) { \
@@ -186,11 +209,32 @@ shared f16vec4 RS[128 * 2];
#endif
// a tiled f16 vector (16 equal rows) as the initial value of an accumulator whose rows are its channels
HAC bias16(uint off, uint stride) { HAC b; coopMatLoad(b, w16.w, off / 4u, stride / 4u, CM); return b; }
#ifdef M_RELAXED
// straight from f32 into FP8, two elements per saturating conversion; the activation in f32 as well
#define SAT2_(dst, i, v) { fe4m3vec2 t_; saturatedConvertEXT(t_, (v)); (dst)[i] = t_.x; (dst)[(i) + 1] = t_.y; }
#define MK8_F8(dst, a) { SAT2_(dst, 0, vec2((a)[0], (a)[1])) SAT2_(dst, 2, vec2((a)[2], (a)[3])) SAT2_(dst, 4, vec2((a)[4], (a)[5])) SAT2_(dst, 6, vec2((a)[6], (a)[7])) }
#define GELU2_(dst, i, h) { const vec2 x_ = vec2((h)[i], (h)[(i) + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); \
const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.08108133), vec2(0.41216981)), vec2(0.5)); SAT2_(dst, i, x_ * g_) }
#define GELU8(dst, h) { GELU2_(dst, 0, h) GELU2_(dst, 2, h) GELU2_(dst, 4, h) GELU2_(dst, 6, h) }
#else
#define MK8_F8(dst, a) { \
const float16_t q0_ = CL8(Q16((a)[0])), q1_ = CL8(Q16((a)[1])), q2_ = CL8(Q16((a)[2])), q3_ = CL8(Q16((a)[3])); \
const float16_t q4_ = CL8(Q16((a)[4])), q5_ = CL8(Q16((a)[5])), q6_ = CL8(Q16((a)[6])), q7_ = CL8(Q16((a)[7])); \
(dst)[0] = floate4m3_t(q0_); (dst)[1] = floate4m3_t(q1_); (dst)[2] = floate4m3_t(q2_); (dst)[3] = floate4m3_t(q3_); \
(dst)[4] = floate4m3_t(q4_); (dst)[5] = floate4m3_t(q5_); (dst)[6] = floate4m3_t(q6_); (dst)[7] = floate4m3_t(q7_); }
#endif
// the normalised tokens into FP8 (inside NM1 / N2: x0..x7 the token, g0 / g1 the gains)
#ifdef M_RELAXED
#define H1N(N, r_) { const float f_ = float(r_); \
SAT2_(h1[N], 0, vec2(x0, x1) * (f_ * vec2(g0.xy))) SAT2_(h1[N], 2, vec2(x2, x3) * (f_ * vec2(g0.zw))) \
SAT2_(h1[N], 4, vec2(x4, x5) * (f_ * vec2(g1.xy))) SAT2_(h1[N], 6, vec2(x6, x7) * (f_ * vec2(g1.zw))) }
#else
#define H1N(N, r_) { \
h1[N][0] = floate4m3_t(CL8(x0 * float16_t(r_ * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(r_ * g0.y))); \
h1[N][2] = floate4m3_t(CL8(x2 * float16_t(r_ * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(r_ * g0.w))); \
h1[N][4] = floate4m3_t(CL8(x4 * float16_t(r_ * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(r_ * g1.y))); \
h1[N][6] = floate4m3_t(CL8(x6 * float16_t(r_ * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(r_ * g1.w))); }
#endif
// NVIDIA's softmax numerator, two neighbouring columns in one 32-bit word (the even column's carry reaches the odd)
void softmax2(float16_t s0, float16_t s1, out float16_t e0, out float16_t e1)
{
@@ -364,18 +408,24 @@ float tanhL(float t)
#endif
void main()
{
const uint mt = gl_LocalInvocationID.x >> 5, lane = gl_LocalInvocationID.x & 31u;
const uint mt0 = (gl_LocalInvocationID.x >> 5) * uint(WPW), lane = gl_LocalInvocationID.x & 31u;
const uint row = lane & 15u, hf = lane >> 4, ch8 = 8u * hf; // one token per lane, its channels 8h..8h+7
const int X0 = 8 * int(gl_WorkGroupID.x) - p.sx, Y0 = 8 * int(gl_WorkGroupID.y) - p.sy;
const uint bxr = 4u * (mt & 1u) + (row & 3u), byr = 4u * (mt >> 1) + (row >> 2);
const int tx = X0 + int(bxr), ty = Y0 + int(byr);
const bool inside = ty >= 0 && ty < p.Hh && tx >= 0 && tx < p.W;
// per window w of this wave (FORW): its index mt in the block and its token's position (WINV)
#define FORW [[unroll]] for (uint w = 0u; w < uint(WPW); ++w)
#define WINV const uint mt = mt0 + w; \
const uint bxr = 4u * (mt & 1u) + (row & 3u), byr = 4u * (mt >> 1) + (row >> 2); \
const int tx = X0 + int(bxr), ty = Y0 + int(byr); \
const bool inside = ty >= 0 && ty < p.Hh && tx >= 0 && tx < p.W; \
const uint iy = uint(mirror(ty, p.Hh)), ix = uint(mirror(tx, p.W));
HAC res[NT]; // the residual stream: rows are channels, columns are tokens
B8 h1[NT]; // the normalised activation, the K operand of every product
HAC RES_[WPW][NT]; // the residual stream: rows are channels, columns are tokens
B8 H1_[WPW][NT]; // the normalised activation, the K operand of every product
#define res RES_[w]
#define h1 H1_[w]
// ------------------------------------------------------------ patch expand: x0 = q8(E) + skip
#if EXPAND
{
FORW {
WINV
const uint W2 = uint(p.W) / 2u, H2 = uint(p.Hh) / 2u, part = UXC / 8u, c0 = (mt * 2u + hf) * part;
const uint lt = uint(mirror((Y0 >> 1) + int(row >> 2), int(H2)) * int(W2) + mirror((X0 >> 1) + int(row & 3u), p.W / 2));
#if PLANES
@@ -383,7 +433,10 @@ void main()
#else
for (uint c = c0; c < c0 + part; ++c) NLI[row * SLO + c] = inb.a[lt * UXC + c];
#endif
SYNC
}
SYNC
FORW {
WINV
// this wave computes sub-position mt of every low-resolution token's 4C columns
HAC ef[NT];
#define EF_INIT(N) ef[N] = bias16(p.o[69] + 16u * (mt * NT + (N)), 4u * UC);
@@ -401,12 +454,13 @@ void main()
EACH_T(EF_K32)
}
EACH_T(EF_STORE)
SYNC
}
SYNC
#endif
#if KIND == 3
// ------------------------------------------------------------ L's input: features, embedding -> x0
{
FORW {
WINV
const float e1 = ep.expo == 0.0 ? 1.0 : ep.expo, e2 = ep.expoPrev == 0.0 ? 1.0 : ep.expoPrev;
const uint li = gl_LocalInvocationID.x;
if (li < 64u)
@@ -445,6 +499,8 @@ void main()
}
#endif
// ------------------------------------------------------------ norm 1: sum of squares, residual and h1
FORW {
WINV
#if EXPAND
const uint lb0 = ((byr >> 1) * 4u + (bxr >> 1)) * S4 + ((byr & 1u) * 2u + (bxr & 1u)) * UC + ch8;
const uint sb0 = (iy * uint(p.W) + ix) * UC + ch8;
@@ -483,47 +539,74 @@ void main()
res[N][2] = float16_t(x2 + d0.z); res[N][3] = float16_t(x3 + d0.w); \
res[N][4] = float16_t(x4 + d1.x); res[N][5] = float16_t(x5 + d1.y); \
res[N][6] = float16_t(x6 + d1.z); res[N][7] = float16_t(x7 + d1.w); \
h1[N][0] = floate4m3_t(CL8(x0 * float16_t(rn * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(rn * g0.y))); \
h1[N][2] = floate4m3_t(CL8(x2 * float16_t(rn * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(rn * g0.w))); \
h1[N][4] = floate4m3_t(CL8(x4 * float16_t(rn * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(rn * g1.y))); \
h1[N][6] = floate4m3_t(CL8(x6 * float16_t(rn * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(rn * g1.w))); }
H1N(N, rn) }
EACH_T(NM1)
}
// ------------------------------------------------------------ attention, one head at a time
for (uint hd = 0u; hd < uint(NH); ++hd)
{
const uint oq = p.o[4u + 4u * hd], ov = p.o[5u + 4u * hd], ob = p.o[6u + 4u * hd], ow = p.o[7u + 4u * hd];
HAC q[2], v[2];
q[0] = HAC(0.0); q[1] = HAC(0.0); v[0] = HAC(0.0); v[1] = HAC(0.0);
#define QV_G(A, B) { \
// Q of every window, then V of every window: each weight tile is loaded once for the wave's windows
A8 QA_[WPW][2];
{
HAC q[WPW][2];
FORW { q[w][0] = HAC(0.0); q[w][1] = HAC(0.0); }
#define Q_G(A, B) { \
const uint kt = A; \
A8 aq0, aq1; \
LDA8(aq0, oq, kt, 32u, 0u); \
LDA8(aq1, oq, kt + 1u, 32u, 0u); \
FORW { K32H(q[w][0], aq0, h1[A], aq1, h1[B]) } \
LDA8(aq0, oq, kt, 32u, 16u); \
LDA8(aq1, oq, kt + 1u, 32u, 16u); \
FORW { K32H(q[w][1], aq0, h1[A], aq1, h1[B]) } }
EACH_T16(Q_G)
FORW {
for (uint m = 0u; m < 2u; ++m)
{
#ifdef M_RELAXED
MK8_F8(QA_[w][m], q[w][m])
#else
const float16_t q0 = CL8(q[w][m][0]), q1 = CL8(q[w][m][1]), q2 = CL8(q[w][m][2]), q3 = CL8(q[w][m][3]);
const float16_t q4 = CL8(q[w][m][4]), q5 = CL8(q[w][m][5]), q6 = CL8(q[w][m][6]), q7 = CL8(q[w][m][7]);
QA_[w][m][0] = floate4m3_t(q0); QA_[w][m][1] = floate4m3_t(q1);
QA_[w][m][2] = floate4m3_t(q2); QA_[w][m][3] = floate4m3_t(q3);
QA_[w][m][4] = floate4m3_t(q4); QA_[w][m][5] = floate4m3_t(q5);
QA_[w][m][6] = floate4m3_t(q6); QA_[w][m][7] = floate4m3_t(q7);
#endif
}
}
}
{
HAC v[WPW][2];
FORW { v[w][0] = HAC(0.0); v[w][1] = HAC(0.0); }
#define V_G(A, B) { \
const uint kt = A; \
A8 aq0, aq1, av0, av1; \
LDA8(aq0, oq, kt, 32u, 0u); \
LDA8(aq1, oq, kt + 1u, 32u, 0u); \
A8 av0, av1; \
LDA8(av0, ov, kt, 32u, 0u); \
LDA8(av1, ov, kt + 1u, 32u, 0u); \
K32H(q[0], aq0, h1[A], aq1, h1[B]) \
K32H(v[0], av0, h1[A], av1, h1[B]) \
LDA8(aq0, oq, kt, 32u, 16u); \
LDA8(aq1, oq, kt + 1u, 32u, 16u); \
FORW { K32H(v[w][0], av0, h1[A], av1, h1[B]) } \
LDA8(av0, ov, kt, 32u, 16u); \
LDA8(av1, ov, kt + 1u, 32u, 16u); \
K32H(q[1], aq0, h1[A], aq1, h1[B]) \
K32H(v[1], av0, h1[A], av1, h1[B]) }
EACH_T16(QV_G)
// scores: q8(Q) against itself, rows are the queries and columns the keys, plus the relative position bias
A8 qa[2]; B8 qb[2];
for (uint m = 0u; m < 2u; ++m)
{
const float16_t q0 = CL8(q[m][0]), q1 = CL8(q[m][1]), q2 = CL8(q[m][2]), q3 = CL8(q[m][3]);
const float16_t q4 = CL8(q[m][4]), q5 = CL8(q[m][5]), q6 = CL8(q[m][6]), q7 = CL8(q[m][7]);
qa[m][0] = floate4m3_t(q0); qb[m][0] = qa[m][0]; qa[m][1] = floate4m3_t(q1); qb[m][1] = qa[m][1];
qa[m][2] = floate4m3_t(q2); qb[m][2] = qa[m][2]; qa[m][3] = floate4m3_t(q3); qb[m][3] = qa[m][3];
qa[m][4] = floate4m3_t(q4); qb[m][4] = qa[m][4]; qa[m][5] = floate4m3_t(q5); qb[m][5] = qa[m][5];
qa[m][6] = floate4m3_t(q6); qb[m][6] = qa[m][6]; qa[m][7] = floate4m3_t(q7); qb[m][7] = qa[m][7];
FORW { K32H(v[w][1], av0, h1[A], av1, h1[B]) } }
EACH_T16(V_G)
// V of every window goes to its shared buffer now, so it does not occupy registers during the scores
FORW {
const uint mt = mt0 + w;
coopMatStore(v[w][0], VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM);
coopMatStore(v[w][1], VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM);
}
// the bias is read column-major, so the accumulator holds S^T: this lane holds the keys 8h..8h+7 of one query
HAC s = bias16(ob, 16u);
K32H(s, qa[0], qb[0], qa[1], qb[1])
}
// scores: q8(Q) against itself, rows are the queries and columns the keys, plus the relative position bias
// (read column-major, so the accumulator holds S^T: this lane holds the keys 8h..8h+7 of one query)
const HAC sbias = bias16(ob, 16u);
B8 OB_[WPW][2];
FORW {
WINV
B8 qb[2];
for (uint m = 0u; m < 2u; ++m) for (int i = 0; i < 8; ++i) qb[m][i] = QA_[w][m][i];
HAC s = sbias;
K32H(s, QA_[w][0], qb[0], QA_[w][1], qb[1])
float16_t e[8];
softmax2(float16_t(s[0]), float16_t(s[1]), e[0], e[1]);
softmax2(float16_t(s[2]), float16_t(s[3]), e[2], e[3]);
@@ -540,8 +623,6 @@ void main()
pb[4] = G2_.x * rc; pb[5] = G2_.y * rc; pb[6] = G3_.x * rc; pb[7] = G3_.y * rc;
// V through this wave's shared buffer, then O^T = V^T x P^T: rows are the head channels, columns the tokens
ACC o[2];
coopMatStore(v[0], VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM);
coopMatStore(v[1], VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM);
A16 vt0, vt1;
coopMatLoad(vt0, VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM);
coopMatLoad(vt1, VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM);
@@ -554,20 +635,22 @@ void main()
o[0] = coopMatMulAdd(vt0, pb, ACC(0.0));
o[1] = coopMatMulAdd(vt1, pb, ACC(0.0));
R16(o[0]); R16(o[1]);
B8 ob8[2];
MK8_F8(ob8[0], o[0]); MK8_F8(ob8[1], o[1]);
MK8_F8(OB_[w][0], o[0]); MK8_F8(OB_[w][1], o[1]);
}
// the output projection of this head's 32 attention channels, straight into the residual stream
#define WO_N(N) { \
A8 a0, a1; \
LDA8(a0, ow, 0u, UC, 16u * (N)); \
LDA8(a1, ow, 1u, UC, 16u * (N)); \
K32H(res[N], a0, ob8[0], a1, ob8[1]) }
FORW { K32H(res[N], a0, OB_[w][0], a1, OB_[w][1]) } }
EACH_T(WO_N)
}
// ------------------------------------------------------------ norm 2
// the native element (T, 8h+i) holds the natural channel tok_xnat(T, 8h+i) = 16T + 4h + (i&3) + 8*((i>>2)&1),
// which lives in my natural-order accumulator at the half ((i>>2)&1) and the element 4h+(i&3). The tree needs
// the pairs in that order, so fetch the (pair 2h, pair 2h+1) of the half ((i>>2)&1) and select by this lane's half.
FORW {
WINV
#if TEXL
#define EV(T, I) float16_t(res[T][I]) // natural order: the lane's own elements
#else
@@ -592,50 +675,55 @@ void main()
const f16vec4 g0 = w16.w[wb], g1 = w16.w[wb + 1u], d0 = w16.w[bb], d1 = w16.w[bb + 1u]; \
const float16_t x0 = res[N][0], x1 = res[N][1], x2 = res[N][2], x3 = res[N][3]; \
const float16_t x4 = res[N][4], x5 = res[N][5], x6 = res[N][6], x7 = res[N][7]; \
h1[N][0] = floate4m3_t(CL8(x0 * float16_t(rn2 * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(rn2 * g0.y))); \
h1[N][2] = floate4m3_t(CL8(x2 * float16_t(rn2 * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(rn2 * g0.w))); \
h1[N][4] = floate4m3_t(CL8(x4 * float16_t(rn2 * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(rn2 * g1.y))); \
h1[N][6] = floate4m3_t(CL8(x6 * float16_t(rn2 * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(rn2 * g1.w))); \
H1N(N, rn2) \
res[N][0] = float16_t(x0 + d0.x); res[N][1] = float16_t(x1 + d0.y); \
res[N][2] = float16_t(x2 + d0.z); res[N][3] = float16_t(x3 + d0.w); \
res[N][4] = float16_t(x4 + d1.x); res[N][5] = float16_t(x5 + d1.y); \
res[N][6] = float16_t(x6 + d1.z); res[N][7] = float16_t(x7 + d1.w); }
EACH_T(N2)
}
// ------------------------------------------------------------ the MLP, 32 hidden channels at a time
for (uint c = 0u; c < UC / 8u; ++c)
{
const uint ow1 = p.o[20u + 3u * c], ob1 = p.o[21u + 3u * c], ow2 = p.o[22u + 3u * c];
HAC hm[2];
hm[0] = bias16(ob1, 32u); hm[1] = bias16(ob1 + 16u, 32u);
HAC hm[WPW][2];
{ const HAC b0_ = bias16(ob1, 32u), b1_ = bias16(ob1 + 16u, 32u); FORW { hm[w][0] = b0_; hm[w][1] = b1_; } }
#define MLP_G(A, B) { \
const uint kt = A; \
A8 a0, a1; \
LDA8(a0, ow1, kt, 32u, 0u); \
LDA8(a1, ow1, kt + 1u, 32u, 0u); \
K32H(hm[0], a0, h1[A], a1, h1[B]) \
FORW { K32H(hm[w][0], a0, h1[A], a1, h1[B]) } \
LDA8(a0, ow1, kt, 32u, 16u); \
LDA8(a1, ow1, kt + 1u, 32u, 16u); \
K32H(hm[1], a0, h1[A], a1, h1[B]) }
FORW { K32H(hm[w][1], a0, h1[A], a1, h1[B]) } }
EACH_T16(MLP_G)
B8 g0, g1;
{ const float16_t v0 = gelu1(hm[0][0]), v1 = gelu1(hm[0][1]), v2 = gelu1(hm[0][2]), v3 = gelu1(hm[0][3]);
const float16_t v4 = gelu1(hm[0][4]), v5 = gelu1(hm[0][5]), v6 = gelu1(hm[0][6]), v7 = gelu1(hm[0][7]);
g0[0] = floate4m3_t(CL8(v0)); g0[1] = floate4m3_t(CL8(v1)); g0[2] = floate4m3_t(CL8(v2)); g0[3] = floate4m3_t(CL8(v3));
g0[4] = floate4m3_t(CL8(v4)); g0[5] = floate4m3_t(CL8(v5)); g0[6] = floate4m3_t(CL8(v6)); g0[7] = floate4m3_t(CL8(v7)); }
{ const float16_t v0 = gelu1(hm[1][0]), v1 = gelu1(hm[1][1]), v2 = gelu1(hm[1][2]), v3 = gelu1(hm[1][3]);
const float16_t v4 = gelu1(hm[1][4]), v5 = gelu1(hm[1][5]), v6 = gelu1(hm[1][6]), v7 = gelu1(hm[1][7]);
g1[0] = floate4m3_t(CL8(v0)); g1[1] = floate4m3_t(CL8(v1)); g1[2] = floate4m3_t(CL8(v2)); g1[3] = floate4m3_t(CL8(v3));
g1[4] = floate4m3_t(CL8(v4)); g1[5] = floate4m3_t(CL8(v5)); g1[6] = floate4m3_t(CL8(v6)); g1[7] = floate4m3_t(CL8(v7)); }
B8 G0_[WPW], G1_[WPW];
FORW {
#ifdef M_RELAXED
GELU8(G0_[w], hm[w][0]) GELU8(G1_[w], hm[w][1])
#else
{ const float16_t v0 = gelu1(hm[w][0][0]), v1 = gelu1(hm[w][0][1]), v2 = gelu1(hm[w][0][2]), v3 = gelu1(hm[w][0][3]);
const float16_t v4 = gelu1(hm[w][0][4]), v5 = gelu1(hm[w][0][5]), v6 = gelu1(hm[w][0][6]), v7 = gelu1(hm[w][0][7]);
G0_[w][0] = floate4m3_t(CL8(v0)); G0_[w][1] = floate4m3_t(CL8(v1)); G0_[w][2] = floate4m3_t(CL8(v2)); G0_[w][3] = floate4m3_t(CL8(v3));
G0_[w][4] = floate4m3_t(CL8(v4)); G0_[w][5] = floate4m3_t(CL8(v5)); G0_[w][6] = floate4m3_t(CL8(v6)); G0_[w][7] = floate4m3_t(CL8(v7)); }
{ const float16_t v0 = gelu1(hm[w][1][0]), v1 = gelu1(hm[w][1][1]), v2 = gelu1(hm[w][1][2]), v3 = gelu1(hm[w][1][3]);
const float16_t v4 = gelu1(hm[w][1][4]), v5 = gelu1(hm[w][1][5]), v6 = gelu1(hm[w][1][6]), v7 = gelu1(hm[w][1][7]);
G1_[w][0] = floate4m3_t(CL8(v0)); G1_[w][1] = floate4m3_t(CL8(v1)); G1_[w][2] = floate4m3_t(CL8(v2)); G1_[w][3] = floate4m3_t(CL8(v3));
G1_[w][4] = floate4m3_t(CL8(v4)); G1_[w][5] = floate4m3_t(CL8(v5)); G1_[w][6] = floate4m3_t(CL8(v6)); G1_[w][7] = floate4m3_t(CL8(v7)); }
#endif
}
#define MLP2_N(N) { \
A8 a0, a1; \
LDA8(a0, ow2, 0u, UC, 16u * (N)); \
LDA8(a1, ow2, 1u, UC, 16u * (N)); \
K32H(res[N], a0, g0, a1, g1) }
FORW { K32H(res[N], a0, G0_[w], a1, G1_[w]) } }
EACH_T(MLP2_N)
}
#if KIND == 4
// ------------------------------------------------------------ L's expansion head: per-pixel reconstruction inputs
{
FORW {
WINV
const uint hb = (mt * 16u + row) * HS + ch8;
for (uint nt = 0u; nt < NT; ++nt)
for (int i = 0; i < 8; ++i) HO[hb + 16u * nt + uint(i)] = float16_t(res[nt][i]);
@@ -661,6 +749,8 @@ void main()
}
#else
// ------------------------------------------------------------ the block output, FP8
FORW {
WINV
const uint dst = uint(ty * p.W + tx) * UC + ch8;
if (inside)
for (uint nt = 0u; nt < NT; ++nt)
@@ -669,15 +759,20 @@ void main()
OUT8(0) OUT8(1) OUT8(2) OUT8(3) OUT8(4) OUT8(5) OUT8(6) OUT8(7)
#undef OUT8
}
}
#endif
// ------------------------------------------------------------ patch merge: 2x2 tokens (4C) -> XC, half resolution
#if MERGE
{
FORW {
WINV
// the 16 merged rows are gathered from the four waves: the row is the merged token of the lane's pixel
const uint mrow = (byr >> 1) * 4u + (bxr >> 1), qq = (byr & 1u) * 2u + (bxr & 1u);
#define MG_W(N) { A8 y; MK8_F8(y, res[N]); coopMatStore(y, NB, 16u * S4 + mrow * S4 + qq * UC + 16u * (N) - row * S4, S4, RM); }
EACH_T(MG_W)
SYNC
}
SYNC
FORW {
WINV
const uint MT = UXC / 16u, NI = (MT + 3u) / 4u;
const int ym = (Y0 >> 1) + int(row >> 2), xm = (X0 >> 1) + int(row & 3u);
const bool mok = ym >= 0 && ym < p.Hh / 2 && xm >= 0 && xm < p.W / 2;
engine/replay.cpp
+3 -2
View File
@@ -270,7 +270,7 @@ struct Image {
auto r = region(w,h); vkCmdCopyBufferToImage(context.command, upload.handle, handle, VK_IMAGE_LAYOUT_GENERAL, 1, &r);
}
};
struct Frame { d4r::InputParams input; d4r::OutputParams output; std::array<d4r::LayerShape,d4r::KLayers> windows; unsigned color, motion, depth; };
struct Frame { d4r::InputParams input; d4r::OutputParams output; std::array<d4r::LayerShape,d4r::KLayers> windows; unsigned color, motion, depth;; bool autoWindows = false; };
int run(const std::string& modelDir, const std::string& caseFile, const std::string& outputDir, int repeats, bool hip) {
#ifndef D4R_ENGINE_HIP
@@ -323,7 +323,7 @@ int run(const std::string& modelDir, const std::string& caseFile, const std::str
>> settings.motionOffset[0] >> settings.motionOffset[1] >> settings.sharpness;
if (!s || f.color >= images.size() || reset > 1) throw std::runtime_error("malformed standard frame declaration");
settings.renderWidth = images[f.color]->w; settings.renderHeight = images[f.color]->h; settings.reset = reset != 0;
auto params = d4r::makeKFrameParams(shape,settings); f.input = params.input; f.output = params.output; f.windows = shape.layers;
auto params = d4r::makeKFrameParams(shape,settings); f.input = params.input; f.output = params.output; f.windows = shape.layers; f.autoWindows = true;
frames.push_back(f);
} else if (!kind.empty()) throw std::runtime_error("unknown case declaration " + kind);
}
@@ -334,6 +334,7 @@ int run(const std::string& modelDir, const std::string& caseFile, const std::str
if (hip) {
for (size_t i = 0; i < frames.size(); ++i) {
const auto expected = d4r::makeKWindows(shape, uint32_t(i));
if (frames[i].autoWindows) { frames[i].windows = expected; continue; } // ordinary frames: NGX's window cycle
for (unsigned layer = 0; layer < d4r::KLayers; ++layer) {
const auto& a = frames[i].windows[layer];
const auto& b = expected[layer];
engine/replay_m.cpp
+3 -1
View File
@@ -149,6 +149,8 @@ int main(int argc, char** argv) try {
}
std::vector<double> times, walls; std::vector<std::vector<double>> stage(6);
const auto now = [] { return std::chrono::duration<double, std::micro>(std::chrono::steady_clock::now().time_since_epoch()).count(); };
// D4R_M_RESET_ALL=1: every frame is a reset frame, as the D3D12 harness passes them without its temporal option
const bool resetAll = std::getenv("D4R_M_RESET_ALL") != nullptr;
for (int pass = 0; pass <= repeats; ++pass)
for (int f = 0; f < frames; ++f) {
char** a = argv + sep + 1 + 5 * f;
@@ -159,7 +161,7 @@ int main(int argc, char** argv) try {
VkBufferImageCopy region{}; region.imageSubresource = {VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1}; region.imageExtent = {rw, rh, 1};
if (pass == 0 || frames > 1)
for (int k = 0; k < 3; ++k) vkCmdCopyBufferToImage(cb, in[k].buffer, images[k]->image, VK_IMAGE_LAYOUT_GENERAL, 1, &region);
d4r::MFrameSettings s; s.jitter[0] = std::atof(a[3]); s.jitter[1] = std::atof(a[4]); s.reset = pass == 0 && f == 0;
d4r::MFrameSettings s; s.jitter[0] = std::atof(a[3]); s.jitter[1] = std::atof(a[4]); s.reset = (pass == 0 && f == 0) || resetAll;
const double t0 = now();
engine.setFrame({color.view, motion.view, depth.view}, s);
vkCmdResetQueryPool(cb, qp, 0, 7);
engine/rr/README.md
+145
View File
@@ -353,3 +353,148 @@ scripts/run_rr_engine_compare.sh`, frames aligned by one):
rounds, a last-bit difference in a few elements. With a non-identity camera, hit distance and mask bound and a
moving scene that showed as 40 - 50 dB on later frames in a synthetic scene whose affected feature is rounding
noise around zero in both paths; no other configuration was affected.
## Against a real NVIDIA card: `D4R_RR_FUSE` (2026-10-13)
Every reference above is the CUDA path through ZLUDA. NVIDIA's own kernels (310.9.1's `rr2_*` PTX, replayed through
plain CUDA on an RTX PRO 4500 Blackwell, driver 580; tools and logs in `build/rr/f/nv/`) give an image 39-42 dB from
that reference on every frame of capF, capF4K and capFQ. Register bisection between the two (`probe_pack.py`,
`probe_cmp.py`: a debug store after every instruction, run on both) found why:
- **Fused multiply-adds.** NVIDIA's compiler contracts a `mul.f16` / `mul.f16x2` whose product only feeds an
`add` / `sub` into one fused multiply-add; the kernels give those instructions no rounding modifier. ZLUDA runs them
as written. `ptx_contract.py` is the rule (single-use product, same block, the first operand when both are
products): with it ZLUDA's result of every network kernel matches the card's (0 of ~5900 4 KiB chunks for most
launches, a handful where an `ex2.approx` / `div.approx` differs by one f32 step). The FP8 products, the e4m3
conversions and every unfiltered texture fetch were already identical. The sites are the residual steps
`out = in + branch * s`, the MLP activation `x (0.5 + y (0.412 - 0.081 |y|))` and the guide's bilinear blend.
Probing a product stops the card from fusing it: compare unprobed results, or probe everything.
- **Filtered fetches of 8-bit textures.** A filtered fetch of the SNORM8 normals returns this card's value (8-bit
weights) rounded to a 16-bit step: `f32(roundEven(v * 32767) / 32767)`, exact on 256 of 256 probed values. The
output filter's UNORM8 fetches differ too; no model for those yet.
- **Approximate instructions.** `lg2.approx` near 1 differs by ~8e-8 absolute, `ex2.approx` / `rcp.approx` /
`div.approx` by one f32 step. Not emulated.
`D4R_RR_FUSE=1 compile_rr2.py CAPTURE OUT --ptx-dir <build/rr/f/nv/mkptx.sh output> --native DLL` builds the model
with the first two: `rr2_m.comp -DRR_FUSE` (`MADD` / `MADD4`), the activation table from `enc1_ref.act` with
`FUSE`, the translated parts from the rewritten modules, encoder 0's normals through `ptx2rr.py --snorm16 528`. All
14 native stages are byte-identical to ZLUDA running the rewritten modules (`build/rr/f/nv/nfuse.sh`), and so is
the whole model on capF's three frames. Presented image against the NVIDIA card, 640x360 -> 1080p:
| frame | default build | `D4R_RR_FUSE` |
|---|---|---|
| 1 (reset) | 40.9 dB, 27.3% of values equal | 85.6 dB, 99.5% equal |
| 2 | 39.2 dB | 43.0 dB |
| 3 | 39.3 dB | 41.8 dB |
Temporal frames stay near 42 dB: the remaining one-step differences (approximate instructions, the UNORM8 fetches)
enter encoder 0's reprojection and the recurrent history amplifies them, as the one-step values of the output stage
did before the fused-convert fix. 1280x720 -> 4K: 5.68 / 6.24 ms (default 5.78 / 6.35). The default build is
unchanged and still byte-identical to the captures.
### This GPU's own arithmetic for preset F: `D4R_RR_RELAXED` (2026-10-13)
Two NVIDIA setups differ from each other as much as a relaxed build differs from the reference: NVIDIA's kernels on
an RTX 4000 Ada (driver 595) and on the Blackwell cards (driver 580) give images 41.0 / 39.0 / 38.9 dB apart on
capF's three frames, although their approximate instructions and filtered texture fetches are bit-identical
(`build/rr/f/nv/math/`); the `D4R_RR_FUSE` build matches the Blackwell result only. So there is no single NVIDIA
image to reproduce, and `compile_rr2.py` now takes `D4R_RR_RELAXED=1|2` like `build_native_model.sh` does for D / E
(`rr_exact.glsl`: f32 accumulators across the 32-deep steps, one rounding into FP8; 2 also leaves the conversions to
the driver). Preset F, 1280x720 -> 4K: 5.40 / 5.99 ms against 5.78 / 6.35 (network 3.45 / 3.77 against 3.83 / 4.14);
40.8 / 39.1 / 39.2 dB from the exact build and 39-41 dB from either NVIDIA card. The default stays exact.
Why only 6%: with the rounding gone the products are about a sixth of a full-resolution stage. dec0's first kernel
at 720p, relaxed (`build/rr/f/nv/exp/t.sh`, cutting parts out): 626 us, of which the two attention heads ~370, the
four MLP branches ~100-160, expansion / residual / store ~120-160, an empty dispatch 10; inside the attention the
softmax table ~30, the bias tiles ~40, weight loads ~12. Its 88 products alone would take ~108 us at the measured
198 T MAC/s. The compiled stage still has 338 binary16 <-> f32 conversions, 176 FP8 conversions and 128 clamps for
44 product instructions (biases and tables are binary16, the activation and softmax tables are indexed by binary16
bits, FP8 conversions saturate): what is left is format traffic and the four waves' shared-memory hand-offs, not
arithmetic. NVIDIA's kernel runs the same stage in 203 us with one 32-thread warp per window.
`D4R_RR_RELAXED=3` adds `RR_NATIVE` (rr2_m.comp): the residual steps and the MLP activation in f32, straight into
FP8, with no binary16 in between and no activation table. 1280x720 -> 4K: 5.17 / 5.72 ms (network 3.24 / 3.48), still
39-41 dB from the exact build and from both NVIDIA cards. dec0's first kernel at that level: 588 us; the barriers
cost nothing; reading the constants (weight tiles, bias tiles, the softmax table) costs 170 us; the products about
108; the rest is the tokens' way through shared memory and the arena, conversions and the expansion
(`build/rr/f/nv/exp/run1.sh`).
What bounds a stage (2026-10-13; `build/rr/f/nv/exp/`). dec0's first kernel at level 3 takes 581 us at 720p and a
wave executes about 2500 instructions (`dyn.py`): 96 products, 356 binary16 <-> f32 conversions, 296 FP8
conversions, 224 clamps of the saturating FP8 conversions, 142 buffer loads, 75 shared-memory operations, about 400
wait / delay markers. What was tried on top of level 3, none of it a gain:
- the clamps removed: -30 us (2.2 ps per instruction and wave), but not usable: without the clamp the hardware
conversion returns NaN above 464 where the clamp gives 448 (`satprobe.comp`);
- f32 copies of the biases: 176 fewer conversions and 78 more loads per wave, slightly slower (587-595 us): a load
costs about three arithmetic instructions;
- the softmax cubic computed, not read from its table: the same time;
- no barriers: the same time;
- one wave per window walking the four token slices, every weight, key and value tile read once
(`rr2_w.comp`, byte-identical to the four-wave build): 27% fewer instructions per window, 30% fewer buffer loads,
60% fewer shared-memory loads, and the same time (600 us); it needs all 256 vector registers (four waves: 192).
So time here is not a function of the instruction count alone, and there is no model yet that predicts both the
clamp result and the one-wave result. The output stage is compute-bound too: 100 / 300 extra integer steps per
thread in the filter add 0.45 / 1.33 ms (`dz.sh`). NVIDIA's kernel runs dec0's first kernel in 203 us.
Shared memory throttles this GPU (2026-10-13, `build/rr/f/nv/cost/occ.py`, `occ2.py`). A dependent chain of f32
multiply-adds runs at full rate while the workgroup declares up to about 4 KiB of `shared` storage per wave; above
that the rate falls in proportion (one wave per workgroup: 6 KiB 1.24x slower, 8 KiB 1.55x, 13 KiB 3.1x, 62 KiB 12x;
four waves: 13 KiB full rate, 26 KiB 1.55x, 52 KiB 3.1x). The register count does not matter (0 to 224 live values:
the same rate). The decoders' first kernels held 21 to 46 KiB (LOW, EX, XE): `rr2_m.comp` now keeps the 4 x 4 low
tokens in KVS and the block's unquantised input in each lane's registers (`xr`), byte-identical on every stage
checked, and the network is 0.10 - 0.13 ms faster at 1280x720 -> 4K: exact build 5.68 / 6.25 ms (was 5.78 / 6.35),
`D4R_RR_RELAXED=3` 5.07 / 5.56 ms (was 5.17 / 5.72). It is also why the first one-wave prototype (13 KiB in one
wave) showed no gain: with 4 KiB (`exp/rr2_w.comp`: the low tokens borrow XT, the expansion borrows KVS, V^T is
stored, read into registers and overwritten by QK) it takes 534 us against the four-wave shader's 581, still
byte-identical to it. Extra arithmetic injected into either shader costs one instruction per cycle on about 128
wave slots (3.8e11 wave-instructions per second); NVIDIA's card measures 7.5e11 (`build/rr/f/nv/sass/rate.cpp`), and
its compiled kernel for this stage is 3,729 instructions per window (`sass/nv/dec0_0.sass`) against about 10,000
executed here. What an instruction costs here (`cost/cost.py`, `cost2.py`): an f32 multiply-add is the unit; a
32-bit integer multiply about seven; converting a pair to FP8 with the clamp about eight; a buffer read about 0.1 ps
per byte and lane (16 bytes: some forty multiply-adds), nearly twice that from a 1 MiB table.
`engine/rr/rr2_w.comp` (2026-10-13) is decoder 0's first kernel for `D4R_RR_RELAXED=3` with one wave per window:
every weight, key and value tile read once per window, the expansion plus skip computed once and kept in registers,
the scores' cross quadrant computed once and reused transposed (368 products for 384), FP8 values converted in
pairs, 4 KiB of shared memory. `compile_rr2.py` uses it for dec0_0 at level 3 (`ONE_WAVE`). The stage alone:
453 us against 581 us at 1280x720, arena byte-identical to `rr2_m.comp` at that level
(`build/rr/f/nv/invest-20261013/REPORT.md`: 20 variants, each change's instruction count and time). In the engine
every image of capF's three frames is identical to the four-wave level 3 model's, and 1280x720 -> 4K is about
5.00 / 5.55 ms (four-wave level 3: 5.10 / 5.61; network 3.05 / 3.31 against 3.15 / 3.39). The shader covers C = 32
only; the other stages still use `rr2_m.comp`.
`rr2_w.comp` also covers encoder 0's block now (KIND 4, embedded by `gen_rr2_enc0.py --one-wave` in NVIDIA's own
32-thread group, the translated guide tail on the same wave; `compile_rr2.py` `ONE_WAVE = {'dec0_0', 'enc0'}` at
`D4R_RR_RELAXED=3`). All 45 images of capF's three frames are identical to the four-wave level 3 model's.
1280x720 -> 4K, four interleaved runs each: four-wave 5.10 / 5.68 ms, dec0_0 on one wave 5.04 / 5.62, encoder 0 as
well 4.98 / 5.59. Decoder 0's second kernel is not converted: it hands r and the MLP sum (8 KiB of binary16) to its
translated tail through shared memory, which a single wave cannot hold without losing the 4 KiB rate.
### Presets D and E against a real NVIDIA card, and level 3 for them (2026-10-14)
NVIDIA's own kernels for presets D and E were replayed on an RTX PRO 4500 Blackwell (drivers 595.91 and 580.178: the
two give identical images; `build/rr/f/nv/pack.py --corpus ~/.cache/d4r-rr-corpus --first luma_convert`, captures
capD / capE3 at 640x360 -> 1280x720, every frame a reset frame). `D4R_RR_RELAXED=3` now also exists for
`build_native_model.sh`: level 2 plus `RR_NATIVE` (`rr_exact.glsl` `ACTQ8`: the MLP activation in f32, two elements
per saturating conversion, no activation table), and from level 2 on the modules are no longer passed through
`spv_exact_convert.py`. Presented image against the NVIDIA card, and the time at 1280x720 -> 4K (reset / next frame):
| build | D, dB | D, ms | E, dB | E, ms |
|---|---|---|---|---|
| exact (identical to the CUDA path through ZLUDA) | 38.5 / 39.0 / 38.5 | 4.4 / 4.7 | 40.0 / 39.9 / 40.0 | 4.34 / 4.54 |
| `D4R_RR_RELAXED=1` | 36.8 / 37.4 / 36.8 | 3.86 / 4.08 | 38.5 / 38.3 / 38.3 | 3.74 / 3.98 |
| `D4R_RR_RELAXED=2` | 36.8 / 37.4 / 36.8 | 3.76 / 3.99 | 38.5 / 38.3 / 38.3 | 3.65 / 3.82 |
| `D4R_RR_RELAXED=3` | 36.7 / 37.4 / 36.7 | 3.65 / 3.95 | 38.5 / 38.3 / 38.3 | 3.51 / 3.70 |
The NVIDIA card runs D at 1280x720 -> 4K in 2.19 / 2.39 ms. Unlike preset F, relaxing costs about 1.6 dB here. What
was measured and not built for D / E (`build/rr/f/nv/dstage.sh`, `dnl.sh`, level 3, 720p):
- shared memory: encoder 1 with 4 / 8 / 16 KiB of extra shared storage takes 287 -> ~305 / ~320 / ~343 us, so the
stages above the limit (decoder 1 at 21 KiB, decoder 0 at 17 KiB) would gain a few per cent each, under 1% of a
frame;
- one wave per window: with every weight tile replaced by one tile loaded once (wrong values, the ceiling of sharing
loads between a window's waves) decoder 0 takes the same time (803-818 us), encoder 1 the same, decoder 1 10% less.
These windows also have 96 key tokens (3 KiB each for the tokens, the keys and the values at C = 32), which does
not fit the 4 KiB a single wave can hold at full rate.
engine/rr/build_native_model.sh
+9 -2
View File
@@ -41,7 +41,13 @@ PRESET="$(awk '$1 == "preset" { print $2 }' "$OUT/rr.txt")"
# keeps the driver from folding a binary16 rounding away without the opaque OR the shaders otherwise need
DEFS=(-DRR_EXACT_CONVERT); [[ "$PRESET" == 4 ]] && DEFS+=(-DRR_PRESET_D)
# D4R_RR_RELAXED=1: one rounding where NVIDIA has two (engine/rr/rr_exact.glsl): faster, not bit-exact
[[ -n "${D4R_RR_RELAXED:-}" ]] && DEFS+=("-DRR_RELAXED=$D4R_RR_RELAXED")
# D4R_RR_RELAXED=3 adds RR_NATIVE: the MLP activation in f32 straight into FP8, no activation table
# Level 3 is the default (the fastest build; about 1.6 dB further from an NVIDIA card for D / E than the exact one);
# D4R_RR_RELAXED=0 builds NVIDIA's roundings exactly (bit-identical to the CUDA path through ZLUDA).
D4R_RR_RELAXED="${D4R_RR_RELAXED:-3}"
if (( D4R_RR_RELAXED > 0 )); then
DEFS+=("-DRR_RELAXED=$(( D4R_RR_RELAXED > 2 ? 2 : D4R_RR_RELAXED ))"); (( D4R_RR_RELAXED >= 3 )) && DEFS+=(-DRR_NATIVE)
fi
# D4R_RR_MIX=1: the 32-deep steps round once through v_fma_mix (the shaders' K32): ~5% faster, not provably bit-exact
[[ "${D4R_RR_MIX:-0}" == 1 ]] && DEFS+=(-DRR_MIX)
glsl() { # glsl OUT.spv SOURCE INCLUDE_DIR...
@@ -49,7 +55,8 @@ glsl() { # glsl OUT.spv SOURCE INCLUDE_DIR...
local dir inc=(); for dir in "$@"; do inc+=("-I$dir"); done
"${CAPPED[@]}" glslangValidator --target-env vulkan1.3 -V "${DEFS[@]}" "${inc[@]}" -o "$out" "$src" > "$WORK/glsl.log" 2>&1 ||
{ grep -E 'ERROR' "$WORK/glsl.log" | head -5; exit 1; }
python3 "$RR/spv_exact_convert.py" "$out" > /dev/null
# from level 2 on the conversions are left to the driver (as compile_rr2.py does for preset F)
(( D4R_RR_RELAXED >= 2 )) || python3 "$RR/spv_exact_convert.py" "$out" > /dev/null
}
native() { # native STAGE WEIGHT_FILE: point rr.txt's stage line at STAGE_m.spv
sed -i -E "s/^stage $1 $1\.spv 32 weights-$1\.bin (.*)/stage $1 $1_m.spv 32 $2 \1/; /^stage $1 /s/ noopt\$//" "$OUT/rr.txt"
engine/rr/dec0_m.comp
+1 -1
View File
@@ -269,7 +269,7 @@ void main()
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i]));
ACTQ8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
engine/rr/enc0_m_tail.glsl
+1 -1
View File
@@ -151,7 +151,7 @@
A8 a0, a1;
LDA(a0, w1, 32u, c, 0u); LDA(a1, w1, 32u, c, 1u);
K32(h, a0, Nb[0], a1, Nb[1])
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i]));
ACTQ8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < 2u; ++m) {
A8 a0, a1;
engine/rr/enc1_ref.py
+15
View File
@@ -111,8 +111,23 @@ def softmax_cubic(x):
return (e * r[:, None]).astype(H)
# D4R_RR_FUSE=1: the arithmetic as an NVIDIA card runs it. NVIDIA's compiler contracts a binary16 multiply whose
# product only feeds an add / sub into a fused multiply-add (the kernels give those no rounding modifier); ZLUDA, and
# so the captures, round the product first. build/rr/f/nv has the measurements (ptx_contract.py is the rule).
FUSE = os.environ.get('D4R_RR_FUSE') == '1'
def fma16(a, b, c):
"""binary16 a * b + c with one rounding (the exact value fits a double)"""
with np.errstate(all='ignore'):
return (np.asarray(a, np.float64) * np.asarray(b, np.float64) + np.asarray(c, np.float64)).astype(H)
def act(x):
y = np.clip(x, H(-2), H(2)).astype(H)
if FUSE:
g = fma16(y, fma16(-np.abs(y), H(0.0810546875), H(0.412109375)), H(0.5))
return (x * g).astype(H)
g = add(H(0.5), (y * (H(0.412109375) - (H(0.0810546875) * np.abs(y)).astype(H)).astype(H)).astype(H))
return (x * g).astype(H)
engine/rr/gen_rr2_enc0.py
+11 -3
View File
@@ -167,8 +167,12 @@ def main():
g = Path(src).read_text().split('\n')
m = next(i for i, l in enumerate(g) if l.startswith('void main()'))
head = '\n'.join(g[:m]).replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;',
'layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;')
# --one-wave (D4R_RR_RELAXED=3): the block as rr2_w.comp, in NVIDIA's own 32-thread group, the tail on the same wave
one_wave = '--one-wave' in sys.argv
head = '\n'.join(g[:m])
if not one_wave:
head = head.replace('layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;',
'layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in;')
if dec: # the tail's tanh from rr2_pack.py's table (ptx2rr.py's RR_TANH_TABLE; its offset from OUT's rr2_stage.h)
v = head.index('\n') + 1
table = re.search(r'#define W_TANH (\d+)u', (Path(out).parent / 'rr2_stage.h').read_text()).group(1)
@@ -210,6 +214,10 @@ def main():
for mo in re.finditer(r'add\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail):
if mo.group(2) in prod:
rmap[var(mo.group(3))] = prod[mo.group(2)]
# the same step in a module rewritten by build/rr/f/nv/ptx_contract.py (D4R_RR_FUSE builds): one fused instruction
for mo in re.finditer(r'fma\.rn\.f16\s+(%rs\d+),\s*(%rs\d+),\s*(%rs\d+),\s*(%rs\d+);', text_tail):
if mo.group(2) in half:
rmap[var(mo.group(4))] = half[mo.group(2)]
if len(rmap) != 2 * 2 * len(last):
raise SystemExit(f'found {len(rmap)} residual registers, expected {4 * len(last)}')
restored |= set(rmap)
@@ -258,7 +266,7 @@ def main():
load.append(f' {name} = ' + (f'uint16BitsToFloat16(uint16_t({v}));' if name in f16vars else f'{v};'))
load.append(' }')
# the native block's source, without its #version line (this shader has its own)
block = (Path(__file__).resolve().parent / 'rr2_m.comp').read_text().split('\n', 1)[1]
block = (Path(__file__).resolve().parent / ('rr2_w.comp' if one_wave else 'rr2_m.comp')).read_text().split('\n', 1)[1]
if split == 'block': # the second pass: the block on the gather pass's record, then the tail on wave 0
recomputed, saved, record, where = split_layout(xf_word, pre, tail, f16vars)
restore = [' ' + (f'{v} = uint16BitsToFloat16(uint16_t(arena.w[rr2_rec + 1024u + {len(saved)}u * l_ + {n}u]));' if v in f16vars
engine/rr/ptx2rr.py
+8
View File
@@ -420,6 +420,7 @@ class RRT(base.T):
def __init__(self, ptx, kernel, block, weights_arg=40, buffers=()):
self.weights_arg = weights_arg
self.buffers = sorted(buffers) # argument offsets of further buffers, each its own binding
self.snorm16 = set() # texture arguments of --snorm16
bx, by, bz = block
if bx * by * bz % 32:
raise base.Err('block size must be a multiple of the 32-lane subgroup')
@@ -621,6 +622,10 @@ class RRT(base.T):
raise base.Err(f'unsupported texture op {op}')
f16 = 'f16' in o
self.E(f'{{ vec4 t_ = {fetch};')
if off in self.snorm16 and o[-1] == 'f32':
# --snorm16: a filtered fetch of an 8-bit SNORM texture on an NVIDIA card is this card's value (8-bit weights)
# rounded to a 16-bit step (measured on enc0's reprojected normals, build/rr/f/nv; D4R_RR_FUSE models)
self.E('{ precise vec4 q_ = roundEven(t_ * 32767.0); precise vec4 d_ = q_ / 32767.0; t_ = d_; }')
for i, d in enumerate(dst):
# a binary16 result is the texel rounded to binary16: through H(HU(...) | SUST_Z) so that the driver cannot
# fold the rounding away when the value is widened again (RR_RUNTIME's note on sust_b)
@@ -837,6 +842,8 @@ def main():
ap.add_argument('--surface', action='append', default=[], metavar='OFFSET=FORMAT',
help="the GLSL image format of a surface whose stores are formatted (sust.p), when it is not the "
"store's own guess, e.g. 144=rgba8_snorm (repeatable)")
ap.add_argument('--snorm16', action='append', type=int, default=[], metavar='OFFSET',
help='texture argument whose filtered fetches are rounded to 16-bit steps, as an NVIDIA card returns them')
ap.add_argument('--weights-arg', type=int, default=40,
help='argument offset of the weight pointer (40 for the DLSS-D kernels, 24 for rr2_*)')
ap.add_argument('--buffer', action='append', type=int, default=[], metavar='OFFSET',
@@ -858,6 +865,7 @@ def main():
def make():
t = RRT(src, ns.kernel, block, ns.weights_arg, ns.buffer)
t.snorm16 = set(ns.snorm16)
t.probes = probes
t.probe_block = tuple(int(x) for x in ns.probe_block.split(','))
t.oob = ns.oob_report
engine/rr/rr2_m.comp
+82 -34
View File
@@ -69,6 +69,35 @@ layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define RND(m) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (m)[i_] = rnd16((m)[i_]);
#define Q8(dst, src) Q8PAIRS(dst, src)
// One 32-deep step onto an accumulator of binary16 values: the reference adds it inside the f32 products and rounds
// RR_FUSE (D4R_RR_FUSE=1): out = in + branch * s and the guide's blend as one fused multiply-add each, which is what
// NVIDIA's compiler makes of those multiply / add pairs on an NVIDIA card; without it the product is rounded first,
// as ZLUDA (the captures) does. The activation table is built to match (enc1_ref.act).
#ifdef RR_FUSE
#define MADD(d, a, b, c) precise float16_t d = fma((a), (b), (c));
#define MADD4(d, a, b, c) precise f16vec4 d = fma((a), (b), (c));
#else
#define MADD(d, a, b, c) precise float16_t d##_p = (a) * (b); precise float16_t d = (c) + d##_p;
#define MADD4(d, a, b, c) precise f16vec4 d##_p = (a) * (b); precise f16vec4 d = d##_p + (c);
#endif
// RR_NATIVE (D4R_RR_RELAXED=3): the residual steps and the MLP activation in f32, straight into FP8 - no binary16 in
// between and no activation table. Not the reference's roundings (about 40 dB from it, like RR_RELAXED).
// (The biases stay binary16: f32 copies were tried for RR_NATIVE and were slower - half the conversions, but one more
// load per accumulator, and a load costs about three arithmetic instructions here.)
#define BV(X, i) RV(w16.w[W_##X / 2u + (i)])
#define BVS(X, stride, n, i) RV(w16.w[(W_##X + (stride) * (n)) / 2u + (i)])
#ifdef RR_NATIVE
#define RVT float
#define E4V(w, b) float(uintBitsToFloate4m3EXT(uint8_t(((w) >> (8u * (b))) & 0xffu)))
#define RESID(d, a, s, c) const float d = fma(float(a), float(s), float(c));
#define Q8R(x) q8s(x)
#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2(float((h)[i_]), float((h)[i_ + 1])); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
#else
#define RVT float16_t
#define E4V(w, b) e4((w), (b))
#define RESID(d, a, s, c) MADD(d, float16_t(a), s, c)
#define Q8R(x) q8h(x)
#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = ACTQ(float16_t((h)[i_]));
#endif
// once (swin_m.comp explains RR_MIX, the one-rounding variant that is not bit-exact by construction)
#ifdef RR_MIX
#define RAC HAC
@@ -82,6 +111,7 @@ layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define K32H(acc, a0, b0, a1, b1) { ACC p_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(acc))); [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (acc)[i_] = float16_t(p_[i_]); }
#endif
const uint HEAD = W_WQK1 - W_WQK0, BRANCH = W_WD1 - W_WD0;
#define BRANCHB BRANCH
shared uint XT[16 * C]; // the window's tokens: [token slot][C FP8 channels in slot order]; later the merge rows
shared uint KVS[2 * 512]; // the current head's QK ([token slot][32]), then V^T ([32][token slot])
@@ -89,9 +119,15 @@ shared uint KVS[2 * 512]; // the current head's QK ([token slot][32]), th
shared uint GT[64 * 32]; // the merge input: [token slot][4 sub-positions x 32 FP8 channels]
#endif
#if KIND == 2
shared uint LOW[4 * CL]; // the 4 x 4 tokens one level down: [row-major][CL FP8 channels]
// Shared memory is what limits how many waves this GPU keeps running: above about 4 KiB per wave (16 KiB for this
// workgroup) the rate falls in proportion (measured: build/rr/f/nv/cost/occ2.py). So the 4 x 4 tokens one level down
// ([row-major][CL FP8 channels]) borrow KVS, which the attention only writes later, and the block's input before its
// FP8 conversion stays in each lane's registers (xr), not in a [token slot][channel slot] array.
#if 4 * CL > 2 * 512
#error rr2_m.comp: LOW no longer fits in KVS
#endif
#define LOW KVS
shared uint EX[16 * C]; // their expansion: [low token][sub-position][C FP8 channels in slot order]
shared float16_t XE[64 * C]; // the block's input before its FP8 conversion: [token slot][channel slot]
#endif
#if KIND == 4
#ifdef RR2_XF_ARENA
@@ -107,6 +143,7 @@ shared uint XF[RR2_GROUP * 64 * 16]; // per window, the features from the gath
#endif
shared float16_t HD[RR2_GROUP * 64 * 4]; // per window, the guide head per token (raster order), for the translated tail
const uint EBRANCH = W_EWD1 - W_EWD0;
#define EBRANCHB EBRANCH
#endif
#if KIND == 5
shared float16_t RSH[RR2_GROUP * 64 * 32], MSH[RR2_GROUP * 64 * 32]; // per window, r and the MLP sum per token (raster order) and natural channel
@@ -165,18 +202,18 @@ void main()
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(fb[kt], XF_SRC, XF_AT + wv * 256u + 4u * kt, 16u, CMJ);
RAC X[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) X[m][i] = RV(w16.w[W_EBU / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (int i = 0; i < 8; ++i) X[m][i] = BV(EBU, 16u * m + 8u * hl + uint(i));
[[unroll]] for (uint br = 0u; br < 2u; ++br) {
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[(W_EBD0 + EBRANCH * br) / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = BVS(EBD0, EBRANCHB, br, 16u * c + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < 4u; kt += 2u) {
A8 a0, a1;
LDA(a0, W_EWD0 + EBRANCH * br, 64u, c, kt); LDA(a1, W_EWD0 + EBRANCH * br, 64u, c, kt + 1u);
K32(h, a0, fb[kt], a1, fb[kt + 1u])
}
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i]));
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
@@ -207,7 +244,7 @@ void main()
[[unroll]] for (uint kt = 0u; kt < 8u; ++kt) coopMatLoad(gb[kt], GT, wv * 512u + 4u * kt, 32u, CMJ);
[[unroll]] for (uint m = 0u; m < MT; ++m) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_B1 / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = BV(B1, 16u * m + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < 8u; kt += 2u) {
A8 w0, w1;
LDA(w0, W_W1, 128u, m, kt); LDA(w1, W_W1, 128u, m, kt + 1u);
@@ -221,6 +258,7 @@ void main()
barrier();
#else
// a decoder's first kernel: wave w expands the 16 low tokens to sub-position w, then x = dq(e) + dq(skip) * s3
float16_t xr[MT][8]; // x of this lane's token (slot 16 wv + col), channel slots 16 m + 8 hl + i
{
const int X0 = x0 >> 1, Y0 = y0 >> 1;
for (uint i = tid; i < 4u * CL; i += 128u) {
@@ -234,7 +272,7 @@ void main()
for (uint mm = 0u; mm < MT; ++mm) {
const uint m = MT * wv + mm;
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = BV(BE, 16u * m + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < CL / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CL, m, kt); LDA(w1, W_WE, CL, m, kt + 1u);
@@ -255,14 +293,27 @@ void main()
const uint e = EX[r * C + s * (C / 4u) + wd];
floate4m3_t c[4];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
precise float16_t ms = e4(sk, b) * s3;
precise float16_t x = e4(e, b) + ms;
XE[tid * C + 4u * wd + b] = x;
MADD(x, e4(sk, b), s3, e4(e, b))
c[b] = q8h(x);
}
XT[tid * (C / 4u) + wd] = pack4(c[0], c[1], c[2], c[3]);
}
}
{ // the same x for the residual step: the bytes of the words 4 m + 2 hl and 4 m + 2 hl + 1 of this lane's token
const int kx = x0 + int(t & 7u), ky = y0 + int(t >> 3u);
const uint tb = (skOff + uint(mirror(ky, ey) * ex + mirror(kx, ex)) * 32u) / 4u;
const uint r = uint(((ky >> 1) - Y0) * 4 + ((kx >> 1) - X0)), s = uint(2 * (ky & 1) + (kx & 1));
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (uint h2 = 0u; h2 < 2u; ++h2) {
const uint wd = 4u * m + 2u * hl + h2;
const uint sk = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
const uint e = EX[r * C + s * (C / 4u) + wd];
[[unroll]] for (uint b = 0u; b < 4u; ++b) {
MADD(x, e4(sk, b), s3, e4(e, b))
xr[m][4u * h2 + b] = x;
}
}
}
barrier();
}
#endif
@@ -270,7 +321,7 @@ void main()
// ---------------------------------------------------------------- attention: the heads' out-projections sum up
RAC AT[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) AT[m][i] = RV(w16.w[W_BO / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (int i = 0; i < 8; ++i) AT[m][i] = BV(BO, 16u * m + 8u * hl + uint(i));
[[dont_unroll]] for (uint hd = 0u; hd < NH; ++hd) {
barrier(); // every wave is done with the previous head's QK / V
for (uint j = 0u; j < 2u; ++j) { // 8 jobs (token tile n, channel tile c)
@@ -302,11 +353,11 @@ void main()
{
float16_t e[4][8];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
A8 k0, k1;
coopMatLoad(k0, KVS, kt * 128u, 8u, RM);
coopMatLoad(k1, KVS, kt * 128u + 4u, 8u, RM);
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
K32H(c, k0, Qb[0], k1, Qb[1])
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(c[i])]));
}
@@ -355,7 +406,7 @@ void main()
}
// ---------------------------------------------------------------- r = x + attention * s1; the MLP's input q8(r)
float16_t rv[MT][8];
RVT rv[MT][8];
B8 Nb[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m) {
#if KIND != 2
@@ -364,31 +415,30 @@ void main()
#endif
[[unroll]] for (int i = 0; i < 8; ++i) {
#if KIND == 2
const float16_t xb = XE[(16u * wv + col) * C + 16u * m + 8u * hl + uint(i)];
const float16_t xb = xr[m][i];
#else
const float16_t xb = e4(xw[i >> 2], uint(i) & 3u);
const RVT xb = E4V(xw[i >> 2], uint(i) & 3u);
#endif
precise float16_t sa = float16_t(AT[m][i]) * s1;
precise float16_t r = xb + sa;
RESID(r, AT[m][i], s1, xb)
rv[m][i] = r;
Nb[m][i] = q8h(r);
Nb[m][i] = Q8R(r);
}
}
RAC ML[MT];
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) ML[m][i] = RV(w16.w[W_BU / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (int i = 0; i < 8; ++i) ML[m][i] = BV(BU, 16u * m + 8u * hl + uint(i));
[[dont_unroll]] for (uint br = 0u; br < NB; ++br) {
const uint w1 = W_WD0 + BRANCH * br, b1 = W_BD0 + BRANCH * br, w2 = W_WU0 + BRANCH * br;
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
RAC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = RV(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = BVS(BD0, BRANCHB, br, 16u * c + 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u);
K32(h, a0, Nb[kt], a1, Nb[kt + 1u])
}
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i]));
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
@@ -402,7 +452,7 @@ void main()
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) {
const uint at = RR2_WIN * 2048u + t * 32u + inter(16u * m + 8u * hl + uint(i));
RSH[at] = rv[m][i];
RSH[at] = float16_t(rv[m][i]);
MSH[at] = float16_t(ML[m][i]);
}
#else
@@ -418,9 +468,8 @@ void main()
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
[[unroll]] for (int i = 0; i < 8; ++i) {
precise float16_t sm = float16_t(ML[m][i]) * s2;
precise float16_t y = rv[m][i] + sm;
Tb[m][i] = q8h(y);
RESID(y, ML[m][i], s2, rv[m][i])
Tb[m][i] = Q8R(y);
}
const uint w0 = pack4(Tb[m][0], Tb[m][1], Tb[m][2], Tb[m][3]), w1 = pack4(Tb[m][4], Tb[m][5], Tb[m][6], Tb[m][7]);
const uint at = 4u * (m & 1u) + 2u * hl;
@@ -435,7 +484,7 @@ void main()
#if GUIDE || KIND == 3 || KIND == 4
{ // the guide head (rows 0-3 on lanes 0-15)
RAC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = RV(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = BV(BH, 8u * hl + uint(i));
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 a0, a1;
LDA(a0, W_WH, C, 0u, kt); LDA(a1, W_WH, C, 0u, kt + 1u);
@@ -457,16 +506,15 @@ void main()
precise float16_t ux = float16_t(1.0) - wx, uy = float16_t(1.0) - wy;
precise f16vec4 p00 = guideAt(gInOff, bx, by, hx) * ux, p10 = guideAt(gInOff, x1, by, hx) * wx;
precise f16vec4 p01 = guideAt(gInOff, bx, y1, hx) * ux, p11 = guideAt(gInOff, x1, y1, hx) * wx;
precise f16vec4 a00 = p00 * uy, a10 = p10 * uy, a01 = p01 * wy, a11 = p11 * wy;
precise f16vec4 u0 = a00 + a10;
precise f16vec4 u1 = u0 + a01;
precise f16vec4 up = u1 + a11;
precise f16vec4 a10 = p10 * uy;
MADD4(u0, p00, f16vec4(uy), a10)
MADD4(u1, p01, f16vec4(wy), u0)
MADD4(up, p11, f16vec4(wy), u1)
const float16_t al = sig(float16_t(hv[3]));
precise f16vec4 gt = up * f16vec4(gate, float16_t(0.0));
precise float16_t ia = float16_t(1.0) - al;
precise f16vec4 pe = guideAt(gEncOff, tx, ty, ex) * ia;
precise f16vec4 pg = gt * al;
precise f16vec4 mx = pe + pg;
MADD4(mx, guideAt(gEncOff, tx, ty, ex), f16vec4(ia), pg)
o4 = min(mx, f16vec4(65504.0));
#else
precise f16vec4 gg = guideAt(gInOff, tx, ty, ex) * f16vec4(gate, float16_t(0.0));
@@ -489,7 +537,7 @@ void main()
const uint o = (mrgOff + uint(Y * hx + X) * 32u) >> 2u;
for (uint m = wv; m < CM / 16u; m += 4u) {
RAC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = RV(w16.w[W_BM / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = BV(BM, 16u * m + 8u * hl + uint(i));
for (uint ks = 0u; ks < C / 8u; ++ks) {
B8 g0, g1; A8 w0, w1;
coopMatLoad(g0, XT, 8u * ks, C, CMJ);
engine/rr/rr2_w.comp
+370
View File
@@ -0,0 +1,370 @@
#version 460
// Decoder 0's first kernel of preset F (rr2_m.comp KIND 2 with C = 32) for D4R_RR_RELAXED=3, with ONE wave per
// 8x8-token window where rr2_m.comp uses four. The wave walks the four 16-token slices itself, so every weight, key
// and value tile is read once per window; the expansion plus skip is computed once and kept in registers; the scores'
// cross quadrant is computed once and reused transposed (the scores are symmetric); FP8 values are converted in pairs.
// Shared memory stays at 4 KiB, above which this GPU runs fewer waves (engine/rr/README.md). The arithmetic is
// rr2_m.comp's with -DRR_RELAXED=2 -DRR_NATIVE and the arena comes out byte-identical to that build
// (build/rr/f/nv/invest-20261013/REPORT.md: 453 us against 581 us at 1280x720). compile_rr2.py uses it for dec0_0.
// KIND 4 is encoder 0's block the same way (the embedding of the gather pass's record, the tokens, the guide head to
// HD), included with RR2_EMBED and RR2_XF_ARENA by the shader engine/rr/gen_rr2_enc0.py --one-wave generates.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_8bit_storage : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#extension GL_EXT_control_flow_attributes : require
#extension GL_GOOGLE_include_directive : require
#ifndef RR2_EMBED
layout(local_size_x = 32) in;
#endif
#include "rr2_stage.h"
#include "rr_exact.glsl"
#if !(KIND == 2 || KIND == 4) || CM != 0 || GUIDE != 0 || C != 32
#error rr2_w.comp: a 32-channel decoder's first kernel, or encoder 0's block
#endif
#if KIND == 4 && !(defined(RR2_EMBED) && defined(RR2_XF_ARENA))
#error rr2_w.comp: encoder 0's block reads the gather pass's record (gen_rr2_enc0.py --split block --one-wave)
#endif
#define MT (C / 16)
#ifdef RR2_EMBED
#define PRM(i) param.w[i]
#define RR2_WGX gl_WorkGroupID.x
#else
layout(set = 0, binding = 0, std430) readonly buffer ParamBlock { uint p[]; } param;
#define PRM(i) param.p[i]
#define RR2_WGX gl_WorkGroupID.x
#endif
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 1, std430) readonly buffer W16 { float16_t w[]; } w16;
layout(set = 0, binding = 1, std430) readonly buffer W16U { uint16_t w[]; } w16u;
layout(set = 0, binding = 2, std430) buffer Arena { uint w[]; } arena;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HAC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define A8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B8 coopmat<floate4m3_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CMJ gl_CooperativeMatrixLayoutColumnMajor
#define LDA(dst, base, k, mt, kt) coopMatLoad((dst), w8.w, (base) + 256u * ((mt) * ((k) / 16u) + (kt)), 16u, RM)
#define LDB(dst, base, n, kt, nt) coopMatLoad((dst), w8.w, (base) + 256u * ((kt) * ((n) / 16u) + (nt)), 16u, RM)
#define K32(acc, a0, b0, a1, b1) { (acc) = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), (acc))); }
#define Q8(dst, src) Q8PAIRS(dst, src)
#define ACT8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2((h)[i_], (h)[i_ + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
const uint HEAD = W_WQK1 - W_WQK0, BRANCH = W_WD1 - W_WD0;
shared uint XT[16 * C];
shared uint KVS[512]; // the expansion (input phase), then per head V^T, then QK
// Shared memory is what limits how many waves this GPU keeps running: above about 4 KiB per wave the rate falls in
// proportion (build/rr/f/nv/cost/occ2.py). So the input phase borrows KVS, which the attention only needs later
// (LOW at word 0, EX at word 512), and the block's unquantised input stays in registers (xr), not in memory.
// (4 KiB in all: the low tokens borrow XT, which is only written after they are in registers; the expansion borrows
// KVS; V^T is stored, read into registers, and then overwritten by QK.)
#if KIND == 2 && (4 * CL > 16 * C || 16 * C > 512)
#error rr2_w.comp: the input phase no longer fits
#endif
#define EX_AT 0u
#if KIND == 4
uint rr2_xf; // the window's record from the gather pass: the index of its first arena word
shared float16_t HD[64 * 4]; // the guide head per token (raster order), for the translated tail
const uint EBRANCH = W_EWD1 - W_EWD0;
#endif
uint pack4(floate4m3_t a, floate4m3_t b, floate4m3_t c, floate4m3_t d)
{
return code(a) | (code(b) << 8u) | (code(c) << 16u) | (code(d) << 24u);
}
f16vec2 e4pair(uint w, uint h) { return f16vec2(uintBitsToFloate4m3EXT(u8vec2(uint8_t(w >> (16u*h)), uint8_t(w >> (16u*h+8u))))); }
float16_t e4(uint w, uint b) { return float16_t(uintBitsToFloate4m3EXT(uint8_t((w >> (8u * b)) & 0xffu))); }
#ifdef RR2_EMBED
void rr2_block()
#else
void main()
#endif
{
const uint lane = gl_SubgroupInvocationID;
const uint hl = lane >> 4u, col = lane & 15u;
RR_ZERO16 = uint16_t(gl_NumWorkGroups.z - 1u);
const int ex = int(PRM(0)), ey = int(PRM(1)), hx = ex / 2, hy = ey / 2;
const uint inOff = PRM(2), skOff = PRM(4), outOff = PRM(8), shift = PRM(20);
const float16_t s1 = hb(PRM(23) >> 16u), s2 = hb(PRM(24) & 0xffffu), s3 = hb(PRM(24) >> 16u);
const int x0 = 8 * int(RR2_WGX) - int(shift & 0xffffu), y0 = 8 * int(gl_WorkGroupID.y) - int(shift >> 16u);
const uint plane = uint(ex * ey) * 32u, lplane = uint(hx * hy) * 32u;
#if KIND == 4
// ---------------------------------------------------------------- the block's input: x = bu + sum over two
// branches of Wu . act(Wd . features + bd); the features are the gather pass's record, read once per slice
B8 qx[4][MT]; // x as FP8, this lane's part per token slice (the residual step decodes it again)
{
B8 fb[4][4];
ACC X[4][MT];
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) coopMatLoad(fb[wv][kt], arena.w, rr2_xf + wv * 256u + 4u * kt, 16u, CMJ);
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) X[wv][m][i] = float(w16.w[W_EBU / 2u + 16u * m + 8u * hl + uint(i)]);
}
[[unroll]] for (uint br = 0u; br < 2u; ++br) {
A8 ewd[2][4], ewu0[MT], ewu1[MT];
[[unroll]] for (uint c = 0u; c < 2u; ++c)
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) LDA(ewd[c][kt], W_EWD0 + EBRANCH * br, 64u, c, kt);
[[unroll]] for (uint m = 0u; m < MT; ++m) {
LDA(ewu0[m], W_EWU0 + EBRANCH * br, 32u, m, 0u); LDA(ewu1[m], W_EWU0 + EBRANCH * br, 32u, m, 1u);
}
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[(W_EBD0 + EBRANCH * br) / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < 4u; kt += 2u) K32(h, ewd[c][kt], fb[wv][kt], ewd[c][kt + 1u], fb[wv][kt + 1u])
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) K32(X[wv][m], ewu0[m], Hb[0], ewu1[m], Hb[1])
}
}
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv)
[[unroll]] for (uint m = 0u; m < MT; ++m) {
Q8(qx[wv][m], X[wv][m])
coopMatStore(qx[wv][m], XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
#else
// ---------------------------------------------------------------- the block's input: expansion plus skip
float16_t xr[4][MT][8]; // this lane's part of x for each token slice: [slice][channel tile][element]
{
const int X0 = x0 >> 1, Y0 = y0 >> 1;
for (uint i = lane; i < 4u * CL; i += 32u) {
const uint r = i / (CL / 4u), wd = i % (CL / 4u);
const int X = mirror(X0 + int(r & 3u), hx), Y = mirror(Y0 + int(r >> 2u), hy);
XT[i] = arena.w[(inOff + (wd >> 3u) * lplane + uint(Y * hx + X) * 32u) / 4u + (wd & 7u)];
}
B8 lb[CL / 16];
[[unroll]] for (uint kt = 0u; kt < CL / 16; ++kt) coopMatLoad(lb[kt], XT, 4u * kt, CL / 4u, CMJ);
for (uint m = 0u; m < 4u * MT; ++m) {
ACC a;
[[unroll]] for (int i = 0; i < 8; ++i) a[i] = float(w16.w[W_BE / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < CL / 16; kt += 2u) {
A8 w0, w1;
LDA(w0, W_WE, CL, m, kt); LDA(w1, W_WE, CL, m, kt + 1u);
K32(a, w0, lb[kt], w1, lb[kt + 1u])
}
B8 q;
Q8(q, a)
coopMatStore(q, KVS, EX_AT + 4u * m, C, CMJ);
}
// the same x for the residual step, in this lane's own layout (token slot 16 wv + col, channel slots
// 16 m + 8 hl + i = bytes of the words 4 m + 2 hl and 4 m + 2 hl + 1)
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
const uint k = inter(16u * wv + col);
const int kx = x0 + int(k & 7u), ky = y0 + int(k >> 3u);
const uint tb = (skOff + uint(mirror(ky, ey) * ex + mirror(kx, ex)) * 32u) / 4u;
const uint r = uint(((ky >> 1) - Y0) * 4 + ((kx >> 1) - X0)), s = uint(2 * (ky & 1) + (kx & 1));
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (uint h2 = 0u; h2 < 2u; ++h2) {
const uint wd = 4u * m + 2u * hl + h2;
const uint sk = arena.w[tb + (wd >> 3u) * (plane >> 2u) + (wd & 7u)];
const uint e = KVS[EX_AT + r * C + s * (C / 4u) + wd];
[[unroll]] for (uint pair = 0u; pair < 2u; ++pair) {
precise f16vec2 x_p = e4pair(sk, pair) * f16vec2(s3);
precise f16vec2 x = e4pair(e, pair) + x_p;
xr[wv][m][4u * h2 + 2u*pair] = x.x;
xr[wv][m][4u * h2 + 2u*pair+1u] = x.y;
}
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
B8 q;
Q8(q, xr[wv][m])
coopMatStore(q, XT, wv * 4u * C + 4u * m, C / 4u, CMJ);
}
}
}
#endif
// ---------------------------------------------------------------- attention
ACC AT[4][MT];
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv)
[[unroll]] for (uint m = 0u; m < MT; ++m)
[[unroll]] for (int i = 0; i < 8; ++i) AT[wv][m][i] = float(w16.w[W_BO / 2u + 16u * m + 8u * hl + uint(i)]);
[[unroll]] for (uint hd = 0u; hd < NH; ++hd) {
// V^T of every token tile first: stored, read back as the product's A operands, then QK takes its place
A8 kt0[4], kt1[4], v0[2][2], v1[2][2], wo0[MT], wo1[MT];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
B8 wvt[MT];
[[unroll]] for (uint kt = 0u; kt < MT; ++kt) LDB(wvt[kt], W_WV0 + HEAD * hd, 32u, kt, c);
[[unroll]] for (uint n = 0u; n < 4u; ++n) {
ACC v = ACC(0.0);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
A8 xa0, xa1;
coopMatLoad(xa0, XT, n * 4u * C + 4u * kt, C / 4u, RM);
coopMatLoad(xa1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, RM);
K32(v, xa0, wvt[kt], xa1, wvt[kt + 1u])
}
A8 va;
Q8(va, v)
coopMatStore(va, KVS, c * 256u + 4u * n, 16u, RM);
}
}
[[unroll]] for (uint c = 0u; c < 2u; ++c)
[[unroll]] for (uint ks = 0u; ks < 2u; ++ks) {
coopMatLoad(v0[c][ks], KVS, c * 256u + 8u * ks, 16u, RM);
coopMatLoad(v1[c][ks], KVS, c * 256u + 8u * ks + 4u, 16u, RM);
}
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
A8 wk[MT];
[[unroll]] for (uint kt = 0u; kt < MT; ++kt) LDA(wk[kt], W_WQK0 + HEAD * hd, C, c, kt);
[[unroll]] for (uint n = 0u; n < 4u; ++n) {
ACC a = ACC(0.0);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) {
B8 xb0, xb1;
coopMatLoad(xb0, XT, n * 4u * C + 4u * kt, C / 4u, CMJ);
coopMatLoad(xb1, XT, n * 4u * C + 4u * kt + 4u, C / 4u, CMJ);
K32(a, wk[kt], xb0, wk[kt + 1u], xb1)
}
A8 ka;
Q8(ka, a)
coopMatStore(ka, KVS, n * 128u + 4u * c, 8u, RM);
}
}
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
coopMatLoad(kt0[kt], KVS, kt * 128u, 8u, RM);
coopMatLoad(kt1[kt], KVS, kt * 128u + 4u, 8u, RM);
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
LDA(wo0[m], W_WO0 + HEAD * hd, 32u, m, 0u); LDA(wo1[m], W_WO0 + HEAD * hd, 32u, m, 1u);
}
// Preserve all query operands before KVS becomes a 32x32 f16 numerator cache.
B8 allQ[4][2];
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv)
[[unroll]] for (uint c = 0u; c < 2u; ++c) coopMatLoad(allQ[wv][c], KVS, wv * 128u + 4u * c, 8u, CMJ);
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
if (wv == 2u) memoryBarrierShared();
B8 Pb[4];
{
float16_t e[4][8];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt) {
if (wv >= 2u && kt < 2u) {
HAC cached;
// Transpose keys/query roles. Stride is measured in uint words.
coopMatLoad(cached, KVS, (wv - 2u) * 256u + kt * 8u, 16u, CMJ);
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = cached[i];
} else {
HAC c;
coopMatLoad(c, w16.w, (W_BIAS0 + HEAD * hd) / 2u + kt * 1024u + 16u * wv, 64u, RM);
ACC p_ = coopMatMulAdd(kt1[kt], allQ[wv][1], coopMatMulAdd(kt0[kt], allQ[wv][0], ACC(c)));
[[unroll]] for (int i = 0; i < 8; ++i) e[kt][i] = hb(uint(w16u.w[W_CUBIC / 2u + bh(float16_t(p_[i]))]));
if (wv < 2u && kt >= 2u) {
HAC cached;
[[unroll]] for (int i = 0; i < 8; ++i) cached[i] = e[kt][i];
coopMatStore(cached, KVS, (kt - 2u) * 256u + wv * 8u, 16u, RM);
}
}
}
f16vec2 ep[4][4];
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) ep[kt][q] = f16vec2(e[kt][2u * q], e[kt][2u * q + 1u]);
precise f16vec2 st[2];
[[unroll]] for (uint tq = 0u; tq < 2u; ++tq) {
const uint q0 = 2u * tq, q1 = q0 + 1u;
precise f16vec2 ev = (ep[0][q0] + ep[0][q1]) + (ep[2][q0] + ep[2][q1]);
precise f16vec2 od = (ep[1][q0] + ep[1][q1]) + (ep[3][q0] + ep[3][q1]);
st[tq] = ev + od;
}
precise f16vec2 loc = st[0] + st[1];
const uint other = subgroupShuffleXor(packFloat2x16(loc), 16u);
precise f16vec2 o = hl == 0u ? loc + unpackFloat2x16(other) : unpackFloat2x16(other) + loc;
precise float16_t rs = o.x + o.y;
const float16_t r = float16_t(1.0 / float(rs));
const f16vec2 r2 = f16vec2(r, r);
[[unroll]] for (uint kt = 0u; kt < 4u; ++kt)
[[unroll]] for (uint q = 0u; q < 4u; ++q) {
precise f16vec2 pv = ep[kt][q] * r2;
{ const fe4m3vec2 t_ = fe4m3vec2(pv); Pb[kt][2u * q] = t_.x; Pb[kt][2u * q + 1u] = t_.y; }
}
}
B8 Ob[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC o = ACC(0.0);
[[unroll]] for (uint ks = 0u; ks < 2u; ++ks) K32(o, v0[c][ks], Pb[2u * ks], v1[c][ks], Pb[2u * ks + 1u])
Q8(Ob[c], o)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) K32(AT[wv][m], wo0[m], Ob[0], wo1[m], Ob[1])
}
}
// ---------------------------------------------------------------- r = x + attention * s1; the MLP
float rv[4][MT][8];
B8 Nb[4][MT];
ACC ML[4][MT];
const float f1 = float(s1), f2 = float(s2);
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv)
[[unroll]] for (uint m = 0u; m < MT; ++m) {
[[unroll]] for (int i = 0; i < 8; ++i) {
#if KIND == 4
const float xb = float(qx[wv][m][i]);
#else
const float xb = float(xr[wv][m][i]);
#endif
const float r = fma(AT[wv][m][i], f1, xb);
rv[wv][m][i] = r;
ML[wv][m][i] = float(w16.w[W_BU / 2u + 16u * m + 8u * hl + uint(i)]);
}
Q8(Nb[wv][m], rv[wv][m])
}
[[unroll]] for (uint br = 0u; br < NB; ++br) {
const uint w1 = W_WD0 + BRANCH * br, b1 = W_BD0 + BRANCH * br, w2 = W_WU0 + BRANCH * br;
A8 wd[2][MT], wu0[MT], wu1[MT];
[[unroll]] for (uint c = 0u; c < 2u; ++c)
[[unroll]] for (uint kt = 0u; kt < MT; ++kt) LDA(wd[c][kt], w1, C, c, kt);
[[unroll]] for (uint m = 0u; m < MT; ++m) { LDA(wu0[m], w2, 32u, m, 0u); LDA(wu1[m], w2, 32u, m, 1u); }
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
B8 Hb[2];
[[unroll]] for (uint c = 0u; c < 2u; ++c) {
ACC h;
[[unroll]] for (int i = 0; i < 8; ++i) h[i] = float(w16.w[b1 / 2u + 16u * c + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) K32(h, wd[c][kt], Nb[wv][kt], wd[c][kt + 1u], Nb[wv][kt + 1u])
ACT8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) K32(ML[wv][m], wu0[m], Hb[0], wu1[m], Hb[1])
}
}
// ---------------------------------------------------------------- y = r + MLP * s2: the stage's tokens
#if KIND == 4
A8 wh[MT];
[[unroll]] for (uint kt = 0u; kt < MT; ++kt) LDA(wh[kt], W_WH, C, 0u, kt);
#endif
[[unroll]] for (uint wv = 0u; wv < 4u; ++wv) {
const uint t = inter(16u * wv + col);
const int tx = x0 + int(t & 7u), ty = y0 + int(t >> 3u);
const bool inside = tx >= 0 && ty >= 0 && tx < ex && ty < ey;
const uint o = (outOff + uint(ty * ex + tx) * 32u) >> 2u;
#if KIND == 4
B8 Tm[MT];
#endif
[[unroll]] for (uint m = 0u; m < MT; ++m) {
floate4m3_t Tb[8];
[[unroll]] for (int i = 0; i < 8; i += 2) {
fe4m3vec2 q; saturatedConvertEXT(q, vec2(fma(ML[wv][m][i], f2, rv[wv][m][i]), fma(ML[wv][m][i+1], f2, rv[wv][m][i+1])));
Tb[i] = q.x; Tb[i+1] = q.y;
}
#if KIND == 4
[[unroll]] for (int i = 0; i < 8; ++i) Tm[m][i] = Tb[i];
#endif
const uint w0 = pack4(Tb[0], Tb[1], Tb[2], Tb[3]), w1 = pack4(Tb[4], Tb[5], Tb[6], Tb[7]);
const uint at = 4u * (m & 1u) + 2u * hl;
if (inside) { arena.w[o + (m >> 1u) * (plane >> 2u) + at] = w0; arena.w[o + (m >> 1u) * (plane >> 2u) + at + 1u] = w1; }
}
#if KIND == 4
{ // the guide head (rows 0-3 on lanes 0-15), for the translated tail
ACC hv;
[[unroll]] for (int i = 0; i < 8; ++i) hv[i] = float(w16.w[W_BH / 2u + 8u * hl + uint(i)]);
[[unroll]] for (uint kt = 0u; kt < MT; kt += 2u) K32(hv, wh[kt], Tm[kt], wh[kt + 1u], Tm[kt + 1u])
if (hl == 0u) [[unroll]] for (uint c = 0u; c < 4u; ++c) HD[t * 4u + c] = float16_t(hv[c]);
}
#endif
}
}
engine/rr/rr_exact.glsl
+7
View File
@@ -110,6 +110,13 @@ float16_t fma16q(float16_t a, float16_t b, float16_t c, inout bool mid)
#define Q8PAIRS(dst, src) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { fe4m3vec2 t_; saturatedConvertEXT(t_, vec2((src)[i_], (src)[i_ + 1])); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
// q8(act(h)) from the packed table (enc1_pack.act_table; the weight buffer's W8U view at W_ACT)
#define ACTQ(h) uintBitsToFloate4m3EXT(w8u.w[W_ACT + bh(h)])
// the activation of a whole accumulator into FP8. RR_NATIVE (D4R_RR_RELAXED=3): in f32, two elements per
// instruction, straight into FP8 - no binary16 in between and no table (not the reference's roundings)
#ifdef RR_NATIVE
#define ACTQ8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; i_ += 2) { const vec2 x_ = vec2((h)[i_], (h)[i_ + 1]); const vec2 y_ = clamp(x_, vec2(-2.0), vec2(2.0)); const vec2 g_ = fma(y_, fma(-abs(y_), vec2(0.0810546875), vec2(0.412109375)), vec2(0.5)); fe4m3vec2 t_; saturatedConvertEXT(t_, x_ * g_); (dst)[i_] = t_.x; (dst)[i_ + 1] = t_.y; }
#else
#define ACTQ8(dst, h) [[unroll]] for (int i_ = 0; i_ < 8; ++i_) (dst)[i_] = ACTQ(float16_t((h)[i_]));
#endif
float16_t act(float16_t x)
{
precise float16_t y = clamp(x, float16_t(-2.0), float16_t(2.0));
engine/rr/swin_m.comp
+1 -1
View File
@@ -342,7 +342,7 @@ void main()
LDA(a0, w1, C, c, kt); LDA(a1, w1, C, c, kt + 1u);
K32(h, a0, Nb[kt], a1, Nb[kt + 1u])
}
[[unroll]] for (int i = 0; i < 8; ++i) Hb[c][i] = ACTQ(float16_t(h[i]));
ACTQ8(Hb[c], h)
}
[[unroll]] for (uint m = 0u; m < MT; ++m) {
A8 a0, a1;
engine/test_k.py
+6 -1
View File
@@ -36,9 +36,14 @@ def defines(e, extra=()):
k = e['kind']
kind = {'plain': 0, 'enc': 1, 'enc0': 1, 'dec': 2, 'dec0': 2}[k]
n16 = e.get('merge_n', 16)
# K_DEFS: extra layer.comp defines. Two token tiles per wave (TPW) in the 32-channel layers by default: the only
# ones with registers to spare for it (engine/README.md); K_TPW=1 turns that off, a -DTPW in K_DEFS / extra wins
env = os.environ.get('K_DEFS', '').split()
given = any(d.startswith('-DTPW=') for d in [*extra, *env])
tpw = [f"-DTPW={os.environ.get('K_TPW', '2')}"] if e['C'] == 32 and not given else []
return [f'-DC={e["C"]}', f'-DNH={e["H"]}', f'-DKIND={kind}', f'-DPOSATTN={int(e["posattn"])}', f'-DEMBED={int(k == "enc0")}',
f'-DHEAD={int(k == "dec0")}', f'-DXC={e["X"] or 16}', f'-DMERGE_N={n16}', f'-DFOLD_GAINS={int(e.get("fold_gains", True))}',
*extra]
*tpw, *extra, *env]
def ulps(m, g):
m, g = m.astype(np.float64), g.astype(np.float64)
engine/test_m.py
+5 -2
View File
@@ -87,9 +87,12 @@ def pack_layer(t, name):
def defines(name):
C, heads, NPM, CIN = mm.LAYERS[name]
kind = {'enc0l': 3, 'dec0l': 4}.get(name, 1 if NPM else 2 if CIN else 0)
# compile_m.py's arithmetic (ZLUDA's lowering); M_TEXL=0 gives the native accuracy kernels' for the network layers
# compile_m.py's arithmetic (ZLUDA's lowering); M_TEXL=0 gives the native accuracy kernels' for the network layers.
# M_DEFS: extra defines for every layer_m.comp stage; M_DEFS_ENC1, M_DEFS_ENC3_TUBE, ...: for one layer
texl = ['-DTEXL=1'] if kind >= 3 or os.environ.get('M_TEXL', '1') != '0' else []
return [f'-DC={C}', f'-DNH={heads}', f'-DKIND={kind}', f'-DXC={NPM or CIN or 16}'] + texl + os.environ.get('M_DEFS', '').split()
# M_WPW=2: two 4x4-token windows per wave in the network layers (layer_m.comp WPW; not L's input / expansion)
wpw = [f"-DWPW={os.environ['M_WPW']}"] if kind < 3 and os.environ.get('M_WPW') else []
return [f'-DC={C}', f'-DNH={heads}', f'-DKIND={kind}', f'-DXC={NPM or CIN or 16}'] + texl + wpw + os.environ.get('M_DEFS', '').split() + os.environ.get('M_DEFS_' + name.upper(), '').split()
def token_codes(buf, off, W, H, C):
packaging/d4r.ini
+2 -2
View File
@@ -81,8 +81,8 @@ Enabled = false
; Another model K folder for games with the HDR flag (a path relative to the d4r folder, or a Windows path
; such as Z:\home\me\k)
ModelDir = auto
; auto / hip = the network on d4r's native HIP layer kernels when the model folder has them (faster),
; vulkan = on the engine's shader
; auto / vulkan = the network on the engine's shader, hip = on d4r's native HIP layer kernels when the model
; folder has them
Network = auto
[Paths]
tools/d4r_nvngx_shim.cpp
+5 -3
View File
@@ -674,8 +674,8 @@ static void load_portable_config()
portable_set("D4R_ENGINE", ini_flag(ini, "engine", "Enabled", 0) ? "1" : "0");
if (const std::string engineModel = ini_value(ini, "engine", "ModelDir"); !engineModel.empty())
portable_set(L"D4R_ENGINE_MODEL_DIR", portable_path(engineModel));
// Network: preset M's network on the native HIP layers ("hip", the default when the model has them: faster) or
// on the engine's Vulkan shader ("vulkan")
// Network: the network on the engine's Vulkan shader ("vulkan", and the default: "auto") or on the native HIP
// layers ("hip", when the model has them)
if (const std::string network = ini_value(ini, "engine", "Network"); network == "hip" || network == "vulkan")
portable_set("D4R_ENGINE_NETWORK", network);
else if (!network.empty() && network != "auto")
@@ -7748,8 +7748,10 @@ static bool create_engine_network(Feature& feature, const std::string& directory
bool stagesAvailable, unsigned int width, unsigned int height, size_t tokenBytes,
size_t headBytes)
{
// the Vulkan network unless [Engine] Network = hip asks for the native layers (2026-10-14: the shader with
// M_RELAXED and two windows per wave is faster than the HIP layers, engine/README.md)
const std::string choice = env_string("D4R_ENGINE_NETWORK");
if (choice == "vulkan")
if (choice != "hip")
return false;
auto unavailable = [&](const char* reason) {
logf("engine backend (%s): HIP network unavailable (%s); using the Vulkan network", preset, reason);