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

engine: SR layers keep their V tiles in a buffer that is dead during the attention

Browse Source 
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 <[email protected]>
This commit is contained in:
itchyfeetleech committed 2026-10-19 21:02:50 +13:00
1 parent 0aab224003
commit 912bef3fb8
1 file changed
+37 -5
Download Patch File Download Diff File Expand all files Collapse all files
engine/layer_m.comp
+37 -5
View File
@@ -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