feat(wp): add resident-device routing for paging

kmbandy · kmbandy · commit d8196031a021 · 2026-07-07T20:54:07.000-04:00
Assisted-by: Codex
diff --git a/common/arg.cpp b/common/arg.cpp
@@ -2731,6 +2731,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.weight_paging_prefetch = true;
         }
     ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_PERPLEXITY}).set_env("LLAMA_ARG_WEIGHT_PAGING_PREFETCH"));
+    add_opt(common_arg(
+        {"--weight-paging-resident-device"}, "<dev|auto>",
+        "device for resident dense weights under WP_RESIDENT_DENSE=1 (default: auto)",
+        [](common_params & params, const std::string & value) {
+            if (value.empty()) {
+                throw std::invalid_argument("invalid value");
+            }
+            params.weight_paging_resident_device = value;
+        }
+    ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI, LLAMA_EXAMPLE_PERPLEXITY}).set_env("LLAMA_ARG_WEIGHT_PAGING_RESIDENT_DEVICE"));
     add_opt(common_arg(
         {"-sm", "--split-mode"}, "{none,layer,row,tensor}",
         "how to split the model across multiple GPUs, one of:\n"
diff --git a/common/common.cpp b/common/common.cpp
@@ -1570,9 +1570,10 @@ struct llama_model_params common_model_params_to_llama(common_params & params) {
     mparams.no_alloc                    = params.no_alloc;
 
     // Weight paging parameters
-    mparams.weight_paging_enabled  = params.weight_paging_enabled;
-    mparams.weight_paging_slots    = params.weight_paging_slots;
-    mparams.weight_paging_prefetch = params.weight_paging_prefetch;
+    mparams.weight_paging_enabled         = params.weight_paging_enabled;
+    mparams.weight_paging_slots           = params.weight_paging_slots;
+    mparams.weight_paging_prefetch        = params.weight_paging_prefetch;
+    mparams.weight_paging_resident_device = params.weight_paging_resident_device.c_str();
 
     return mparams;
 }
diff --git a/common/common.h b/common/common.h
@@ -587,6 +587,7 @@ struct common_params {
     bool    weight_paging_enabled = false;  // enable NVMe→VRAM demand paging for model weights
     int32_t weight_paging_slots   = -1;     // number of VRAM slots for weight paging (-1 = auto)
     bool    weight_paging_prefetch = false;  // enable async prefetch of next layer
+    std::string weight_paging_resident_device = "auto"; // dense-resident device name or auto
 
     bool single_turn       = false; // single turn chat conversation
 
diff --git a/include/llama.h b/include/llama.h
@@ -337,6 +337,7 @@ extern "C" {
         bool weight_paging_enabled;    // enable NVMe→VRAM demand paging for model weights
         int32_t weight_paging_slots;   // number of VRAM slots for weight paging (-1 = auto)
         bool weight_paging_prefetch;   // enable async prefetch of next layer
+        const char * weight_paging_resident_device; // dense-resident device name or "auto"
     };
 
     struct llama_sampler_seq_config {
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
@@ -1273,6 +1273,54 @@ static bool wp_is_routed_expert(const char * name) {
            std::strstr(p, "ffn_down_exps.") != nullptr;
 }
 
+static int wp_find_device_index_by_name(const std::vector<llama_device> & devices, const char * name) {
+    if (name == nullptr || name[0] == '\0') {
+        return -1;
+    }
+    for (size_t i = 0; i < devices.size(); ++i) {
+        if (std::strcmp(ggml_backend_dev_name(devices[i].dev), name) == 0) {
+            return (int) i;
+        }
+    }
+    return -1;
+}
+
+static int wp_select_paging_device_index(const llama_model_params & params, const std::vector<llama_device> & devices) {
+    if (devices.empty()) {
+        return -1;
+    }
+    if (params.main_gpu >= 0 && params.main_gpu < (int) devices.size()) {
+        return params.main_gpu;
+    }
+    return 0;
+}
+
+static int wp_select_resident_device_index(const llama_model_params & params,
+                                           const std::vector<llama_device> & devices,
+                                           int paging_idx) {
+    if (devices.empty()) {
+        return -1;
+    }
+    const char * requested = params.weight_paging_resident_device
+        ? params.weight_paging_resident_device
+        : "auto";
+    if (std::strcmp(requested, "auto") != 0) {
+        int idx = wp_find_device_index_by_name(devices, requested);
+        if (idx >= 0) {
+            return idx;
+        }
+        LLAMA_LOG_WARN("%s: resident device '%s' not found; falling back to paging device\n",
+                       __func__, requested);
+        return paging_idx;
+    }
+    for (size_t i = 0; i < devices.size(); ++i) {
+        if ((int) i != paging_idx) {
+            return (int) i;
+        }
+    }
+    return paging_idx;
+}
+
 bool llama_model_base::load_tensors(llama_model_loader & ml) {
     const auto & split_mode   = params.split_mode;
     const auto & use_mlock    = params.use_mlock;
@@ -1308,6 +1356,43 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
         pimpl->gpu_buft_list.emplace(dev.dev, std::move(buft_list));
     }
 
+    std::string wp_expert_override_pattern;
+    std::string wp_dense_override_pattern;
+    std::vector<llama_model_tensor_buft_override> wp_tensor_buft_overrides;
+    ggml_backend_buffer_type_t wp_paging_buft = nullptr;
+    ggml_backend_buffer_type_t wp_resident_buft = nullptr;
+    bool wp_device_router_enabled = false;
+
+    if (params.weight_paging_enabled && wp_resident_dense_enabled() && !devices.empty()) {
+        const int paging_idx = wp_select_paging_device_index(params, devices);
+        const int resident_idx = wp_select_resident_device_index(params, devices, paging_idx);
+        if (paging_idx >= 0 && resident_idx >= 0) {
+            wp_paging_buft = ggml_backend_dev_buffer_type(devices[paging_idx].dev);
+            wp_resident_buft = ggml_backend_dev_buffer_type(devices[resident_idx].dev);
+        }
+        if (wp_paging_buft != nullptr && wp_resident_buft != nullptr) {
+            wp_expert_override_pattern = "ffn_(up|gate|down)_exps\\.";
+            wp_dense_override_pattern = ".*";
+            wp_tensor_buft_overrides.push_back({ wp_expert_override_pattern.c_str(), wp_paging_buft });
+            wp_tensor_buft_overrides.push_back({ wp_dense_override_pattern.c_str(),  wp_resident_buft });
+            if (params.tensor_buft_overrides != nullptr) {
+                for (const auto * o = params.tensor_buft_overrides; o->pattern != nullptr; ++o) {
+                    wp_tensor_buft_overrides.push_back(*o);
+                }
+            }
+            wp_tensor_buft_overrides.push_back({ nullptr, nullptr });
+            ml.tensor_buft_overrides = wp_tensor_buft_overrides.data();
+            wp_device_router_enabled = true;
+            LLAMA_LOG_INFO("%s: WP_RESIDENT_DENSE router: paging=%s (%s), resident=%s (%s)\n",
+                           __func__,
+                           ggml_backend_dev_name(devices[paging_idx].dev), ggml_backend_buft_name(wp_paging_buft),
+                           ggml_backend_dev_name(devices[resident_idx].dev), ggml_backend_buft_name(wp_resident_buft));
+        } else {
+            LLAMA_LOG_WARN("%s: WP_RESIDENT_DENSE router disabled: could not resolve paging/resident bufts\n",
+                           __func__);
+        }
+    }
+
     ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
     if (cpu_dev == nullptr) {
         throw std::runtime_error(format("%s: no CPU backend found", __func__));
@@ -1760,10 +1845,21 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
                 }
             }
 
-            bool ctx_has_weights = false;
+            bool ctx_has_paged_weights = false;
             for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
                 const auto * weight = ml.get_weight(ggml_get_name(t));
                 if (!weight) { continue; }
+                if (wp_device_router_enabled && ctx_buft != nullptr) {
+                    const char * n = ggml_get_name(t);
+                    ggml_backend_buffer_type_t expected_buft =
+                        wp_is_routed_expert(n) ? wp_paging_buft : wp_resident_buft;
+                    if (expected_buft != nullptr && ctx_buft != expected_buft) {
+                        throw std::runtime_error(format(
+                            "weight-paging: tensor %s routed to %s, expected %s",
+                            n, ggml_backend_buft_name(ctx_buft),
+                            ggml_backend_buft_name(expected_buft)));
+                    }
+                }
                 // MAD-88: skip huge non-block tensors from paging. token_embd
                 // and output are needed every forward pass (vocab embed +
                 // logits), so paging them just adds SSD churn for no win.
@@ -1835,13 +1931,13 @@ bool llama_model_base::load_tensors(llama_model_loader & ml) {
                     if (weight_pager) {
                         weight_pager->weight_tensor_ptrs.push_back(t);
                     }
+                    ctx_has_paged_weights = true;
                 }
-                ctx_has_weights = true;
             }
 
-            // Record the buft once per ctx that owns weight tensors.
+            // Record the buft once per ctx that owns paged weight tensors.
             // De-dupe so the multi-device guard reads a clean set.
-            if (ctx_has_weights && ctx_buft != nullptr && weight_pager) {
+            if (ctx_has_paged_weights && ctx_buft != nullptr && weight_pager) {
                 bool seen = false;
                 for (auto * existing : weight_pager->weight_bufts) {
                     if (existing == ctx_buft) { seen = true; break; }
@@ -2778,6 +2874,7 @@ llama_model_params llama_model_default_params() {
         /*.weight_paging_enabled       =*/ false,
         /*.weight_paging_slots         =*/ -1,
         /*.weight_paging_prefetch      =*/ false,
+        /*.weight_paging_resident_device =*/ "auto",
     };
 
     return result;
diff --git a/src/llama.cpp b/src/llama.cpp
@@ -174,7 +174,14 @@ static bool init_weight_pager(llama_model & model, llama_model_loader & ml, cons
     // 3. Build the new pager: catalog + init.
     model.wp_pager = std::make_unique<wp::WeightPager>();
     for (const auto & info : ml.weight_page_infos) {
-        model.wp_pager->add_page(info.name, info.file_idx, info.offset, info.size, info.n_experts);
+        ggml_backend_buffer_type_t page_buft = buft;
+        for (ggml_tensor * t : model.weight_pager->weight_tensor_ptrs) {
+            if (t != nullptr && info.name == ggml_get_name(t) && t->buffer != nullptr) {
+                page_buft = ggml_backend_buffer_get_type(t->buffer);
+                break;
+            }
+        }
+        model.wp_pager->add_page(info.name, info.file_idx, info.offset, info.size, info.n_experts, page_buft);
     }
     LLAMA_LOG_INFO("%s: catalog populated: %d page entries (source had %zu)\n",
                    __func__, model.wp_pager->n_pages(), ml.weight_page_infos.size());
diff --git a/src/weight-pager/wp-eval-cb.cpp b/src/weight-pager/wp-eval-cb.cpp
@@ -784,7 +784,7 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
                                     // src0->buffer = nullptr for paged tensors, so without
                                     // this patch we NULL-deref before reaching the
                                     // routing-aware dispatcher gate.
-                                    ggml_backend_buffer_t pool_buf = pager->pool_buf();
+                                    ggml_backend_buffer_t pool_buf = pager->pool_buf(weight_page);
                                     if (t->src[0]->buffer == nullptr && pool_buf != nullptr) {
                                         t->src[0]->buffer = pool_buf;
                                     }
@@ -1093,7 +1093,6 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
 
     // Step 2: page each one in (waiting on prefetch if in flight, sync
     // fallback otherwise) and patch the matching src tensors.
-    ggml_backend_buffer_t pool_buf = pager->pool_buf();
     int  patches_this_op = 0;
     int  views_this_op   = 0;
 
@@ -1123,6 +1122,7 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
 #endif
 
         const std::string & page_name = pager->page_meta(page_idx).tensor_name;
+        ggml_backend_buffer_t pool_buf = pager->pool_buf(page_idx);
 
         // Patch every src whose direct name OR view_src's name matches
         // this page.
diff --git a/src/weight-pager/wp-pager.cpp b/src/weight-pager/wp-pager.cpp
@@ -150,16 +150,28 @@ int WeightPager::register_pinned(const std::string & name, void * device_ptr, si
 }
 
 int WeightPager::add_page(const std::string & name, uint16_t file_idx,
-                          uint64_t file_offset, size_t size, int n_experts) {
+                          uint64_t file_offset, size_t size, int n_experts,
+                          ggml_backend_buffer_type_t buft) {
+    const int n_before = catalog_.size();
+    int first = -1;
     // Non-MoE / per-expert tensor: add as-is.
     if (n_experts <= 1) {
-        return catalog_.add(name, file_idx, file_offset, size);
+        first = catalog_.add(name, file_idx, file_offset, size);
+    } else {
+        // Consolidated MoE tensor: register N sub-pages, one per expert.
+        // Returns the index of the FIRST sub-page (subsequent experts are at
+        // sequential indices). Per-expert size is the consolidated size
+        // divided by n_experts; per-expert offset is base_offset + e * size_e.
+        first = catalog_.add_consolidated_experts(name, file_idx, file_offset, size, n_experts);
+    }
+    const int n_after = catalog_.size();
+    if ((int) page_buft_.size() < n_after) {
+        page_buft_.resize((size_t) n_after, nullptr);
     }
-    // Consolidated MoE tensor: register N sub-pages, one per expert.
-    // Returns the index of the FIRST sub-page (subsequent experts are at
-    // sequential indices). Per-expert size is the consolidated size
-    // divided by n_experts; per-expert offset is base_offset + e * size_e.
-    return catalog_.add_consolidated_experts(name, file_idx, file_offset, size, n_experts);
+    for (int i = n_before; i < n_after; ++i) {
+        page_buft_[(size_t) i] = buft;
+    }
+    return first;
 }
 
 bool WeightPager::init(const Config &             cfg,
@@ -176,20 +188,11 @@ bool WeightPager::init(const Config &             cfg,
         return false;
     }
 
-    // Multi-GPU guard (B-P7). Phase 1 is single-device by explicit design.
-    // This is the *only* defence — the rest of the pager assumes a single
-    // pool on a single device.
     if (devices_used.size() > 1) {
-        LLAMA_LOG_ERROR(
-            "wp::WeightPager::init: multi-device configurations are not supported by the "
-            "weight pager (got %zu devices). Use --device with a single ROCm/CUDA index "
-            "for paging, or run without --weight-paging.\n",
-            devices_used.size());
-        // Caller passed fds; close them so they don't leak.
-        for (int fd : fds) {
-            if (fd >= 0) close(fd);
-        }
-        return false;
+        LLAMA_LOG_WARN(
+            "wp::WeightPager::init: model uses %zu devices; pager pool remains on configured "
+            "paging device %d and per-page buffer selection is enabled\n",
+            devices_used.size(), device_idx);
     }
     if (!devices_used.empty() && devices_used.front() != device_idx) {
         LLAMA_LOG_WARN("wp::WeightPager::init: device mismatch (used=%d, configured=%d)\n",
@@ -379,6 +382,7 @@ void WeightPager::shutdown() {
         }
     }
     page_async_event_.clear();
+    page_buft_.clear();
 
     // Tear down in reverse construction order.
     prefetch_.shutdown();
@@ -419,6 +423,11 @@ void WeightPager::shutdown() {
     initialized_ = false;
 }
 
+ggml_backend_buffer_t WeightPager::pool_buf(int page_idx) const {
+    (void) page_idx;
+    return pool_.vram_buf();
+}
+
 void WeightPager::restore_disable_graphs_env_() {
     if (!env_disable_graphs_forced_) {
         return;
diff --git a/src/weight-pager/wp-pager.h b/src/weight-pager/wp-pager.h
@@ -104,7 +104,8 @@ class WeightPager {
                  uint16_t            file_idx,
                  uint64_t            file_offset,
                  size_t              size,
-                 int                 n_experts = 1);
+                 int                 n_experts = 1,
+                 ggml_backend_buffer_type_t buft = nullptr);
 
     // Initialise the pool, transport, file-io layer, and prefetch scheduler.
     //
@@ -248,6 +249,7 @@ class WeightPager {
     // Backing buffer for the pool — used by the eval-cb adapter when
     // patching tensor->buffer (B-P4 requires a valid ggml backend buffer).
     ggml_backend_buffer_t pool_buf() const { return pool_.vram_buf(); }
+    ggml_backend_buffer_t pool_buf(int page_idx) const;
 
     // Slot-and-page metadata (read-only public view).
     const PageMeta & page_meta(int page_idx) const { return catalog_.at(page_idx); }
@@ -317,6 +319,7 @@ class WeightPager {
     std::vector<bool> cross_layer_prefetch_candidate_;
     std::vector<std::chrono::steady_clock::time_point> prefetch_started_at_;
     std::vector<int> page_async_event_;
+    std::vector<ggml_backend_buffer_type_t> page_buft_;
     // Reverse map: slot_idx -> page_idx (or -1 if free). Used by the
     // eviction callback to clear page_to_slot_ / page_loaded_ correctly.
     std::vector<int> slot_to_page_;
diff --git a/tests/test-weight-pager.cpp b/tests/test-weight-pager.cpp
@@ -1017,6 +1017,41 @@ static int test_pool_alloc_skips_pinned_in_eviction() {
     return fails;
 }
 
+static int test_pool_allocator_two_pools_independent() {
+    int fails = 0;
+    ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type();
+    EXPECT(buft != nullptr, "cpu buffer_type available");
+    if (!buft) return fails;
+
+    wp::PoolAllocator pool_a;
+    wp::PoolAllocator pool_b;
+    EXPECT(pool_a.init(buft, /*n_slots=*/2, /*slot_size=*/128), "pool A init");
+    EXPECT(pool_b.init(buft, /*n_slots=*/3, /*slot_size=*/256), "pool B init");
+
+    int evict_a = -1;
+    int evict_b = -1;
+    pool_a.set_eviction_callback([&](int slot) { evict_a = slot; });
+    pool_b.set_eviction_callback([&](int slot) { evict_b = slot; });
+
+    EXPECT_EQ_INT(pool_a.alloc_slot(), 0, "pool A first slot");
+    EXPECT_EQ_INT(pool_a.alloc_slot(), 1, "pool A second slot");
+    EXPECT_EQ_INT(pool_b.alloc_slot(), 0, "pool B first slot");
+    EXPECT_EQ_INT(pool_b.alloc_slot(), 1, "pool B second slot");
+    EXPECT_EQ_INT(pool_b.alloc_slot(), 2, "pool B third slot");
+
+    pool_a.pin_slot(0);
+    EXPECT_EQ_INT(pool_a.alloc_slot(), 1, "pool A evicts only its unpinned slot");
+    EXPECT_EQ_INT(evict_a, 1, "pool A eviction callback");
+    EXPECT_EQ_INT(evict_b, -1, "pool B not evicted by pool A pressure");
+    EXPECT(pool_a.is_pinned(0), "pool A pin remains local");
+    EXPECT(!pool_b.is_pinned(0), "pool B pin state remains independent");
+
+    EXPECT_EQ_INT(pool_b.alloc_slot(), 0, "pool B evicts its own LRU slot");
+    EXPECT_EQ_INT(evict_b, 0, "pool B eviction callback");
+    pool_a.unpin_slot(0);
+    return fails;
+}
+
 static int test_pool_alloc_returns_neg1_when_all_pinned() {
     int fails = 0;
     ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type();
@@ -1676,6 +1711,7 @@ int main() {
         { "read_mem_available_bytes", test_read_mem_available_bytes },
         { "is_uma_device_smoke",      test_is_uma_device_smoke      },
         { "pool_allocator",     test_pool_allocator     },
+        { "pool_allocator_two_pools_independent", test_pool_allocator_two_pools_independent },
         { "pool_size_class_packs_small_pages", test_pool_size_class_packs_small_pages },
         { "pool_size_class_pin_skip",          test_pool_size_class_pin_skip          },
         { "pool_pin_basic",                       test_pool_pin_basic                       },
