From 42f7e8948b686d86392c8299d66790565350f875 Mon Sep 17 00:00:00 2001 From: Daniele Zannotti Date: Thu, 17 Sep 2026 23:08:09 +0100 Subject: [PATCH] cuda: enable the BF16 WMMA matmul path only for qwen4exp mmb_enabled() was true for every gfx1151 device, and mmb_quant_type() claims almost every quant type, so the BF16 WMMA path took MUL_MAT/MUL_MAT_ID away from MMQ for every model. Its tiles and fusions are tuned for the qwen4exp shapes (the hardcoded 320/10240 gate-mix shapes among them); elsewhere it loses. Measured on gfx1151, ROCm 7.2.1, llama-bench -ngl 99 -fa on, prefill t/s: model ub mmb on this PR pre-#63 (cfe6bb14e) Signal-3.8-27B Q4_K_XL 512 119.6 467.7 461.2 Signal-3.8-27B Q4_K_XL 16384 123.3 421.8 416.1 Qwen3.8-27B GSQ-RCO IQ3_S 512 116.3 438.8 434.9 gemma-4-26B-A4B Q4_0 2048 1502.2 2087.6 2083.0 gpt-oss-20b MXFP4 2048 2066.2 2406.6 2411.3 Dense qwen35 paid 3.4-3.9x, MoE 14-28%. Token generation is unchanged in every case (the path only engages from 512 tokens up, so it never runs during decode). The llama layer now opts a model in by architecture through the backend's get_proc_address, so qwen4exp keeps the path it was tuned for and nothing else pays for it. A process that loads several models shares the flag, last load wins; that is noted at the definition. Not measured: qwen4exp itself, for want of a checkpoint on this box. The gain there is pwilkin's (#63) and this does not change that path. Co-Authored-By: Claude Opus 5 (1M context) --- ggml/src/ggml-cuda/ggml-cuda.cu | 8 ++++++++ ggml/src/ggml-cuda/mmb.cu | 18 +++++++++++++++++- ggml/src/ggml-cuda/mmb.cuh | 2 ++ src/llama.cpp | 11 +++++++++++ 4 files changed, 38 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 255fcf8294b..b0140687764 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -7481,8 +7481,16 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t GGML_UNUSED(reg); } +// enable the BF16 WMMA matmul path (mmb); the llama layer calls this per model, by architecture +static void ggml_backend_cuda_set_mmb_enabled(bool enable) { + ggml_cuda_mmb_set_opt_in(enable); +} + static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { GGML_UNUSED(reg); + if (strcmp(name, "ggml_backend_cuda_set_mmb_enabled") == 0) { + return (void *)ggml_backend_cuda_set_mmb_enabled; + } if (strcmp(name, "ggml_backend_comm_init") == 0) { return (void *)ggml_backend_cuda_comm_init; } diff --git a/ggml/src/ggml-cuda/mmb.cu b/ggml/src/ggml-cuda/mmb.cu index 4fc7e744097..0f73f3fdb27 100644 --- a/ggml/src/ggml-cuda/mmb.cu +++ b/ggml/src/ggml-cuda/mmb.cu @@ -1,4 +1,6 @@ #include "mmb.cuh" + +#include #include "unary.cuh" #include #include @@ -679,8 +681,17 @@ static const uint16_t * mmb_shadow_lookup(const ggml_tensor * w) { auto it = g_mmb_shadow.find(w->data); return it == g_mmb_shadow.end() ? nullptr : it->second; } -// the BF16 WMMA kernels are RDNA3.5 (gfx1151) work: other devices keep the MMQ/MMVQ paths +// The BF16 WMMA kernels are RDNA3.5 (gfx1151) work, and their tiles and fusions are tuned for the +// qwen4exp shapes. On other architectures they take MUL_MATs away from MMQ and lose: dense qwen35 +// prefill measured 3.4-3.6x slower on gfx1151 at every ubatch from 512 to 16384, and MoE 1.1-1.2x. +// So the llama layer opts a model in by arch (ggml_backend_cuda_set_mmb_enabled); default is off. +// A process that loads several models shares this flag: the last load wins. +static std::atomic g_mmb_opt_in{false}; + bool mmb_enabled() { + if (!g_mmb_opt_in.load(std::memory_order_relaxed)) { + return false; + } const int id = ggml_cuda_get_device(); return GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[id].cc); } @@ -697,6 +708,11 @@ bool mmb_glu() { return true; } } // namespace +// opted in per model by the llama layer, by architecture (see the note on mmb_enabled above) +void ggml_cuda_mmb_set_opt_in(bool enable) { + g_mmb_opt_in.store(enable, std::memory_order_relaxed); +} + const uint16_t * ggml_cuda_mmb_cache_lookup(const ggml_tensor * t) { const ggml_tensor * root = mmb_root(t); for (auto & e : g_mmb_slots) if (e.buf && e.root == root && e.data == t->data) return e.buf->get(); diff --git a/ggml/src/ggml-cuda/mmb.cuh b/ggml/src/ggml-cuda/mmb.cuh index 431785828ce..09f4e0a0bee 100644 --- a/ggml/src/ggml-cuda/mmb.cuh +++ b/ggml/src/ggml-cuda/mmb.cuh @@ -1,6 +1,8 @@ #pragma once #include "common.cuh" // Quantized-weight BF16 WMMA GEMM on gfx1151, from 512 tokens up. +// opted in per model by the llama layer (see ggml_backend_cuda_set_mmb_enabled): tuned for qwen4exp +void ggml_cuda_mmb_set_opt_in(bool enable); bool ggml_cuda_mmb_supported_mm (const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst); bool ggml_cuda_mmb_supported_mmid(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * dst); void ggml_cuda_mul_mat_mmb (ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); diff --git a/src/llama.cpp b/src/llama.cpp index ad8e443882a..c3683237532 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -344,6 +344,17 @@ static std::pair llama_model_load(struct gguf_context * meta model->hparams.vocab_only = params.vocab_only; model->hparams.no_alloc = params.no_alloc; + // The CUDA/HIP BF16 WMMA matmul path is tuned for the qwen4exp shapes; on other architectures it + // takes MUL_MATs away from MMQ and loses (dense qwen35 prefill measured 3.4x slower on gfx1151). + // Opt in by architecture rather than by device. + for (size_t i = 0; i < ggml_backend_reg_count(); ++i) { + auto * set_mmb_fn = (void (*)(bool)) ggml_backend_reg_get_proc_address( + ggml_backend_reg_get(i), "ggml_backend_cuda_set_mmb_enabled"); + if (set_mmb_fn) { + set_mmb_fn(model->arch == LLM_ARCH_QWEN4EXP); + } + } + try { model->load_hparams(ml); } catch(const std::exception & e) {