// Hardware-verified on Intel ARL (Mesa 25.2.8) + lavapipe (Mesa 25.2.8 LLVMpipe)
// — 8/8 PASS for the 8-token fixture (see packages/inference/README.md
// "Verification matrix" and reports/porting/2026-05-09-w4/vulkan-turbo-fix.md).
//
// turbo3 KV cache dequant + Q·K dot product (Vulkan compute shader, GLSL 450).
//
// Ports buun-llama-cpp's CUDA dequantize_turbo3_0 + the inline FA dot product
// from ggml/src/ggml-cuda/turbo-quant-cuda.cuh (commit 6575873e9c) and the
// flash-attention vec path that consumes it (fattn-vec.cuh).
//
// Block layout (matches block_turbo3_0 in ggml-common.h, 14 bytes):
//     half  norm                 // [0..1]   fp16 corrected group norm
//     u8    qs[8]                // [2..9]   QK_TURBO3/4 = 8 bytes, low 2 bits per element
//     u8    signs[4]             // [10..13] QK_TURBO3/8 = 4 bytes, hi 1 bit per element
//
// Four 32-element blocks form one 128-element rotation group. The graph
// pre-rotates the queries with the same FWHT used at quantize time, so this
// shader does NOT apply an inverse rotation — it just decodes 128 centroids
// and accumulates the dot product against a pre-rotated Q vector.
//
// Dispatch: one workgroup per (n_kv_block, n_head). 32 invocations cooperate
// to dequant 128 elements (4 per invocation) and reduce a partial dot product.

#version 450
#extension GL_EXT_shader_16bit_storage  : require
#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require
#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require

layout(local_size_x = 32, local_size_y = 1, local_size_z = 1) in;

// --- Q vector: pre-rotated, fp32, length head_dim (we assume 128 here). ---
layout(std430, binding = 0) readonly buffer QBuf {
    float q[];          // [n_head * 128]
};

// --- K cache: tightly packed block_turbo3_0 stream. ---
// Vulkan std430 buffers are aligned to 4 bytes; we pack 4 blocks (4 * 14 = 56
// bytes) interleaved as a flat byte stream and index it as uint chunks.
layout(std430, binding = 1) readonly buffer KBuf {
    uint k_packed[];    // [(n_kv * head_dim) / sizeof(block_turbo3_0) bytes / 4]
};

// --- Output: scalar score per (n_kv_block, n_head). ---
layout(std430, binding = 2) writeonly buffer OutBuf {
    float scores[];
};

// --- Push constants. ---
layout(push_constant) uniform Push {
    uint head_dim;          // must be 128 for this DRAFT
    uint n_kv;              // sequence length
    uint kv_stride_blocks;  // blocks per KV row (head_dim / 32 = 4 for d=128)
    uint q_head;            // head index into Q
    uint head_offset_bytes; // per-head byte offset into k_packed
} push;

// --- Centroids (Lloyd-Max for N(0, 1/128)). CUDA: d_turbo_centroids_3bit. ---
const float TURBO_CENTROIDS_3BIT[8] = float[8](
    -0.190685, -0.117832, -0.065717, -0.021460,
     0.021460,  0.065717,  0.117832,  0.190685
);

// fp16 → fp32 (manual; KBuf is interpreted as raw bytes). Mirrors __half2float.
float fp16_to_fp32(uint h16) {
    uint sign = (h16 & 0x8000u) << 16;
    uint exp  = (h16 >> 10) & 0x1Fu;
    uint mant = h16 & 0x3FFu;
    uint u;
    if (exp == 0u) {
        if (mant == 0u) {
            u = sign;
        } else {
            // Subnormal — rare for norms, but handle for correctness.
            uint e = 1u;
            while ((mant & 0x400u) == 0u) { mant <<= 1; e += 1u; }
            mant &= 0x3FFu;
            u = sign | ((127u - 15u - e + 1u) << 23) | (mant << 13);
        }
    } else if (exp == 0x1Fu) {
        u = sign | 0x7F800000u | (mant << 13);
    } else {
        u = sign | ((exp + 127u - 15u) << 23) | (mant << 13);
    }
    return uintBitsToFloat(u);
}

// Read one byte from k_packed at absolute byte offset `b`.
uint read_byte(uint b) {
    uint w = k_packed[b >> 2];
    return (w >> ((b & 3u) * 8u)) & 0xFFu;
}

// Read a 16-bit little-endian halfword starting at byte offset `b`.
uint read_u16(uint b) {
    return read_byte(b) | (read_byte(b + 1u) << 8u);
}

// 32-thread tree reduction over a workgroup-shared scratch. Driver-portable:
// does NOT depend on any specific subgroup size, unlike subgroupAdd (which is
// only correct when a single 32-lane subgroup covers the whole workgroup).
// On Intel ARL minSubgroupSize=8, so subgroupAdd would silently under-reduce.
shared float partials[32];

float reduce_sum_32(float v, uint tid) {
    partials[tid] = v;
    barrier();
    for (uint stride = 16u; stride > 0u; stride >>= 1) {
        if (tid < stride) {
            partials[tid] += partials[tid + stride];
        }
        barrier();
    }
    return partials[0];
}

void main() {
    uint tid = gl_LocalInvocationID.x;
    uint kv_idx = gl_WorkGroupID.x;
    if (kv_idx >= push.n_kv) return;

    // Each thread owns 4 consecutive elements. That range is fully inside one
    // 32-element block, so hoist the per-block norm + packed code bytes out of
    // the inner loop (mirrors the Metal port and avoids repeated raw byte
    // loads on native Vulkan drivers).
    uint group_byte_off = push.head_offset_bytes + kv_idx * (push.kv_stride_blocks * 14u);
    uint base_elem = tid * 4u;
    uint block = base_elem >> 5;              // 0..3
    uint within0 = base_elem & 31u;           // 0,4,8,...,28
    uint blk_off = group_byte_off + block * 14u;
    float norm = fp16_to_fp32(read_u16(blk_off));
    uint qs_byte = read_byte(blk_off + 2u + (within0 >> 2));
    uint sb = read_byte(blk_off + 2u + 8u + (within0 >> 3));

    float acc = 0.0;
    for (uint local = 0u; local < 4u; local += 1u) {
        uint elem = base_elem + local;        // 0..127
        uint within = within0 + local;        // 0..31 within block
        uint low2 = (qs_byte >> ((within & 3u) * 2u)) & 0x3u;
        uint hi1 = (sb >> (within & 7u)) & 0x1u;

        uint idx = low2 | (hi1 << 2);

        // Centroid * per-block norm; multiply by pre-rotated Q[elem].
        float k_val = TURBO_CENTROIDS_3BIT[idx] * norm;
        float q_val = q[push.q_head * push.head_dim + elem];
        acc = fma(q_val, k_val, acc);
    }

    float sum = reduce_sum_32(acc, tid);
    if (tid == 0u) {
        scores[push.q_head * push.n_kv + kv_idx] = sum;
    }
}