split out q4_k/q5_k cm2 shader to fix Ampere regression

This commit is contained in:
Ruben Ortlam
2026-09-03 13:13:38 +02:00
parent 6537ccf99d
commit 842771698a
3 changed files with 30 additions and 43 deletions
+15
View File
@@ -4781,11 +4781,19 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
}
#endif
for (const auto type : non_lut_quant_types) {
// regression in unified shader on Ampere
if (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K) {
continue;
}
auto& tc = ((type >= GGML_TYPE_Q2_K && type <= GGML_TYPE_Q6_K) || type == GGML_TYPE_TQ2_0) ? tc_mmq_k : tc_mmq;
spec_fn_t qs = [&, type](const std::vector<uint32_t>& wt, bool a) { return ggml_vk_mul_mm_cm2_spec(wt, a, (uint32_t)type); };
create_mm_pipelines({type, GGML_TYPE_F16, false, true}, tc, "matmul_quant_f16_f16acc", matmul_quant_f16_f16acc_cm2_len, matmul_quant_f16_f16acc_cm2_data, sizeof(vk_mat_mat_push_constants), 3, qs, true);
create_mm_pipelines({type, GGML_TYPE_F16, false, false}, tc, "matmul_quant_f16", matmul_quant_f16_cm2_len, matmul_quant_f16_cm2_data, sizeof(vk_mat_mat_push_constants), 3, qs, true);
}
create_mm_pipelines({GGML_TYPE_Q4_K, GGML_TYPE_F16, false, true}, tc_mmq_k, "matmul_q4_k_f16_f16acc", matmul_q4_k_f16_f16acc_cm2_len, matmul_q4_k_f16_f16acc_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true);
create_mm_pipelines({GGML_TYPE_Q4_K, GGML_TYPE_F16, false, false}, tc_mmq_k, "matmul_q4_k_f16", matmul_q4_k_f16_cm2_len, matmul_q4_k_f16_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true);
create_mm_pipelines({GGML_TYPE_Q5_K, GGML_TYPE_F16, false, true}, tc_mmq_k, "matmul_q5_k_f16_f16acc", matmul_q5_k_f16_f16acc_cm2_len, matmul_q5_k_f16_f16acc_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true);
create_mm_pipelines({GGML_TYPE_Q5_K, GGML_TYPE_F16, false, false}, tc_mmq_k, "matmul_q5_k_f16", matmul_q5_k_f16_cm2_len, matmul_q5_k_f16_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true);
#define X_CM2(TYPE, tstr) \
{ auto tc = filter_tc(tc_mmq, TYPE, false); \
if (!tc.empty()) { \
@@ -4817,10 +4825,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
}
#endif
for (const auto type : non_lut_quant_types) {
if (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K) {
continue;
}
spec_fn_t qs_id = [&, type](const std::vector<uint32_t>& wt, bool a) { return ggml_vk_mul_mm_cm2_spec(wt, a, (uint32_t)type); };
create_mm_pipelines({type, GGML_TYPE_F16, true, true}, tc_mmqid, "matmul_id_subgroup_quant_f16_f16acc", matmul_id_subgroup_quant_f16_f16acc_cm2_len, matmul_id_subgroup_quant_f16_f16acc_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, qs_id, true);
create_mm_pipelines({type, GGML_TYPE_F16, true, false}, tc_mmqid, "matmul_id_subgroup_quant_f16", matmul_id_subgroup_quant_f16_cm2_len, matmul_id_subgroup_quant_f16_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, qs_id, true);
}
create_mm_pipelines({GGML_TYPE_Q4_K, GGML_TYPE_F16, true, true}, tc_mmqid, "matmul_id_subgroup_q4_k_f16_f16acc", matmul_id_subgroup_q4_k_f16_f16acc_cm2_len, matmul_id_subgroup_q4_k_f16_f16acc_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true);
create_mm_pipelines({GGML_TYPE_Q4_K, GGML_TYPE_F16, true, false}, tc_mmqid, "matmul_id_subgroup_q4_k_f16", matmul_id_subgroup_q4_k_f16_cm2_len, matmul_id_subgroup_q4_k_f16_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true);
create_mm_pipelines({GGML_TYPE_Q5_K, GGML_TYPE_F16, true, true}, tc_mmqid, "matmul_id_subgroup_q5_k_f16_f16acc", matmul_id_subgroup_q5_k_f16_f16acc_cm2_len, matmul_id_subgroup_q5_k_f16_f16acc_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true);
create_mm_pipelines({GGML_TYPE_Q5_K, GGML_TYPE_F16, true, false}, tc_mmqid, "matmul_id_subgroup_q5_k_f16", matmul_id_subgroup_q5_k_f16_cm2_len, matmul_id_subgroup_q5_k_f16_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true);
#define X_CM2_ID(TYPE, tstr) \
{ auto tc = filter_tc(tc_mmqid, TYPE, true); \
if (!tc.empty()) { \
@@ -438,21 +438,14 @@ layout(buffer_reference, std430, buffer_reference_align = 16) buffer decodeBufQ4
// function that stores into shared memory.
// Q4_K and Q5_K have the same encoding of scales, so everything is shared except
// the part that fetches from the structure (which has a different block layout).
#if defined(DATA_A_Q4_K) || defined(DATA_A_Q5_K) || defined(MULMAT_QUANT)
#if defined(DATA_A_Q4_K) || defined(DATA_A_Q5_K)
const uint shAscales_stride = (BM + 2);
// 1 scale per 32 elements -> 8 scales per block, per row
#if defined(MULMAT_QUANT)
// In the unified shader the quant type is a runtime spec constant, so only reserve
// the full scales tile when Q4_K/Q5_K is actually selected. Otherwise size it down
// to 8 (1 triggers an Nvidia compiler bug) so it costs no meaningful shared memory.
shared vec2 shAscales[(MmTypeA == GGML_TYPE_Q4_K || MmTypeA == GGML_TYPE_Q5_K) ? 8 * shAscales_stride : 8];
#else
shared vec2 shAscales[8 * shAscales_stride];
#endif
uvec4 row_v;
#endif
#if defined(DATA_A_Q4_K) || defined(MULMAT_QUANT)
#if defined(DATA_A_Q4_K)
layout (binding = 0) readonly buffer A_Q4_K_128 {block_q4_K_packed128 data_a_q4_k_packed128[];};
void fetch_scalesQ4_K(uint ir_BM, uint pos_a, uint stride_a, uint block_k, uint tid, bool in_bounds)
@@ -463,13 +456,13 @@ void fetch_scalesQ4_K(uint ir_BM, uint pos_a, uint stride_a, uint block_k, uint
uint tid_row = tid / tids_per_row;
uint row = ir_BM + tid_row;
uint block_index = pos_a + row * stride_a + (block_k / QUANT_K_Q4_K);
uint block_index = pos_a + row * stride_a + (block_k / QUANT_K);
if (in_bounds || row < p.M) {
row_v = data_a_q4_k_packed128[block_index].q4k[0];
}
}
#endif
#if defined(DATA_A_Q5_K) || defined(MULMAT_QUANT)
#if defined(DATA_A_Q5_K)
layout (binding = 0) readonly buffer A_Q5_K_128 {block_q5_K_packed128 data_a_q5_k_packed128[];};
void fetch_scalesQ5_K(uint ir_BM, uint pos_a, uint stride_a, uint block_k, uint tid, bool in_bounds)
@@ -480,14 +473,14 @@ void fetch_scalesQ5_K(uint ir_BM, uint pos_a, uint stride_a, uint block_k, uint
uint tid_row = tid / tids_per_row;
uint row = ir_BM + tid_row;
uint block_index = pos_a + row * stride_a + (block_k / QUANT_K_Q5_K);
uint block_index = pos_a + row * stride_a + (block_k / QUANT_K);
if (in_bounds || row < p.M) {
row_v = data_a_q5_k_packed128[block_index].q5k[0];
}
}
#endif
#if defined(DATA_A_Q4_K) || defined(DATA_A_Q5_K) || defined(MULMAT_QUANT)
#if defined(DATA_A_Q4_K) || defined(DATA_A_Q5_K)
void store_scalesQ4_K(uint tid)
{
barrier();
@@ -530,26 +523,6 @@ void store_scalesQ4_K(uint tid)
}
#endif
#if defined(MULMAT_QUANT)
// Runtime dispatch for the unified shader: pick the right fetch based on the spec
// constant type. Only Q4_K/Q5_K use the shared-scales path; other types are no-ops.
void fetch_scales_cm2(uint ir_BM, uint pos_a, uint stride_a, uint block_k, uint tid, bool in_bounds)
{
if (MmTypeA == GGML_TYPE_Q4_K) {
fetch_scalesQ4_K(ir_BM, pos_a, stride_a, block_k, tid, in_bounds);
} else if (MmTypeA == GGML_TYPE_Q5_K) {
fetch_scalesQ5_K(ir_BM, pos_a, stride_a, block_k, tid, in_bounds);
}
}
void store_scales_cm2(uint tid)
{
if (MmTypeA == GGML_TYPE_Q4_K || MmTypeA == GGML_TYPE_Q5_K) {
store_scalesQ4_K(tid);
}
}
#endif
#endif
float16_t dequantFuncQ4_K(const in decodeBufQ4_K bl, const in uint blockCoords[2], const in uint coordInBlock[2])
@@ -561,7 +534,7 @@ float16_t dequantFuncQ4_K(const in decodeBufQ4_K bl, const in uint blockCoords[2
const uint b = (idx & 0x20) >> 5; // 0,1
const uint is = (idx & 0xE0) >> 5; // 0..7
#if defined(IS_MUL_MM2) && (defined(DATA_A_Q4_K) || defined(MULMAT_QUANT))
#if defined(IS_MUL_MM2) && defined(DATA_A_Q4_K)
vec2 v = shAscales[is * shAscales_stride + (blockCoords[0] % BM)];
float d = v.x;
float m = v.y;
@@ -608,7 +581,7 @@ f16vec4 dequantFuncQ4_K_v(const in decodeBufQ4_K bl, const in uint blockCoords[2
const uint is = idx >> 5; // 0..7
#if defined(IS_MUL_MM2) && (defined(DATA_A_Q4_K) || defined(MULMAT_QUANT))
#if defined(IS_MUL_MM2) && defined(DATA_A_Q4_K)
vec2 v = shAscales[is * shAscales_stride + (blockCoords[0] % BM)];
float d = v.x;
float m = v.y;
@@ -674,7 +647,7 @@ float16_t dequantFuncQ5_K(const in decodeBufQ5_K bl, const in uint blockCoords[2
const uint b = (idx & 0x20) >> 5; // 0,1
const uint is = (idx & 0xE0) >> 5; // 0..7
#if defined(IS_MUL_MM2) && (defined(DATA_A_Q5_K) || defined(MULMAT_QUANT))
#if defined(IS_MUL_MM2) && defined(DATA_A_Q5_K)
vec2 v = shAscales[is * shAscales_stride + (blockCoords[0] % BM)];
float d = v.x;
float m = v.y;
@@ -725,7 +698,7 @@ f16vec4 dequantFuncQ5_K_v(const in decodeBufQ5_K bl, const in uint blockCoords[2
const uint idx = coordInBlock[1];
const uint is = idx >> 5;
#if defined(IS_MUL_MM2) && (defined(DATA_A_Q5_K) || defined(MULMAT_QUANT))
#if defined(IS_MUL_MM2) && defined(DATA_A_Q5_K)
vec2 v = shAscales[is * shAscales_stride + (blockCoords[0] % BM)];
float d = v.x;
float m = v.y;
@@ -1491,9 +1464,3 @@ f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords
#elif defined(DATA_A_F32)
#define dequantFuncA dequantFuncF32
#endif
// Unified shader: dispatch the Q4_K/Q5_K shared-scales path at runtime on MmTypeA.
#if defined(MULMAT_QUANT)
#define fetch_scales fetch_scales_cm2
#define store_scales store_scales_cm2
#endif
@@ -647,6 +647,11 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c
#endif
continue;
}
// dedicated shader needed due to regression on Ampere
if (coopmat2 && (tname == "q4_k" || tname == "q5_k")) {
string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc);
}
}
// Quant shader: one SPIR-V for all quant types, selected via MmTypeA spec constant