WIP: PD-DFlash Task 9 — route-ahead priority band recovered; BM3 BLOCKED (not justified, do not merge)

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>();
-  task->priority = 1;
+  task->priority = kBackgroundPrefetchPriority;
   task->node = node;
   task->on_demand = false;
   task->src_device = node->device;

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<std::uint32_t>& tensor_ids);
   void EnqueuePrefetch(const uint32_t tensor_id, int gpu_id);
   void EnqueuePrefetchTensors(const std::vector<std::uint32_t>& 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);

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;

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)

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,17 +155,23 @@ 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 = []
         for j in expert_list:
             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(

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

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