kompute: op_mul_mat_q4_k permutted support

Signed-off-by: Sergio Lopez <slp@redhat.com>
This commit is contained in:
Sergio Lopez 2024-11-25 17:52:52 +01:00
parent 1b8afa88dc
commit 2ac1d0ee3e
2 changed files with 28 additions and 13 deletions

View file

@ -1079,19 +1079,26 @@ static void ggml_vk_mul_mat_q4_k(
const std::shared_ptr<kp::Tensor>& inB, const std::shared_ptr<kp::Tensor>& inB,
const std::shared_ptr<kp::Tensor>& out, const std::shared_ptr<kp::Tensor>& out,
uint32_t inAOff, uint32_t inBOff, uint32_t outOff, uint32_t inAOff, uint32_t inBOff, uint32_t outOff,
int32_t ne00, int32_t ne01, int32_t ne02, int32_t ne10, int32_t ne00, int32_t ne01, int32_t ne02,
int32_t ne11, int32_t ne12, int32_t ne13, int32_t ne0, int32_t ne10, int32_t ne11, int32_t ne12, int32_t ne13,
int32_t ne1, int32_t r2, int32_t r3 int32_t ne0, int32_t ne1,
uint32_t nb01, uint32_t nb02, uint32_t nb03,
uint32_t nb11, uint32_t nb12, uint32_t nb13,
uint32_t r2, uint32_t r3
) { ) {
const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q4_k_comp_spv, const static auto spirv = getSpirvShader(kp::shader_data::op_mul_mat_q4_k_comp_spv,
kp::shader_data::op_mul_mat_q4_k_comp_spv_len); kp::shader_data::op_mul_mat_q4_k_comp_spv_len);
struct PushConstants { struct PushConstants {
uint32_t inAOff, inBOff, outOff; uint32_t inAOff, inBOff, outOff;
int32_t ne00, ne10, ne0, ne1, ne01, ne02, ne12, r2, r3; int32_t ne00, ne10, ne0, ne1, ne01, ne02, ne12;
uint32_t nb01, nb02, nb03, nb11, nb12, nb13;
uint32_t r2, r3;
} pushConsts { } pushConsts {
0, 0, 0, inAOff, safe_divide(inBOff, 4), safe_divide(outOff, 4),
ne00, ne10, ne0, ne1, ne01, ne02, ne12, r2, r3 ne00, ne10, ne0, ne1, ne01, ne02, ne12,
nb01, nb02, nb03, nb11, nb12, nb13,
r2, r3
}; };
std::shared_ptr<kp::Algorithm> s_algo = nullptr; std::shared_ptr<kp::Algorithm> s_algo = nullptr;
@ -1705,7 +1712,8 @@ static void ggml_vk_graph_compute(struct ggml_kompute_context * ctx, struct ggml
case GGML_TYPE_Q4_K: case GGML_TYPE_Q4_K:
ggml_vk_mul_mat_q4_k( ggml_vk_mul_mat_q4_k(
seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst, seq, id_src0, id_src1, id_dst, off_src0, off_src1, off_dst,
ne00, ne01, ne02, ne10, ne11, ne12, ne13, ne0, ne1, ne12/ne02, ne13/ne03 ne00, ne01, ne02, ne10, ne11, ne12, ne13, ne0, ne1,
nb01, nb02, nb03, nb11, nb12, nb13, r2, r3
); );
break; break;
case GGML_TYPE_Q6_K: case GGML_TYPE_Q6_K:

View file

@ -24,8 +24,14 @@ layout (push_constant) uniform parameter {
int ne01; int ne01;
int ne02; int ne02;
int ne12; int ne12;
int r2; uint nb01;
int r3; uint nb02;
uint nb03;
uint nb11;
uint nb12;
uint nb13;
uint r2;
uint r3;
} pcs; } pcs;
void main() { void main() {
@ -50,10 +56,11 @@ void main() {
const uint i12 = im%pcs.ne12; const uint i12 = im%pcs.ne12;
const uint i13 = im/pcs.ne12; const uint i13 = im/pcs.ne12;
const uint offset0 = (i12/pcs.r2)*(nb*pcs.ne01) + (i13/pcs.r3)*(nb*pcs.ne01*pcs.ne02); const uint offset0 = first_row*(pcs.nb01/SIZE_OF_BLOCK) + (i12/pcs.r2)*(pcs.nb02/SIZE_OF_BLOCK) + (i13/pcs.r3)*(pcs.nb03/SIZE_OF_BLOCK);
const uint offset1 = r1*pcs.nb11 + (i12 )*pcs.nb12 + (i13 )*pcs.nb13;
const uint xblk = ib_row + offset0 + pcs.inAOff; const uint xblk = offset0 + pcs.inAOff;
const uint y = r1*pcs.ne10 + im*pcs.ne00*pcs.ne1 + pcs.inBOff; const uint y = (offset1 / 4) + pcs.inBOff;
float yl[16]; float yl[16];
float yh[16]; float yh[16];
@ -74,7 +81,7 @@ void main() {
} }
for (int row = 0; row < N_DST; row++) { for (int row = 0; row < N_DST; row++) {
uint row_idx = row * nb; uint row_idx = row * (pcs.nb01 / SIZE_OF_BLOCK);
uint16_t sc_0 = u8BufToU16(inA[blk_idx + row_idx].scales, iq * 2 + 0); uint16_t sc_0 = u8BufToU16(inA[blk_idx + row_idx].scales, iq * 2 + 0);
uint16_t sc_1 = u8BufToU16(inA[blk_idx + row_idx].scales, iq * 2 + 2); uint16_t sc_1 = u8BufToU16(inA[blk_idx + row_idx].scales, iq * 2 + 2);