From f5ac635473301c0b9b538cbdc024d753b223be01 Mon Sep 17 00:00:00 2001 From: Jared Van Bortel Date: Thu, 25 Jan 2024 11:27:11 -0500 Subject: [PATCH] kompute : fix q8_0 mmv, 41 -> 28 failures --- ggml-kompute.cpp | 48 +++----------- kompute-shaders/common.comp | 8 +++ kompute-shaders/op_mul_mat_q8_0.comp | 95 ++++++++++++++++------------ 3 files changed, 73 insertions(+), 78 deletions(-) diff --git a/ggml-kompute.cpp b/ggml-kompute.cpp index f45902e36..030cb7a23 100644 --- a/ggml-kompute.cpp +++ b/ggml-kompute.cpp @@ -935,44 +935,6 @@ static void ggml_vk_mul_mat_f16( seq.record(s_algo); } -static void ggml_vk_mul_mat_q8_0(kp::Sequence& seq, - const std::shared_ptr& inA, - const std::shared_ptr& inB, - const std::shared_ptr& out, - uint32_t inAOff, uint32_t inBOff, uint32_t outOff, - int32_t ne00, int32_t ne01, - uint32_t nb01, uint32_t nb02, - int32_t ne11, int32_t ne12, - uint32_t nb11, uint32_t nb12, - int32_t ne0, int32_t ne1) { - const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q8_0_comp_spv, - kp::shader_data::op_mul_mat_q8_0_comp_spv_len); - struct PushConstants { - uint32_t inAOff, inBOff, outOff; - int32_t ne00; - uint32_t nb01, nb02; - uint32_t nb11, nb12; - int32_t ne0, ne1; - } pushConsts { - inAOff, safe_divide(inBOff, 4), safe_divide(outOff, 4), - ne00, nb01, nb02, nb11, nb12, ne0, ne1, - }; - - std::shared_ptr s_algo = nullptr; - if (!komputeManager()->hasAlgorithm(__func__)) { - const uint32_t local_x = ggml_vk_current_device().subgroupSize; - s_algo = komputeManager()->algorithm(__func__, s_kompute_context->pool.get(), {inA, inB, out}, spirv, {unsigned(ne01), unsigned(ne11), unsigned(ne12)}, {local_x}, {pushConsts}); - } else { - s_algo = komputeManager()->getAlgorithm(__func__); - s_algo->setTensors({inA, inB, out}); - s_algo->setWorkgroup({unsigned(ne01), unsigned(ne11), unsigned(ne12)}); - s_algo->setPushConstants({pushConsts}); - s_algo->updateDescriptors(s_kompute_context->pool.get()); - } - seq.record(s_algo); -} - - static void ggml_vk_mul_mat_mat_f32(kp::Sequence& seq, const std::shared_ptr& inA, const std::shared_ptr& inB, @@ -1079,6 +1041,14 @@ static void ggml_vk_mul_mat_q4_1(Args&&... args) { ggml_vk_mul_mat_impl(spirv, "q4_1", 1/*We access blocks unaligned*/, std::forward(args)...); } +template +static void ggml_vk_mul_mat_q8_0(Args&&... args) { + const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q8_0_comp_spv, + kp::shader_data::op_mul_mat_q8_0_comp_spv_len); + + ggml_vk_mul_mat_impl(spirv, "q8_0", 1/*We access blocks unaligned*/, std::forward(args)...); +} + static void ggml_vk_mul_mat_q6_k( kp::Sequence& seq, const std::shared_ptr& inA, @@ -1618,7 +1588,7 @@ void ggml_vk_graph_compute(struct ggml_kompute_context * ctx, struct ggml_cgraph case GGML_TYPE_Q8_0: ggml_vk_mul_mat_q8_0( seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst, - ne00, ne01, nb01, nb02, ne11, ne12, nb11, nb12, ne0, ne1 + ne00, ne01, ne02, ne10, ne11, ne12, ne13, ne0, ne1, r2, r3 ); break; case GGML_TYPE_Q4_0: diff --git a/kompute-shaders/common.comp b/kompute-shaders/common.comp index 0df6db7d0..0918657d5 100644 --- a/kompute-shaders/common.comp +++ b/kompute-shaders/common.comp @@ -95,3 +95,11 @@ mat4 dequantize_q6_k(const block_q6_k xb, uint il) { } return reg; } + + +#define QK8_0 32 +// struct block_q8_0 { +// float16_t d; // delta +// int8_t qs[QK8_0]; // quants +// }; +#define sizeof_block_q8_0 34 diff --git a/kompute-shaders/op_mul_mat_q8_0.comp b/kompute-shaders/op_mul_mat_q8_0.comp index 1c4ddbb08..34d015e90 100644 --- a/kompute-shaders/op_mul_mat_q8_0.comp +++ b/kompute-shaders/op_mul_mat_q8_0.comp @@ -2,55 +2,72 @@ #include "common.comp" -#define BLOCKS_IN_QUANT QK8_0 -#define SIZE_OF_BLOCK sizeof_block_q8_0 -#define N_ROWS 4 +#include "op_mul_mv_q_n_pre.comp" -layout(local_size_x_id = 0) in; -layout(local_size_y = 1) in; -layout(local_size_z = 1) in; - -layout (binding = 0) readonly buffer tensorInA { uint8_t inA[]; }; -layout (binding = 1) readonly buffer tensorInB { float inB[]; }; -layout (binding = 2) writeonly buffer tensorOut { float out_[]; }; - -layout (push_constant) uniform parameter { - uint inAOff; - uint inBOff; - uint outOff; - int ne00; - int ne10; - int ne0; - int ne1; - int ne01; - int gqa; -} pcs; - -#define ELS_PER_BLOCK 32 #define SIZE_OF_D 2 -#define BLOCK_SIZE (ELS_PER_BLOCK + SIZE_OF_D) + +#define N_DST 4 // each SIMD group works on 4 rows +#define N_SIMDGROUP 2 // number of SIMD groups in a thread group +#define N_SIMDWIDTH 32 // assuming SIMD group size is 32 + +#define NB_Q8_0 8 void main() { + // NB: hack to make compatible with AMD GPUs that have a subgroup size of 64 + if (gl_SubgroupInvocationID > 31) + return; + + const int nr = N_DST; + const int nsg = N_SIMDGROUP; + const int nw = N_SIMDWIDTH; + + const int nb = pcs.ne00/QK8_0; const uint r0 = gl_WorkGroupID.x; const uint r1 = gl_WorkGroupID.y; const uint im = gl_WorkGroupID.z; - const uint x = r0 * (pcs.ne00/ELS_PER_BLOCK) * BLOCK_SIZE + pcs.inAOff; // Based from inA - const uint y = r1 * pcs.ne10 + pcs.inBOff; // based from inB + const uint first_row = (r0 * nsg + gl_SubgroupID) * nr; - float sumf = 0.0f; - for (uint i = gl_SubgroupInvocationID.x; i < pcs.ne00; i += gl_SubgroupSize) { - const uint block_number = i / ELS_PER_BLOCK; - const uint block_offset = block_number * BLOCK_SIZE; - const float d = u8BufToFloat16(inA, x + block_offset); - const uint position_in_block = i % ELS_PER_BLOCK; - const int q = int8_t(inA[x+block_offset+SIZE_OF_D+position_in_block]); - const float dq = d * q; - sumf += dq * float(inB[y+i]); + const uint i12 = im%pcs.ne12; + const uint i13 = im/pcs.ne12; + + const uint offset0 = first_row * nb + (i12/pcs.r2)*(nb*pcs.ne01) + (i13/pcs.r3)*(nb*pcs.ne01*pcs.ne02); + + const uint x = offset0*sizeof_block_q8_0 + pcs.inAOff; // Based from inA + const uint y = r1*pcs.ne10 + im*pcs.ne00*pcs.ne1 + pcs.inBOff; // based from inB + + float yl[NB_Q8_0]; + float sumf[N_DST]={0.f, 0.f, 0.f, 0.f}; + + const uint ix = gl_SubgroupInvocationID.x/4; + const uint il = gl_SubgroupInvocationID.x%4; + + uint yb = y + ix * QK8_0 + NB_Q8_0*il; + + // each thread in a SIMD group deals with NB_Q8_0 quants at a time + for (uint ib = ix; ib < nb; ib += nw/4) { + for (int i = 0; i < NB_Q8_0; ++i) { + yl[i] = inB[yb + i]; + } + + for (int row = 0; row < nr; row++) { + const uint block_offset = (ib+row*nb) * sizeof_block_q8_0; + float sumq = 0.f; + for (int iq = 0; iq < NB_Q8_0; ++iq) { + const int8_t qs_iq = int8_t(inA[x + block_offset + SIZE_OF_D + NB_Q8_0*il + iq]); + sumq += qs_iq * yl[iq]; + } + const float16_t d = u8BufToFloat16(inA, x + block_offset); + sumf[row] += sumq*d; + } + + yb += NB_Q8_0 * nw; } - const float all_sum = subgroupAdd(sumf); - if (subgroupElect()) { - out_[im*pcs.ne1*pcs.ne0 + r1*pcs.ne0 + r0 + pcs.outOff] = all_sum; + for (int row = 0; row < nr; ++row) { + const float tot = subgroupAdd(sumf[row]); + if (subgroupElect() && first_row + row < pcs.ne01) { + out_[r1*pcs.ne0 + im*pcs.ne0*pcs.ne1 + first_row + row] = tot; + } } }