#version 460
// d4r engine, preset M (rrlite): one Swin block over 4x4-token windows, register-resident.
//
// The cooperative matrix is opaque in GLSL, but RADV's gfx1201 lowering has a fixed lane/slot <-> element
// mapping, the same for every 16x16x16 use. It was proven on this GPU with the probes in
// build/engine-quality/rm-probe (each lane writes its own slot values into an operand and the resulting image is
// read back):
//
//   RowMajor    sources and destinations are indexed [row][col]
//   ColumnMajor sources and destinations are indexed [col][row]
//   accumulator and B operand: element i of lane l is (row = 8*(l>>4) + i, col = l & 15)
//   A operand:                 element i of lane l is (row = l & 15,        col = 8*(l>>4) + i)
//
// Every product here is therefore weights x activations (A = weights, B = activations): a product's rows are the
// output channels and its columns the tokens, and the per-lane elements of a result ARE the per-lane elements of
// the next product's B operand. Cooperative-matrix subscripts read and write those elements in place, so the
// attention and the MLP chain never leave the registers. Only the patch expand (decoders) and the patch merge
// (encoders) move data between waves, and only those use shared memory and a workgroup barrier.
//
// NVIDIA's arithmetic (the semantics of kernels/m/swin_common.h K32_MIX, which the captured launches agree with)
// rounds the accumulator to binary16 after every logical 32 products; the 16x16x16 cells have no K32 form, so a
// logical k32 is two K16 steps summed in f32 and rounded once. RADV folds a register-only f32 -> f16 -> f32
// conversion away, so the rounding of P + the previous accumulator passes through a per-invocation binary16
// scratch slot (no cross-lane traffic: a subgroup memory barrier is enough). The tile loops are written out one
// tile at a time because a cooperative matrix indexed by a loop variable is not a register any more.
//
// The two norms do NOT use the natural element order of this shader's accumulators: the native register element
// (tile T, half hf, element I) holds the natural channel tok_xnat(T, 8hf+I) = 16T + 8*((I>>2)&1) + 4hf + (I&3),
// and tok_rsq's summation tree groups the lanes' squares in that order. Norm 1 therefore reads the input at that
// channel (XVALN, the absolute one, not the lane's own half) and norm 2 rebuilds the native order from the
// accumulator registers with a pair-aligned half-wave exchange. Without this the ten layers reach only 41-46 dB
// against the capture even though every product, weight map and bias is already the native one.
//
// Build-time shape: C, NH, KIND (0 plain, 1 encoder with 2x2 patch merge, 2 decoder with patch expand + skip),
// XC (merged channels or low-resolution channels). One workgroup is an 8x8-token block: four waves, one per
// 4x4-token window. Bindings and the parameter block are unchanged; the FP8 weight matrices are the packed
// fixed-16 tiles of test_m.pack8 (the f16 vectors stay row-major).
#extension GL_KHR_cooperative_matrix : require
#extension GL_KHR_memory_scope_semantics : require
#extension GL_EXT_shader_explicit_arithmetic_types : require
#extension GL_EXT_float_e4m3 : require
#extension GL_KHR_shader_subgroup_basic : require
#extension GL_KHR_shader_subgroup_shuffle : require
layout(local_size_x = 128) in;
// o: 0 g1, 1 bo, 2 g2, 3 b2 (f16); per head h at 4 + 4h: q, v (FP8), bias (f16), wo (FP8);
//    per MLP chunk c at 20 + 3c: w1 (FP8), b1 (f16), w2 (FP8); 68, 69: merge or expand weights (FP8) and bias (f16)
layout(set = 0, binding = 0, std430) readonly buffer Pb { int W; int Hh; int sx; int sy; uint o[72]; } p;
layout(set = 0, binding = 1, std430) readonly buffer W8 { floate4m3_t w[]; } w8;
layout(set = 0, binding = 2, std430) readonly buffer W16 { f16vec4 w[]; } w16;
layout(set = 0, binding = 3, std430) readonly buffer Ib { floate4m3_t a[]; } inb;      // [H][W][C]; decoders: [H/2][W/2][XC]
layout(set = 0, binding = 4, std430) readonly buffer Sb { floate4m3_t a[]; } skb;      // decoders: skip [H][W][C]
layout(set = 0, binding = 5, std430) buffer O0 { floate4m3_t a[]; } out0;              // [H][W][C]
layout(set = 0, binding = 6, std430) buffer O1 { floate4m3_t a[]; } out1;              // encoders: merged [H/2][W/2][XC]
#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 A16 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseA>
#define B16 coopmat<float16_t, gl_ScopeSubgroup, 16, 16, gl_MatrixUseB>
#define RM gl_CooperativeMatrixLayoutRowMajor
#define CM gl_CooperativeMatrixLayoutColumnMajor
const uint UC = uint(C), NT = uint(C) / 16u, UXC = uint(XC);
const uint SLO = UXC + 8u;                  // row stride of the low-resolution token image
const uint S4 = 4u * UC + 8u;               // row stride of the 4C rows (patch expand / merge)
const uint VWD = 16u;                       // row width of the per-wave V transpose buffer, in f16
int mirror(int v, int n) { v = abs(v); return min(v, 2 * n - 2 - v); }
#define Q16(v) float16_t(v)
#define CL8(v) clamp(v, float16_t(-448.0), float16_t(448.0))
// NVIDIA's accumulator rounding: one binary16 rounding of the whole tile after every logical 32 products.
// RADV folds a register-only f32->f16->f32 conversion away (measured: half of all values and every tie
// midpoint unchanged), so the rounding has to pass through binary16 memory. Each invocation owns 8 halves of
// its own scratch: no cross-lane traffic, so a subgroup fence is enough (no workgroup barrier).
shared f16vec4 RS[128 * 2];
#define R16(m) { \
    RS[gl_LocalInvocationID.x * 2u + 0u] = f16vec4(float16_t((m)[0]), float16_t((m)[1]), float16_t((m)[2]), float16_t((m)[3])); \
    RS[gl_LocalInvocationID.x * 2u + 1u] = f16vec4(float16_t((m)[4]), float16_t((m)[5]), float16_t((m)[6]), float16_t((m)[7])); \
    subgroupMemoryBarrierShared(); \
    const f16vec4 r0_ = RS[gl_LocalInvocationID.x * 2u + 0u], r1_ = RS[gl_LocalInvocationID.x * 2u + 1u]; \
    (m)[0] = float(r0_.x); (m)[1] = float(r0_.y); (m)[2] = float(r0_.z); (m)[3] = float(r0_.w); \
    (m)[4] = float(r1_.x); (m)[5] = float(r1_.y); (m)[6] = float(r1_.z); (m)[7] = float(r1_.w); }
