elizaos--eliza/plugins/plugin-local-inference/native/vulkan/polar_get_rows.comp

// DRAFT: COMPILE-VERIFIED on Linux. Sister kernel to polar.comp.
// Wave-4-D: match polar.comp's threadgroup-cooperative residual add and
// Hadamard butterfly so get_rows fallback has the same parallel shape as the
// dot-product path.
// Ports `kernel_get_rows_q4_polar` from metal/polar.metal — decode one
// PolarQuant block to fp32 head_dim vector.
//
// Same 5-step decode as polar.comp's mul_mv but writes the dequantized
// vector directly instead of dotting against q.

#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;

layout(std430, binding = 0) readonly buffer KBlock {
    uint k_packed[];     // single 82-byte block
};

layout(std430, binding = 1) writeonly buffer Out {
    float out_buf[];     // QK_POLAR floats
};

layout(push_constant) uniform Push {
    uint head_dim;       // must equal QK_POLAR (128)
    uint use_qjl;        // 0 / 1
} push;

const uint  QK_POLAR             = 128u;
const float POLAR_QJL_CORRECTION_MAGNITUDE = 0.5;
const float POLAR_QJL_INV_SQRT_QK = 0.08838834764831845;
const float POLAR_INV_QK = 1.0 / 128.0;

const float POLAR_Q4_CENTROIDS[16] = float[16](
    -2.754354807, -2.093562707, -1.643041510, -1.279739752,
    -0.962640978, -0.672392117, -0.397897103, -0.131757782,
     0.131757782,  0.397897103,  0.672392117,  0.962640978,
     1.279739752,  1.643041510,  2.093562707,  2.754354807
);

const float POLAR_QJL_SIGNS[128] = float[128](
    -1.0, -1.0,  1.0, -1.0, -1.0, -1.0, -1.0, -1.0,
     1.0,  1.0, -1.0,  1.0, -1.0, -1.0,  1.0,  1.0,
    -1.0, -1.0,  1.0, -1.0,  1.0,  1.0, -1.0, -1.0,
    -1.0,  1.0, -1.0,  1.0,  1.0,  1.0, -1.0, -1.0,
    -1.0, -1.0,  1.0, -1.0,  1.0, -1.0,  1.0, -1.0,
    -1.0, -1.0, -1.0,  1.0, -1.0,  1.0,  1.0,  1.0,
     1.0,  1.0, -1.0,  1.0, -1.0, -1.0,  1.0,  1.0,
     1.0,  1.0, -1.0, -1.0, -1.0,  1.0, -1.0,  1.0,
     1.0, -1.0,  1.0, -1.0,  1.0,  1.0,  1.0,  1.0,
    -1.0, -1.0, -1.0, -1.0,  1.0, -1.0, -1.0, -1.0,
     1.0, -1.0, -1.0,  1.0,  1.0,  1.0,  1.0, -1.0,
    -1.0,  1.0, -1.0,  1.0,  1.0, -1.0,  1.0,  1.0,
     1.0, -1.0, -1.0, -1.0,  1.0,  1.0, -1.0,  1.0,
    -1.0,  1.0,  1.0, -1.0, -1.0,  1.0, -1.0, -1.0,
     1.0, -1.0,  1.0, -1.0, -1.0,  1.0, -1.0, -1.0,
     1.0,  1.0,  1.0, -1.0,  1.0, -1.0, -1.0,  1.0
);

uint read_byte(uint b) {
    uint w = k_packed[b >> 2];
    return (w >> ((b & 3u) * 8u)) & 0xFFu;
}

uint read_u16(uint b) {
    return read_byte(b) | (read_byte(b + 1u) << 8u);
}

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 {
            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);
}

shared float buf[128];

void polar_hadamard_inplace_tg32(uint tid) {
    for (uint h = 1u; h < QK_POLAR; h <<= 1) {
        uint p0 = tid;
        uint p1 = tid + 32u;
        uint twoh = h << 1;
        uint b0 = (p0 / h) * twoh;
        uint o0 = p0 - (p0 / h) * h;
        uint b1 = (p1 / h) * twoh;
        uint o1 = p1 - (p1 / h) * h;
        uint j0 = b0 + o0;
        uint j1 = b1 + o1;
        float a0 = buf[j0];
        float c0 = buf[j0 + h];
        float a1 = buf[j1];
        float c1 = buf[j1 + h];
        buf[j0]     = a0 + c0;
        buf[j0 + h] = a0 - c0;
        buf[j1]     = a1 + c1;
        buf[j1 + h] = a1 - c1;
        barrier();
    }
}

void main() {
    uint tid = gl_LocalInvocationID.x;
    if (push.head_dim != QK_POLAR) return;

    // Step 1+2: unpack codes.
    for (uint b = tid; b < QK_POLAR / 2u; b += 32u) {
        uint byte = read_byte(2u + b);
        buf[2u * b]      = POLAR_Q4_CENTROIDS[byte & 0x0Fu];
        buf[2u * b + 1u] = POLAR_Q4_CENTROIDS[(byte >> 4) & 0x0Fu];
    }
    barrier();

    if (push.use_qjl != 0u) {
        uint  bit    = read_byte(2u + 64u) & 1u;
        float sign_v = bit != 0u ? 1.0 : -1.0;
        float mag    = POLAR_QJL_CORRECTION_MAGNITUDE * POLAR_QJL_INV_SQRT_QK;
        float scaled = sign_v * mag;
        for (uint i = tid; i < QK_POLAR; i += 32u) {
            buf[i] += scaled * POLAR_QJL_SIGNS[i];
        }
        barrier();
    }

    polar_hadamard_inplace_tg32(tid);

    uint norm16 = read_u16(0u);
    float scale = fp16_to_fp32(norm16) * POLAR_INV_QK;
    for (uint i = tid; i < QK_POLAR; i += 32u) {
        out_buf[i] = buf[i] * scale;
    }
}