Merge branch 'master' into xsn/common_json

common/common.cpp

@@ -1294,34 +1294,11 @@ common_init_result::common_init_result(common_params & params, bool model_only)
     if (params.fit_params) {
         COM_TRC("%s", "fitting params to device memory ...\n");
         COM_TRC("%s", "(for bugs during this step try to reproduce them with -fit off, or provide --verbose logs if the bug only occurs with -fit on)\n");
-
-        // the draft context is created from the same base params and follows the main context, fit both together
-        const bool has_draft = params.speculative.has_dft();
-        const bool spec_mtp  = std::find(params.speculative.types.begin(), params.speculative.types.end(),
-            COMMON_SPECULATIVE_TYPE_DRAFT_MTP) != params.speculative.types.end();
-
-        common_params params_dft = common_base_params_to_speculative(params);
-
-        auto mparams_dft = common_model_params_to_llama(params_dft);
-        auto cparams_dft = common_context_params_to_llama(params_dft);
-        if (spec_mtp) {
-            cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
-        }
-        cparams_dft.n_rs_seq = 0;
-
-        const common_fit_extra_model extra = {
-            /*.path_model   =*/ params_dft.model.path.c_str(),
-            /*.mparams      =*/ &mparams_dft,
-            /*.cparams      =*/ &cparams_dft,
-            /*.shares_model =*/ !has_draft, // an MTP context runs on the weights of the main model
-        };
-
         common_fit_params(params.model.path.c_str(), &mparams, &cparams,
             params.tensor_split,
             params.tensor_buft_overrides.data(),
             params.fit_params_target.data(),
             params.fit_params_min_ctx,
-            has_draft || spec_mtp ? &extra : nullptr,
             params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
     }
 

common/fit.cpp

@@ -178,7 +178,7 @@ common_device_memory_data_vec common_get_device_memory_data(
 static void common_params_fit_impl(
         const char * path_model, struct llama_model_params * mparams, struct llama_context_params * cparams,
         float * tensor_split, struct llama_model_tensor_buft_override * tensor_buft_overrides,
-        size_t * margins_s, uint32_t n_ctx_min, const common_fit_extra_model * extra, enum ggml_log_level log_level) {
+        size_t * margins_s, uint32_t n_ctx_min, enum ggml_log_level log_level) {
     if (mparams->split_mode == LLAMA_SPLIT_MODE_TENSOR) {
         throw common_params_fit_exception("llama_params_fit is not implemented for SPLIT_MODE_TENSOR, abort");
     }
@@ -191,92 +191,10 @@ static void common_params_fit_impl(
     uint32_t hp_nct = 0; // hparams.n_ctx_train
     uint32_t hp_nex = 0; // hparams.n_expert
 
-    // with non-unified kv, we need to take into account n_streams
-    // for example, if memory can hold more than model's trained context size, we must extend the n_ctx to hold enough n_streams
-    const uint32_t n_streams  = cparams->kv_unified ? 1 : std::max<uint32_t>(1, cparams->n_seq_max);
-    const bool     n_ctx_auto = cparams->n_ctx == 0;
-
-    dmds_t   dmds_extra;       // memory of the extra model, laid out on the devices of the main model
-    uint32_t n_ctx_extra = 0;  // context that memory was measured at
-
-    // the extra model competes for the same memory as the main model, add it to every measurement
-    // its memory is measured again whenever the context it follows changes
-    auto add_extra_memory = [&](dmds_t & dmds) {
-        if (extra == nullptr) {
-            return;
-        }
-
-        if (dmds_extra.empty() || n_ctx_extra != cparams->n_ctx) {
-            std::vector<ggml_backend_dev_t> devs_extra;
-            uint32_t ngl_extra = 0;
-            uint32_t nct_extra = 0;
-            uint32_t nex_extra = 0;
-
-            extra->cparams->n_ctx = cparams->n_ctx;
-
-            LOG_TRC("%s: getting device memory data for the extra model at a context size of %" PRIu32 ":\n",
-                __func__, cparams->n_ctx);
-
-            dmds_t measured;
-            try {
-                measured = common_get_device_memory_data_impl(
-                    extra->path_model, extra->mparams, extra->cparams, devs_extra, ngl_extra, nct_extra, nex_extra, log_level);
-            } catch (const std::runtime_error & e) {
-                // the extra model is optional, fit the main model alone rather than giving up
-                LOG_WRN("%s: failed to measure the memory of the extra model, fitting without it: %s\n", __func__, e.what());
-                dmds_extra = dmds_t(devs.size() + 1);
-                n_ctx_extra = cparams->n_ctx;
-                return;
-            }
-
-            dmds_extra = dmds_t(devs.size() + 1);
-            dmds_extra.back().mb = measured.back().mb;
-            for (size_t je = 0; je < devs_extra.size(); je++) {
-                for (size_t id = 0; id < devs.size(); id++) {
-                    if (devs_extra[je] == devs[id]) {
-                        dmds_extra[id].mb.model   += measured[je].mb.model;
-                        dmds_extra[id].mb.context += measured[je].mb.context;
-                        dmds_extra[id].mb.compute += measured[je].mb.compute;
-                        break;
-                    }
-                }
-            }
-            if (extra->shares_model) {
-                for (llama_device_memory_data & dmd : dmds_extra) {
-                    dmd.mb.model = 0;
-                }
-            }
-
-            n_ctx_extra = cparams->n_ctx;
-        }
-
-        for (size_t id = 0; id < dmds.size(); id++) {
-            dmds[id].mb.model   += dmds_extra[id].mb.model;
-            dmds[id].mb.context += dmds_extra[id].mb.context;
-            dmds[id].mb.compute += dmds_extra[id].mb.compute;
-        }
-    };
-
     // step 1: get data for default parameters and check whether any changes are necessary in the first place
 
     LOG_TRC("%s: getting device memory data for initial parameters:\n", __func__);
-    dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
-
-    // saturate instead of overflowing, this also preserves the UINT32_MAX sentinel of n_ctx_min:
-    const uint32_t n_ctx_max       = (uint32_t) std::min<uint64_t>(uint64_t(hp_nct)    * n_streams, UINT32_MAX);
-    const uint32_t n_ctx_min_total = (uint32_t) std::min<uint64_t>(uint64_t(n_ctx_min) * n_streams, UINT32_MAX);
-
-    // llama_context would use only hp_nct in total for n_ctx == 0, resolve the context before measuring anything else:
-    if (n_ctx_auto) {
-        cparams->n_ctx = n_ctx_max;
-        if (n_streams > 1) {
-            LOG_TRC("%s: context size unset and KV cache not unified -> using %" PRIu32 " for %" PRIu32 " sequences:\n",
-                __func__, n_ctx_max, n_streams);
-            dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
-        }
-    }
-    add_extra_memory(dmds_full);
-
+    const dmds_t dmds_full = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
     const size_t nd = devs.size(); // number of devices
 
     std::vector<int64_t> margins; // this function uses int64_t rather than size_t for memory sizes to more conveniently handle deficits
@@ -389,8 +307,8 @@ static void common_params_fit_impl(
                     "%s: cannot meet free memory targets on all devices, need to use %" PRId64 " MiB less in total\n",
                     __func__, -global_surplus/MiB);
             }
-            if (n_ctx_auto) {
-                if (n_ctx_max > n_ctx_min_total) {
+            if (cparams->n_ctx == 0) {
+                if (hp_nct > n_ctx_min) {
                     int64_t sum_used_target = sum_free;
                     if (nd == 0) {
                         sum_used_target -= margins[0];
@@ -410,9 +328,8 @@ static void common_params_fit_impl(
                     }
 
                     int64_t sum_projected_used_min_ctx = 0;
-                    cparams->n_ctx = n_ctx_min_total;
-                    dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
-                    add_extra_memory(dmds_min_ctx);
+                    cparams->n_ctx = n_ctx_min;
+                    const dmds_t dmds_min_ctx = common_get_device_memory_data_impl(path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
                     if (nd == 0) {
                         sum_projected_used_min_ctx = dmds_min_ctx.back().mb.total();
                     } else {
@@ -422,16 +339,14 @@ static void common_params_fit_impl(
                     }
                     if (sum_used_target > sum_projected_used_min_ctx) {
                         // linear interpolation between minimum and maximum context size:
-                        cparams->n_ctx += (n_ctx_max - n_ctx_min_total) * (sum_used_target - sum_projected_used_min_ctx)
+                        cparams->n_ctx += (hp_nct - n_ctx_min) * (sum_used_target - sum_projected_used_min_ctx)
                             / (sum_projected_used - sum_projected_used_min_ctx);
-                        // round down context for CUDA backend, keep it divisible by the number of streams:
-                        const uint32_t align = 256 * n_streams;
-                        cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % align, n_ctx_min_total);
+                        cparams->n_ctx = std::max(cparams->n_ctx - cparams->n_ctx % 256, n_ctx_min); // round down context for CUDA backend
 
-                        const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (n_ctx_max - n_ctx_min_total);
-                        const int64_t memory_reduction = (n_ctx_max - cparams->n_ctx) * bytes_per_ctx;
+                        const int64_t bytes_per_ctx = (sum_projected_used - sum_projected_used_min_ctx) / (hp_nct - n_ctx_min);
+                        const int64_t memory_reduction = (hp_nct - cparams->n_ctx) * bytes_per_ctx;
                         LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n",
-                            __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB);
+                            __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB);
                         if (nd <= 1) {
                             LOG_TRC("%s: entire model can be fit by reducing context\n", __func__);
                             return;
@@ -440,14 +355,14 @@ static void common_params_fit_impl(
                     } else {
                         const int64_t memory_reduction = sum_projected_used - sum_projected_used_min_ctx;
                         LOG_TRC("%s: context size reduced from %" PRIu32 " to %" PRIu32 " -> need %" PRId64 " MiB less memory in total\n",
-                            __func__, n_ctx_max, cparams->n_ctx, memory_reduction/MiB);
+                            __func__, hp_nct, cparams->n_ctx, memory_reduction/MiB);
                     }
                 } else {
                     if (n_ctx_min == UINT32_MAX) {
-                        LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, n_ctx_max);
+                        LOG_TRC("%s: user has requested full context size of %" PRIu32 " -> no change\n", __func__, hp_nct);
                     } else {
                         LOG_TRC("%s: default model context size is %" PRIu32 " which is <= the min. context size of %" PRIu32 " -> no change\n",
-                            __func__, n_ctx_max, n_ctx_min_total);
+                            __func__, hp_nct, n_ctx_min);
                     }
                 }
             } else {
@@ -592,9 +507,8 @@ static void common_params_fit_impl(
         llama_model_params mparams_copy = *mparams;
         set_ngl_tensor_split_tbo(ngl_per_device, overflow_bufts, mparams_copy);
 
-        dmds_t dmd_nl = common_get_device_memory_data_impl(
+        const dmds_t dmd_nl = common_get_device_memory_data_impl(
             path_model, &mparams_copy, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
-        add_extra_memory(dmd_nl);
 
         LOG_TRC("%s: memory for test allocation by device:\n", func_name);
         for (size_t id = 0; id < nd; id++) {
@@ -621,9 +535,8 @@ static void common_params_fit_impl(
         mparams->tensor_buft_overrides = tensor_buft_overrides;
 
         LOG_TRC("%s: getting device memory data with all MoE tensors moved to system memory:\n", __func__);
-        dmds_t dmds_cpu_moe = common_get_device_memory_data_impl(
+        const dmds_t dmds_cpu_moe = common_get_device_memory_data_impl(
             path_model, mparams, cparams, devs, hp_ngl, hp_nct, hp_nex, log_level);
-        add_extra_memory(dmds_cpu_moe);
 
         for (size_t id = 0; id < nd; id++) {
             global_surplus_cpu_moe += dmds_cpu_moe[id].free;
@@ -883,12 +796,11 @@ enum common_params_fit_status common_fit_params(
         llama_model_tensor_buft_override * tensor_buft_overrides,
         size_t * margins,
         uint32_t n_ctx_min,
-        const common_fit_extra_model * extra,
         ggml_log_level log_level) {
     const int64_t t0_us = llama_time_us();
     common_params_fit_status status = COMMON_PARAMS_FIT_STATUS_SUCCESS;
     try {
-        common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, extra, log_level);
+        common_params_fit_impl(path_model, mparams, cparams, tensor_split, tensor_buft_overrides, margins, n_ctx_min, log_level);
         LOG_TRC("%s: successfully fit params to free device memory\n", __func__);
     } catch (const common_params_fit_exception & e) {
         LOG_WRN("%s: failed to fit params to free device memory: %s\n", __func__, e.what());

common/fit.h

@@ -11,16 +11,6 @@ enum common_params_fit_status {
     COMMON_PARAMS_FIT_STATUS_ERROR   = 2, // a hard error occurred, e.g. because no model could be found at the specified path
 };
 
-// a second model that shares the devices of the main model, e.g. a draft model
-//   - its context follows the context of the main model, so its memory is measured again whenever that context changes
-//   - shares_model tells the fit that the weights are already counted in the main model, as for an MTP context
-struct common_fit_extra_model {
-    const char * path_model;
-    llama_model_params * mparams;
-    llama_context_params * cparams;
-    bool shares_model;
-};
-
 // fits mparams and cparams to free device memory (assumes system memory is unlimited)
 //   - returns true if the parameters could be successfully modified to fit device memory
 //   - this function is NOT thread safe because it modifies the global llama logger state
@@ -34,7 +24,6 @@ common_params_fit_status common_fit_params(
    llama_model_tensor_buft_override * tensor_buft_overrides, // writable buffer for overrides, needs at least llama_max_tensor_buft_overrides elements
                              size_t * margins,               // margins of memory to leave per device in bytes
                            uint32_t   n_ctx_min,             // minimum context size to set when trying to reduce memory use
-      const common_fit_extra_model * extra,                  // model to fit alongside the main one, nullptr if there is none
                      ggml_log_level   log_level);            // minimum log level to print during fitting, lower levels go to debug log
 
 // print estimated memory to stdout

common/speculative.cpp

@@ -2388,9 +2388,6 @@ common_speculative_init_result::common_speculative_init_result(
         cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
     }
 
-    // the draft context holds as many tokens per sequence as the target context
-    cparams.n_ctx = llama_n_ctx(ctx_tgt);
-
     // note: for small models maybe we can set this to the maximum possible draft from all speculative types
     //       the extra memory for small models is likely negligible?
     cparams.n_rs_seq  = 0;

tools/fit-params/fit-params.cpp

@@ -33,7 +33,6 @@ int llama_fit_params(int argc, char ** argv) {
     if (!params.fit_params_print) {
         const common_params_fit_status status = common_fit_params(params.model.path.c_str(), &mparams, &cparams,
                 params.tensor_split, params.tensor_buft_overrides.data(), params.fit_params_target.data(), params.fit_params_min_ctx,
-                nullptr,
                 params.verbosity >= LOG_LEVEL_DEBUG ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
         if (status != COMMON_PARAMS_FIT_STATUS_SUCCESS) {
             LOG_ERR("%s: failed to fit CLI arguments to free memory, exiting...\n", __func__);

tools/llama-bench/llama-bench.cpp

@@ -2294,7 +2294,6 @@ int llama_bench(int argc, char ** argv) {
                 fit_overrides.data(),
                 margins.data(),
                 inst.fit_min_ctx,
-                nullptr,
                 params.verbose ? GGML_LOG_LEVEL_DEBUG : GGML_LOG_LEVEL_ERROR);
        }
 

tools/mtmd/clip-graph.h

@@ -137,15 +137,9 @@ struct clip_graph {
             int il,
             ggml_tensor * sinks = nullptr) const;
 
-    // implementation of the 2D RoPE using two ggml_rope_ext calls
-    //
-    // unlike GGML_ROPE_TYPE_VISION which forces NEOX ordering, this rotates adjacent pairs (normal ordering)
-    //
-    // example:
-    //  given a single head with size = 8 --> [00000000]
-    //  dims [0, 4) rotate with pos_a, dims [4, 8) rotate with pos_b --> [aaaabbbb]
-    //  interleave_freq = false --> both halves use the same inv_freq set (like GGML_ROPE_TYPE_VISION)
-    //  interleave_freq = true  --> first half uses even inv_freq, second half uses odd inv_freq (used by pixtral)
+    // implementation of the 2D RoPE without adding a new op in ggml
+    // this is not efficient (use double the memory), but works on all backends
+    // TODO: there was a more efficient which relies on ggml_view and ggml_rope_ext_inplace, but the rope inplace does not work well with non-contiguous tensors ; we should fix that and revert back to the original implementation in https://github.com/ggml-org/llama.cpp/pull/13065
     ggml_tensor * build_rope_2d(
         ggml_context * ctx0,
         ggml_tensor * cur,

tools/mtmd/clip.cpp

@@ -819,6 +819,8 @@ ggml_tensor * clip_graph::build_attn(
 }
 
 // implementation of the 2D RoPE without adding a new op in ggml
+// this is not efficient (use double the memory), but works on all backends
+// TODO: there was a more efficient which relies on ggml_view and ggml_rope_ext_inplace, but the rope inplace does not work well with non-contiguous tensors ; we should fix that and revert back to the original implementation in https://github.com/ggml-org/llama.cpp/pull/13065
 ggml_tensor * clip_graph::build_rope_2d(
     ggml_context * ctx0,
     ggml_tensor * cur,
@@ -827,7 +829,9 @@ ggml_tensor * clip_graph::build_rope_2d(
     const float freq_base,
     const bool interleave_freq
 ) {
-    const int64_t n_dim = cur->ne[0];
+    const int64_t n_dim  = cur->ne[0];
+    const int64_t n_head = cur->ne[1];
+    const int64_t n_pos  = cur->ne[2];
 
     // for example, if we have cur tensor of shape (n_dim=8, n_head, n_pos)
     // we will have a list of 4 inv_freq: 1e-0, 1e-1, 1e-2, 1e-3
@@ -841,30 +845,46 @@ ggml_tensor * clip_graph::build_rope_2d(
                                 ? std::pow(freq_base, (float)-2/n_dim)
                                 : 1.0;
 
-    // first half, dims [0, n_dim/2)
-    cur = ggml_rope_ext(
-        ctx0,
-        cur,
-        pos_a,      // positions
-        nullptr,    // freq factors
-        n_dim/2,    // n_dims
-        0, 0, freq_base,
-        1.0f, 0.0f, 1.0f, 0.0f, 0.0f
-    );
+    // first half
+    ggml_tensor * first;
+    {
+        first = ggml_view_3d(ctx0, cur,
+            n_dim/2, n_head, n_pos,
+            cur->nb[1],
+            cur->nb[2],
+            0);
+        first = ggml_rope_ext(
+            ctx0,
+            first,
+            pos_a,      // positions
+            nullptr,    // freq factors
+            n_dim/2,    // n_dims
+            0, 0, freq_base,
+            1.0f, 0.0f, 1.0f, 0.0f, 0.0f
+        );
+    }
 
-    // second half, dims [n_dim/2, n_dim)
-    cur = ggml_rope_ext(
-        ctx0,
-        cur,
-        pos_b,      // positions
-        nullptr,    // freq factors
-        n_dim/2,    // n_dims
-        0, 0, freq_base,
-        freq_scale_odd,
-        0.0f, 1.0f, 0.0f, 0.0f
-    );
-    cur = ggml_rope_set_offset(cur, n_dim/2);
+    // second half
+    ggml_tensor * second;
+    {
+        second = ggml_view_3d(ctx0, cur,
+            n_dim/2, n_head, n_pos,
+            cur->nb[1],
+            cur->nb[2],
+            n_dim/2 * ggml_element_size(cur));
+        second = ggml_rope_ext(
+            ctx0,
+            second,
+            pos_b,      // positions
+            nullptr,    // freq factors
+            n_dim/2,    // n_dims
+            0, 0, freq_base,
+            freq_scale_odd,
+            0.0f, 1.0f, 0.0f, 0.0f
+        );
+    }
 
+    cur = ggml_concat(ctx0, first, second, 0);
     return cur;
 }
 

tools/mtmd/models/gemma4v.cpp

@@ -44,31 +44,51 @@ ggml_cgraph * clip_graph_gemma4v::build() {
 
     // similar to build_rope_2d, but use neox ordering
     auto add_pos = [&](ggml_tensor * cur, const clip_layer &) {
-        const int64_t n_dim = cur->ne[0];
+        const int64_t n_dim  = cur->ne[0];
+        const int64_t n_head = cur->ne[1];
+        const int64_t n_pos  = cur->ne[2];
 
-        // first half, dims [0, n_dim/2)
-        cur = ggml_rope_ext(
-            ctx0,
-            cur,
-            pos_x,      // positions
-            nullptr,    // freq factors
-            n_dim/2,    // n_dims
-            GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta,
-            1.0f, 0.0f, 1.0f, 0.0f, 0.0f
-        );
+        // first half
+        ggml_tensor * first;
+        {
+            first = ggml_view_4d(ctx0, cur,
+                n_dim/2, n_head, n_pos, n_batch,
+                cur->nb[1],
+                cur->nb[2],
+                cur->nb[3],
+                0);
+            first = ggml_rope_ext(
+                ctx0,
+                first,
+                pos_x,      // positions
+                nullptr,    // freq factors
+                n_dim/2,    // n_dims
+                GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta,
+                1.0f, 0.0f, 1.0f, 0.0f, 0.0f
+            );
+        }
 
-        // second half, dims [n_dim/2, n_dim)
-        cur = ggml_rope_ext(
-            ctx0,
-            cur,
-            pos_y,      // positions
-            nullptr,    // freq factors
-            n_dim/2,    // n_dims
-            GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta,
-            1.0f, 0.0f, 1.0f, 0.0f, 0.0f
-        );
-        cur = ggml_rope_set_offset(cur, n_dim/2);
+        // second half
+        ggml_tensor * second;
+        {
+            second = ggml_view_4d(ctx0, cur,
+                n_dim/2, n_head, n_pos, n_batch,
+                cur->nb[1],
+                cur->nb[2],
+                cur->nb[3],
+                n_dim/2 * ggml_element_size(cur));
+            second = ggml_rope_ext(
+                ctx0,
+                second,
+                pos_y,      // positions
+                nullptr,    // freq factors
+                n_dim/2,    // n_dims
+                GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta,
+                1.0f, 0.0f, 1.0f, 0.0f, 0.0f
+            );
+        }
 
+        cur = ggml_concat(ctx0, first, second, 0);
         return cur;
     };
 

tools/mtmd/models/minimax-m3.cpp

@@ -2,22 +2,30 @@
 
 ggml_tensor * clip_graph_minimax_m3::apply_rope(
         ggml_tensor * x, ggml_tensor * pos_h, ggml_tensor * pos_w) {
-    const int dh  = (int) x->ne[0];
-    const int axd = 2 * ((2 * (dh / 2) / 3) / 2);
+    const int64_t Hn  = x->ne[1];
+    const int64_t P   = x->ne[2];
+    const size_t  es  = ggml_element_size(x);
+    const int     dh  = (int) x->ne[0];
+    const int     axd = 2 * ((2 * (dh / 2) / 3) / 2);
 
+    GGML_ASSERT(x->nb[0] == es);
     GGML_ASSERT(3 * axd <= dh);
 
     const float th  = hparams.rope_theta;
 
     // layout of x is [t, h, w, pad]
     // t is unrotated, h and w are rotated, pad is unrotated
-    x = ggml_rope_ext(ctx0, x, pos_h, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
-    x = ggml_rope_set_offset(x, axd);
+    // note: everything from n_dims onward untouched, so w and pad are rotated in one call.
+    auto sl = [&](int off, int n) {
+        return ggml_cont(ctx0, ggml_view_3d(ctx0, x, n, Hn, P, x->nb[1], x->nb[2], (size_t) off * es));
+    };
+    ggml_tensor * t = sl(0,       axd);
+    ggml_tensor * h = sl(axd,     axd);
+    ggml_tensor * w = sl(2 * axd, dh - 2 * axd); // w + pad
 
-    x = ggml_rope_ext(ctx0, x, pos_w, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
-    x = ggml_rope_set_offset(x, 2 * axd);
-
-    return x;
+    h = ggml_rope_ext(ctx0, h, pos_h, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+    w = ggml_rope_ext(ctx0, w, pos_w, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+    return ggml_concat(ctx0, ggml_concat(ctx0, t, h, 0), w, 0);
 }
 
 ggml_cgraph * clip_graph_minimax_m3::build() {

tools/server/server-context.cpp

@@ -1038,7 +1038,62 @@ private:
             }
         }
 
-        // note: the draft / MTP context is fitted together with the target model, see common_fit_extra_model
+        // optionally reserve VRAM for the draft / MTP context before fitting the target model
+        if (params_base.fit_params) {
+            if (has_spec) {
+                // MTP draft context lives on the target model, only context+compute are new
+                bool measure_model_bytes = has_draft;
+
+                common_params params_dft = common_base_params_to_speculative(params_base);
+
+                auto mparams_dft = common_model_params_to_llama(params_dft);
+                auto cparams_dft = common_context_params_to_llama(params_dft);
+                if (spec_mtp) {
+                    cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP;
+                }
+                cparams_dft.n_rs_seq = 0;
+
+                std::vector<ggml_backend_dev_t> devs;
+                uint32_t hp_ngl = 0;
+                uint32_t hp_nct = 0;
+                uint32_t hp_nex = 0;
+                try {
+                    auto dmd = common_get_device_memory_data(
+                        params_dft.model.path.c_str(), &mparams_dft, &cparams_dft,
+                        devs, hp_ngl, hp_nct, hp_nex, GGML_LOG_LEVEL_ERROR);
+
+                    GGML_ASSERT(!params_base.fit_params_target.empty());
+                    size_t total = 0;
+
+                    std::vector<ggml_backend_dev_t> tgt_devices = params.devices;
+
+                    if (tgt_devices.empty()) {
+                        for(size_t i = 0; i < ggml_backend_dev_count(); ++i) {
+                           tgt_devices.push_back(ggml_backend_dev_get(i));
+                        }
+                    }
+
+                    for (size_t j = 0; j < devs.size(); ++j) {
+                        const size_t bytes = (measure_model_bytes ? dmd[j].model : 0) + dmd[j].context + dmd[j].compute;
+                        total += bytes;
+                        for (size_t i = 0; i < tgt_devices.size(); i++) {
+                            if (tgt_devices[i] == devs[j]) {
+                                SRV_DBG("[spec] adding %.2f MiB to fit_params_target for device %s\n",
+                                        bytes / (1024.0 * 1024.0), ggml_backend_dev_name(devs[j]));
+                                params_base.fit_params_target[i] += bytes;
+                                break;
+                            }
+                        }
+                    }
+                    SRV_TRC("[spec] estimated memory usage of %s is %.2f MiB\n",
+                            has_draft ? "draft model" : "MTP context",
+                            total / (1024.0 * 1024.0));
+                } catch (const std::exception & e) {
+                    SRV_WRN("[spec] failed to measure %s memory: %s\n",
+                            has_draft ? "draft model" : "MTP context", e.what());
+                }
+            }
+        }
 
         // attach a progress callback
         {