vulkan: raise the hoisted row-id limit for mul_mat_id from 256 to 512 experts (#28501)

* vulkan: raise the hoisted row-id limit for mul_mat_id to 512 experts

The expert-count shader (count_experts.comp) sizes its shared arrays
with BLOCK_SIZE, which is 256. Because of that, row-id hoisting is
switched off for any model with more than 256 experts, and every
mul_mat_id workgroup has to rescan the whole ids tensor on its own.
Qwen3.8-Flash-Next has 512 experts and was quietly running on that
slow path.

This change sizes the arrays with a separate MAX_EXPERTS constant (512),
clears them in a loop instead of one entry per thread, and raises the
matching limit on the host side.

On Strix Halo at batch 2048 the expert matmuls drop from 12.5 to 9.5 ms
(iq3_s) and from 14.0 to 7.5 ms (iq4_nl) per op, and prompt processing
gets about 19 % faster at 8k tokens. test-backend-ops MUL_MAT_ID passes
(891/891) with new 512-expert test cases.

Assisted-by: Claude Fable 5.1

* vulkan: raise the hoisted row-id limit for mul_mat_id to 1024 experts

Follow-up to review feedback: 1024 matches LLAMA_MAX_EXPERTS instead of
stopping at 512. The three shared arrays in count_experts.comp grow to
3 * 1024 * 4 = 12 KiB, which fits the 16 KiB that Vulkan guarantees for
maxComputeSharedMemorySize.

Adds mul_mat_id test cases at 1024 experts alongside the existing 512
ones. test-backend-ops MUL_MAT_ID passes on Vulkan (RADV, Strix Halo,
Radeon 8060S): 889/889.
This commit is contained in:
drluoto
2026-09-18 09:00:15 +02:00
committed by GitHub
parent 972d2313bc
commit 5c53396b89
3 changed files with 21 additions and 6 deletions
+3 -1
View File
@@ -7037,7 +7037,9 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context&
// n_as counts, n_as offsets, one total, then one packed row id per (expert, token).
// Hoisting requires 16-bit indices for the packing and a table that fits one binding.
const uint64_t hoisted_row_id_words = 2 * n_as + 1 + nei0 * nei1;
const bool hoist_row_ids = n_as <= 256 && nei0 <= 0xffff && nei1 <= 0xffff &&
// 1024 matches MAX_EXPERTS in count_experts.comp and LLAMA_MAX_EXPERTS. It costs
// 3 * 1024 * 4 = 12 KiB of shared memory, within the 16 KiB Vulkan guarantees.
const bool hoist_row_ids = n_as <= 1024 && nei0 <= 0xffff && nei1 <= 0xffff &&
hoisted_row_id_words * sizeof(uint32_t) <=
ctx->device->properties.limits.maxStorageBufferRange;
@@ -30,9 +30,14 @@ layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in;
layout (binding = 0) readonly buffer A {uint data_a[];};
layout (binding = 1) writeonly buffer D {uint data_d[];};
shared uint vals[BLOCK_SIZE];
shared uint offsets[BLOCK_SIZE];
shared uint cursors[BLOCK_SIZE];
// Upper bound on n_experts for the hoisted row-id path. Must match the limit in
// ggml_vk_mul_mat_id_q_f16 (hoist_row_ids). The non-hoisted reduction below only
// needs BLOCK_SIZE entries.
#define MAX_EXPERTS 1024
shared uint vals[MAX_EXPERTS];
shared uint offsets[MAX_EXPERTS];
shared uint cursors[MAX_EXPERTS];
// data_d layout when p.hoist_row_ids is set:
// [0, n_experts) per-expert row count
@@ -46,8 +51,8 @@ void main() {
const uint tid = gl_LocalInvocationID.x;
if (p.hoist_row_ids != 0) {
if (tid < p.n_experts) {
vals[tid] = 0;
for (uint e = tid; e < p.n_experts; e += BLOCK_SIZE) {
vals[e] = 0;
}
barrier();
+8
View File
@@ -10077,6 +10077,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
// gpt-oss issue with Vulkan mmq_id
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_MXFP4, GGML_TYPE_F32, 32, 2, false, 2880, 32, 2880));
// more than 256 experts (hoisted row-id path): 512 as in Qwen3.8-Flash-Next,
// and 1024 at the LLAMA_MAX_EXPERTS limit
for (int n : {1, 5, 64, 300}) {
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ3_S, GGML_TYPE_F32, 512, 10, false, 128, n, 512));
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_0, GGML_TYPE_F32, 512, 10, false, 256, n, 128));
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_IQ3_S, GGML_TYPE_F32, 1024, 10, false, 128, n, 512));
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_0, GGML_TYPE_F32, 1024, 10, false, 256, n, 128));
}
test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q4_0, GGML_TYPE_F32, 32, 2, false, 2880, 32, 2880));
// multiple blocks per row: exercises the block-stride loop and the