d4r: official NVIDIA DLSS on AMD GPUs under Linux and Proton

kernels/k/dltss_pwin_dec0_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_pos.h"
+POS_DECODER(dltss_pwin_dec0_layer, 2, 32, 64, 40, 48, 2)

kernels/k/dltss_pwin_dec1_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_layer.h"
+PWIN_DECODER(dltss_pwin_dec1_layer, 2, 64, 64)

kernels/k/dltss_pwin_dec2_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_layer.h"
+PWIN_DECODER(dltss_pwin_dec2_layer, 2, 64, 96)

kernels/k/dltss_pwin_dec3_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_DECODER(dltss_pwin_dec3_layer, 4, 96, 128, 2)

kernels/k/dltss_pwin_dec4_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_DECODER(dltss_pwin_dec4_layer, 4, 128, 160, 4)

kernels/k/dltss_pwin_dec5_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_PLAIN(dltss_pwin_dec5_layer, 8, 160, 4)

kernels/k/dltss_pwin_enc0_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_pos.h"
+POS_ENCODER(dltss_pwin_enc0_layer, 2, 32, 64, 16, 2)

kernels/k/dltss_pwin_enc1_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_layer.h"
+PWIN_ENCODER_POS(dltss_pwin_enc1_layer, 2, 64, 64, true)

kernels/k/dltss_pwin_enc2_layer.hip

@@ -0,0 +1,3 @@
+#define PWIN_F32ACC
+#include "pwin_layer.h"
+PWIN_ENCODER(dltss_pwin_enc2_layer, 2, 64, 96)

kernels/k/dltss_pwin_enc3_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_ENCODER(dltss_pwin_enc3_layer, 4, 96, 128, 2)

kernels/k/dltss_pwin_enc4_layer.hip

@@ -0,0 +1,2 @@
+#include "pwin_wide.h"
+WIDE_ENCODER(dltss_pwin_enc4_layer, 4, 128, 160, 4)

kernels/k/pwin_common.h

