diff --git a/CMakeLists.txt b/CMakeLists.txt index fbbb46bbf..df6b53dce 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -482,6 +482,7 @@ if (LLAMA_KOMPUTE) kompute/op_mul_mat_mat_f16.comp kompute/op_mul_mat_mat_f32.comp kompute/op_mul_mat_mat_q4_0.comp + kompute/op_mul_mat_mat_q4_1.comp kompute/op_mul_mat_mat_q8_0.comp kompute/op_mul_mat_mat_q6_k.comp kompute/op_mul_mat_f16.comp @@ -517,6 +518,7 @@ if (LLAMA_KOMPUTE) shaderop_mul_mat_mat_f16.h shaderop_mul_mat_mat_f32.h shaderop_mul_mat_mat_q4_0.h + shaderop_mul_mat_mat_q4_1.h shaderop_mul_mat_mat_q8_0.h shaderop_mul_mat_mat_q6_k.h shaderop_mul_mat_f16.h diff --git a/ggml-vulkan.cpp b/ggml-vulkan.cpp index 488683ec3..56f15310d 100644 --- a/ggml-vulkan.cpp +++ b/ggml-vulkan.cpp @@ -30,6 +30,7 @@ #include "shaderop_mul_mat_mat_f32.h" #include "shaderop_mul_mat_mat_f16.h" #include "shaderop_mul_mat_mat_q4_0.h" +#include "shaderop_mul_mat_mat_q4_1.h" #include "shaderop_mul_mat_mat_q8_0.h" #include "shaderop_mul_mat_mat_q6_k.h" #include "shaderop_getrows_f16.h" @@ -1214,6 +1215,14 @@ void ggml_vk_mul_mat_mat_q4_0(Args&&... args) { ggml_vk_mul_mat_mat_q4_x(spirv, std::forward(args)...); } +template +void ggml_vk_mul_mat_mat_q4_1(Args&&... args) { + const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_mat_q4_1_comp_spv, + kp::shader_data::op_mul_mat_mat_q4_1_comp_spv_len); + + ggml_vk_mul_mat_mat_q4_x(spirv, std::forward(args)...); +} + void ggml_vk_mul_mat_q4_x(const std::vector& spirv, uint32_t block_size, kp::Sequence& seq, const std::shared_ptr& inA, const std::shared_ptr& inB, @@ -1659,6 +1668,16 @@ void ggml_vk_graph_compute(struct ggml_kompute_context * ctx, struct ggml_cgraph nb11, nb12, nb1, nb2); break; + case GGML_TYPE_Q4_1: + ggml_vk_mul_mat_mat_q4_1(seq, + id_src0, id_src1, id_dst, + off_src0, off_src1, off_dst, + ne00, ne01, ne02, + nb01, nb02, + ne11, ne12, + nb11, nb12, + nb1, nb2); + break; case GGML_TYPE_Q8_0: ggml_vk_mul_mat_mat_q8_0(seq, id_src0, id_src1, id_dst, diff --git a/kompute/op_mul_mat_mat_q4_1.comp b/kompute/op_mul_mat_mat_q4_1.comp new file mode 100644 index 000000000..d7fbc96db --- /dev/null +++ b/kompute/op_mul_mat_mat_q4_1.comp @@ -0,0 +1,73 @@ +/** + * Copyright (c) 2023 Nomic, Inc. All rights reserved. + * + * This software is licensed under the terms of the Software for Open Models + * License (SOM), version 1.0, as detailed in the LICENSE_SOM.txt file. A copy + * of this license should accompany this software. Except as expressly granted + * in the SOM license, all rights are reserved by Nomic, Inc. + */ + +#version 450 + +#include "common.comp" + +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_EXT_debug_printf : enable + +layout(local_size_x = 32) 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 ne01; + int ne02; + int ne11; + int ne12; + uint nb01; + uint nb02; + uint nb11; + uint nb12; + uint nb1; + uint nb2; +} +pcs; + +#define ELS_PER_BLOCK 32 +#define M_OFFSET 2 +#define QS_OFFSET 4 +#define BLOCK_SIZE ((ELS_PER_BLOCK / 2) + QS_OFFSET) + +void main() { + uvec3 gid = gl_GlobalInvocationID; + + uint bc_ab = pcs.ne12 > pcs.ne02 ? gid.z / (pcs.ne12 / pcs.ne02) : gid.z; + uint bc_ba = pcs.ne02 > pcs.ne12 ? gid.z / (pcs.ne02 / pcs.ne12) : gid.z; + + + const uint x = (gid.x*pcs.nb01 + bc_ab*pcs.nb02) + pcs.inAOff; // Based from inA + const uint y = (gid.y*pcs.nb11 + bc_ba*pcs.nb12) / 4 + pcs.inBOff; // based from inB + float sum = 0.0f; + for (uint i = 0; i < pcs.ne00; i+=ELS_PER_BLOCK) { + for (uint j = 0; j < ELS_PER_BLOCK / 2; j++) { + 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 float m = u8BufToFloat16(inA, x + block_offset + M_OFFSET); + const uint byte_position_in_block = j; + const int q0 = (inA[x+block_offset+QS_OFFSET+byte_position_in_block] & 0x0F); + const int q1 = (inA[x+block_offset+QS_OFFSET+byte_position_in_block] >> 4); + const float dq0 = (d * q0) + m; + const float dq1 = (d * q1) + m; + sum += (dq0 * float(inB[y+i+j])) + \ + (dq1 * float(inB[y+i+j+(ELS_PER_BLOCK/2)])); + } + } + + out_[gid.z*(pcs.nb2/4) + gid.y*(pcs.nb1/4) + gid.x + pcs.outOff] = sum; +}