#version 450

#extension GL_EXT_control_flow_attributes : enable
#extension GL_EXT_shader_16bit_storage : require

#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require

#ifdef FLOAT16
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
#extension GL_EXT_shader_subgroup_extended_types_float16 : require
#endif

#ifdef MMQ
#extension GL_EXT_integer_dot_product : require
#extension GL_KHR_shader_subgroup_clustered : require

#include "mul_mmq_shmem_types.glsl"
#endif

#extension GL_KHR_shader_subgroup_shuffle : enable
#extension GL_KHR_shader_subgroup_vote : enable

#include "types.glsl"
#include "dot_product_funcs.glsl"
#include "flash_attn_base.glsl"
#include "flash_attn_dequant.glsl"

const uint32_t HSK_per_thread = HSK / D_split;
const uint32_t HSV_per_thread = HSV / D_split;

const uint32_t rows_per_thread = Br / row_split;
const uint32_t cols_per_iter = WorkGroupSize / D_split / row_split;
const uint32_t cols_per_thread = Bc / cols_per_iter;
const uint32_t num_subgroups = SubGroupSize == 0 ? 0 : WorkGroupSize / SubGroupSize;


layout (binding = 0) readonly buffer Q {float data_q[];};
layout (binding = 0) readonly buffer QV4 {vec4 data_qv4[];};
layout (binding = 1) readonly buffer K {float16_t data_k[];};
layout (binding = 1) readonly buffer KV4 {f16vec4 data_kv4[];};
layout (binding = 2) readonly buffer V {float16_t data_v[];};
layout (binding = 2) readonly buffer VV4 {f16vec4 data_vv4[];};
layout (binding = 3) readonly buffer M {float16_t data_m[];};

// If SubGroupSize is set to 0 then only use shmem reductions
const uint32_t tmpsh_size = (SubGroupSize > 0) ? (row_split == 1 ? num_subgroups * D_split : num_subgroups) : WorkGroupSize;
shared float tmpsh[tmpsh_size];
shared FLOAT_TYPEV4 tmpshv4[tmpsh_size];

const uint32_t masksh_stride = Br + 1;
shared FLOAT_TYPE masksh[Bc * masksh_stride];

#ifndef MMQ
const uint32_t qf_stride = HSK / 4 + 1;
shared FLOAT_TYPEV4 Qf[Br * qf_stride];
#else

const uint32_t qf_stride = HSK / 32;
shared block_b_cache Qf[Br * qf_stride];
#endif

#ifndef MMQ
const uint32_t D = HSK > HSV ? HSK : HSV;
#else
const uint32_t D = HSV;
#endif
const uint32_t kvsh_stride = D / 4 + 1;
shared FLOAT_TYPEV4 kvsh[SHMEM_STAGING != 0 ? Bc * kvsh_stride : 1];

#ifdef MMQ

shared block_a_cache kblocksh[SHMEM_STAGING != 0 ? Bc * qf_stride : 1];
#endif

shared vec4 occupancy_limiter[LIMIT_OCCUPANCY_SHMEM > 0 ? LIMIT_OCCUPANCY_SHMEM : 1];

#ifdef MMQ
#include "flash_attn_mmq_funcs.glsl"
#endif

void main() {
#ifdef NEEDS_INIT_IQ_SHMEM
    if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) {
        init_iq_shmem(gl_WorkGroupSize);
    }
#endif

    init_indices();

    const uint32_t tid = gl_LocalInvocationIndex;
    const uint32_t threads_per_rowgroup = gl_WorkGroupSize.x / row_split;
    const uint32_t row_tid = gl_LocalInvocationIndex / threads_per_rowgroup;
    const uint32_t rowgroup_tid = gl_LocalInvocationIndex % threads_per_rowgroup;
    const uint32_t d_tid = gl_LocalInvocationIndex % D_split;
    const uint32_t col_tid = (gl_LocalInvocationIndex % threads_per_rowgroup) / D_split;

    if (LIMIT_OCCUPANCY_SHMEM > 0) {
        // This just exists to avoid the occupancy_limiter array getting optimized out
        occupancy_limiter[tid] = vec4(tid);

        barrier();

        if (occupancy_limiter[tid] == vec4(99999.0)) {
            data_ov4[0] = D_TYPEV4(occupancy_limiter[tid]);
        }
    }