@@ -0,0 +1,188 @@
+// Native RDNA3 (gfx11 wave32 WMMA) kernels for the DLSS 4 (preset K) DltssPaddedWinLayer blocks.
+// Semantics: kernels/tools/pwin_model.py (verified against the PTX interpreter, kernels/tools/ptxsim.py).
+//
+// Layout conventions (transposed WMMA, as in the preset M kernels): a GEMM Y = X W is computed as
+// Y^T = W^T X^T with the weights as the WMMA A operand (lane l: output channel l & 15, 16 K values)
+// and the activations as the B operand (lane l: token l & 15, 16 K values). The f32 result then has
+// lane l = token l & 15 and VGPR i = output channel 2i + (l >> 4) of the 16-channel tile ("D^T").
+#pragma once
+#include <hip/hip_runtime.h>
+#include <stdint.h>
+
+#pragma clang fp contract(off)
+
+typedef _Float16 half_t;
+typedef _Float16 hv2 __attribute__((ext_vector_type(2)));
+typedef _Float16 h16 __attribute__((ext_vector_type(16)));
+typedef float f8v __attribute__((ext_vector_type(8)));
+typedef uint32_t u8v __attribute__((ext_vector_type(8)));
+typedef uint32_t u4v __attribute__((ext_vector_type(4)));
+
+__device__ __forceinline__ uint32_t lane_id()
+{
+    return __builtin_amdgcn_mbcnt_lo(~0u, 0u);
+}
+
+__device__ __forceinline__ f8v wmma(u8v a, u8v b, f8v c)
+{
+    return __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(__builtin_bit_cast(h16, a), __builtin_bit_cast(h16, b), c);
+}
+
+// one k16 step with the f16 accumulator of NVIDIA's f16 wmma (rounded after the step)
+// PWIN_F32ACC: keep the accumulator in f32 through the chain (rounded to f16 where the values are used)
+__device__ __forceinline__ f8v mma16(u8v a, u8v b, f8v c)
+{
+    f8v d = wmma(a, b, c);
+#ifndef PWIN_F32ACC
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        d[i] = (float)(half_t)d[i];
+#endif
+    return d;
+}
+
+__device__ __forceinline__ f8v splat8(float v)
+{
+    return (f8v){v, v, v, v, v, v, v, v};
+}
+
+__device__ __forceinline__ uint32_t other_half(uint32_t v)
+{
+    return __builtin_amdgcn_permlanex16(v, v, 0x76543210u, 0xfedcba98u, false, false);
+}
+
+__device__ __forceinline__ uint32_t pack2(half_t lo, half_t hi)
+{
+    return __builtin_bit_cast(uint32_t, (hv2){lo, hi});
+}
+
+__device__ __forceinline__ half_t lo16(uint32_t v)
+{
+    return __builtin_bit_cast(hv2, v)[0];
+}
+__device__ __forceinline__ half_t hi16(uint32_t v)
+{
+    return __builtin_bit_cast(hv2, v)[1];
+}
+
+// D^T f16 values (own[i] = channel 2i + hf of token l & 15) -> B operand of the same token (all 16
+// channels in order): the other parity comes from lane l ^ 16
+__device__ __forceinline__ u8v operand_from_dt(const half_t own[8])
+{
+    const uint32_t hf = lane_id() >> 4;
+    uint32_t mine[4], theirs[4];
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+    {
+        mine[j] = pack2(own[2 * j], own[2 * j + 1]);
+        theirs[j] = other_half(mine[j]);
+    }
+#ifndef PWIN_NO_PERM
+    // v_perm_b32(theirs, mine): bytes 0-3 = mine, 4-7 = theirs. Even j: the low halves, odd j: the high
+    // halves, in the order (mine, theirs) for hf = 0 and (theirs, mine) for hf = 1.
+    const uint32_t sel_lo = hf ? 0x01000504u : 0x05040100u, sel_hi = hf ? 0x03020706u : 0x07060302u;
+    u8v r;
+#pragma unroll
+    for (int j = 0; j < 8; ++j)
+        r[j] = __builtin_amdgcn_perm(theirs[j >> 1], mine[j >> 1], (j & 1) ? sel_hi : sel_lo);
+    return r;
+#else
+    u8v r;
+#pragma unroll
+    for (int j = 0; j < 8; ++j)
+    {
+        const half_t a = (j & 1) ? hi16(mine[j >> 1]) : lo16(mine[j >> 1]);
+        const half_t b = (j & 1) ? hi16(theirs[j >> 1]) : lo16(theirs[j >> 1]);
+        r[j] = hf ? pack2(b, a) : pack2(a, b);
+    }
+    return r;
+#endif
+}
+
+__device__ __forceinline__ u8v operand_from_f8(f8v d)
+{
+    half_t h[8];
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        h[i] = (half_t)d[i];
+    return operand_from_dt(h);
+}
+
+// element c (0..15) of a B operand (16 halves)
+__device__ __forceinline__ half_t op_get(const u8v& v, int c)
+{
+    return (c & 1) ? hi16(v[c >> 1]) : lo16(v[c >> 1]);
+}
+
+__device__ __forceinline__ u8v lds_row16(const half_t* p)
+{
+    const u4v a = *(const u4v*)p, b = *(const u4v*)(p + 8);
+    return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
+}
+
+// Global activation traffic (tokens in, rows out): PWIN_NT_LOAD / PWIN_NT_STORE mark it non-temporal so
+// it does not evict the weight images from the caches. Same data either way.
+__device__ __forceinline__ u8v gload_row16(const half_t* p)
+{
+#ifdef PWIN_NT_LOAD
+    const u4v a = __builtin_nontemporal_load((const u4v*)p), b = __builtin_nontemporal_load((const u4v*)(p + 8));
+#else
+    const u4v a = *(const u4v*)p, b = *(const u4v*)(p + 8);
+#endif
+    return (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
+}
+__device__ __forceinline__ void gstore16(half_t* p, u4v v)
+{
+#ifdef PWIN_NT_STORE
+    __builtin_nontemporal_store(v, (u4v*)p);
+#else
+    *(u4v*)p = v;
+#endif
+}
+__device__ __forceinline__ void gstore_row16(half_t* dst, const u8v& v)
+{
+    gstore16(dst, (u4v){v[0], v[1], v[2], v[3]});
+    gstore16(dst + 8, (u4v){v[4], v[5], v[6], v[7]});
+}
+
+__device__ __forceinline__ void block_sync()
+{
+    __syncthreads();
+}
+
+__device__ __forceinline__ int mirror(int v, int n)
+{
+    v = v < 0 ? -v : v;
+    return v < 2 * n - 2 - v ? v : 2 * n - 2 - v;
+}
+
+// ------------------------------------------------ pre-swizzled f16 weight fragments (NVIDIA)
+// logical B[k][n] of one 16x16 fragment (512 bytes)
+__device__ __forceinline__ int frag_offset(int k, int n)
+{
+    return 64 * (n & 7) + 16 * ((k & 7) >> 1) + 8 * (n >> 3) + 4 * (k >> 3) + 2 * (k & 1);
+}
+
+// A-operand image of a K x N weight: tile (kt, nt) = 16 lanes x 8 dwords (lane = output channel)
+struct WDesc
+{
+    int src;       // byte offset of the fragments in the weight buffer
+    int kt_stride; // bytes between k tiles
+    int nt_stride; // bytes between n tiles
+    int KT, NT;    // tiles
+    int dst;       // u8v index of the image in the prep buffer
+};
+
+// weight operand tile `tile` of an A-operand image (16 lanes x 32 bytes): lanes l and l + 16 need the same
+// row, so each half-wave loads one 16-byte half and the halves are swapped across (halves L2 traffic)
+__device__ __forceinline__ u8v wtile(const u8v* __restrict__ img, int tile)
+{
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const u4v mine = ((const u4v*)&img[tile * 16 + m])[hf];
+    u4v other;
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+        other[j] = other_half(mine[j]);
+    const u4v lo = hf ? other : mine, hi = hf ? mine : other;
+    return (u8v){lo[0], lo[1], lo[2], lo[3], hi[0], hi[1], hi[2], hi[3]};
+}

kernels/k/pwin_layer.h

@@ -0,0 +1,837 @@
+// Native DltssPaddedWinLayer (DLSS 4, preset K) template: 8x8 windows, H heads x 32, C channels.
+// One block per window, 4 waves (one 16-token tile each). Semantics: kernels/tools/pwin_model.py (swin_core,
+// pwin_encoder); weight layout: pwin_model.Layout.
+#pragma once
+#include "pwin_common.h"
+#ifndef PWIN_VGPR
+#define PWIN_VGPR
+#endif
+
+struct PwinParams
+{
+    int W, H;              // 0: token grid
+    const uint8_t* in;     // 8
+    const uint8_t* skip;   // 16 (decoders)
+    uint8_t* out24;        // 24: encoders: merged output; decoders: output
+    uint8_t* out32;        // 32: encoders: full-resolution output
+    uint64_t pad40, pad48; // 40
+    int sx, sy;            // 56
+    const uint8_t* w;      // 64
+    uint8_t rest[176 - 72];
+};
+static_assert(sizeof(PwinParams) == 176, "param block");
+
+template <int H_, int C_, int COUT_, bool POS_ = false> struct PwinLayout
+{
+    static constexpr bool POS = POS_; // position-only attention (trait flag 2): the table is P itself
+    static constexpr int H = H_, C = C_, COUT = COUT_, KT = C / 16, NCH = C / 32, NMLP = C / 8;
+    static constexpr int G1 = 0, QKV = 2 * C, HEAD = 192 * C, BIAS = QKV + HEAD * H, BO = BIAS + 8192 * H,
+                         WO = BO + 2 * C, G2 = WO + 64 * C * H, B2 = G2 + 2 * C, W1 = B2 + 2 * C, B1 = W1 + 8 * C * C,
+                         W2 = B1 + 8 * C, PM = W2 + 8 * C * C;
+    // patch merge: NVIDIA's warps own COUT / H channels each, allocated as whole 16-column tiles
+    static constexpr int NPW = COUT / H, NPA = (NPW + 15) / 16 * 16, PMT = H * NPA / 16, PMB = PM + 8 * C * NPA * H;
+    // prep image tiles (16 lanes of u8v each)
+    static constexpr int T_QKV = 0;                           // [h][j][kt KT][nt 2]
+    static constexpr int T_WO = T_QKV + H * 3 * KT * 2;        // [ks 2H][nt KT]
+    static constexpr int T_W1 = T_WO + 2 * H * KT;             // [c][kt KT][nt 2]
+    static constexpr int T_W2 = T_W1 + NMLP * KT * 2;          // [c][ks 2][nt KT]
+    static constexpr int T_PM = T_W2 + NMLP * 2 * KT;          // [ks C/4][nt COUT/16]
+    static constexpr int T_END = T_PM + (C / 4) * PMT;
+    static constexpr int NBIAS = H * 4 * 4 * 32;               // u4v per lane
+    static constexpr int PREP_ITEMS = T_END * 16 + NBIAS;
+
+    __device__ static int src_of(int tile)
+    {
+        if (tile < T_WO)
+        {
+            const int nt = tile & 1, kt = (tile >> 1) % KT, j = (tile / (2 * KT)) % 3, h = tile / (6 * KT);
+            return QKV + HEAD * h + 2048 * (3 * (kt >> 1) + j) + 512 * (kt & 1) + 1024 * nt;
+        }
+        if (tile < T_W1)
+        {
+            const int t = tile - T_WO, nt = t % KT, ks = t / KT;
+            return WO + 64 * C * (ks >> 1) + 512 * (ks & 1) + 1024 * nt;
+        }
+        if (tile < T_W2)
+        {
+            const int t = tile - T_W1, nt = t & 1, kt = (t >> 1) % KT, c = t / (2 * KT);
+            return W1 + 64 * C * c + 512 * kt + 32 * C * nt;
+        }
+        if (tile < T_PM)
+        {
+            const int t = tile - T_W2, nt = t % KT, ks = (t / KT) & 1, c = t / (2 * KT);
+            return W2 + 64 * C * c + 512 * ks + 1024 * nt;
+        }
+        const int t = tile - T_PM, nt = t % PMT, ks = t / PMT;
+        return PM + 512 * ks + 128 * C * nt; // allocated tiles are contiguous
+    }
+};
+
+__device__ __forceinline__ half_t wload(const uint8_t* w, int byte)
+{
+    return *(const half_t*)(w + byte);
+}
+
+// channels 16 nt + 2i + hf (i = 0..7) of an f16 vector at byte offset `byte` (16 nt already added)
+__device__ __forceinline__ void dvec8(const uint8_t* w, int byte, uint32_t hf, half_t out[8])
+{
+    const u4v a = *(const u4v*)(w + byte), b = *(const u4v*)(w + byte + 16);
+    const u8v v = (u8v){a[0], a[1], a[2], a[3], b[0], b[1], b[2], b[3]};
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+        out[i] = hf ? hi16(v[i]) : lo16(v[i]);
+}
+
+__device__ __forceinline__ u8v img_tile(const uint8_t* w, int src, int n)
+{
+    u8v r;
+#pragma unroll
+    for (int j = 0; j < 8; ++j)
+        r[j] = pack2(wload(w, src + frag_offset(2 * j, n)), wload(w, src + frag_offset(2 * j + 1, n)));
+    return r;
+}
+
+template <class L> __device__ void pwin_prep_at(const PwinParams& p, u8v* img, u4v* bias, int idx)
+{
+    const uint8_t* w = p.w;
+    if (idx < L::T_END * 16)
+    {
+        img[idx] = img_tile(w, L::src_of(idx >> 4), idx & 15);
+        return;
+    }
+    const int b = idx - L::T_END * 16;
+    if (b >= L::NBIAS)
+        return;
+    if constexpr (L::POS)
+    {
+        // B operand of P V: lane = query 16 qt + l16, 16 keys 16 kt .. 16 kt + 15 (u8v per lane; 16 lanes)
+        if (b >= L::H * 4 * 4 * 16)
+            return;
+        const int l16 = b & 15, kt = (b >> 4) & 3, qt = (b >> 6) & 3, h = b >> 8, r = l16;
+        u8v v;
+#pragma unroll
+        for (int j = 0; j < 8; ++j)
+        {
+            half_t e[2];
+#pragma unroll
+            for (int u = 0; u < 2; ++u)
+            {
+                const int c = 2 * j + u;
+                const int off = 512 * (4 * kt + qt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3);
+                e[u] = wload(w, L::BIAS + 8192 * h + off);
+            }
+            v[j] = pack2(e[0], e[1]);
+        }
+        ((u8v*)bias)[b] = v;
+        return;
+    }
+    const int lane = b & 31, kt = (b >> 5) & 3, qt = (b >> 7) & 3, h = b >> 9;
+    const int q = 16 * qt + (lane & 15);
+    half_t v[8];
+#pragma unroll
+    for (int i = 0; i < 8; ++i)
+    {
+        const int key = 16 * kt + 2 * i + (lane >> 4), r = q & 15, c = key & 15;
+        const int off = 512 * (4 * kt + qt) + 64 * (r & 7) + 16 * ((c & 7) >> 1) + 2 * (c & 1) + 4 * (r >> 3) + 8 * (c >> 3);
+        v[i] = wload(w, L::BIAS + 8192 * h + off);
+    }
+    bias[b] = (u4v){pack2(v[0], v[1]), pack2(v[2], v[3]), pack2(v[4], v[5]), pack2(v[6], v[7])};
+}
+
+template <class L> __device__ void pwin_prep(const PwinParams& p, u8v* img, u4v* bias)
+{
+    pwin_prep_at<L>(p, img, bias, blockIdx.x * 128 + threadIdx.x);
+}
+
+__device__ __forceinline__ half_t hmin(half_t a, half_t b) { return a < b ? a : b; }
+__device__ __forceinline__ half_t hmax(half_t a, half_t b) { return a > b ? a : b; }
+
+__device__ __forceinline__ half_t softmax_e(half_t s)
+{
+    const half_t hi = __builtin_bit_cast(half_t, (uint16_t)0x4D69), lo = __builtin_bit_cast(half_t, (uint16_t)0xCD69);
+    const half_t a = __builtin_bit_cast(half_t, (uint16_t)0x91D5), c = __builtin_bit_cast(half_t, (uint16_t)0xC50D);
+    const half_t t = hmax(hmin(s, hi), lo);
+    const half_t u = t * t;
+    const half_t v = __builtin_fmaf16(u, a, (half_t)1.0f);
+    const half_t x = __builtin_fmaf16(t, v, c);
+    return (half_t)__builtin_amdgcn_exp2f((float)x);
+}
+
+__device__ __forceinline__ half_t gelu(half_t x)
+{
+    const half_t k1 = (half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
+    const half_t c = hmin(hmax(x, (half_t)-2.0f), (half_t)2.0f);
+    const half_t ac = c < (half_t)0.0f ? -c : c;
+    const half_t u = k1 - k2 * ac;
+    return x * ((half_t)0.5f + c * u);
+}
+
+// gelu of 8 accumulator values as 4 packed f16 pairs (the same f16 operations as gelu(), element-wise)
+__device__ __forceinline__ void gelu8(const f8v& d, half_t g[8])
+{
+    const hv2 k1 = (hv2)(half_t)__builtin_bit_cast(float, 0x3ED306EBu), k2 = (hv2)(half_t)__builtin_bit_cast(float, 0x3DA60DD6u);
+#pragma unroll
+    for (int j = 0; j < 4; ++j)
+    {
+        const hv2 x = (hv2){(half_t)d[2 * j], (half_t)d[2 * j + 1]};
+        const hv2 c = __builtin_elementwise_min(__builtin_elementwise_max(x, (hv2)(half_t)-2.0f), (hv2)(half_t)2.0f);
+        const hv2 u = k1 - k2 * __builtin_elementwise_abs(c);
+        const hv2 r = x * ((hv2)(half_t)0.5f + c * u);
+        g[2 * j] = r[0];
+        g[2 * j + 1] = r[1];
+    }
+}
+
+// store the D^T tile of token m as 16 contiguous channels (lanes of one half-wave write)
+__device__ __forceinline__ void store_row16(half_t* dst, const u8v& v)
+{
+    *(u4v*)dst = (u4v){v[0], v[1], v[2], v[3]};
+    *(u4v*)(dst + 8) = (u4v){v[4], v[5], v[6], v[7]};
+}
+
+#ifdef PWIN_STAMP
+// per-block stage timestamps of wave 0 (s_memtime), slot k
+__device__ uint64_t g_stamp[16384][16];
+#define STAMP(k)                                                                                                   \
+    do                                                                                                             \
+    {                                                                                                              \
+        if (threadIdx.z == 0 && lane_id() == 0)                                                                    \
+        {                                                                                                          \
+            const int bid = blockIdx.y * gridDim.x + blockIdx.x;                                                   \
+            if (bid < 16384)                                                                                       \
+                g_stamp[bid][k] = __builtin_readcyclecounter();                                                    \
+        }                                                                                                          \
+    } while (0)
+#else
+#define STAMP(k) do {} while (0)
+#endif
+#ifdef PWIN_DBG
+__device__ float g_dbg[12][64][160];
+#define DBG_ON (blockIdx.x == DBX && blockIdx.y == DBY)
+#endif
+// The Swin core for one window: x0 rows of this wave's 16 tokens are in act (f16, C per token).
+// On return act holds the block output y. hb: normalised input rows; K/V of one head at a time in kl/vt.
+template <class L>
+__device__ void pwin_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
+                          half_t (*act)[L::C], half_t (*hb)[L::C], half_t (*kl)[32], half_t (*vt)[64])
+{
+    constexpr int C = L::C, KT = L::KT, H = L::H;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z;
+    const uint8_t* w = p.w;
+    half_t* arow = act[16 * wv + m];
+    half_t* hrow = hb[16 * wv + m];
+
+    STAMP(1);
+    // ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb
+    float ss = 0.0f;
+    for (int kt = 0; kt < KT; ++kt)
+    {
+        const u8v xv = lds_row16(arow + 16 * kt);
+#pragma unroll
+        for (int c = 0; c < 16; ++c)
+        {
+            const half_t v = op_get(xv, c);
+            ss += (float)(half_t)(v * v);
+        }
+    }
+    const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
+    for (int kt = (int)hf; kt < KT; kt += 2)
+    {
+        const u8v xv = lds_row16(arow + 16 * kt);
+        u8v hv;
+        const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt); // 16 contiguous gammas
+#pragma unroll
+        for (int j = 0; j < 8; ++j)
+            hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
+                          op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
+        store_row16(hrow + 16 * kt, hv);
+    }
+    __builtin_amdgcn_wave_barrier();
+#ifdef PWIN_DBG
+    if (DBG_ON && !hf)
+        for (int c = 0; c < C; ++c)
+            g_dbg[0][16 * wv + m][c] = (float)hrow[c];
+#endif
+
+    if constexpr (L::POS)
+        block_sync(); // V of all 64 tokens is computed by every wave from hb
+    STAMP(2);
+    // ---- attention, one head at a time; the output projection accumulates head by head
+    f8v acc[KT];
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+{
+
+    half_t vv_[8];
+
+    dvec8(w, L::BO + 32 * nt, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        acc[nt][i] = (float)vv_[i];
+
+}
+    for (int h = 0; h < H; ++h)
+    {
+        // position-only attention: V^T A operands straight from D = h1 Wv (lane = dim, 16 keys per tile)
+        u8v vtop[2][4];
+        if constexpr (L::POS)
+        {
+            // each wave computes the V^T operands of its own key tile; shared through LDS, one slot per head
+            __shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                f8v d = splat8(0.0f);
+                for (int ks = 0; ks < KT; ++ks)
+                    d = mma16(lds_row16(&hb[16 * wv + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
+                const u8v vo = operand_from_f8(d);
+                if (!hf)
+                    vsh[h][nt][wv][m] = vo;
+            }
+            block_sync();
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+#pragma unroll
+                for (int kt = 0; kt < 4; ++kt)
+                    vtop[nt][kt] = vsh[h][nt][kt][m];
+        }
+        u8v pop[4];
+        if constexpr (L::POS)
+        {
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+                pop[kt] = ((const u8v*)bias)[((h * 4 + wv) * 4 + kt) * 16 + m];
+        }
+        else
+        {
+        u8v qop[2];
+        // six independent chains (Q, K, V x 2 n-tiles), each over kt in order
+        f8v dq[6];
+#pragma unroll
+        for (int u = 0; u < 6; ++u)
+            dq[u] = splat8(0.0f);
+        for (int kt = 0; kt < KT; ++kt)
+        {
+            const u8v row = lds_row16(hrow + 16 * kt);
+#pragma unroll
+            for (int u = 0; u < 6; ++u)
+                dq[u] = mma16(img[(L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m], row, dq[u]);
+        }
+#pragma unroll
+        for (int j = 0; j < 3; ++j)
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                const f8v d = dq[2 * j + nt];
+                if (j == 0)
+                    qop[nt] = operand_from_f8(d);
+                else if (j == 1)
+                {
+                    const u8v kop = operand_from_f8(d);
+                    if (!hf)
+                        store_row16(&kl[16 * wv + m][16 * nt], kop);
+                }
+                else
+                {
+#pragma unroll
+                    for (int i = 0; i < 8; ++i)
+                        vt[16 * nt + 2 * i + hf][16 * wv + m] = (half_t)d[i];
+                }
+            }
+        block_sync();
+        half_t e[4][8];
+        float rs = 0.0f;
+#pragma unroll
+        for (int kt = 0; kt < 4; ++kt)
+        {
+            const u4v bi = bias[((h * 4 + wv) * 4 + kt) * 32 + l];
+            f8v d;
+#pragma unroll
+            for (int i = 0; i < 4; ++i)
+            {
+                d[2 * i] = (float)lo16(bi[i]);
+                d[2 * i + 1] = (float)hi16(bi[i]);
+            }
+#pragma unroll
+            for (int ds = 0; ds < 2; ++ds)
+                d = mma16(lds_row16(&kl[16 * kt + m][16 * ds]), qop[ds], d);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                e[kt][i] = softmax_e((half_t)d[i]);
+                rs += (float)e[kt][i];
+            }
+        }
+        rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
+        const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
+#pragma unroll
+        for (int kt = 0; kt < 4; ++kt)
+        {
+            half_t pv[8];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                pv[i] = e[kt][i] * rc;
+            pop[kt] = operand_from_dt(pv);
+        }
+        }
+        u8v oop[2];
+#pragma unroll
+        for (int nt = 0; nt < 2; ++nt)
+        {
+            f8v d = splat8(0.0f);
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+            {
+                if constexpr (L::POS)
+                    d = mma16(vtop[nt][kt], pop[kt], d);
+                else
+                    d = mma16(lds_row16(&vt[16 * nt + m][16 * kt]), pop[kt], d);
+            }
+            oop[nt] = operand_from_f8(d);
+#ifdef PWIN_DBG
+            if (DBG_ON)
+                for (int i = 0; i < 8; ++i)
+                    g_dbg[2 + h][16 * wv + m][16 * nt + 2 * i + hf] = (float)(half_t)d[i];
+#endif
+        }
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+#pragma unroll
+            for (int ks = 0; ks < 2; ++ks)
+                acc[nt] = mma16(img[(L::T_WO + (2 * h + ks) * KT + nt) * 16 + m], oop[ks], acc[nt]);
+        if constexpr (!L::POS)
+            block_sync(); // kl / vt are reused by the next head
+    }
+
+#ifdef PWIN_DBG
+    if (DBG_ON)
+        for (int nt = 0; nt < KT; ++nt)
+            for (int i = 0; i < 8; ++i)
+                g_dbg[1][16 * wv + m][16 * nt + 2 * i + hf] = acc[nt][i];
+#endif
+    STAMP(3);
+    // ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+    {
+        const u8v xv = lds_row16(arow + 16 * nt);
+        half_t mv[8], g2[8], b2[8];
+        dvec8(w, L::G2 + 32 * nt, hf, g2);
+        dvec8(w, L::B2 + 32 * nt, hf, b2);
+#pragma unroll
+        for (int i = 0; i < 8; ++i)
+        {
+            const half_t x1 = (half_t)acc[nt][i] + op_get(xv, 2 * i + hf);
+            mv[i] = x1 * g2[i];
+            acc[nt][i] = (float)(half_t)(x1 + b2[i]);
+        }
+        const u8v mo = operand_from_dt(mv);
+        if (!hf)
+            store_row16(hrow + 16 * nt, mo);
+    }
+    __builtin_amdgcn_wave_barrier();
+
+    // ---- MLP: C/8 chunks of 32 hidden
+    for (int c = 0; c < L::NMLP; ++c)
+    {
+        u8v gop[2];
+#pragma unroll
+        for (int hn = 0; hn < 2; ++hn)
+        {
+            f8v d;
+{
+
+    half_t vv_[8];
+
+    dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        d[i] = (float)vv_[i];
+
+}
+            for (int kt = 0; kt < KT; ++kt)
+                d = mma16(img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m], lds_row16(hrow + 16 * kt), d);
+            half_t g[8];
+            gelu8(d, g);
+            gop[hn] = operand_from_dt(g);
+        }
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+#pragma unroll
+            for (int ks = 0; ks < 2; ++ks)
+                acc[nt] = mma16(img[(L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m], gop[ks], acc[nt]);
+    }
+    STAMP(4);
+    // y -> act (rows of this wave)
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+    {
+        const u8v yo = operand_from_f8(acc[nt]);
+        if (!hf)
+            store_row16(arow + 16 * nt, yo);
+    }
+    __builtin_amdgcn_wave_barrier();
+}
+
+// first layer: token embedding x0 = relu(in (CIN) W_e + b_e) in front of the core (core weights at CB)
+template <int H_, int C_, int COUT_, bool POS_, int CIN_> struct PwinEmbLayout : PwinLayout<H_, C_, COUT_, POS_>
+{
+    using B = PwinLayout<H_, C_, COUT_, POS_>;
+    static constexpr int CIN = CIN_, KE = CIN / 16, CB = CIN ? 2 * C_ + 2 * CIN * C_ : 0;
+    static constexpr int T_EMB = B::T_END, T_END = T_EMB + KE * (C_ / 16);
+    static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
+};
+
+template <class L> __device__ void pwin_emb_prep(const PwinParams& p, u8v* img, u4v* bias)
+{
+    const int idx = blockIdx.x * 128 + threadIdx.x;
+    if (idx >= L::T_EMB * 16 && idx < L::T_END * 16)
+    {
+        const int t = (idx >> 4) - L::T_EMB, nt = t % (L::C / 16), ks = t / (L::C / 16);
+        img[idx] = img_tile(p.w, 2 * L::C + 512 * ks + 512 * L::KE * nt, idx & 15);
+        return;
+    }
+    PwinParams q = p;
+    q.w = p.w + L::CB;
+    // core tiles and the bias / P images (indices after T_END are shifted down to the base layout's)
+    const int base_idx = idx < L::T_EMB * 16 ? idx : idx - (L::T_END - L::T_EMB) * 16;
+    if (idx >= L::T_EMB * 16 && idx < L::T_END * 16)
+        return;
+    pwin_prep_at<typename L::B>(q, img, bias, base_idx);
+}
+
+// encoder: core + full-resolution output + 2x2 patch merge
+template <class L>
+__device__ void pwin_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    PwinParams p = p0;
+    if constexpr (L::CB != 0)
+        p.w = p0.w + L::CB;
+    constexpr int C = L::C, COUT = L::COUT;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    __shared__ __attribute__((aligned(16))) half_t kl[64][32];
+    __shared__ __attribute__((aligned(16))) half_t vt[32][64];
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
+    const int bx = blockIdx.x, by = blockIdx.y;
+
+    STAMP(0);
+    // input rows (mirror padded) -> act (first layer: embedding + ReLU)
+    {
+        const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
+        if constexpr (L::CIN != 0)
+        {
+            const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
+            u8v xe[L::KE];
+#pragma unroll
+            for (int ks = 0; ks < L::KE; ++ks)
+                xe[ks] = gload_row16(xin + 16 * ks);
+#pragma unroll
+            for (int nt = 0; nt < C / 16; ++nt)
+            {
+                f8v d;
+                {
+                    half_t vv_[8];
+                    dvec8(p0.w, 32 * nt, hf, vv_);
+#pragma unroll
+                    for (int i = 0; i < 8; ++i)
+                        d[i] = (float)vv_[i];
+                }
+#pragma unroll
+                for (int ks = 0; ks < L::KE; ++ks)
+                    d = mma16(img[(L::T_EMB + ks * (C / 16) + nt) * 16 + m], xe[ks], d);
+                half_t xo[8];
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    xo[i] = hmax((half_t)d[i], (half_t)0.0f);
+                const u8v xv = operand_from_dt(xo);
+                if (!hf)
+                    store_row16(&act[tok][16 * nt], xv);
+            }
+        }
+        else
+        {
+            const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
+            for (int kt = (int)hf; kt < L::KT; kt += 2)
+                store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
+        }
+    }
+    __builtin_amdgcn_wave_barrier();
+    pwin_core<L>(p, img, bias, act, hb, kl, vt);
+
+    STAMP(5);
+    // full-resolution output
+    const int Y0 = 8 * by - p.sy + ty, X0 = 8 * bx - p.sx + tx;
+    if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
+    {
+        half_t* yout = (half_t*)p.out32 + (size_t)(Y0 * p.W + X0) * C;
+        for (int kt = (int)hf; kt < L::KT; kt += 2)
+            gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
+    }
+    block_sync();
+
+    STAMP(6);
+    // patch merge: 16 merged tokens, K = 4 sub-tokens x C, COUT/16 n-tiles over the 4 waves
+    const int my = m >> 2, mx = m & 3;
+    const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
+    const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
+    half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * COUT;
+    const uint8_t* w = p.w;
+    for (int vt = wv; vt < L::PMT; vt += 4)
+    {
+        f8v d;
+{
+
+    half_t vv_[8];
+
+    dvec8(w, L::PMB + 32 * vt, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        d[i] = (float)vv_[i];
+
+    STAMP(7);
+}
+        for (int ks = 0; ks < C / 4; ++ks)
+        {
+            const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
+            const int src = 8 * (2 * my + dy) + 2 * mx + dx;
+            d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&act[src][16 * (ks % L::KT)]), d);
+        }
+        const u8v mo = operand_from_f8(d);
+        // allocated column 16 vt + c belongs to original warp g, local column 16 u + c
+        const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
+        half_t* dst = mout + L::NPW * g + 16 * u;
+        if (!hf && minb)
+        {
+            gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
+            if (valid >= 16)
+                gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
+        }
+    }
+}
+
+#define PWIN_ENCODER_EMB(NAME, H, C, COUT, POS, CIN)                                                               \
+    using NAME##_L = PwinEmbLayout<H, C, COUT, POS, CIN>;                                                                       \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4;                                          \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_emb_prep<NAME##_L>(p, g_img, g_bias);                                                                     \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(128) PWIN_VGPR NAME(PwinParams p)                                 \
+    {                                                                                                              \
+        pwin_encoder<NAME##_L>(p, g_img, g_bias);                                                                  \
+    }
+#define PWIN_ENCODER(NAME, H, C, COUT) PWIN_ENCODER_EMB(NAME, H, C, COUT, false, 0)
+#define PWIN_ENCODER_POS(NAME, H, C, COUT, POS) PWIN_ENCODER_EMB(NAME, H, C, COUT, POS, 0)
+
+// bottleneck (dec5): core only, output at the same resolution to +24
+template <class L>
+__device__ void pwin_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    __shared__ __attribute__((aligned(16))) half_t kl[64][32];
+    __shared__ __attribute__((aligned(16))) half_t vt[32][64];
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, tok = 16 * wv + m, ty = tok >> 3, tx = tok & 7;
+    const int bx = blockIdx.x, by = blockIdx.y;
+    {
+        const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
+        const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
+        for (int kt = (int)hf; kt < L::KT; kt += 2)
+            store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
+    }
+    __builtin_amdgcn_wave_barrier();
+    pwin_core<L>(p, img, bias, act, hb, kl, vt);
+    const int Y0 = 8 * by - p.sy + ty, X0 = 8 * bx - p.sx + tx;
+    if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
+    {
+        half_t* yout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * C;
+        for (int kt = (int)hf; kt < L::KT; kt += 2)
+            gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
+    }
+}
+
+#define PWIN_KERNEL(NAME, H, C, COUT, BODY)                                                                        \
+    using NAME##_L = PwinLayout<H, C, COUT>;                                                                       \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4;                                          \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_prep<NAME##_L>(p, g_img, g_bias);                                                                     \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(128) NAME(PwinParams p)                                           \
+    {                                                                                                              \
+        BODY<NAME##_L>(p, g_img, g_bias);                                                                          \
+    }
+
+// ---------------------------------------------------------------- decoders (patch expand + skip)
+template <int H_, int C_, int CL_, bool POS_ = false, int NOUT_ = 0, int NOUTA_ = 0>
+struct PwinDecLayout : PwinLayout<H_, C_, 16, POS_>
+{
+    using B = PwinLayout<H_, C_, 16, POS_>;
+    // last layer: output head out = y W_h + b_h (NOUTA computed, NOUT stored); b_h at B::PM, W_h after it
+    static constexpr int NOUT = NOUT_, NOUTA = NOUTA_, HB = B::PM, HW = B::PM + 2 * NOUTA_;
+    static constexpr int CL = CL_, KL = CL / 16, EXPN = 4 * C_ / 16; // expand: K CL, N 4C
+    static constexpr int EXPB = 2 * CL * 4 * C_, CB = EXPB + 8 * C_;  // expand bias, core weights base
+    static constexpr int T_EXP = B::T_PM;                              // expand tiles replace the merge tiles
+    static constexpr int T_HEAD = T_EXP + KL * EXPN;
+    static constexpr int T_END = T_HEAD + B::KT * (NOUTA_ / 16);
+    static constexpr int PREP_ITEMS = T_END * 16 + B::NBIAS;
+};
+
+template <class L> __device__ void pwin_dec_prep(const PwinParams& p, u8v* img, u4v* bias)
+{
+    const int idx = blockIdx.x * 128 + threadIdx.x;
+    if (idx >= L::T_EXP * 16 && idx < L::T_HEAD * 16)
+    {
+        const int t = (idx >> 4) - L::T_EXP, nt = t % L::EXPN, ks = t / L::EXPN;
+        img[idx] = img_tile(p.w, 512 * ks + 32 * L::CL * nt, idx & 15);
+        return;
+    }
+    if (idx >= L::T_HEAD * 16 && idx < L::T_END * 16)
+    {
+        const int t = (idx >> 4) - L::T_HEAD, nt = t % (L::NOUTA / 16), ks = t / (L::NOUTA / 16);
+        img[idx] = img_tile(p.w + L::CB, L::HW + 512 * ks + 512 * L::KT * nt, idx & 15);
+        return;
+    }
+    PwinParams q = p;
+    q.w = p.w + L::CB;
+    // core tiles lie below T_EXP; bias / P images follow the base layout's (unused) merge tiles
+    pwin_prep_at<typename L::B>(q, img, bias, idx < L::T_EXP * 16 ? idx : L::B::T_END * 16 + (idx - L::T_END * 16));
+}
+
+template <class L>
+__device__ void pwin_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    constexpr int KVR = L::POS ? 1 : 64, KVC = L::POS ? 1 : 32;
+    __shared__ __attribute__((aligned(16))) half_t kl[KVR][32];
+    __shared__ __attribute__((aligned(16))) half_t vt[KVC][64];
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
+
+    // ---- patch expand: wave wv produces sub-token q = wv (dy, dx) of all 16 low-resolution tokens
+    {
+        const int W2 = p.W / 2, H2 = p.H / 2;
+        const int ly = m >> 2, lx = m & 3;
+        const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
+        const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
+        const int dy = wv >> 1, dx = wv & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
+        const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
+        const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
+        for (int nt = 0; nt < C / 16; ++nt)
+        {
+            const int gnt = wv * (C / 16) + nt; // expand output column tile
+            f8v d;
+{
+
+    half_t vv_[8];
+
+    dvec8(p.w, L::EXPB + 32 * gnt, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        d[i] = (float)vv_[i];
+
+}
+            for (int ks = 0; ks < L::KL; ++ks)
+                d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], gload_row16(xl + 16 * ks), d);
+            const u8v sv = gload_row16(sk + 16 * nt);
+            half_t xo[8];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                xo[i] = (half_t)d[i] + op_get(sv, 2 * i + hf);
+            const u8v xv = operand_from_dt(xo);
+            if (!hf)
+                store_row16(&act[t][16 * nt], xv);
+        }
+    }
+    block_sync();
+    PwinParams q = p;
+    q.w = p.w + L::CB;
+    pwin_core<L>(q, img, bias, act, hb, kl, vt);
+
+    const int tok = 16 * wv + m, Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
+    const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
+    if constexpr (L::NOUTA != 0)
+    {
+        half_t* hout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * L::NOUT;
+#pragma unroll
+        for (int nt = 0; nt < L::NOUTA / 16; ++nt)
+        {
+            f8v d;
+{
+
+    half_t vv_[8];
+
+    dvec8(q.w, L::HB + 32 * nt, hf, vv_);
+
+#pragma unroll
+
+    for (int i = 0; i < 8; ++i)
+
+        d[i] = (float)vv_[i];
+
+}
+#pragma unroll
+            for (int kt = 0; kt < L::KT; ++kt)
+                d = mma16(img[(L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m], lds_row16(&act[tok][16 * kt]), d);
+            const u8v ho = operand_from_f8(d);
+            const int valid = L::NOUT - 16 * nt;
+            if (!hf && inb)
+            {
+                gstore16(hout + 16 * nt, (u4v){ho[0], ho[1], ho[2], ho[3]});
+                if (valid >= 16)
+                    gstore16(hout + 16 * nt + 8, (u4v){ho[4], ho[5], ho[6], ho[7]});
+            }
+        }
+    }
+    else if (inb)
+    {
+        half_t* yout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * C;
+        for (int kt = (int)hf; kt < L::KT; kt += 2)
+            gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
+    }
+}
+
+#define PWIN_DECODER_EX(NAME, H, C, CL, POS, NOUT, NOUTA)                                                          \
+    using NAME##_L = PwinDecLayout<H, C, CL, POS, NOUT, NOUTA>;                                                                      \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4;                                          \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_dec_prep<NAME##_L>(p, g_img, g_bias);                                                                 \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(128) PWIN_VGPR NAME(PwinParams p)                                 \
+    {                                                                                                              \
+        pwin_decoder<NAME##_L>(p, g_img, g_bias);                                                                  \
+    }
+
+#define PWIN_DECODER(NAME, H, C, CL) PWIN_DECODER_EX(NAME, H, C, CL, false, 0, 0)