// one logical k32, exactly as the native K32_MIX: P = A0*B0 + A1*B1 from a zero seed, then one binary16
// rounding of P + the previous accumulator (the native v_fma_mix), through the scratch so it is not folded.
// Packed fixed-16 weight tiles (test_m.pack8): matrix W[K][N] as [K/16][N/16][N16][K16], so one 16x16 A tile
// is 256 contiguous bytes at base + 256*(kTile*(N/16) + nOffset/16), read RowMajor with stride 16. Same A
// operand values as the old plain W[K][N] ColumnMajor stride-N load, but a plain (non-transposed) load.
#define LDA8(dst, base, ks, stride, moff) coopMatLoad((dst), w8.w, (base) + 256u * (uint(ks) * (uint(stride) / 16u) + (uint(moff) / 16u)), 16u, RM)

#define K32H(dst, a0, b0, a1, b1) { \
    ACC P_ = coopMatMulAdd((a1), (b1), coopMatMulAdd((a0), (b0), ACC(0.0))); \
    RS[gl_LocalInvocationID.x * 2u + 0u] = f16vec4( \
        float16_t(P_[0] + float((dst)[0])), float16_t(P_[1] + float((dst)[1])), \
        float16_t(P_[2] + float((dst)[2])), float16_t(P_[3] + float((dst)[3]))); \
    RS[gl_LocalInvocationID.x * 2u + 1u] = f16vec4( \
        float16_t(P_[4] + float((dst)[4])), float16_t(P_[5] + float((dst)[5])), \
        float16_t(P_[6] + float((dst)[6])), float16_t(P_[7] + float((dst)[7]))); \
    subgroupMemoryBarrierShared(); \
    const f16vec4 u0_ = RS[gl_LocalInvocationID.x * 2u + 0u], u1_ = RS[gl_LocalInvocationID.x * 2u + 1u]; \
    (dst)[0] = u0_.x; (dst)[1] = u0_.y; (dst)[2] = u0_.z; (dst)[3] = u0_.w; \
    (dst)[4] = u1_.x; (dst)[5] = u1_.y; (dst)[6] = u1_.z; (dst)[7] = u1_.w; }