#define tile_row(r) (row_tid * rows_per_thread + (r))

    uint32_t q_offset = gqa_iq1*p.nb01 + (iq2*p.nb02 + iq3*p.nb03) / 4;

    [[unroll]] for (uint32_t idx = 0; idx < Br * HSK / 4; idx += gl_WorkGroupSize.x) {
        uint32_t d = (idx + tid) % (HSK / 4);
        uint32_t r = (idx + tid) / (HSK / 4);
        const bool is_in_bounds = r < Br && d < HSK / 4 && i * Br + r < N;
#ifndef MMQ
        if (is_in_bounds) {
            Qf[r * qf_stride + d] = FLOAT_TYPEV4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale);
        }
#else
        const uint buf_ib = r * qf_stride + d / 8;
        const uint buf_iqs = d % 8;

        vec4 vals = is_in_bounds ? data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale : vec4(0.0f);
        const vec4 abs_vals = abs(vals);

        const float thread_max = max(max(abs_vals.x, abs_vals.y), max(abs_vals.z, abs_vals.w));
        const float amax = subgroupClusteredMax(thread_max, 8);
        const float qd = amax / 127.0f;
        const float qd_inv = qd != 0.0f ? 1.0f / qd : 0.0f;
        vals = round(vals * qd_inv);

        Qf[buf_ib].qs[buf_iqs] = pack32(i8vec4(vals));

        // Q8_0 K only needs (qd, _); the asymmetric Q4_*/Q5_* family also stores
        // the row-sum scaled by qd, used in k_dot_correction.
        if (FaTypeK == GGML_TYPE_Q8_0) {
            if (buf_iqs == 0) {
                Qf[buf_ib].ds = FLOAT_TYPEV2(qd, 0.0f);
            }
        } else {
            const float thread_sum = vals.x + vals.y + vals.z + vals.w;
            const float sum = subgroupClusteredAdd(thread_sum, 8);

            if (buf_iqs == 0) {
                Qf[buf_ib].ds = FLOAT_TYPEV2(qd, sum * qd);
            }
        }
