offgrid-sys 0.2.1

#version 450

#extension GL_EXT_control_flow_attributes : require
#extension GL_EXT_shader_16bit_storage : require
#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require

#define BINDING_IDX_K 0u

#include "types.glsl"
#include "fa_types.glsl"
#define FaTypeV GGML_TYPE_F32

layout(constant_id = 0) const uint FaTypeK = GGML_TYPE_F32;
layout(constant_id = 1) const uint FaBlockBytesK = 4;

#include "flash_attn_dequant.glsl"

// one workgroup scores a tile of BK keys against BT tokens: the keys are staged once,
// the queries one head at a time, and each invocation owns one key for TPI tokens
#define HEAD_SIZE 128
#define WG_SIZE   256
#define BK        64
#define BT        8
#define TS        (WG_SIZE / BK)
#define TPI       (BT / TS)

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

layout(binding = 0) readonly buffer QBuf { float q[]; };
layout(binding = 1) readonly buffer KBufF16 { float16_t k_f16[]; };
layout(binding = 1) readonly buffer KBufF32 { float k_f32[]; };
layout(binding = 1) readonly buffer KBufBF16 { uint16_t k_bf16[]; };
layout(binding = 2) readonly buffer WBuf { float weights[]; };
layout(binding = 3) readonly buffer MBuf { float16_t mask[]; };
layout(binding = 4) writeonly buffer DstBuf { float dst[]; };

layout(push_constant) uniform PushConstants {
    uint n_kv;
    uint n_heads;
    uint n_tokens;
    uint n_masks;
    uint q_nb1;
    uint q_nb2;
    uint q_nb3;
    uint k_nb2;
    uint k_nb3;
    uint w_nb1;
    uint w_nb3;
    uint m_nb1;
    uint m_nb3;
    uint d_nb1;
    uint d_nb3;
};

// the row padding keeps the keys of consecutive invocations in distinct banks
shared f16vec4 k_tile[BK][HEAD_SIZE / 4 + 1];
shared f16vec4 q_tile[BT][HEAD_SIZE / 4];
shared float   w_tile[BT];

void main() {
    const uint tid = gl_LocalInvocationID.x;
    const uint ik0 = gl_WorkGroupID.x * BK;
    const uint t0  = gl_WorkGroupID.y * BT;
    const uint s   = gl_WorkGroupID.z;

    if (fa_type_needs_shmem(FaTypeK)) {
        init_iq_shmem(gl_WorkGroupSize);
    }

    // k strides come in as bytes, so scale them down to the view being indexed
    const uint k_block_elems = fa_block_elems(FaTypeK);
    const uint k_elem_bytes = FaBlockBytesK / k_block_elems;

    // stage the key tile four elements at a time, rows past n_kv are zero
    [[unroll]] for (uint i = tid; i < BK * HEAD_SIZE / 4; i += WG_SIZE) {
        const uint r = i / (HEAD_SIZE / 4);
        const uint c4 = i % (HEAD_SIZE / 4);

        vec4 v = vec4(0.0);
        if (ik0 + r < n_kv) {
            const uint k_offset = (ik0 + r) * k_nb2 + s * k_nb3;
            const uint e = k_offset / k_elem_bytes + c4 * 4;

            if (FaTypeK == GGML_TYPE_F16) {
                v = vec4(k_f16[e], k_f16[e + 1], k_f16[e + 2], k_f16[e + 3]);
            } else if (FaTypeK == GGML_TYPE_F32) {
                v = vec4(k_f32[e], k_f32[e + 1], k_f32[e + 2], k_f32[e + 3]);
            } else if (FaTypeK == GGML_TYPE_BF16) {
                v = bf16_to_fp32(uvec4(k_bf16[e], k_bf16[e + 1], k_bf16[e + 2], k_bf16[e + 3]));
            } else {
                v = dequantize4((c4 * 4) / k_block_elems, (c4 * 4) % k_block_elems, k_offset / FaBlockBytesK, BINDING_IDX_K);
            }
        }

        k_tile[r][c4] = f16vec4(v);
    }

    const uint kl = tid % BK;
    const uint tl = tid / BK;

    float score[TPI];
    [[unroll]] for (uint j = 0; j < TPI; ++j) {
        score[j] = 0.0;
    }

    for (uint h = 0; h < n_heads; ++h) {
        // the previous head is fully consumed and, on the first pass, the key tile is complete
        barrier();

        [[unroll]] for (uint i = tid; i < BT * HEAD_SIZE / 4; i += WG_SIZE) {
            const uint r = i / (HEAD_SIZE / 4);
            const uint c4 = i % (HEAD_SIZE / 4);
            const uint t = t0 + r;

            vec4 v = vec4(0.0);
            if (t < n_tokens) {
                const uint q_base = h * q_nb1 + t * q_nb2 + s * q_nb3 + c4 * 4;
                v = vec4(q[q_base], q[q_base + 1], q[q_base + 2], q[q_base + 3]);
            }
            q_tile[r][c4] = f16vec4(v);
        }

        if (tid < BT) {
            const uint t = t0 + tid;
            w_tile[tid] = t < n_tokens ? weights[h + t * w_nb1 + s * w_nb3] : 0.0;
        }

        barrier();

        float qk[TPI];
        [[unroll]] for (uint j = 0; j < TPI; ++j) {
            qk[j] = 0.0;
        }

        [[unroll]] for (uint c4 = 0; c4 < HEAD_SIZE / 4; ++c4) {
            const f16vec4 kv = k_tile[kl][c4];
            [[unroll]] for (uint j = 0; j < TPI; ++j) {
                const f16vec4 qv = q_tile[tl + j * TS][c4];
                qk[j] += float(dot(kv, qv));
            }
        }

        [[unroll]] for (uint j = 0; j < TPI; ++j) {
            score[j] += max(qk[j], 0.0) * w_tile[tl + j * TS];
        }
    }

    const uint ik = ik0 + kl;
    if (ik >= n_kv) {
        return;
    }

    [[unroll]] for (uint j = 0; j < TPI; ++j) {
        const uint t = t0 + tl + j * TS;
        if (t < n_tokens) {
            const uint mask_offset = ik + t * m_nb1 + (s % n_masks) * m_nb3;
            dst[ik + t * d_nb1 + s * d_nb3] = score[j] + float(mask[mask_offset]);
        }
    }
}