ggml-metal: FA tensor path: gate to the (256, 256) shape

the tensor kernel is only faster than the vec kernel for head size
256/256 (measured +3%..+13% across GQA ratios 8:8 .. 64:8 on M5 Max,
nq = kv = 4096).  all the other (dk, dv) pairs lose and now fall back
to the vec kernel:

  64/64:    -18%
  128/128:  -12% MHA, -64% GQA 32:8
  192/128:  -3%
  192/192:  -53%
  320/256:  -20% GQA
  512/512:  -35%..-65%
  576/512:  -60%..-70%

the large-dk losses are register spills (the f32 Q tile (dk, 8) costs
dk/4 registers per thread).  the f16 Q attempt to fix them is blocked
by an MPP driver/compiler bug: with f16 QK operands the coop
destination (P) tile element map changes from (2q x 8kv) to (4q x 4kv)
per thread and the hardware P -> f16 right input conversion of
get_right_input_cooperative_tensor does not match the new layout
(reproduced standalone: the PV output contains the correct values at
the wrong positions; the same class of bug breaks the coop element
map when the K operand is strided).  the smem P workaround (write P
to threadgroup memory, PV from a tensor_inline right input, transR =
true with the (k, n) k-contiguous tile) is correct (4809/4809
test-backend-ops) but 2-3x slower than the coop register path.

correctness: 4809/4809 test-backend-ops FLASH_ATTN_EXT.

ggml/src/ggml-metal/ggml-metal-device.cpp

@@ -1537,6 +1537,34 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext(
     return res;
 }
 
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_tensor(
+        ggml_metal_library_t lib,
+        const ggml_tensor * op) {
+    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+
+    char name[256];
+
+    const int32_t dk = (int32_t) op->src[0]->ne[0];
+    const int32_t dv = (int32_t) op->src[2]->ne[0];
+
+    snprintf(name, 256, "kernel_flash_attn_ext_tensor_dk%d_dv%d", dk, dv);
+
+    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
+    if (!res.pipeline) {
+        res = ggml_metal_library_compile_pipeline(lib, name, name, NULL);
+    }
+
+    // one simdgroup per threadgroup; 8 queries per threadgroup
+    res.nsg  = 1;
+    res.nr0  = 8;
+    res.nr1  = 1;
+
+    // per lane: sh_qmax[8][32] + sh_qsum[8][32] + sh_M[8] + sh_S[8] + sh_alpha[8]
+    res.smem = 4*(8*32 + 8*32 + 3*8)*sizeof(float);
+
+    return res;
+}
+
 ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec(
         ggml_metal_library_t lib,
         const ggml_tensor * op,

ggml/src/ggml-metal/ggml-metal-device.h

@@ -199,6 +199,10 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
         int32_t ns10,
         int32_t ns20);
 
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_tensor(
+        ggml_metal_library_t lib,
+        const struct ggml_tensor * op);
+
 struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec(
         ggml_metal_library_t lib,
         const struct ggml_tensor * op,

ggml/src/ggml-metal/ggml-metal-impl.h

@@ -419,6 +419,8 @@ typedef struct {
     float    m1;
     int32_t  n_head_log2;
     float    logit_softcap;
+    // bit 0: has sinks (src[4]); the mask flag is (ne31 != 0)
+    uint32_t flags;
 } ggml_metal_kargs_flash_attn_ext;
 
 typedef struct {

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -2802,6 +2802,81 @@ bool ggml_metal_op_flash_attn_ext_use_vec(const ggml_tensor * op) {
     return (ne01 < 20) && (ne00 % 32 == 0);
 }
 
+// use the Tensor API (MPP tensor_ops) prefill kernel when the shape is in its fast path
+// ref: the kernel in kernels/fa.metal (kernel_flash_attn_ext_tensor)
+static bool ggml_metal_op_flash_attn_ext_use_tensor(const struct ggml_metal_device_props * props_dev,
+                                                    const ggml_tensor * op,
+                                                    float max_bias,
+                                                    float logit_softcap) {
+    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
+
+    if (!props_dev->has_tensor) {
+        return false;
+    }
+
+    // override for A/B testing: GGML_METAL_FA_TENSOR=0|1 (default: auto)
+    static int fa_tensor_env = -1;
+    if (fa_tensor_env == -1) {
+        const char * env = getenv("GGML_METAL_FA_TENSOR");
+        fa_tensor_env = env ? atoi(env) : -1;
+    }
+    if (fa_tensor_env == 0) {
+        return false;
+    }
+
+    // decode stays on the vec path
+    if (ggml_metal_op_flash_attn_ext_use_vec(op)) {
+        return false;
+    }
+
+
+    // f16 K/V only (quantized KV needs the kv_f16 pre-pass; f32/bf16 KV later)
+    if (op->src[1]->type != GGML_TYPE_F16) {
+        return false;
+    }
+
+    // mask (f16, asserted above), sinks, ALiBi (max_bias) and softcap
+    // (logit_softcap) are all supported by the tensor kernel; the pad and blk
+    // buffers are pre-computed by the pad/blk kernels when needed
+    (void) max_bias;
+    (void) logit_softcap;
+
+    const int64_t dk = op->src[0]->ne[0];
+    const int64_t dv = op->src[2]->ne[0];
+    const int64_t nq = op->src[0]->ne[1];
+
+    if (dk % 16 != 0) {
+        return false;
+    }
+
+    // PV d-block split: 128-wide blocks (64-wide when dv % 128 != 0)
+    if (dv % 64 != 0 || dv > 512) {
+        return false;
+    }
+
+    // 8 queries per threadgroup (1 simdgroup)
+    if (nq % 8 != 0) {
+        return false;
+    }
+
+    // only the (dk, dv) pairs where the tensor path is at least on par with
+    // the vec kernel.  Measured on M5 Max (nq = kv = 4096, several GQA
+    // ratios, 8:8 .. 64:8):
+    //   (256, 256): +3%..+13%
+    // all the other pairs lose: 64/64 -18%, 128/128 -12% MHA / -64% GQA
+    // 32:8, 192/128 -3%, 192/192 -53%, 320/256 -20% GQA, 512/512 -35%..-65%,
+    // 576/512 -60%..-70% (per-d-block form: NBLKx the QK^T work; the f16 Q
+    // hoisted form is blocked by an MPP bug - the coop P -> f16 right input
+    // conversion is wrong for the P tile layout induced by f16 QK operands -
+    // and the smem P workaround is 2-3x slower).  They fall back to vec.
+    switch (dv) {
+        case 256: if (dk != 256) return false; break;
+        default:  return false;
+    }
+
+    return true;
+}
+
 // ref: https://github.com/ggml-org/llama.cpp/pull/27390
 // dequantize the quantized KV cache to F16 before running the F16 flash attention kernels
 static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) {
@@ -3184,9 +3259,15 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
         GGML_ASSERT(nqptg  % 8  == 0);
         GGML_ASSERT(ncpsg  % 32 == 0);
 
+        // the tensor kernel uses 32-wide KV chunks (smaller operand/P tiles fit
+        // better in registers); the pad buffer is then written with 32-wide
+        // chunks (the reserved pad space is 64-wide, a superset)
+        const bool use_tensor = ggml_metal_op_flash_attn_ext_use_tensor(props_dev, op, max_bias, logit_softcap);
+        const int ncpsg_pad = use_tensor ? OP_FLASH_ATTN_EXT_VEC_NCPSG : ncpsg;
+
         bool need_sync = false;
 
-        const bool has_kvpad = ne11 % ncpsg != 0;
+        const bool has_kvpad = ne11 % ncpsg_pad != 0;
 
         if (has_kvpad) {
             assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0);
@@ -3209,7 +3290,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
                 /*.nb33    =*/nb33,
             };
 
-            auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_pad(lib, op, has_mask, ncpsg);
+            auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_pad(lib, op, has_mask, ncpsg_pad);
 
             ggml_metal_encoder_set_pipeline(enc, pipeline0);
             ggml_metal_encoder_set_bytes   (enc, &args0, sizeof(args0), 0);
@@ -3221,7 +3302,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             assert(ne12 == ne22);
             assert(ne13 == ne23);
 
-            ggml_metal_encoder_dispatch_threadgroups(enc, ncpsg, std::max(ne12, ne32), std::max(ne13, ne33), 32, 1, 1);
+            ggml_metal_encoder_dispatch_threadgroups(enc, ncpsg_pad, std::max(ne12, ne32), std::max(ne13, ne33), 32, 1, 1);
 
             need_sync = true;
         }