#endif
    }
    barrier();

    FLOAT_TYPEV4 Of[rows_per_thread][HSV_per_thread / 4];
    [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
        [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
            Of[r][d] = FLOAT_TYPEV4(0.0);
        }
    }

    float Lf[rows_per_thread], Mf[rows_per_thread];

    // Use -FLT_MAX/2 rather than -inf to reduce the possibility of NaNs, e.g. when computing Mold-M.
    const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF);

    [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
        Lf[r] = 0;
        Mf[r] = NEG_FLT_MAX_OVER_2;
    }

    ACC_TYPE slope[rows_per_thread];
    [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
        slope[r] = ACC_TYPE(1.0);
    }

    // ALiBi
    if (p.max_bias > 0.0f) {
        [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
            slope[r] = perElemOpComputeSlope(tile_row(r), col_tid, ACC_TYPE(0), iq2);
        }
    }

    const uint32_t mo_stride = CEIL_DIV(KV, 16 * Bc);
    // mo_offset will point to the tile starting at row i*Br and col 0
    uint32_t mo_offset = mo_stride * i;

    // FaBlockBytesK/V == 2 for f16, 16 for f32, ggml block byte size for quants.
    uint32_t k_offset = (ik2*p.nb12 + ik3*p.nb13) / FaBlockBytesK;
    uint32_t v_offset = (iv2*p.nb22 + iv3*p.nb23) / FaBlockBytesV;
    uint32_t m_offset = gqa_iq1*KV;
    if (p.nem2 != 1 || p.nem3 != 1) {
        m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV;
        mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride;
    }

    uint32_t mask_opt = 0;
    uint32_t mask_opt_idx = ~0;
    uint32_t mask_opt_bits = 0;

    [[dont_unroll]]
    for (uint32_t j = start_j; j < end_j; ++j) {
        if (MASK_ENABLE) {
            if (USE_MASK_OPT && mask_opt_idx != j / 16) {
                mask_opt_idx = j / 16;
                mask_opt = data_mask_opt[mo_offset + mask_opt_idx];
            }
            mask_opt_bits = (mask_opt >> ((j % 16) * 2)) & 0x3;
            if (mask_opt_bits == MASK_OPT_ALL_NEG_INF) {
                // skip this block
                continue;
            }
            // Only load if the block is not all zeros
            if (mask_opt_bits != MASK_OPT_ALL_ZERO) {
                bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0;

                float max_mask = NEG_FLT_MAX_OVER_2;
                barrier();
                [[unroll]] for (uint32_t idx = 0; idx < Bc * Br; idx += gl_WorkGroupSize.x) {
                    uint32_t c = (idx + tid) % Bc;
                    uint32_t r = (idx + tid) / Bc;
                    if (idx + tid < Bc * Br) {
                        if ((!KV_bounds_check || j * Bc + c < KV) && (!nem1_bounds_check || i * Br + r < p.nem1)) {
                            FLOAT_TYPE m = FLOAT_TYPE(data_m[m_offset + (i * Br + r) * m_stride + (j * Bc + c)]);
                            masksh[c * masksh_stride + r] = m;
                            max_mask = max(max_mask, float(m));
                        } else {
                            masksh[c * masksh_stride + r] = FLOAT_TYPE(0);
                        }
                    }
                }
                // skip the block if the mask is entirely -inf
                bool all_less = subgroupAll(max_mask <= NEG_FLT_MAX_OVER_2);
                barrier();
                if (gl_SubgroupInvocationID == 0) {
                    tmpsh[gl_SubgroupID] = all_less ? NEG_FLT_MAX_OVER_2 : 0.0f;
                }
                barrier();
                [[unroll]] for (uint s = 0; s < gl_NumSubgroups; ++s) {
                    max_mask = max(max_mask, tmpsh[s]);
                }
                if (max_mask <= NEG_FLT_MAX_OVER_2) {
                    continue;
                }
            }
        }

        ACC_TYPE Sf[rows_per_thread][cols_per_thread];
        [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
            [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
                Sf[r][c] = ACC_TYPE(0.0);
            }
        }

        if (SHMEM_STAGING != 0) {
            barrier();
#ifndef MMQ
            [[unroll]] for (uint32_t idx = 0; idx < Bc * HSK / 4; idx += gl_WorkGroupSize.x) {
                uint32_t d = (idx + tid) % (HSK / 4);
                uint32_t c = (idx + tid) / (HSK / 4);
                if (idx + gl_WorkGroupSize.x <= Bc * HSK / 4 || c < Bc) {
                    FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0);
                    if (!KV_bounds_check || j * Bc + c < KV) {
                        if (USE_DECODE_K) {
                            uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE_K + 4 * d;
                            uint ib = coord / BLOCK_SIZE_K;
                            uint iqs = (coord % BLOCK_SIZE_K);
                            K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K);
                        } else {
                            K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]);
                        }
                    }

                    kvsh[c * kvsh_stride + d] = K_Tf;
                }
            }
#else // MMQ
            const uint ints_per_block = 8u / fa_quant_r_mmq(FaTypeK);
            const uint quant_iters = Bc * HSK / 32 * ints_per_block;
            [[unroll]] for (uint32_t idx = 0; idx < quant_iters; idx += gl_WorkGroupSize.x) {
                const uint32_t iqs = (idx + tid) % ints_per_block;
                const uint32_t ib = (idx + tid) / ints_per_block;
                const uint32_t c = ib / (HSK / 32);
                const uint32_t block = ib % (HSK / 32);
                if (idx + gl_WorkGroupSize.x <= quant_iters || c < Bc) {
                    const uint buf_ib = c * qf_stride + block;
                    if (!KV_bounds_check || j * Bc + c < KV) {
                        const uint global_ib = (j * Bc + c) * k_stride + block;
                        k_block_to_shmem(buf_ib, global_ib, iqs, k_offset);
                    } else {
                        k_block_to_shmem_zero(buf_ib, iqs);
                    }
                }
            }
#endif // MMQ
            barrier();
        }

