engine: this GPU's own arithmetic and shared weight tiles as the builders' defaults
- 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
1 parent
8eb7a94533
commit
f0b55a850a
24 files changed
+1108
-254
No files matched your search
@@ -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.
|
||||
@@ -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"
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
+82
-34
@@ -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);
|
||||
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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));
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in new issue
Block a user