Skip to content

Commit 42c1825

Browse files
kmbandyclaude
kmbandy
and
claude
committed
feat(MAD-186): port hipfire weight-pager intelligence (7 stories)
Implements seven hipfire-style pager features on top of the WeightPager runtime, all gated by env vars so default behaviour is byte-identical to the pre-MAD-186 path: MAD-231 slot pin/refcount wp-pool, wp-pager Protect in-flight slots from eviction during ops. Refcounted via PoolAllocator pin_slot/unpin_slot. Conservative window spanning ensure -> next-eval-cb covers compute-stream reads. Diagnostic: lru_walk_pinned_skips. MAD-232 posix_fadvise K-layer lookahead wp-file-io, wp-pager One advise call per upcoming paged tensor in [block+1, block+k]. WP_FADVISE_LOOKAHEAD env-gated. Walks the catalog via compute_advise_ranges (pure, unit-testable). MAD-233 cross-layer N+K MoE prefetch wp-eval-cb Sister-prefetch the current layer's expert set, plus a cross-layer block looking K layers ahead. WP_NEXT_LAYER_PREFETCH_K env-gated. Composes with MAD-235 batched submit. MAD-234 UMA/APU MemAvailable pre-flight wp-pool Refuse-with-actionable-error policy. is_uma_device combines hipDeviceAttributeIntegrated with a gfx-string prefix fallback for gfx115x / gfx1103 / gfx1102 / gfx1100 / gfx90c / gfx940. Pre-flight clamp: pool > MemAvailable - 2 GiB -> hard error. MAD-235 batched io_uring submit wp-file-io, wp-prefetch Submit N SQEs and one io_uring_submit. All-or-nothing semantics: any sub-request failure rolls back every reservation. WP_BATCH_PREFETCH env-gated; falls back to per-page on false. MAD-236 always-resident pinning wp-page-catalog, wp-pager register_pinned + add_pinned: tensors whose bytes already live in caller-owned VRAM (e.g. token_embd) are first-class catalog entries. ensure() short-circuits to resident_ptr; no slot allocation, no disk read. Telemetry counters for pinned bytes + page count. MAD-237 popularity counter + hot-set wp-pool Two-pass LRU walk under WP_HOT_HIT_THRESHOLD > 0: A) evict LRU among unpinned AND hit_count <= threshold B) if A empty, fall through to unpinned LRU. Periodic decay (halve hit counts every 1024 evictions) keeps counters bounded and lets cold-becomes-hot demote over time. Default threshold=0 -> pure LRU, byte-identical to pre-237. Tests: 19 new unit tests covering algorithm correctness on the CPU buffer-type (PoolAllocator pin/eviction-skip/hot-protection, FileIOLayer advise + batched submit + rollback, PageCatalog pinned mixing, UMA prefix table). Real-HIP integration validation deferred to a follow-up R9700 run. Bug discovered during test validation: compute_advise_ranges filtered consolidated parents by `m.size == 0` based on a stale assumption - add_consolidated_experts stores total_size on the parent (it's load- bearing for other consumers' name lookups). Switched filter to the authoritative `is_consolidated` flag; defence-in-depth size==0 guard kept for malformed catalog entries. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
1 parent c4ebc52 commit 42c1825

12 files changed

Lines changed: 1942 additions & 17 deletions

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

Lines changed: 185 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ extern "C++" void * ggml_cuda_get_wp_compute_stream();
1616

1717
#include <cstdlib> // getenv
1818
#include <cstring>
19+
#include <limits> // numeric_limits — MAD-232 advise sentinel
1920
#include <unordered_map>
2021
#include <unordered_set>
2122
#include <vector>
@@ -66,6 +67,22 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
6667
ggml_cuda_take_routed_expert_ptrs();
6768
#endif
6869