#ifndef MMQ
        // More d iterations means Q register caching becomes relevant
        // Few iterations means the additional registers needed are worse than the speed-up from caching
        if (HSK_per_thread / 4 > 4) {
            [[unroll]] for (uint32_t d = 0; d < HSK_per_thread / 4; ++d) {
                FLOAT_TYPEV4 Q_cache[rows_per_thread];
                [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                    Q_cache[r] = Qf[tile_row(r) * qf_stride + d * D_split + d_tid];
                }

                [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
                    if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) {
                        continue;
                    }

                    FLOAT_TYPEV4 K_Tf;
                    if (SHMEM_STAGING != 0) {
                        K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)];
                    } else if (USE_DECODE_K) {
                        uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid);
                        uint ib = coord / BLOCK_SIZE_K;
                        uint iqs = (coord % BLOCK_SIZE_K);
                        K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K);
                    } else {
                        K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]);
                    }
                    [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                        Sf[r][c] = dot_product(Q_cache[r], K_Tf, Sf[r][c]);
                    }
                }
            }
        } else {
            [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
                if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) {
                    continue;
                }

                [[unroll]] for (uint32_t d = 0; d < HSK_per_thread / 4; ++d) {
                    FLOAT_TYPEV4 K_Tf;
                    if (SHMEM_STAGING != 0) {
                        K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)];
                    } else if (USE_DECODE_K) {
                        uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid);
                        uint ib = coord / BLOCK_SIZE_K;
                        uint iqs = (coord % BLOCK_SIZE_K);
                        K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K);
                    } else {
                        K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]);
                    }
                    [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                        Sf[r][c] = dot_product(Qf[tile_row(r) * qf_stride + d * D_split + d_tid], K_Tf, Sf[r][c]);
                    }
                }
            }
        }
#else // MMQ
        const uint hsk4 = HSK_per_thread / 4;
        const uint d_per_step = (hsk4 % 8 == 0) ? 8 :
                                (hsk4 % 4 == 0) ? 4 :
                                (hsk4 % 2 == 0) ? 2 : 1;

        [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
            if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) {
                continue;
            }

            [[unroll]] for (uint32_t d_block = 0; d_block < HSK_per_thread / 4; d_block += d_per_step) {
                int32_t k_quants[d_per_step];
                ACC_TYPEV2 k_dm;

                // Q4_*/Q5_* take the block-8 fast path when one step covers a full
                // block; Q8_0 always goes through the per-int get_k_qs* helpers
                // (its qs is byte-packed, not nibble-packed).
                const bool block8_fast = (d_per_step == 8) && (FaTypeK != GGML_TYPE_Q8_0);

                if (SHMEM_STAGING != 0) {
                    const uint k_block_idx = (d_tid * (HSK_per_thread / 4) + d_block) / 8;
                    const uint buf_ib = (c * cols_per_iter + col_tid) * qf_stride + k_block_idx;
                    k_dm = ACC_TYPEV2(kblocksh[buf_ib].dm);

                    if (block8_fast) {
                        const bool has_qh = (FaTypeK == GGML_TYPE_Q5_0) || (FaTypeK == GGML_TYPE_Q5_1);
                        [[unroll]] for (uint32_t d = 0; d < 4; d++) {
                            uint vui = kblocksh[buf_ib].qs[d];
                            k_quants[d    ] = int32_t( vui       & 0x0F0F0F0F);
                            k_quants[d + 4] = int32_t((vui >> 4) & 0x0F0F0F0F);
                            if (has_qh) {
                                uint qh_lo = (kblocksh[buf_ib].qh >> (d * 4)) & 0xF;
                                uint qh_hi = (kblocksh[buf_ib].qh >> (d * 4 + 16)) & 0xF;
                                k_quants[d    ] |= int32_t((qh_lo * 0x02040810u) & 0x10101010u);
                                k_quants[d + 4] |= int32_t((qh_hi * 0x02040810u) & 0x10101010u);
                            }
                        }
                    } else {
                        [[unroll]] for (uint32_t d = 0; d < d_per_step; d++) {
                            k_quants[d] = get_k_qs_shmem(buf_ib, (d_tid * (HSK_per_thread / 4) + d_block) % 8 + d);
                        }
                    }
                } else {
                    const uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d_tid * (HSK_per_thread / 4) + d_block);
                    const uint ib = coord / BLOCK_SIZE_K;
                    const uint iqs = (coord % BLOCK_SIZE_K);

                    k_dm = ACC_TYPEV2(get_k_scale(ib, k_offset));

                    if (block8_fast) {
                        fa_k_qs_block8 blk = get_k_qs_block8(ib, k_offset);
                        [[unroll]] for (uint32_t d = 0; d < 8; d++) {
                            k_quants[d] = blk.qs[d];
                        }
                    } else {
                        [[unroll]] for (uint32_t d = 0; d < d_per_step; d++) {
                            k_quants[d] = get_k_qs(ib, iqs + d * 4, k_offset);
                        }
                    }
                }

                [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                    const uint qib = tile_row(r) * qf_stride + (d_tid * (HSK_per_thread / 4) + d_block) / 8;
                    const uint qiqs = (d_tid * (HSK_per_thread / 4) + d_block) % 8;

                    int32_t acc = 0;
                    [[unroll]] for (uint32_t d = 0; d < d_per_step; d++) {
                        acc += dotPacked4x8EXT(Qf[qib].qs[qiqs + d], k_quants[d]);
                    }

                    Sf[r][c] += ACC_TYPE(acc) * ACC_TYPE(Qf[qib].ds.x) * k_dm.x;
                    if ((d_tid * (HSK_per_thread / 4) + d_block) % 8 == 0) {
                        Sf[r][c] += k_dot_correction(qib, k_dm);
                    }
                }
            }
        }
