From ec60a4dd2864556ec4a10072b392fedfa753758a Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Wed, 29 Jul 2026 02:55:21 -0700 Subject: [PATCH] llama: let the backend pick n_ubatch, and default to 2048 on RDNA3.5 The 512-token default physical batch predates these iGPUs. On gfx1151 a 512-token ubatch leaves the matrix cores short of work and re-reads the weights more times than necessary, so prompt processing gives up throughput on some architectures while gaining it on others. Rather than hardcode a number, add a way for a backend to state a preference. Adds LLAMA_N_UBATCH_AUTO as the new llama_context_default_params() value for n_ubatch. When set, llama_context queries each device for an optional "ggml_backend_get_optimal_ubatch" proc address; devices with no preference return 0 and are ignored, and if several disagree the smallest preference wins. Anything other than LLAMA_N_UBATCH_AUTO is passed through untouched, so an explicit -ub 512 still means 512. The CUDA/HIP backend returns 2048 for RDNA3.5 and 0 elsewhere, so no other architecture changes behaviour. llama-bench now reports the resolved value via llama_n_ubatch() instead of the requested one, since the request may be "auto". Measured on gfx1151 (Radeon 8060S), Q4_K_M, pp2048, -r 30 on an idle GPU. This is not a uniform win and the default is chosen on the strength of the MoE cases: GLM-4.7-Flash 1013 -> 1132 +11.7% gemma-4-26B-A4B-it 1527 -> 1668 +9.3% Qwen2.5-7B-Instruct 1541 -> 1588 +3.1% Qwen2.5-0.5B-Instruct 15606 -> 16021 +2.7% gemma-2b-it 4317 -> 4384 +1.6% gemma-3-12b-it 857 -> 866 +1.0% Qwen2.5-3B-Instruct 3218 -> 3230 +0.4% Llama-3.1-8B-Instruct 1459 -> 1422 -2.5% Qwen3.6-27B 386 -> 379 -1.8% Qwen3-8B 1433 -> 1389 -3.0% Qwen3-4B 2447 -> 2353 -3.9% Qwen3-1.7B 5410 -> 5125 -5.3% Qwen3.5-0.8B 9206 -> 8097 -12.1% Across 20 model/repo combinations: 5 improve by more than 1%, 4 are flat, and 11 regress by more than 1%. Median -1.8%, mean -0.8%. The large wins are concentrated in MoE models and the largest regressions in small dense models, where a 2048-token ubatch is a substantial fraction of the whole prompt. --- common/arg.cpp | 2 +- common/common.cpp | 2 +- common/common.h | 2 +- ggml/include/ggml-backend.h | 2 ++ ggml/src/ggml-cuda/ggml-cuda.cu | 13 +++++++++++++ include/llama.h | 5 ++++- src/llama-context.cpp | 28 ++++++++++++++++++++++++++-- tools/llama-bench/llama-bench.cpp | 9 +++++---- 8 files changed, 53 insertions(+), 10 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 8a78291658eb..b6bcb9a30ab9 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1432,7 +1432,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex ).set_env("LLAMA_ARG_BATCH")); add_opt(common_arg( {"-ub", "--ubatch-size"}, "N", - string_format("physical maximum batch size (default: %d)", params.n_ubatch), + "physical maximum batch size (default: auto, chosen by the backend)", [](common_params & params, int value) { params.n_ubatch = value; } diff --git a/common/common.cpp b/common/common.cpp index 8f13217ab442..71ddeb7d925f 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1595,7 +1595,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.n_rs_seq = params.speculative.need_n_rs_seq(); cparams.n_outputs_max = std::max(params.n_outputs_max, 0); cparams.n_batch = params.n_batch; - cparams.n_ubatch = params.n_ubatch; + cparams.n_ubatch = params.n_ubatch < 0 ? LLAMA_N_UBATCH_AUTO : params.n_ubatch; cparams.n_threads = params.cpuparams.n_threads; cparams.n_threads_batch = params.cpuparams_batch.n_threads == -1 ? params.cpuparams.n_threads : params.cpuparams_batch.n_threads; diff --git a/common/common.h b/common/common.h index 153531700886..4d437cdd67eb 100644 --- a/common/common.h +++ b/common/common.h @@ -450,7 +450,7 @@ struct common_params { int32_t n_predict = -1; // max. number of new tokens to predict, -1 == no limit int32_t n_ctx = 0; // context size, 0 == context the model was trained with int32_t n_batch = 2048; // logical batch size for prompt processing (must be >=32 to use BLAS) - int32_t n_ubatch = 512; // physical batch size for prompt processing (must be >=32 to use BLAS) + int32_t n_ubatch = -1; // physical batch size for prompt processing (must be >=32 to use BLAS; -1 = backend decides) int32_t n_keep = 0; // number of tokens to keep from initial prompt int32_t n_chunks = -1; // max number of chunks to process (-1 = unlimited) int32_t n_parallel = 1; // number of parallel sequences to decode diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index 2924fdbe9884..e08e8d35488d 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -215,6 +215,8 @@ extern "C" { typedef ggml_backend_buffer_type_t * (*ggml_backend_dev_get_extra_bufts_t)(ggml_backend_dev_t device); // Set the abort callback for the backend typedef void (*ggml_backend_set_abort_callback_t)(ggml_backend_t backend, ggml_abort_callback abort_callback, void * abort_callback_data); + // Physical batch size the device prefers for prompt processing, or 0 for "no preference" + typedef uint32_t (*ggml_backend_get_optimal_ubatch_t)(ggml_backend_dev_t device); // Get a list of feature flags supported by the backend (returns a NULL-terminated array) struct ggml_backend_feature { const char * name; diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 67f8a93248d1..06c0cdb57272 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5400,8 +5400,21 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t GGML_UNUSED(reg); } +// RDNA3.5 (gfx1150/gfx1151) prefers a larger physical batch than the 512 default. Returns 0 on +// every other architecture, meaning "no opinion, keep the caller's default". +static uint32_t ggml_backend_cuda_get_optimal_ubatch(ggml_backend_dev_t dev) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + + const int cc = ggml_cuda_info().devices[ctx->device].cc; + + return GGML_CUDA_CC_IS_RDNA3_5(cc) ? 2048 : 0; +} + 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_get_optimal_ubatch") == 0) { + return (void *)ggml_backend_cuda_get_optimal_ubatch; + } if (strcmp(name, "ggml_backend_comm_init") == 0) { return (void *)ggml_backend_cuda_comm_init; } diff --git a/include/llama.h b/include/llama.h index a311ac202357..0f1ba6d503c4 100644 --- a/include/llama.h +++ b/include/llama.h @@ -38,6 +38,9 @@ #define LLAMA_TOKEN_NULL -1 +// llama_context_params::n_ubatch - let the backend pick the physical batch size +#define LLAMA_N_UBATCH_AUTO 0xFFFFFFFF + #define LLAMA_FILE_MAGIC_GGLA 0x67676c61u // 'ggla' #define LLAMA_FILE_MAGIC_GGSN 0x6767736eu // 'ggsn' #define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq' @@ -340,7 +343,7 @@ extern "C" { struct llama_context_params { uint32_t n_ctx; // text context, 0 = from model uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode - uint32_t n_ubatch; // physical maximum batch size + uint32_t n_ubatch; // physical maximum batch size (LLAMA_N_UBATCH_AUTO = ask the backend) uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5edfc85abfda..ac42fcd880cd 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -229,7 +229,31 @@ llama_context::llama_context( // with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch; - cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch); + // LLAMA_N_UBATCH_AUTO: let the devices state a preferred physical batch size. Devices that have + // no preference are ignored; if several disagree the smallest preference wins. + uint32_t n_ubatch_req = params.n_ubatch; + if (n_ubatch_req == LLAMA_N_UBATCH_AUTO) { + n_ubatch_req = 512; + + bool have_hint = false; + for (const auto & dev : model.devices) { + auto * fn = (ggml_backend_get_optimal_ubatch_t) ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(dev.dev), "ggml_backend_get_optimal_ubatch"); + if (fn == nullptr) { + continue; + } + + const uint32_t hint = fn(dev.dev); + if (hint == 0) { + continue; + } + + n_ubatch_req = have_hint ? std::min(n_ubatch_req, hint) : hint; + have_hint = true; + } + } + + cparams.n_ubatch = std::min(cparams.n_batch, n_ubatch_req == 0 ? params.n_batch : n_ubatch_req); cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max; @@ -3419,7 +3443,7 @@ llama_context_params llama_context_default_params() { llama_context_params result = { /*.n_ctx =*/ 512, /*.n_batch =*/ 2048, - /*.n_ubatch =*/ 512, + /*.n_ubatch =*/ LLAMA_N_UBATCH_AUTO, /*.n_seq_max =*/ 1, /*.n_rs_seq =*/ 0, /*.n_outputs_max =*/ 0, diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index c96508854566..a74620f97410 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -381,7 +381,7 @@ static const cmd_params cmd_params_defaults = { /* n_pg */ {}, /* n_depth */ { 0 }, /* n_batch */ { 2048 }, - /* n_ubatch */ { 512 }, + /* n_ubatch */ { -1 }, /* type_k */ { GGML_TYPE_F16 }, /* type_v */ { GGML_TYPE_F16 }, /* n_threads */ { common_cpu_get_num_math() }, @@ -453,7 +453,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -pg (default: %s)\n", join(transform_to_str(cmd_params_defaults.n_pg, pair_str), ",").c_str()); printf(" -d, --n-depth (default: %s)\n", join(cmd_params_defaults.n_depth, ",").c_str()); printf(" -b, --batch-size (default: %s)\n", join(cmd_params_defaults.n_batch, ",").c_str()); - printf(" -ub, --ubatch-size (default: %s)\n", join(cmd_params_defaults.n_ubatch, ",").c_str()); + printf(" -ub, --ubatch-size (default: auto)\n"); printf(" -ctk, --cache-type-k (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_k, ggml_type_name), ",").c_str()); printf(" -ctv, --cache-type-v (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_v, ggml_type_name), ",").c_str()); printf(" -t, --threads (default: %s)\n", join(cmd_params_defaults.n_threads, ",").c_str()); @@ -1253,7 +1253,7 @@ struct cmd_params_instance { cparams.n_ctx = n_prompt + n_gen + n_depth; cparams.n_batch = n_batch; - cparams.n_ubatch = n_ubatch; + cparams.n_ubatch = n_ubatch < 0 ? LLAMA_N_UBATCH_AUTO : n_ubatch; cparams.type_k = type_k; cparams.type_v = type_v; cparams.offload_kqv = !no_kv_offload; @@ -1463,7 +1463,8 @@ struct test { model_size = llama_model_size(lmodel); model_n_params = llama_model_n_params(lmodel); n_batch = inst.n_batch; - n_ubatch = inst.n_ubatch; + // inst.n_ubatch may be -1 (auto), so report what the context actually resolved to + n_ubatch = llama_n_ubatch(ctx); n_threads = inst.n_threads; cpu_mask = inst.cpu_mask; cpu_strict = inst.cpu_strict;