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

// Multi-token-per-workgroup variant of qjl.comp (the QJL attention-score hot
// path). Identical math; one workgroup processes TOKENS_PER_WG consecutive
// tokens along the t axis serially in a 32-thread loop, amortising launch tax.
// TOKENS_PER_WG is a SPIR-V specialization constant (constant_id = 0); default
// 1 == qjl.comp behaviour. local_size stays 32 (portable shared-memory tree
// reduction, no subgroup assumptions).
//
// Dispatch shape:
//     grid = (n_heads, ceil(n_tokens / TOKENS_PER_WG), 1)
//     tg   = (32, 1, 1)

#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(constant_id = 0) const uint TOKENS_PER_WG = 1u;

layout(std430, binding = 0) readonly buffer QSketch { float q_sketch[]; };
layout(std430, binding = 1) readonly buffer KPacked { uint k_packed[]; };
layout(std430, binding = 2) writeonly buffer ScoresOut { float scores[]; };

layout(push_constant) uniform Push {
    uint n_heads;
    uint n_kv_heads;
    uint n_tokens;
    uint proj_dim;       // must equal 256
} push;

const uint QJL_PROJECTION_DIM = 256u;
const uint QJL_PACKED_BYTES   = 32u;
const uint QJL_BLOCK_BYTES    = 34u;
const float QJL_SCORE_SCALE = 1.2533141373155003 / 256.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 qjl_bf16_to_fp32(uint b16) {
    return uintBitsToFloat(b16 << 16);
}

shared float reduce_scratch[32];

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

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

    uint gqa = push.n_heads / push.n_kv_heads;
    uint h_k = h_q / gqa;

    uint byte_i = tid;
    uint base   = byte_i * 8u;
    uint q_off  = h_q * QJL_PROJECTION_DIM + base;
    vec4 q0 = vec4(q_sketch[q_off + 0u], q_sketch[q_off + 1u],
                   q_sketch[q_off + 2u], q_sketch[q_off + 3u]);
    vec4 q1 = vec4(q_sketch[q_off + 4u], q_sketch[q_off + 5u],
                   q_sketch[q_off + 6u], q_sketch[q_off + 7u]);

    uint t_base = gl_WorkGroupID.y * TOKENS_PER_WG;
    for (uint b = 0u; b < TOKENS_PER_WG; ++b) {
        uint t = t_base + b;
        if (t >= push.n_tokens) return;

        uint blk_off = (h_k * push.n_tokens + t) * QJL_BLOCK_BYTES;
        uint bits = read_byte(blk_off + byte_i);
        vec4 s0 = vec4(
            float(int(((bits >> 0) & 1u) << 1) - 1),
            float(int(((bits >> 1) & 1u) << 1) - 1),
            float(int(((bits >> 2) & 1u) << 1) - 1),
            float(int(((bits >> 3) & 1u) << 1) - 1));
        vec4 s1 = vec4(
            float(int(((bits >> 4) & 1u) << 1) - 1),
            float(int(((bits >> 5) & 1u) << 1) - 1),
            float(int(((bits >> 6) & 1u) << 1) - 1),
            float(int(((bits >> 7) & 1u) << 1) - 1));
        vec4 acc4 = fma(q0, s0, fma(q1, s1, vec4(0.0)));
        float acc = acc4.x + acc4.y + acc4.z + acc4.w;

        float sum = reduce_sum_32(acc, tid);
        if (tid == 0u) {
            uint norm16 = read_u16(blk_off + QJL_PACKED_BYTES);
            float norm_k = qjl_bf16_to_fp32(norm16);
            scores[h_q * push.n_tokens + t] = QJL_SCORE_SCALE * norm_k * sum;
        }
    }
}