From bbd488c42abbcf725e7c713bc184eb0ca4a73986 Mon Sep 17 00:00:00 2001 From: Daniel Varga Date: Fri, 18 Sep 2026 10:46:02 +0200 Subject: [PATCH] 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 --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 14 ++++-- .../vulkan-shaders/mul_mmq_funcs.glsl | 49 +++++++++++++++++++ .../vulkan-shaders/mul_mmq_shmem_types.glsl | 6 +++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 2 +- 4 files changed, 67 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8777c340a..da0e24fcd 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -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 diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl index 24da4f715..5fc4d3db4 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl @@ -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; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_shmem_types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_shmem_types.glsl index 2b7adcb6c..7632e457a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_shmem_types.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_shmem_types.glsl @@ -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 { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index d3f425968..73cef00b0 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -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