feat(wp): routing-boundary pre-pass (mark_routing_boundaries/is_routing_break)

Assisted-by: Codex

Author: kmbandy

src/weight-pager/wp-pager.cpp

@@ -1,6 +1,8 @@
 #include "wp-pager.h"
 
 #include "ggml-backend.h"
+#include "ggml.h"
+#include "../../ggml/src/ggml-impl.h"
 #include "llama-impl.h"  // LLAMA_LOG_*
 #include "wp-eval-cb.h"
 
@@ -512,6 +514,25 @@ bool WeightPager::batch_safe() const {
     return stats_.evictions == 0 && pool_.size_class_slots_enabled() && !catalog_.has_experts();
 }
 
+void WeightPager::mark_routing_boundaries(const struct ggml_cgraph * gf) {
+    if (gf == nullptr || gf->n_nodes <= 0) { return; }
+    const void * first = gf->nodes[0];
+    const void * last  = gf->nodes[gf->n_nodes - 1];
+    if (routing_sig_.n_nodes == gf->n_nodes && routing_sig_.first == first && routing_sig_.last == last) { return; }
+    routing_break_tensors_.clear();
+    for (int i = 0; i < gf->n_nodes; ++i) {
+        struct ggml_tensor * node = gf->nodes[i];
+        if (node->op != GGML_OP_MUL_MAT_ID) { continue; }
+        routing_break_tensors_.insert(node);
+        struct ggml_tensor * ids = node->src[2];
+        while (ids != nullptr && ids->view_src != nullptr) { ids = ids->view_src; }
+        if (ids != nullptr) { routing_break_tensors_.insert(ids); }
+    }
+    routing_sig_.n_nodes = gf->n_nodes;
+    routing_sig_.first = first;
+    routing_sig_.last  = last;
+}
+
 int WeightPager::loaded_pages() const {
     int n = 0;
     for (bool loaded : page_loaded_) {

src/weight-pager/wp-pager.h

@@ -31,6 +31,7 @@
 #include <memory>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 struct ggml_backend_buffer;
@@ -39,6 +40,9 @@ typedef struct ggml_backend_buffer * ggml_backend_buffer_t;
 struct ggml_backend_buffer_type;
 typedef struct ggml_backend_buffer_type * ggml_backend_buffer_type_t;
 
+struct ggml_cgraph;
+struct ggml_tensor;
+
 namespace wp {
 
 class WeightPager {
@@ -147,6 +151,10 @@ class WeightPager {
     bool   async_ensure_enabled()               const { return async_ensure_enabled_; }
     const Stats & stats() const;
     bool   batch_safe() const;
+    void   mark_routing_boundaries(const struct ggml_cgraph * gf);
+    bool   is_routing_break(const struct ggml_tensor * t) const {
+        return routing_break_tensors_.count(t) != 0;
+    }
     uint64_t sync_fallback_count() const { return stats_.sync_fallbacks; }
     int    loaded_pages() const;
     int    pending_prefetches() const { return prefetch_.pending(); }
@@ -298,6 +306,12 @@ class WeightPager {
     // Stored as slots, not pages, so unpin releases the same refcounts even
     // if a page later moves to another slot.
     std::unordered_map<const void *, std::vector<int>> graph_pin_slots_;
+    std::unordered_set<const struct ggml_tensor *> routing_break_tensors_;
+    struct {
+        int n_nodes = -1;
+        const void * first = nullptr;
+        const void * last = nullptr;
+    } routing_sig_;
 
     // Shared pinned staging buffer for page_in_sync_. Allocated once at
     // init, sized to max_page_size, reused across every sync page-in.

tests/test-weight-pager.cpp

@@ -12,6 +12,7 @@
 #include "weight-pager/wp-pool.h"
 
 #include "ggml-backend.h"
+#include "ggml.h"
 
 #include <cerrno>
 #include <cstdint>
@@ -1511,6 +1512,28 @@ static int test_is_uma_device_smoke() {
     return fails;
 }
 
+static int test_routing_boundary_prepass() {
+    int fails = 0;
+    struct ggml_init_params ip = { /*.mem_size=*/ 16*1024*1024, /*.mem_buffer=*/ nullptr, /*.no_alloc=*/ true };
+    struct ggml_context * ctx = ggml_init(ip);
+    struct ggml_tensor * ids_producer = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 8);
+    ggml_set_name(ids_producer, "ids_producer");
+    struct ggml_tensor * ids_view = ggml_view_1d(ctx, ids_producer, 8, 0);
+    struct ggml_tensor * as = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 4, 4, 2);
+    struct ggml_tensor * b  = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, 8);
+    struct ggml_tensor * mmid = ggml_mul_mat_id(ctx, as, b, ids_view);
+    ggml_set_name(mmid, "mmid");
+    struct ggml_cgraph * gf = ggml_new_graph(ctx);
+    ggml_build_forward_expand(gf, mmid);
+    wp::WeightPager pager;
+    pager.mark_routing_boundaries(gf);
+    if (!pager.is_routing_break(mmid))         { fprintf(stderr, "FAIL: mmid not marked\n"); fails++; }
+    if (!pager.is_routing_break(ids_producer)) { fprintf(stderr, "FAIL: ids producer (view root) not marked\n"); fails++; }
+    if (pager.is_routing_break(b))             { fprintf(stderr, "FAIL: unrelated tensor marked\n"); fails++; }
+    ggml_free(ctx);
+    return fails;
+}
+
 // ---------------------------------------------------------------------------
 // main
 // ---------------------------------------------------------------------------
@@ -1556,6 +1579,7 @@ int main() {
         { "pool_hot_threshold_protects",         test_pool_hot_threshold_protects_in_eviction },
         { "pool_hot_fallback_when_all_hot",      test_pool_hot_fallback_when_all_hot      },
         { "pool_default_threshold_zero_lru",     test_pool_default_threshold_zero_is_pure_lru },
+        { "routing_boundary_prepass",            test_routing_boundary_prepass            },
     };
 
     for (const auto & t : tests) {