// a tiled f16 vector (16 equal rows) as the initial value of an accumulator whose rows are its channels
HAC bias16(uint off, uint stride) { HAC b; coopMatLoad(b, w16.w, off / 4u, stride / 4u, CM); return b; }
#define MK8_F8(dst, a) { \
    const float16_t q0_ = CL8(Q16((a)[0])), q1_ = CL8(Q16((a)[1])), q2_ = CL8(Q16((a)[2])), q3_ = CL8(Q16((a)[3])); \
    const float16_t q4_ = CL8(Q16((a)[4])), q5_ = CL8(Q16((a)[5])), q6_ = CL8(Q16((a)[6])), q7_ = CL8(Q16((a)[7])); \
    (dst)[0] = floate4m3_t(q0_); (dst)[1] = floate4m3_t(q1_); (dst)[2] = floate4m3_t(q2_); (dst)[3] = floate4m3_t(q3_); \
    (dst)[4] = floate4m3_t(q4_); (dst)[5] = floate4m3_t(q5_); (dst)[6] = floate4m3_t(q6_); (dst)[7] = floate4m3_t(q7_); }
// NVIDIA's softmax numerator, two neighbouring columns in one 32-bit word (the even column's carry reaches the odd)
void softmax2(float16_t s0, float16_t s1, out float16_t e0, out float16_t e1)
{
    const float16_t lim = float16_t(0.55615234375);
    const float16_t cs = float16_t(0.0972222164273262 / 5.656854152679443);
    const float16_t t0 = clamp(s0 * cs, -lim, lim), t1 = clamp(s1 * cs, -lim, lim);
    const float16_t p0 = fma(t0, fma(t0, -t0, float16_t(0.92730712890625)), float16_t(1.375));
    const float16_t p1 = fma(t1, fma(t1, -t1, float16_t(0.92730712890625)), float16_t(1.375));
    const uint w = ((uint(float16BitsToUint16(p0)) | (uint(float16BitsToUint16(p1)) << 16)) << 5) + 0x7FF88000u;
    e0 = uint16BitsToFloat16(uint16_t(w & 0xFFFFu)); e1 = uint16BitsToFloat16(uint16_t(w >> 16));
}
float16_t gelu1(float16_t x)
{
    const float16_t c = clamp(x, float16_t(-2.0), float16_t(2.0));
    const float16_t ac = abs(c);
    precise float16_t inner = float16_t(0.41216981) - float16_t(0.08108133) * ac;
    precise float16_t outer = float16_t(0.5) + c * inner;
    return x * outer;
}
#if C == 64
#define EACH_T(X) X(0) X(1) X(2) X(3)
#define EACH_T16(X) X(0, 1) X(2, 3)
#elif C == 96
#define EACH_T(X) X(0) X(1) X(2) X(3) X(4) X(5)
#define EACH_T16(X) X(0, 1) X(2, 3) X(4, 5)
#else
#define EACH_T(X) X(0) X(1) X(2) X(3) X(4) X(5) X(6) X(7)
#define EACH_T16(X) X(0, 1) X(2, 3) X(4, 5) X(6, 7)
#endif

// ------------------------------------------------------------------ native arithmetic order
// The accuracy build (kernels/build.sh D4R_PREFER_ACCURACY=1 -> -DD4R_ACCURACY -> SWIN_EXACT, and
// D4R_NATIVE_FP8=1 -> -DD4R_FP8_WMMA -> SWIN_A8/K32_MIX) runs the token-lane block of kernels/m/swin_tok.h.
// Its two non-obvious orderings are reproduced here in binary16:
//  * tok_rsq: per tile pp = a*a + b*b over the pairs (0,1)-(2,3) and (4,5)-(6,7), then a tree that splits the
//    tiles by parity, then a + other(a) across the half-waves, then u[0] + u[1], then the two components.
//  * the softmax: v1[k] = wg[k] + other(wg[k]) elementwise first, then (v1[0]+v1[1]) + (v1[2]+v1[3]).
// f16 arithmetic is marked precise so nothing is contracted into an FMA (the native kernel builds with
// #pragma clang fp contract(off)).
#define SH16(v) subgroupShuffleXor(packHalf2x16(vec2(float((v).x), float((v).y))), 16u)
#define UN16(p) f16vec2(float16_t(unpackHalf2x16(p).x), float16_t(unpackHalf2x16(p).y))
#if C == 64
#define RSQ_TREE(OUT) { \
    f16vec2 P0_[NT], P1_[NT]; \
    EACH_T(RSQ_TILE) \
    precise f16vec2 W00_ = P0_[0] + P0_[2], W01_ = P1_[0] + P1_[2]; \
    precise f16vec2 W10_ = P0_[1] + P0_[3], W11_ = P1_[1] + P1_[3]; \
    precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \
    precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
    precise f16vec2 S_ = U0_ + U1_; \
    OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); }
