diff --git a/benchmarks/dflash/_serving_measure.py b/benchmarks/dflash/_serving_measure.py index 1d41f974..0c4c714f 100644 --- a/benchmarks/dflash/_serving_measure.py +++ b/benchmarks/dflash/_serving_measure.py @@ -130,11 +130,14 @@ def measure_configuration( slo_ms=args.slo_ms, warnings=warnings, ) + effective_h2d_gbps = args.measured_h2d_gbps + if effective_h2d_gbps is None and getattr(args, "probe_h2d", False): + effective_h2d_gbps = _probe_h2d_gbps() cost_terms = _collect_cost_terms( baseline=baseline, speculator=speculator, engine=engine, - measured_h2d_gbps=args.measured_h2d_gbps, + measured_h2d_gbps=effective_h2d_gbps, warnings=warnings, ) return make_observation_row( @@ -268,6 +271,42 @@ def _route_ahead_snapshot(speculator: Any) -> Tuple[float, Optional[int]]: return coverage, (int(wasted) if wasted is not None else None) +def _probe_h2d_gbps( + nbytes: int = 256 * 1024 * 1024, iters: int = 20 +) -> Optional[float]: + import torch + + if not torch.cuda.is_available(): + return None + try: + elements = max(1, nbytes // 2) + host = torch.empty(elements, dtype=torch.float16, pin_memory=True) + device = torch.empty(elements, dtype=torch.float16, device="cuda:0") + torch.cuda.synchronize() + device.copy_(host, non_blocking=True) + torch.cuda.synchronize() + started = time.perf_counter() + for _ in range(max(1, iters)): + device.copy_(host, non_blocking=True) + torch.cuda.synchronize() + elapsed = max(time.perf_counter() - started, 1e-9) + except Exception: + return None + moved = float(elements * 2 * max(1, iters)) + return moved / elapsed / 1_000_000_000.0 + + +def _native_wasted_prefetch_bytes(engine: Any) -> Optional[float]: + prefetcher = getattr(engine, "expert_prefetcher", None) + getter = getattr(prefetcher, "wasted_prefetch_bytes", None) + if callable(getter): + try: + return float(getter()) + except Exception: + return None + return None + + def _extract_float(source: Any, names: Tuple[str, ...]) -> Optional[float]: for name in names: value = getattr(source, name, None) @@ -300,11 +339,16 @@ def _collect_metrics( per_round_latency = elapsed / rounds coverage, wasted_bytes = _route_ahead_snapshot(speculator) if wasted_bytes is None: - warnings.append( - "wasted_prefetch_bytes unavailable from RouteAheadStats; a route-" - "ahead configuration on offloaded experts must report real bytes" - ) - wasted_bytes = 0 + native_wasted = _native_wasted_prefetch_bytes(engine) + if native_wasted is not None: + wasted_bytes = int(native_wasted) + else: + warnings.append( + "wasted_prefetch_bytes unavailable from RouteAheadStats; a " + "route-ahead configuration on offloaded experts must report " + "real bytes" + ) + wasted_bytes = 0 hit_rate = _extract_float( getattr(engine, "expert_prefetcher", engine), @@ -358,6 +402,10 @@ def _collect_cost_terms( warnings: List[str], ) -> Dict[str, Any]: coverage, wasted_bytes = _route_ahead_snapshot(speculator) + if wasted_bytes is None: + native_wasted = _native_wasted_prefetch_bytes(engine) + if native_wasted is not None: + wasted_bytes = int(native_wasted) terms: Dict[str, Any] = { "route_ahead_coverage": coverage, "wasted_prefetch_bytes": wasted_bytes, diff --git a/benchmarks/dflash/pd_dflash_serving.py b/benchmarks/dflash/pd_dflash_serving.py index f9ea3721..ba9af83a 100644 --- a/benchmarks/dflash/pd_dflash_serving.py +++ b/benchmarks/dflash/pd_dflash_serving.py @@ -15,8 +15,10 @@ Design contract (frozen here, cross-checked by the aggregator): * baselines are exactly B0-B3 with the design-doc §8 semantics; -* the required generalization targets are ``Qwen/Qwen3-Coder-30B-A3B`` and - ``openai/gpt-oss-20b`` with their ``z-lab`` DFlash drafts; +* the required generalization targets are + ``Qwen/Qwen3-Coder-30B-A3B-Instruct`` and ``openai/gpt-oss-20b`` with their + ``z-lab`` DFlash drafts (the bare ``Qwen/Qwen3-Coder-30B-A3B`` repo returns + HTTP 404 and never existed; ``-Instruct`` is the real ``qwen3_moe`` repo); * block sizes are 8 and 16, concurrency sweeps 1..32; and * every emitted observation carries the full ``REQUIRED_METRICS`` schema, plus the byte-accurate route-ahead ``wasted_prefetch_bytes`` from the instrumented @@ -50,12 +52,12 @@ } REQUIRED_MODELS: Tuple[str, ...] = ( - "Qwen/Qwen3-Coder-30B-A3B", + "Qwen/Qwen3-Coder-30B-A3B-Instruct", "openai/gpt-oss-20b", ) REQUIRED_DRAFTS: Dict[str, str] = { - "Qwen/Qwen3-Coder-30B-A3B": "z-lab/Qwen3-Coder-30B-A3B-DFlash", + "Qwen/Qwen3-Coder-30B-A3B-Instruct": "z-lab/Qwen3-Coder-30B-A3B-DFlash", "openai/gpt-oss-20b": "z-lab/gpt-oss-20b-DFlash", } @@ -251,6 +253,7 @@ class RunnerArgs: requests: int warmup_rounds: int measured_h2d_gbps: Optional[float] + probe_h2d: bool slo_ms: Optional[float] seed: int device_memory_ratio: float @@ -294,6 +297,15 @@ def build_arg_parser() -> argparse.ArgumentParser: parser.add_argument("--requests", type=int, default=64) parser.add_argument("--warmup-rounds", type=int, default=5) parser.add_argument("--measured-h2d-gbps", type=float, default=None) + parser.add_argument( + "--probe-h2d", + action="store_true", + help=( + "when --measured-h2d-gbps is unset, measure host->device " + "bandwidth on the visible GPU and use it for the §7 hide " + "inequality (never a theoretical PCIe figure)" + ), + ) parser.add_argument("--slo-ms", type=float, default=None) parser.add_argument("--seed", type=int, default=1408) parser.add_argument( @@ -323,6 +335,7 @@ def parse_args(argv: Optional[Sequence[str]] = None) -> RunnerArgs: requests=parsed.requests, warmup_rounds=parsed.warmup_rounds, measured_h2d_gbps=parsed.measured_h2d_gbps, + probe_h2d=parsed.probe_h2d, slo_ms=parsed.slo_ms, seed=parsed.seed, device_memory_ratio=parsed.device_memory_ratio, diff --git a/core/model/model_topology.cpp b/core/model/model_topology.cpp index 494d7da9..f04582e5 100644 --- a/core/model/model_topology.cpp +++ b/core/model/model_topology.cpp @@ -372,6 +372,24 @@ ArcherTopologyHandle::GetNodeVisitCounts() { return node_visit_counts; } +std::tuple +ArcherTopologyHandle::GetResidentAndWastedBytes() { + std::lock_guard lock(mutex_); + std::int64_t resident = 0; + std::int64_t wasted = 0; + for (auto& stage : pipeline_.stages) { + for (auto& node_body : stage->nodes) { + auto& node = node_body->node; + if (node == nullptr) continue; + if (node->device.is_cuda()) { + resident += node->byte_size; + } + wasted += node->byte_size * static_cast(node->unused_count); + } + } + return std::make_tuple(resident, wasted); +} + std::vector ArcherTopologyHandle::GetChildVisitCounts() { std::lock_guard lock(mutex_); int num_layers = 0; diff --git a/core/model/model_topology.h b/core/model/model_topology.h index beaee25e..7dc6b969 100644 --- a/core/model/model_topology.h +++ b/core/model/model_topology.h @@ -185,6 +185,7 @@ class ArcherTopologyHandle : public base::noncopyable { void DisableTrace() noexcept { trace_enabled_ = false; } std::vector> GetNodeVisitCounts(); + std::tuple GetResidentAndWastedBytes(); std::vector GetChildVisitCounts(); void SetNodeVisitCounts(const std::vector& visit_counts); void SetChildVisitCounts(const std::vector& visit_counts); diff --git a/core/parallel/expert_dispatcher.cpp b/core/parallel/expert_dispatcher.cpp index be39ce06..f3b525a1 100644 --- a/core/parallel/expert_dispatcher.cpp +++ b/core/parallel/expert_dispatcher.cpp @@ -233,6 +233,8 @@ void ExpertDispatcher::Enqueue(CallArgs& args) { exec_args.out_dtype = c10::typeMetaToScalarType(hidden_states_.dtype()); exec_args.evict = false; exec_args.hit = true; + cache_hit_count_.fetch_add(1, std::memory_order_relaxed); + cache_access_count_.fetch_add(1, std::memory_order_relaxed); // transfer_event = nullptr: expert is already on GPU, no H2D wait needed // module_->SetTensorsFromIds(expert_node->node->tensor_ids); @@ -296,6 +298,39 @@ void ExpertDispatcher::ClearExpertCacheCounts() { expert_node->node->incache_visit_count = 0; } } + cache_hit_count_.store(0, std::memory_order_relaxed); + cache_access_count_.store(0, std::memory_order_relaxed); +} + +std::int64_t ExpertDispatcher::GetCacheOccupancyBytes() { + std::int64_t total = 0; + for (auto& gpu_cache : cached_experts_) { + for (auto key : gpu_cache) { + int layer_idx = static_cast(key >> 32); + int expert_idx = static_cast(key & 0xFFFFFFFF); + if (expert_idx < 0 || expert_idx >= static_cast(experts_.size())) { + continue; + } + if (layer_idx < 0 || + layer_idx >= static_cast(experts_[expert_idx].size())) { + continue; + } + auto& expert_node = experts_[expert_idx][layer_idx]; + if (expert_node && expert_node->node) { + total += expert_node->node->byte_size; + } + } + } + return total; +} + +double ExpertDispatcher::GetCacheHitRate() const { + std::uint64_t access = cache_access_count_.load(std::memory_order_relaxed); + if (access == 0) { + return 0.0; + } + std::uint64_t hits = cache_hit_count_.load(std::memory_order_relaxed); + return static_cast(hits) / static_cast(access); } // void ExpertDispatcher::GPUThreadFunc(int gpu_id) { @@ -388,6 +423,8 @@ void ExpertDispatcher::GPUFetchFunc(int gpu_id) { exec_args.out_dtype = c10::typeMetaToScalarType(hidden_states_.dtype()); exec_args.evict = false; exec_args.hit = true; + cache_hit_count_.fetch_add(1, std::memory_order_relaxed); + cache_access_count_.fetch_add(1, std::memory_order_relaxed); exec_args.transfer_event = nullptr; exec_queue_[gpu_id].Push(exec_args); continue; @@ -537,6 +574,8 @@ void ExpertDispatcher::GPUFetchFunc(int gpu_id) { exec_args.out_dtype = c10::typeMetaToScalarType(hidden_states_.dtype()); exec_args.evict = gpu_overload_[gpu_id].load(std::memory_order_acquire); exec_args.hit = cache_hit; + cache_access_count_.fetch_add(1, std::memory_order_relaxed); + if (cache_hit) cache_hit_count_.fetch_add(1, std::memory_order_relaxed); exec_args.transfer_event = transfer_done; // std::lock_guard lock(exec_mutex_[gpu_id]); // exec_queue_[gpu_id].emplace_back(std::move(exec_args)); diff --git a/core/parallel/expert_dispatcher.h b/core/parallel/expert_dispatcher.h index 39b08cc6..47d3b69a 100644 --- a/core/parallel/expert_dispatcher.h +++ b/core/parallel/expert_dispatcher.h @@ -108,6 +108,9 @@ class ExpertDispatcher : public base::noncopyable { const std::vector& tensor_ids, std::string jit_path); void ClearExpertCacheCounts(); + // Read-only observability accessors; neither alters routing/dispatch. + std::int64_t GetCacheOccupancyBytes(); + double GetCacheHitRate() const; void SetExpectedQueue(int expected_pending = 0) { pending_.store(expected_pending); } @@ -150,6 +153,10 @@ class ExpertDispatcher : public base::noncopyable { std::atomic pending_; + // Passive counters for GetCacheHitRate(); reset by ClearExpertCacheCounts(). + std::atomic cache_hit_count_{0}; + std::atomic cache_access_count_{0}; + std::mutex pending_mutex_; std::condition_variable pending_cv_; diff --git a/core/prefetch/archer_prefetch_handle.cpp b/core/prefetch/archer_prefetch_handle.cpp index b3bae148..a5e52fb1 100644 --- a/core/prefetch/archer_prefetch_handle.cpp +++ b/core/prefetch/archer_prefetch_handle.cpp @@ -341,6 +341,14 @@ torch::Tensor ArcherPrefetchHandle::GetHitRate() { return trace; } +std::int64_t ArcherPrefetchHandle::GetExpertOccupancyBytes() { + return std::get<0>(kTopologyHandle->GetResidentAndWastedBytes()); +} + +std::int64_t ArcherPrefetchHandle::GetWastedPrefetchBytes() { + return std::get<1>(kTopologyHandle->GetResidentAndWastedBytes()); +} + void ArcherPrefetchHandle::SetTrace(const torch::Tensor& trace) { if (trace.dim() != 3 || !trace.is_contiguous() || !trace.is_cpu()) { DLOG_ERROR("Trace should be a contiguous 3D tensor on CPU"); diff --git a/core/prefetch/archer_prefetch_handle.h b/core/prefetch/archer_prefetch_handle.h index c18a8f84..dbf2c7e1 100644 --- a/core/prefetch/archer_prefetch_handle.h +++ b/core/prefetch/archer_prefetch_handle.h @@ -43,6 +43,8 @@ class ArcherPrefetchHandle { torch::Tensor GetTrace(); torch::Tensor GetHitRate(); + std::int64_t GetExpertOccupancyBytes(); + std::int64_t GetWastedPrefetchBytes(); void SetTrace(const torch::Tensor& trace); void TraceRequest(const std::uint64_t request_id, const TensorID tensor_id); void SetTopology(const std::vector< diff --git a/core/python/py_archer_prefetch.cpp b/core/python/py_archer_prefetch.cpp index 6e6b016c..ffa33383 100644 --- a/core/python/py_archer_prefetch.cpp +++ b/core/python/py_archer_prefetch.cpp @@ -53,6 +53,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { // ArcherPrefetchHandle::GetTrace) .def("get_hit_rate", (torch::Tensor(ArcherPrefetchHandle::*)()) & ArcherPrefetchHandle::GetHitRate) + .def("get_expert_occupancy_bytes", + &ArcherPrefetchHandle::GetExpertOccupancyBytes) + .def("get_wasted_prefetch_bytes", + &ArcherPrefetchHandle::GetWastedPrefetchBytes) .def("set_trace", (void(ArcherPrefetchHandle::*)(const torch::Tensor&)) & ArcherPrefetchHandle::SetTrace) // .def("trace_request", @@ -111,6 +115,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("notify_fetch_start", &ExpertDispatcher::NotifyFetchStart) .def("clear_expert_cache_counts", &ExpertDispatcher::ClearExpertCacheCounts) + .def("get_cache_occupancy_bytes", + &ExpertDispatcher::GetCacheOccupancyBytes) + .def("get_cache_hit_rate", &ExpertDispatcher::GetCacheHitRate) .def("set_scales", &ExpertDispatcher::SetScales, "Store fp8 block scales for dequant-on-copy (fp8-in-store path)"); } diff --git a/moe_infinity/memory/expert_prefetcher.py b/moe_infinity/memory/expert_prefetcher.py index ad1b8d98..b6d9b908 100644 --- a/moe_infinity/memory/expert_prefetcher.py +++ b/moe_infinity/memory/expert_prefetcher.py @@ -27,10 +27,26 @@ IOProfiler = None +def _hit_rate_from_visit_counts(counts: Any) -> Optional[float]: + if counts is None or not hasattr(counts, "numel"): + return None + try: + if counts.numel() == 0 or counts.dim() != 2 or counts.shape[1] < 4: + return None + visit = float(counts[:, 0].sum().item()) + hit = float(counts[:, 3].sum().item()) + except Exception: + return None + if visit <= 0.0: + return None + return hit / visit + + class ExpertPrefetcher(object): cache_file_rd: Optional[Any] = None first_k_dense_replace: int = 0 archer_engine: Any + expert_dispatcher: Optional[Any] = None expert_tensor_map: dict[tuple[int, int], int] expert_nbytes_map: dict[tuple[int, int], int] @@ -40,6 +56,7 @@ def __init__(self, config: PretrainedConfig): parse_moe_param(config) ) self.archer_engine: Optional[Any] = None + self.expert_dispatcher: Optional[Any] = None self.expert_tensor_map: Dict[Tuple[int, int], int] = {} self.expert_nbytes_map: Dict[Tuple[int, int], int] = {} self._last_speculative_prediction: Set[int] = set() @@ -49,6 +66,89 @@ def set_archer_engine(self, archer_engine: Any): _expert_prefetcher = archer_engine self.archer_engine = archer_engine + @property + def num_offloaded_experts(self) -> int: + engine = self.archer_engine + checker = ( + getattr(engine, "is_tensor_offloaded", None) + if engine is not None + else None + ) + if not callable(checker): + return len(self.expert_tensor_map) + count = 0 + for tensor_id in self.expert_tensor_map.values(): + try: + if checker(int(tensor_id)): + count += 1 + except Exception: + continue + return count + + def get_hit_rate(self) -> float: + dispatcher = self.expert_dispatcher + if dispatcher is not None: + getter = getattr(dispatcher, "get_cache_hit_rate", None) + if callable(getter): + try: + rate = float(getter()) + except Exception: + rate = 0.0 + if rate: + return rate + engine = self.archer_engine + getter = ( + getattr(engine, "get_hit_rate", None) + if engine is not None + else None + ) + if callable(getter): + try: + counts = getter() + except Exception: + counts = None + rate = _hit_rate_from_visit_counts(counts) + if rate is not None: + return rate + return 0.0 + + def expert_occupancy_bytes(self) -> float: + total = 0.0 + dispatcher = self.expert_dispatcher + if dispatcher is not None: + getter = getattr(dispatcher, "get_cache_occupancy_bytes", None) + if callable(getter): + try: + total += float(getter()) + except Exception: + pass + engine = self.archer_engine + getter = ( + getattr(engine, "get_expert_occupancy_bytes", None) + if engine is not None + else None + ) + if callable(getter): + try: + total += float(getter()) + except Exception: + pass + return total + + def wasted_prefetch_bytes(self) -> float: + engine = self.archer_engine + getter = ( + getattr(engine, "get_wasted_prefetch_bytes", None) + if engine is not None + else None + ) + if callable(getter): + try: + return float(getter()) + except Exception: + return 0.0 + return 0.0 + def prefetch_experts_list(self, layer_id: int, expert_list: List[int]): if self.archer_engine is None: return diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index 66c75a32..bd8f6fa2 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -14,7 +14,7 @@ import re import tempfile import warnings -from typing import Callable, Dict, Type, Union +from typing import Callable, Dict, Optional, Type, Union import torch import transformers @@ -461,6 +461,41 @@ class OffloadEngine(object): request_id = 0 config = {} + @property + def num_offloaded_experts(self) -> int: + prefetcher = getattr(self, "expert_prefetcher", None) + if prefetcher is None: + return 0 + return int(getattr(prefetcher, "num_offloaded_experts", 0)) + + @property + def expert_cache_hit_rate(self) -> float: + prefetcher = getattr(self, "expert_prefetcher", None) + getter = getattr(prefetcher, "get_hit_rate", None) + return float(getter()) if callable(getter) else 0.0 + + @property + def expert_occupancy_bytes(self) -> float: + prefetcher = getattr(self, "expert_prefetcher", None) + getter = getattr(prefetcher, "expert_occupancy_bytes", None) + return float(getter()) if callable(getter) else 0.0 + + @property + def kv_occupancy_bytes(self) -> Optional[float]: + manager = getattr(self, "kv_cache_manager", None) + if manager is None: + return None + for name in ("occupancy_bytes", "resident_bytes", "kv_cache_bytes"): + value = getattr(manager, name, None) + if callable(value): + try: + value = value() + except Exception: + value = None + if isinstance(value, (int, float)) and not isinstance(value, bool): + return float(value) + return None + def __init__( self, capacity, @@ -472,6 +507,7 @@ def __init__( ): self.offload_exemption = set() self.expert_modules = [] + self.kv_cache_manager = kv_cache_manager self.ckpt_files = [] @@ -1040,6 +1076,9 @@ def archer_from_pretrained(cls, *args, **kwargs): parse_expert_type(self.config), self.archer_config.num_threads, ) + self.expert_prefetcher.expert_dispatcher = ( + self.expert_dispatcher + ) for name, param in model.named_parameters(recurse=True): # remove base_model_prefix from self.name_id_map diff --git a/tests/python/dflash/test_native_stat_accessors.py b/tests/python/dflash/test_native_stat_accessors.py new file mode 100644 index 00000000..ed05ea4c --- /dev/null +++ b/tests/python/dflash/test_native_stat_accessors.py @@ -0,0 +1,152 @@ +"""CPU-only tests for the PD-DFlash native offloaded-expert stat accessors. + +Exercises the Python surface the §8 runner probes (``_serving_measure`` and +``ExpertPrefetcher``) with fakes standing in for the native ``_store`` engine, +so the wiring is verified without a GPU or a built extension. The native C++ +getters themselves are covered by the on-hardware §8 re-run. +""" + +from __future__ import annotations + +import torch + +from benchmarks.dflash._serving_measure import ( + _count_offloaded_experts, + _native_wasted_prefetch_bytes, + _probe_h2d_gbps, +) +from moe_infinity.memory.expert_prefetcher import ( + ExpertPrefetcher, + _hit_rate_from_visit_counts, +) + + +def _bare_prefetcher(**attrs): + prefetcher = object.__new__(ExpertPrefetcher) + prefetcher.archer_engine = None + prefetcher.expert_dispatcher = None + prefetcher.expert_tensor_map = {} + for key, value in attrs.items(): + setattr(prefetcher, key, value) + return prefetcher + + +class _FakeHandle: + def __init__(self, *, offloaded=(), occupancy=0, wasted=0, counts=None): + self._offloaded = set(offloaded) + self._occupancy = occupancy + self._wasted = wasted + self._counts = counts + + def is_tensor_offloaded(self, tensor_id): + return int(tensor_id) in self._offloaded + + def get_expert_occupancy_bytes(self): + return self._occupancy + + def get_wasted_prefetch_bytes(self): + return self._wasted + + def get_hit_rate(self): + return self._counts + + +class _FakeDispatcher: + def __init__(self, *, hit_rate=0.0, occupancy=0): + self._hit_rate = hit_rate + self._occupancy = occupancy + + def get_cache_hit_rate(self): + return self._hit_rate + + def get_cache_occupancy_bytes(self): + return self._occupancy + + +def test_num_offloaded_experts_counts_offloaded_tensor_ids(): + handle = _FakeHandle(offloaded=(10, 11, 13)) + prefetcher = _bare_prefetcher( + archer_engine=handle, + expert_tensor_map={(0, 0): 10, (0, 1): 11, (0, 2): 12, (1, 0): 13}, + ) + assert prefetcher.num_offloaded_experts == 3 + + +def test_num_offloaded_experts_falls_back_to_map_len_without_checker(): + prefetcher = _bare_prefetcher( + archer_engine=object(), + expert_tensor_map={(0, 0): 1, (0, 1): 2}, + ) + assert prefetcher.num_offloaded_experts == 2 + + +def test_hit_rate_prefers_nonzero_dispatcher_signal(): + prefetcher = _bare_prefetcher( + archer_engine=_FakeHandle(counts=torch.tensor([[10, 0, 0, 9]])), + expert_dispatcher=_FakeDispatcher(hit_rate=0.75), + ) + assert prefetcher.get_hit_rate() == 0.75 + + +def test_hit_rate_falls_back_to_topology_visit_counts(): + counts = torch.tensor([[10, 6, 4, 4], [6, 4, 2, 3]], dtype=torch.int64) + prefetcher = _bare_prefetcher(archer_engine=_FakeHandle(counts=counts)) + assert abs(prefetcher.get_hit_rate() - (7.0 / 16.0)) < 1e-9 + + +def test_hit_rate_zero_when_no_signal(): + prefetcher = _bare_prefetcher(archer_engine=_FakeHandle(counts=None)) + assert prefetcher.get_hit_rate() == 0.0 + + +def test_hit_rate_from_visit_counts_handles_empty_and_no_visits(): + assert _hit_rate_from_visit_counts(None) is None + assert _hit_rate_from_visit_counts(torch.empty((0, 4))) is None + assert _hit_rate_from_visit_counts(torch.tensor([[0, 0, 0, 0]])) is None + + +def test_expert_occupancy_sums_dispatcher_and_handle(): + prefetcher = _bare_prefetcher( + archer_engine=_FakeHandle(occupancy=2048), + expert_dispatcher=_FakeDispatcher(occupancy=4096), + ) + assert prefetcher.expert_occupancy_bytes() == 6144.0 + + +def test_wasted_prefetch_bytes_reads_handle_getter(): + prefetcher = _bare_prefetcher(archer_engine=_FakeHandle(wasted=512)) + assert prefetcher.wasted_prefetch_bytes() == 512.0 + + +class _FakeEngine: + def __init__(self, prefetcher=None, num_offloaded=None): + if prefetcher is not None: + self.expert_prefetcher = prefetcher + if num_offloaded is not None: + self.num_offloaded_experts = num_offloaded + + +def test_count_offloaded_experts_reads_native_engine_attr(): + assert _count_offloaded_experts(_FakeEngine(num_offloaded=7)) == 7 + + +def test_count_offloaded_experts_falls_back_to_nbytes_map(): + class _P: + expert_nbytes_map = {(0, 0): 10, (0, 1): 20} + + assert _count_offloaded_experts(_FakeEngine(prefetcher=_P())) == 2 + + +def test_native_wasted_prefetch_bytes_reads_prefetcher(): + prefetcher = _bare_prefetcher(archer_engine=_FakeHandle(wasted=768)) + engine = _FakeEngine(prefetcher=prefetcher) + assert _native_wasted_prefetch_bytes(engine) == 768.0 + + +def test_native_wasted_prefetch_bytes_none_without_getter(): + assert _native_wasted_prefetch_bytes(_FakeEngine()) is None + + +def test_probe_h2d_gbps_returns_none_or_positive(): + result = _probe_h2d_gbps(nbytes=1 << 20, iters=2) + assert result is None or (isinstance(result, float) and result > 0.0) diff --git a/tests/python/dflash/test_pd_dflash_serving_contract.py b/tests/python/dflash/test_pd_dflash_serving_contract.py index 4c263d87..d38571b3 100644 --- a/tests/python/dflash/test_pd_dflash_serving_contract.py +++ b/tests/python/dflash/test_pd_dflash_serving_contract.py @@ -88,7 +88,7 @@ def test_require_offloaded_refuses_resident_b0_b1_b2(): def test_observation_row_requires_full_metric_schema(): row = make_observation_row( - model="Qwen/Qwen3-Coder-30B-A3B", + model="Qwen/Qwen3-Coder-30B-A3B-Instruct", draft="z-lab/Qwen3-Coder-30B-A3B-DFlash", baseline="B1", block_size=16, @@ -165,7 +165,7 @@ def test_parse_args_defaults_cover_the_full_matrix(): args = parse_args( [ "--model", - "Qwen/Qwen3-Coder-30B-A3B", + "Qwen/Qwen3-Coder-30B-A3B-Instruct", "--draft", "z-lab/Qwen3-Coder-30B-A3B-DFlash", "--offload-dir", diff --git a/tests/python/dflash/test_pd_dflash_serving_gpu.py b/tests/python/dflash/test_pd_dflash_serving_gpu.py index b6d97efe..f531f994 100644 --- a/tests/python/dflash/test_pd_dflash_serving_gpu.py +++ b/tests/python/dflash/test_pd_dflash_serving_gpu.py @@ -33,12 +33,12 @@ def test_dry_run_contract_matches_required_matrix(capsys): assert main(["--dry-run-contract"]) == 0 contract = json.loads(capsys.readouterr().out) - assert "Qwen/Qwen3-Coder-30B-A3B" in contract["models"] + assert "Qwen/Qwen3-Coder-30B-A3B-Instruct" in contract["models"] assert "openai/gpt-oss-20b" in contract["models"] assert set(contract["models"]) == set(REQUIRED_MODELS) assert ( - contract["drafts"]["Qwen/Qwen3-Coder-30B-A3B"] + contract["drafts"]["Qwen/Qwen3-Coder-30B-A3B-Instruct"] == "z-lab/Qwen3-Coder-30B-A3B-DFlash" ) assert (