#version 460
// d4r engine, preset K: one Swin block (8x8-token window) written from the layer's mathematics, not from NVIDIA's
// kernel. Weights come from the engine's own model file (extract_k.py): plain row-major matrices, with the two
// per-channel gains folded into Q/K/V and fc1.
// Build-time shape: C, NH, KIND (0 plain, 1 encoder with patch merge, 2 decoder with patch expand + skip),
// POSATTN (position-only attention), EMBED (enc0: 16-channel input embedding + ReLU), HEAD (dec0: output head),
// XC (merged channels or low-resolution channels), MERGE_N (XC rounded up to 16).
//  - Four waves per window, one per 16-token tile. The residual stream of a tile stays in registers (one f32
//    accumulator per 16 channels) from the input to the output.
//  - Shared memory only stages what has to be read back as a matrix operand or by another wave: one buffer BF and
//    the Q/K/V regions. It is kept small on purpose: the kernel waits on memory a lot, so its speed depends on how
//    many windows fit on a compute unit (64 KB of shared memory each).
//  - Rows are padded: with power-of-two row strides the 16 rows of an operand load land in the same LDS banks.
//  - Global arrays use f16vec4 units. Shared token tiles are row-major except enc0, which uses padded
//    column-major scalar-f16 arrays so cooperative-matrix stores can pack adjacent token values.
//  - FOLD_GAINS=0 keeps the per-channel gains explicit instead of folding them into weights.
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
#ifndef PADX
#define PADX 8
#endif
#ifndef FOLD_GAINS
#define FOLD_GAINS 1
#endif
#ifndef PADT
#define PADT 8
#endif
#ifndef PADS
#define PADS 8
#endif
#ifndef ABL_HEADS
#define ABL_HEADS uint(NH)
#endif
#ifndef ABL_CHUNKS
#define ABL_CHUNKS (UC / 8u)
#endif
layout(local_size_x = 128) 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;
layout(set = 0, binding = 3, std430) readonly buffer Sb { f16vec4 a[]; } skb;
layout(set = 0, binding = 4, std430) buffer O0 { f16vec4 a[]; } out0;
layout(set = 0, binding = 5, std430) buffer O1 { f16vec4 a[]; } out1;
#define ACC coopmat<float, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define HACC coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseAccumulator>
#define MA coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define MB coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CMJ gl_CooperativeMatrixLayoutColumnMajor
#define LD(m, arr, off, stride) coopMatLoad(m, arr, (off) / 4u, (stride) / 4u, RM)
#define ST(m, arr, off, stride) coopMatStore(HACC(m), arr, (off) / 4u, (stride) / 4u, RM)
#define SYNC memoryBarrierShared(); barrier();
const uint UC = uint(C), NT = uint(C) / 16u, HC = uint(C) / 2u;
const uint SX = UC + uint(PADX);            // row stride of the token layout (64 rows of C)
const uint S4 = 4u * UC + 8u;               // row stride of the 16 x 4C layout (patch expand / merge)
const uint SQ = 32u + uint(PADT);           // row stride of Q, K, V, attention output and MLP hidden (64 rows each)
const uint RQ = 64u * SQ;                   // one such region
const uint SLO = uint(XC) + 8u, SEM = 24u, SHD = 56u, SMG = uint(MERGE_N) + 8u;   // low-res input, embedding input, head, merged
#if POSATTN
#define NREG 1
#define VOFF 0u
#else
#define NREG 2
#define VOFF RQ
#endif
// TT must also hold the token layout (decoders: x0), the 16 x 4C layout (encoders: merge input) or the head's output
#if KIND == 2
#define TT_ALT (16 * (C + PADX))
#elif KIND == 1
#define TT_ALT (4 * (4 * C + 8))
#else
#define TT_ALT 0
#endif
#if HEAD
#define TT_HEAD (16 * 56)
#else
#define TT_HEAD 0
#endif
#define TT_REG (NREG * 16 * (32 + PADT))
#define MAX2(a, b) ((a) > (b) ? (a) : (b))
// BF: this tile's rows hold, in turn, the normalised input, Q, the scores / probabilities, the attention output, the
// MLP input and the block output; as a whole it is also the 16 x 4C scratch. TT: K | V shared by the four waves; its
// first region doubles as MLP hidden and scratch.
// enc0's 32-channel, position-only block benefits from column-major shared tiles. All addressing uses
// cooperative-matrix layouts rather than hardware fragment indices; other layers retain their existing layout.
#ifndef COLUMN_MAJOR_ENC0
#define COLUMN_MAJOR_ENC0 (EMBED && C == 32)
#endif
#if COLUMN_MAJOR_ENC0 && (!EMBED || C != 32 || !POSATTN || KIND != 1 || HEAD)
#error Column-major shared tiles require the 32-channel positional enc0 block
#endif
#if COLUMN_MAJOR_ENC0
const uint SP = 72u; // 64 token rows plus padding; patch scratch is transposed with 16 rows
shared float16_t BF[32 * 72];
shared float16_t TT[32 * 72];
#define LOAD_B_TOKEN(m,r,c) coopMatLoad(m,BF,(c)*SP+(r),SP,CMJ)
#define STORE_B_TOKEN(m,r,c) coopMatStore(HACC(m),BF,(c)*SP+(r),SP,CMJ)
#define LOAD_T_TOKEN(m,r,c,s) coopMatLoad(m,TT,(c)*SP+(r),SP,CMJ)
#define STORE_T_TOKEN(m,r,c,s) coopMatStore(HACC(m),TT,(c)*SP+(r),SP,CMJ)
#define STORE_B_PATCH(m,r,c,s) coopMatStore(HACC(m),BF,(c)*16u+(r),16u,CMJ)
#define LOAD_T_RAW(m,off,s) coopMatLoad(m,TT,off,s,RM)
f16vec4 readBToken(uint r,uint c) {uint i=c*SP+r;return f16vec4(BF[i],BF[i+SP],BF[i+2u*SP],BF[i+3u*SP]);}
f16vec4 readTToken(uint r,uint c,uint stride) {uint i=c*SP+r;return f16vec4(TT[i],TT[i+SP],TT[i+2u*SP],TT[i+3u*SP]);}
f16vec4 readBPatch(uint r,uint c,uint stride) {uint i=c*16u+r;return f16vec4(BF[i],BF[i+16u],BF[i+32u],BF[i+48u]);}
void writeBToken(uint r,uint c,f16vec4 v) {uint i=c*SP+r;BF[i]=v.x;BF[i+SP]=v.y;BF[i+2u*SP]=v.z;BF[i+3u*SP]=v.w;}
void writeTToken(uint r,uint c,uint stride,f16vec4 v) {uint i=c*SP+r;TT[i]=v.x;TT[i+SP]=v.y;TT[i+2u*SP]=v.z;TT[i+3u*SP]=v.w;}
void writeTRaw(uint off,f16vec4 v) {TT[off]=v.x;TT[off+1u]=v.y;TT[off+2u]=v.z;TT[off+3u]=v.w;}
#else
shared f16vec4 BF[16 * (C + PADX)];
shared f16vec4 TT[MAX2(TT_REG, MAX2(TT_ALT, TT_HEAD))];
#define LOAD_B_TOKEN(m,r,c) LD(m,BF,(r)*SX+(c),SX)
#define STORE_B_TOKEN(m,r,c) ST(m,BF,(r)*SX+(c),SX)
#define LOAD_T_TOKEN(m,r,c,s) LD(m,TT,(r)*(s)+(c),s)
#define STORE_T_TOKEN(m,r,c,s) ST(m,TT,(r)*(s)+(c),s)
#define STORE_B_PATCH(m,r,c,s) ST(m,BF,(r)*(s)+(c),s)
#define LOAD_T_RAW(m,off,s) LD(m,TT,off,s)
f16vec4 readBToken(uint r,uint c) {return BF[(r*SX+c)/4u];}
f16vec4 readTToken(uint r,uint c,uint stride) {return TT[(r*stride+c)/4u];}
f16vec4 readBPatch(uint r,uint c,uint stride) {return BF[(r*stride+c)/4u];}
void writeBToken(uint r,uint c,f16vec4 v) {BF[(r*SX+c)/4u]=v;}
void writeTToken(uint r,uint c,uint stride,f16vec4 v) {TT[(r*stride+c)/4u]=v;}
void writeTRaw(uint off,f16vec4 v) {TT[off/4u]=v;}
#endif
int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
int tdiv2(int v) { return v < 0 ? -((-v) / 2) : v / 2; }
ACC bias_tile(uint off, uint stride) { HACC b; LD(b, wb.w, off, stride); return ACC(b); }   // vectors are stored as 16 equal rows
f16vec4 softmax_exp(f16vec4 s)
{
    const f16vec4 t = clamp(s, f16vec4(uint16BitsToFloat16(uint16_t(0xCD69))), f16vec4(uint16BitsToFloat16(uint16_t(0x4D69))));
    const f16vec4 v = fma(t * t, f16vec4(uint16BitsToFloat16(uint16_t(0x91D5))), f16vec4(1.0));
    return f16vec4(exp2(vec4(fma(t, v, f16vec4(uint16BitsToFloat16(uint16_t(0xC50D)))))));
}
f16vec4 gelu(f16vec4 x)
{
    const f16vec4 c = clamp(x, f16vec4(-2.0), f16vec4(2.0));
    // the reference rounds after every binary16 operation (its compiler has fp contraction off): keep the
    // same steps, or a contracted fma changes the last bit of most outputs
    precise f16vec4 u = f16vec4(0.41216981) - f16vec4(0.08108133) * abs(c);
    precise f16vec4 y = f16vec4(0.5) + c * u;
    return x * y;
}
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 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
    // ------------------------------------------------------------ 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 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]);
        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]);
        }
        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);
        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]);
        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);
        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);
        }
        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
    {
        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
    {
        // 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);
            residual[nt] = x;
            acc[nt] = bias_tile(p.o[1] + nt * 16u, UC);
        }
        vec4 s4 = vec4(0.0);
        f16vec4 x[C / 8];
        for (uint c = 0u; c < HC; c += 4u) {
            x[c / 4u] = READ_X0(tok,ch0+c);
#if FOLD_GAINS
            const vec4 f = vec4(x[c / 4u]); s4 += f * f;
#else
            s4 += vec4(x[c / 4u] * x[c / 4u]);
#endif
        }
        float s = s4.x + s4.y + s4.z + s4.w;
        s += subgroupShuffleXor(s, 16u);
        const float16_t r = float16_t(inversesqrt(float(float16_t(s))));
        for (uint c = 0u; c < HC; c += 4u) {
#if FOLD_GAINS
            writeBToken(tok,ch0+c,x[c/4u]*r);
#else
            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);
    // ------------------------------------------------------------ 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);
                }
