From 2b94398ed7cf8d9bc5f6635737182db7d58a6dc7 Mon Sep 17 00:00:00 2001
From: Lumpiasty <arek.dzski@gmail.com>
Date: Thu, 23 Jul 2026 18:36:51 +0200
Subject: [PATCH] common : fractional -ncmoe for tensor-granularity expert
 placement

Extends --n-cpu-moe to accept a fractional layer count. The integer part
offloads whole layers as before; the fractional part offloads a subset of the
boundary layer expert tensors (gate, then up), keeping down_proj resident.
This realizes ATSInfer tensor-granularity static placement, giving sub-layer
control over expert VRAM residency. Placement-only, lossless.

Assisted-by: Claude
---
 common/arg.cpp  | 12 +++++++-----
 common/common.h | 22 ++++++++++++++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/common/arg.cpp b/common/arg.cpp
index 1f5b5a8ab..5a0c463ef 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -2575,15 +2575,17 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
     ).set_env("LLAMA_ARG_CPU_MOE"));
     add_opt(common_arg(
         {"-ncmoe", "--n-cpu-moe"}, "N",
-        "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU",
-        [](common_params & params, int value) {
-            if (value < 0) {
+        "keep the Mixture of Experts (MoE) weights of the first N layers in the CPU; "
+        "fractional N offloads part of the boundary layer at tensor granularity",
+        [](common_params & params, const std::string & value) {
+            const double n = std::stod(value);
+            if (n < 0) {
                 throw std::invalid_argument("invalid value");
             }
-            for (int i = 0; i < value; ++i) {
+            for (const std::string & re : llm_ffn_exps_cpu_block_regexes(n)) {
                 // keep strings alive and avoid leaking memory by storing them in a static vector
                 static std::list<std::string> buft_overrides;
-                buft_overrides.push_back(llm_ffn_exps_block_regex(i));
+                buft_overrides.push_back(re);
                 params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()});
             }
         }
diff --git a/common/common.h b/common/common.h
index bffc1767a..1b68019d3 100644
--- a/common/common.h
+++ b/common/common.h
@@ -1080,6 +1080,28 @@ inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() {
     return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() };
 }
 
+// ATSInfer-style tensor-granularity static placement of MoE expert weights.
+// Offloads the expert weights of the first floor(n) layers to the CPU, plus a
+// subset of the boundary layer's three expert tensors for the fractional part.
+// Expert tensors are dropped from the GPU in ascending performance-density
+// order (gate, then up), keeping the higher-value down_proj resident longest.
+inline std::vector<std::string> llm_ffn_exps_cpu_block_regexes(double n_cpu_moe) {
+    std::vector<std::string> regexes;
+    const int n_full = n_cpu_moe > 0 ? (int) n_cpu_moe : 0;
+    for (int i = 0; i < n_full; ++i) {
+        regexes.push_back(llm_ffn_exps_block_regex(i));
+    }
+    const int k = (int) ((n_cpu_moe - n_full) * 3.0 + 0.5);
+    if (k >= 3) {
+        regexes.push_back(llm_ffn_exps_block_regex(n_full));
+    } else if (k == 2) {
+        regexes.push_back(string_format("blk\\.%d\\.ffn_(gate|up)_(ch|)exps", n_full));
+    } else if (k == 1) {
+        regexes.push_back(string_format("blk\\.%d\\.ffn_gate_(ch|)exps", n_full));
+    }
+    return regexes;
+}
+
 //
 // training utils
 //