#elif C == 96
#define RSQ_TREE(OUT) { \
    f16vec2 P0_[NT], P1_[NT]; \
    EACH_T(RSQ_TILE) \
    precise f16vec2 W00_ = (P0_[0] + P0_[2]) + P0_[4], W01_ = (P1_[0] + P1_[2]) + P1_[4]; \
    precise f16vec2 W10_ = (P0_[1] + P0_[3]) + P0_[5], W11_ = (P1_[1] + P1_[3]) + P1_[5]; \
    precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \
    precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
    precise f16vec2 S_ = U0_ + U1_; \
    OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); }
#else
#define RSQ_TREE(OUT) { \
    f16vec2 P0_[NT], P1_[NT]; \
    EACH_T(RSQ_TILE) \
    precise f16vec2 W00_ = (P0_[0] + P0_[2]) + (P0_[4] + P0_[6]), W01_ = (P1_[0] + P1_[2]) + (P1_[4] + P1_[6]); \
    precise f16vec2 W10_ = (P0_[1] + P0_[3]) + (P0_[5] + P0_[7]), W11_ = (P1_[1] + P1_[3]) + (P1_[5] + P1_[7]); \
    precise f16vec2 A0_ = W00_ + W10_, A1_ = W01_ + W11_; \
    precise f16vec2 U0_ = A0_ + UN16(SH16(A0_)), U1_ = A1_ + UN16(SH16(A1_)); \
    precise f16vec2 S_ = U0_ + U1_; \
    OUT = float16_t(1.0 / sqrt(float(float16_t((S_.x + S_.y) + float16_t(0.0001220703125))))); }
#endif
#define RSQ_TILE(T) { \
    precise f16vec2 A_ = f16vec2(EV(T, 0), EV(T, 1)), B_ = f16vec2(EV(T, 2), EV(T, 3)); \
    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
#if KIND == 1
shared floate4m3_t NB[32 * S4];             // 16 merged rows, preceded by 16 rows of slack for the gather
#endif
#if KIND == 2
shared floate4m3_t NLI[16 * SLO];           // the block's 16 low-resolution tokens
shared floate4m3_t NEX[16 * S4];            // their expanded 4C columns
#endif
#define SYNC memoryBarrierShared(); barrier();
void main()
{
    const uint mt = gl_LocalInvocationID.x >> 5, lane = gl_LocalInvocationID.x & 31u;
    const uint row = lane & 15u, hf = lane >> 4, ch8 = 8u * hf;     // one token per lane, its channels 8h..8h+7
    const int X0 = 8 * int(gl_WorkGroupID.x) - p.sx, Y0 = 8 * int(gl_WorkGroupID.y) - p.sy;
    const uint bxr = 4u * (mt & 1u) + (row & 3u), byr = 4u * (mt >> 1) + (row >> 2);
    const int tx = X0 + int(bxr), ty = Y0 + int(byr);
    const bool inside = ty >= 0 && ty < p.Hh && tx >= 0 && tx < p.W;
    const uint iy = uint(mirror(ty, p.Hh)), ix = uint(mirror(tx, p.W));
    HAC res[NT];                                // the residual stream: rows are channels, columns are tokens
    B8 h1[NT];                                  // the normalised activation, the K operand of every product
    // ------------------------------------------------------------ patch expand: x0 = q8(E) + skip
#if KIND == 2
    {
        const uint W2 = uint(p.W) / 2u, H2 = uint(p.Hh) / 2u, part = UXC / 8u, c0 = (mt * 2u + hf) * part;
        const uint lo = uint(mirror((Y0 >> 1) + int(row >> 2), int(H2)) * int(W2) + mirror((X0 >> 1) + int(row & 3u), p.W / 2)) * UXC;
        for (uint c = c0; c < c0 + part; ++c) NLI[row * SLO + c] = inb.a[lo + c];
        SYNC
        // this wave computes sub-position mt of every low-resolution token's 4C columns
        HAC ef[NT];
#define EF_INIT(N) ef[N] = bias16(p.o[69] + 16u * (mt * NT + (N)), 4u * UC);
#define EF_K32(N) { A8 xa0, xa1; \
    LDA8(xa0, p.o[68], kt, 4u * UC, 16u * (mt * NT + (N))); \
    LDA8(xa1, p.o[68], kt + 1u, 4u * UC, 16u * (mt * NT + (N))); \
    K32H(ef[N], xa0, b0, xa1, b1) }
#define EF_STORE(N) { A8 ey; MK8_F8(ey, ef[N]); coopMatStore(ey, NEX, mt * UC + 16u * (N), S4, RM); }
        EACH_T(EF_INIT)
        for (uint kt = 0u; kt < UXC / 16u; kt += 2u)
        {
            B8 b0, b1;
            coopMatLoad(b0, NLI, kt * 16u, SLO, CM);
            coopMatLoad(b1, NLI, (kt + 1u) * 16u, SLO, CM);
            EACH_T(EF_K32)
        }
        EACH_T(EF_STORE)
        SYNC
    }