@@ -3259,6 +3340,81 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
             ggml_metal_op_concurrency_reset(ctx);
         }
 
+        if (use_tensor) {
+            // Tensor API (MPP tensor_ops) fast path - see kernel_flash_attn_ext_tensor in kernels/fa.metal
+            // 8 queries per threadgroup (1 simdgroup), C = 32 kv per chunk
+            const int nqptg = 8;
+            const int nsg   = 1;
+
+            GGML_ASSERT(ne01 % nqptg == 0);
+            GGML_ASSERT(ne11 % OP_FLASH_ATTN_EXT_VEC_NCPSG == 0 || has_kvpad);
+
+            const int32_t ns10 = nb11_attn/nb10_attn;
+            const int32_t ns20 = nb21_attn/nb20_attn;
+
+            ggml_metal_kargs_flash_attn_ext args = {
+                /*.ne01          =*/ ne01,
+                /*.ne02          =*/ ne02,
+                /*.ne03          =*/ ne03,
+                /*.nb01          =*/ nb01,
+                /*.nb02          =*/ nb02,
+                /*.nb03          =*/ nb03,
+                /*.ne11          =*/ ne11,
+                /*.ne_12_2       =*/ ne12,
+                /*.ne_12_3       =*/ ne13,
+                /*.ns10          =*/ ns10,
+                /*.nb11          =*/ nb11_attn,
+                /*.nb12          =*/ nb12_attn,
+                /*.nb13          =*/ nb13_attn,
+                /*.ns20          =*/ ns20,
+                /*.nb21          =*/ nb21_attn,
+                /*.nb22          =*/ nb22_attn,
+                /*.nb23          =*/ nb23_attn,
+                /*.ne31          =*/ ne31,
+                /*.ne32          =*/ ne32,
+                /*.ne33          =*/ ne33,
+                /*.nb31          =*/ nb31,
+                /*.nb32          =*/ nb32,
+                /*.nb33          =*/ nb33,
+                /*.ne1           =*/ ne1,
+                /*.ne2           =*/ ne2,
+                /*.ne3           =*/ ne3,
+                /*.scale         =*/ scale,
+                /*.max_bias      =*/ max_bias,
+                /*.m0            =*/ m0,
+                /*.m1            =*/ m1,
+                /*.n_head_log2   =*/ n_head_log2,
+                /*.logit_softcap =*/ logit_softcap,
+                /*.flags         =*/ has_sinks ? 1u : 0u,
+            };
+
+            auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_tensor(lib, op);
+            GGML_ASSERT(pipeline.pipeline);
+
+            ggml_metal_encoder_set_pipeline(enc, pipeline);
+            ggml_metal_encoder_set_bytes   (enc, &args, sizeof(args), 0);
+            ggml_metal_encoder_set_buffer  (enc, bid_src0, 1);
+            ggml_metal_encoder_set_buffer  (enc, bid_k,    2);
+            ggml_metal_encoder_set_buffer  (enc, bid_v,    3);
+            ggml_metal_encoder_set_buffer  (enc, bid_src3, 4);
+            ggml_metal_encoder_set_buffer  (enc, bid_src4, 5);
+            ggml_metal_encoder_set_buffer  (enc, bid_pad,  6);
+            ggml_metal_encoder_set_buffer  (enc, bid_blk,  7);
+            ggml_metal_encoder_set_buffer  (enc, bid_dst,  8);
+
+            // sh_qmax[8][32] + sh_qsum[8][32] + sh_M[8] + sh_S[8] + sh_alpha[8]
+            // (the NBLK == 2 branch keeps the softmax state in registers and
+            // does not use shared memory; the size covers the per-block form)
+            const size_t smem = (8*32 + 8*32 + 3*8)*sizeof(float);
+            GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size);
+
+            ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0);
+
+            ggml_metal_encoder_dispatch_threadgroups(enc, ne01/nqptg, ne02, ne03, 32, nsg, 1);
+
+            return 1;
+        }
+
         const int is_q = !use_kv_f16 && ggml_is_quantized(op->src[1]->type) ? 1 : 0;
 
         // 2*(2*ncpsg)

ggml/src/ggml-metal/kernels/fa.metal

