From 912bef3fb806d4e19395947ceaf139d3fe9a1800 Mon Sep 17 00:00:00 2001 From: itchyfeetleech <169226071+itchyfeetleech@users.noreply.github.com> Date: Mon, 19 Oct 2026 17:02:04 +1300 Subject: [PATCH] engine: SR layers keep their V tiles in a buffer that is dead during the attention layer_m.comp's V tiles (VSH, 4 KiB, per wave and head, accessed only as matrices and only during the attention) had their own shared memory. Shared memory caps how many waves this GPU runs, and ACO gives a shader the registers that many waves leave, so the network layers ran at 3 - 6 waves per SIMD with registers to spare. VSH now takes the place of a buffer that is dead during the attention: - the merge rows NB (enc1, enc2, L's input stage), written after it; - the expansion NEX (dec1, dec2), read before it; - HO (L's expansion stage), written after it. One barrier each before the buffer's own use (none was needed while the buffers were apart). The tube keeps its own 4 KiB. | layer | shared memory | vector registers | waves per SIMD | | enc1 | 13 -> 9 KiB | 256 -> 216 | 4 -> 7 | | enc2 | 17 -> 13 KiB | 256 | 3 -> 4 | | dec2 | 13 -> 9 KiB | 256 -> 240 | 4 -> 6 | | dec1 | 10 -> 6 KiB | 240 -> 192 | 6 -> 8 | | L input | 16 -> 12 KiB | 192 -> 144 | 8 -> 10 | | L expansion | 13 -> 9 KiB | 192 -> 96 | 8 -> 14 | Output byte-identical (M and L, six temporal frames at 1280x720 -> 1440p and 1114x626 -> 1080p). RX 9070 XT, d4r-m 720p -> 1440p, with the previous commit: M 1583 -> 1566 us (network 967 -> 951), L 2349 -> 2293 us (input 479 -> 453, expansion 401 -> 382). The tube on one window per wave instead of two (192 registers, 8 waves) is 9 % slower and stays as it is. Signed-off-by: itchyfeetleech <169226071+itchyfeetleech@users.noreply.github.com> --- engine/layer_m.comp | 42 +++++++++++++++++++++++++++++++++++++----- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/engine/layer_m.comp b/engine/layer_m.comp index 8ead4dc..36bd42c 100644 --- a/engine/layer_m.comp +++ b/engine/layer_m.comp @@ -366,7 +366,6 @@ float16_t gelu1(float16_t x) precise f16vec2 C_ = f16vec2(EV(T, 4), EV(T, 5)), D_ = f16vec2(EV(T, 6), EV(T, 7)); \ P0_[T] = (A_ * A_) + (B_ * B_); P1_[T] = (C_ * C_) + (D_ * D_); } -shared f16vec4 VSH[4 * 2 * 16 * VWD / 4]; // per wave and head: the V tiles, transposed through shared memory #define MERGE (KIND == 1 || KIND == 3) #define EXPAND (KIND == 2 || KIND == 4) #if MERGE @@ -385,6 +384,34 @@ shared float16_t X0S[64 * XS]; // the embedding's output x0, the bl const uint HS = 40u; shared float16_t HO[64 * HS]; // the block's f16 output, window-major, for the per-pixel head #endif +// VSH: per wave and head the V tiles, transposed through shared memory (4 KiB, 512 bytes a tile). It is only accessed +// as matrices, and only during the attention, so it takes the place of a buffer that is dead then: the merge rows +// (written after the attention), the expansion (read before it) or, for KIND 4, the head's input (written after it). +// Shared memory caps how many waves this GPU runs, and ACO gives a shader the registers that number of waves leaves. +#if MERGE +#define VSH NB +#define VSH_UNIT 1u +#if 32 * (4 * C + 8) < 4096 +#error layer_m.comp: the V tiles no longer fit in NB +#endif +#elif KIND == 2 +#define VSH NEX +#define VSH_UNIT 1u +#if 16 * (4 * C + 8) < 4096 +#error layer_m.comp: the V tiles no longer fit in NEX +#endif +#elif KIND == 4 +#define VSH HO +#define VSH_UNIT 2u +#if 64 * 40 * 2 < 4096 +#error layer_m.comp: the V tiles no longer fit in HO +#endif +#else +shared floate4m3_t VSH[4 * 2 * 16 * VWD * 2]; +#define VSH_UNIT 1u +#endif +#define VSH_AT(t) ((t) * (512u / VSH_UNIT)) // tile t, in the buffer's elements +#define VSH_STRIDE (VWD * 2u / VSH_UNIT) // a row of VWD binary16 #define SYNC memoryBarrierShared(); barrier(); #if KIND == 4 // NVIDIA's f32 activations of the head (rrlite_dec0_4x4's output code): a sigmoid rounded to binary16, and a @@ -552,6 +579,9 @@ void main() H1N(N, rn) } EACH_T(NM1) } +#if KIND == 2 + SYNC // every wave has read the block's input from NEX: the V tiles take its place +#endif // ------------------------------------------------------------ attention, one head at a time for (uint hd = 0u; hd < uint(NH); ++hd) { @@ -603,8 +633,8 @@ void main() // 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); + coopMatStore(v[w][0], VSH, VSH_AT(mt * 2u + 0u), VSH_STRIDE, RM); + coopMatStore(v[w][1], VSH, VSH_AT(mt * 2u + 1u), VSH_STRIDE, RM); } } // scores: q8(Q) against itself, rows are the queries and columns the keys, plus the relative position bias @@ -634,8 +664,8 @@ void main() // 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]; A16 vt0, vt1; - coopMatLoad(vt0, VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM); - coopMatLoad(vt1, VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM); + coopMatLoad(vt0, VSH, VSH_AT(mt * 2u + 0u), VSH_STRIDE, RM); + coopMatLoad(vt1, VSH, VSH_AT(mt * 2u + 1u), VSH_STRIDE, RM); #if TEXL // P V is the block's one product of two f16 operands, and its f32 sum over the 16 keys is not exact: ZLUDA's // lowering of m16n8k16 puts key pair (j & 1) * 4 + 2 h + (j >> 1) in K dword j of half h (swin_tok.h @@ -732,6 +762,7 @@ void main() } #if KIND == 4 // ------------------------------------------------------------ L's expansion head: per-pixel reconstruction inputs + SYNC // every wave is done with the V tiles in HO FORW { WINV const uint hb = (mt * 16u + row) * HS + ch8; @@ -773,6 +804,7 @@ void main() #endif // ------------------------------------------------------------ patch merge: 2x2 tokens (4C) -> XC, half resolution #if MERGE + SYNC // every wave is done with the V tiles in NB FORW { WINV // the 16 merged rows are gathered from the four waves: the row is the merged token of the lane's pixel