diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f4fa37d2cc92..f3995be01e97 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -387,6 +387,9 @@ static void ggml_vk_destroy_buffer(vk_buffer& buf); static void ggml_vk_synchronize(ggml_backend_vk_context * ctx); static constexpr uint32_t mul_mat_vec_max_cols = 8; +// Selected types additionally get mat-vec pipelines for 9..mul_mat_vec_ext_cols +// columns (held in the *_x pipeline arrays). Currently only enabled on BC-250. +static constexpr uint32_t mul_mat_vec_ext_cols = 16; static constexpr uint32_t p021_max_gqa_ratio = 8; enum vk_device_architecture { @@ -926,6 +929,9 @@ struct vk_device_struct { vk_pipeline pipeline_dequant_transpose[GGML_TYPE_COUNT]; // fused dequant+transpose for FA quant-KV vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; + // column-extended variants (cols mul_mat_vec_max_cols+1 .. mul_mat_vec_ext_cols), only populated for opted-in types + vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32_x[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; + vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32_x[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT]; vk_pipeline pipeline_dequant_mul_mat_vec_q8_1_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; @@ -1154,6 +1160,18 @@ struct vk_device_struct { } }; +// BC-250 = AMD Cyan Skillfish APU: RDNA1 (gfx1013), no coopmat, no int dot product. +static bool ggml_vk_is_bc250(const vk_device& device) { + return device->vendor_id == VK_VENDOR_ID_AMD && device->architecture == AMD_RDNA1 && + device->properties.deviceID == 0x13fe && !device->coopmat_support; +} + +// Types that get column-extended mat-vec pipelines (>mul_mat_vec_max_cols). BC-250 only. +static bool ggml_vk_mmv_ext(const vk_device& device, ggml_type a_type) { + return ggml_vk_is_bc250(device) && + (a_type == GGML_TYPE_PQ2_0 || a_type == GGML_TYPE_BF16 || a_type == GGML_TYPE_Q8_0); +} + void vk_command_pool::init(vk_device& device, vk_queue *q_) { cmd_buffers.clear(); q = q_; @@ -5257,10 +5275,28 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { #define OCP_DMMV_DATA(NAME, REDUC) NAME ## _data[REDUC] #endif + // BC-250 (gfx1013): RDNA1, no coopmat/int-dot - the dedicated PQ2_0 mat-vec runs on a + // 32-wide workgroup with a per-column-count rows table. + static constexpr uint32_t bc250_pq2_mmv_rows_def[mul_mat_vec_ext_cols] = { + 8, 8, 8, 8, 16, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, 8, + }; + const bool bc250_pq2_mmv = ggml_vk_is_bc250(device); + const uint32_t mmv_pq2_wg = 32; + uint32_t mmv_pq2_rows[mul_mat_vec_ext_cols]; + for (uint32_t c = 0; c < mul_mat_vec_ext_cols; ++c) { + mmv_pq2_rows[c] = bc250_pq2_mmv ? bc250_pq2_mmv_rows_def[c] : 2 * rm_stdq; + } + const uint32_t mmv_pq2_sg = bc250_pq2_mmv ? 32 : force_subgroup_size; + for (uint32_t w = 0; w < DMMV_WG_SIZE_COUNT; ++w) { const uint32_t wg_size_subgroup = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size : (subgroup_size * 4); const uint32_t wg_size_subgroup16 = (w == DMMV_WG_SIZE_SUBGROUP) ? subgroup_size16 : (subgroup_size16 * 4); + // BC-250: fixed wg size for PQ2_0 (both w slots get the same tuned pipeline) + const uint32_t pq2_wg = bc250_pq2_mmv ? mmv_pq2_wg : wg_size_subgroup; + const shader_reduction_mode reduc_pq2 = !use_subgroups ? SHADER_REDUCTION_MODE_SHMEM : + (pq2_wg <= mmv_pq2_sg ? SHADER_REDUCTION_MODE_SUBGROUP : SHADER_REDUCTION_MODE_HYBRID); + const shader_reduction_mode reduc = (use_subgroups && w == DMMV_WG_SIZE_SUBGROUP) ? SHADER_REDUCTION_MODE_SUBGROUP : (use_subgroups && w == DMMV_WG_SIZE_LARGE) ? SHADER_REDUCTION_MODE_HYBRID : SHADER_REDUCTION_MODE_SHMEM; @@ -5275,7 +5311,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f32_f32", arr_dmmv_bf16_f32_f32_len[reduc], arr_dmmv_bf16_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q1_0][i], "mul_mat_vec_q1_0_f32_f32", arr_dmmv_q1_0_f32_f32_len[reduc], arr_dmmv_q1_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_PTQ1_0][i], "mul_mat_vec_ptq1_0_f32_f32", arr_dmmv_ptq1_0_f32_f32_len[reduc], arr_dmmv_ptq1_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); - ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_PQ2_0][i], "mul_mat_vec_pq2_0_f32_f32", arr_dmmv_pq2_0_f32_f32_len[reduc], arr_dmmv_pq2_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_PQ2_0][i], "mul_mat_vec_pq2_0_f32_f32", arr_dmmv_pq2_0_f32_f32_len[reduc_pq2], arr_dmmv_pq2_0_f32_f32_data[reduc_pq2], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {mmv_pq2_rows[i], 1, 1}, {pq2_wg, mmv_pq2_rows[i], i+1}, 1, true, use_subgroups, mmv_pq2_sg); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_f32_f32", arr_dmmv_q2_0_f32_f32_len[reduc], arr_dmmv_q2_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f32_f32", arr_dmmv_q4_0_f32_f32_len[reduc], arr_dmmv_q4_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f32_f32", arr_dmmv_q4_1_f32_f32_len[reduc], arr_dmmv_q4_1_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); @@ -5305,7 +5341,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f16_f32", arr_dmmv_bf16_f16_f32_len[reduc], arr_dmmv_bf16_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, i+1}, 1, false, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q1_0][i], "mul_mat_vec_q1_0_f16_f32", arr_dmmv_q1_0_f16_f32_len[reduc], arr_dmmv_q1_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_PTQ1_0][i], "mul_mat_vec_ptq1_0_f16_f32", arr_dmmv_ptq1_0_f16_f32_len[reduc], arr_dmmv_ptq1_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); - ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_PQ2_0][i], "mul_mat_vec_pq2_0_f16_f32", arr_dmmv_pq2_0_f16_f32_len[reduc], arr_dmmv_pq2_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_PQ2_0][i], "mul_mat_vec_pq2_0_f16_f32", arr_dmmv_pq2_0_f16_f32_len[reduc_pq2], arr_dmmv_pq2_0_f16_f32_data[reduc_pq2], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {mmv_pq2_rows[i], 1, 1}, {pq2_wg, mmv_pq2_rows[i], i+1}, 1, true, use_subgroups, mmv_pq2_sg); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q2_0][i], "mul_mat_vec_q2_0_f16_f32", arr_dmmv_q2_0_f16_f32_len[reduc], arr_dmmv_q2_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_0][i], "mul_mat_vec_q4_0_f16_f32", arr_dmmv_q4_0_f16_f32_len[reduc], arr_dmmv_q4_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_1][i], "mul_mat_vec_q4_1_f16_f32", arr_dmmv_q4_1_f16_f32_len[reduc], arr_dmmv_q4_1_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); @@ -5360,6 +5396,20 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { #endif // GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT } + // BC-250: extended mat-vec columns (mul_mat_vec_max_cols+1 .. mul_mat_vec_ext_cols) - + // dedicated PQ2_0 shader; generic kernels for BF16/Q8_0 (activation-path types). + if (bc250_pq2_mmv) { + for (uint32_t i = 0; i < mul_mat_vec_max_cols; ++i) { + const uint32_t ncol = i + 1 + mul_mat_vec_max_cols; + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32_x[w][GGML_TYPE_PQ2_0][i], "mul_mat_vec_pq2_0_f32_f32", arr_dmmv_pq2_0_f32_f32_len[reduc_pq2], arr_dmmv_pq2_0_f32_f32_data[reduc_pq2], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {mmv_pq2_rows[ncol - 1], 1, 1}, {pq2_wg, mmv_pq2_rows[ncol - 1], ncol}, 1, true, use_subgroups, mmv_pq2_sg); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32_x[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f32_f32", arr_dmmv_bf16_f32_f32_len[reduc], arr_dmmv_bf16_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, ncol}, 1, false, use_subgroups, force_subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32_x[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f32_f32", arr_dmmv_q8_0_f32_f32_len[reduc], arr_dmmv_q8_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, ncol}, 1, true, use_subgroups, force_subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32_x[w][GGML_TYPE_PQ2_0][i], "mul_mat_vec_pq2_0_f16_f32", arr_dmmv_pq2_0_f16_f32_len[reduc_pq2], arr_dmmv_pq2_0_f16_f32_data[reduc_pq2], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {mmv_pq2_rows[ncol - 1], 1, 1}, {pq2_wg, mmv_pq2_rows[ncol - 1], ncol}, 1, true, use_subgroups, mmv_pq2_sg); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32_x[w][GGML_TYPE_BF16][i], "mul_mat_vec_bf16_f16_f32", arr_dmmv_bf16_f16_f32_len[reduc], arr_dmmv_bf16_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2, 1, 1}, {wg_size_subgroup, 2, ncol}, 1, false, use_subgroups, force_subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32_x[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f16_f32", arr_dmmv_q8_0_f16_f32_len[reduc], arr_dmmv_q8_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, ncol}, 1, true, use_subgroups, force_subgroup_size); + } + } + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_F32 ], "mul_mat_vec_id_f32_f32", arr_dmmv_id_f32_f32_f32_len[reduc], arr_dmmv_id_f32_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1, 1, 1}, {wg_size_subgroup, 1}, 1, false, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_F16 ], "mul_mat_vec_id_f16_f32", arr_dmmv_id_f16_f32_f32_len[reduc], arr_dmmv_id_f16_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {wg_size_subgroup, 2}, 1, false, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_BF16], "mul_mat_vec_id_bf16_f32", arr_dmmv_id_bf16_f32_f32_len[reduc], arr_dmmv_id_bf16_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2, 1, 1}, {wg_size_subgroup, 2}, 1, false, use_subgroups, force_subgroup_size); @@ -7842,7 +7892,17 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type, uint32_t num_cols, uint32_t m, uint32_t k) { VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); GGML_ASSERT(b_type == GGML_TYPE_F32 || b_type == GGML_TYPE_F16 || b_type == GGML_TYPE_Q8_1); - GGML_ASSERT(num_cols >= 1 && num_cols <= mul_mat_vec_max_cols); + GGML_ASSERT(num_cols >= 1 && num_cols <= mul_mat_vec_ext_cols); + + if (num_cols > mul_mat_vec_max_cols) { + // Extended columns exist only for opted-in types (currently BC-250); the + // q8_1-activation path has no extended variants. + if (b_type == GGML_TYPE_Q8_1 || !ggml_vk_mmv_ext(ctx->device, a_type)) { + return nullptr; + } + return b_type == GGML_TYPE_F32 ? ctx->device->pipeline_dequant_mul_mat_vec_f32_f32_x[DMMV_WG_SIZE_SUBGROUP][a_type][num_cols-1-mul_mat_vec_max_cols] + : ctx->device->pipeline_dequant_mul_mat_vec_f16_f32_x[DMMV_WG_SIZE_SUBGROUP][a_type][num_cols-1-mul_mat_vec_max_cols]; + } if (b_type == GGML_TYPE_Q8_1) { switch (a_type) { @@ -10165,7 +10225,8 @@ static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, c ggml_vk_mul_mat_vec_nc_f16_f32(ctx, subctx, cgraph, node_idx); // mul_mat_vec supports batching ne12*ne13 when ne11==1, or treating ne11 as the batch size (up to four) // when ne12 and ne13 are one. - } else if ((dst->ne[1] == 1 || (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1)) && + } else if ((dst->ne[1] == 1 || (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1) || + (dst->ne[1] <= mul_mat_vec_ext_cols && src1->ne[2] * src1->ne[3] == 1 && ggml_vk_mmv_ext(ctx->device, src0->type))) && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16 || ggml_is_quantized(src0->type))) { ggml_vk_mul_mat_vec_q_f16(ctx, subctx, cgraph, node_idx); } else { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_pq2_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_pq2_0.comp new file mode 100644 index 000000000000..3cf337ffdce0 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vec_pq2_0.comp @@ -0,0 +1,106 @@ +#version 450 + +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require + +#include "mul_mat_vec_base.glsl" + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +FLOAT_TYPE temp[NUM_COLS][NUM_ROWS]; + +// PQ2_0 dedicated mat-vec. +// block_pq2_0 = { float16 d; uint8 qs[32]; } // 34 B per 128 weights (QUANT_K=128) +// qs byte b packs elements 4*b..4*b+3 at bit levels 0,2,4,6; w = (q - 1) * d. +// Thread map: 16 lanes per block (tpb=16), lane itid reads u16 qs[itid] covering +// elements 8*itid..8*itid+7 - always 2-byte aligned despite the 34-byte stride. +// For workgroups smaller than 16, each thread replays for the extra lanes +// (lane = itid + w*tpb), so the shader stays correct down to wg=4. +// +// Inner loop uses the identity w = (q - 1) * d => d * (sum(q*y) - sum(y)): +// per (block, row) the 2-bit fields are decoded once (outside the column loop) +// and per column the B vector is fetched once for all rows. + +// Accumulate one 8-element span of one block into temp[*][0..num_rows). +// lane selects the u16 word (elements 8*lane..8*lane+7). +void calc_span(const uint a_offset, const uint b_offset, const uint lane, + const uint i, const uint num_blocks_per_row, + const uint first_row, const uint num_rows) { + vec4 q0[NUM_ROWS], q1[NUM_ROWS]; + FLOAT_TYPE dv[NUM_ROWS]; + + [[unroll]] for (uint n = 0; n < NUM_ROWS; ++n) { + if (n >= num_rows) { break; } + const uint ib = a_offset + (first_row + n) * num_blocks_per_row + i; + const uint q16 = uint(data_a_packed16[ib].qs[lane]); + q0[n] = vec4(uvec4(q16, q16 >> 2u, q16 >> 4u, q16 >> 6u) & 3u); + q1[n] = vec4(uvec4(q16 >> 8u, q16 >> 10u, q16 >> 12u, q16 >> 14u) & 3u); + dv[n] = FLOAT_TYPE(data_a[ib].d); + } + + const uint y_idx = i * QUANT_K + lane * 8u; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + const uint bo = j * p.batch_stride_b + b_offset + y_idx; + const vec4 bv0 = vec4(data_b_v4[bo / 4u]); + const vec4 bv1 = vec4(data_b_v4[bo / 4u + 1u]); + const FLOAT_TYPE sumb = (bv0.x + bv0.y) + (bv0.z + bv0.w) + + (bv1.x + bv1.y) + (bv1.z + bv1.w); + [[unroll]] for (uint n = 0; n < NUM_ROWS; ++n) { + if (n >= num_rows) { break; } + const FLOAT_TYPE sumq = dot(q0[n], bv0) + dot(q1[n], bv1); + temp[j][n] = fma(dv[n], sumq - sumb, temp[j][n]); + } + } +} + +void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { + uint a_offset, b_offset, d_offset; + get_offsets(a_offset, b_offset, d_offset); + + const uint num_blocks_per_row = p.ncols / QUANT_K; + + // 16 threads cooperate on one block (one u16 word each). it_size = number + // of blocks processed per iteration; leftover threads (wg%16) stay idle. + const uint tpb = min(gl_WorkGroupSize.x, 16u); + const uint it_size = gl_WorkGroupSize.x / tpb; + const uint tid = gl_LocalInvocationID.x; + const uint itid = tid % tpb; + const uint ix = tid / tpb; + + [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { + [[unroll]] for (uint n = 0; n < NUM_ROWS; ++n) { + temp[j][n] = FLOAT_TYPE(0); + } + } + + for (uint i0 = 0; i0 < num_blocks_per_row; i0 += it_size) { + const uint i = i0 + ix; + if (ix < it_size && i < num_blocks_per_row) { + // replay for lanes >= tpb when the workgroup is narrower than 16 + [[unroll]] for (uint w = 0; w < 4; ++w) { + const uint lane = itid + w * tpb; + if (lane >= 16u) { break; } + calc_span(a_offset, b_offset, lane, i, num_blocks_per_row, + first_row, num_rows); + } + } + } + + reduce_result(temp, d_offset, first_row, num_rows, tid); +} + +void main() { + const uint first_row = NUM_ROWS * (gl_WorkGroupID.x + gl_NumWorkGroups.x * gl_WorkGroupID.z); + + // do NUM_ROWS at a time, unless there aren't enough remaining rows + if (first_row + NUM_ROWS <= p.stride_d) { + compute_outputs(first_row, NUM_ROWS); + } else { + if (first_row >= p.stride_d) { + return; + } + compute_outputs(first_row, min(NUM_ROWS, p.stride_d - first_row)); + } +} 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 f2239186d922..961fa52cdd53 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -740,7 +740,7 @@ void process_shaders() { for (const auto& tname : type_names) { // mul mat vec std::string data_a_key = "DATA_A_" + to_uppercase(tname); - std::string shader = (string_ends_with(tname, "_k") || string_starts_with(tname, "iq1_") || string_starts_with(tname, "iq2_") || string_starts_with(tname, "iq3_") || tname == "tq2_0") ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; + std::string shader = (string_ends_with(tname, "_k") || string_starts_with(tname, "iq1_") || string_starts_with(tname, "iq2_") || string_starts_with(tname, "iq3_") || tname == "tq2_0" || tname == "pq2_0") ? "mul_mat_vec_" + tname + ".comp" : "mul_mat_vec.comp"; string_to_spv("mul_mat_vec_" + tname + "_f32_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float"}, {"B_TYPEV2", "vec2"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}})); string_to_spv("mul_mat_vec_" + tname + "_f16_f32", shader, merge_maps(base_dict, {{data_a_key, "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV2", "f16vec2"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}})); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 72e580d895c0..55d630e15d31 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9310,6 +9310,19 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_PQ2_0, GGML_TYPE_F32, 4, 2, false, 70, n, 2048)); } + // PQ2_0 extended mat-vec columns 9..16 (used on BC-250): tails, batch, and a production ffn shape + for (int64_t n : {9, 10, 11, 12, 13, 14, 15, 16}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, 67, n, 1024, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, 67, n, 5120, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, 17408, n, 5120, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_PQ2_0, GGML_TYPE_F32, 64, n, 2048, {2, 3}, {1, 1})); + } + // BF16/Q8_0 also get extended mat-vec columns on BC-250 + for (int64_t n : {9, 12, 16}) { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_BF16, GGML_TYPE_F32, 128, n, 1024, {1, 1}, {1, 1})); + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 128, n, 1024, {1, 1}, {1, 1})); + } + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1})); test_cases.emplace_back(new test_mul_mat(GGML_TYPE_MXFP4, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1}));