#endif
    // ------------------------------------------------------------ norm 1: sum of squares, residual and h1
#if KIND == 2
    const uint lb0 = ((byr >> 1) * 4u + (bxr >> 1)) * S4 + ((byr & 1u) * 2u + (bxr & 1u)) * UC + ch8;
    const uint sb0 = (iy * uint(p.W) + ix) * UC + ch8;
#define XVAL(C) (float16_t(NEX[lb0 + (C)]) + float16_t(skb.a[sb0 + (C)]))
#else
    const uint ib0 = (iy * uint(p.W) + ix) * UC + ch8;
#define XVAL(C) float16_t(inb.a[ib0 + (C)])
#endif
// The native's register element (tile T, half hf, element I) holds the natural channel tok_xnat(T, 8hf+I) =
// 16T + 8*((I>>2)&1) + 4hf + (I&3) (kernels/m/swin_tok.h tok_rsq on the register layout the input plane implies).
// This shader's own accumulators and A rows keep natural channels, so h1 and the residual use XVAL (the lane's
// own half); the norm tree, however, groups per the native element order, so it reads the absolute channel.
#if KIND == 2
#define XVALN(C) (float16_t(NEX[(lb0 - ch8) + (C)]) + float16_t(skb.a[(sb0 - ch8) + (C)]))
#else
#define XVALN(C) float16_t(inb.a[(iy * uint(p.W) + ix) * UC + (C)])
#endif
#define EV(T, I) XVALN(16u * (T) + 8u * ((uint(I) >> 2) & 1u) + 4u * hf + (uint(I) & 3u))
    float16_t rn; RSQ_TREE(rn)
