mirror of
https://github.com/ggerganov/llama.cpp.git
synced 2024-12-29 04:44:34 +00:00
q6k mm works
This commit is contained in:
Aaron Miller
cebtenzzre
parent
d5741c07a5
commit
b78a94bc6d
@ -483,6 +483,7 @@ if (LLAMA_KOMPUTE)
|
||||
kompute/op_mul_mat_mat_f32.comp
|
||||
kompute/op_mul_mat_mat_q4_0.comp
|
||||
kompute/op_mul_mat_mat_q8_0.comp
|
||||
kompute/op_mul_mat_mat_q6_k.comp
|
||||
kompute/op_mul_mat_f16.comp
|
||||
kompute/op_mul_mat_q8_0.comp
|
||||
kompute/op_mul_mat_q4_0.comp
|
||||
@ -517,6 +518,7 @@ if (LLAMA_KOMPUTE)
|
||||
shaderop_mul_mat_mat_f32.h
|
||||
shaderop_mul_mat_mat_q4_0.h
|
||||
shaderop_mul_mat_mat_q8_0.h
|
||||
shaderop_mul_mat_mat_q6_k.h
|
||||
shaderop_mul_mat_f16.h
|
||||
shaderop_mul_mat_q8_0.h
|
||||
shaderop_mul_mat_q4_0.h
|
||||
|
||||
@ -31,6 +31,7 @@
|
||||
#include "shaderop_mul_mat_mat_f16.h"
|
||||
#include "shaderop_mul_mat_mat_q4_0.h"
|
||||
#include "shaderop_mul_mat_mat_q8_0.h"
|
||||
#include "shaderop_mul_mat_mat_q6_k.h"
|
||||
#include "shaderop_getrows_f16.h"
|
||||
#include "shaderop_getrows_q4_0.h"
|
||||
#include "shaderop_getrows_q4_1.h"
|
||||
@ -1109,6 +1110,54 @@ void ggml_vk_mul_mat_mat_q8_0(
|
||||
seq.record<kp::OpAlgoDispatch>(s_algo);
|
||||
}
|
||||
|
||||
void ggml_vk_mul_mat_mat_q6_k(
|
||||
kp::Sequence& seq,
|
||||
const std::shared_ptr<kp::Tensor>& inA,
|
||||
const std::shared_ptr<kp::Tensor>& inB,
|
||||
const std::shared_ptr<kp::Tensor>& out,
|
||||
uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
|
||||
int32_t ne00, int32_t ne01, int32_t ne02,
|
||||
uint32_t nb01, uint32_t nb02,
|
||||
int32_t ne11, int32_t ne12,
|
||||
uint32_t nb11, uint32_t nb12,
|
||||
uint32_t nb1, uint32_t nb2) {
|
||||
const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_mat_q6_k_comp_spv,
|
||||
kp::shader_data::op_mul_mat_mat_q6_k_comp_spv_len);
|
||||
struct PushConstants {
|
||||
uint32_t inAOff, inBOff, outOff;
|
||||
int32_t ne00, ne01, ne02, ne11, ne12;
|
||||
uint32_t nb01, nb02;
|
||||
uint32_t nb11, nb12;
|
||||
uint32_t nb1, nb2;
|
||||
} pushConsts {
|
||||
inAOff, safe_divide(inBOff, 4), safe_divide(outOff, 4),
|
||||
ne00, ne01, ne02, ne11, ne12,
|
||||
nb01, nb02, nb11, nb12,
|
||||
nb1, nb2
|
||||
};
|
||||
|
||||
std::shared_ptr<kp::Algorithm> s_algo = nullptr;
|
||||
if (!komputeManager()->hasAlgorithm(__func__)) {
|
||||
s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(),
|
||||
{inA, inB, out}, spirv,
|
||||
{unsigned(ne01)/32,
|
||||
unsigned(ne11),
|
||||
unsigned(std::max(ne12, ne02))
|
||||
},
|
||||
{},
|
||||
{pushConsts});
|
||||
} else {
|
||||
s_algo = komputeManager()->getAlgorithm(__func__);
|
||||
s_algo->setTensors({inA, inB, out});
|
||||
s_algo->setWorkgroup({unsigned(ne01)/32,
|
||||
unsigned(ne11),
|
||||
unsigned(std::max(ne12, ne02)),
|
||||
});
|
||||
s_algo->setPushConstants<PushConstants>({pushConsts});
|
||||
s_algo->updateDescriptors(s_kompute_context->pool.get());
|
||||
}
|
||||
seq.record<kp::OpAlgoDispatch>(s_algo);
|
||||
}
|
||||
|
||||
void ggml_vk_mul_mat_mat_q4_x(const std::vector<uint32_t>& spirv,
|
||||
kp::Sequence& seq,
|
||||
@ -1138,7 +1187,7 @@ void ggml_vk_mul_mat_mat_q4_x(const std::vector<uint32_t>& spirv,
|
||||
if (!komputeManager()->hasAlgorithm(__func__)) {
|
||||
s_algo = komputeManager()->algorithm<float, PushConstants>(__func__, s_kompute_context->pool.get(),
|
||||
{inA, inB, out}, spirv,
|
||||
{unsigned(ne01),
|
||||
{unsigned(ne01)/32,
|
||||
unsigned(ne11),
|
||||
unsigned(std::max(ne12, ne02))},
|
||||
{},
|
||||
@ -1619,6 +1668,16 @@ void ggml_vk_graph_compute(struct ggml_kompute_context * ctx, struct ggml_cgraph
|
||||
ne11, ne12,
|
||||
nb11, nb12,
|
||||
nb1, nb2);
|
||||
break;
|
||||
case GGML_TYPE_Q6_K:
|
||||
ggml_vk_mul_mat_mat_q6_k(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;
|
||||
default: {
|
||||
fprintf(stderr, "%s: %s: Unsupported quantization for M*M: %u/%u\n", __func__, ggml_op_name(dst->op), src0t, src1t);
|
||||
|
||||
88
kompute/op_mul_mat_mat_q6_k.comp
Normal file
88
kompute/op_mul_mat_mat_q6_k.comp
Normal file
@ -0,0 +1,88 @@
|
||||
/**
|
||||
* 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 256 //QK_K
|
||||
#define QH_OFFSET (ELS_PER_BLOCK / 2)
|
||||
#define QSCALES_OFFSET (QH_OFFSET + (ELS_PER_BLOCK / 4))
|
||||
#define SCALE_SCALE_OFFSET (QSCALES_OFFSET + (ELS_PER_BLOCK / 16))
|
||||
#define BLOCK_SIZE (SCALE_SCALE_OFFSET + 2)
|
||||
|
||||
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;
|
||||
const uint n_blocks = pcs.ne00 / ELS_PER_BLOCK;
|
||||
// this is pretty much all lifted right from dequantize_row_q6_K
|
||||
uint outoff = 0;
|
||||
for (uint i = 0; i < n_blocks; i++) {
|
||||
const uint block_number = i;
|
||||
const uint block_offset = block_number * BLOCK_SIZE;
|
||||
const float scales_d = u8BufToFloat16(inA, x + block_offset + SCALE_SCALE_OFFSET);
|
||||
uint qloff = block_offset;
|
||||
uint qhoff = block_offset + QH_OFFSET;
|
||||
uint scoff = block_offset + QSCALES_OFFSET;
|
||||
for (int n = 0; n < 256; n += 128) {
|
||||
for (int l = 0; l < 32; ++l) {
|
||||
int is = l/16;
|
||||
const int q1 = int((inA[x + qloff + l + 0] & 0xF) | (((inA[x + qhoff + l] >> 0) & 3) << 4)) - 32;
|
||||
const int q2 = int((inA[x + qloff + l + 32] & 0xF) | (((inA[x + qhoff + l] >> 2) & 3) << 4)) - 32;
|
||||
const int q3 = int((inA[x + qloff + l + 0] >> 4) | (((inA[x + qhoff + l] >> 4) & 3) << 4)) - 32;
|
||||
const int q4 = int((inA[x + qloff + l + 32] >> 4) | (((inA[x + qhoff + l] >> 6) & 3) << 4)) - 32;
|
||||
sum += inB[y + outoff + l + 0] * scales_d * int8_t(inA[x + scoff + is + 0]) * q1;
|
||||
sum += inB[y + outoff + l + 32] * scales_d * int8_t(inA[x + scoff + is + 2]) * q2;
|
||||
sum += inB[y + outoff + l + 64] * scales_d * int8_t(inA[x + scoff + is + 4]) * q3;
|
||||
sum += inB[y + outoff + l + 96] * scales_d * int8_t(inA[x + scoff + is + 6]) * q4;
|
||||
}
|
||||
outoff += 128;
|
||||
qloff += 64;
|
||||
qhoff += 32;
|
||||
scoff += 8;
|
||||
}
|
||||
}
|
||||
|
||||
out_[gid.z*(pcs.nb2/4) + gid.y*(pcs.nb1/4) + gid.x + pcs.outOff] = sum;
|
||||
}
|
||||
Loading…
Reference in New Issue
Block a user