#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);
                }
#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);
        }
        SYNC
#if !POSATTN
        {
            ACC s[4];
            for (uint nt = 0u; nt < 4u; ++nt) { HACC tb; LD(tb, wb.w, otab + mt * 1024u + nt * 16u, 64u); s[nt] = ACC(tb); }
            MA qa[2]; LOAD_B_TOKEN(qa[0],mt*16u,0u); LOAD_B_TOKEN(qa[1],mt*16u,16u);
            for (uint kt = 0u; kt < 2u; ++kt)
            {
                MB b[4];
                for (uint nt = 0u; nt < 4u; ++nt) coopMatLoad(b[nt], TT, (nt * 16u * SQ + kt * 16u) / 4u, SQ / 4u, CMJ);     // K transposed
                for (uint nt = 0u; nt < 4u; ++nt) s[nt] = coopMatMulAdd(qa[kt], b[nt], s[nt]);
            }
            for (uint nt = 0u; nt < 4u; ++nt) STORE_B_TOKEN(s[nt],mt*16u,nt*16u);        // over Q, which is in registers
            // softmax over the token's 64 scores: 32 per lane
            f16vec4 e[8];
            vec4 sum = vec4(0.0);
            for (uint j = 0u; j < 8u; ++j) { e[j] = softmax_exp(readBToken(tok,hf*32u+j*4u)); sum += vec4(e[j]); }
            float tot = sum.x + sum.y + sum.z + sum.w;
            tot += subgroupShuffleXor(tot, 16u);
            const float16_t rc = float16_t(1.0 / float(float16_t(tot)));
            for (uint j = 0u; j < 8u; ++j) writeBToken(tok,hf*32u+j*4u,e[j]*rc);
        }
