Commit d20426bbb for llama.cpp
commit d20426bbbe32dde9c352e954d305105b0e3dd1f8
Author: Jeff Bolz <jbolz@nvidia.com>
Date: Sat Nov 14 16:36:05 2026 -0600
vulkan: optimize matrix multiplication for narrow N (#31990)
* vulkan: optimize matrix multiplication for narrow N
Add new "narrow" tile configs for coopmat1/coopmat2 that are tuned for N<=32.
* Attempt to handle wave64
* add n=16 tile for cm1
* apply RDNA split_k heuristic to narrow_n generally
---------
Co-authored-by: Ruben Ortlam <rortlam@redhat.com>
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
index e72afc6ae..ac9d05e6e 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h
@@ -270,12 +270,14 @@ struct vk_matmul_pipeline_pair {
vk_pipeline unaligned;
vk_pipeline aligned;
uint32_t align;
+ uint32_t narrow_n_max = 0;
};
struct vk_tile_config {
std::vector<uint32_t> warptile;
std::array<uint32_t, 3> wg_denoms;
uint32_t align;
+ uint32_t narrow_n_max = 0;
};
using matmul_tile_selector_t = std::function<uint32_t(
@@ -1049,6 +1051,7 @@ struct vk_device_struct {
bool allow_sysmem_fallback;
bool disable_graph_optimize;
bool xe_fa_support {};
+ bool disable_narrow_mul_mat;
std::unique_ptr<vk_memory_logger> memory_logger;
@@ -1237,6 +1240,7 @@ inline std::mutex vk_memory_logger::log_mutex;
class vk_perf_logger {
public:
+ explicit vk_perf_logger(const vk_device & device) : device(device) {}
void print_timings(bool force = false);
@@ -1271,6 +1275,7 @@ class vk_perf_logger {
}
private:
+ vk_device device;
std::map<std::string, std::vector<uint64_t>> timings;
std::map<std::string, std::vector<uint64_t>> flops;
uint32_t print_count {};
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index bb5aec365..b692f2f8b 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -1614,6 +1614,11 @@ static vk_xe_fa_tuning_params get_xe_fa_tuning_params(const vk_device_struct * d
return params;
}
+// CM2 uses narrow tiles for 5 <= N <= 32; CM1 uses them for 9 <= N <= 32.
+static uint32_t ggml_vk_mul_mat_vec_max_cols(const vk_device & device) {
+ return device->coopmat2 && !device->disable_narrow_mul_mat ? 4u : mul_mat_vec_max_cols;
+}
+
static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector<uint32_t>& warptile, bool mul_mat_id, ggml_type src0_type) {
uint32_t lut_size = 0;
@@ -1662,7 +1667,7 @@ static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vec
const uint32_t load_bufs = (warptile[1] + warptile[2]) * (warptile[3] + bank_conflict_offset) * type_size;
const uint32_t mmid_row_ids = mul_mat_id ? (warptile[2] * 2 * sizeof(uint16_t)) : 0;
- const uint32_t coopmat_stage = device->coopmat_support ? warptile[7] * warptile[8] / warps * sizeof(float) : 0;
+ const uint32_t coopmat_stage = device->coopmat_support ? warptile[7] * warptile[8] * warps * sizeof(float) : 0;
const uint32_t total_size = load_bufs + mmid_row_ids + coopmat_stage + lut_size;
const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize;
@@ -2418,14 +2423,19 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
const bool first_call = vec.empty();
for (size_t i = 0; i < tile_configs.size(); i++) {
const auto& tc = tile_configs[i];
+ if (tc.narrow_n_max != 0 && (key.mul_mat_id || device->disable_narrow_mul_mat)) {
+ continue;
+ }
// Intel coopmat1 pins the required subgroup size to each warptile's WARP element.
- const uint32_t rsgs = pin_subgroup_to_warp ? tc.warptile[WARP_SIZE_IDX] : required_subgroup_size;
- const bool rfs = require_full_subgroups || pin_subgroup_to_warp;
+ const bool pin_sg = pin_subgroup_to_warp || (!device->coopmat2 && tc.narrow_n_max != 0);
+ const uint32_t rsgs = pin_sg ? tc.warptile[WARP_SIZE_IDX] : required_subgroup_size;
+ const bool rfs = require_full_subgroups || pin_sg;
if (first_call) {
vk_matmul_pipeline_pair pair{};
pair.align = tc.align;
+ pair.narrow_n_max = tc.narrow_n_max;
std::string suffix = "_" + std::to_string(i);
pair.unaligned = std::make_shared<vk_pipeline_struct>();
if (create_aligned) {
@@ -2462,12 +2472,60 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
enabled[1] = is_id ? device->mul_mat_id_m[type] : device->mul_mat_m[type];
enabled[2] = is_id ? device->mul_mat_id_l[type] : device->mul_mat_l[type];
}
- for (size_t i = 0; i < configs.size() && i < 3; i++) {
- if (enabled[i]) result.push_back(configs[i]);
+ for (size_t i = 0; i < configs.size(); i++) {
+ if (configs[i].narrow_n_max != 0 || (i < 3 && enabled[i])) result.push_back(configs[i]);
}
return result;
};
+ auto add_narrow_n_cm1_tile = [&](std::vector<vk_tile_config>& configs, ggml_type type, uint32_t bn) {
+ constexpr uint32_t bm = 32, bk = 32;
+ const uint32_t tm = device->coopmat_m, tn = device->coopmat_n, tk = device->coopmat_k;
+
+ const bool rdna = device->architecture == AMD_RDNA3 || device->architecture == AMD_RDNA4;
+ const uint32_t base_subgroup_size = rdna && device->subgroup_size_control ? device->subgroup_min_size : mul_mat_subgroup_size;
+ const uint32_t subgroup_size = std::max(base_subgroup_size, tm);
+
+ const uint32_t wm = subgroup_size > 32 ? 32 : 16;
+ const uint32_t wn = std::min(bn, subgroup_size > 64 ? 32u : 16u);
+ const uint32_t subgroup_count = (bm / wm) * (bn / wn);
+ const uint32_t workgroup_size = subgroup_count * subgroup_size;
+
+ // Two-element loads need at least BK/2 invocations; eight-element loads must fit both tiles.
+ if (workgroup_size < bk / 2 || workgroup_size > std::min(bm, bn) * bk / 8) {
+ return;
+ }
+
+ // Native matrices must tile each subgroup's output region and the K tile.
+ if (tm == 0 || tn == 0 || tk == 0 || wm % tm != 0 || wn % tn != 0 || bk % tk != 0) {
+ return;
+ }
+ const bool subgroup_size_supported = device->subgroup_size_control ?
+ subgroup_size >= device->subgroup_min_size && subgroup_size <= device->subgroup_max_size :
+ subgroup_size == device->subgroup_size;
+
+ // Tail stores use lane%TM for rows and lane/TM for columns of one native matrix.
+ if (!subgroup_size_supported || subgroup_size % tm != 0 || subgroup_size > tm * tn) {
+ return;
+ }
+ if (subgroup_size > device->properties.limits.maxComputeWorkGroupInvocations / subgroup_count ||
+ subgroup_size > device->properties.limits.maxComputeWorkGroupSize[0] / subgroup_count) {
+ return;
+ }
+
+ const std::vector<uint32_t> wt = {workgroup_size, bm, bn, bk, wm, wn, 2, tm, tn, tk, subgroup_size};
+ if (ggml_vk_matmul_shmem_support(device, wt, false, type)) {
+ configs.push_back({wt, {bm, bn, 1}, bk, bn});
+ }
+ };
+ auto add_narrow_n_cm1 = [&](std::vector<vk_tile_config>& configs, ggml_type type, bool is_id) {
+ if (is_id || configs.empty()) {
+ return;
+ }
+ add_narrow_n_cm1_tile(configs, type, 16);
+ add_narrow_n_cm1_tile(configs, type, 32);
+ };
+
std::vector<vk_tile_config> tc_mm = {{s_warptile, s_wg_denoms, s_align}, {m_warptile, m_wg_denoms, m_align}, {l_warptile, l_wg_denoms, l_align}};
std::vector<vk_tile_config> tc_mmq = {{s_warptile_mmq, s_mmq_wg_denoms, s_align}, {m_warptile_mmq, m_mmq_wg_denoms, m_align}, {l_warptile_mmq, l_mmq_wg_denoms, l_align}};
@@ -2486,6 +2544,13 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
std::vector<vk_tile_config> tc_mmq_k = {{s_warptile_mmq_k, s_mmq_wg_denoms_k, s_align}, {m_warptile_mmq_k, m_mmq_wg_denoms_k, m_align}, {l_warptile_mmq_k, l_mmq_wg_denoms_k, l_align}};
std::vector<vk_tile_config> tc_mmqid = {{s_warptile_mmqid, s_mmqid_wg_denoms, s_align}, {m_warptile_mmqid, m_mmqid_wg_denoms, m_align}, {l_warptile_mmqid, l_mmqid_wg_denoms, l_align}};
+ const vk_tile_config narrow_n16 = {{128, 32, 16, 64, 0}, {32, 16, 1}, 64, 16};
+ const vk_tile_config narrow_n32 = {{128, 32, 32, 128, 0}, {32, 32, 1}, 128, 32};
+ for (auto * tc : {&tc_mm, &tc_mmq, &tc_mmq_k}) {
+ tc->push_back(narrow_n16);
+ tc->push_back(narrow_n32);
+ }
+
spec_fn_t cm2_spec = [&](const std::vector<uint32_t>& wt, bool a) { return ggml_vk_mul_mm_cm2_spec(wt, a); };
// F16 x F16
@@ -2596,12 +2661,14 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
auto cm1_create = [&](vk_matmul_pipeline_key key, const std::vector<vk_tile_config>& tc_base,
const std::string& name, size_t len, const void* data, uint32_t pc_size, uint32_t pc) {
auto tc = filter_tc(tc_base, key.type_a, key.mul_mat_id);
+ add_narrow_n_cm1(tc, key.type_a, key.mul_mat_id);
if (!tc.empty()) create_mm_pipelines(key, tc, name, len, data, pc_size, pc, cm1_spec, false, true, 0, true, cm1_pin);
};
auto cm1_create_quant = [&](vk_matmul_pipeline_key key, const std::vector<vk_tile_config>& tc_base,
const std::string& name, size_t len, const void* data, uint32_t pc_size, uint32_t pc) {
spec_fn_t qs = [&, type_a=key.type_a](const std::vector<uint32_t>& wt, bool a) { return ggml_vk_mul_mm_spec_quant(wt, a, (uint32_t)type_a); };
auto tc = filter_tc(tc_base, key.type_a, key.mul_mat_id);
+ add_narrow_n_cm1(tc, key.type_a, key.mul_mat_id);
if (!tc.empty()) create_mm_pipelines(key, tc, name, len, data, pc_size, pc, qs, false, true, 0, true, cm1_pin);
};
// int8 MMQ helper: per-type cm1 shader, warptile passed as-is (carries DEVICE_ARCH in
@@ -3023,8 +3090,13 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
device->matmul_tile_selector = [](uint32_t m, uint32_t n, uint32_t /*k*/, uint32_t shader_core_count,
const std::vector<vk_matmul_pipeline_pair>& configs) -> uint32_t {
if (configs.size() <= 1) return 0;
+ for (uint32_t i = 0; i < configs.size(); ++i) {
+ if (configs[i].narrow_n_max != 0 && n <= configs[i].narrow_n_max) return i;
+ }
uint32_t last = (uint32_t)configs.size() - 1;
- if (configs.size() == 2) {
+ while (last > 0 && configs[last].narrow_n_max != 0) --last;
+ if (last == 0) return 0;
+ if (last == 1) {
uint32_t crossover = configs[0].unaligned->wg_denoms[1];
return (n > crossover) ? 1 : 0;
}
@@ -3058,10 +3130,15 @@ void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) {
device->matmul_tile_selector = [](uint32_t m, uint32_t n, uint32_t /*k*/, uint32_t /*shader_core_count*/,
const std::vector<vk_matmul_pipeline_pair>& configs) -> uint32_t {
if (configs.size() <= 1) return 0;
- if (m <= 32 || n <= 32) return 0;
- if (configs.size() == 2) return 1;
+ for (uint32_t i = 0; i < configs.size(); ++i) {
+ if (configs[i].narrow_n_max != 0 && n <= configs[i].narrow_n_max) return i;
+ }
+ uint32_t last = (uint32_t)configs.size() - 1;
+ while (last > 0 && configs[last].narrow_n_max != 0) --last;
+ if (last == 0 || m <= 32 || n <= 32) return 0;
+ if (last == 1) return 1;
if (m <= 64 || n <= 64) return 1;
- return (uint32_t)configs.size() - 1;
+ return last;
};
device->matmul_id_tile_selector = device->matmul_tile_selector;
}
@@ -4280,6 +4357,7 @@ vk_device ggml_vk_get_device(size_t idx) {
const char* GGML_VK_DISABLE_GRAPH_OPTIMIZE = getenv("GGML_VK_DISABLE_GRAPH_OPTIMIZE");
device->disable_graph_optimize = GGML_VK_DISABLE_GRAPH_OPTIMIZE != nullptr;
+ device->disable_narrow_mul_mat = getenv("GGML_VK_DISABLE_NARROW_MUL_MAT") != nullptr;
bool fp16_storage = false;
bool fp16_compute = false;
@@ -5786,7 +5864,7 @@ void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) {
}
if (vk_perf_logger_enabled) {
- ctx->perf_logger = std::unique_ptr<vk_perf_logger>(new vk_perf_logger());
+ ctx->perf_logger = std::unique_ptr<vk_perf_logger>(new vk_perf_logger(ctx->device));
}
#ifdef GGML_VULKAN_CHECK_RESULTS
@@ -6183,7 +6261,7 @@ void deferred_memset(void * dst, uint32_t val, size_t size, std::vector<vk_stagi
}
}
-static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, uint32_t m, uint32_t n, uint32_t k, bool disable_split_k, const vk_pipeline& pipeline) {
+static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, uint32_t m, uint32_t n, uint32_t k, bool disable_split_k, const vk_pipeline& pipeline, bool narrow_n) {
VK_LOG_DEBUG("ggml_vk_guess_split_k(" << m << ", " << n << ", " << k << ", " << disable_split_k << ")");
if (disable_split_k) {
@@ -6191,9 +6269,9 @@ static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, uint32_t m,
}
const bool rdna = ctx->device->architecture == AMD_RDNA3 || ctx->device->architecture == AMD_RDNA4;
- uint32_t split_k = 1;
- if (rdna && ctx->device->shader_core_count != 0 && k >= 2048 && n <= 32) {
- // Small n gives too few workgroups to hide latency, aim for 8 per CU. Small tiles are one subgroup each.
+ uint32_t split_k = narrow_n && ctx->device->shader_core_count == 0 ? 8u : 1u;
+ if ((narrow_n || rdna) && ctx->device->shader_core_count != 0 && k >= 2048 && n <= 32) {
+ // Small n gives too few workgroups to hide latency, aim for 8 per CU.
const uint32_t tiles = CEIL_DIV(m, pipeline->wg_denoms[0]) * CEIL_DIV(n, pipeline->wg_denoms[1]);
if (tiles < 4 * ctx->device->shader_core_count) {
split_k = std::min(8u, CEIL_DIV(8 * ctx->device->shader_core_count, tiles));
@@ -6201,7 +6279,7 @@ static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, uint32_t m,
}
// On RDNA, also split when n <= 64 only partly fills one tile column.
- if (split_k == 1 && ctx->device->shader_core_count != 0 && (n >= pipeline->wg_denoms[1] || (rdna && n <= 64))) {
+ if (!narrow_n && split_k == 1 && ctx->device->shader_core_count != 0 && (n >= pipeline->wg_denoms[1] || (rdna && n <= 64))) {
// If k is 'large' and the SMs will fill less than halfway, use split_k.
uint32_t m_tiles = CEIL_DIV(m, pipeline->wg_denoms[0]);
uint32_t n_tiles = CEIL_DIV(n, pipeline->wg_denoms[1]);
@@ -6222,21 +6300,10 @@ static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, uint32_t m,
}
}
- // ggml_vk_matmul will align the splits to be a multiple of 256.
- // If this rounded up size would cause the last split to be empty,
- // then reduce the split count.
- while (true) {
- if (split_k == 1) {
- break;
- }
- uint32_t k_split = CEIL_DIV(k, split_k);
- k_split = ROUNDUP_POW2(k_split, 256);
- if (k_split * (split_k - 1) < k) {
- break;
- }
- split_k--;
+ // Each split starts on a 256-value boundary; keep the last split nonempty.
+ while (split_k > 1 && ROUNDUP_POW2(CEIL_DIV(uint64_t(k), split_k), 256) * (split_k - 1) >= k) {
+ --split_k;
}
-
return split_k;
}
@@ -6755,7 +6822,8 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub
const uint64_t y_ne = padded_n * ne10 * ne12 * ne13;
const uint64_t d_ne = ggml_nelements(dst);
- const uint32_t split_k = ggml_vk_guess_split_k(ctx, ne01, ne11, ne10, disable_split_k, pipeline);
+ const bool narrow_n = mmp_map->back().narrow_n_max != 0 && ne11 <= mmp_map->back().narrow_n_max;
+ const uint32_t split_k = ggml_vk_guess_split_k(ctx, ne01, ne11, ne10, disable_split_k, pipeline, narrow_n);
const uint64_t qx_sz = ggml_type_size(src0->type) * x_ne / ggml_blck_size(src0->type);
const uint64_t qy_sz = ggml_type_size(src1->type) * y_ne / ggml_blck_size(src1->type);
@@ -7735,7 +7803,7 @@ void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const st
ggml_vk_mul_mat_vec_q_f16(ctx, subctx, cgraph, node_idx, true);
// 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] <= ggml_vk_mul_mat_vec_max_cols(ctx->device) && src1->ne[2] * src1->ne[3] == 1)) &&
(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 {
@@ -17170,7 +17238,7 @@ std::string vk_perf_logger::get_node_fusion_name(const ggml_tensor * node, const
const uint64_t k = node->src[1]->ne[0];
const uint64_t batch = node->ne[2] * node->ne[3];
std::string name = ggml_op_name(node->op);
- if ((node->op == GGML_OP_MUL_MAT && n <= mul_mat_vec_max_cols) ||
+ if ((node->op == GGML_OP_MUL_MAT && n <= ggml_vk_mul_mat_vec_max_cols(device)) ||
(node->op == GGML_OP_MUL_MAT_ID && node->src[2]->ne[1] == 1)) {
name += "_VEC";
}
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
index 10dda7768..dad9260b9 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp
@@ -41,8 +41,10 @@ layout (constant_id = 2) const uint BN = 64;
layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working with a quant
layout (constant_id = 4) const bool enable_smaller_matrices = false;
-const uint BNover2 = BN / 2;
-const uint BNover4 = enable_smaller_matrices ? (BN / 4) : (BN / 2);
+// Keep tail matrix types valid for each workgroup size.
+const uint BNmin = BLOCK_SIZE <= 128 ? 16 : 32;
+const uint BNover2 = BN / 2 >= BNmin ? BN / 2 : BNmin;
+const uint BNover4 = enable_smaller_matrices ? (BN / 4 >= BNmin ? BN / 4 : BNmin) : BNover2;
layout (constant_id = 5) const uint ALIGNED = 0;
layout (constant_id = 6) const uint subgroup_size = 32;
#define GROUP_A_BUDGET (8u * 1024u * 1024u)
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index c328072bc..3ff5ced0b 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -13760,6 +13760,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
}
}
+ for (int n : {1, 2, 3, 4, 5, 6, 7, 8, 9, 16, 17, 24, 25, 32, 33}) {
+ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_K, GGML_TYPE_F32, 4096, n, 14336, {1, 1}, {1, 1}));
+ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_K, GGML_TYPE_F32, 14336, n, 4096, {1, 1}, {1, 1}));
+ }
+
// Q4_K multi-column mat-vec
for (int64_t m : {4096, 6144, 6272, 14336}) {
for (int bs : {1, 2, 3, 4, 8}) {