mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-19 09:15:04 +02:00
vulkan: add IQ3_S MMQ matmul kernels (#28822)
* vulkan: add IQ3_S MMQ matmul kernels * Make block_a_to_shmem do 2-byte loads (110 bytes is divisible by 2) * Align the check, IQ3_S is also using K tile size
This commit is contained in:
@@ -1484,10 +1484,14 @@ static bool ggml_vk_matmul_int_shmem_support(const vk_device& device, const std:
|
||||
case GGML_TYPE_Q4_K: block_a_size = std430_size({{16, 4}, {fp2_size, fp2_align}}); break; // qs[4] + dm(vec2)
|
||||
case GGML_TYPE_Q5_K: block_a_size = std430_size({{32, 4}, {fp2_size, fp2_align}}); break; // qs[8] + dm(vec2)
|
||||
case GGML_TYPE_Q6_K: block_a_size = std430_size({{32, 4}, {fp2_size, fp2_align}}); break; // qs[8] + d_scales(vec2)
|
||||
case GGML_TYPE_IQ3_S: block_a_size = std430_size({{32, 4}, {fp_size, fp_align}}); break; // qs[8] + d
|
||||
default:
|
||||
return false;
|
||||
}
|
||||
|
||||
// IQ3_S also copies its 512-entry grid into shared memory (types.glsl, init_iq_shmem)
|
||||
const uint32_t lut_size = (src0_type == GGML_TYPE_IQ3_S) ? 4*512 : 0;
|
||||
|
||||
// block_b_cache: { int32_t qs[8]; FLOAT_TYPEV2 ds; }
|
||||
const uint32_t block_b_size = std430_size({{32, 4}, {fp2_size, fp2_align}});
|
||||
|
||||
@@ -1503,7 +1507,7 @@ static bool ggml_vk_matmul_int_shmem_support(const vk_device& device, const std:
|
||||
const uint32_t warps = warptile[0] / warptile[10];
|
||||
const uint32_t ballots_sh = mul_mat_id ? (warps * 4u * (uint32_t)sizeof(uint32_t)) : 0u;
|
||||
|
||||
const uint32_t total_size = buf_a_size + buf_b_size + mmid_row_ids + ballots_sh;
|
||||
const uint32_t total_size = buf_a_size + buf_b_size + mmid_row_ids + ballots_sh + lut_size;
|
||||
const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize;
|
||||
|
||||
VK_LOG_DEBUG("ggml_vk_matmul_int_shmem_support(warptile=(" << warptile[0] << "," << warptile[1] << "," << warptile[2] << "), "
|
||||
@@ -1780,10 +1784,10 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
}
|
||||
|
||||
// The q8_1 mmq path has its own (larger) shmem layout, check it separately.
|
||||
// K-quants use the _int_k warptiles, others use _int.
|
||||
// K-quants and IQ3_S use the _int_k warptiles, others use _int.
|
||||
const bool is_k_quant = (t == GGML_TYPE_Q2_K || t == GGML_TYPE_Q3_K ||
|
||||
t == GGML_TYPE_Q4_K || t == GGML_TYPE_Q5_K ||
|
||||
t == GGML_TYPE_Q6_K);
|
||||
t == GGML_TYPE_Q6_K || t == GGML_TYPE_IQ3_S);
|
||||
const auto & s_int = is_k_quant ? s_warptile_mmq_int_k : s_warptile_mmq_int;
|
||||
const auto & m_int = is_k_quant ? m_warptile_mmq_int_k : m_warptile_mmq_int;
|
||||
const auto & l_int = is_k_quant ? l_warptile_mmq_int_k : l_warptile_mmq_int;
|
||||
@@ -2447,6 +2451,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
sg_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q4_k_q8_1", matmul_q4_k_q8_1_len, matmul_q4_k_q8_1_data, sizeof(vk_mat_mat_push_constants), 3);
|
||||
sg_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q5_k_q8_1", matmul_q5_k_q8_1_len, matmul_q5_k_q8_1_data, sizeof(vk_mat_mat_push_constants), 3);
|
||||
sg_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q6_k_q8_1", matmul_q6_k_q8_1_len, matmul_q6_k_q8_1_data, sizeof(vk_mat_mat_push_constants), 3);
|
||||
sg_create_mmq({GGML_TYPE_IQ3_S, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_iq3_s_q8_1", matmul_iq3_s_q8_1_len, matmul_iq3_s_q8_1_data, sizeof(vk_mat_mat_push_constants), 3);
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -2483,6 +2488,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
sg_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_subgroup_q4_k_q8_1", matmul_id_subgroup_q4_k_q8_1_len, matmul_id_subgroup_q4_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16);
|
||||
sg_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_subgroup_q5_k_q8_1", matmul_id_subgroup_q5_k_q8_1_len, matmul_id_subgroup_q5_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16);
|
||||
sg_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_subgroup_q6_k_q8_1", matmul_id_subgroup_q6_k_q8_1_len, matmul_id_subgroup_q6_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16);
|
||||
sg_create_mmq({GGML_TYPE_IQ3_S, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_subgroup_iq3_s_q8_1", matmul_id_subgroup_iq3_s_q8_1_len, matmul_id_subgroup_iq3_s_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16);
|
||||
}
|
||||
#endif
|
||||
} else {
|
||||
@@ -2518,6 +2524,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
sg_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_q4_k_q8_1", matmul_id_q4_k_q8_1_len, matmul_id_q4_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count);
|
||||
sg_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_q5_k_q8_1", matmul_id_q5_k_q8_1_len, matmul_id_q5_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count);
|
||||
sg_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_q6_k_q8_1", matmul_id_q6_k_q8_1_len, matmul_id_q6_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count);
|
||||
sg_create_mmq({GGML_TYPE_IQ3_S, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_iq3_s_q8_1", matmul_id_iq3_s_q8_1_len, matmul_id_iq3_s_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
@@ -2554,6 +2561,7 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
|
||||
sg_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q4_k_q8_1", matmul_q4_k_q8_1_fp32_len, matmul_q4_k_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3);
|
||||
sg_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q5_k_q8_1", matmul_q5_k_q8_1_fp32_len, matmul_q5_k_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3);
|
||||
sg_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q6_k_q8_1", matmul_q6_k_q8_1_fp32_len, matmul_q6_k_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3);
|
||||
sg_create_mmq({GGML_TYPE_IQ3_S, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_iq3_s_q8_1", matmul_iq3_s_q8_1_fp32_len, matmul_iq3_s_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3);
|
||||
}
|
||||
#endif
|
||||
|
||||
|
||||
@@ -454,6 +454,55 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
|
||||
}
|
||||
#endif
|
||||
|
||||
#if defined(DATA_A_IQ3_S)
|
||||
// 2-byte loads for IQ3_S blocks (110 bytes)
|
||||
void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
|
||||
const uint ib_k = ib / 8;
|
||||
const uint ib32 = ib % 8;
|
||||
|
||||
// grid indices for qs[2 * iqs] and qs[2 * iqs + 1]
|
||||
const uint qs = uint(data_a_packed16[ib_k].qs[ib32 * 4 + iqs]);
|
||||
// their two high index bits
|
||||
const uint qh = uint(data_a_packed16[ib_k].qh[ib32 / 2]) >> ((ib32 & 1) * 8 + 2 * iqs);
|
||||
// one sign bit per value, 8 values
|
||||
const uint signs = uint(data_a_packed16[ib_k].signs[ib32 * 2 + iqs / 2]) >> ((iqs & 1) * 8);
|
||||
|
||||
// grid holds 4 values of 1..15, one per byte
|
||||
const ivec4 vals0 = ivec4(unpack8(iq3s_grid[( qs & 0xFF) | ((qh & 1) << 8)]));
|
||||
const ivec4 vals1 = ivec4(unpack8(iq3s_grid[((qs >> 8) & 0xFF) | ((qh & 2) << 7)]));
|
||||
|
||||
// negate with (v ^ -s) - -s to avoid branches
|
||||
const ivec4 m0 = -(ivec4(signs, signs >> 1, signs >> 2, signs >> 3) & 1);
|
||||
const ivec4 m1 = -(ivec4(signs >> 4, signs >> 5, signs >> 6, signs >> 7) & 1);
|
||||
|
||||
buf_a[buf_ib].qs[2 * iqs ] = pack32(i8vec4((vals0 ^ m0) - m0));
|
||||
buf_a[buf_ib].qs[2 * iqs + 1] = pack32(i8vec4((vals1 ^ m1) - m1));
|
||||
|
||||
if (iqs == 0) {
|
||||
const uint scale = (uint(data_a_packed16[ib_k].scales[ib32 / 4]) >> ((ib32 & 3) * 4)) & 0xF;
|
||||
|
||||
buf_a[buf_ib].d = FLOAT_TYPE(float(data_a_packed16[ib_k].d) * float(1 + 2 * scale));
|
||||
}
|
||||
}
|
||||
|
||||
void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
|
||||
cache_a[reg_ib].d = buf_a[buf_ib].d;
|
||||
|
||||
[[unroll]] for (uint iqs = 0; iqs < 8; iqs++) {
|
||||
cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs];
|
||||
}
|
||||
}
|
||||
|
||||
ACC_TYPE mmq_dot_product(const uint ib_a) {
|
||||
int32_t q_sum = 0;
|
||||
[[unroll]] for (uint iqs = 0; iqs < 8; iqs++) {
|
||||
q_sum += dotPacked4x8EXT(cache_a[ib_a].qs[iqs], cache_b.qs[iqs]);
|
||||
}
|
||||
|
||||
return ACC_TYPE(float(cache_a[ib_a].d) * float(cache_b.ds.x) * float(q_sum));
|
||||
}
|
||||
#endif
|
||||
|
||||
void block_b_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
|
||||
if (is_in_bounds) {
|
||||
const uint ib_outer = ib / 4;
|
||||
|
||||
@@ -59,6 +59,12 @@ struct block_a_cache {
|
||||
int32_t qs[8];
|
||||
FLOAT_TYPE d;
|
||||
};
|
||||
#elif defined(DATA_A_IQ3_S)
|
||||
#define QUANT_R_MMQ 2
|
||||
struct block_a_cache {
|
||||
int32_t qs[8];
|
||||
FLOAT_TYPE d;
|
||||
};
|
||||
#elif defined(DATA_A_Q2_K)
|
||||
#define QUANT_R_MMQ 4
|
||||
struct block_a_cache {
|
||||
|
||||
@@ -624,7 +624,7 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c
|
||||
};
|
||||
|
||||
#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT)
|
||||
if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4")) {
|
||||
if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4" || tname == "iq3_s")) {
|
||||
string_to_spv(shader_name + "_" + tname + "_q8_1", "mul_mmq.comp", merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"D_TYPE", "float"},}), fp16, coopmat, coopmat2, f16acc);
|
||||
}
|
||||
#endif
|
||||
|
||||
Reference in New Issue
Block a user