70+
// MAD-231: drain pins from the PREVIOUS op's pages now that the GPU
71+
// has had a full eval-cb-cycle of latency to finish reading them.
72+
// Conservative: this window is a superset of the actual GPU read
73+
// (the previous op was dispatched on the compute stream and any
74+
// subsequent prefetch H2D is enqueued on the same stream → ordered),
75+
// but the pin set is small (handful of pages per op) so the over-
76+
// protection is free.
77+
//
78+
// Single-threaded: ggml's scheduler dispatches eval_cb on one thread,
79+
// so no synchronization is needed for s_pinned_pages_prev_op.
80+
static std::vector<int> s_pinned_pages_prev_op;
81+
for (int prev_page : s_pinned_pages_prev_op) {
82+
pager->unpin_page(prev_page);
83+
}
84+
s_pinned_pages_prev_op.clear();
85+
6986
// Diagnostic: detect MUL_MAT_ID ops and check whether their weight
7087
// source is a consolidated MoE parent. This is the entry point for
7188
// routing-aware paging (MAD-88 Phase 2 part 2). Currently informational
@@ -224,9 +241,35 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
224241
// already reserved + in-flight, and
225242
// wait_for completion instead of doing
226243
// a fresh sync pread.
227-
for (int e : active) {
228-
const int sub_page_idx = weight_page + 1 + e;
229-
pager->prefetch_page(sub_page_idx);
244+
//
245+
// MAD-235: prefer the atomic batch path
246+
// (one io_uring_submit syscall for the
247+
// whole expert set). WP_BATCH_PREFETCH=0
248+
// reverts to per-expert loop for A/B
249+
// measurement / regression rollback.
250+
static int s_batch_prefetch_env = -1;
251+
if (s_batch_prefetch_env < 0) {
252+
const char * env = std::getenv("WP_BATCH_PREFETCH");
253+
s_batch_prefetch_env = (env != nullptr && env[0] == '0') ? 0 : 1;
254+
}
255+
bool batch_ok = false;
256+
if (s_batch_prefetch_env != 0) {
257+
std::vector<int> active_pages;
258+
active_pages.reserve(active.size());
259+
for (int e : active) {
260+
active_pages.push_back(weight_page + 1 + e);
261+
}
262+
batch_ok = pager->prefetch_pages_batch(active_pages);
263+
}
264+
if (!batch_ok) {
265+
// Either batch was disabled OR scheduler
266+
// refused (queue full / capacity tight).
267+
// Fall back to per-expert prefetch — best-
268+
// effort, ignores individual failures.
269+
for (int e : active) {
270+
const int sub_page_idx = weight_page + 1 + e;
271+
pager->prefetch_page(sub_page_idx);
272+
}
230273
}
231274
// Drive the prefetch state machine forward
232275
// so submitted reads get out the door
@@ -250,6 +293,11 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
250293
first_active_slot = slot;
251294
}
252295
++n_ensures;
296+
// MAD-231: pin the slot so a later prefetch
297+
// alloc_slot in this same eval_cb can't evict
298+
// it. Unpinned in the NEXT eval_cb (above).
299+
pager->pin_page(sub_page_idx);
300+
s_pinned_pages_prev_op.push_back(sub_page_idx);
253301
}
254302
}
255303
// Safety: fill INACTIVE expert slots with a non-null
@@ -354,6 +402,91 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
354402
pager->prefetch_page(sister_sub);
355403
}
356404
}
405+
406+
// MAD-233 — cross-layer N+K MoE expert prefetch.
407+
//
408+
// Project the CURRENT layer's active expert set forward to
409+
// layers [block+1, block+K]. We can't know the true active
410+
// set for future layers (the router decides per layer), but
411+
// empirical Qwen3-MoE has 40-50% expert reuse across
412+
// consecutive tokens — locality enough that prefetching
413+
// the same indices means most ensure()s for the next layer
414+
// are cache hits, and the rest at least overlap NVMe I/O
415+
// with this layer's compute (5 ms per MoE layer).
416+
//
417+
// Safety: MAD-231 slot pinning protects this layer's
418+
// in-flight slots from eviction by the prefetch alloc_slot.
419+
// Without that, the prefetch could evict a slot the current
420+
// op is still reading. With MAD-231 the worst case is
421+
// "prefetch couldn't find an unpinned slot, batch refused,
422+
// we don't get the win" — never corruption.
423+
//
424+
// WP_NEXT_LAYER_PREFETCH_K env tunes lookahead (default 1,
425+
// 0 disables). Cache the parent list per source-parent for
426+
// O(1) reuse across all 3 MUL_MAT_IDs of a layer.
427+
static int s_next_k = -1;
428+
if (s_next_k < 0) {
429+
const char * env = std::getenv("WP_NEXT_LAYER_PREFETCH_K");
430+
s_next_k = env ? std::atoi(env) : 1;
431+
if (s_next_k < 0) s_next_k = 0;
432+
LLAMA_LOG_INFO("[wp::eval_cb] WP_NEXT_LAYER_PREFETCH_K=%d "
433+
"(0=disabled, requires MAD-231 pinning)\n",
434+
s_next_k);
435+
}
436+
if (s_next_k > 0) {
437+
// Cache: weight_page -> vector of consolidated parent
438+
// indices for blocks [my_block+1, my_block+s_next_k].
439+
static std::unordered_map<int, std::vector<int>> s_next_layer_parents_cache;
440+
auto next_it = s_next_layer_parents_cache.find(weight_page);
441+
if (next_it == s_next_layer_parents_cache.end()) {
442+
std::vector<int> next_parents;
443+
const int my_block = meta.block_idx;
444+
for (int i = 0; i < pager->n_pages(); ++i) {
445+
const auto & p = pager->page_meta(i);
446+
if (!p.is_consolidated) continue;
447+
if (p.block_idx <= my_block) continue;
448+
if (p.block_idx > my_block + s_next_k) continue;
449+
next_parents.push_back(i);
450+
}
451+
next_it = s_next_layer_parents_cache.emplace(
452+
weight_page, std::move(next_parents)).first;
453+
}
454+
455+
// Build a single batched prefetch covering every
456+
// (future_parent, active_expert) pair. The batch path
457+
// (MAD-235) collapses to one io_uring_submit syscall
458+
// and either queues the whole set atomically or refuses
459+
// — refusal falls back to per-page prefetch which is
460+
// best-effort.
461+
if (!next_it->second.empty()) {
462+
std::vector<int> future_pages;
463+
future_pages.reserve(next_it->second.size() * active.size());
464+
for (int future_parent : next_it->second) {
465+
for (int e : active) {
466+
if (e < 0 || e >= n_subs) continue;
467+
future_pages.push_back(future_parent + 1 + e);
468+
}
469+
}
470+
if (!future_pages.empty()) {
471+
const bool batch_ok = pager->prefetch_pages_batch(future_pages);
472+
if (!batch_ok) {
473+
// Per-page fallback (best-effort, ignores
474+
// individual failures — eviction will sort
475+
// it out and ensure() on the next layer
476+
// falls back to sync if nothing landed).
477+
for (int fp : future_pages) {
478+
pager->prefetch_page(fp);
479+
}
480+
}
481+
if (g_debug.mmid_consolidated <= 4) {
482+
LLAMA_LOG_INFO("[wp::eval_cb] cross-layer prefetch: %zu pages "
483+
"(parents=%zu, K=%d, batch=%d)\n",
484+
future_pages.size(), next_it->second.size(),
485+
s_next_k, (int) batch_ok);
486+
}
487+
}
488+
}
489+
}
357490
}
358491
}
359492
}
@@ -476,6 +609,49 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
476609
}
477610
++g_debug.ops_with_pages;
478611

