diff --git a/core/prefetch/archer_prefetch_handle.cpp b/core/prefetch/archer_prefetch_handle.cpp index a5e52fb..7ae955f 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 dbf2c7e..ec4f986 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 832beae..aec3658 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 ffa3338..e6f3db7 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 b6d9b90..d1ef328 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 b6d953e..9249922 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 1f67f3a..09d6757 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