kernels/k/pwin_pos.h

@@ -0,0 +1,433 @@
+// Position-only-attention layers (trait flag 2: enc0, enc1, dec0) with MT token tiles per wave, so every
+// weight operand loaded from L2 is used MT times (4 / MT waves per window). Same math as pwin_core.
+#pragma once
+#include "pwin_layer.h"
+
+template <class L, int MT>
+__device__ void pos_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
+                         half_t (*act)[L::C], half_t (*hb)[L::C])
+{
+    constexpr int C = L::C, KT = L::KT, H = L::H, NW = 4 / MT;
+    static_assert(L::POS, "position-only attention layers");
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z;
+    const uint8_t* w = p.w;
+
+    // ---- L2 norm of every tile of this wave -> hb
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+    {
+        const half_t* arow = act[16 * (wv * MT + mi) + m];
+        half_t* hrow = hb[16 * (wv * MT + mi) + m];
+        float ss = 0.0f;
+#pragma unroll
+        for (int kt = 0; kt < KT; ++kt)
+        {
+            const u8v xv = lds_row16(arow + 16 * kt);
+#pragma unroll
+            for (int c = 0; c < 16; ++c)
+            {
+                const half_t v = op_get(xv, c);
+                ss += (float)(half_t)(v * v);
+            }
+        }
+        const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
+        for (int kt = (int)hf; kt < KT; kt += 2)
+        {
+            const u8v xv = lds_row16(arow + 16 * kt);
+            const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
+            u8v hv;
+#pragma unroll
+            for (int j = 0; j < 8; ++j)
+                hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
+                              op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
+            store_row16(hrow + 16 * kt, hv);
+        }
+    }
+    block_sync(); // V of all 64 tokens is computed from hb
+    STAMP(2);
+
+    // ---- attention: O = P V per head (P from the table), output projection accumulated head by head
+    f8v acc[MT][KT];
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+    {
+        half_t vv[8];
+        dvec8(w, L::BO + 32 * nt, hf, vv);
+#pragma unroll
+        for (int mi = 0; mi < MT; ++mi)
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                acc[mi][nt][i] = (float)vv[i];
+    }
+#ifdef POS_VSHARE
+    // V^T operands of every head: each wave computes those of its own key tiles, shared through LDS
+    // (one slot per head, so a head's slot is never rewritten while it is read)
+    __shared__ __attribute__((aligned(16))) u8v vsh[H][2][4][16];
+#endif
+    for (int h = 0; h < H; ++h)
+    {
+        u8v vtop[2][4];
+#ifdef POS_VSHARE
+#pragma unroll
+        for (int nt = 0; nt < 2; ++nt)
+#pragma unroll
+            for (int mi = 0; mi < MT; ++mi)
+            {
+                const int kt = wv * MT + mi;
+                f8v d = splat8(0.0f);
+#pragma unroll
+                for (int ks = 0; ks < KT; ++ks)
+                    d = mma16(lds_row16(&hb[16 * kt + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
+                const u8v vo = operand_from_f8(d);
+                if (!hf)
+                    vsh[h][nt][kt][m] = vo;
+            }
+        block_sync();
+#pragma unroll
+        for (int nt = 0; nt < 2; ++nt)
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+                vtop[nt][kt] = vsh[h][nt][kt][m];
+#else
+#pragma unroll
+        for (int nt = 0; nt < 2; ++nt)
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+            {
+                f8v d = splat8(0.0f);
+#pragma unroll
+                for (int ks = 0; ks < KT; ++ks)
+                    d = mma16(lds_row16(&hb[16 * kt + m][16 * ks]), img[(L::T_QKV + ((h * 3 + 2) * KT + ks) * 2 + nt) * 16 + m], d);
+                vtop[nt][kt] = operand_from_f8(d);
+            }
+#endif
+        u8v oop[MT][2];
+#pragma unroll
+        for (int mi = 0; mi < MT; ++mi)
+        {
+            const int qt = wv * MT + mi;
+            u8v pop[4];
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+                pop[kt] = ((const u8v*)bias)[((h * 4 + qt) * 4 + kt) * 16 + m];
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                f8v d = splat8(0.0f);
+#pragma unroll
+                for (int kt = 0; kt < 4; ++kt)
+                    d = mma16(vtop[nt][kt], pop[kt], d);
+                oop[mi][nt] = operand_from_f8(d);
+            }
+        }
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+#pragma unroll
+            for (int ks = 0; ks < 2; ++ks)
+            {
+                const u8v wo = img[(L::T_WO + (2 * h + ks) * KT + nt) * 16 + m];
+#pragma unroll
+                for (int mi = 0; mi < MT; ++mi)
+                    acc[mi][nt] = mma16(wo, oop[mi][ks], acc[mi][nt]);
+            }
+    }
+
+    STAMP(3);
+    // ---- residual, MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
+#pragma unroll
+    for (int nt = 0; nt < KT; ++nt)
+    {
+        half_t g2[8], b2[8];
+        dvec8(w, L::G2 + 32 * nt, hf, g2);
+        dvec8(w, L::B2 + 32 * nt, hf, b2);
+#pragma unroll
+        for (int mi = 0; mi < MT; ++mi)
+        {
+            const int tok = 16 * (wv * MT + mi) + m;
+            const u8v xv = lds_row16(&act[tok][16 * nt]);
+            half_t mv[8];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                const half_t x1 = (half_t)acc[mi][nt][i] + op_get(xv, 2 * i + hf);
+                mv[i] = x1 * g2[i];
+                acc[mi][nt][i] = (float)(half_t)(x1 + b2[i]);
+            }
+            const u8v mo = operand_from_dt(mv);
+            if (!hf)
+                store_row16(&hb[tok][16 * nt], mo);
+        }
+    }
+    __builtin_amdgcn_wave_barrier();
+
+    STAMP(4);
+    // ---- MLP: C/8 chunks of 32 hidden; each weight operand serves the MT tiles
+    for (int c = 0; c < L::NMLP; ++c)
+    {
+        u8v gop[MT][2];
+#pragma unroll
+        for (int hn = 0; hn < 2; ++hn)
+        {
+            half_t bb[8];
+            dvec8(w, L::B1 + 64 * c + 32 * hn, hf, bb);
+            f8v d[MT];
+#pragma unroll
+            for (int mi = 0; mi < MT; ++mi)
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[mi][i] = (float)bb[i];
+#pragma unroll
+            for (int kt = 0; kt < KT; ++kt)
+            {
+                const u8v w1 = img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m];
+#pragma unroll
+                for (int mi = 0; mi < MT; ++mi)
+                    d[mi] = mma16(w1, lds_row16(&hb[16 * (wv * MT + mi) + m][16 * kt]), d[mi]);
+            }
+#pragma unroll
+            for (int mi = 0; mi < MT; ++mi)
+            {
+                half_t g[8];
+                gelu8(d[mi], g);
+                gop[mi][hn] = operand_from_dt(g);
+            }
+        }
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+#pragma unroll
+            for (int ks = 0; ks < 2; ++ks)
+            {
+                const u8v w2 = img[(L::T_W2 + (c * 2 + ks) * KT + nt) * 16 + m];
+#pragma unroll
+                for (int mi = 0; mi < MT; ++mi)
+                    acc[mi][nt] = mma16(w2, gop[mi][ks], acc[mi][nt]);
+            }
+    }
+    STAMP(5);
+    // y -> act
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+#pragma unroll
+        for (int nt = 0; nt < KT; ++nt)
+        {
+            const u8v yo = operand_from_f8(acc[mi][nt]);
+            if (!hf)
+                store_row16(&act[16 * (wv * MT + mi) + m][16 * nt], yo);
+        }
+    block_sync();
+}
+
+// encoder (enc0 with embedding, enc1): core + full-resolution output + patch merge
+template <class L, int MT>
+__device__ void pos_encoder(const PwinParams& p0, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C, NW = 4 / MT;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    PwinParams p = p0;
+    if constexpr (L::CB != 0)
+        p.w = p0.w + L::CB;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
+    STAMP(0);
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+    {
+        const int tok = 16 * (wv * MT + mi) + m, ty = tok >> 3, tx = tok & 7;
+        const int X = mirror(8 * bx - p.sx + tx, p.W), Y = mirror(8 * by - p.sy + ty, p.H);
+        if constexpr (L::CIN != 0)
+        {
+            const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * L::CIN;
+            u8v xe[L::KE];
+#pragma unroll
+            for (int ks = 0; ks < L::KE; ++ks)
+                xe[ks] = gload_row16(xin + 16 * ks);
+#pragma unroll
+            for (int nt = 0; nt < C / 16; ++nt)
+            {
+                half_t vv[8];
+                dvec8(p0.w, 32 * nt, hf, vv);
+                f8v d;
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[i] = (float)vv[i];
+#pragma unroll
+                for (int ks = 0; ks < L::KE; ++ks)
+                    d = mma16(img[(L::T_EMB + ks * (C / 16) + nt) * 16 + m], xe[ks], d);
+                half_t xo[8];
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    xo[i] = hmax((half_t)d[i], (half_t)0.0f);
+                const u8v xv = operand_from_dt(xo);
+                if (!hf)
+                    store_row16(&act[tok][16 * nt], xv);
+            }
+        }
+        else
+        {
+            const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
+            for (int kt = (int)hf; kt < L::KT; kt += 2)
+                store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
+        }
+    }
+    __builtin_amdgcn_wave_barrier();
+    STAMP(1);
+    pos_core<L, MT>(p, img, bias, act, hb);
+    STAMP(6);
+
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+    {
+        const int tok = 16 * (wv * MT + mi) + m;
+        const int Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
+        if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
+        {
+            half_t* yout = (half_t*)p.out32 + (size_t)(Y0 * p.W + X0) * C;
+            for (int kt = (int)hf; kt < L::KT; kt += 2)
+                gstore_row16(yout + 16 * kt, lds_row16(&act[tok][16 * kt]));
+        }
+    }
+    STAMP(7);
+    const int my = m >> 2, mx = m & 3;
+    const int Ym = (8 * by - p.sy) / 2 + my, Xm = (8 * bx - p.sx) / 2 + mx;
+    const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
+    half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * L::COUT;
+    for (int vt = wv; vt < L::PMT; vt += NW)
+    {
+        half_t vv[8];
+        dvec8(p.w, L::PMB + 32 * vt, hf, vv);
+        f8v d;
+#pragma unroll
+        for (int i = 0; i < 8; ++i)
+            d[i] = (float)vv[i];
+        for (int ks = 0; ks < C / 4; ++ks)
+        {
+            const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
+            const int src = 8 * (2 * my + dy) + 2 * mx + dx;
+            d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&act[src][16 * (ks % L::KT)]), d);
+        }
+        const u8v mo = operand_from_f8(d);
+        const int g = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
+        half_t* dst = mout + L::NPW * g + 16 * u;
+        if (!hf && minb)
+        {
+            gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
+            if (valid >= 16)
+                gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
+        }
+    }
+    STAMP(8);
+}
+
+// dec0: patch expand + skip, core, output head
+template <class L, int MT>
+__device__ void pos_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C, NW = 4 / MT;
+    __shared__ __attribute__((aligned(16))) half_t act[64][C];
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, bx = blockIdx.x, by = blockIdx.y;
+    {
+        const int W2 = p.W / 2, H2 = p.H / 2, ly = m >> 2, lx = m & 3;
+        const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
+        const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
+        u8v xls[L::KL];
+#pragma unroll
+        for (int ks = 0; ks < L::KL; ++ks)
+            xls[ks] = gload_row16(xl + 16 * ks);
+        for (int q = wv; q < 4; q += NW)
+        {
+            const int dy = q >> 1, dx = q & 1, t = 8 * (2 * ly + dy) + 2 * lx + dx;
+            const int SX = mirror(8 * bx - p.sx + (t & 7), p.W), SY = mirror(8 * by - p.sy + (t >> 3), p.H);
+            const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
+#pragma unroll
+            for (int nt = 0; nt < C / 16; ++nt)
+            {
+                const int gnt = q * (C / 16) + nt;
+                half_t vv[8];
+                dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
+                f8v d;
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[i] = (float)vv[i];
+#pragma unroll
+                for (int ks = 0; ks < L::KL; ++ks)
+                    d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], xls[ks], d);
+                const u8v sv = gload_row16(sk + 16 * nt);
+                half_t xo[8];
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    xo[i] = (half_t)d[i] + op_get(sv, 2 * i + hf);
+                const u8v xv = operand_from_dt(xo);
+                if (!hf)
+                    store_row16(&act[t][16 * nt], xv);
+            }
+        }
+    }
+    block_sync();
+    PwinParams q = p;
+    q.w = p.w + L::CB;
+    pos_core<L, MT>(q, img, bias, act, hb);
+#pragma unroll
+    for (int mi = 0; mi < MT; ++mi)
+    {
+        const int tok = 16 * (wv * MT + mi) + m, Y0 = 8 * by - p.sy + (tok >> 3), X0 = 8 * bx - p.sx + (tok & 7);
+        const bool inb = Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W;
+        half_t* hout = (half_t*)p.out24 + (size_t)(Y0 * p.W + X0) * L::NOUT;
+#pragma unroll
+        for (int nt = 0; nt < L::NOUTA / 16; ++nt)
+        {
+            half_t vv[8];
+            dvec8(q.w, L::HB + 32 * nt, hf, vv);
+            f8v d;
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                d[i] = (float)vv[i];
+#pragma unroll
+            for (int kt = 0; kt < L::KT; ++kt)
+                d = mma16(img[(L::T_HEAD + kt * (L::NOUTA / 16) + nt) * 16 + m], lds_row16(&act[tok][16 * kt]), d);
+            const u8v ho = operand_from_f8(d);
+            const int valid = L::NOUT - 16 * nt;
+            if (!hf && inb)
+            {
+                gstore16(hout + 16 * nt, (u4v){ho[0], ho[1], ho[2], ho[3]});
+                if (valid >= 16)
+                    gstore16(hout + 16 * nt + 8, (u4v){ho[4], ho[5], ho[6], ho[7]});
+            }
+        }
+    }
+}
+
+#define POS_ENCODER(NAME, H, C, COUT, CIN, MT)                                                                     \
+    using NAME##_L = PwinEmbLayout<H, C, COUT, true, CIN>;                                                         \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT;                                     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_emb_prep<NAME##_L>(p, g_img, g_bias);                                                                 \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p)                                 \
+    {                                                                                                              \
+        pos_encoder<NAME##_L, MT>(p, g_img, g_bias);                                                               \
+    }
+
+#define POS_DECODER(NAME, H, C, CL, NOUT, NOUTA, MT)                                                               \
+    using NAME##_L = PwinDecLayout<H, C, CL, true, NOUT, NOUTA>;                                                   \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 / MT;                                     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        pwin_dec_prep<NAME##_L>(p, g_img, g_bias);                                                                 \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(32 * 4 / MT) PWIN_VGPR NAME(PwinParams p)                                 \
+    {                                                                                                              \
+        pos_decoder<NAME##_L, MT>(p, g_img, g_bias);                                                               \
+    }