612+
// MAD-232: posix_fadvise(WILLNEED) for the next K layers' paged tensors.
613+
// Warms NVMe→page-cache while THIS layer's compute runs, so by the time
614+
// the eval-cb reaches layer N+1's ensure() the bytes are already in
615+
// page cache (memcpy at ~10 GB/s vs cold pread at ~500 MB/s QD=1).
616+
//
617+
// Idempotent: the kernel deduplicates overlapping advise hints, and our
618+
// sentinel `s_last_advised_block` skips re-issuing for the same boundary.
619+
//
620+
// Wrap detection: when block_idx drops below `s_last_advised_block` we
621+
// assume a new forward pass began (the scheduler revisits block 0). Reset
622+
// the sentinel so we re-advise from the current block. This handles both
623+
// decode steps and any future op-reordering the scheduler does.
624+
//
625+
// RAM cost: each advised range is `size` bytes of page-cache pressure.
626+
// Tune via WP_FADVISE_LOOKAHEAD (default 2, 0 disables).
627+
{
628+
static int s_advise_k = -2; // -2 = unread env
629+
static int s_last_advised_block = -1;
630+
if (s_advise_k == -2) {
631+
const char * env = std::getenv("WP_FADVISE_LOOKAHEAD");
632+
s_advise_k = env ? std::atoi(env) : 2;
633+
if (s_advise_k < 0) s_advise_k = 0;
634+
LLAMA_LOG_INFO("[wp::eval_cb] WP_FADVISE_LOOKAHEAD=%d (0=disabled)\n", s_advise_k);
635+
}
636+
if (s_advise_k > 0) {
637+
int min_block = std::numeric_limits<int>::max();
638+
for (int j = 0; j < n_page_indices; ++j) {
639+
const int b = pager->page_meta(page_indices[j]).block_idx;
640+
if (b >= 0 && b < min_block) min_block = b;
641+
}
642+
if (min_block != std::numeric_limits<int>::max()) {
643+
if (min_block < s_last_advised_block) {
644+
// New forward pass — sentinel wraps. Re-advise from current.
645+
s_last_advised_block = -1;
646+
}
647+
if (min_block > s_last_advised_block) {
648+
pager->advise_layer_lookahead(min_block, s_advise_k);
649+
s_last_advised_block = min_block;
650+
}
651+
}
652+
}
653+
}
654+
479655
// Step 2: page each one in (waiting on prefetch if in flight, sync
480656
// fallback otherwise) and patch the matching src tensors.
481657
ggml_backend_buffer_t pool_buf = pager->pool_buf();
@@ -494,6 +670,12 @@ bool weight_pager_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) {
494670
// keeps debugging signal local to the failing op.
495671
continue;
496672
}
673+
// MAD-231: pin the slot so a subsequent prefetch alloc_slot in
674+
// tick() (or in a later op's pre-cb) cannot evict it while the
675+
// GPU is still reading from it. Unpinned at the top of the NEXT
676+
// eval_cb invocation.
677+
pager->pin_page(page_idx);
678+
s_pinned_pages_prev_op.push_back(page_idx);
497679

498680
const std::string & page_name = pager->page_meta(page_idx).tensor_name;
499681

0 commit comments

Comments
 (0)