#undef EV
#define NM1(N) { \
    const uint wb = (p.o[0] + 16u * (N) + ch8) / 4u, bb = (p.o[1] + 16u * (N) + ch8) / 4u; \
    const f16vec4 g0 = w16.w[wb], g1 = w16.w[wb + 1u], d0 = w16.w[bb], d1 = w16.w[bb + 1u]; \
    const float16_t x0 = XVAL(16u * (N) + 0u), x1 = XVAL(16u * (N) + 1u), x2 = XVAL(16u * (N) + 2u), x3 = XVAL(16u * (N) + 3u); \
    const float16_t x4 = XVAL(16u * (N) + 4u), x5 = XVAL(16u * (N) + 5u), x6 = XVAL(16u * (N) + 6u), x7 = XVAL(16u * (N) + 7u); \
    res[N][0] = float16_t(x0 + d0.x); res[N][1] = float16_t(x1 + d0.y); \
    res[N][2] = float16_t(x2 + d0.z); res[N][3] = float16_t(x3 + d0.w); \
    res[N][4] = float16_t(x4 + d1.x); res[N][5] = float16_t(x5 + d1.y); \
    res[N][6] = float16_t(x6 + d1.z); res[N][7] = float16_t(x7 + d1.w); \
    h1[N][0] = floate4m3_t(CL8(x0 * float16_t(rn * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(rn * g0.y))); \
    h1[N][2] = floate4m3_t(CL8(x2 * float16_t(rn * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(rn * g0.w))); \
    h1[N][4] = floate4m3_t(CL8(x4 * float16_t(rn * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(rn * g1.y))); \
    h1[N][6] = floate4m3_t(CL8(x6 * float16_t(rn * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(rn * g1.w))); }
    EACH_T(NM1)
    // ------------------------------------------------------------ attention, one head at a time
    for (uint hd = 0u; hd < uint(NH); ++hd)
    {
        const uint oq = p.o[4u + 4u * hd], ov = p.o[5u + 4u * hd], ob = p.o[6u + 4u * hd], ow = p.o[7u + 4u * hd];
        HAC q[2], v[2];
        q[0] = HAC(0.0); q[1] = HAC(0.0); v[0] = HAC(0.0); v[1] = HAC(0.0);
#define QV_G(A, B) { \
        const uint kt = A; \
        A8 aq0, aq1, av0, av1; \
        LDA8(aq0, oq, kt, 32u, 0u); \
        LDA8(aq1, oq, kt + 1u, 32u, 0u); \
        LDA8(av0, ov, kt, 32u, 0u); \
        LDA8(av1, ov, kt + 1u, 32u, 0u); \
        K32H(q[0], aq0, h1[A], aq1, h1[B]) \
        K32H(v[0], av0, h1[A], av1, h1[B]) \
        LDA8(aq0, oq, kt, 32u, 16u); \
        LDA8(aq1, oq, kt + 1u, 32u, 16u); \
        LDA8(av0, ov, kt, 32u, 16u); \
        LDA8(av1, ov, kt + 1u, 32u, 16u); \
        K32H(q[1], aq0, h1[A], aq1, h1[B]) \
        K32H(v[1], av0, h1[A], av1, h1[B]) }
        EACH_T16(QV_G)
        // scores: q8(Q) against itself, rows are the queries and columns the keys, plus the relative position bias
        A8 qa[2]; B8 qb[2];
        for (uint m = 0u; m < 2u; ++m)
        {
            const float16_t q0 = CL8(q[m][0]), q1 = CL8(q[m][1]), q2 = CL8(q[m][2]), q3 = CL8(q[m][3]);
            const float16_t q4 = CL8(q[m][4]), q5 = CL8(q[m][5]), q6 = CL8(q[m][6]), q7 = CL8(q[m][7]);
            qa[m][0] = floate4m3_t(q0); qb[m][0] = qa[m][0]; qa[m][1] = floate4m3_t(q1); qb[m][1] = qa[m][1];
            qa[m][2] = floate4m3_t(q2); qb[m][2] = qa[m][2]; qa[m][3] = floate4m3_t(q3); qb[m][3] = qa[m][3];
            qa[m][4] = floate4m3_t(q4); qb[m][4] = qa[m][4]; qa[m][5] = floate4m3_t(q5); qb[m][5] = qa[m][5];
            qa[m][6] = floate4m3_t(q6); qb[m][6] = qa[m][6]; qa[m][7] = floate4m3_t(q7); qb[m][7] = qa[m][7];
        }
        // the bias is read column-major, so the accumulator holds S^T: this lane holds the keys 8h..8h+7 of one query
        HAC s = bias16(ob, 16u);
        K32H(s, qa[0], qb[0], qa[1], qb[1])
        float16_t e[8];
        softmax2(float16_t(s[0]), float16_t(s[1]), e[0], e[1]);
        softmax2(float16_t(s[2]), float16_t(s[3]), e[2], e[3]);
        softmax2(float16_t(s[4]), float16_t(s[5]), e[4], e[5]);
        softmax2(float16_t(s[6]), float16_t(s[7]), e[6], e[7]);
        precise f16vec2 G0_ = f16vec2(e[0], e[1]), G1_ = f16vec2(e[2], e[3]);
        precise f16vec2 G2_ = f16vec2(e[4], e[5]), G3_ = f16vec2(e[6], e[7]);
        precise f16vec2 V0_ = G0_ + UN16(SH16(G0_)), V1_ = G1_ + UN16(SH16(G1_));
        precise f16vec2 V2_ = G2_ + UN16(SH16(G2_)), V3_ = G3_ + UN16(SH16(G3_));
        precise f16vec2 VV_ = (V0_ + V1_) + (V2_ + V3_);
        const float16_t rc = float16_t(1.0 / float(VV_.x + VV_.y));
        B16 pb;
        pb[0] = G0_.x * rc; pb[1] = G0_.y * rc; pb[2] = G1_.x * rc; pb[3] = G1_.y * rc;
        pb[4] = G2_.x * rc; pb[5] = G2_.y * rc; pb[6] = G3_.x * rc; pb[7] = G3_.y * rc;
        // V through this wave's shared buffer, then O^T = V^T x P^T: rows are the head channels, columns the tokens
        ACC o[2];
        coopMatStore(v[0], VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM);
        coopMatStore(v[1], VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM);
        A16 vt0, vt1;
        coopMatLoad(vt0, VSH, (mt * 2u + 0u) * 64u, VWD / 4u, RM);
        coopMatLoad(vt1, VSH, (mt * 2u + 1u) * 64u, VWD / 4u, RM);
        o[0] = coopMatMulAdd(vt0, pb, ACC(0.0));
        o[1] = coopMatMulAdd(vt1, pb, ACC(0.0));
        R16(o[0]); R16(o[1]);
        B8 ob8[2];
        MK8_F8(ob8[0], o[0]); MK8_F8(ob8[1], o[1]);
        // the output projection of this head's 32 attention channels, straight into the residual stream
#define WO_N(N) { \
        A8 a0, a1; \
        LDA8(a0, ow, 0u, UC, 16u * (N)); \
        LDA8(a1, ow, 1u, UC, 16u * (N)); \
        K32H(res[N], a0, ob8[0], a1, ob8[1]) }
        EACH_T(WO_N)
    }
    // ------------------------------------------------------------ norm 2
    // the native element (T, 8h+i) holds the natural channel tok_xnat(T, 8h+i) = 16T + 4h + (i&3) + 8*((i>>2)&1),
    // which lives in my natural-order accumulator at the half ((i>>2)&1) and the element 4h+(i&3). The tree needs
    // the pairs in that order, so fetch the (pair 2h, pair 2h+1) of the half ((i>>2)&1) and select by this lane's half.
    f16vec2 XA_[NT], XB_[NT], XC_[NT], XD_[NT];
    for (uint t_ = 0u; t_ < NT; ++t_)
    {
        const f16vec2 p0_ = f16vec2(float16_t(res[t_][0]), float16_t(res[t_][1]));
        const f16vec2 p1_ = f16vec2(float16_t(res[t_][2]), float16_t(res[t_][3]));
        const f16vec2 p2_ = f16vec2(float16_t(res[t_][4]), float16_t(res[t_][5]));
        const f16vec2 p3_ = f16vec2(float16_t(res[t_][6]), float16_t(res[t_][7]));
        const f16vec2 o0_ = UN16(SH16(p0_)), o1_ = UN16(SH16(p1_)), o2_ = UN16(SH16(p2_)), o3_ = UN16(SH16(p3_));
        if (hf == 0u) { XA_[t_] = p0_; XB_[t_] = p1_; XC_[t_] = o0_; XD_[t_] = o1_; }
        else          { XA_[t_] = o2_; XB_[t_] = o3_; XC_[t_] = p2_; XD_[t_] = p3_; }
    }
#define EV(T, I) ((uint(I) < 2u) ? XA_[T][uint(I) & 1u] : (uint(I) < 4u) ? XB_[T][uint(I) & 1u] : \
                  (uint(I) < 6u) ? XC_[T][uint(I) & 1u] : XD_[T][uint(I) & 1u])
    float16_t rn2; RSQ_TREE(rn2)