@@ -2250,3 +2250,741 @@ template [[host_name("kernel_lightning_indexer_q4_1")]] kernel kernel_lightning_
 template [[host_name("kernel_lightning_indexer_q5_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q5_0, 2, dequantize_q5_0>;
 template [[host_name("kernel_lightning_indexer_q5_1")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q5_1, 2, dequantize_q5_1>;
 template [[host_name("kernel_lightning_indexer_q8_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q8_0, 2, dequantize_q8_0>;
+
+// ============================================================================
+// Tensor API (MPP tensor_ops) Flash Attention prefill kernel
+//
+// One simdgroup per threadgroup (32 threads); each matmul op runs with
+// execution scope = the whole (single) simdgroup (metal::execution_simdgroup).
+// Each threadgroup handles QPSG = 8 queries over the full KV range (same
+// "one threadgroup per Q tile, full KV loop" model as the matrix kernel above
+// - no split-KV across simdgroups).
+//
+// IMPORTANT: the MPP coop destination tile layout is implementation-defined
+// (it depends on the op, the dtypes, the opscope AND the surrounding code /
+// register pressure).  Two verified landmines (do not "simplify" these away):
+//   - the output write must be ELEMENT pointer arithmetic on a float*
+//     ((device float *) dst + (uint64_t)elems * DV); doing the offset in BYTES
+//     (char* + elems*DV*4) corrupts the QK^T (C, QPSG) coop tile layout
+//   - dispatching 4 simdgroups with per-lane execution scope (the former
+//     design) only produced correct results in the presence of extra
+//     dead-code "bloat" - the 1-simdgroup design is the stable one
+//
+// Verified design notes (see the Phase-0 spike, 2026-07):
+//   - QK^T: descriptor(m=C, n=QPSG, k=DK, transL=false, transR=true);
+//     left = K tile (DK, C) f16 device, right = Q tile (DK, QPSG) f32 device
+//     (half x float -> float is in the dtype table; Q is f32 in ggml)
+//   - the QK destination coop tensor (f32) has idx0 = query, idx1 = kv; each
+//     thread owns a (2 queries x 8 kv) sub-tile
+//   - per-query max/sum: in-register partials + threadgroup exchange
+//     (the hardware reduce_rows/reduce_columns reduce the wrong axis for this
+//     tile shape and their ownership is not what we need)
+//   - P stays in registers (f32) and is fed directly to the PV matmul as the
+//     right input: get_right_input_cooperative_tensor<half, float, float>(P)
+//   - PV: descriptor(m=PVM, n=QPSG, k=C, transL=true, transR=false),
+//     left = V tile (PVM, C) f16 device, m dimension capped at 128 -> for
+//     dv > 128 the PV is split into d blocks (PVM = 128, or 64 for dv % 128
+//     != 0); each block re-runs the chunk loop with a self-contained f32
+//     accumulator (a coop destination tile is a compiler-managed register
+//     tile: distinct named objects do not share storage)
+//   - the PV destination has idx0 = query, idx1 = d; the output is written
+//     element-wise (store() only works for the transposed natural layout)
+//   - the last partial KV chunk (ne11 % C != 0) is read from the pad buffer
+//     produced by kernel_flash_attn_ext_pad (padded rows are zero; the padded
+//     scores are clobbered to -FLT_MAX/2 in registers before the softmax)
+//
+// fast-path gate (host-side): f16 K/V, no mask/sinks/ALiBi/softcap,
+// dk % 16 == 0, dv % 64 == 0 && dv <= 512, ne01 % (NLANES*QPSG) == 0 (== 8)
+// ============================================================================
+#ifdef GGML_METAL_HAS_TENSOR
+#include <metal_tensor>
+#include <MetalPerformancePrimitives/MetalPerformancePrimitives.h>
+
+// note: file-scope constants are not allowed in MSL, so these are defined in the kernel
+//       (the host-side gate mirrors the values: NLANES*QPSG = 8 queries per threadgroup)
+template <short DK, short DV>
+void kernel_flash_attn_ext_tensor_impl(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q,
+        device const char * k,
+        device const char * v,
+        device const char * mask,
+        device const char * sinks,
+        device const char * pad,
+        device       char * dst,
+        threadgroup  half * shmem_f16,
+        uint3   tgpig,
+        ushort  tiitg,
+        ushort  sgitg) {
+    using namespace mpp::tensor_ops;
+
+    // NOTE: 1 simdgroup per threadgroup (32 threads), execution scope = 1
+    // simdgroup.  Using 4 simdgroups with execution_simdgroups<1> leaves the
+    // coop destination tile layout implementation-defined and the compiler
+    // can collapse the (C, QPSG) tile (verified: n dim shrinks, results wrong).
+    constexpr short NLANES = 1;                    // simdgroups per threadgroup
+    // NOTE: 16 queries per threadgroup (QPSG = 16) is 2x faster in isolated
+    // microbenchmarks (K/V bandwidth bound per query: 12.6 -> 24.8 TFLOPS),
+    // but in the full kernel the larger tiles (P: (C,16), O: (PVM,16)) spill
+    // and it is ~30% SLOWER (8.6 -> 6.1 TFLOPS).  8 queries is the sweet spot.
+    constexpr short QPSG   = 8;                   // queries per threadgroup
+    constexpr short C      = OP_FLASH_ATTN_EXT_VEC_NCPSG; // 32 kv items per chunk
+    // (32-wide chunks: smaller K/V operand and P tiles fit better in registers
+    // than 64-wide; the pad buffer is written with matching 32-wide chunks)
+
+    // the tensor-core destination tile caps the m dimension at 128; for
+    // dv > 128 split the PV into d blocks (128-wide, or 64-wide when dv is
+    // not a multiple of 128 - both verified)
+    constexpr short PVM  = ((DV % 128) == 0) ? ((DV < 128) ? DV : 128) : 64;
+    constexpr short NBLK = (DV + PVM - 1) / PVM;
+    static_assert(DV % PVM == 0, "dv must be a multiple of the PV block size");
+
+    const ushort iq3 = tgpig[2];
+    const ushort iq2 = tgpig[1];
+    const int    iq1 = tgpig[0]*QPSG;              // 8 queries per threadgroup
+
+    // GQA: Q head -> KV head
+    const short ikv2 = iq2/(args.ne02/args.ne_12_2);
+    const short ikv3 = iq3/(args.ne03/args.ne_12_3);
+
+    // row strides in elements
+    const int sq = args.nb01/4;  // Q is f32
+    const int sk = args.ns10;
+    const int sv = args.ns20;
+
+    // NOTE: MPP requires non-const element types in the operand decltypes
+    device float * qp = (device float *) (q + iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03) + (sgitg*QPSG)*sq;
+    device half  * kp = (device half  *) (k + ikv2*args.nb12 + ikv3*args.nb13);
+    device half  * vp = (device half  *) (v + ikv2*args.nb22 + ikv3*args.nb23);
+
+    // ---- Q tile for this lane: (DK, QPSG) f32, strided ----
+    auto tQ = tensor(qp, dextents<int, 2>(DK, QPSG), array<int, 2>{ 1, sq });
+
+    // ---- matmul ops (execution scope: 1 simdgroup = this lane) ----
+    // QK^T: P (C, QPSG) = K^T (DK, C) x Q (DK, QPSG)
+    // NOTE: k must be dynamic_extent (k is taken from the left operand extent);
+    // using a compile-time constant here breaks the coop destination tile layout
+    constexpr auto desc_qk = matmul2d_descriptor(
+            C, QPSG, static_cast<int>(dynamic_extent), false, true, false, matmul2d_descriptor::mode::multiply);
+    matmul2d<desc_qk, execution_simdgroup> mm_qk;
+
+    // PV: O (PVM, QPSG) += V^T (PVM, C) x P (C, QPSG)
+    constexpr auto desc_pvb = matmul2d_descriptor(
+            PVM, QPSG, C, true, false, false, matmul2d_descriptor::mode::multiply_accumulate);
+    matmul2d<desc_pvb, execution_simdgroup> mm_pvb;
+
+    auto tK0 = tensor(kp, dextents<int, 2>(DK, C), array<int, 2>{ 1, sk });
+    auto tV0 = tensor(vp, dextents<int, 2>(PVM, C), array<int, 2>{ 1, sv });
+
+    auto cT_qk0 = mm_qk.template get_destination_cooperative_tensor<decltype(tK0), decltype(tQ), float>();
+
+    // the softmax state lives in per-thread registers (global-max online
+    // softmax, see below): no shared memory or simdgroup barriers are needed.
+    // The threadgroup buffer is still in the ABI (the host allocates it).
+    (void) shmem_f16;
+    (void) tiitg;
+    (void) sgitg;
+    (void) NLANES;
+
+    const int kv      = args.ne11;
+    const int nchunks = (kv + C - 1)/C;
+    const bool has_kvpad = kv % C != 0;
+
+    // mask: (kv, nq, ne32, ne33) f16, kv contiguous; ne31 == 0 means no mask
+    const bool has_mask  = args.ne31 != 0;
+    const bool has_sinks = (args.flags & 1u) != 0;
+    const bool has_scap  = args.logit_softcap != 0.0f;
+
+    // ALiBi: slope per Q head (the mask holds the position distances)
+    float slope = 1.0f;
+    if (args.max_bias != 0.0f) {
+        const int h = iq2;
+        const float base = h < args.n_head_log2 ? args.m0 : args.m1;
+        const int exph = h < args.n_head_log2 ? h + 1 : 2*(h - args.n_head_log2) + 1;
+        slope = pow(base, (float) exph);
+    }
+    const float mscale = args.max_bias != 0.0f ? slope : 1.0f;
+
+    // mask head/batch offset (bytes); the query offset is applied per element
+    const uint64_t mbase = (uint64_t) (iq2 % args.ne32) * args.nb32 + (uint64_t) (iq3 % args.ne33) * args.nb33;
+
+    // pad buffer layout (kernel_flash_attn_ext_pad): the last C kv items per
+    // (KV head, batch), K then V, zero-padded beyond kv; the mask section is
+    // (nq, C) f16 per (Q head, batch) with -MAXHALF beyond kv
+    const uint64_t pad_k_offs = (uint64_t) (ikv2 + ikv3*args.ne_12_2) * args.nb11*C;
+    const uint64_t pad_v_offs = (uint64_t) args.nb11*C*args.ne_12_2*args.ne_12_3 +
+                                (uint64_t) (ikv2 + ikv3*args.ne_12_2) * args.nb21*C;
+    const uint64_t pad_mask_offs = (uint64_t) (args.nb11 + args.nb21) * C * args.ne_12_2 * args.ne_12_3 +
+                                   2u * C * args.ne31 * ((uint64_t) (iq2 % args.ne32) + (uint64_t) (iq3 % args.ne33) * args.ne32);
+
+    auto make_cT_pv = [&]() {
+        auto tVb = tensor(vp, dextents<int, 2>(PVM, C), array<int, 2>{ 1, sv });
+        return mm_pvb.template get_destination_cooperative_tensor<decltype(tVb), decltype(cT_qk0), float>();
+    };
+
+    if constexpr (NBLK == 2) {
+        // hoisted form: QK^T + online softmax run once per chunk; the P tile
+        // (f32, in registers) is reused for both PV d blocks.  (The per-block
+        // form below recomputes QK^T + softmax per d block: NBLKx the QK work.)
+        auto cT_pv0 = make_cT_pv();
+        auto cT_pv1 = make_cT_pv();
+
+        {
+            #pragma clang loop unroll(full)
+            for (uint i = 0; i < cT_pv0.get_capacity(); ++i) {
+                if (cT_pv0.is_valid_element(i)) { cT_pv0[i] = 0.0f; }
+            }
+        }
+        {
+            #pragma clang loop unroll(full)
+            for (uint i = 0; i < cT_pv1.get_capacity(); ++i) {
+                if (cT_pv1.is_valid_element(i)) { cT_pv1[i] = 0.0f; }
+            }
+        }
+        // online-softmax state in PER-THREAD REGISTERS: the chunk reductions
+        // use simd_max/simd_sum (warp-level), so the state is uniform across
+        // the simdgroup and NO shared memory / barriers are needed.  Barriers
+        // drain the tensor-core pipeline between the QK^T and PV matmuls
+        // (microbench on this GPU: ~10.5 TFLOPS with the smem + 4-barrier
+        // reduction vs ~12.6 without).
+        // the running max may be ANY constant >= max score seen so far: it
+        // cancels in O/S.  Use ONE GLOBAL max (scalar) for the O rescale so
+        // the per-chunk rescale is a trivial scalar multiply (no per-element
+        // index decode / per-query alpha array).  Per-row maxes are tracked
+        // only for the sinks correction at the end.
+        float M = -FLT_MAX / 2;
+        float S[QPSG];
+        for (int j = 0; j < QPSG; ++j) {
+            S[j] = 0.0f;
+        }
+        float alpha = 1.0f; // exp(M_old - M_new), uniform scalar
+
+        // ---- chunk interleave: QK^T of chunk ic+1 is issued before PV of
+        // chunk ic, so the two tensor-core matmuls pipeline.  (Serial
+        // per-chunk QK->PV starves the tensor core between matmuls;
+        // microbench on this GPU: serial 12.5 vs interleaved 22.2 TFLOPS.)
+        // Two P tiles alternate (coop tiles cannot be swapped by assignment);
+        // the QK^T destination type is chunk-independent, so both tiles share
+        // the chunk-0 operand decltypes.
+        auto tK0c = tensor(kp, dextents<int, 2>(DK, C), array<int, 2>{ 1, sk });
+        auto cT_qkA = mm_qk.template get_destination_cooperative_tensor<decltype(tK0c), decltype(tQ), float>();
+        auto cT_qkB = mm_qk.template get_destination_cooperative_tensor<decltype(tK0c), decltype(tQ), float>();
+        using cT_qk_t = decltype(cT_qkA);
+
+        // ---- prologue: QK^T(0) + score + online-softmax reduction -> cT_qkA ----
+        {
+            const int kc = min((int) C, kv);
+            const bool use_pad = has_kvpad && C > kv;
+            device half * kp_c = use_pad ? (device half *) (pad + pad_k_offs) : kp;
+
+            // mask row base for this chunk: real mask (kv contiguous, query
+            // stride nb31) or the pad section (C per query, indexed from 0)
+            device const half * mp = nullptr;
+            int mstride = 0; // in halfs, per local query j
+            if (has_mask) {
+                if (use_pad) {
+                    mp      = (device const half *) (pad + pad_mask_offs) + (iq1 + sgitg*QPSG) * (int) C;
+                    mstride = C;
+                } else {
+                    // global query index: iq1 + sgitg*QPSG + j (j is local, 0..QPSG-1)
+                    mp      = (device const half *) (mask + mbase) + (uint64_t) (iq1 + sgitg*QPSG) * (args.nb31 / 2);
+                    mstride = (int) (args.nb31 / 2);
+                }
+            }
+
+            auto tK = tensor(kp_c, dextents<int, 2>(DK, C), array<int, 2>{ 1, sk });
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_qkA.get_capacity(); ++i) {
+                    if (cT_qkA.is_valid_element(i)) { cT_qkA[i] = 0.0f; }
+                }
+            }
+            mm_qk.run(tK, tQ, cT_qkA);
+
+            // scale in registers; clobber padded kv (idx1 >= kc) to -inf
+            float lmax[QPSG], lsum[QPSG];
+            #pragma clang loop unroll(full)
+            for (int j = 0; j < QPSG; ++j) { lmax[j] = -FLT_MAX / 2; lsum[j] = 0.0f; }
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_qkA.get_capacity(); ++i) {
+                    if (!cT_qkA.is_valid_element(i)) { continue; }
+                    auto idx = cT_qkA.get_multidimensional_index(i);
+                    const int q = (int) idx[0];
+                    float s;
+                    if ((int) idx[1] >= kc) {
+                        s = -FLT_MAX / 2;
+                    } else {
+                        s = cT_qkA[i]*args.scale;
+                        if (has_scap) { s = args.logit_softcap * tanh(s); }
+                        if (has_mask) { s += (float) mp[(uint) q * mstride + (uint) idx[1]] * mscale; }
+                    }
+                    cT_qkA[i] = s;
+                    if (s > lmax[q]) { lmax[q] = s; }
+                }
+            }
+
+            // warp-level reduction: ONE global max (scalar) + per-row sums.
+            // (The max cancels in O/S; a global max keeps the per-chunk O
+            // rescale a trivial scalar multiply.)
+            float lmax_g = -FLT_MAX / 2;
+            #pragma clang loop unroll(full)
+            for (int j = 0; j < QPSG; ++j) {
+                lmax_g = max(lmax_g, simd_max(lmax[j]));
+            }
+            {
+                const float m_old = M;
+                M = max(M, lmax_g);
+                alpha = exp(m_old - M);
+                #pragma clang loop unroll(full)
+                for (int j = 0; j < QPSG; ++j) {
+                    S[j] = S[j]*alpha;
+                }
+            }
+
+            // ---- exp in registers, partial sums ----
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_qkA.get_capacity(); ++i) {
+                    if (!cT_qkA.is_valid_element(i)) { continue; }
+                    auto idx = cT_qkA.get_multidimensional_index(i);
+                    const int q = (int) idx[0];
+                    const float p = exp(cT_qkA[i] - M);
+                    cT_qkA[i] = p; // P in registers (f32), later used as PV right input
+                    lsum[q] += p;
+                }
+            }
+            #pragma clang loop unroll(full)
+            for (int j = 0; j < QPSG; ++j) {
+                S[j] += simd_sum(lsum[j]);
+            }
+        }
+
+        for (int ic = 0; ic < nchunks; ++ic) {
+            // ---- current-chunk context (PV only) ----
+            const int k0 = ic*C;
+            const bool use_pad = has_kvpad && k0 + C > kv;
+            device half * vp_c = use_pad ? (device half *) (pad + pad_v_offs) : vp + (uint) k0*sv;
+            auto tV0 = tensor(vp_c, dextents<int, 2>(PVM, C), array<int, 2>{ 1, sv });
+            auto tV1 = tensor(vp_c + PVM, dextents<int, 2>(PVM, C), array<int, 2>{ 1, sv });
+
+            // ---- rescale both O accumulators (per element: query = idx0) ----
+            // reads alpha_ic (written by the reduction of chunk ic); skip the
+            // first chunk: alpha_0 = exp(-FLT_MAX/2 - m_new) = 0
+            // rescale only when the running max actually changed (alpha < 1):
+            // with random/causal scores M stabilizes after a few chunks and
+            // the pass is skipped for the rest (uniform scalar branch - no
+            // divergence).  The pass itself does not overlap with the tensor
+            // core (it serializes between QK^T and PV), so skipping matters.
+            if (ic > 0 && alpha != 1.0f) {
+                // scalar alpha (global max): a trivial multiply, no index decode
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_pv0.get_capacity(); ++i) {
+                    if (cT_pv0.is_valid_element(i)) { cT_pv0[i] *= alpha; }
+                }
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_pv1.get_capacity(); ++i) {
+                    if (cT_pv1.is_valid_element(i)) { cT_pv1[i] *= alpha; }
+                }
+            }
+
+            // ---- early QK^T(ic+1) + score + reduction -> the other P tile ----
+            // issued before PV(ic) so the two matmuls pipeline on the tensor core
+            if (ic + 1 < nchunks) {
+                const int k1 = (ic + 1)*C;
+                const int kc = min((int) C, kv - k1);
+                const bool use_pad1 = has_kvpad && k1 + C > kv;
+                device half * kp_c = use_pad1 ? (device half *) (pad + pad_k_offs) : kp + (uint) k1*sk;
+
+                device const half * mp = nullptr;
+                int mstride = 0; // in halfs, per local query j
+                if (has_mask) {
+                    if (use_pad1) {
+                        mp      = (device const half *) (pad + pad_mask_offs) + (iq1 + sgitg*QPSG) * (int) C;
+                        mstride = C;
+                    } else {
+                        // global query index: iq1 + sgitg*QPSG + j (j is local, 0..QPSG-1)
+                        mp      = (device const half *) (mask + (uint64_t) k1 * 2 + mbase)
+                                + (uint64_t) (iq1 + sgitg*QPSG) * (args.nb31 / 2);
+                        mstride = (int) (args.nb31 / 2);
+                    }
+                }
+
+                thread cT_qk_t & cT_qkN = (ic & 1) ? cT_qkA : cT_qkB;
+                auto tK = tensor(kp_c, dextents<int, 2>(DK, C), array<int, 2>{ 1, sk });
+                {
+                    #pragma clang loop unroll(full)
+                    for (uint i = 0; i < cT_qkN.get_capacity(); ++i) {
+                        if (cT_qkN.is_valid_element(i)) { cT_qkN[i] = 0.0f; }
+                    }
+                }
+                mm_qk.run(tK, tQ, cT_qkN);
+
+                // scale in registers; clobber padded kv (idx1 >= kc) to -inf
+                float lmax[QPSG], lsum[QPSG];
+                #pragma clang loop unroll(full)
+                for (int j = 0; j < QPSG; ++j) { lmax[j] = -FLT_MAX / 2; lsum[j] = 0.0f; }
+                {
+                    #pragma clang loop unroll(full)
+                    for (uint i = 0; i < cT_qkN.get_capacity(); ++i) {
+                        if (!cT_qkN.is_valid_element(i)) { continue; }
+                        auto idx = cT_qkN.get_multidimensional_index(i);
+                        const int q = (int) idx[0];
+                        float s;
+                        if ((int) idx[1] >= kc) {
+                            s = -FLT_MAX / 2;
+                        } else {
+                            s = cT_qkN[i]*args.scale;
+                            if (has_scap) { s = args.logit_softcap * tanh(s); }
+                            if (has_mask) { s += (float) mp[(uint) q * mstride + (uint) idx[1]] * mscale; }
+                        }
+                        cT_qkN[i] = s;
+                        if (s > lmax[q]) { lmax[q] = s; }
+                    }
+                }
+
+                // warp-level reduction: one global max (scalar) + per-row sums
+                float lmax_g = -FLT_MAX / 2;
+                #pragma clang loop unroll(full)
+                for (int j = 0; j < QPSG; ++j) {
+                    lmax_g = max(lmax_g, simd_max(lmax[j]));
+                }
+                {
+                    const float m_old = M;
+                    M = max(M, lmax_g);
+                    alpha = exp(m_old - M);
+                    #pragma clang loop unroll(full)
+                    for (int j = 0; j < QPSG; ++j) {
+                        S[j] = S[j]*alpha;
+                    }
+                }
+
+                // ---- exp in registers, partial sums ----
+                {
+                    #pragma clang loop unroll(full)
+                    for (uint i = 0; i < cT_qkN.get_capacity(); ++i) {
+                        if (!cT_qkN.is_valid_element(i)) { continue; }
+                        auto idx = cT_qkN.get_multidimensional_index(i);
+                        const int q = (int) idx[0];
+                        const float p = exp(cT_qkN[i] - M);
+                        cT_qkN[i] = p; // P in registers (f32), later used as PV right input
+                        lsum[q] += p;
+                    }
+                }
+                #pragma clang loop unroll(full)
+                for (int j = 0; j < QPSG; ++j) {
+                    S[j] += simd_sum(lsum[j]);
+                }
+            }
+
+            // ---- PV(ic): P(ic) (in the current P tile) as right input, both d blocks ----
+            {
+                thread cT_qk_t & cT_qkC = (ic & 1) ? cT_qkB : cT_qkA;
+                auto cT_pr = mm_pvb.template get_right_input_cooperative_tensor<half, float, float>(cT_qkC);
+                mm_pvb.run(tV0, cT_pr, cT_pv0);
+                mm_pvb.run(tV1, cT_pr, cT_pv1);
+            }
+            simdgroup_barrier(mem_flags::mem_threadgroup);
+        }
+
+        // ---- final: sinks, O /= S, element-wise write for both d blocks ----
+        {
+            // sinks: a virtual sink position with score sinks[iq2] and no O
+            // contribution.  O and S are both relative to the final global
+            // max M, so the correction uses M: M' = max(M, sink),
+            // S' = S*exp(M - M') + exp(sink - M'), O' = O*exp(M - M').
+            // (M' is uniform across queries: the sink is per head, not per
+            // query, but the per-query loop keeps the write form simple.)
+            float alpha_final[QPSG];
+            float s_sink = 0.0f;
+            if (has_sinks) {
+                s_sink = ((device const float *) sinks)[iq2];
+            }
+            const float m2 = has_sinks ? max(M, s_sink) : M;
+            const float a_final = has_sinks ? exp(M - m2) : 1.0f;
+            for (int j = 0; j < QPSG; ++j) {
+                alpha_final[j] = a_final;
+                if (has_sinks) {
+                    S[j] = S[j]*a_final + exp(s_sink - m2);
+                }
+            }
+
+            // output layout is (DV, heads, batch, batch3) with DV innermost:
+            //   offset = (iq3*ne2*ne1 + iq2 + (iq1 + sgitg*QPSG + j)*ne1)*DV + d
+            // NOTE: the write must use ELEMENT pointer arithmetic on a float*
+            // (cast dst first, then add element offsets).  Doing the offset in
+            // BYTES (char* + element*DV) corrupts the coop tile register layout
+            // of the matmul ops above (verified: QK^T (C,QPSG) tile collapses).
+            device float * op = (device float *) dst +
+                    (uint64_t) (iq3*args.ne2*args.ne1 + iq2 + (iq1 + sgitg*QPSG)*args.ne1) * DV;
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_pv0.get_capacity(); ++i) {
+                    if (!cT_pv0.is_valid_element(i)) { continue; }
+                    auto idx = cT_pv0.get_multidimensional_index(i);
+                    const int j = (int) idx[0]; // query
+                    const int d = (int) idx[1]; // head dim (block 0)
+                    if (j >= 0 && j < QPSG && d >= 0 && d < (int) PVM && iq1 + j < args.ne01) {
+                        const float s = S[j];
+                        op[(uint64_t) j*args.ne1*DV + (uint) d] = (s == 0.0f) ? 0.0f : cT_pv0[i]*alpha_final[j]/s;
+                    }
+                }
+            }
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_pv1.get_capacity(); ++i) {
+                    if (!cT_pv1.is_valid_element(i)) { continue; }
+                    auto idx = cT_pv1.get_multidimensional_index(i);
+                    const int j = (int) idx[0]; // query
+                    const int d = PVM + (int) idx[1]; // head dim (block 1)
+                    if (j >= 0 && j < QPSG && d >= 0 && d < (int) DV && iq1 + j < args.ne01) {
+                        const float s = S[j];
+                        op[(uint64_t) j*args.ne1*DV + (uint) d] = (s == 0.0f) ? 0.0f : cT_pv1[i]*alpha_final[j]/s;
+                    }
+                }
+            }
+
+        }
+    } else {
+        // NBLK != 2 (NBLK == 1, or NBLK == 4 for dv = 512): per-d-block form.
+        // For NBLK == 4, carrying 4 live PV destination tiles would overflow
+        // the register file, so QK^T + softmax are recomputed per d block
+        // (correct, but NBLKx the QK work).
+        for (int b = 0; b < NBLK; ++b) {
+            auto cT_pv = make_cT_pv();
+
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_pv.get_capacity(); ++i) {
+                    if (cT_pv.is_valid_element(i)) { cT_pv[i] = 0.0f; }
+                }
+            }
+        // online-softmax state in per-thread registers (uniform across the
+        // simdgroup after simd_max/simd_sum): one GLOBAL running max (scalar)
+        // so the O rescale is a trivial scalar multiply; the rescale does not
+        // overlap the tensor core, so it is skipped when the max is unchanged
+        float M = -FLT_MAX / 2;
+        float S[QPSG];
+        for (int j = 0; j < QPSG; ++j) { S[j] = 0.0f; }
+        float alpha = 1.0f;
+
+        for (int ic = 0; ic < nchunks; ++ic) {
+            const int k0 = ic*C;
+            const int kc = min((int) C, kv - k0);
+
+            // the last partial chunk is read from the pad buffer (padded with 0);
+            // the pad buffer holds only the last C kv items per (KV head, batch),
+            // so its rows are indexed from 0 (not k0)
+            const bool use_pad = has_kvpad && k0 + C > kv;
+            device half * kp_c = use_pad ? (device half *) (pad + pad_k_offs) : kp + (uint) k0*sk;
+            device half * vp_c = use_pad ? (device half *) (pad + pad_v_offs) : vp + (uint) k0*sv;
+
+            // mask row base for this chunk: real mask (kv contiguous, query
+            // stride nb31) or the pad section (C per query, indexed from 0)
+            device const half * mp = nullptr;
+            int mstride = 0; // in halfs, per local query j
+            if (has_mask) {
+                if (use_pad) {
+                    mp      = (device const half *) (pad + pad_mask_offs) + (iq1 + sgitg*QPSG) * (int) C;
+                    mstride = C;
+                } else {
+                    // global query index: iq1 + sgitg*QPSG + j (j is local, 0..QPSG-1)
+                    mp      = (device const half *) (mask + (uint64_t) k0 * 2 + mbase)
+                            + (uint64_t) (iq1 + sgitg*QPSG) * (args.nb31 / 2);
+                    mstride = (int) (args.nb31 / 2);
+                }
+            }
+
+            auto tK  = tensor(kp_c, dextents<int, 2>(DK, C), array<int, 2>{ 1, sk });
+            // block b covers d = b*PVM .. b*PVM + PVM - 1 (d is the innermost dim of V)
+            auto tVb = tensor(vp_c + (uint) b*PVM, dextents<int, 2>(PVM, C), array<int, 2>{ 1, sv });
+
+            // ---- QK^T ----
+            auto cT_qk = mm_qk.template get_destination_cooperative_tensor<decltype(tK), decltype(tQ), float>();
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_qk.get_capacity(); ++i) {
+                    if (cT_qk.is_valid_element(i)) { cT_qk[i] = 0.0f; }
+                }
+            }
+            mm_qk.run(tK, tQ, cT_qk);
+
+            // scale in registers; clobber padded kv (idx1 >= kc) to -inf
+            float lmax[QPSG], lsum[QPSG];
+            #pragma clang loop unroll(full)
+            for (int j = 0; j < QPSG; ++j) { lmax[j] = -FLT_MAX / 2; lsum[j] = 0.0f; }
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_qk.get_capacity(); ++i) {
+                    if (!cT_qk.is_valid_element(i)) { continue; }
+                    auto idx = cT_qk.get_multidimensional_index(i);
+                    const int q = (int) idx[0];
+                    float s;
+                    if ((int) idx[1] >= kc) {
+                        s = -FLT_MAX / 2;
+                    } else {
+                        s = cT_qk[i]*args.scale;
+                        if (has_scap) { s = args.logit_softcap * tanh(s); }
+                        if (has_mask) { s += (float) mp[(uint) q * mstride + (uint) idx[1]] * mscale; }
+                    }
+                    cT_qk[i] = s;
+                    if (s > lmax[q]) { lmax[q] = s; }
+                }
+            }
+
+            // warp-level reduction: one global max (scalar) + per-row sums
+            float lmax_g = -FLT_MAX / 2;
+            #pragma clang loop unroll(full)
+            for (int j = 0; j < QPSG; ++j) {
+                lmax_g = max(lmax_g, simd_max(lmax[j]));
+            }
+            {
+                const float m_old = M;
+                M = max(M, lmax_g);
+                alpha = exp(m_old - M);
+                #pragma clang loop unroll(full)
+                for (int j = 0; j < QPSG; ++j) {
+                    S[j] = S[j]*alpha;
+                }
+            }
+
+            // ---- exp in registers, partial sums ----
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_qk.get_capacity(); ++i) {
+                    if (!cT_qk.is_valid_element(i)) { continue; }
+                    auto idx = cT_qk.get_multidimensional_index(i);
+                    const int q = (int) idx[0];
+                    const float p = exp(cT_qk[i] - M);
+                    cT_qk[i] = p; // P in registers (f32), later used as PV right input
+                    lsum[q] += p;
+                }
+            }
+            #pragma clang loop unroll(full)
+            for (int j = 0; j < QPSG; ++j) {
+                S[j] += simd_sum(lsum[j]);
+            }
+
+            // ---- rescale O accumulator: scalar alpha, skipped when M is
+            // unchanged (the pass does not overlap the tensor core) ----
+            if (ic > 0 && alpha != 1.0f) {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_pv.get_capacity(); ++i) {
+                    if (cT_pv.is_valid_element(i)) { cT_pv[i] *= alpha; }
+                }
+            }
+
+            // ---- PV: P (in registers) as right input ----
+            {
+                auto cT_pr = mm_pvb.template get_right_input_cooperative_tensor<half, float, float>(cT_qk);
+                mm_pvb.run(tVb, cT_pr, cT_pv);
+            }
+        }
+
+        // ---- final: sinks, O /= S, element-wise write for this d block ----
+        {
+            // sinks: O and S are both relative to the final global max M, so
+            // the correction uses M (uniform across queries)
+            float alpha_final;
+            float s_sink = 0.0f;
+            if (has_sinks) {
+                s_sink = ((device const float *) sinks)[iq2];
+            }
+            const float m2 = has_sinks ? max(M, s_sink) : M;
+            alpha_final = has_sinks ? exp(M - m2) : 1.0f;
+            if (has_sinks) {
+                for (int j = 0; j < QPSG; ++j) {
+                    S[j] = S[j]*alpha_final + exp(s_sink - m2);
+                }
+            }
+
+            // output layout is (DV, heads, batch, batch3) with DV innermost:
+            //   offset = (iq3*ne2*ne1 + iq2 + (iq1 + sgitg*QPSG + j)*ne1)*DV + d
+            // (this lane handles queries iq1 + sgitg*QPSG .. + QPSG-1)
+            // NOTE: the write must use ELEMENT pointer arithmetic on a float*
+            // (cast dst first, then add element offsets).  Doing the offset in
+            // BYTES (char* + element*DV) corrupts the coop tile register layout
+            // of the matmul ops above (verified: QK^T (C,QPSG) tile collapses).
+            device float * op = (device float *) dst +
+                    (uint64_t) (iq3*args.ne2*args.ne1 + iq2 + (iq1 + sgitg*QPSG)*args.ne1) * DV;
+            {
+                #pragma clang loop unroll(full)
+                for (uint i = 0; i < cT_pv.get_capacity(); ++i) {
+                    if (!cT_pv.is_valid_element(i)) { continue; }
+                    auto idx = cT_pv.get_multidimensional_index(i);
+                    const int j = (int) idx[0]; // query
+                    const int d = b*PVM + (int) idx[1]; // head dim
+                    if (j >= 0 && j < QPSG && d >= 0 && d < (int) DV && iq1 + j < args.ne01) {
+                        const float s = S[j];
+                        op[(uint64_t) j*args.ne1*DV + (uint) d] = (s == 0.0f) ? 0.0f : cT_pv[i]*alpha_final/s;
+                    }
+                }
+            }
+
+        }
+        }
+    }
+}
+
+template <short DK, short DV>
+kernel void kernel_flash_attn_ext_tensor(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q,
+        device const char * k,
+        device const char * v,
+        device const char * mask,
+        device const char * sinks,
+        device const char * pad,
+        device const char * blk,
+        device       char * dst,
+        threadgroup  half * shmem_f16 [[threadgroup(0)]],
+        uint3   tgpig[[threadgroup_position_in_grid]],
+        ushort  tiitg[[thread_index_in_threadgroup]],
+        ushort  sgitg[[simdgroup_index_in_threadgroup]]) {
+    (void) blk;
+
+    kernel_flash_attn_ext_tensor_impl<DK, DV>(args, q, k, v, mask, sinks, pad, dst, shmem_f16, tgpig, tiitg, sgitg);
+}
+
+// head size instantiations (f16 K/V only - see the host-side fast-path gate)
+template [[host_name("kernel_flash_attn_ext_tensor_dk64_dv64"  )]]  kernel void kernel_flash_attn_ext_tensor<  64,  64>(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q, device const char * k, device const char * v, device const char * mask,
+        device const char * sinks, device const char * pad, device const char * blk,
+        device       char * dst, threadgroup half * shmem_f16, uint3 tgpig, ushort tiitg, ushort sgitg);
+template [[host_name("kernel_flash_attn_ext_tensor_dk128_dv128")]]  kernel void kernel_flash_attn_ext_tensor< 128, 128>(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q, device const char * k, device const char * v, device const char * mask,
+        device const char * sinks, device const char * pad, device const char * blk,
+        device       char * dst, threadgroup half * shmem_f16, uint3 tgpig, ushort tiitg, ushort sgitg);
+template [[host_name("kernel_flash_attn_ext_tensor_dk192_dv128")]]  kernel void kernel_flash_attn_ext_tensor< 192, 128>(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q, device const char * k, device const char * v, device const char * mask,
+        device const char * sinks, device const char * pad, device const char * blk,
+        device       char * dst, threadgroup half * shmem_f16, uint3 tgpig, ushort tiitg, ushort sgitg);
+template [[host_name("kernel_flash_attn_ext_tensor_dk192_dv192")]]  kernel void kernel_flash_attn_ext_tensor< 192, 192>(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q, device const char * k, device const char * v, device const char * mask,
+        device const char * sinks, device const char * pad, device const char * blk,
+        device       char * dst, threadgroup half * shmem_f16, uint3 tgpig, ushort tiitg, ushort sgitg);
+template [[host_name("kernel_flash_attn_ext_tensor_dk256_dv256")]]  kernel void kernel_flash_attn_ext_tensor< 256, 256>(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q, device const char * k, device const char * v, device const char * mask,
+        device const char * sinks, device const char * pad, device const char * blk,
+        device       char * dst, threadgroup half * shmem_f16, uint3 tgpig, ushort tiitg, ushort sgitg);
+template [[host_name("kernel_flash_attn_ext_tensor_dk320_dv256")]]  kernel void kernel_flash_attn_ext_tensor< 320, 256>(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q, device const char * k, device const char * v, device const char * mask,
+        device const char * sinks, device const char * pad, device const char * blk,
+        device       char * dst, threadgroup half * shmem_f16, uint3 tgpig, ushort tiitg, ushort sgitg);
+template [[host_name("kernel_flash_attn_ext_tensor_dk512_dv512")]]  kernel void kernel_flash_attn_ext_tensor< 512, 512>(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q, device const char * k, device const char * v, device const char * mask,
+        device const char * sinks, device const char * pad, device const char * blk,
+        device       char * dst, threadgroup half * shmem_f16, uint3 tgpig, ushort tiitg, ushort sgitg);
+template [[host_name("kernel_flash_attn_ext_tensor_dk576_dv512")]]  kernel void kernel_flash_attn_ext_tensor< 576, 512>(
+        constant ggml_metal_kargs_flash_attn_ext & args,
+        device const char * q, device const char * k, device const char * v, device const char * mask,
+        device const char * sinks, device const char * pad, device const char * blk,
+        device       char * dst, threadgroup half * shmem_f16, uint3 tgpig, ushort tiitg, ushort sgitg);
+#endif // GGML_METAL_HAS_TENSOR