#endif // MMQ

        [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
            // Compute sum across the D_split
            [[unroll]] for (uint s = D_split / 2; s > 0; s >>= 1) {
                [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                    Sf[r][c] += subgroupShuffleXor(Sf[r][c], s);
                }
            }
        }

        if (LOGIT_SOFTCAP) {
            [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
                    Sf[r][c] = ACC_TYPE(p.logit_softcap * tanh(Sf[r][c]));
                }
            }
        }

        if (MASK_ENABLE && mask_opt_bits != MASK_OPT_ALL_ZERO) {
            [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
                [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                    FLOAT_TYPE mvf = masksh[(c * cols_per_iter + col_tid) * masksh_stride + tile_row(r)];

                    Sf[r][c] += slope[r]*mvf;
                }
            }
        }

        float eMf[rows_per_thread];
        [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
            float rowmaxf = NEG_FLT_MAX_OVER_2;
            [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
                if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) {
                    continue;
                }
                rowmaxf = max(rowmaxf, float(Sf[r][c]));
            }
            rowmaxf += FATTN_KQ_MAX_OFFSET;
            float Moldf = Mf[r];

            // M = max(rowmax, Mold)
            // P = e^(S - M)
            // eM = e^(Mold - M)
            Mf[r] = max(rowmaxf, Moldf);
            eMf[r] = exp(Moldf - Mf[r]);
            Lf[r] = eMf[r]*Lf[r];
        }

        [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
            [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                Of[r][d] = FLOAT_TYPE(eMf[r]) * Of[r][d];
            }
        }

        if (SHMEM_STAGING != 0) {
            barrier();
            [[unroll]] for (uint32_t idx = 0; idx < Bc * HSV / 4; idx += gl_WorkGroupSize.x) {
                uint32_t d = (idx + tid) % (HSV / 4);
                uint32_t c = (idx + tid) / (HSV / 4);
                if (idx + gl_WorkGroupSize.x <= Bc * HSV / 4 || c < Bc) {
                    FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0);
                    if (!KV_bounds_check || j * Bc + c < KV) {
                        if (USE_DECODE_V) {
                            uint coord = (j * Bc + c) * v_stride * BLOCK_SIZE_V + 4 * d;
                            uint ib = coord / BLOCK_SIZE_V;
                            uint iqs = (coord % BLOCK_SIZE_V);
                            V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V);
                        } else {
                            V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c) * v_stride / 4 + d]);
                        }
                    }

                    kvsh[c * kvsh_stride + d] = V_Tf;
                }
            }
            barrier();
        }

        [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) {
            if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) {
                continue;
            }

            FLOAT_TYPE Pf[rows_per_thread];
            [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                Pf[r] = FLOAT_TYPE(exp(float(Sf[r][c]) - Mf[r]));
                Lf[r] += Pf[r];
            }

            [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
                FLOAT_TYPEV4 Vf;
                if (SHMEM_STAGING != 0) {
                    Vf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)];
                } else if (USE_DECODE_V) {
                    uint coord = (j * Bc + c * cols_per_iter + col_tid) * v_stride * BLOCK_SIZE_V + 4 * (d * D_split + d_tid);
                    uint ib = coord / BLOCK_SIZE_V;
                    uint iqs = (coord % BLOCK_SIZE_V);
                    Vf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V);
                } else {
                    Vf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * v_stride / 4 + d * D_split + d_tid]);
                }
                [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                    Of[r][d] += FLOAT_TYPEV4(Pf[r] * Vf);
                }
            }
        }
    }

    // prevent race on tmpsh
    barrier();

    // reduce across threads

    [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
        float rowmaxf = Mf[r];

        // Compute max across the row
        if (SubGroupSize > 0) {
            [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) {
                rowmaxf = max(rowmaxf, subgroupShuffleXor(rowmaxf, s));
            }
            if (row_split == 1) {
                // Reduce inside workgroup with shmem
                barrier();
                if (gl_SubgroupInvocationID == d_tid) {
                    tmpsh[gl_SubgroupID * D_split + d_tid] = rowmaxf;
                }
                barrier();
                rowmaxf = tmpsh[d_tid];
                [[unroll]] for (uint32_t s = 1; s < num_subgroups; ++s) {
                    rowmaxf = max(rowmaxf, tmpsh[s * D_split + d_tid]);
                }
            }
        } else {
            barrier();
            tmpsh[tid] = rowmaxf;
            barrier();
            [[unroll]] for (int s = int(threads_per_rowgroup) / 2; s >= D_split; s >>= 1) {
                if (rowgroup_tid < s) {
                    tmpsh[tid] = max(tmpsh[tid], tmpsh[tid ^ s]);
                }
                barrier();
            }
            rowmaxf = tmpsh[row_tid * threads_per_rowgroup + d_tid];
        }

        float Moldf = Mf[r];

        // M = max(rowmax, Mold)
        // eM = e^(Mold - M)
        Mf[r] = max(rowmaxf, Moldf);
        float eMf = exp(Moldf - Mf[r]);

        Lf[r] = eMf*Lf[r];

        // Compute sum across the row
        if (SubGroupSize > 0) {
            [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) {
                Lf[r] += subgroupShuffleXor(Lf[r], s);
            }
            if (row_split == 1) {
                barrier();
                if (gl_SubgroupInvocationID == d_tid) {
                    tmpsh[gl_SubgroupID * D_split + d_tid] = Lf[r];
                }
                barrier();
                Lf[r] = tmpsh[d_tid];
                [[unroll]] for (uint32_t s = 1; s < num_subgroups; ++s) {
                    Lf[r] += tmpsh[s * D_split + d_tid];
                }
            }
        } else {
            barrier();
            tmpsh[tid] = Lf[r];
            barrier();
            [[unroll]] for (int s = int(threads_per_rowgroup) / 2; s >= D_split; s >>= 1) {
                if (rowgroup_tid < s) {
                    tmpsh[tid] = tmpsh[tid] + tmpsh[tid ^ s];
                }
                barrier();
            }
            Lf[r] = tmpsh[row_tid * threads_per_rowgroup + d_tid];
        }

        [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
            Of[r][d] = FLOAT_TYPE(eMf) * Of[r][d];

            if (SubGroupSize > 0) {
                [[unroll]] for (uint s = D_split; s < SubGroupSize; s *= 2) {
                    if (!OLD_AMD_WINDOWS) {
                        Of[r][d] += subgroupShuffleXor(Of[r][d], s);
                    } else {
                        // Something about f16vec4 subgroupShuffleXor is broken on AMD Windows RDNA2 and below.
                        // Shuffle full vec4 as workaround.
                        // See https://github.com/ggml-org/llama.cpp/issues/19881#issuecomment-3958643697
                        Of[r][d] += FLOAT_TYPEV4(subgroupShuffleXor(vec4(Of[r][d]), s));
                    }
                }
                if (row_split == 1) {
                    barrier();
                    if (gl_SubgroupInvocationID == d_tid) {
                        tmpshv4[gl_SubgroupID * D_split + d_tid] = Of[r][d];
                    }
                    barrier();
                    Of[r][d] = tmpshv4[d_tid];
                    [[unroll]] for (uint32_t s = 1; s < num_subgroups; ++s) {
                        Of[r][d] += tmpshv4[s * D_split + d_tid];
                    }
                }
            } else {
                barrier();
                tmpshv4[tid] = Of[r][d];
                barrier();
                [[unroll]] for (int s = int(threads_per_rowgroup) / 2; s >= D_split; s >>= 1) {
                    if (rowgroup_tid < s) {
                        Of[r][d] += tmpshv4[tid ^ s];
                        tmpshv4[tid] = Of[r][d];
                    }
                    barrier();
                }
                Of[r][d] = tmpshv4[row_tid * threads_per_rowgroup + d_tid];
            }
        }
    }


    // If there is split_k, then the split_k resolve shader does the final
    // division by L. Store the intermediate O value and per-row m and L values.
    if (p.k_num > 1) {
        if (p.gqa_ratio > 1) {
            // note: O and Q have swapped coord 1,2.
            uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)) / 4;

            [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                const uint row = tile_row(r);
                if (row < N) {
                    [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
                        gqaStore(row, d * D_split + d_tid, Of[r][d], o_offset, iq2, N);
                    }
                }
            }

            o_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3));
            [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                const uint row = tile_row(r);
                if (row < N) {
                    perElemOpStoreCol0(row, 0u, ACC_TYPE(Lf[r]), o_offset, iq2, N);
                    perElemOpStoreCol0(row, 0u, ACC_TYPE(Mf[r]), o_offset + p.ne1, iq2, N);
                }
            }
        } else {
            [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
                const uint row = tile_row(r);
                const uint global_row = i * Br + row;

                if (global_row < N) {
                    uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)) / 4;

                    [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
                        data_ov4[o_offset + iq2 * HSV/4 + d * D_split + d_tid] = D_TYPEV4(Of[r][d]);
                    }
                }

                if (global_row < N && d_tid == 0 && col_tid == 0) {
                    uint32_t lm_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3));
                    data_o[lm_offset + iq2] = D_TYPE(Lf[r]);
                    data_o[lm_offset + p.ne1 + iq2] = D_TYPE(Mf[r]);
                }
            }
        }
        return;
    }

    if ((p.mask_n_head_log2 & SINK_ENABLE_BIT) != 0) {
        [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
            float sink = perElemOpGetSink(tile_row(r), 0u, ACC_TYPE(0), iq2);

            float ms = 1.0f;
            float vs = 1.0f;

            if (sink > Mf[r]) {
                ms = exp(Mf[r] - sink);

                [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
                    Of[r][d] *= FLOAT_TYPE(ms);
                }
            } else {
                vs = exp(sink - Mf[r]);
            }

            Lf[r] = Lf[r]*ms + vs;
        }
    }

    float Lfrcp[rows_per_thread];
    [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
        Lfrcp[r] = (Lf[r] == 0.0) ? 0.0 : (1.0 / Lf[r]);
    }

    [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
        [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
            Of[r][d] *= FLOAT_TYPE(Lfrcp[r]);
#if defined(FLOAT_TYPE_MAX)
            Of[r][d] = clamp(Of[r][d], -FLOAT_TYPE_MAX, FLOAT_TYPE_MAX);
#endif
        }
    }

    uint32_t o_offset = (gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV) / 4;

    if (p.gqa_ratio > 1) {
        [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
            const uint row = tile_row(r);
            if (row < N) {
                [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
                    gqaStore(row, d * D_split + d_tid, Of[r][d], o_offset, iq2, N);
                }
            }
        }
    } else {
        [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) {
            const uint row = tile_row(r);
            if (i * Br + row < N) {
                [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) {
                    data_ov4[o_offset + (iq2 * HSV + (i * Br + row) * p.ne1 * HSV) / 4 + d * D_split + d_tid] = D_TYPEV4(Of[r][d]);
                }
            }
        }
    }
}