feat(amx): add --amx toggle; prefer CPU 'extra' with GPU host+mmap when enabled

- CLI/server/bench: --amx (presence=enabled) -> mparams.amx_enable_mmap
- Loader: with mmap + GPU host buft, prefer CPU 'extra' if supported (AMX repack), else fallback
- llama-bench: add --amx flag to match CLI/server behavior

Author: Gadflyii

common/arg.cpp

@@ -2538,6 +2538,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.use_mmap = false;
         }
     ).set_env("LLAMA_ARG_NO_MMAP"));
+add_opt(common_arg(
+    {"--amx"},
+    "enable AMX-aware CPU repack when mmap is on and a GPU host buffer would be used; prefers CPU \"extra\" buffer types (e.g., AMX) for weights on CPU.",
+    [](common_params & params) {
+        params.amx_enable_mmap = true;
+    }
+));
+
     add_opt(common_arg(
         {"--numa"}, "TYPE",
         "attempt optimizations that help on some NUMA systems\n"

common/common.cpp

@@ -1109,28 +1109,42 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
         mparams.n_gpu_layers = params.n_gpu_layers;
     }
 
-    mparams.main_gpu        = params.main_gpu;
-    mparams.split_mode      = params.split_mode;
-    mparams.tensor_split    = params.tensor_split;
-    mparams.use_mmap        = params.use_mmap;
-    mparams.use_mlock       = params.use_mlock;
-    mparams.check_tensors   = params.check_tensors;
+    mparams.main_gpu      = params.main_gpu;
+    mparams.split_mode    = params.split_mode;
+
+    // NOTE: common_params::tensor_split is a C-array (float [LLAMA_MAX_DEVICES])
+    // Upstream expects a pointer to the first element – do NOT use .data().
+    mparams.tensor_split  = params.tensor_split;
+
+    mparams.use_mmap      = params.use_mmap;
+    mparams.use_mlock     = params.use_mlock;
+    mparams.check_tensors = params.check_tensors;
+
+    // Keep upstream policy: disable extra buffer types when --no-extra-bufts is set
     mparams.use_extra_bufts = !params.no_extra_bufts;
 
+    // NEW: forward the AMX toggle from CLI into model params
+    mparams.amx_enable_mmap = params.amx_enable_mmap;
+
+    // Preserve upstream sentinel handling for KV overrides
     if (params.kv_overrides.empty()) {
         mparams.kv_overrides = NULL;
     } else {
-        GGML_ASSERT(params.kv_overrides.back().key[0] == 0 && "KV overrides not terminated with empty key");
+        GGML_ASSERT(params.kv_overrides.back().key[0] == 0 &&
+            "KV overrides not terminated with empty key");
         mparams.kv_overrides = params.kv_overrides.data();
     }
 
+    // Preserve upstream sentinel handling for tensor buffer overrides
     if (params.tensor_buft_overrides.empty()) {
         mparams.tensor_buft_overrides = NULL;
     } else {
-        GGML_ASSERT(params.tensor_buft_overrides.back().pattern == nullptr && "Tensor buffer overrides not terminated with empty pattern");
+        GGML_ASSERT(params.tensor_buft_overrides.back().pattern == nullptr &&
+            "Tensor buffer overrides not terminated with empty pattern");
         mparams.tensor_buft_overrides = params.tensor_buft_overrides.data();
     }
 
+    // Keep upstream progress callback wiring
     mparams.progress_callback           = params.load_progress_callback;
     mparams.progress_callback_user_data = params.load_progress_callback_user_data;
 

common/common.h

@@ -392,6 +392,8 @@ struct common_params {
     bool check_tensors     = false; // validate tensor data
     bool no_op_offload     = false; // globally disable offload host tensor operations to device
     bool no_extra_bufts    = false; // disable extra buffer types (used for weight repacking)
+    bool amx_enable_mmap  = false; // prefer CPU "extra" buffers when GPU host+mmap is chosen (enable AMX)
+
 
     bool single_turn       = false; // single turn chat conversation
 

include/llama.h

@@ -296,6 +296,7 @@ extern "C" {
         bool use_mlock;       // force system to keep model in RAM
         bool check_tensors;   // validate model tensor data
         bool use_extra_bufts; // use extra buffer types (used for weight repacking)
+        bool amx_enable_mmap; // prefer CPU 'extra' buffers with GPU host+mmap (enable AMX repack on CPU)
     };
 
     // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations

src/llama-model.cpp

@@ -2288,24 +2288,56 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
                 }
             }
 
-            // avoid using a host buffer when using mmap
-            auto * buft_dev = ggml_backend_buft_get_device(buft);
-            if (ml.use_mmap && buft_dev && buft == ggml_backend_dev_host_buffer_type(buft_dev)) {
-                auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
-                if (!cpu_dev) {
-                    throw std::runtime_error("no CPU backend found");
-                }
-                buft = ggml_backend_dev_buffer_type(cpu_dev);
+// avoid using a host buffer when using mmap
+auto * buft_dev = ggml_backend_buft_get_device(buft);
+if (ml.use_mmap && buft_dev && buft == ggml_backend_dev_host_buffer_type(buft_dev)) {
+    auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
+    if (!cpu_dev) {
+        throw std::runtime_error("no CPU backend found");
+    }
+
+    // If enabled, prefer CPU "extra" (AMX) buffer types for weights on CPU; else use CPU default
+    ggml_backend_buffer_type_t cpu_default_buft = ggml_backend_dev_buffer_type(cpu_dev);
+    const bool prefer_cpu_extra = params.amx_enable_mmap;
+
+    if (!prefer_cpu_extra) {
+        buft = cpu_default_buft;
+    } else {
+        ggml_backend_buffer_type_t chosen = nullptr;
+
+        // Iterate available buffer types, skipping device-host buffer types
+        for (const auto & cur : *buft_list) {
+            ggml_backend_dev_t           cur_dev  = cur.first;
+            ggml_backend_buffer_type_t   cur_buft = cur.second;
+
+            if (cur_dev && cur_buft == ggml_backend_dev_host_buffer_type(cur_dev)) {
+                continue;
             }
 
-            if (buft != buft_list->front().second) {
-                n_moved_tensors++;
-                if (!first_moved_tensor) {
-                    first_moved_tensor = t_meta;
-                    first_moved_from_buft = buft_list->front().second;
-                    first_moved_to_buft   = buft;
+            // Prefer CPU "extra" (non-default) if supported for this tensor/op
+            if (cur_dev == cpu_dev && cur_buft != cpu_default_buft) {
+                if (weight_buft_supported(hparams, t_meta, op, cur_buft, cur_dev)) {
+                    chosen = cur_buft;
+                    break;
                 }
             }
+        }
+
+        buft = chosen ? chosen : cpu_default_buft;
+    }
+}
+
+
+// (keep your existing moved-tensors accounting exactly as-is)
+if (buft != buft_list->front().second) {
+    n_moved_tensors++;
+    if (!first_moved_tensor) {
+        first_moved_tensor   = t_meta;
+        first_moved_from_buft = buft_list->front().second;
+        first_moved_to_buft   = buft;
+    }
+}
+
 
             ggml_context * ctx = ctx_for_buft(buft);
 
@@ -19642,6 +19674,7 @@ llama_model_params llama_model_default_params() {
         /*.use_mlock                   =*/ false,
         /*.check_tensors               =*/ false,
         /*.use_extra_bufts             =*/ true,
+        /*.amx_enable_mmap             =*/ false,
     };
 
     return result;