#endif
        {
            ACC o0 = ACC(0.0), o1 = ACC(0.0);
            // two K steps at a time: few live operands keep the register count, and with it the occupancy, in check
            for (uint k2 = 0u; k2 < 4u; k2 += 2u)
            {
                MA a[2]; MB b[4];
                for (uint i = 0u; i < 2u; ++i)
                {
                    const uint kt = k2 + i;
#if POSATTN
                    LD(a[i], wb.w, otab + mt * 1024u + kt * 16u, 64u);
#else
                    LOAD_B_TOKEN(a[i],mt*16u,kt*16u);
#endif
                    LOAD_T_TOKEN(b[2u*i],VOFF/SQ+kt*16u,0u,SQ); LOAD_T_TOKEN(b[2u*i+1u],VOFF/SQ+kt*16u,16u,SQ);
                }
                for (uint i = 0u; i < 2u; ++i) { o0 = coopMatMulAdd(a[i], b[2u * i], o0); o1 = coopMatMulAdd(a[i], b[2u * i + 1u], o1); }
            }
            // 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]));
                }
            }
        }
        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)
    {
        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);
            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); }
            }
            STORE_T_TOKEN(h0,mt*16u,0u,SQ); STORE_T_TOKEN(h1,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]));
            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]));
            }
        }
        for (uint nt = 0u; nt < NT; ++nt) STORE_B_TOKEN(acc[nt],mt*16u,nt*16u);
    }
    // ------------------------------------------------------------ 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);
#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);
        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]);
        }
        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);
#else
    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));
        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)
            {
                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[i] = coopMatMulAdd(a[j], b[j * NI + i], mg[i]);
        }
        // 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);
    }
#endif
}