feat(wp): WP_PAGED_BATCH flag + batch_safe pinnability-governed under it

Assisted-by: Codex

Author: kmbandy

src/weight-pager/wp-eval-cb.cpp

@@ -27,6 +27,14 @@ extern "C++" void *                ggml_cuda_get_wp_compute_stream();
 
 namespace wp {
 
+bool wp_paged_batch_enabled() {
+    static const bool enabled = []() {
+        const char * v = std::getenv("WP_PAGED_BATCH");
+        return v != nullptr && std::strcmp(v, "1") == 0;
+    }();
+    return enabled;
+}
+
 namespace {
 // Diagnostic counters. Logged only when WP_EVAL_DEBUG=1 is set in the
 // environment. First few ops get verbose output; afterwards we suppress

src/weight-pager/wp-eval-cb.h

@@ -27,6 +27,8 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data);
 // Drain eval-callback state associated with pager before the pager tears down.
 void weight_pager_eval_cb_reset(WeightPager * pager);
 
+bool wp_paged_batch_enabled();
+
 // Diagnostic (WP_PROFILE_EVAL=1, default off): print the total host-side wall
 // time spent inside weight_pager_eval_cb over the run, split into the Step-2
 // ensure/patch phase vs page-resolution+other. Lets us see whether paged-decode

src/weight-pager/wp-pager.cpp

@@ -511,6 +511,11 @@ bool WeightPager::batch_safe() const {
     // routed-expert TLS / pinned-slot lifetime isn't isolated -> near-null expert-pointer
     // GPU fault. Until the MoE-batching redesign lands (routing op must break the range
     // *before* it), keep MoE on the per-op sync path.
+    if (wp_paged_batch_enabled()) {
+        // Under WP_PAGED_BATCH, batching is governed by live pinnability and
+        // routing-boundary breaks in wp-eval-cb.cpp, not a static eviction count.
+        return pool_.size_class_slots_enabled();
+    }
     return stats_.evictions == 0 && pool_.size_class_slots_enabled() && !catalog_.has_experts();
 }
 

tests/test-weight-pager.cpp

@@ -6,6 +6,7 @@
 // runtime — they no-op compile-out under non-HIP builds.
 
 #include "weight-pager/wp-page-catalog.h"
+#include "weight-pager/wp-eval-cb.h"
 #include "weight-pager/wp-file-io.h"
 #include "weight-pager/wp-host-tier.h"
 #include "weight-pager/wp-pager.h"   // compute_advise_ranges / AdviseRange
@@ -1534,6 +1535,14 @@ static int test_routing_boundary_prepass() {
     return fails;
 }
 
+static int test_wp_paged_batch_flag_default_off() {
+    int fails = 0;
+    ScopedEnv guard("WP_PAGED_BATCH");
+    unsetenv("WP_PAGED_BATCH");
+    if (wp::wp_paged_batch_enabled()) { fprintf(stderr, "FAIL: WP_PAGED_BATCH must default OFF\n"); fails++; }
+    return fails;
+}
+
 // ---------------------------------------------------------------------------
 // main
 // ---------------------------------------------------------------------------
@@ -1580,6 +1589,7 @@ int main() {
         { "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            },
+        { "wp_paged_batch_flag_default_off",     test_wp_paged_batch_flag_default_off     },
     };
 
     for (const auto & t : tests) {