#undef EV
#define N2(N) { \
    const uint wb = (p.o[2] + 16u * (N) + ch8) / 4u, bb = (p.o[3] + 16u * (N) + ch8) / 4u; \
    const f16vec4 g0 = w16.w[wb], g1 = w16.w[wb + 1u], d0 = w16.w[bb], d1 = w16.w[bb + 1u]; \
    const float16_t x0 = res[N][0], x1 = res[N][1], x2 = res[N][2], x3 = res[N][3]; \
    const float16_t x4 = res[N][4], x5 = res[N][5], x6 = res[N][6], x7 = res[N][7]; \
    h1[N][0] = floate4m3_t(CL8(x0 * float16_t(rn2 * g0.x))); h1[N][1] = floate4m3_t(CL8(x1 * float16_t(rn2 * g0.y))); \
    h1[N][2] = floate4m3_t(CL8(x2 * float16_t(rn2 * g0.z))); h1[N][3] = floate4m3_t(CL8(x3 * float16_t(rn2 * g0.w))); \
    h1[N][4] = floate4m3_t(CL8(x4 * float16_t(rn2 * g1.x))); h1[N][5] = floate4m3_t(CL8(x5 * float16_t(rn2 * g1.y))); \
    h1[N][6] = floate4m3_t(CL8(x6 * float16_t(rn2 * g1.z))); h1[N][7] = floate4m3_t(CL8(x7 * float16_t(rn2 * g1.w))); \
    res[N][0] = float16_t(x0 + d0.x); res[N][1] = float16_t(x1 + d0.y); \
    res[N][2] = float16_t(x2 + d0.z); res[N][3] = float16_t(x3 + d0.w); \
    res[N][4] = float16_t(x4 + d1.x); res[N][5] = float16_t(x5 + d1.y); \
    res[N][6] = float16_t(x6 + d1.z); res[N][7] = float16_t(x7 + d1.w); }
    EACH_T(N2)
    // ------------------------------------------------------------ the MLP, 32 hidden channels at a time
    for (uint c = 0u; c < UC / 8u; ++c)
    {
        const uint ow1 = p.o[20u + 3u * c], ob1 = p.o[21u + 3u * c], ow2 = p.o[22u + 3u * c];
        HAC hm[2];
        hm[0] = bias16(ob1, 32u); hm[1] = bias16(ob1 + 16u, 32u);
#define MLP_G(A, B) { \
        const uint kt = A; \
        A8 a0, a1; \
        LDA8(a0, ow1, kt, 32u, 0u); \
        LDA8(a1, ow1, kt + 1u, 32u, 0u); \
        K32H(hm[0], a0, h1[A], a1, h1[B]) \
        LDA8(a0, ow1, kt, 32u, 16u); \
        LDA8(a1, ow1, kt + 1u, 32u, 16u); \
        K32H(hm[1], a0, h1[A], a1, h1[B]) }
        EACH_T16(MLP_G)
        B8 g0, g1;
        { const float16_t v0 = gelu1(hm[0][0]), v1 = gelu1(hm[0][1]), v2 = gelu1(hm[0][2]), v3 = gelu1(hm[0][3]);
          const float16_t v4 = gelu1(hm[0][4]), v5 = gelu1(hm[0][5]), v6 = gelu1(hm[0][6]), v7 = gelu1(hm[0][7]);
          g0[0] = floate4m3_t(CL8(v0)); g0[1] = floate4m3_t(CL8(v1)); g0[2] = floate4m3_t(CL8(v2)); g0[3] = floate4m3_t(CL8(v3));
          g0[4] = floate4m3_t(CL8(v4)); g0[5] = floate4m3_t(CL8(v5)); g0[6] = floate4m3_t(CL8(v6)); g0[7] = floate4m3_t(CL8(v7)); }
        { const float16_t v0 = gelu1(hm[1][0]), v1 = gelu1(hm[1][1]), v2 = gelu1(hm[1][2]), v3 = gelu1(hm[1][3]);
          const float16_t v4 = gelu1(hm[1][4]), v5 = gelu1(hm[1][5]), v6 = gelu1(hm[1][6]), v7 = gelu1(hm[1][7]);
          g1[0] = floate4m3_t(CL8(v0)); g1[1] = floate4m3_t(CL8(v1)); g1[2] = floate4m3_t(CL8(v2)); g1[3] = floate4m3_t(CL8(v3));
          g1[4] = floate4m3_t(CL8(v4)); g1[5] = floate4m3_t(CL8(v5)); g1[6] = floate4m3_t(CL8(v6)); g1[7] = floate4m3_t(CL8(v7)); }
