llama: let the backend pick n_ubatch, and default to 2048 on RDNA3.5

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;
         }

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;

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

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;

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;
     }

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)

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,

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 <pp,tg>                                 (default: %s)\n", join(transform_to_str(cmd_params_defaults.n_pg, pair_str), ",").c_str());
     printf("  -d, --n-depth <n>                           (default: %s)\n", join(cmd_params_defaults.n_depth, ",").c_str());
     printf("  -b, --batch-size <n>                        (default: %s)\n", join(cmd_params_defaults.n_batch, ",").c_str());
-    printf("  -ub, --ubatch-size <n>                      (default: %s)\n", join(cmd_params_defaults.n_ubatch, ",").c_str());
+    printf("  -ub, --ubatch-size <n>                      (default: auto)\n");
     printf("  -ctk, --cache-type-k <t>                    (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_k, ggml_type_name), ",").c_str());
     printf("  -ctv, --cache-type-v <t>                    (default: %s)\n", join(transform_to_str(cmd_params_defaults.type_v, ggml_type_name), ",").c_str());
     printf("  -t, --threads <n>                           (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;