From 82d0089979b08d14b4abadc5bb717fc332e09e06 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sun, 16 Aug 2026 16:54:03 +0000 Subject: [PATCH] perf(prefetch): recover route-ahead priority band (BM3 blocked, NOT shipped) Re-implements the Task 9 route-ahead priority-band candidate that PR #161 implemented+built+validated locally then reverted (Python-only shipped) because its DFlash draft was absent. No original add/revert commit exists (verified via git log/pickaxe/fsck/stash) -- the revert was a working-tree discard -- so this restores the candidate per the plan spec to make it reviewable. C++ (core/prefetch, core/python): - task_scheduler.h: name the bands kOnDemandPriority=0, kRouteAheadPriority=1, kBackgroundPrefetchPriority=2 (NUM_PRIORITY stays 20). - EnqueuePrefetch (ordinary/background prefetch) now enqueues at background (2); EnqueuePrefetchTensors + the prefetch_tensors binding default to the dedicated route-ahead band (1). Python (expert_prefetcher.py): - add ExpertPrefetcher.route_ahead_priority (the BM3 knob, mirrors the native constants). Explicit route-ahead prefetch_experts_list issues at that band; legacy speculative_prefetch issues at background. Tests: - native GPU smoke asserts all three bands + the knob issue without raising; CPU mock tests assert explicit=route-ahead / legacy=background band mapping. BM3 was NOT run to a ship/no-ship verdict -- blocked by two issues on dev that are independent of this candidate (a 1-line EnqueuePrefetch priority + a Python knob; the native smoke swept all three bands without hanging): (A) bench_prefetch_priority._reset_cache() calls the terminal clean_up_resources() between arms, deadlocking the offload engine; the next forward hangs in fetch_tensors (model_offload.py:1670). No JSON. (B) exposed_fetch_seconds is uninstrumented, so _exposed_fetch_seconds() returns 0.0 for every arm and bm3_decision.exposed_fetch_improved can never be true -- ship_priority_band is structurally unmeasurable on dev. Per the plan's rule (no C++ ships without its paired benchmark proving the exposed window is closed) the band is NOT shipped. No false win either way. Validated (not a ship claim): SM120 build clean (0 errors); native priority smoke 5 passed on offloaded gpt-oss-20b; 60 CPU decision/mock tests passed; gpt-oss offload no-regression 9 passed. --- core/prefetch/archer_prefetch_handle.cpp | 2 +- core/prefetch/archer_prefetch_handle.h | 3 +- core/prefetch/task_scheduler.h | 7 +++ core/python/py_archer_prefetch.cpp | 2 +- moe_infinity/memory/expert_prefetcher.py | 20 +++++-- .../python/dflash/test_prefetch_native_gpu.py | 54 +++++++++++++++++++ .../dflash/test_speculative_prefetch.py | 51 +++++++++++++++++- 7 files changed, 131 insertions(+), 8 deletions(-) diff --git a/core/prefetch/archer_prefetch_handle.cpp b/core/prefetch/archer_prefetch_handle.cpp index a5e52fb1..7ae955f4 100644 --- a/core/prefetch/archer_prefetch_handle.cpp +++ b/core/prefetch/archer_prefetch_handle.cpp @@ -239,7 +239,7 @@ void ArcherPrefetchHandle::EnqueuePrefetch(const uint32_t tensor_id, auto node = kTopologyHandle->GetNodeFromTensorID(tensor_id); auto task = std::make_shared(); - task->priority = 1; + task->priority = kBackgroundPrefetchPriority; task->node = node; task->on_demand = false; task->src_device = node->device; diff --git a/core/prefetch/archer_prefetch_handle.h b/core/prefetch/archer_prefetch_handle.h index dbf2c7e1..ec4f9866 100644 --- a/core/prefetch/archer_prefetch_handle.h +++ b/core/prefetch/archer_prefetch_handle.h @@ -8,6 +8,7 @@ #include "aio/archer_tensor_handle.h" #include "model/model_topology.h" #include "parallel/expert_dispatcher.h" +#include "prefetch/task_scheduler.h" class ArcherPrefetchHandle { public: @@ -29,7 +30,7 @@ class ArcherPrefetchHandle { void ReplaceCacheCandidates(const std::vector& tensor_ids); void EnqueuePrefetch(const uint32_t tensor_id, int gpu_id); void EnqueuePrefetchTensors(const std::vector& tensor_ids, - std::uint32_t priority = 1); + std::uint32_t priority = kRouteAheadPriority); void OffloadTensor(torch::Tensor& tensor, const std::uint32_t tensor_id); void RegisterTensor(torch::Tensor& tensor, const std::uint32_t tensor_id); diff --git a/core/prefetch/task_scheduler.h b/core/prefetch/task_scheduler.h index 832beaef..aec3658b 100644 --- a/core/prefetch/task_scheduler.h +++ b/core/prefetch/task_scheduler.h @@ -23,6 +23,13 @@ #define NUM_PRIORITY 20UL +// Priority bands, lowest value serviced first (GPUThreadFunc scans queue 0 up). +// Band 0 is also the on-demand queue: dedup sweeps skip index 0, so route-ahead +// work parked at 0 keeps on-demand's non-preemptible semantics. +constexpr std::uint32_t kOnDemandPriority = 0; +constexpr std::uint32_t kRouteAheadPriority = 1; +constexpr std::uint32_t kBackgroundPrefetchPriority = 2; + struct Task { bool on_demand = false; NodePtr node; diff --git a/core/python/py_archer_prefetch.cpp b/core/python/py_archer_prefetch.cpp index ffa33383..e6f3db70 100644 --- a/core/python/py_archer_prefetch.cpp +++ b/core/python/py_archer_prefetch.cpp @@ -91,7 +91,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { &ArcherPrefetchHandle::GetNodeDefaultDevice) .def("get_node_device", &ArcherPrefetchHandle::GetNodeDevice) .def("prefetch_tensors", &ArcherPrefetchHandle::EnqueuePrefetchTensors, - py::arg("tensor_ids"), py::arg("priority") = 1) + py::arg("tensor_ids"), py::arg("priority") = kRouteAheadPriority) .def("replace_cache_candidates", &ArcherPrefetchHandle::ReplaceCacheCandidates) .def("enqueue_prefetch", &ArcherPrefetchHandle::EnqueuePrefetch) diff --git a/moe_infinity/memory/expert_prefetcher.py b/moe_infinity/memory/expert_prefetcher.py index b6d9b908..d1ef3289 100644 --- a/moe_infinity/memory/expert_prefetcher.py +++ b/moe_infinity/memory/expert_prefetcher.py @@ -12,6 +12,11 @@ from moe_infinity.utils import parse_moe_param +# Native prefetch priority bands; must mirror core/prefetch/task_scheduler.h. +ON_DEMAND_PRIORITY = 0 +ROUTE_AHEAD_PRIORITY = 1 +BACKGROUND_PREFETCH_PRIORITY = 2 + try: import nvtx # type: ignore[reportMissingTypeStubs] except ImportError: @@ -45,6 +50,7 @@ def _hit_rate_from_visit_counts(counts: Any) -> Optional[float]: class ExpertPrefetcher(object): cache_file_rd: Optional[Any] = None first_k_dense_replace: int = 0 + route_ahead_priority: int = ROUTE_AHEAD_PRIORITY archer_engine: Any expert_dispatcher: Optional[Any] = None expert_tensor_map: dict[tuple[int, int], int] @@ -149,7 +155,12 @@ def wasted_prefetch_bytes(self) -> float: return 0.0 return 0.0 - def prefetch_experts_list(self, layer_id: int, expert_list: List[int]): + def prefetch_experts_list( + self, + layer_id: int, + expert_list: List[int], + priority: Optional[int] = None, + ): if self.archer_engine is None: return tensor_ids = [] @@ -157,9 +168,10 @@ def prefetch_experts_list(self, layer_id: int, expert_list: List[int]): tensor_ids.append(self.expert_tensor_map[(layer_id, j)]) if not tensor_ids: return + band = self.route_ahead_priority if priority is None else priority batched_issue = getattr(self.archer_engine, "prefetch_tensors", None) if callable(batched_issue): - batched_issue(tensor_ids) + batched_issue(tensor_ids, priority=band) return for tensor_id in tensor_ids: gpu_id = self.archer_engine.get_node_default_device([tensor_id]) @@ -274,7 +286,9 @@ def speculative_prefetch( ::-1 ].tolist() - self.prefetch_experts_list(next_layer, topk_indices) + self.prefetch_experts_list( + next_layer, topk_indices, priority=BACKGROUND_PREFETCH_PRIORITY + ) self._last_speculative_prediction = set(topk_indices) def correct_prefetch( diff --git a/tests/python/dflash/test_prefetch_native_gpu.py b/tests/python/dflash/test_prefetch_native_gpu.py index b6d953e7..9249922c 100644 --- a/tests/python/dflash/test_prefetch_native_gpu.py +++ b/tests/python/dflash/test_prefetch_native_gpu.py @@ -117,3 +117,57 @@ def test_native_batched_prefetch_experts_list_uses_batched_path( if layer == some_layer ) offloaded_prefetcher.prefetch_experts_list(some_layer, experts) + + +def test_native_priority_bands_accept_each_service_class_reverse_order( + offloaded_prefetcher, +) -> None: + from moe_infinity.memory.expert_prefetcher import ( + BACKGROUND_PREFETCH_PRIORITY, + ON_DEMAND_PRIORITY, + ROUTE_AHEAD_PRIORITY, + ) + + engine = offloaded_prefetcher.archer_engine + tensor_ids = _saturated_ids(offloaded_prefetcher) + assert tensor_ids, "no offloaded expert tensors to issue" + + for priority in ( + BACKGROUND_PREFETCH_PRIORITY, + ROUTE_AHEAD_PRIORITY, + ON_DEMAND_PRIORITY, + ): + assert engine.prefetch_tensors(tensor_ids, priority) is None + torch.cuda.synchronize() + + +def test_native_route_ahead_priority_knob_issues_each_band( + offloaded_prefetcher, +) -> None: + from moe_infinity.memory.expert_prefetcher import ( + BACKGROUND_PREFETCH_PRIORITY, + ON_DEMAND_PRIORITY, + ROUTE_AHEAD_PRIORITY, + ) + + layers = sorted( + {layer for layer, _e in offloaded_prefetcher.expert_tensor_map} + ) + some_layer = layers[0] + experts = sorted( + expert + for layer, expert in offloaded_prefetcher.expert_tensor_map + if layer == some_layer + ) + original = offloaded_prefetcher.route_ahead_priority + try: + for band in ( + BACKGROUND_PREFETCH_PRIORITY, + ROUTE_AHEAD_PRIORITY, + ON_DEMAND_PRIORITY, + ): + offloaded_prefetcher.route_ahead_priority = band + offloaded_prefetcher.prefetch_experts_list(some_layer, experts) + torch.cuda.synchronize() + finally: + offloaded_prefetcher.route_ahead_priority = original diff --git a/tests/python/dflash/test_speculative_prefetch.py b/tests/python/dflash/test_speculative_prefetch.py index 1f67f3a6..09d67577 100644 --- a/tests/python/dflash/test_speculative_prefetch.py +++ b/tests/python/dflash/test_speculative_prefetch.py @@ -26,7 +26,12 @@ import pytest import torch -from moe_infinity.memory.expert_prefetcher import ExpertPrefetcher +from moe_infinity.memory.expert_prefetcher import ( + BACKGROUND_PREFETCH_PRIORITY, + ON_DEMAND_PRIORITY, + ROUTE_AHEAD_PRIORITY, + ExpertPrefetcher, +) # Hand-checked legacy fixture: [2 tokens, 4 experts] # mean over tokens = [2.0, 3.5, 3.0, 0.5] -> top-2 = experts [1, 2] @@ -185,7 +190,9 @@ def _make_prefetcher_without_batch(num_layers: int = 8, num_experts: int = 8): def test_prefetch_experts_list_batches_one_native_call_when_available(): prefetcher, engine = _make_prefetcher(num_layers=8, num_experts=8) prefetcher.prefetch_experts_list(3, [3, 1, 7]) - engine.prefetch_tensors.assert_called_once_with([303, 301, 307]) + engine.prefetch_tensors.assert_called_once_with( + [303, 301, 307], priority=ROUTE_AHEAD_PRIORITY + ) engine.enqueue_prefetch.assert_not_called() @@ -202,3 +209,43 @@ def test_prefetch_experts_list_empty_batch_calls_neither_path(): prefetcher.prefetch_experts_list(3, []) engine.prefetch_tensors.assert_not_called() engine.enqueue_prefetch.assert_not_called() + + +# --------------------------------------------------------------------------- +# priority bands (plan Task 9 Step 5): explicit route-ahead vs legacy background +# --------------------------------------------------------------------------- + + +def _issued_priority(engine: MagicMock) -> int: + return engine.prefetch_tensors.call_args.kwargs["priority"] + + +def test_explicit_route_ahead_issues_on_the_route_ahead_band(): + prefetcher, engine = _make_prefetcher(num_layers=8, num_experts=8) + prefetcher.speculative_prefetch( + 1, expert_ids=[3, 1, 7], prefetch_layer_id=5 + ) + assert _enqueued_tensor_ids(engine) == [503, 501, 507] + assert _issued_priority(engine) == ROUTE_AHEAD_PRIORITY + + +def test_legacy_router_logits_issues_on_the_background_band(): + prefetcher, engine = _make_prefetcher(num_layers=8, num_experts=4) + prefetcher.speculative_prefetch(2, LOGITS) + assert _enqueued_tensor_ids(engine) == [301, 302] + assert _issued_priority(engine) == BACKGROUND_PREFETCH_PRIORITY + + +def test_correct_prefetch_issues_on_the_route_ahead_band(): + prefetcher, engine = _make_prefetcher(num_layers=8, num_experts=8) + prefetcher._last_speculative_prediction = {1} + prefetcher.correct_prefetch(3, [3, 1, 7]) + assert _enqueued_tensor_ids(engine) == [303, 307] + assert _issued_priority(engine) == ROUTE_AHEAD_PRIORITY + + +def test_route_ahead_priority_knob_sweeps_the_issued_band(): + prefetcher, engine = _make_prefetcher(num_layers=8, num_experts=8) + prefetcher.route_ahead_priority = ON_DEMAND_PRIORITY + prefetcher.prefetch_experts_list(3, [3, 1, 7]) + assert _issued_priority(engine) == ON_DEMAND_PRIORITY