kernels/k/pwin_wide.h

@@ -0,0 +1,455 @@
+// Wide variant of the DltssPaddedWinLayer template for the deep, latency-bound levels (few windows):
+// 4 * NG waves per window. Wave wv = threadIdx.z works on token tile t = wv & 3 and channel group
+// g = wv >> 2: heads g, g + NG, ... (one head per group and round), output tiles nt = g + NG j of the
+// output projection and MLP, and MLP hidden chunks c = NG rr + g. Every output element goes through
+// the same f16 operations in the same order as pwin_layer.h (bit-identical results).
+//
+// Shared memory: hb (64 x C) = normalised input, then MLP input, then the block output y; R = union of
+// x0 (64 x C), the K / V images of the NG heads of an attention round, their outputs O, and the double
+// buffered MLP hidden chunks of a round.
+#pragma once
+#include "pwin_layer.h"
+
+template <class L, int NG> struct WideDims
+{
+    // NA: heads per attention round (their K / V images share R)
+    static constexpr int NA = NG < 4 ? NG : 4;
+    static constexpr int C = L::C, KT = L::KT, H = L::H, NTL = (KT + NG - 1) / NG, ROUNDS = (H + NA - 1) / NA;
+    static constexpr int X0_BYTES = 64 * C * 2, KV_BYTES = NA * 2 * 64 * 32 * 2;
+    // MLP hidden chunks: one per group and round, double-buffered while that fits
+    static constexpr int NB = NG <= 4 ? 2 : 1, MROUNDS = (L::NMLP + NG - 1) / NG, HID_BYTES = NB * NG * 64 * 32 * 2;
+    static constexpr int R_BYTES = X0_BYTES > KV_BYTES ? (X0_BYTES > HID_BYTES ? X0_BYTES : HID_BYTES)
+                                                       : (KV_BYTES > HID_BYTES ? KV_BYTES : HID_BYTES);
+    static_assert(!L::POS, "position-only attention layers use pwin_pos.h");
+};
+
+// Core for one window. On entry R holds x0 (64 x C); on return hb holds the block output y.
+template <class L, int NG>
+__device__ void wide_core(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias,
+                          half_t (*hb)[L::C], uint8_t* R)
+{
+    using D = WideDims<L, NG>;
+    constexpr int C = L::C, KT = L::KT, H = L::H, NTL = D::NTL;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
+    const uint8_t* w = p.w;
+    half_t(*act)[C] = (half_t(*)[C])R;
+    const half_t* arow = act[tok];
+    half_t* hrow = hb[tok];
+
+    // ---- L2 norm: h1 = x0 * (rsqrt(sum x0^2) * gamma1) -> hb (all groups of a token tile compute the sum)
+    float ss = 0.0f;
+    for (int kt = 0; kt < KT; ++kt)
+    {
+        const u8v xv = lds_row16(arow + 16 * kt);
+#pragma unroll
+        for (int c = 0; c < 16; ++c)
+        {
+            const half_t v = op_get(xv, c);
+            ss += (float)(half_t)(v * v);
+        }
+    }
+    const half_t r16 = (half_t)__builtin_amdgcn_rsqf((float)(half_t)ss);
+    for (int kt = 2 * g + (int)hf; kt < KT; kt += 2 * NG)
+    {
+        const u8v xv = lds_row16(arow + 16 * kt);
+        u8v hv;
+        const u8v gv = lds_row16((const half_t*)(w + L::G1) + 16 * kt);
+#pragma unroll
+        for (int j = 0; j < 8; ++j)
+            hv[j] = pack2(op_get(xv, 2 * j) * (half_t)(r16 * op_get(gv, 2 * j)),
+                          op_get(xv, 2 * j + 1) * (half_t)(r16 * op_get(gv, 2 * j + 1)));
+        store_row16(hrow + 16 * kt, hv);
+    }
+    // residual input of the owned output tiles (D^T order: channel 16 nt + 2 i + hf)
+    half_t x0[NTL][8];
+#pragma unroll
+    for (int j = 0; j < NTL; ++j)
+    {
+        const int nt = g + NG * j;
+        if (nt < KT)
+        {
+            const u8v xv = lds_row16(arow + 16 * nt);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                x0[j][i] = op_get(xv, 2 * i + hf);
+        }
+    }
+    block_sync(); // hb complete, R free
+
+    // ---- attention: NG heads per round; the output projection accumulates head by head
+    f8v acc[NTL];
+#pragma unroll
+    for (int j = 0; j < NTL; ++j)
+    {
+        const int nt = g + NG * j;
+        if (nt < KT)
+        {
+            half_t vv[8];
+            dvec8(w, L::BO + 32 * nt, hf, vv);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                acc[j][i] = (float)vv[i];
+        }
+    }
+    half_t(*kl)[64][32] = (half_t(*)[64][32])R;                   // [group][token][dim]
+    constexpr int NA = D::NA;
+    half_t(*vt)[32][64] = (half_t(*)[32][64])(R + NA * 64 * 32 * 2); // [group][dim][token]
+    half_t(*ob)[64][32] = (half_t(*)[64][32])R;                   // [group][token][dim] (after K / V)
+    for (int r = 0; r < D::ROUNDS; ++r)
+    {
+        const int h = g < NA ? r * NA + g : H;
+        u8v qop[2], oop[2];
+        if (h < H)
+        {
+            // six independent chains (Q, K, V x 2 n-tiles), each over kt in order
+            f8v dq[6];
+#pragma unroll
+            for (int u = 0; u < 6; ++u)
+                dq[u] = splat8(0.0f);
+            for (int kt = 0; kt < KT; ++kt)
+            {
+                const u8v row = lds_row16(hrow + 16 * kt);
+#pragma unroll
+                for (int u = 0; u < 6; ++u)
+                    dq[u] = mma16(img[(L::T_QKV + ((h * 3 + (u >> 1)) * KT + kt) * 2 + (u & 1)) * 16 + m], row, dq[u]);
+            }
+#pragma unroll
+            for (int j = 0; j < 3; ++j)
+#pragma unroll
+                for (int nt = 0; nt < 2; ++nt)
+                {
+                    const f8v d = dq[2 * j + nt];
+                    if (j == 0)
+                        qop[nt] = operand_from_f8(d);
+                    else if (j == 1)
+                    {
+                        const u8v kop = operand_from_f8(d);
+                        if (!hf)
+                            store_row16(&kl[g][tok][16 * nt], kop);
+                    }
+                    else
+                    {
+#pragma unroll
+                        for (int i = 0; i < 8; ++i)
+                            vt[g][16 * nt + 2 * i + hf][tok] = (half_t)d[i];
+                    }
+                }
+        }
+        block_sync();
+        if (h < H)
+        {
+            half_t e[4][8];
+            float rs = 0.0f;
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+            {
+                const u4v bi = bias[((h * 4 + t) * 4 + kt) * 32 + l];
+                f8v d;
+#pragma unroll
+                for (int i = 0; i < 4; ++i)
+                {
+                    d[2 * i] = (float)lo16(bi[i]);
+                    d[2 * i + 1] = (float)hi16(bi[i]);
+                }
+#pragma unroll
+                for (int ds = 0; ds < 2; ++ds)
+                    d = mma16(lds_row16(&kl[g][16 * kt + m][16 * ds]), qop[ds], d);
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                {
+                    e[kt][i] = softmax_e((half_t)d[i]);
+                    rs += (float)e[kt][i];
+                }
+            }
+            rs += __builtin_bit_cast(float, other_half(__builtin_bit_cast(uint32_t, rs)));
+            const half_t rc = (half_t)__builtin_amdgcn_rcpf((float)(half_t)rs);
+            u8v pop[4];
+#pragma unroll
+            for (int kt = 0; kt < 4; ++kt)
+            {
+                half_t pv[8];
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    pv[i] = e[kt][i] * rc;
+                pop[kt] = operand_from_dt(pv);
+            }
+#pragma unroll
+            for (int nt = 0; nt < 2; ++nt)
+            {
+                f8v d = splat8(0.0f);
+#pragma unroll
+                for (int kt = 0; kt < 4; ++kt)
+                    d = mma16(lds_row16(&vt[g][16 * nt + m][16 * kt]), pop[kt], d);
+                oop[nt] = operand_from_f8(d);
+            }
+        }
+        block_sync(); // K / V of the round consumed
+        if (h < H && !hf)
+        {
+            store_row16(&ob[g][tok][0], oop[0]);
+            store_row16(&ob[g][tok][16], oop[1]);
+        }
+        block_sync();
+        for (int jh = 0; jh < NA; ++jh)
+        {
+            const int hh = r * NA + jh;
+            if (hh >= H)
+                break;
+            const u8v o0 = lds_row16(&ob[jh][tok][0]), o1 = lds_row16(&ob[jh][tok][16]);
+#pragma unroll
+            for (int j = 0; j < NTL; ++j)
+            {
+                const int nt = g + NG * j;
+                if (nt < KT)
+                {
+                    acc[j] = mma16(img[(L::T_WO + (2 * hh) * KT + nt) * 16 + m], o0, acc[j]);
+                    acc[j] = mma16(img[(L::T_WO + (2 * hh + 1) * KT + nt) * 16 + m], o1, acc[j]);
+                }
+            }
+        }
+        if (r + 1 < D::ROUNDS)
+            block_sync(); // ob is overwritten by the next round's K / V
+    }
+
+    // ---- residual; MLP init acc = x1 + b2, MLP input m = x1 * g2 (-> hb)
+#pragma unroll
+    for (int j = 0; j < NTL; ++j)
+    {
+        const int nt = g + NG * j;
+        if (nt < KT)
+        {
+            half_t mv[8], g2[8], b2[8];
+            dvec8(w, L::G2 + 32 * nt, hf, g2);
+            dvec8(w, L::B2 + 32 * nt, hf, b2);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+            {
+                const half_t x1 = (half_t)acc[j][i] + x0[j][i];
+                mv[i] = x1 * g2[i];
+                acc[j][i] = (float)(half_t)(x1 + b2[i]);
+            }
+            const u8v mo = operand_from_dt(mv);
+            if (!hf)
+                store_row16(hrow + 16 * nt, mo);
+        }
+    }
+    block_sync();
+
+    // ---- MLP: rounds of NG chunks of 32 hidden (double-buffered in R)
+    half_t(*hid)[NG][64][32] = (half_t(*)[NG][64][32])R;
+    for (int rr = 0; rr < D::MROUNDS; ++rr)
+    {
+        const int c = rr * NG + g;
+#pragma unroll
+        for (int hn = 0; hn < 2; ++hn)
+        {
+            if (c >= L::NMLP)
+                break;
+            f8v d;
+            {
+                half_t vv[8];
+                dvec8(w, L::B1 + 64 * c + 32 * hn, hf, vv);
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[i] = (float)vv[i];
+            }
+            for (int kt = 0; kt < KT; ++kt)
+                d = mma16(img[(L::T_W1 + (c * KT + kt) * 2 + hn) * 16 + m], lds_row16(hrow + 16 * kt), d);
+            half_t gg[8];
+            gelu8(d, gg);
+            const u8v gop = operand_from_dt(gg);
+            if (!hf)
+                store_row16(&hid[rr % D::NB][g][tok][16 * hn], gop);
+        }
+        block_sync();
+        for (int jc = 0; jc < NG; ++jc)
+        {
+            const int cc = rr * NG + jc;
+            if (cc >= L::NMLP)
+                break;
+            const u8v h0 = lds_row16(&hid[rr % D::NB][jc][tok][0]), h1 = lds_row16(&hid[rr % D::NB][jc][tok][16]);
+#pragma unroll
+            for (int j = 0; j < NTL; ++j)
+            {
+                const int nt = g + NG * j;
+                if (nt < KT)
+                {
+                    acc[j] = mma16(img[(L::T_W2 + (cc * 2) * KT + nt) * 16 + m], h0, acc[j]);
+                    acc[j] = mma16(img[(L::T_W2 + (cc * 2 + 1) * KT + nt) * 16 + m], h1, acc[j]);
+                }
+            }
+        }
+        if (D::NB == 1 && rr + 1 < D::MROUNDS)
+            block_sync();
+    }
+    // y -> hb (the last reads of hb were the W1 operands before the final barrier)
+#pragma unroll
+    for (int j = 0; j < NTL; ++j)
+    {
+        const int nt = g + NG * j;
+        if (nt < KT)
+        {
+            const u8v yo = operand_from_f8(acc[j]);
+            if (!hf)
+                store_row16(hrow + 16 * nt, yo);
+        }
+    }
+    block_sync();
+}
+
+// x0 rows of a window (mirror padded) -> R
+template <class L, int NG> __device__ void wide_load_x0(const PwinParams& p, uint8_t* R)
+{
+    constexpr int C = L::C;
+    half_t(*act)[C] = (half_t(*)[C])R;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
+    const int X = mirror(8 * (int)blockIdx.x - p.sx + (tok & 7), p.W), Y = mirror(8 * (int)blockIdx.y - p.sy + (tok >> 3), p.H);
+    const half_t* xin = (const half_t*)p.in + (size_t)(Y * p.W + X) * C;
+    for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
+        store_row16(&act[tok][16 * kt], gload_row16(xin + 16 * kt));
+}
+
+// y rows of the window's own tokens -> dst (C per token)
+template <class L, int NG> __device__ void wide_store_y(const PwinParams& p, half_t (*hb)[L::C], uint8_t* dst)
+{
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, t = wv & 3, g = wv >> 2, tok = 16 * t + m;
+    const int Y0 = 8 * (int)blockIdx.y - p.sy + (tok >> 3), X0 = 8 * (int)blockIdx.x - p.sx + (tok & 7);
+    if (Y0 >= 0 && Y0 < p.H && X0 >= 0 && X0 < p.W)
+    {
+        half_t* yout = (half_t*)dst + (size_t)(Y0 * p.W + X0) * L::C;
+        for (int kt = 2 * g + (int)hf; kt < L::KT; kt += 2 * NG)
+            gstore_row16(yout + 16 * kt, lds_row16(&hb[tok][16 * kt]));
+    }
+}
+
+template <class L, int NG>
+__device__ void wide_encoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C, COUT = L::COUT;
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    __shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
+    wide_load_x0<L, NG>(p, R);
+    block_sync();
+    wide_core<L, NG>(p, img, bias, hb, R);
+    wide_store_y<L, NG>(p, hb, p.out32);
+
+    // patch merge: 16 merged tokens, K = 4 sub-tokens x C, the allocated n-tiles over the waves
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z;
+    const int my = m >> 2, mx = m & 3;
+    const int Ym = (8 * (int)blockIdx.y - p.sy) / 2 + my, Xm = (8 * (int)blockIdx.x - p.sx) / 2 + mx;
+    const bool minb = Ym >= 0 && Ym < p.H / 2 && Xm >= 0 && Xm < p.W / 2;
+    half_t* mout = (half_t*)p.out24 + (size_t)(Ym * (p.W / 2) + Xm) * COUT;
+    const uint8_t* w = p.w;
+    for (int vt = wv; vt < L::PMT; vt += 4 * NG)
+    {
+        f8v d;
+        {
+            half_t vv[8];
+            dvec8(w, L::PMB + 32 * vt, hf, vv);
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                d[i] = (float)vv[i];
+        }
+        for (int ks = 0; ks < C / 4; ++ks)
+        {
+            const int sub = ks / L::KT, dy = sub >> 1, dx = sub & 1;
+            const int src = 8 * (2 * my + dy) + 2 * mx + dx;
+            d = mma16(img[(L::T_PM + ks * L::PMT + vt) * 16 + m], lds_row16(&hb[src][16 * (ks % L::KT)]), d);
+        }
+        const u8v mo = operand_from_f8(d);
+        const int gw = vt / (L::NPA / 16), u = vt % (L::NPA / 16), valid = L::NPW - 16 * u;
+        half_t* dst = mout + L::NPW * gw + 16 * u;
+        if (!hf && minb)
+        {
+            gstore16(dst, (u4v){mo[0], mo[1], mo[2], mo[3]});
+            if (valid >= 16)
+                gstore16(dst + 8, (u4v){mo[4], mo[5], mo[6], mo[7]});
+        }
+    }
+}
+
+// bottleneck (dec5): core only, output at the same resolution to +24
+template <class L, int NG>
+__device__ void wide_plain(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    __shared__ __attribute__((aligned(16))) half_t hb[64][L::C];
+    __shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
+    wide_load_x0<L, NG>(p, R);
+    block_sync();
+    wide_core<L, NG>(p, img, bias, hb, R);
+    wide_store_y<L, NG>(p, hb, p.out24);
+}
+
+template <class L, int NG>
+__device__ void wide_decoder(const PwinParams& p, const u8v* __restrict__ img, const u4v* __restrict__ bias)
+{
+    constexpr int C = L::C;
+    __shared__ __attribute__((aligned(16))) half_t hb[64][C];
+    __shared__ __attribute__((aligned(16))) uint8_t R[WideDims<L, NG>::R_BYTES];
+    half_t(*act)[C] = (half_t(*)[C])R;
+    const uint32_t l = lane_id(), m = l & 15, hf = l >> 4;
+    const int wv = threadIdx.z, q = wv & 3, g = wv >> 2, bx = blockIdx.x, by = blockIdx.y;
+
+    // ---- patch expand: waves of sub-token q = (dy, dx) produce its tiles nt = g + NG j for the 16 low-res tokens
+    {
+        const int W2 = p.W / 2, H2 = p.H / 2;
+        const int ly = m >> 2, lx = m & 3;
+        const int LX = mirror((8 * bx - p.sx) / 2 + lx, W2), LY = mirror((8 * by - p.sy) / 2 + ly, H2);
+        const half_t* xl = (const half_t*)p.in + (size_t)(LY * W2 + LX) * L::CL;
+        const int dy = q >> 1, dx = q & 1, tt = 8 * (2 * ly + dy) + 2 * lx + dx;
+        const int SX = mirror(8 * bx - p.sx + (tt & 7), p.W), SY = mirror(8 * by - p.sy + (tt >> 3), p.H);
+        const half_t* sk = (const half_t*)p.skip + (size_t)(SY * p.W + SX) * C;
+        for (int nt = g; nt < C / 16; nt += NG)
+        {
+            const int gnt = q * (C / 16) + nt;
+            f8v d;
+            {
+                half_t vv[8];
+                dvec8(p.w, L::EXPB + 32 * gnt, hf, vv);
+#pragma unroll
+                for (int i = 0; i < 8; ++i)
+                    d[i] = (float)vv[i];
+            }
+            for (int ks = 0; ks < L::KL; ++ks)
+                d = mma16(img[(L::T_EXP + ks * L::EXPN + gnt) * 16 + m], gload_row16(xl + 16 * ks), d);
+            const u8v sv = gload_row16(sk + 16 * nt);
+            half_t xo[8];
+#pragma unroll
+            for (int i = 0; i < 8; ++i)
+                xo[i] = (half_t)d[i] + op_get(sv, 2 * i + hf);
+            const u8v xv = operand_from_dt(xo);
+            if (!hf)
+                store_row16(&act[tt][16 * nt], xv);
+        }
+    }
+    block_sync();
+    PwinParams qp = p;
+    qp.w = p.w + L::CB;
+    wide_core<L, NG>(qp, img, bias, hb, R);
+    static_assert(L::NOUTA == 0, "output head: use pwin_layer.h");
+    wide_store_y<L, NG>(p, hb, p.out24);
+}
+
+#define WIDE_KERNEL(NAME, PREP, BODY, NG, ...)                                                                  \
+    using NAME##_L = __VA_ARGS__;                                                                                     \
+    __device__ u8v g_img[NAME##_L::T_END * 16];                                                                    \
+    __device__ u4v g_bias[NAME##_L::NBIAS];                                                                        \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_blocks = (NAME##_L::PREP_ITEMS + 127) / 128;     \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_block_z = 4 * (NG);                                   \
+    extern "C" __device__ __attribute__((used)) uint32_t d4r_prep_key_offset = 64; /* prep reads only p.w */       \
+    extern "C" __global__ void __launch_bounds__(128) NAME##_prep(PwinParams p)                                    \
+    {                                                                                                              \
+        PREP<NAME##_L>(p, g_img, g_bias);                                                                          \
+    }                                                                                                              \
+    extern "C" __global__ void __launch_bounds__(128 * (NG)) PWIN_VGPR NAME(PwinParams p)                          \
+    {                                                                                                              \
+        BODY<NAME##_L, NG>(p, g_img, g_bias);                                                                      \
+    }
+
+#define WIDE_ENCODER(NAME, H, C, COUT, NG) WIDE_KERNEL(NAME, pwin_prep, wide_encoder, NG, PwinLayout<H, C, COUT>)
+#define WIDE_PLAIN(NAME, H, C, NG) WIDE_KERNEL(NAME, pwin_prep, wide_plain, NG, PwinLayout<H, C, 16>)
+#define WIDE_DECODER(NAME, H, C, CL, NG) WIDE_KERNEL(NAME, pwin_dec_prep, wide_decoder, NG, PwinDecLayout<H, C, CL>)