#define MLP2_N(N) { \
        A8 a0, a1; \
        LDA8(a0, ow2, 0u, UC, 16u * (N)); \
        LDA8(a1, ow2, 1u, UC, 16u * (N)); \
        K32H(res[N], a0, g0, a1, g1) }
        EACH_T(MLP2_N)
    }
    // ------------------------------------------------------------ the block output, FP8
    const uint dst = uint(ty * p.W + tx) * UC + ch8;
    if (inside)
        for (uint nt = 0u; nt < NT; ++nt)
        {
#define OUT8(I) out0.a[dst + 16u * nt + uint(I)] = floate4m3_t(CL8(res[nt][I]));
            OUT8(0) OUT8(1) OUT8(2) OUT8(3) OUT8(4) OUT8(5) OUT8(6) OUT8(7)
#undef OUT8
        }
    // ------------------------------------------------------------ patch merge: 2x2 tokens (4C) -> XC, half resolution
#if KIND == 1
    {
        // the 16 merged rows are gathered from the four waves: the row is the merged token of the lane's pixel
        const uint mrow = (byr >> 1) * 4u + (bxr >> 1), qq = (byr & 1u) * 2u + (bxr & 1u);
#define MG_W(N) { A8 y; MK8_F8(y, res[N]); coopMatStore(y, NB, 16u * S4 + mrow * S4 + qq * UC + 16u * (N) - row * S4, S4, RM); }
        EACH_T(MG_W)
        SYNC
        const uint MT = UXC / 16u, NI = (MT + 3u) / 4u;
        const int ym = (Y0 >> 1) + int(row >> 2), xm = (X0 >> 1) + int(row & 3u);
        const bool mok = ym >= 0 && ym < p.Hh / 2 && xm >= 0 && xm < p.W / 2;
        const uint md = uint(ym * (p.W / 2) + xm) * UXC + ch8;
        for (uint i = 0u; i < NI; ++i)
        {
            const uint u = mt + 4u * i;
            if (u < MT)
            {
                HAC y = bias16(p.o[69] + 16u * u, UXC);
                for (uint kt = 0u; kt < 4u * UC / 16u; kt += 2u)
                {
                    B8 b0, b1; coopMatLoad(b0, NB, 16u * S4 + kt * 16u, S4, CM);
                    coopMatLoad(b1, NB, 16u * S4 + (kt + 1u) * 16u, S4, CM);
                    A8 a0, a1; LDA8(a0, p.o[68], kt, UXC, 16u * u);
                    LDA8(a1, p.o[68], kt + 1u, UXC, 16u * u);
                    K32H(y, a0, b0, a1, b1)
                }
                if (mok)
                {
#define MGO(I) out1.a[md + 16u * u + uint(I)] = floate4m3_t(CL8(y[I]));
                    MGO(0) MGO(1) MGO(2) MGO(3) MGO(4) MGO(5) MGO(6) MGO(7)
#undef MGO
                }
            }
        }
    }
#endif
}