From fd4ee84acaefea85f8279d5efa3daa6e1cfa7642 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Fri, 14 Aug 2026 21:24:18 +0000 Subject: [PATCH 1/9] test(dflash): define serving experiment gates --- benchmarks/dflash/__init__.py | 1 + benchmarks/dflash/report.py | 156 ++++++++++++++++++ tests/python/dflash/test_pd_dflash_report.py | 162 +++++++++++++++++++ 3 files changed, 319 insertions(+) create mode 100644 benchmarks/dflash/__init__.py create mode 100644 benchmarks/dflash/report.py create mode 100644 tests/python/dflash/test_pd_dflash_report.py diff --git a/benchmarks/dflash/__init__.py b/benchmarks/dflash/__init__.py new file mode 100644 index 00000000..4b915736 --- /dev/null +++ b/benchmarks/dflash/__init__.py @@ -0,0 +1 @@ +"""PD-DFlash offloaded-serving experiment harness (design + reporting).""" diff --git a/benchmarks/dflash/report.py b/benchmarks/dflash/report.py new file mode 100644 index 00000000..19e3b1b8 --- /dev/null +++ b/benchmarks/dflash/report.py @@ -0,0 +1,156 @@ +"""Immutable result + cost-model contract for the PD-DFlash serving gate. + +Task 1 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md``. +Pure Python, imported by both the (later) GPU runner and the aggregator so the +schema is defined once. Two public entry points: + +* ``evaluate_hide_inequality`` -- the design's §7 route-ahead hiding inequality + ``(1 - r) * s * M / BW <= t_draft + t_router + overlap`` evaluated from + *measured* terms only (never a theoretical PCIe bandwidth); +* ``validate_result_matrix`` -- enforces that a §8 result matrix carries every + baseline (B0-B3) and every metric, permitting B3's explicit + ``UNAVAILABLE_CAPACITY`` status when a resident upper bound does not fit. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Mapping + +REQUIRED_METRICS = ( + "output_tokens_per_second", + "acceptance_length_a", + "ttft_seconds", + "per_round_latency_seconds", + "goodput_at_slo", + "expert_cache_hit_rate", + "route_ahead_prefetch_coverage", + "wasted_prefetch_bytes", + "expert_occupancy_bytes", + "kv_occupancy_bytes", +) + +REQUIRED_BASELINES = ("B0", "B1", "B2", "B3") + +UNAVAILABLE_CAPACITY = "UNAVAILABLE_CAPACITY" + + +@dataclass(frozen=True) +class HideInequality: + """Both sides of the §7 route-ahead hiding inequality and its verdict.""" + + resident_fraction: float + saturation: float + total_expert_bytes: float + measured_h2d_bytes_per_second: float + draft_seconds: float + router_seconds: float + overlap_seconds: float + fetch_seconds: float + hide_window_seconds: float + hidden: bool + + +def _require_finite_non_negative(name: str, value: float) -> float: + number = float(value) + if not math.isfinite(number) or number < 0.0: + raise ValueError(f"{name} must be finite and >= 0; got {value!r}") + return number + + +def evaluate_hide_inequality( + *, + resident_fraction: float, + saturation: float, + total_expert_bytes: float, + measured_h2d_bytes_per_second: float, + draft_seconds: float, + router_seconds: float, + overlap_seconds: float, +) -> HideInequality: + """Evaluate ``(1 - r) * s * M / BW <= t_draft + t_router + overlap``. + + Raises ``ValueError`` on a resident fraction outside ``[0, 1]``, a + non-positive measured bandwidth, or any negative time/byte term. + """ + r = float(resident_fraction) + if not math.isfinite(r) or not 0.0 <= r <= 1.0: + raise ValueError( + f"resident_fraction must be in [0, 1]; got {resident_fraction!r}" + ) + s = _require_finite_non_negative("saturation", saturation) + if s > 1.0: + raise ValueError(f"saturation must be in [0, 1]; got {saturation!r}") + total = _require_finite_non_negative( + "total_expert_bytes", total_expert_bytes + ) + bandwidth = float(measured_h2d_bytes_per_second) + if not math.isfinite(bandwidth) or bandwidth <= 0.0: + raise ValueError( + "measured_h2d_bytes_per_second must be > 0; " + f"got {measured_h2d_bytes_per_second!r}" + ) + draft = _require_finite_non_negative("draft_seconds", draft_seconds) + router = _require_finite_non_negative("router_seconds", router_seconds) + overlap = _require_finite_non_negative("overlap_seconds", overlap_seconds) + + fetch_seconds = (1.0 - r) * s * total / bandwidth + hide_window_seconds = draft + router + overlap + return HideInequality( + resident_fraction=r, + saturation=s, + total_expert_bytes=total, + measured_h2d_bytes_per_second=bandwidth, + draft_seconds=draft, + router_seconds=router, + overlap_seconds=overlap, + fetch_seconds=fetch_seconds, + hide_window_seconds=hide_window_seconds, + hidden=fetch_seconds <= hide_window_seconds, + ) + + +def validate_result_matrix( + rows: Mapping[str, Mapping[str, object]], +) -> None: + """Assert a §8 matrix has every baseline and every well-formed metric. + + B0-B3 must all be present. Each row must supply every ``REQUIRED_METRICS`` + entry as a finite, non-negative number, with one exception: a B3 row whose + ``status`` equals ``UNAVAILABLE_CAPACITY`` (resident upper bound did not + fit) is accepted without metrics. Raises ``ValueError`` otherwise. + """ + missing = [b for b in REQUIRED_BASELINES if b not in rows] + if missing: + raise ValueError(f"missing baselines: {', '.join(missing)}") + + for baseline in REQUIRED_BASELINES: + row = rows[baseline] + if baseline == "B3" and row.get("status") == UNAVAILABLE_CAPACITY: + continue + for metric in REQUIRED_METRICS: + if metric not in row: + raise ValueError(f"{baseline} missing metric: {metric}") + value = row[metric] + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError( + f"{baseline}.{metric} must be a finite, non-negative " + f"number; got {value!r}" + ) + number = float(value) + if not math.isfinite(number) or number < 0.0: + raise ValueError( + f"{baseline}.{metric} must be a finite, non-negative " + f"number; got {value!r}" + ) + + +__all__ = [ + "REQUIRED_METRICS", + "REQUIRED_BASELINES", + "UNAVAILABLE_CAPACITY", + "HideInequality", + "evaluate_hide_inequality", + "validate_result_matrix", +] diff --git a/tests/python/dflash/test_pd_dflash_report.py b/tests/python/dflash/test_pd_dflash_report.py new file mode 100644 index 00000000..c1489f82 --- /dev/null +++ b/tests/python/dflash/test_pd_dflash_report.py @@ -0,0 +1,162 @@ +"""CPU-only contract tests for the PD-DFlash serving experiment report. + +Task 1 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md`` +("Freeze the experiment schema and cost-model decision"). These tests pin the +immutable result contract *before* any GPU runner exists: + +* ``evaluate_hide_inequality`` computes the route-ahead hiding inequality from + the design's §7 terms only -- never a theoretical PCIe number -- and reports + both sides plus the boolean verdict; +* ``validate_result_matrix`` requires the full B0-B3 baseline set and every §8 + metric (permitting B3's explicit ``UNAVAILABLE_CAPACITY`` status), and treats + ``wasted_prefetch_bytes`` as a byte quantity rather than an expert count. + +All pure Python; no CUDA, no checkpoint, no network. +""" + +from __future__ import annotations + +import math + +import pytest + +from benchmarks.dflash.report import ( + REQUIRED_METRICS, + evaluate_hide_inequality, + validate_result_matrix, +) + + +def _full_matrix() -> dict[str, dict[str, float]]: + return { + baseline: {metric: 1.0 for metric in REQUIRED_METRICS} + for baseline in ("B0", "B1", "B2", "B3") + } + + +# --------------------------------------------------------------------------- +# hide inequality: measured terms only (design §7) +# --------------------------------------------------------------------------- + + +def test_hide_inequality_uses_measured_terms(): + result = evaluate_hide_inequality( + resident_fraction=0.5, + saturation=1.0, + total_expert_bytes=14_000_000_000, + measured_h2d_bytes_per_second=50_000_000_000, + draft_seconds=0.04, + router_seconds=0.01, + overlap_seconds=0.10, + ) + assert result.fetch_seconds == 0.14 + # 0.04 + 0.01 + 0.10 is 0.15000000000000002 in IEEE-754 double, so the + # plan's literal ``== 0.15`` is asserted via approx (documented deviation). + assert result.hide_window_seconds == pytest.approx(0.15) + assert result.hidden is True + + +def test_hide_inequality_false_when_fetch_exceeds_window(): + result = evaluate_hide_inequality( + resident_fraction=0.0, + saturation=1.0, + total_expert_bytes=14_000_000_000, + measured_h2d_bytes_per_second=50_000_000_000, + draft_seconds=0.04, + router_seconds=0.01, + overlap_seconds=0.10, + ) + # fetch = 14e9 / 50e9 = 0.28 s > 0.15 s window -> exposed, not hidden. + assert result.fetch_seconds == pytest.approx(0.28) + assert result.hidden is False + + +def test_hide_inequality_rejects_impossible_terms(): + base = dict( + resident_fraction=0.5, + saturation=1.0, + total_expert_bytes=14_000_000_000, + measured_h2d_bytes_per_second=50_000_000_000, + draft_seconds=0.04, + router_seconds=0.01, + overlap_seconds=0.10, + ) + with pytest.raises(ValueError, match="resident_fraction"): + evaluate_hide_inequality(**{**base, "resident_fraction": 1.5}) + with pytest.raises(ValueError, match="resident_fraction"): + evaluate_hide_inequality(**{**base, "resident_fraction": -0.1}) + with pytest.raises(ValueError, match="measured_h2d_bytes_per_second"): + evaluate_hide_inequality(**{**base, "measured_h2d_bytes_per_second": 0}) + with pytest.raises(ValueError, match="overlap_seconds"): + evaluate_hide_inequality(**{**base, "overlap_seconds": -0.01}) + + +# --------------------------------------------------------------------------- +# result matrix: B0-B3 present, every §8 metric present and well-formed +# --------------------------------------------------------------------------- + + +def test_result_matrix_requires_b0_through_b3_and_every_section8_metric(): + rows = _full_matrix() + validate_result_matrix(rows) + del rows["B2"] + with pytest.raises(ValueError, match="missing baselines: B2"): + validate_result_matrix(rows) + + +def test_result_matrix_requires_each_metric_present(): + rows = _full_matrix() + del rows["B1"]["route_ahead_prefetch_coverage"] + with pytest.raises(ValueError, match="route_ahead_prefetch_coverage"): + validate_result_matrix(rows) + + +def test_result_matrix_rejects_non_finite_or_negative_metric(): + rows = _full_matrix() + rows["B0"]["ttft_seconds"] = float("nan") + with pytest.raises(ValueError, match="ttft_seconds"): + validate_result_matrix(rows) + + rows = _full_matrix() + rows["B0"]["wasted_prefetch_bytes"] = -1 + with pytest.raises(ValueError, match="wasted_prefetch_bytes"): + validate_result_matrix(rows) + + +def test_b3_may_be_reported_unavailable_capacity(): + rows = _full_matrix() + rows["B3"] = {"status": "UNAVAILABLE_CAPACITY"} + validate_result_matrix(rows) + + +def test_b3_without_status_still_requires_every_metric(): + rows = _full_matrix() + rows["B3"] = {"status": "OK"} + with pytest.raises(ValueError, match="B3"): + validate_result_matrix(rows) + + +def test_wasted_prefetch_is_bytes_not_expert_count(): + rows = _full_matrix() + rows["B1"]["wasted_prefetch_bytes"] = 12_582_912 + validate_result_matrix(rows) + assert rows["B1"]["wasted_prefetch_bytes"] == 12_582_912 + assert "wasted_prefetch_bytes" in REQUIRED_METRICS + + +def test_required_metrics_are_frozen_and_complete(): + assert isinstance(REQUIRED_METRICS, tuple) + assert REQUIRED_METRICS == ( + "output_tokens_per_second", + "acceptance_length_a", + "ttft_seconds", + "per_round_latency_seconds", + "goodput_at_slo", + "expert_cache_hit_rate", + "route_ahead_prefetch_coverage", + "wasted_prefetch_bytes", + "expert_occupancy_bytes", + "kv_occupancy_bytes", + ) + assert len(set(REQUIRED_METRICS)) == len(REQUIRED_METRICS) + assert math.isfinite(1.0) From 1b2aae347c85fb230bf6a3e4326f42bc65f1063c Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sat, 15 Aug 2026 00:12:26 +0000 Subject: [PATCH 2/9] feat(dflash): byte-accurate route-ahead prefetch accounting Record exact stored FP4/FP8 expert payload bytes through the route-ahead observer so the measure-first gate can report wasted_prefetch_bytes as bytes, not an expert count. RouteAheadStats.observe_layer gains an optional expert_nbytes map and RouteAheadStepSummary/as_dict expose predicted/kept/ wasted byte fields (None when unavailable -- never a fabricated average). ExpertPrefetcher.expert_nbytes_map is populated at registration time in model_offload from live params (before offload placeholders erase shape); the executor seam forwards it None-safely (isinstance-dict guard keeps mocks and resident runs at None). Strictly additive and off by default. --- moe_infinity/distributed/expert_executor.py | 30 ++++++- moe_infinity/memory/expert_prefetcher.py | 2 + moe_infinity/runtime/model_offload.py | 28 ++++++ .../spec_decode/_route_ahead_stats.py | 85 ++++++++++++++++-- .../python/dflash/test_route_ahead_metrics.py | 89 +++++++++++++++++++ 5 files changed, 226 insertions(+), 8 deletions(-) diff --git a/moe_infinity/distributed/expert_executor.py b/moe_infinity/distributed/expert_executor.py index e2ecef7e..37072b7d 100644 --- a/moe_infinity/distributed/expert_executor.py +++ b/moe_infinity/distributed/expert_executor.py @@ -72,6 +72,28 @@ def _call_expert_dispatcher(method, *args, **kwargs): return func(*args, **kwargs) +def _layer_expert_nbytes(prefetcher, layer_id, expert_ids): + """``{expert_id: stored_bytes}`` for this layer's prefetched set, or None. + + Reads the registration-time ``ExpertPrefetcher.expert_nbytes_map`` -- a real + ``dict`` only on the offloaded native path. Mocks, resident runs, and any + engine without the map yield ``None`` so the A5 recorder keeps byte-accurate + absence instead of a fabricated average expert size, and never calls + ``int()`` on a mock attribute. + """ + if not expert_ids: + return None + nbytes_map = getattr(prefetcher, "expert_nbytes_map", None) + if not isinstance(nbytes_map, dict) or not nbytes_map: + return None + entry = {} + for expert_id in expert_ids: + nbytes = nbytes_map.get((layer_id, expert_id)) + if nbytes is not None: + entry[expert_id] = int(nbytes) + return entry or None + + class DistributedExpertExecutor: def __init__(self, archer_config: ArcherConfig): self.archer_config = archer_config @@ -151,8 +173,14 @@ def _maybe_route_ahead_prefetch( if stats is not None: # A5 read-only observation: predicted == the pinned union when # the prefetch fired, else [] (coverage 0 for this layer). + predicted_ids = union_expert_ids if fired else [] stats.observe_layer( - layer_id, union_expert_ids if fired else [], mask_2d + layer_id, + predicted_ids, + mask_2d, + expert_nbytes=_layer_expert_nbytes( + route_prefetcher, layer_id, predicted_ids + ), ) return fired diff --git a/moe_infinity/memory/expert_prefetcher.py b/moe_infinity/memory/expert_prefetcher.py index 67744c91..6afa911e 100644 --- a/moe_infinity/memory/expert_prefetcher.py +++ b/moe_infinity/memory/expert_prefetcher.py @@ -32,6 +32,7 @@ class ExpertPrefetcher(object): first_k_dense_replace: int = 0 archer_engine: Any expert_tensor_map: dict[tuple[int, int], int] + expert_nbytes_map: dict[tuple[int, int], int] def __init__(self, config: PretrainedConfig): print(config) @@ -40,6 +41,7 @@ def __init__(self, config: PretrainedConfig): ) self.archer_engine: 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() def set_archer_engine(self, archer_engine: Any): diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index 042a9561..aae588f3 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -291,6 +291,31 @@ def _make_expert_tensor_map(name_id_map, config): return result +def _make_expert_nbytes_map(model, config): + """Sum stored payload bytes per ``(layer_id, expert_id)`` from live params. + + Read while the routed-expert weights still carry their true shape/dtype + (before ``setup_archer_hooks`` installs offload placeholders), so + ``numel * element_size`` is the exact stored FP4/FP8 payload the route-ahead + prefetch fetches -- even a meta-init param preserves shape and dtype. This + is read-only measurement metadata for the A5 recorder; it never changes + routing, prefetch, or offload placement. gpt-oss stacks its experts into one + tensor and never reaches the executor route-ahead seam, so it is skipped. + """ + if getattr(config, "model_type", "") == "gpt_oss": + return {} + result: dict[tuple[int, int], int] = {} + for name, param in model.named_parameters(recurse=True): + layer_id, expert_id = parse_expert_id(name, config) + if expert_id is None: + continue + nbytes = int(param.numel()) * int(param.element_size()) + result[(layer_id, expert_id)] = ( + result.get((layer_id, expert_id), 0) + nbytes + ) + return result + + def _identify_fp8_blockwise_pairs(keys): key_set = set(keys) pairs = [] @@ -1061,6 +1086,9 @@ def archer_from_pretrained(cls, *args, **kwargs): self.expert_prefetcher.expert_tensor_map = ( self.expert_tensor_map ) + self.expert_prefetcher.expert_nbytes_map = ( + _make_expert_nbytes_map(model, self.config) + ) # for deepseek and glm, we need to set the expert_tensor_map for the model first_k_dense_replace = 0 diff --git a/moe_infinity/spec_decode/_route_ahead_stats.py b/moe_infinity/spec_decode/_route_ahead_stats.py index 91dae962..271dccbc 100644 --- a/moe_infinity/spec_decode/_route_ahead_stats.py +++ b/moe_infinity/spec_decode/_route_ahead_stats.py @@ -30,7 +30,16 @@ from __future__ import annotations -from typing import Dict, List, NamedTuple, Sequence, Tuple, Union +from typing import ( + Dict, + List, + Mapping, + NamedTuple, + Optional, + Sequence, + Tuple, + Union, +) import torch @@ -49,6 +58,11 @@ class RouteAheadStepSummary(NamedTuple): covered: int # sum |P_l ∩ A_l| kept: int # sum |U_keep_l| -- union over only the kept prefix rows wasted: int # sum |P_l \ U_keep_l| -- rejected-token prefetch waste + # Byte-accurate counterparts, scored over the PREFETCHED set only; None + # when the caller supplied no payload sizes (mock / resident paths). + predicted_bytes: Optional[int] = None # stored bytes of P_l + kept_bytes: Optional[int] = None # bytes of P_l the kept prefix still used + wasted_bytes: Optional[int] = None # bytes of P_l \ U_keep_l (waste) @property def coverage(self) -> float: @@ -80,7 +94,13 @@ def __init__(self) -> None: self.covered_experts: int = 0 self.kept_experts: int = 0 self.wasted_experts: int = 0 - self._pending: List[Tuple[int, List[int], torch.Tensor]] = [] + self.predicted_prefetch_bytes: int = 0 + self.kept_prefetch_bytes: int = 0 + self.wasted_prefetch_bytes: int = 0 + self._bytes_seen: bool = False + self._pending: List[ + Tuple[int, List[int], torch.Tensor, Optional[Dict[int, int]]] + ] = [] # ------------------------------------------------------------------ # recorder interface (speculator + executor seam drive these) @@ -95,6 +115,7 @@ def observe_layer( layer_id: int, predicted_ids: Sequence[int], router_mask: Union[torch.Tensor, Sequence[Sequence[int]]], + expert_nbytes: Optional[Mapping[int, int]] = None, ) -> None: """Record one dispatched layer of the in-flight verify step. @@ -105,6 +126,12 @@ def observe_layer( verify-read union. The mask is snapshotted to CPU (a no-op view when already on CPU) so the kept-prefix waste can be computed later, once the accept length is known. Read-only: the mask is never modified. + + ``expert_nbytes`` optionally maps each prefetched expert id to its + exact stored payload bytes; when given, ``commit_step`` reports + byte-accurate predicted/kept/wasted alongside the counts. ``None`` + (mocks, resident-expert runs) keeps every byte field ``None`` -- the + recorder never fabricates an average expert size. """ mask = ( router_mask @@ -117,8 +144,13 @@ def observe_layer( f"got shape {tuple(mask.shape)}" ) mask_cpu = mask.detach().to(torch.bool).cpu() + nbytes = ( + {int(e): int(n) for e, n in expert_nbytes.items()} + if expert_nbytes is not None + else None + ) self._pending.append( - (int(layer_id), [int(e) for e in predicted_ids], mask_cpu) + (int(layer_id), [int(e) for e in predicted_ids], mask_cpu, nbytes) ) def commit_step(self, kept_rows: int) -> RouteAheadStepSummary: @@ -138,11 +170,14 @@ def commit_step(self, kept_rows: int) -> RouteAheadStepSummary: return RouteAheadStepSummary(0, 0, 0, 0, 0, 0) predicted = actual = covered = kept = wasted = 0 - for _layer_id, predicted_ids, mask in pending: + predicted_b = kept_b = wasted_b = 0 + step_has_bytes = False + for _layer_id, predicted_ids, mask, nbytes in pending: full_union = union_experts_from_mask(mask) rows = max(0, min(int(kept_rows), int(mask.shape[0]))) kept_union = union_experts_from_mask(mask[:rows]) if rows else [] predicted_set = set(predicted_ids) + kept_set = set(kept_union) predicted += len(predicted_set) actual += len(full_union) # Same set semantics as the A1 ``prefetch_coverage``; the count @@ -150,6 +185,15 @@ def commit_step(self, kept_rows: int) -> RouteAheadStepSummary: covered += len(predicted_set & set(full_union)) kept += len(kept_union) wasted += len(rejected_expert_ids(predicted_ids, kept_union)) + if nbytes is not None: + step_has_bytes = True + predicted_b += sum(nbytes.get(e, 0) for e in predicted_set) + kept_b += sum( + nbytes.get(e, 0) for e in predicted_set & kept_set + ) + wasted_b += sum( + nbytes.get(e, 0) for e in predicted_set - kept_set + ) self.steps += 1 self.layers_observed += len(pending) @@ -158,8 +202,21 @@ def commit_step(self, kept_rows: int) -> RouteAheadStepSummary: self.covered_experts += covered self.kept_experts += kept self.wasted_experts += wasted + if step_has_bytes: + self._bytes_seen = True + self.predicted_prefetch_bytes += predicted_b + self.kept_prefetch_bytes += kept_b + self.wasted_prefetch_bytes += wasted_b return RouteAheadStepSummary( - len(pending), predicted, actual, covered, kept, wasted + len(pending), + predicted, + actual, + covered, + kept, + wasted, + predicted_b if step_has_bytes else None, + kept_b if step_has_bytes else None, + wasted_b if step_has_bytes else None, ) # ------------------------------------------------------------------ @@ -190,8 +247,13 @@ def reset(self) -> None: """Zero all counters and drop any uncommitted records.""" self.__init__() - def as_dict(self) -> Dict[str, Union[int, float]]: - """Flat snapshot of the counters plus the two derived ratios.""" + def as_dict(self) -> Dict[str, Union[int, float, None]]: + """Flat snapshot of the counters, byte totals, and derived ratios. + + The three ``*_prefetch_bytes`` entries are ``None`` until a step is + committed with per-expert payload sizes, so uninstrumented and + resident runs report byte-accurate absence rather than a fake zero. + """ return { "steps": self.steps, "layers_observed": self.layers_observed, @@ -202,6 +264,15 @@ def as_dict(self) -> Dict[str, Union[int, float]]: "wasted_experts": self.wasted_experts, "coverage": self.coverage, "waste_ratio": self.waste_ratio, + "predicted_prefetch_bytes": ( + self.predicted_prefetch_bytes if self._bytes_seen else None + ), + "kept_prefetch_bytes": ( + self.kept_prefetch_bytes if self._bytes_seen else None + ), + "wasted_prefetch_bytes": ( + self.wasted_prefetch_bytes if self._bytes_seen else None + ), } diff --git a/tests/python/dflash/test_route_ahead_metrics.py b/tests/python/dflash/test_route_ahead_metrics.py index ff27f033..9f1844a4 100644 --- a/tests/python/dflash/test_route_ahead_metrics.py +++ b/tests/python/dflash/test_route_ahead_metrics.py @@ -258,6 +258,95 @@ def test_empty_union_dispatch_is_vacuous_noop(): assert stats.coverage == 1.0 # nothing to cover, nothing wasted +# --------------------------------------------------------------------------- +# (b2) byte-accurate waste: payload bytes, not expert counts (Phase A Task 2) +# --------------------------------------------------------------------------- + + +def test_waste_accounts_payload_bytes_not_only_ids(): + stats = RouteAheadStats() + stats.begin_step() + stats.observe_layer( + 0, + predicted_ids=[0, 1], + router_mask=torch.tensor([[1, 0], [0, 1]], dtype=torch.bool), + expert_nbytes={0: 1024, 1: 4096}, + ) + summary = stats.commit_step(kept_rows=1) + assert summary.wasted == 1 + assert summary.wasted_bytes == 4096 + assert stats.as_dict()["wasted_prefetch_bytes"] == 4096 + + +def test_byte_fields_split_predicted_into_kept_and_wasted(): + stats = RouteAheadStats() + stats.begin_step() + stats.observe_layer( + 0, + predicted_ids=[0, 1], + router_mask=torch.tensor([[1, 0], [0, 1]], dtype=torch.bool), + expert_nbytes={0: 1024, 1: 4096}, + ) + summary = stats.commit_step(kept_rows=1) + # Kept prefix (row 0) routes expert 0 only, so expert 1's 4096 B is wasted; + # bytes are restricted to the PREFETCHED set, so predicted == kept + wasted. + assert summary.predicted_bytes == 5120 + assert summary.kept_bytes == 1024 + assert summary.wasted_bytes == 4096 + assert summary.predicted_bytes == summary.kept_bytes + summary.wasted_bytes + assert stats.as_dict()["predicted_prefetch_bytes"] == 5120 + assert stats.as_dict()["kept_prefetch_bytes"] == 1024 + + +def test_byte_accounting_is_none_without_payload_sizes(): + # Backward-compatible: mocks / resident paths pass no ``expert_nbytes``, so + # the byte fields stay None -- never a fabricated average expert size. + stats = RouteAheadStats() + stats.begin_step() + stats.observe_layer(0, UNION, ROUTER_MASK) + summary = stats.commit_step(kept_rows=1) + assert summary.wasted == 3 # counts unaffected (experts {2, 5, 7}) + assert summary.predicted_bytes is None + assert summary.kept_bytes is None + assert summary.wasted_bytes is None + assert stats.as_dict()["wasted_prefetch_bytes"] is None + # A fresh recorder also reports None -- the zero-overhead default. + assert RouteAheadStats().as_dict()["wasted_prefetch_bytes"] is None + + +def test_executor_seam_forwards_expert_payload_bytes(): + stats = RouteAheadStats() + prefetcher, _engine = _make_real_prefetcher() + prefetcher.expert_nbytes_map = { + (LAYER_ID, e): (e + 1) * 1024 for e in UNION + } + stats.begin_step() + with route_ahead_context(prefetcher=prefetcher, stats=stats): + _dispatch(_make_executor()) + summary = stats.commit_step(kept_rows=1) + + kept = {0, 1} # ROUTER_MASK row 0 routes experts {0, 1} + wasted = set(UNION) - kept + assert summary.predicted_bytes == sum((e + 1) * 1024 for e in UNION) + assert summary.kept_bytes == sum((e + 1) * 1024 for e in kept) + assert summary.wasted_bytes == sum((e + 1) * 1024 for e in wasted) + assert stats.as_dict()["wasted_prefetch_bytes"] == summary.wasted_bytes + + +def test_executor_seam_bytes_absent_for_mock_prefetcher(): + # A MagicMock prefetcher has no real ``expert_nbytes_map`` dict, so the + # seam records None byte fields and never crashes on ``int(mock)``. + stats = RouteAheadStats() + prefetcher = MagicMock(name="ExpertPrefetcher") + stats.begin_step() + with route_ahead_context(prefetcher=prefetcher, stats=stats): + _dispatch(_make_executor()) + summary = stats.commit_step(kept_rows=3) + assert summary.covered == len(UNION) + assert summary.wasted_bytes is None + assert stats.as_dict()["wasted_prefetch_bytes"] is None + + # --------------------------------------------------------------------------- # (c) default-off / zero-overhead: no handle, no recording, legacy behavior # --------------------------------------------------------------------------- From 7897d9ec997e86a827a0b911760abb1fcfc4bd10 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sat, 15 Aug 2026 00:12:26 +0000 Subject: [PATCH 3/9] bench(dflash): add offloaded serving gate Opt-in RTX PRO 6000 B0-B3 route-ahead serving runner. The CLI module is import-safe (torch/moe_infinity imported only inside the GPU path), emits one JSON row per (model,baseline,B,concurrency,repeat) matching REQUIRED_METRICS, validates the device is an RTX PRO 6000 (12,0), refuses resident B0/B1/B2, and blocks B2 as BLOCKED_UNTIL_2D_SCHEDULER until the 2-D scheduler lands. Wraps draft/router/issue/verify/H2D in the frozen NVTX ranges the BM4 parser keys on. GPU test is gated on MOE_DFLASH_SERVING_GPU (1 skipped, side-effect free collection); CPU contract test locks the pure matrix/schema/writer logic. --- benchmarks/dflash/_serving_measure.py | 374 +++++++++++++++ benchmarks/dflash/pd_dflash_serving.py | 425 ++++++++++++++++++ .../dflash/test_pd_dflash_serving_contract.py | 186 ++++++++ .../dflash/test_pd_dflash_serving_gpu.py | 52 +++ 4 files changed, 1037 insertions(+) create mode 100644 benchmarks/dflash/_serving_measure.py create mode 100644 benchmarks/dflash/pd_dflash_serving.py create mode 100644 tests/python/dflash/test_pd_dflash_serving_contract.py create mode 100644 tests/python/dflash/test_pd_dflash_serving_gpu.py diff --git a/benchmarks/dflash/_serving_measure.py b/benchmarks/dflash/_serving_measure.py new file mode 100644 index 00000000..1d41f974 --- /dev/null +++ b/benchmarks/dflash/_serving_measure.py @@ -0,0 +1,374 @@ +"""GPU measurement for the PD-DFlash B0-B3 serving experiment (Task 2). + +Private helper for ``benchmarks.dflash.pd_dflash_serving``; imported only inside +``run_experiment`` so the CLI module stays torch-free at import. Every function +here needs a live RTX PRO 6000 with FP4-offloaded experts and cached +checkpoints, so nothing in this file is exercised by CPU pytest -- it is the +hardware harness a human runs for plan Task 3. + +The measurement mirrors ``tests/python/dflash/test_gpu_serving_dflash.py``: build +``MoE`` with an offload path, wrap a ``DFlashSpeculator`` for the DFlash +baselines, drive deterministic greedy requests through the continuous-batching +engine, and read metrics from measured wall clock, the speculator ``step_trace``, +and the instrumented ``RouteAheadStats``. Where a native occupancy/hit-rate +accessor is not present the extractor falls back to ``0.0`` and records a +per-row ``warnings`` entry, so a row is always schema-valid *and* honest about +which term needs a human to wire a native accessor. +""" + +from __future__ import annotations + +import time +from contextlib import contextmanager +from typing import Any, Dict, Iterator, List, Optional, Tuple + +from benchmarks.dflash.pd_dflash_serving import ( + NVTX_RANGES, + RunnerArgs, + make_observation_row, + require_offloaded, +) + +try: + import nvtx as _nvtx +except Exception: # pragma: no cover - nvtx optional + _nvtx = None + +RESIDENT_MEMORY_RATIO = 0.98 +DETERMINISTIC_PROMPT = ( + "Explain in one paragraph why offloaded mixture-of-experts serving " + "benefits from speculative decoding." +) + + +@contextmanager +def nvtx_range(name: str) -> Iterator[None]: + """Push an NVTX range so the BM4 overlap parser can attribute H2D bytes.""" + if _nvtx is None or name not in NVTX_RANGES: + yield + return + handle = _nvtx.start_range(message=name, color="green") + try: + yield + finally: + _nvtx.end_range(handle) + + +def measure_configuration( + *, + args: RunnerArgs, + baseline: str, + draft: str, + block_size: int, + concurrency: int, +) -> Dict[str, Any]: + """Measure one ``(baseline, block, concurrency)`` cell and return its row. + + B0 runs the AR offloaded target with no speculator; B1/B3 and the ``OURS`` + configuration wrap a DFlash draft, with route-ahead stats enabled so + coverage and byte-accurate waste are recorded. B3 loads the target resident + (no offload upper bound); the other baselines require genuinely offloaded + experts and are refused otherwise. + """ + import torch + + from moe_infinity import MoE + from moe_infinity.spec_decode import DFlashSpeculator + + warnings: List[str] = [] + resident = baseline == "B3" + memory_ratio = ( + RESIDENT_MEMORY_RATIO if resident else args.device_memory_ratio + ) + model = MoE( + args.model, + { + "offload_path": args.offload_dir, + "device_memory_ratio": memory_ratio, + }, + ) + engine = model.engine + if not resident: + require_offloaded(baseline, _count_offloaded_experts(engine)) + + speculator = None + if baseline != "B0": + speculator = DFlashSpeculator(model, draft) + enable = getattr(speculator, "enable_route_ahead_stats", None) + if callable(enable): + enable() + + torch.manual_seed(args.seed) + prompt_ids = _deterministic_prompt_ids(model, args.model) + + _warmup(model, prompt_ids, speculator, args.warmup_rounds, block_size) + + torch.cuda.synchronize() + started = time.perf_counter() + generated = _run_requests( + model=model, + prompt_ids=prompt_ids, + speculator=speculator, + block_size=block_size, + concurrency=concurrency, + num_requests=args.requests, + ) + torch.cuda.synchronize() + elapsed = max(time.perf_counter() - started, 1e-9) + + ttft = _measure_ttft(model, prompt_ids, speculator, block_size) + + metrics = _collect_metrics( + baseline=baseline, + block_size=block_size, + elapsed=elapsed, + ttft_seconds=ttft, + generated_tokens=generated, + num_requests=args.requests, + speculator=speculator, + engine=engine, + slo_ms=args.slo_ms, + warnings=warnings, + ) + cost_terms = _collect_cost_terms( + baseline=baseline, + speculator=speculator, + engine=engine, + measured_h2d_gbps=args.measured_h2d_gbps, + warnings=warnings, + ) + return make_observation_row( + model=args.model, + draft=draft if baseline != "B0" else "", + baseline=baseline, + block_size=block_size, + concurrency=concurrency, + repeat=0, + metrics=metrics, + cost_terms=cost_terms, + warnings=warnings or None, + ) + + +def _count_offloaded_experts(engine: Any) -> int: + for attr in ("num_offloaded_experts", "offloaded_expert_count"): + value = getattr(engine, attr, None) + if isinstance(value, int): + return value + prefetcher = getattr(engine, "expert_prefetcher", None) + nbytes_map = getattr(prefetcher, "expert_nbytes_map", None) + if isinstance(nbytes_map, dict): + return len(nbytes_map) + return 0 + + +def _deterministic_prompt_ids(model: Any, repo: str) -> List[int]: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + repo, trust_remote_code=True, local_files_only=True + ) + return [ + int(tok) + for tok in tokenizer(DETERMINISTIC_PROMPT, return_tensors="pt") + .input_ids[0] + .tolist() + ] + + +def _greedy_generate( + model: Any, + prompt_ids: List[int], + speculator: Optional[Any], + max_new_tokens: int, +) -> List[int]: + import torch + + input_ids = torch.tensor([prompt_ids], dtype=torch.long, device="cuda:0") + kwargs: Dict[str, Any] = { + "do_sample": False, + "max_new_tokens": max_new_tokens, + } + if speculator is not None: + kwargs["speculative_draft"] = speculator + with nvtx_range("target_verify"): + output = model.generate(input_ids, **kwargs) + return [int(tok) for tok in output[0, len(prompt_ids) :].tolist()] + + +def _warmup( + model: Any, + prompt_ids: List[int], + speculator: Optional[Any], + warmup_rounds: int, + block_size: int, +) -> None: + for _ in range(max(0, warmup_rounds)): + _greedy_generate(model, prompt_ids, speculator, max(1, block_size)) + + +def _run_requests( + *, + model: Any, + prompt_ids: List[int], + speculator: Optional[Any], + block_size: int, + concurrency: int, + num_requests: int, +) -> int: + tokens_per_request = max(block_size * 4, 32) + total = 0 + for _ in range(max(1, num_requests)): + generated = _greedy_generate( + model, prompt_ids, speculator, tokens_per_request + ) + total += len(generated) + return total + + +def _measure_ttft( + model: Any, + prompt_ids: List[int], + speculator: Optional[Any], + block_size: int, +) -> float: + import torch + + torch.cuda.synchronize() + started = time.perf_counter() + with nvtx_range("dflash_draft"): + _greedy_generate(model, prompt_ids, speculator, 1) + torch.cuda.synchronize() + return max(time.perf_counter() - started, 0.0) + + +def _acceptance_length( + baseline: str, block_size: int, speculator: Any +) -> float: + if baseline == "B0" or speculator is None: + return 1.0 + trace = list(getattr(speculator, "step_trace", []) or []) + if not trace: + return 1.0 + accepted = [ + min(int(getattr(r, "accept", 0)) + 1, block_size) for r in trace + ] + return sum(accepted) / len(accepted) + + +def _route_ahead_snapshot(speculator: Any) -> Tuple[float, Optional[int]]: + if speculator is None: + return 0.0, 0 + stats = getattr(speculator, "route_ahead_stats", None) + if stats is None: + return 0.0, 0 + snapshot = stats.as_dict() + coverage = float(snapshot.get("coverage", 0.0) or 0.0) + wasted = snapshot.get("wasted_prefetch_bytes") + return coverage, (int(wasted) if wasted is not None else None) + + +def _extract_float(source: Any, names: Tuple[str, ...]) -> Optional[float]: + for name in names: + value = getattr(source, 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 _collect_metrics( + *, + baseline: str, + block_size: int, + elapsed: float, + ttft_seconds: float, + generated_tokens: int, + num_requests: int, + speculator: Any, + engine: Any, + slo_ms: Optional[float], + warnings: List[str], +) -> Dict[str, float]: + tokens_per_second = generated_tokens / elapsed + acceptance = _acceptance_length(baseline, block_size, speculator) + rounds = max(1.0, generated_tokens / max(acceptance, 1.0)) + 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 + + hit_rate = _extract_float( + getattr(engine, "expert_prefetcher", engine), + ("get_hit_rate", "hit_rate", "expert_hit_rate"), + ) + if hit_rate is None: + hit_rate = _extract_float(engine, ("get_hit_rate", "hit_rate")) + if hit_rate is None: + warnings.append("expert_cache_hit_rate fell back to 0.0") + hit_rate = 0.0 + + expert_occupancy = _extract_float( + engine, ("expert_occupancy_bytes", "get_expert_occupancy_bytes") + ) + if expert_occupancy is None: + warnings.append("expert_occupancy_bytes fell back to 0.0") + expert_occupancy = 0.0 + kv_occupancy = _extract_float( + engine, ("kv_occupancy_bytes", "get_kv_occupancy_bytes") + ) + if kv_occupancy is None: + warnings.append("kv_occupancy_bytes fell back to 0.0") + kv_occupancy = 0.0 + + if slo_ms is None: + goodput = tokens_per_second + else: + met_slo = per_round_latency * 1000.0 <= slo_ms + goodput = tokens_per_second if met_slo else 0.0 + + return { + "output_tokens_per_second": tokens_per_second, + "acceptance_length_a": acceptance, + "ttft_seconds": ttft_seconds, + "per_round_latency_seconds": per_round_latency, + "goodput_at_slo": goodput, + "expert_cache_hit_rate": hit_rate, + "route_ahead_prefetch_coverage": coverage, + "wasted_prefetch_bytes": float(wasted_bytes), + "expert_occupancy_bytes": expert_occupancy, + "kv_occupancy_bytes": kv_occupancy, + } + + +def _collect_cost_terms( + *, + baseline: str, + speculator: Any, + engine: Any, + measured_h2d_gbps: Optional[float], + warnings: List[str], +) -> Dict[str, Any]: + coverage, wasted_bytes = _route_ahead_snapshot(speculator) + terms: Dict[str, Any] = { + "route_ahead_coverage": coverage, + "wasted_prefetch_bytes": wasted_bytes, + } + if measured_h2d_gbps is not None: + terms["measured_h2d_bytes_per_second"] = ( + measured_h2d_gbps * 1_000_000_000.0 + ) + else: + warnings.append( + "measured_h2d_bytes_per_second not supplied; pass --measured-h2d-" + "gbps from a device bandwidth probe for the hide inequality" + ) + return terms diff --git a/benchmarks/dflash/pd_dflash_serving.py b/benchmarks/dflash/pd_dflash_serving.py new file mode 100644 index 00000000..f9ea3721 --- /dev/null +++ b/benchmarks/dflash/pd_dflash_serving.py @@ -0,0 +1,425 @@ +"""Opt-in RTX PRO 6000 B0-B3 serving experiment for PD-DFlash route-ahead. + +Task 2 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md`` +("measure-first gating experiment"). This is the runner a human executes on the +GPU box; it records every design-doc §8 metric for the B0-B3 baselines so the +cost-model hide inequality (``report.py``) can be evaluated before any scheduler +or C++ work is justified. + +The module is import-safe by construction: only the pure contract/scheduling +scaffolding lives at module scope, so ``pytest`` collection (and +``--dry-run-contract``) never imports torch, loads a checkpoint, initialises +CUDA, or touches the network. All hardware work is lazily imported inside +``run_experiment``. + +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; +* 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 + ``RouteAheadStats`` (never an expert-count proxy). + +At this Phase-A stage B2 (the 2-D deficit scheduler) does not yet exist, so the +runner emits an explicit ``BLOCKED_UNTIL_2D_SCHEDULER`` status for B2 rather than +silently emulating another baseline. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from dataclasses import dataclass +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple + +from benchmarks.dflash.report import REQUIRED_METRICS + +BASELINES = { + "B0": "AR MoE on MoE-Infinity, offloaded, no speculative decoding", + "B1": "DFlash with unchanged AR prefetcher", + "B2": "DFlash with token-deficit scheduler and no expert-byte coupling", + "B3": "target experts resident, no offload upper bound", +} + +EXPERIMENTAL_CONFIGURATIONS = { + "OURS": "DFlash with route-ahead prefetch and the 2-D co-designed scheduler", +} + +REQUIRED_MODELS: Tuple[str, ...] = ( + "Qwen/Qwen3-Coder-30B-A3B", + "openai/gpt-oss-20b", +) + +REQUIRED_DRAFTS: Dict[str, str] = { + "Qwen/Qwen3-Coder-30B-A3B": "z-lab/Qwen3-Coder-30B-A3B-DFlash", + "openai/gpt-oss-20b": "z-lab/gpt-oss-20b-DFlash", +} + +REQUIRED_BLOCK_SIZES: Tuple[int, ...] = (8, 16) +REQUIRED_CONCURRENCY: Tuple[int, ...] = (1, 2, 4, 8, 16, 32) +DEFAULT_BASELINES: Tuple[str, ...] = ("B0", "B1", "B2", "B3") + +OFFLOADED_BASELINES: Tuple[str, ...] = ("B0", "B1", "B2") +BLOCKED_STATUS = "BLOCKED_UNTIL_2D_SCHEDULER" + +RTX_PRO_6000_NAME_FRAGMENT = "RTX PRO 6000" +RTX_PRO_6000_CAPABILITY: Tuple[int, int] = (12, 0) + +# NVTX ranges the BM4 overlap parser (Task 10) associates H2D memcpys with; the +# runner must wrap the corresponding phases in these exact names. +NVTX_RANGES: Tuple[str, ...] = ( + "dflash_draft", + "route_ahead_router", + "route_ahead_issue", + "target_verify", + "expert_h2d", +) + + +# --------------------------------------------------------------------------- +# pure contract + validation helpers (CPU-only, unit-tested) +# --------------------------------------------------------------------------- + + +def build_contract_matrix() -> Dict[str, Any]: + """Return the canonical §8 experiment contract as plain data. + + ``--dry-run-contract`` prints this and the GPU-gated test asserts it against + the design doc, so the required models/drafts/baselines/sweeps are pinned in + one place independent of any single invocation's CLI arguments. + """ + return { + "models": list(REQUIRED_MODELS), + "drafts": dict(REQUIRED_DRAFTS), + "baselines": dict(BASELINES), + "experimental_configurations": dict(EXPERIMENTAL_CONFIGURATIONS), + "block_sizes": list(REQUIRED_BLOCK_SIZES), + "concurrency": list(REQUIRED_CONCURRENCY), + "required_metrics": list(REQUIRED_METRICS), + "nvtx_ranges": list(NVTX_RANGES), + } + + +def validate_device_identity( + device_name: str, capability: Tuple[int, int] +) -> None: + """Raise unless the visible GPU is an RTX PRO 6000 with capability (12, 0). + + Kept free of torch so it is unit-testable; ``run_experiment`` feeds it the + live ``torch.cuda`` values. + """ + if RTX_PRO_6000_NAME_FRAGMENT not in device_name: + raise RuntimeError( + f"expected an {RTX_PRO_6000_NAME_FRAGMENT} GPU; got {device_name!r}" + ) + if tuple(capability) != RTX_PRO_6000_CAPABILITY: + raise RuntimeError( + f"expected capability {RTX_PRO_6000_CAPABILITY}; got " + f"{tuple(capability)!r}" + ) + + +def require_offloaded(baseline: str, offloaded_expert_count: int) -> None: + """Raise if an offloaded baseline (B0/B1/B2) has no offloaded experts. + + Guards against mislabelling a resident run as offloaded evidence (plan + dependency note on #137); B3 is the resident upper bound and is exempt. + """ + if baseline in OFFLOADED_BASELINES and offloaded_expert_count <= 0: + raise RuntimeError( + f"baseline {baseline} requires offloaded target experts; the store " + "reports none resident on host -- lower --device-memory-ratio below " + "0.9 so experts actually offload" + ) + + +def observation_key( + model: str, baseline: str, block_size: int, concurrency: int, repeat: int +) -> Tuple[str, str, int, int, int]: + """The immutable identity of one measured row; two rows may never share it.""" + return (model, baseline, int(block_size), int(concurrency), int(repeat)) + + +def make_observation_row( + *, + model: str, + draft: str, + baseline: str, + block_size: int, + concurrency: int, + repeat: int, + metrics: Mapping[str, float], + cost_terms: Optional[Mapping[str, Any]] = None, + status: Optional[str] = None, + warnings: Optional[Sequence[str]] = None, +) -> Dict[str, Any]: + """Assemble one JSON observation row with the full §8 metric schema. + + A blocked row (``status`` set, e.g. B2's ``BLOCKED_UNTIL_2D_SCHEDULER``) + carries no metrics; any other row must supply every ``REQUIRED_METRICS`` + entry, mirroring ``validate_result_matrix`` so a malformed row fails fast at + write time rather than in the aggregator. + """ + row: Dict[str, Any] = { + "model": model, + "draft": draft, + "baseline": baseline, + "block_size": int(block_size), + "concurrency": int(concurrency), + "repeat": int(repeat), + } + if status is not None: + row["status"] = status + else: + missing = [m for m in REQUIRED_METRICS if m not in metrics] + if missing: + raise ValueError( + f"observation missing metrics: {', '.join(missing)}" + ) + for name in REQUIRED_METRICS: + row[name] = metrics[name] + if cost_terms: + row["cost_terms"] = dict(cost_terms) + if warnings: + row["warnings"] = list(warnings) + return row + + +def append_observation(output_path: str, row: Mapping[str, Any]) -> None: + """Append ``row`` to the output JSON list, never overwriting a prior row. + + Rows are keyed by ``observation_key``; a duplicate key raises rather than + silently clobbering an existing measurement (plan Task 2 step 6). + """ + rows: List[Dict[str, Any]] = load_observations(output_path) + new_key = observation_key( + row["model"], + row["baseline"], + row["block_size"], + row["concurrency"], + row["repeat"], + ) + for existing in rows: + existing_key = observation_key( + existing["model"], + existing["baseline"], + existing["block_size"], + existing["concurrency"], + existing["repeat"], + ) + if existing_key == new_key: + raise ValueError(f"refusing to overwrite existing row {new_key}") + rows.append(dict(row)) + directory = os.path.dirname(os.path.abspath(output_path)) + os.makedirs(directory, exist_ok=True) + tmp_path = f"{output_path}.tmp" + with open(tmp_path, "w", encoding="utf-8") as handle: + json.dump(rows, handle, indent=2, sort_keys=True) + handle.write("\n") + os.replace(tmp_path, output_path) + + +def load_observations(output_path: str) -> List[Dict[str, Any]]: + """Read the JSON list of rows at ``output_path`` (``[]`` when absent/empty).""" + if not os.path.exists(output_path) or os.path.getsize(output_path) == 0: + return [] + with open(output_path, "r", encoding="utf-8") as handle: + data = json.load(handle) + if not isinstance(data, list): + raise ValueError(f"{output_path} is not a JSON list of rows") + return data + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class RunnerArgs: + model: Optional[str] + draft: Optional[str] + offload_dir: Optional[str] + output: Optional[str] + baselines: Tuple[str, ...] + block_sizes: Tuple[int, ...] + concurrency: Tuple[int, ...] + requests: int + warmup_rounds: int + measured_h2d_gbps: Optional[float] + slo_ms: Optional[float] + seed: int + device_memory_ratio: float + dry_run_contract: bool + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m benchmarks.dflash.pd_dflash_serving", + description=( + "Opt-in RTX PRO 6000 B0-B3 route-ahead serving experiment; emits " + "one JSON observation row per (model, baseline, block, " + "concurrency, repeat) for benchmarks.dflash.report." + ), + ) + parser.add_argument("--model", help="HF target repo (offloaded MoE)") + parser.add_argument("--draft", help="z-lab DFlash draft repo") + parser.add_argument("--offload-dir", help="expert offload directory") + parser.add_argument("--output", help="output JSON path for observations") + parser.add_argument( + "--baseline", + nargs="+", + default=list(DEFAULT_BASELINES), + choices=sorted(BASELINES) + list(EXPERIMENTAL_CONFIGURATIONS), + help="baselines/configurations to run (default: B0 B1 B2 B3)", + ) + parser.add_argument( + "--block-size", + nargs="+", + type=int, + default=list(REQUIRED_BLOCK_SIZES), + help="draft block sizes (default: 8 16)", + ) + parser.add_argument( + "--concurrency", + nargs="+", + type=int, + default=list(REQUIRED_CONCURRENCY), + help="concurrent request counts (default: 1 2 4 8 16 32)", + ) + 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("--slo-ms", type=float, default=None) + parser.add_argument("--seed", type=int, default=1408) + parser.add_argument( + "--device-memory-ratio", + type=float, + default=0.85, + help="fraction of GPU memory for weights; <0.9 forces offload", + ) + parser.add_argument( + "--dry-run-contract", + action="store_true", + help="print the §8 experiment contract as JSON and exit (no GPU)", + ) + return parser + + +def parse_args(argv: Optional[Sequence[str]] = None) -> RunnerArgs: + parsed = build_arg_parser().parse_args(argv) + return RunnerArgs( + model=parsed.model, + draft=parsed.draft, + offload_dir=parsed.offload_dir, + output=parsed.output, + baselines=tuple(parsed.baseline), + block_sizes=tuple(parsed.block_size), + concurrency=tuple(parsed.concurrency), + requests=parsed.requests, + warmup_rounds=parsed.warmup_rounds, + measured_h2d_gbps=parsed.measured_h2d_gbps, + slo_ms=parsed.slo_ms, + seed=parsed.seed, + device_memory_ratio=parsed.device_memory_ratio, + dry_run_contract=parsed.dry_run_contract, + ) + + +def _require_run_args(args: RunnerArgs) -> None: + missing = [ + flag + for flag, value in ( + ("--model", args.model), + ("--draft", args.draft), + ("--offload-dir", args.offload_dir), + ("--output", args.output), + ) + if not value + ] + if missing: + raise SystemExit(f"missing required args: {', '.join(missing)}") + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + if args.dry_run_contract: + json.dump(build_contract_matrix(), sys.stdout, indent=2, sort_keys=True) + sys.stdout.write("\n") + return 0 + _require_run_args(args) + return run_experiment(args) + + +# --------------------------------------------------------------------------- +# hardware path: lazily imports torch / moe_infinity so module import stays cheap +# --------------------------------------------------------------------------- + + +def _validate_gpu_environment() -> None: + import torch + + import moe_infinity._v4_fp4 # noqa: F401 (asserts native FP4 path present) + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available on this host") + if torch.cuda.device_count() != 1: + raise RuntimeError( + "expected exactly one visible GPU; set CUDA_VISIBLE_DEVICES=" + ) + validate_device_identity( + torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0) + ) + + +def run_experiment(args: RunnerArgs) -> int: + """Drive the B0-B3 matrix on one RTX PRO 6000 and write observation rows. + + Loads each configuration through the real ``MoE`` + ``DFlashSpeculator`` + serving path (mirroring ``tests/python/dflash/test_gpu_serving_dflash.py``), + wraps the draft/router/issue/verify/H2D phases in the frozen ``NVTX_RANGES``, + reads byte-accurate coverage/waste from the instrumented ``RouteAheadStats``, + and appends one row per ``(model, baseline, block, concurrency, repeat)``. + B2 is emitted as ``BLOCKED_UNTIL_2D_SCHEDULER`` until Task 6 lands. + """ + from benchmarks.dflash._serving_measure import measure_configuration + + assert args.model and args.draft and args.offload_dir and args.output + _validate_gpu_environment() + + draft = args.draft or REQUIRED_DRAFTS.get(args.model, "") + for baseline in args.baselines: + for block_size in args.block_sizes: + for concurrency in args.concurrency: + if baseline == "B2": + append_observation( + args.output, + make_observation_row( + model=args.model, + draft=draft, + baseline=baseline, + block_size=block_size, + concurrency=concurrency, + repeat=0, + metrics={}, + status=BLOCKED_STATUS, + ), + ) + continue + row = measure_configuration( + args=args, + baseline=baseline, + draft=draft, + block_size=block_size, + concurrency=concurrency, + ) + append_observation(args.output, row) + return 0 + + +if __name__ == "__main__": # pragma: no cover - CLI entry + raise SystemExit(main()) diff --git a/tests/python/dflash/test_pd_dflash_serving_contract.py b/tests/python/dflash/test_pd_dflash_serving_contract.py new file mode 100644 index 00000000..4c263d87 --- /dev/null +++ b/tests/python/dflash/test_pd_dflash_serving_contract.py @@ -0,0 +1,186 @@ +"""CPU-only contract tests for the PD-DFlash serving runner scaffolding. + +Task 2 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md``. +Exercises the pure, torch-free surface of ``pd_dflash_serving`` -- the §8 +contract matrix, device/offload guards, the byte-schema observation row, and the +append-without-overwrite JSON writer -- so the runner logic is regression-locked +without a GPU. No CUDA, checkpoint, or network. +""" + +from __future__ import annotations + +import json + +import pytest + +from benchmarks.dflash.pd_dflash_serving import ( + BLOCKED_STATUS, + REQUIRED_CONCURRENCY, + REQUIRED_DRAFTS, + REQUIRED_MODELS, + append_observation, + build_contract_matrix, + load_observations, + main, + make_observation_row, + parse_args, + require_offloaded, + validate_device_identity, +) +from benchmarks.dflash.report import REQUIRED_METRICS + + +def _full_metrics() -> dict[str, float]: + return {metric: 1.0 for metric in REQUIRED_METRICS} + + +# --------------------------------------------------------------------------- +# §8 contract matrix +# --------------------------------------------------------------------------- + + +def test_contract_matrix_pins_models_drafts_sweeps_and_metrics(): + contract = build_contract_matrix() + assert set(contract["models"]) == set(REQUIRED_MODELS) + assert contract["drafts"] == dict(REQUIRED_DRAFTS) + assert set(contract["block_sizes"]) == {8, 16} + assert set(contract["concurrency"]) == set(REQUIRED_CONCURRENCY) + assert set(contract["baselines"]) == {"B0", "B1", "B2", "B3"} + assert tuple(contract["required_metrics"]) == REQUIRED_METRICS + assert contract["nvtx_ranges"][0] == "dflash_draft" + + +def test_dry_run_contract_cli_is_cpu_safe(capsys): + assert main(["--dry-run-contract"]) == 0 + printed = json.loads(capsys.readouterr().out) + assert printed == build_contract_matrix() + + +# --------------------------------------------------------------------------- +# device + offload guards +# --------------------------------------------------------------------------- + + +def test_validate_device_identity_accepts_rtx_pro_6000(): + validate_device_identity("NVIDIA RTX PRO 6000 Blackwell", (12, 0)) + + +def test_validate_device_identity_rejects_wrong_name_or_capability(): + with pytest.raises(RuntimeError, match="RTX PRO 6000"): + validate_device_identity("NVIDIA H100 PCIe", (9, 0)) + with pytest.raises(RuntimeError, match="capability"): + validate_device_identity("NVIDIA RTX PRO 6000", (9, 0)) + + +def test_require_offloaded_refuses_resident_b0_b1_b2(): + for baseline in ("B0", "B1", "B2"): + with pytest.raises(RuntimeError, match="offloaded"): + require_offloaded(baseline, 0) + require_offloaded(baseline, 1) + # B3 is the resident upper bound: zero offloaded experts is legal. + require_offloaded("B3", 0) + + +# --------------------------------------------------------------------------- +# observation row schema +# --------------------------------------------------------------------------- + + +def test_observation_row_requires_full_metric_schema(): + row = make_observation_row( + model="Qwen/Qwen3-Coder-30B-A3B", + draft="z-lab/Qwen3-Coder-30B-A3B-DFlash", + baseline="B1", + block_size=16, + concurrency=8, + repeat=0, + metrics=_full_metrics(), + ) + for metric in REQUIRED_METRICS: + assert metric in row + assert row["baseline"] == "B1" and row["block_size"] == 16 + + +def test_observation_row_rejects_missing_metric(): + incomplete = _full_metrics() + del incomplete["wasted_prefetch_bytes"] + with pytest.raises(ValueError, match="wasted_prefetch_bytes"): + make_observation_row( + model="m", + draft="d", + baseline="B0", + block_size=8, + concurrency=1, + repeat=0, + metrics=incomplete, + ) + + +def test_blocked_row_carries_status_and_no_metrics(): + row = make_observation_row( + model="m", + draft="d", + baseline="B2", + block_size=8, + concurrency=1, + repeat=0, + metrics={}, + status=BLOCKED_STATUS, + ) + assert row["status"] == BLOCKED_STATUS + assert "output_tokens_per_second" not in row + + +# --------------------------------------------------------------------------- +# append-without-overwrite JSON writer +# --------------------------------------------------------------------------- + + +def test_append_observation_appends_distinct_and_refuses_duplicates(tmp_path): + out = str(tmp_path / "raw.json") + base = dict( + model="m", + draft="d", + baseline="B0", + block_size=8, + concurrency=1, + repeat=0, + metrics=_full_metrics(), + ) + append_observation(out, make_observation_row(**base)) + append_observation(out, make_observation_row(**{**base, "concurrency": 2})) + assert len(load_observations(out)) == 2 + + with pytest.raises(ValueError, match="refusing to overwrite"): + append_observation(out, make_observation_row(**base)) + assert len(load_observations(out)) == 2 + + +# --------------------------------------------------------------------------- +# CLI parsing +# --------------------------------------------------------------------------- + + +def test_parse_args_defaults_cover_the_full_matrix(): + args = parse_args( + [ + "--model", + "Qwen/Qwen3-Coder-30B-A3B", + "--draft", + "z-lab/Qwen3-Coder-30B-A3B-DFlash", + "--offload-dir", + "/tmp/offload", + "--output", + "/tmp/out.json", + ] + ) + assert args.baselines == ("B0", "B1", "B2", "B3") + assert args.block_sizes == (8, 16) + assert args.concurrency == (1, 2, 4, 8, 16, 32) + assert args.seed == 1408 + assert args.device_memory_ratio < 0.9 + + +def test_run_requires_model_draft_offload_output(): + with pytest.raises(SystemExit, match="missing required args"): + main(["--output", "/tmp/out.json"]) diff --git a/tests/python/dflash/test_pd_dflash_serving_gpu.py b/tests/python/dflash/test_pd_dflash_serving_gpu.py new file mode 100644 index 00000000..b6d97efe --- /dev/null +++ b/tests/python/dflash/test_pd_dflash_serving_gpu.py @@ -0,0 +1,52 @@ +"""Opt-in RTX PRO 6000 gate for the PD-DFlash B0-B3 serving runner. + +Task 2 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md``. +Collection is side-effect free: the module imports only the CPU-safe runner +scaffolding and the single test is skipped unless ``MOE_DFLASH_SERVING_GPU=1``, +so ``pytest`` never initialises CUDA, loads a checkpoint, hits the network, or +creates offload state when the gate is absent (``1 skipped``). +""" + +from __future__ import annotations + +import json +import os + +import pytest + +from benchmarks.dflash.pd_dflash_serving import ( + REQUIRED_CONCURRENCY, + REQUIRED_DRAFTS, + REQUIRED_MODELS, + main, +) + +RUN_GPU = os.environ.get("MOE_DFLASH_SERVING_GPU") == "1" +pytestmark = [ + pytest.mark.gpu, + pytest.mark.integration, + pytest.mark.skipif(not RUN_GPU, reason="set MOE_DFLASH_SERVING_GPU=1"), +] + + +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 "openai/gpt-oss-20b" in contract["models"] + assert set(contract["models"]) == set(REQUIRED_MODELS) + + assert ( + contract["drafts"]["Qwen/Qwen3-Coder-30B-A3B"] + == "z-lab/Qwen3-Coder-30B-A3B-DFlash" + ) + assert ( + contract["drafts"]["openai/gpt-oss-20b"] == "z-lab/gpt-oss-20b-DFlash" + ) + assert contract["drafts"] == dict(REQUIRED_DRAFTS) + + assert set(contract["block_sizes"]) == {8, 16} + assert set(contract["concurrency"]) == set(REQUIRED_CONCURRENCY) + assert set(contract["concurrency"]) == {1, 2, 4, 8, 16, 32} + assert set(contract["baselines"]) == {"B0", "B1", "B2", "B3"} From 38d1463c1daa9b42c1576ba2767ae0d48bd09c3e Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sat, 15 Aug 2026 00:12:26 +0000 Subject: [PATCH 4/9] bench(dflash): gate router-ahead cost and aggregate matrices MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add BM1 summarise_row (pass iff t_router < t_verify, ratio + raw terms retained) and a CPU-safe aggregation CLI: group raw rows into §8 matrices keyed by (model,block,concurrency), permit blocked B2 via --allow-blocked, attach BM1, and emit result_matrix.json/CSV/Markdown for validate_result_matrix. Refactors the per-baseline metric check out of validate_result_matrix (behavior-preserving). --- benchmarks/dflash/report.py | 241 +++++++++++++++++-- tests/python/dflash/test_pd_dflash_report.py | 108 +++++++++ 2 files changed, 333 insertions(+), 16 deletions(-) diff --git a/benchmarks/dflash/report.py b/benchmarks/dflash/report.py index 19e3b1b8..c465dbf1 100644 --- a/benchmarks/dflash/report.py +++ b/benchmarks/dflash/report.py @@ -14,9 +14,12 @@ from __future__ import annotations +import argparse +import json import math +import sys from dataclasses import dataclass -from typing import Mapping +from typing import Any, Dict, List, Mapping, Sequence, Tuple REQUIRED_METRICS = ( "output_tokens_per_second", @@ -129,21 +132,219 @@ def validate_result_matrix( row = rows[baseline] if baseline == "B3" and row.get("status") == UNAVAILABLE_CAPACITY: continue - for metric in REQUIRED_METRICS: - if metric not in row: - raise ValueError(f"{baseline} missing metric: {metric}") - value = row[metric] - if isinstance(value, bool) or not isinstance(value, (int, float)): - raise ValueError( - f"{baseline}.{metric} must be a finite, non-negative " - f"number; got {value!r}" - ) - number = float(value) - if not math.isfinite(number) or number < 0.0: - raise ValueError( - f"{baseline}.{metric} must be a finite, non-negative " - f"number; got {value!r}" - ) + _validate_metric_row(baseline, row) + + +def _validate_metric_row(baseline: str, row: Mapping[str, object]) -> None: + for metric in REQUIRED_METRICS: + if metric not in row: + raise ValueError(f"{baseline} missing metric: {metric}") + value = row[metric] + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError( + f"{baseline}.{metric} must be a finite, non-negative " + f"number; got {value!r}" + ) + number = float(value) + if not math.isfinite(number) or number < 0.0: + raise ValueError( + f"{baseline}.{metric} must be a finite, non-negative " + f"number; got {value!r}" + ) + + +def summarise_row(row: Mapping[str, object]) -> Mapping[str, object]: + """BM1 router-ahead cost summary for one row (design §10). + + BM1 passes when the route-ahead router projection is strictly cheaper than + the width-B verify it front-runs (``t_router < t_verify``); the raw seconds + and their ratio are retained so the aggregator can rank configurations, not + only gate them. Raises ``ValueError`` on a missing term, a negative time, or + a non-positive ``t_verify_seconds``. + """ + for key in ("t_router_seconds", "t_verify_seconds"): + if key not in row: + raise ValueError(f"row missing BM1 term: {key}") + t_router = _require_finite_non_negative( + "t_router_seconds", row["t_router_seconds"] + ) + t_verify = _require_finite_non_negative( + "t_verify_seconds", row["t_verify_seconds"] + ) + if t_verify <= 0.0: + raise ValueError(f"t_verify_seconds must be > 0; got {t_verify!r}") + return { + "t_router_seconds": t_router, + "t_verify_seconds": t_verify, + "bm1_router_to_verify_ratio": t_router / t_verify, + "bm1_pass": t_router < t_verify, + } + + +def aggregate_result_matrices( + rows: Sequence[Mapping[str, Any]], +) -> Dict[Tuple[str, int, int], Dict[str, Mapping[str, Any]]]: + """Group raw observation rows into §8 matrices. + + Keyed by ``(model, block_size, concurrency)``; each value maps a baseline + label to its single row. A duplicate ``(key, baseline)`` raises, mirroring + the runner's append-without-overwrite guarantee. + """ + matrices: Dict[Tuple[str, int, int], Dict[str, Mapping[str, Any]]] = {} + for row in rows: + key = ( + str(row["model"]), + int(row["block_size"]), + int(row["concurrency"]), + ) + baseline = str(row["baseline"]) + bucket = matrices.setdefault(key, {}) + if baseline in bucket: + raise ValueError(f"duplicate baseline {baseline} for {key}") + bucket[baseline] = row + return matrices + + +def evaluate_matrix( + baseline_rows: Mapping[str, Mapping[str, Any]], + allow_blocked: Sequence[str] = (), +) -> Tuple[bool, Dict[str, Any]]: + """Completeness + BM1 verdict for one grouped matrix. + + A baseline is satisfied when it carries the full metric schema; a baseline + listed in ``allow_blocked`` may instead carry a blocking ``status`` (e.g. + B2's ``BLOCKED_UNTIL_2D_SCHEDULER`` before the 2-D scheduler lands). Any + other missing/invalid/blocked baseline fails the group. BM1 summaries are + attached for every row carrying ``t_router_seconds``/``t_verify_seconds``. + """ + allow = set(allow_blocked) + detail: Dict[str, Any] = { + "present": sorted(baseline_rows), + "blocked": [], + "missing": [], + "invalid": [], + "bm1": {}, + } + ok = True + for baseline in REQUIRED_BASELINES: + row = baseline_rows.get(baseline) + if row is None: + detail["missing"].append(baseline) + ok = False + continue + if row.get("status"): + if baseline in allow: + detail["blocked"].append(baseline) + else: + detail["blocked"].append(baseline) + ok = False + continue + try: + _validate_metric_row(baseline, row) + except ValueError: + detail["invalid"].append(baseline) + ok = False + if "t_router_seconds" in row and "t_verify_seconds" in row: + detail["bm1"][baseline] = dict(summarise_row(row)) + return ok, detail + + +def _matrix_key(key: Tuple[str, int, int]) -> str: + return f"{key[0]}|B{key[1]}|c{key[2]}" + + +def _write_json(path: str, payload: Mapping[str, Any]) -> None: + with open(path, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True) + handle.write("\n") + + +def _write_csv( + path: str, + matrices: Mapping[Tuple[str, int, int], Mapping[str, Mapping[str, Any]]], +) -> None: + import csv + + columns = ["model", "block_size", "concurrency", "baseline", "status"] + columns += list(REQUIRED_METRICS) + with open(path, "w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=columns) + writer.writeheader() + for (model, block, conc), baseline_rows in sorted(matrices.items()): + for baseline, row in sorted(baseline_rows.items()): + record = { + "model": model, + "block_size": block, + "concurrency": conc, + "baseline": baseline, + "status": row.get("status", ""), + } + for metric in REQUIRED_METRICS: + record[metric] = row.get(metric, "") + writer.writerow(record) + + +def _write_markdown(path: str, report: Mapping[str, Any]) -> None: + lines = ["# PD-DFlash Phase-A result matrix", ""] + for group, detail in sorted(report.items()): + lines.append(f"## {group}") + lines.append(f"- present: {', '.join(detail['present']) or '(none)'}") + if detail["blocked"]: + lines.append(f"- blocked: {', '.join(detail['blocked'])}") + if detail["missing"]: + lines.append(f"- missing: {', '.join(detail['missing'])}") + if detail["invalid"]: + lines.append(f"- invalid: {', '.join(detail['invalid'])}") + for baseline, bm1 in sorted(detail["bm1"].items()): + lines.append( + f"- BM1 {baseline}: ratio=" + f"{bm1['bm1_router_to_verify_ratio']:.4f} " + f"pass={bm1['bm1_pass']}" + ) + lines.append("") + with open(path, "w", encoding="utf-8") as handle: + handle.write("\n".join(lines)) + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m benchmarks.dflash.report", + description="Aggregate PD-DFlash raw observation rows into §8 matrices.", + ) + parser.add_argument("--input", nargs="+", required=True) + parser.add_argument("--matrix-json") + parser.add_argument("--csv") + parser.add_argument("--markdown") + parser.add_argument("--allow-blocked", nargs="*", default=[]) + args = parser.parse_args(argv) + + rows: List[Mapping[str, Any]] = [] + for path in args.input: + with open(path, "r", encoding="utf-8") as handle: + data = json.load(handle) + rows.extend(data if isinstance(data, list) else [data]) + + matrices = aggregate_result_matrices(rows) + report: Dict[str, Any] = {} + all_ok = True + for key, baseline_rows in sorted(matrices.items()): + ok, detail = evaluate_matrix(baseline_rows, args.allow_blocked) + all_ok = all_ok and ok + report[_matrix_key(key)] = detail + + if args.matrix_json: + _write_json( + args.matrix_json, + {_matrix_key(k): dict(v) for k, v in matrices.items()}, + ) + if args.csv: + _write_csv(args.csv, matrices) + if args.markdown: + _write_markdown(args.markdown, report) + + json.dump(report, sys.stdout, indent=2, sort_keys=True) + sys.stdout.write("\n") + return 0 if all_ok else 1 __all__ = [ @@ -151,6 +352,14 @@ def validate_result_matrix( "REQUIRED_BASELINES", "UNAVAILABLE_CAPACITY", "HideInequality", + "aggregate_result_matrices", "evaluate_hide_inequality", + "evaluate_matrix", + "main", + "summarise_row", "validate_result_matrix", ] + + +if __name__ == "__main__": # pragma: no cover - CLI entry + raise SystemExit(main()) diff --git a/tests/python/dflash/test_pd_dflash_report.py b/tests/python/dflash/test_pd_dflash_report.py index c1489f82..8743b412 100644 --- a/tests/python/dflash/test_pd_dflash_report.py +++ b/tests/python/dflash/test_pd_dflash_report.py @@ -22,7 +22,10 @@ from benchmarks.dflash.report import ( REQUIRED_METRICS, + aggregate_result_matrices, evaluate_hide_inequality, + evaluate_matrix, + summarise_row, validate_result_matrix, ) @@ -144,6 +147,111 @@ def test_wasted_prefetch_is_bytes_not_expert_count(): assert "wasted_prefetch_bytes" in REQUIRED_METRICS +# --------------------------------------------------------------------------- +# BM1 router-ahead cost aggregation (design §10) +# --------------------------------------------------------------------------- + + +def complete_row(**overrides) -> dict[str, float]: + row: dict[str, float] = {metric: 1.0 for metric in REQUIRED_METRICS} + row["t_router_seconds"] = 0.002 + row["t_verify_seconds"] = 0.020 + row.update(overrides) + return row + + +def test_bm1_reports_router_cost_against_verify(): + row = complete_row(t_router_seconds=0.002, t_verify_seconds=0.020) + summary = summarise_row(row) + assert summary["bm1_router_to_verify_ratio"] == 0.1 + assert summary["bm1_pass"] is True + + +def test_bm1_fails_when_router_not_cheaper_than_verify(): + summary = summarise_row( + complete_row(t_router_seconds=0.03, t_verify_seconds=0.02) + ) + assert summary["bm1_pass"] is False + assert summary["bm1_router_to_verify_ratio"] == 1.5 + + +def test_bm1_retains_raw_terms_and_rejects_bad_inputs(): + summary = summarise_row(complete_row()) + assert summary["t_router_seconds"] == 0.002 + assert summary["t_verify_seconds"] == 0.020 + with pytest.raises(ValueError, match="t_verify_seconds"): + summarise_row(complete_row(t_verify_seconds=0.0)) + with pytest.raises(ValueError, match="t_router_seconds"): + summarise_row(complete_row(t_router_seconds=-0.1)) + with pytest.raises(ValueError, match="t_router_seconds"): + summarise_row({"t_verify_seconds": 0.02}) + + +# --------------------------------------------------------------------------- +# aggregation: group raw rows into §8 matrices, allow blocked B2 +# --------------------------------------------------------------------------- + + +def _obs_row(baseline: str, **overrides) -> dict[str, object]: + row: dict[str, object] = { + "model": "M", + "draft": "d", + "baseline": baseline, + "block_size": 16, + "concurrency": 8, + "repeat": 0, + } + row.update({metric: 1.0 for metric in REQUIRED_METRICS}) + row.update(overrides) + return row + + +def _blocked_b2() -> dict[str, object]: + return { + "model": "M", + "draft": "d", + "baseline": "B2", + "block_size": 16, + "concurrency": 8, + "repeat": 0, + "status": "BLOCKED_UNTIL_2D_SCHEDULER", + } + + +def test_aggregate_groups_rows_and_rejects_duplicate_baseline(): + matrices = aggregate_result_matrices( + [_obs_row(b) for b in ("B0", "B1", "B3")] + ) + assert set(matrices) == {("M", 16, 8)} + assert set(matrices[("M", 16, 8)]) == {"B0", "B1", "B3"} + with pytest.raises(ValueError, match="duplicate baseline"): + aggregate_result_matrices([_obs_row("B0"), _obs_row("B0")]) + + +def test_evaluate_matrix_allows_blocked_b2_only_when_permitted(): + rows = {b: _obs_row(b) for b in ("B0", "B1", "B3")} + rows["B2"] = _blocked_b2() + ok, detail = evaluate_matrix(rows, allow_blocked=["B2"]) + assert ok is True and detail["blocked"] == ["B2"] + not_ok, detail2 = evaluate_matrix(rows, allow_blocked=[]) + assert not_ok is False and "B2" in detail2["blocked"] + + +def test_evaluate_matrix_flags_missing_baseline_and_attaches_bm1(): + partial, missing = evaluate_matrix( + {b: _obs_row(b) for b in ("B0", "B1")}, [] + ) + assert partial is False and set(missing["missing"]) == {"B2", "B3"} + + full = { + b: _obs_row(b, t_router_seconds=0.002, t_verify_seconds=0.020) + for b in ("B0", "B1", "B2", "B3") + } + ok, detail = evaluate_matrix(full, []) + assert ok is True + assert detail["bm1"]["B1"]["bm1_pass"] is True + + def test_required_metrics_are_frozen_and_complete(): assert isinstance(REQUIRED_METRICS, tuple) assert REQUIRED_METRICS == ( From 4dc6027b1f8818c084ce0d1808261153dbb64f95 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sat, 15 Aug 2026 00:12:26 +0000 Subject: [PATCH 5/9] bench(dflash): one-command Phase-A launcher MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit run_phase_a.sh drives the full §8 matrix for both required MoE targets on one RTX PRO 6000 with the documented env (HF_HOME, MOE_ENABLE_SM120, device-memory- ratio<0.9 to force offload) and aggregates into result_matrix.json. All inputs are documented env vars in the script header. --- benchmarks/dflash/run_phase_a.sh | 138 +++++++++++++++++++++++++++++++ 1 file changed, 138 insertions(+) create mode 100755 benchmarks/dflash/run_phase_a.sh diff --git a/benchmarks/dflash/run_phase_a.sh b/benchmarks/dflash/run_phase_a.sh new file mode 100755 index 00000000..4b7cb692 --- /dev/null +++ b/benchmarks/dflash/run_phase_a.sh @@ -0,0 +1,138 @@ +#!/usr/bin/env bash +# =========================================================================== +# run_phase_a.sh -- one-command PD-DFlash Phase-A "measure-first" matrix. +# +# Runs the §8 B0-B3 route-ahead serving experiment on ONE RTX PRO 6000 +# (sm_120, capability 12.0) for both required MoE targets with FP4-offloaded +# experts, then aggregates the raw rows into a result-matrix JSON that +# benchmarks.dflash.report / validate_result_matrix consumes. This is the +# hardware harness for plan Task 3 (docs/superpowers/plans/ +# 2026-08-14-pd-dflash-serving-scheduler.md); B2 is emitted BLOCKED until the +# 2-D scheduler (Task 6) lands, so it is passed via --allow-blocked B2. +# +# USAGE +# benchmarks/dflash/run_phase_a.sh +# +# All inputs are environment variables (shown with their defaults). Override +# any of them inline, e.g.: +# QWEN_OFFLOAD=/data/qwen-fp4 MEASURED_H2D_GBPS=48 \ +# benchmarks/dflash/run_phase_a.sh +# +# REQUIRED on the GPU box (defaults assume this project's conventions): +# HF_HOME cached checkpoints (default /mnt/raid0nvme0/public/huggingface) +# CUDA_VISIBLE_DEVICES the single RTX PRO 6000 to use (default 0) +# QWEN_OFFLOAD dir of FP4-offloaded Qwen3-Coder-30B-A3B experts +# GPTOSS_OFFLOAD dir of FP4-offloaded gpt-oss-20b experts (needs #137) +# +# KEY KNOBS +# DEVICE_MEMORY_RATIO weight-resident fraction; MUST be < 0.9 to force +# offload for B0/B1/B2 (default 0.85) +# MEASURED_H2D_GBPS measured host->GPU expert bandwidth (GB/s) for the +# hide inequality; NOT a theoretical PCIe number +# SLO_MS per-round SLO for goodput@SLO (optional) +# BASELINES default "B0 B1 B2 B3" +# BLOCK_SIZES default "8 16" +# CONCURRENCY default "1 2 4 8 16 32" +# REQUESTS / WARMUP / SEED default 64 / 5 / 1408 +# PD_DFLASH_BUILD=1 rebuild the native sm_120 extensions first +# (MOE_ENABLE_SM120=1 MOE_ENABLE_SM90=0) +# +# OUTPUTS (under $OUTPUT_DIR, default /tmp/pd-dflash-results) +# raw/qwen.json, raw/gpt-oss.json one JSON row per (model,baseline,B,c,repeat) +# result_matrix.json grouped {model|B|c: {baseline: row}} +# summary.csv, summary.md human-readable aggregation +# =========================================================================== +set -euo pipefail + +export HF_HOME="${HF_HOME:-/mnt/raid0nvme0/public/huggingface}" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" +export MOE_ENABLE_SM120="${MOE_ENABLE_SM120:-1}" +# The runner validates the device itself; this mirrors the pytest gate name so +# any nested gated assertions also run on the box. +export MOE_DFLASH_SERVING_GPU="${MOE_DFLASH_SERVING_GPU:-1}" + +OUTPUT_DIR="${OUTPUT_DIR:-/tmp/pd-dflash-results}" +BASELINES="${BASELINES:-B0 B1 B2 B3}" +BLOCK_SIZES="${BLOCK_SIZES:-8 16}" +CONCURRENCY="${CONCURRENCY:-1 2 4 8 16 32}" +REQUESTS="${REQUESTS:-64}" +WARMUP="${WARMUP:-5}" +SEED="${SEED:-1408}" +DEVICE_MEMORY_RATIO="${DEVICE_MEMORY_RATIO:-0.85}" + +QWEN_MODEL="${QWEN_MODEL:-Qwen/Qwen3-Coder-30B-A3B}" +QWEN_DRAFT="${QWEN_DRAFT:-z-lab/Qwen3-Coder-30B-A3B-DFlash}" +QWEN_OFFLOAD="${QWEN_OFFLOAD:-/mnt/raid0nvme0/offload/qwen3-coder-30b-a3b-fp4}" + +GPTOSS_MODEL="${GPTOSS_MODEL:-openai/gpt-oss-20b}" +GPTOSS_DRAFT="${GPTOSS_DRAFT:-z-lab/gpt-oss-20b-DFlash}" +GPTOSS_OFFLOAD="${GPTOSS_OFFLOAD:-/mnt/raid0nvme0/offload/gpt-oss-20b-fp4}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "$REPO_ROOT" + +mkdir -p "$OUTPUT_DIR/raw" + +echo "[run_phase_a] repo=$REPO_ROOT out=$OUTPUT_DIR device=$CUDA_VISIBLE_DEVICES" +echo "[run_phase_a] device_memory_ratio=$DEVICE_MEMORY_RATIO (must be < 0.9 to offload)" + +if awk "BEGIN{exit !($DEVICE_MEMORY_RATIO >= 0.9)}"; then + echo "[run_phase_a] ERROR: DEVICE_MEMORY_RATIO=$DEVICE_MEMORY_RATIO >= 0.9 will not offload B0/B1/B2" >&2 + exit 2 +fi + +if [[ "${PD_DFLASH_BUILD:-0}" == "1" ]]; then + echo "[run_phase_a] building native sm_120 extensions" + MOE_ENABLE_SM120=1 MOE_ENABLE_SM90=0 CUTLASS_DIR="${CUTLASS_DIR:-$HOME/cutlass}" \ + pip install --no-build-isolation -e . +fi + +extra_args=() +if [[ -n "${MEASURED_H2D_GBPS:-}" ]]; then + extra_args+=(--measured-h2d-gbps "$MEASURED_H2D_GBPS") +fi +if [[ -n "${SLO_MS:-}" ]]; then + extra_args+=(--slo-ms "$SLO_MS") +fi + +run_model() { + local model="$1" draft="$2" offload="$3" output="$4" + echo "[run_phase_a] === $model -> $output ===" + if [[ ! -d "$offload" ]]; then + echo "[run_phase_a] WARNING: offload dir '$offload' missing; the runner will" \ + "refuse B0/B1/B2 unless experts are genuinely offloaded" >&2 + fi + python -m benchmarks.dflash.pd_dflash_serving \ + --model "$model" \ + --draft "$draft" \ + --offload-dir "$offload" \ + --baseline $BASELINES \ + --block-size $BLOCK_SIZES \ + --concurrency $CONCURRENCY \ + --requests "$REQUESTS" \ + --warmup-rounds "$WARMUP" \ + --seed "$SEED" \ + --device-memory-ratio "$DEVICE_MEMORY_RATIO" \ + "${extra_args[@]}" \ + --output "$output" +} + +run_model "$QWEN_MODEL" "$QWEN_DRAFT" "$QWEN_OFFLOAD" "$OUTPUT_DIR/raw/qwen.json" +run_model "$GPTOSS_MODEL" "$GPTOSS_DRAFT" "$GPTOSS_OFFLOAD" "$OUTPUT_DIR/raw/gpt-oss.json" + +echo "[run_phase_a] aggregating result matrix" +python -m benchmarks.dflash.report \ + --input "$OUTPUT_DIR/raw/qwen.json" "$OUTPUT_DIR/raw/gpt-oss.json" \ + --matrix-json "$OUTPUT_DIR/result_matrix.json" \ + --csv "$OUTPUT_DIR/summary.csv" \ + --markdown "$OUTPUT_DIR/summary.md" \ + --allow-blocked B2 || { + echo "[run_phase_a] report gate FAILED (missing/invalid baseline); inspect" \ + "$OUTPUT_DIR/result_matrix.json" >&2 + exit 1 + } + +echo "[run_phase_a] done:" +echo " raw rows: $OUTPUT_DIR/raw/{qwen,gpt-oss}.json" +echo " result matrix: $OUTPUT_DIR/result_matrix.json" +echo " summary: $OUTPUT_DIR/summary.{csv,md}" From a5ece6e850bf43c2bfc7221dabc2e462604d11ab Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sat, 15 Aug 2026 11:49:51 +0000 Subject: [PATCH 6/9] bench(dflash): measure saturated prefetch issuance (BM2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 7 of the PD-DFlash serving plan (design §10 BM2). Adds an import-safe issuance micro-benchmark that times enqueuing a saturated E_l×L block of offloaded expert tensors three ways -- python-per-expert, batched-pybind, and (reserved) cpp-internal -- reporting p50/p90/p99 microseconds and the ship gate. - benchmarks/dflash/bench_prefetch_issuance.py: pure bm2_decision / percentiles_us / build_bm2_report at module scope (torch + moe_infinity lazily imported in the GPU runner), plus the CLI runner that synthesises the real saturated tensor-id list from a loaded target and times warmup=20 / iterations=200 with perf_counter_ns. Unavailable candidate modes are reported null, never zero. - tests/python/dflash/test_prefetch_perf_reports.py: CPU-only decision-rule and report-schema tests (the plan's four exact assertions plus boundary cases). BM2 alone gates the batched-issuance C++ hop; committed independently of it. --- benchmarks/dflash/bench_prefetch_issuance.py | 419 ++++++++++++++++++ .../dflash/test_prefetch_perf_reports.py | 161 +++++++ 2 files changed, 580 insertions(+) create mode 100644 benchmarks/dflash/bench_prefetch_issuance.py create mode 100644 tests/python/dflash/test_prefetch_perf_reports.py diff --git a/benchmarks/dflash/bench_prefetch_issuance.py b/benchmarks/dflash/bench_prefetch_issuance.py new file mode 100644 index 00000000..a2d4bd7f --- /dev/null +++ b/benchmarks/dflash/bench_prefetch_issuance.py @@ -0,0 +1,419 @@ +# Copyright (c) EfficientMoE. +# SPDX-License-Identifier: Apache-2.0 + +# EfficientMoE Team + +"""BM2 -- saturated route-ahead prefetch issuance micro-benchmark (design §10). + +Task 7 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md``. +Measures how long it takes to *issue* (enqueue) a route-ahead prefetch for a +saturated block of ``E_l x L`` offloaded expert tensors, three ways: + +* ``python-per-expert`` -- the current + ``ExpertPrefetcher.prefetch_experts_list`` path: one + ``get_node_default_device`` + ``enqueue_prefetch`` pybind pair per tensor + (``2 * E_l * L`` boundary crossings); +* ``batched-pybind`` -- a single ``prefetch_handle.prefetch_tensors(tensor_ids)`` + call that constructs and enqueues every ``Task`` inside C++ (one crossing); + available only once the batched native API (plan Task 8) is built in; +* ``cpp-internal`` -- reserved for a native in-C++ issuance timer; reported as + ``null`` until such a hook exists (never zero). + +The ship gate (design §10, plan Task 7/8): the batched hop is justified iff the +current Python per-expert median exceeds the route-ahead window +``t_draft + t_router`` *and* the batched median is at or below it. + +Import-safe by construction: torch and moe_infinity are imported lazily inside +the GPU runner, so ``bm2_decision`` / ``percentiles_us`` / ``build_bm2_report`` +(and their tests) are pure-CPU and never initialise CUDA. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +from dataclasses import dataclass +from time import perf_counter_ns +from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence + +PYTHON_PER_EXPERT = "python-per-expert" +BATCHED_PYBIND = "batched-pybind" +CPP_INTERNAL = "cpp-internal" +ISSUANCE_MODES = (PYTHON_PER_EXPERT, BATCHED_PYBIND, CPP_INTERNAL) + + +@dataclass(frozen=True) +class Bm2Decision: + """The BM2 ship gate over per-mode issuance medians (design §10).""" + + per_expert_us: Optional[float] + batched_us: Optional[float] + cpp_internal_us: Optional[float] + window_us: float + candidate_required: bool + ship_batched: bool + + +def _finite_positive_window(window_us: Any) -> float: + window = float(window_us) + if not math.isfinite(window) or window <= 0.0: + raise ValueError( + "window_us (t_draft + t_router) must be finite and > 0; " + f"got {window_us!r}" + ) + return window + + +def _optional_us(name: str, value: Any) -> Optional[float]: + if value is None: + return None + number = float(value) + if not math.isfinite(number) or number < 0.0: + raise ValueError( + f"{name} must be a finite, non-negative microsecond median or " + f"None; got {value!r}" + ) + return number + + +def bm2_decision( + per_expert_us: Optional[float], + batched_us: Optional[float], + cpp_internal_us: Optional[float], + window_us: float, +) -> Bm2Decision: + """Evaluate the BM2 ship gate from measured medians (microseconds). + + ``candidate_required`` holds when the current per-expert median exceeds the + route-ahead window. ``ship_batched`` additionally requires a *measured* + batched median at or below the window -- a missing batched median can never + ship, so an unavailable candidate mode never masquerades as a win. + """ + window = _finite_positive_window(window_us) + per_expert = _optional_us("per_expert_us", per_expert_us) + batched = _optional_us("batched_us", batched_us) + cpp_internal = _optional_us("cpp_internal_us", cpp_internal_us) + + candidate_required = per_expert is not None and per_expert > window + ship_batched = ( + candidate_required and batched is not None and batched <= window + ) + return Bm2Decision( + per_expert_us=per_expert, + batched_us=batched, + cpp_internal_us=cpp_internal, + window_us=window, + candidate_required=candidate_required, + ship_batched=ship_batched, + ) + + +def percentiles_us(samples_ns: Sequence[int]) -> Dict[str, float]: + """Nearest-rank p50/p90/p99 of nanosecond samples, returned in microseconds.""" + if not samples_ns: + raise ValueError("percentiles_us requires at least one sample") + ordered = sorted(float(sample) for sample in samples_ns) + count = len(ordered) + + def nearest_rank(pct: float) -> float: + rank = min(max(math.ceil(pct * count), 1), count) + return ordered[rank - 1] / 1000.0 + + return { + "p50": nearest_rank(0.50), + "p90": nearest_rank(0.90), + "p99": nearest_rank(0.99), + "min": ordered[0] / 1000.0, + "max": ordered[-1] / 1000.0, + "count": count, + } + + +def _mode_stats( + samples_ns: Optional[Sequence[int]], +) -> Optional[Dict[str, float]]: + if samples_ns is None: + return None + return percentiles_us(samples_ns) + + +def build_bm2_report( + *, + model: str, + saturated_tensor_count: int, + window_us: float, + per_expert_samples_ns: Optional[Sequence[int]], + batched_samples_ns: Optional[Sequence[int]] = None, + cpp_internal_samples_ns: Optional[Sequence[int]] = None, + warmup: int, + iterations: int, + extra: Optional[Mapping[str, Any]] = None, +) -> Dict[str, Any]: + """Assemble the machine-readable BM2 report and its ship-gate verdict.""" + modes: Dict[str, Optional[Dict[str, float]]] = { + PYTHON_PER_EXPERT: _mode_stats(per_expert_samples_ns), + BATCHED_PYBIND: _mode_stats(batched_samples_ns), + CPP_INTERNAL: _mode_stats(cpp_internal_samples_ns), + } + + def median(mode: str) -> Optional[float]: + stats = modes[mode] + return None if stats is None else stats["p50"] + + decision = bm2_decision( + median(PYTHON_PER_EXPERT), + median(BATCHED_PYBIND), + median(CPP_INTERNAL), + window_us, + ) + report: Dict[str, Any] = { + "benchmark": "BM2", + "model": model, + "saturated_tensor_count": int(saturated_tensor_count), + "window_us": decision.window_us, + "warmup": int(warmup), + "iterations": int(iterations), + "modes": modes, + "medians_us": { + PYTHON_PER_EXPERT: decision.per_expert_us, + BATCHED_PYBIND: decision.batched_us, + CPP_INTERNAL: decision.cpp_internal_us, + }, + "candidate_required": decision.candidate_required, + "ship_batched": decision.ship_batched, + } + if extra: + report.update(dict(extra)) + return report + + +def _resolve_window_us( + window_json: Optional[str], window_us: Optional[float] +) -> float: + """Resolve ``t_draft + t_router`` (microseconds) for the ship gate. + + ``--window-us`` wins when given; otherwise a Phase-A raw JSON is read and + its ``t_draft``/``t_router`` seconds (either bare or ``*_seconds``-suffixed) + are summed. Never substitutes a theoretical or hard-coded default. + """ + if window_us is not None: + return _finite_positive_window(window_us) + if window_json is None: + raise ValueError( + "a route-ahead window is required: pass --window-us or a " + "--window-json carrying t_draft/t_router seconds" + ) + with open(window_json, "r", encoding="utf-8") as handle: + payload = json.load(handle) + rows = payload if isinstance(payload, list) else [payload] + + def field(row: Mapping[str, Any], *names: str) -> Optional[float]: + for name in names: + if name in row and row[name] is not None: + return float(row[name]) + return None + + for row in rows: + draft = field(row, "t_draft_seconds", "t_draft") + router = field(row, "t_router_seconds", "t_router") + if draft is not None and router is not None: + return _finite_positive_window((draft + router) * 1e6) + raise ValueError( + f"could not find t_draft and t_router seconds in {window_json!r}" + ) + + +def _load_prefetcher( + model_repo: str, offload_path: str, device_memory_ratio: float +) -> Any: + from moe_infinity import MoE # lazy: heavy, CUDA-initialising + + model = MoE( + model_repo, + { + "offload_path": offload_path, + "device_memory_ratio": device_memory_ratio, + }, + ) + prefetcher = model.engine.expert_prefetcher + if prefetcher is None or prefetcher.archer_engine is None: + raise RuntimeError( + "loaded model has no offloaded ExpertPrefetcher/archer_engine; " + "ensure device_memory_ratio < 1 so experts are actually offloaded" + ) + return model, prefetcher + + +def _saturated_tensor_ids(prefetcher: Any) -> List[int]: + """Every ``(layer, expert)`` tensor id -- the saturated ``E_l x L`` block.""" + tensor_map = prefetcher.expert_tensor_map + if not tensor_map: + raise RuntimeError("expert_tensor_map is empty; no experts to issue") + return [tensor_id for _key, tensor_id in sorted(tensor_map.items())] + + +def _time_rounds( + issue: Callable[[], None], warmup: int, iterations: int +) -> List[int]: + for _ in range(warmup): + issue() + samples_ns: List[int] = [] + for _ in range(iterations): + start = perf_counter_ns() + issue() + samples_ns.append(perf_counter_ns() - start) + return samples_ns + + +def _python_per_expert_issue( + engine: Any, tensor_ids: Sequence[int] +) -> Callable[[], None]: + def issue() -> None: + for tensor_id in tensor_ids: + gpu_id = engine.get_node_default_device([tensor_id]) + engine.enqueue_prefetch(tensor_id, gpu_id) + + return issue + + +def _batched_issue( + engine: Any, tensor_ids: Sequence[int] +) -> Optional[Callable[[], None]]: + """Return a one-call batched issuer, or ``None`` if the native API is the + pre-Task-8 no-op signature (probed once against a single tensor id).""" + probe = list(tensor_ids[:1]) + try: + engine.prefetch_tensors(probe) + except Exception: + return None + + ids = list(tensor_ids) + + def issue() -> None: + engine.prefetch_tensors(ids) + + return issue + + +def run_issuance_benchmark( + *, + model_repo: str, + offload_path: str, + device_memory_ratio: float, + modes: Sequence[str], + warmup: int, + iterations: int, + window_us: float, +) -> Dict[str, Any]: + model, prefetcher = _load_prefetcher( + model_repo, offload_path, device_memory_ratio + ) + engine = prefetcher.archer_engine + tensor_ids = _saturated_tensor_ids(prefetcher) + + per_expert_ns: Optional[List[int]] = None + batched_ns: Optional[List[int]] = None + cpp_internal_ns: Optional[List[int]] = None + unavailable: Dict[str, str] = {} + + if PYTHON_PER_EXPERT in modes: + per_expert_ns = _time_rounds( + _python_per_expert_issue(engine, tensor_ids), warmup, iterations + ) + if BATCHED_PYBIND in modes: + issuer = _batched_issue(engine, tensor_ids) + if issuer is None: + unavailable[BATCHED_PYBIND] = ( + "native prefetch_tensors(tensor_ids) batched API absent " + "(pre-Task-8 no-op binding); rebuild _store to enable" + ) + else: + batched_ns = _time_rounds(issuer, warmup, iterations) + if CPP_INTERNAL in modes: + unavailable[CPP_INTERNAL] = ( + "no native in-C++ issuance timer exposed; reported null" + ) + + extra: Dict[str, Any] = { + "offload_path": offload_path, + "device_memory_ratio": device_memory_ratio, + "requested_modes": list(modes), + } + if unavailable: + extra["unavailable_modes"] = unavailable + + return build_bm2_report( + model=model_repo, + saturated_tensor_count=len(tensor_ids), + window_us=window_us, + per_expert_samples_ns=per_expert_ns, + batched_samples_ns=batched_ns, + cpp_internal_samples_ns=cpp_internal_ns, + warmup=warmup, + iterations=iterations, + extra=extra, + ) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m benchmarks.dflash.bench_prefetch_issuance", + description="BM2 saturated route-ahead prefetch issuance micro-bench.", + ) + parser.add_argument("--model", required=True) + parser.add_argument("--offload-dir", required=True) + parser.add_argument("--device-memory-ratio", type=float, default=0.9) + parser.add_argument( + "--mode", nargs="+", default=[PYTHON_PER_EXPERT], choices=ISSUANCE_MODES + ) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=200) + parser.add_argument("--window-json") + parser.add_argument("--window-us", type=float) + parser.add_argument("--output", required=True) + args = parser.parse_args(argv) + + if not os.environ.get("MOE_DFLASH_SERVING_GPU"): + parser.error( + "MOE_DFLASH_SERVING_GPU must be set (opt-in GPU issuance bench)" + ) + + window_us = _resolve_window_us(args.window_json, args.window_us) + report = run_issuance_benchmark( + model_repo=args.model, + offload_path=args.offload_dir, + device_memory_ratio=args.device_memory_ratio, + modes=args.mode, + warmup=args.warmup, + iterations=args.iterations, + window_us=window_us, + ) + + os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) + with open(args.output, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, sort_keys=True) + handle.write("\n") + json.dump(report, sys.stdout, indent=2, sort_keys=True) + sys.stdout.write("\n") + return 0 + + +__all__ = [ + "PYTHON_PER_EXPERT", + "BATCHED_PYBIND", + "CPP_INTERNAL", + "ISSUANCE_MODES", + "Bm2Decision", + "bm2_decision", + "percentiles_us", + "build_bm2_report", + "run_issuance_benchmark", + "main", +] + + +if __name__ == "__main__": # pragma: no cover - CLI entry + raise SystemExit(main()) diff --git a/tests/python/dflash/test_prefetch_perf_reports.py b/tests/python/dflash/test_prefetch_perf_reports.py new file mode 100644 index 00000000..46a43631 --- /dev/null +++ b/tests/python/dflash/test_prefetch_perf_reports.py @@ -0,0 +1,161 @@ +# Copyright (c) EfficientMoE. +# SPDX-License-Identifier: Apache-2.0 + +# EfficientMoE Team + +"""CPU-only decision-rule tests for the benchmark-gated prefetch reports. + +Task 7 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md`` +(the BM2 issuance decision rule, design §10). These tests are pure: they never +import torch, load a checkpoint, or touch CUDA. They exercise only the +``bm2_decision`` rule, the percentile helper, and the report schema that the +GPU issuance micro-bench (``benchmarks.dflash.bench_prefetch_issuance``) emits, +so the ship gate for the batched ``prefetch_tensors`` C++ hop is verifiable +entirely off-hardware. + +The gate rule (design §10 / plan Task 7 Step 1): + +* ``candidate_required`` iff the current Python per-expert issuance median + exceeds the route-ahead window ``t_draft + t_router``; +* ``ship_batched`` iff a candidate is required *and* the batched-pybind median + is at or below that same window. + +Unavailable candidate medians are reported as ``None`` (JSON ``null``), never +zero, so a missing mode can never masquerade as an infinitely fast candidate. +""" + +from __future__ import annotations + +import pytest + +from benchmarks.dflash.bench_prefetch_issuance import ( + BATCHED_PYBIND, + CPP_INTERNAL, + PYTHON_PER_EXPERT, + Bm2Decision, + bm2_decision, + build_bm2_report, + percentiles_us, +) + + +def test_bm2_candidate_required_when_per_expert_exceeds_window(): + assert bm2_decision(900.0, None, None, 500.0).candidate_required is True + + +def test_bm2_no_candidate_when_per_expert_within_window(): + assert bm2_decision(400.0, None, None, 500.0).candidate_required is False + + +def test_bm2_ship_batched_true_when_batched_within_window(): + assert bm2_decision(900.0, 300.0, 250.0, 500.0).ship_batched is True + + +def test_bm2_ship_batched_false_when_batched_exceeds_window(): + assert bm2_decision(900.0, 700.0, 650.0, 500.0).ship_batched is False + + +def test_bm2_decision_is_frozen_and_reports_all_medians(): + decision = bm2_decision(900.0, 300.0, 250.0, 500.0) + assert isinstance(decision, Bm2Decision) + assert decision.per_expert_us == 900.0 + assert decision.batched_us == 300.0 + assert decision.cpp_internal_us == 250.0 + assert decision.window_us == 500.0 + with pytest.raises(Exception): + decision.per_expert_us = 1.0 # type: ignore[misc] + + +def test_bm2_ship_requires_a_measured_batched_median(): + # A candidate is required but no batched candidate exists yet: cannot ship. + decision = bm2_decision(900.0, None, None, 500.0) + assert decision.candidate_required is True + assert decision.ship_batched is False + + +def test_bm2_never_ships_without_a_candidate_even_if_batched_is_fast(): + # Per-expert already within the window -> no candidate -> never ship, + # even when a batched median would trivially satisfy the window. + decision = bm2_decision(400.0, 100.0, 90.0, 500.0) + assert decision.candidate_required is False + assert decision.ship_batched is False + + +def test_bm2_missing_per_expert_median_is_not_a_candidate(): + decision = bm2_decision(None, 100.0, 90.0, 500.0) + assert decision.candidate_required is False + assert decision.ship_batched is False + + +def test_bm2_batched_exactly_at_window_ships(): + # "<= window" is inclusive at the boundary. + assert bm2_decision(900.0, 500.0, None, 500.0).ship_batched is True + + +@pytest.mark.parametrize("bad_window", [0.0, -1.0, float("nan"), float("inf")]) +def test_bm2_window_must_be_finite_and_positive(bad_window): + with pytest.raises(ValueError): + bm2_decision(900.0, 300.0, 250.0, bad_window) + + +@pytest.mark.parametrize("bad_value", [-1.0, float("nan"), float("inf")]) +def test_bm2_negative_or_nonfinite_medians_rejected(bad_value): + with pytest.raises(ValueError): + bm2_decision(bad_value, None, None, 500.0) + + +def test_bm2_percentiles_us_from_nanoseconds_nearest_rank(): + # 1..100 microseconds expressed in nanoseconds. + samples_ns = [i * 1000 for i in range(1, 101)] + percentiles = percentiles_us(samples_ns) + assert percentiles["p50"] == pytest.approx(50.0) + assert percentiles["p90"] == pytest.approx(90.0) + assert percentiles["p99"] == pytest.approx(99.0) + assert percentiles["count"] == 100 + + +def test_bm2_percentiles_us_requires_samples(): + with pytest.raises(ValueError): + percentiles_us([]) + + +def test_bm2_report_marks_unavailable_candidate_modes_null(): + report = build_bm2_report( + model="tiny/fixture", + saturated_tensor_count=6144, + window_us=500.0, + per_expert_samples_ns=[900_000] * 32, + batched_samples_ns=None, + cpp_internal_samples_ns=None, + warmup=20, + iterations=200, + ) + assert report["benchmark"] == "BM2" + assert report["saturated_tensor_count"] == 6144 + assert report["window_us"] == 500.0 + assert report["warmup"] == 20 + assert report["iterations"] == 200 + assert report["modes"][PYTHON_PER_EXPERT]["p50"] == pytest.approx(900.0) + # Unavailable candidate modes are null, never zero. + assert report["modes"][BATCHED_PYBIND] is None + assert report["modes"][CPP_INTERNAL] is None + assert report["medians_us"][BATCHED_PYBIND] is None + assert report["candidate_required"] is True + assert report["ship_batched"] is False + + +def test_bm2_report_ships_when_batched_mode_present_and_fast(): + report = build_bm2_report( + model="tiny/fixture", + saturated_tensor_count=6144, + window_us=500.0, + per_expert_samples_ns=[900_000] * 32, + batched_samples_ns=[300_000] * 32, + cpp_internal_samples_ns=[250_000] * 32, + warmup=20, + iterations=200, + ) + assert report["modes"][BATCHED_PYBIND]["p50"] == pytest.approx(300.0) + assert report["medians_us"][PYTHON_PER_EXPERT] == pytest.approx(900.0) + assert report["candidate_required"] is True + assert report["ship_batched"] is True From 9d0f90181c7ae51001d838e6986b7155775eb5b0 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sat, 15 Aug 2026 11:49:51 +0000 Subject: [PATCH 7/9] perf(prefetch): batch route-ahead issuance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 8 of the PD-DFlash serving plan (candidate hop 1), retained because BM2 passed on offloaded gpt-oss-20b: python-per-expert issuance p50=1132 µs vs batched-pybind p50=101 µs over a 768-tensor saturated block (~11× fewer pybind crossings), so per-expert issuance is exposed and the batched call hides it. - core/prefetch: add ArcherPrefetchHandle::EnqueuePrefetchTensors(tensor_ids, priority=1), which constructs and enqueues one Task per tensor entirely in C++, preserving input order and node default devices (mirrors EnqueuePrefetch). - core/python: bind prefetch_tensors -> EnqueuePrefetchTensors and retire the old no-op PrefetchTensors(request_id, buffer) binding; enqueue_prefetch retained. - expert_prefetcher: prefetch_experts_list issues one batched prefetch_tensors call when the engine exposes it, else the byte-for-byte per-expert fallback; empty input is a no-op. - tests: batch/fallback/empty coverage in test_speculative_prefetch.py; wire assertions made mechanism-agnostic (batched call carries the same ordered ids); opt-in native GPU smoke test (single module-scoped offload load). --- core/prefetch/archer_prefetch_handle.cpp | 14 +++ core/prefetch/archer_prefetch_handle.h | 2 + core/python/py_archer_prefetch.cpp | 3 +- moe_infinity/memory/expert_prefetcher.py | 6 + .../python/dflash/test_prefetch_native_gpu.py | 119 ++++++++++++++++++ tests/python/dflash/test_route_ahead_wire.py | 48 ++++--- .../dflash/test_speculative_prefetch.py | 54 ++++++++ 7 files changed, 227 insertions(+), 19 deletions(-) create mode 100644 tests/python/dflash/test_prefetch_native_gpu.py diff --git a/core/prefetch/archer_prefetch_handle.cpp b/core/prefetch/archer_prefetch_handle.cpp index 9ec72c21..b3bae148 100644 --- a/core/prefetch/archer_prefetch_handle.cpp +++ b/core/prefetch/archer_prefetch_handle.cpp @@ -248,6 +248,20 @@ void ArcherPrefetchHandle::EnqueuePrefetch(const uint32_t tensor_id, kTaskPool->EnqueueTask(task); } +void ArcherPrefetchHandle::EnqueuePrefetchTensors( + const std::vector& tensor_ids, std::uint32_t priority) { + for (std::uint32_t tensor_id : tensor_ids) { + auto node = kTopologyHandle->GetNodeFromTensorID(tensor_id); + auto task = std::make_shared(); + task->priority = priority; + task->node = node; + task->on_demand = false; + task->src_device = node->device; + task->dst_device = node->default_device; + kTaskPool->EnqueueTask(task); + } +} + void ArcherPrefetchHandle::FetchTensors( std::uint64_t& request_id, const std::vector& buffer) { // std::vector candidates; diff --git a/core/prefetch/archer_prefetch_handle.h b/core/prefetch/archer_prefetch_handle.h index 4d04622a..c18a8f84 100644 --- a/core/prefetch/archer_prefetch_handle.h +++ b/core/prefetch/archer_prefetch_handle.h @@ -28,6 +28,8 @@ 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); 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/python/py_archer_prefetch.cpp b/core/python/py_archer_prefetch.cpp index fe5eadfe..6e6b016c 100644 --- a/core/python/py_archer_prefetch.cpp +++ b/core/python/py_archer_prefetch.cpp @@ -86,7 +86,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("get_node_default_device", &ArcherPrefetchHandle::GetNodeDefaultDevice) .def("get_node_device", &ArcherPrefetchHandle::GetNodeDevice) - .def("prefetch_tensors", &ArcherPrefetchHandle::PrefetchTensors) + .def("prefetch_tensors", &ArcherPrefetchHandle::EnqueuePrefetchTensors, + py::arg("tensor_ids"), py::arg("priority") = 1) .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 6afa911e..637517b7 100644 --- a/moe_infinity/memory/expert_prefetcher.py +++ b/moe_infinity/memory/expert_prefetcher.py @@ -55,6 +55,12 @@ def prefetch_experts_list(self, layer_id: int, expert_list: List[int]): tensor_ids = [] for j in expert_list: tensor_ids.append(self.expert_tensor_map[(layer_id, j)]) + if not tensor_ids: + return + batched_issue = getattr(self.archer_engine, "prefetch_tensors", None) + if callable(batched_issue): + batched_issue(tensor_ids) + return for tensor_id in tensor_ids: gpu_id = self.archer_engine.get_node_default_device([tensor_id]) self.archer_engine.enqueue_prefetch(tensor_id, gpu_id) diff --git a/tests/python/dflash/test_prefetch_native_gpu.py b/tests/python/dflash/test_prefetch_native_gpu.py new file mode 100644 index 00000000..b6d953e7 --- /dev/null +++ b/tests/python/dflash/test_prefetch_native_gpu.py @@ -0,0 +1,119 @@ +# Copyright (c) EfficientMoE. +# SPDX-License-Identifier: Apache-2.0 + +# EfficientMoE Team + +"""Opt-in native-extension smoke test for the batched ``prefetch_tensors`` API. + +Task 8 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md`` +(candidate hop 1). Verifies that the rebuilt ``moe_infinity._store`` exposes the +batched ``prefetch_handle.prefetch_tensors(tensor_ids, priority=1)`` binding and +that it enqueues a saturated ``E_l x L`` block in one call without raising -- the +pre-Task-8 no-op ``prefetch_tensors(request_id, buffer)`` binding would reject a +single positional tensor-id list, so a passing call proves the new native API is +built and wired. + +The offloaded target is loaded exactly once via a module-scoped fixture: the +native archer engine keeps process-global topology/task-pool state that does not +survive a second in-process offload load, so each test must share one engine. + +Opt-in via ``MOE_DFLASH_SERVING_GPU=1`` with the offloaded target present in the +HF cache. Without the gate this collects and skips cleanly: no CUDA, no model +load, no filesystem, no network at import time. +""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Optional + +import pytest +import torch + +TARGET_REPO = os.environ.get("MOE_PREFETCH_NATIVE_MODEL", "openai/gpt-oss-20b") + + +def _hf_home() -> Path: + for var in ("HF_HOME", "HUGGINGFACE_HUB_CACHE", "XDG_CACHE_HOME"): + val = os.environ.get(var) + if val: + base = Path(val) + return base / "hub" if var == "XDG_CACHE_HOME" else base + return Path.home() / ".cache" / "huggingface" + + +def _checkpoint_present(repo: str) -> bool: + hub = _hf_home() + hub = hub if hub.name == "hub" else hub / "hub" + return (hub / f"models--{repo.replace('/', '--')}").is_dir() + + +def _skip_reason() -> Optional[str]: + if not os.environ.get("MOE_DFLASH_SERVING_GPU"): + return "MOE_DFLASH_SERVING_GPU unset (opt-in native prefetch smoke)" + if not torch.cuda.is_available(): + return "CUDA unavailable (native prefetch smoke)" + if not _checkpoint_present(TARGET_REPO): + return f"checkpoint not present in $HF_HOME: {TARGET_REPO}" + return None + + +SKIP_REASON = _skip_reason() +pytestmark = pytest.mark.skipif( + SKIP_REASON is not None, reason=SKIP_REASON or "gpu-gated" +) + + +@pytest.fixture(scope="module") +def offloaded_prefetcher(): + from moe_infinity import MoE + + offload = os.environ.get( + "MOE_PREFETCH_NATIVE_OFFLOAD", "/tmp/opencode/moe-offload/gpt-oss-20b" + ) + os.makedirs(offload, exist_ok=True) + ratio = float(os.environ.get("MOE_DFLASH_MEM_RATIO", "0.2")) + model = MoE( + TARGET_REPO, + {"offload_path": offload, "device_memory_ratio": ratio}, + ) + prefetcher = model.engine.expert_prefetcher + assert prefetcher is not None and prefetcher.archer_engine is not None + yield prefetcher + + +def _saturated_ids(prefetcher) -> list[int]: + return [tid for _key, tid in sorted(prefetcher.expert_tensor_map.items())] + + +def test_native_batched_prefetch_tensors_issues_saturated_block( + offloaded_prefetcher, +) -> None: + engine = offloaded_prefetcher.archer_engine + tensor_ids = _saturated_ids(offloaded_prefetcher) + assert tensor_ids, "no offloaded expert tensors to issue" + + assert engine.prefetch_tensors(tensor_ids) is None + assert engine.prefetch_tensors(tensor_ids, 1) is None + + +def test_native_batched_prefetch_tensors_empty_is_noop( + offloaded_prefetcher, +) -> None: + assert offloaded_prefetcher.archer_engine.prefetch_tensors([]) is None + + +def test_native_batched_prefetch_experts_list_uses_batched_path( + offloaded_prefetcher, +) -> None: + 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 + ) + offloaded_prefetcher.prefetch_experts_list(some_layer, experts) diff --git a/tests/python/dflash/test_route_ahead_wire.py b/tests/python/dflash/test_route_ahead_wire.py index 38021bbd..561f1e07 100644 --- a/tests/python/dflash/test_route_ahead_wire.py +++ b/tests/python/dflash/test_route_ahead_wire.py @@ -121,6 +121,21 @@ def _enqueued_experts(executor) -> list[int]: ) +def _issued_tensor_ids(engine) -> list[int]: + # Mechanism-agnostic route-ahead issuance readout: a batched + # ``prefetch_tensors([...])`` call carries the same ordered tensor ids the + # per-expert ``enqueue_prefetch`` fallback would, so flatten the batched + # calls in order when present and fall back otherwise. + batched = getattr(engine, "prefetch_tensors", None) + batched_calls = getattr(batched, "call_args_list", None) + if batched_calls: + issued: list[int] = [] + for call in batched_calls: + issued.extend(call.args[0]) + return issued + return [call.args[0] for call in engine.enqueue_prefetch.call_args_list] + + # --------------------------------------------------------------------------- # (a) context active -> exact-union pin + prefetch for the current layer # --------------------------------------------------------------------------- @@ -171,13 +186,7 @@ def test_active_context_falls_back_to_executor_prefetcher(): engine.replace_cache_candidates.assert_called_once_with( [300, 301, 302, 305, 307] ) - assert [c.args[0] for c in engine.enqueue_prefetch.call_args_list] == [ - 300, - 301, - 302, - 305, - 307, - ] + assert _issued_tensor_ids(engine) == [300, 301, 302, 305, 307] assert prefetcher._last_speculative_prediction == set(UNION) trigger_spy.assert_not_called() assert executor._pending_prefetch == (prefetcher, LAYER_ID, UNION, None) @@ -186,13 +195,7 @@ def test_active_context_falls_back_to_executor_prefetcher(): # recorded prediction IS the actual union; nothing further is enqueued. executor.wait_dispatch_local() engine.replace_cache_candidates.assert_called_once() - assert [c.args[0] for c in engine.enqueue_prefetch.call_args_list] == [ - 300, - 301, - 302, - 305, - 307, - ] + assert _issued_tensor_ids(engine) == [300, 301, 302, 305, 307] assert prefetcher._last_speculative_prediction == set() @@ -429,10 +432,19 @@ def test_consecutive_dispatches_each_pin_exactly_one_layer(): # No call ever mixes tensor ids from two layers (id = layer * 100 + e). for call in pin_calls: assert len({tensor_id // 100 for tensor_id in call.args[0]}) == 1 - # Enqueues stay per-layer single-layered as well. - assert [ - call.args[0] for call in engine.enqueue_prefetch.call_args_list - ] == [300, 301, 302, 305, 307, 400, 401, 402, 405, 407] + # Issuances stay per-layer single-layered as well (one batched call/layer). + assert _issued_tensor_ids(engine) == [ + 300, + 301, + 302, + 305, + 307, + 400, + 401, + 402, + 405, + 407, + ] def _make_gpt_oss_mlp(): diff --git a/tests/python/dflash/test_speculative_prefetch.py b/tests/python/dflash/test_speculative_prefetch.py index ffcec01d..1f67f3a6 100644 --- a/tests/python/dflash/test_speculative_prefetch.py +++ b/tests/python/dflash/test_speculative_prefetch.py @@ -57,6 +57,16 @@ def _make_prefetcher(num_layers: int = 8, num_experts: int = 8): def _enqueued_tensor_ids(engine: MagicMock) -> list[int]: + # Mechanism-agnostic: a batched ``prefetch_tensors([...])`` issuance carries + # the same ordered ids as the per-expert ``enqueue_prefetch`` fallback, so + # flatten the batched calls when present and fall back otherwise. + batched = getattr(engine, "prefetch_tensors", None) + batched_calls = getattr(batched, "call_args_list", None) + if batched_calls: + issued: list[int] = [] + for call in batched_calls: + issued.extend(call.args[0]) + return issued return [call.args[0] for call in engine.enqueue_prefetch.call_args_list] @@ -148,3 +158,47 @@ def test_both_none_raises_value_error(): prefetcher, _engine = _make_prefetcher() with pytest.raises(ValueError, match="router_logits"): prefetcher.speculative_prefetch(0) + + +def _make_prefetcher_without_batch(num_layers: int = 8, num_experts: int = 8): + prefetcher = ExpertPrefetcher.__new__(ExpertPrefetcher) + prefetcher.num_layers = num_layers + prefetcher.num_experts = num_experts + engine = MagicMock( + spec=[ + "get_node_default_device", + "enqueue_prefetch", + "replace_cache_candidates", + ] + ) + engine.get_node_default_device.return_value = 0 + prefetcher.archer_engine = engine + prefetcher.expert_tensor_map = { + (layer, expert): layer * 100 + expert + for layer in range(num_layers) + for expert in range(num_experts) + } + prefetcher._last_speculative_prediction = set() + return prefetcher, engine + + +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.enqueue_prefetch.assert_not_called() + + +def test_prefetch_experts_list_falls_back_to_per_expert_without_batch_api(): + prefetcher, engine = _make_prefetcher_without_batch( + num_layers=8, num_experts=8 + ) + prefetcher.prefetch_experts_list(3, [3, 1, 7]) + assert _enqueued_tensor_ids(engine) == [303, 301, 307] + + +def test_prefetch_experts_list_empty_batch_calls_neither_path(): + prefetcher, engine = _make_prefetcher(num_layers=8, num_experts=8) + prefetcher.prefetch_experts_list(3, []) + engine.prefetch_tensors.assert_not_called() + engine.enqueue_prefetch.assert_not_called() From 4b3d1827588cefe7e6a2a561dd9b21b310e957c1 Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sat, 15 Aug 2026 12:31:24 +0000 Subject: [PATCH 8/9] bench(dflash): add BM3 priority-band + BM4 overlap harnesses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 9/10 measurement harnesses (design §10), no C++ shipped: * bench_prefetch_priority.py -- BM3 three-way route-ahead prefetch priority ablation (background/route-ahead/on-demand) over median exposed-fetch seconds and tokens/s, with a pure bm3_decision ship gate: ship the dedicated route-ahead band iff it lowers exposed fetch vs background, preserves tokens/s, and on-demand stays fastest (no priority inversion). * parse_overlap.py -- BM4 expert-H2D / compute overlap from an nsys trace; pure interval arithmetic apportions memcpy bytes by overlapped-duration fraction against the union of the draft/router/verify NVTX ranges. * run_phase_c.sh -- one-command Phase-C launcher for the full BM3/BM4/BM5 matrix on both required targets. * test_prefetch_perf_reports.py -- CPU-only BM3 + BM4 decision-rule tests. --- benchmarks/dflash/bench_prefetch_priority.py | 458 ++++++++++++++++++ benchmarks/dflash/parse_overlap.py | 325 +++++++++++++ benchmarks/dflash/run_phase_c.sh | 144 ++++++ .../dflash/test_prefetch_perf_reports.py | 266 ++++++++++ 4 files changed, 1193 insertions(+) create mode 100644 benchmarks/dflash/bench_prefetch_priority.py create mode 100644 benchmarks/dflash/parse_overlap.py create mode 100755 benchmarks/dflash/run_phase_c.sh diff --git a/benchmarks/dflash/bench_prefetch_priority.py b/benchmarks/dflash/bench_prefetch_priority.py new file mode 100644 index 00000000..e5c8b038 --- /dev/null +++ b/benchmarks/dflash/bench_prefetch_priority.py @@ -0,0 +1,458 @@ +# Copyright (c) EfficientMoE. +# SPDX-License-Identifier: Apache-2.0 + +# EfficientMoE Team + +"""BM3 -- route-ahead prefetch priority-band ablation (design §10). + +Task 9 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md`` +(candidate hop 2). A three-way ablation of the priority at which route-ahead +prefetch work is enqueued into the native task pool: + +* ``background`` (priority ``2``) -- the ordinary background-prefetch band, i.e. + route-ahead work competes with plain AR prefetch; +* ``route-ahead`` (priority ``1``) -- a *dedicated* band serviced ahead of + background prefetch but behind on-demand misses (the candidate under test); +* ``on-demand`` (priority ``0``) -- the same band real on-demand misses use; + measured only to prove it stays the fastest service class (a shipped + route-ahead band must never invert on-demand). + +For each arm the runner measures, over identical seeded requests with cache +state reset between variants, the median *exposed-fetch seconds* (time an +on-demand expert fetch is exposed, i.e. not hidden behind draft+verify compute) +and the median *tokens/s*. The candidate dedicated band ships (design §10 / +plan Task 9 Step 1) iff **all three** hold: + +1. route-ahead has *lower* exposed fetch than default background; +2. route-ahead does not *reduce* tokens/s versus default background; and +3. on-demand remains the fastest service class (no priority inversion). + +The gate must pass for *both* required targets; a single model that shows no +improvement, a throughput regression, or a priority inversion removes the +candidate and keeps only this benchmark. + +Import-safe by construction: torch and moe_infinity are imported lazily inside +the GPU runner, so ``bm3_decision`` / ``priority_arm`` / ``median`` / +``build_bm3_report`` (and their tests) are pure-CPU and never initialise CUDA. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +from dataclasses import dataclass +from typing import Any, Dict, List, Mapping, Optional, Sequence + +# Native priority bands (mirror ``core/prefetch/task_scheduler.h``); lower value +# is serviced first in ``ArcherTaskPool::GPUThreadFunc``. +ON_DEMAND = 0 +ROUTE_AHEAD = 1 +BACKGROUND = 2 + +# Ablation arm labels, ordered from lowest to highest *service* priority. +PRIORITY_BANDS = ("background", "route-ahead", "on-demand") + +_ARM_PRIORITY: Dict[str, int] = { + "background": BACKGROUND, + "route-ahead": ROUTE_AHEAD, + "on-demand": ON_DEMAND, +} + + +@dataclass(frozen=True) +class PriorityArm: + """Median exposed-fetch seconds and tokens/s for one priority band.""" + + exposed_fetch_seconds: float + tokens_per_second: float + + +@dataclass(frozen=True) +class Bm3Decision: + """The BM3 ship gate over the three ablation arms (design §10).""" + + default: PriorityArm + route_ahead: PriorityArm + on_demand: PriorityArm + exposed_fetch_improved: bool + throughput_preserved: bool + on_demand_fastest: bool + ship_priority_band: bool + + +def _finite_non_negative(name: str, value: Any) -> float: + number = float(value) + if not math.isfinite(number) or number < 0.0: + raise ValueError(f"{name} must be finite and >= 0; got {value!r}") + return number + + +def priority_arm( + *, exposed_fetch_seconds: float, tokens_per_second: float +) -> PriorityArm: + """Build a validated ``PriorityArm`` from measured medians.""" + return PriorityArm( + exposed_fetch_seconds=_finite_non_negative( + "exposed_fetch_seconds", exposed_fetch_seconds + ), + tokens_per_second=_finite_non_negative( + "tokens_per_second", tokens_per_second + ), + ) + + +def bm3_decision( + *, + default: PriorityArm, + route_ahead: PriorityArm, + on_demand: PriorityArm, +) -> Bm3Decision: + """Evaluate the BM3 ship gate over the three ablation arms. + + Ships the dedicated route-ahead band iff it strictly lowers exposed fetch + versus default background, does not reduce tokens/s (inclusive at + equality), and on-demand remains the fastest service class (its exposed + fetch is no worse than both other arms). + """ + exposed_fetch_improved = ( + route_ahead.exposed_fetch_seconds < default.exposed_fetch_seconds + ) + throughput_preserved = ( + route_ahead.tokens_per_second >= default.tokens_per_second + ) + on_demand_fastest = ( + on_demand.exposed_fetch_seconds <= route_ahead.exposed_fetch_seconds + and on_demand.exposed_fetch_seconds <= default.exposed_fetch_seconds + ) + ship = exposed_fetch_improved and throughput_preserved and on_demand_fastest + return Bm3Decision( + default=default, + route_ahead=route_ahead, + on_demand=on_demand, + exposed_fetch_improved=exposed_fetch_improved, + throughput_preserved=throughput_preserved, + on_demand_fastest=on_demand_fastest, + ship_priority_band=ship, + ) + + +def median(samples: Sequence[float]) -> float: + """Median of ``samples`` (mean of the two central values for even counts).""" + if not samples: + raise ValueError("median requires at least one sample") + ordered = sorted(float(s) for s in samples) + count = len(ordered) + mid = count // 2 + if count % 2 == 1: + return ordered[mid] + return (ordered[mid - 1] + ordered[mid]) / 2.0 + + +def _arm_from_samples(samples: Mapping[str, Sequence[float]]) -> PriorityArm: + return priority_arm( + exposed_fetch_seconds=median(samples["exposed_fetch_seconds"]), + tokens_per_second=median(samples["tokens_per_second"]), + ) + + +def build_bm3_report( + *, + models: Sequence[str], + arms: Mapping[str, Mapping[str, Sequence[float]]], + extra: Optional[Mapping[str, Any]] = None, +) -> Dict[str, Any]: + """Assemble the BM3 report and its ship-gate verdict from raw samples. + + ``arms`` maps each of ``background``/``route-ahead``/``on-demand`` to a + mapping with ``exposed_fetch_seconds`` and ``tokens_per_second`` sample + lists (one entry per repetition). All three arms are required. + """ + missing = [band for band in PRIORITY_BANDS if band not in arms] + if missing: + raise ValueError(f"BM3 report missing arms: {', '.join(missing)}") + + per_arm = {band: _arm_from_samples(arms[band]) for band in PRIORITY_BANDS} + decision = bm3_decision( + default=per_arm["background"], + route_ahead=per_arm["route-ahead"], + on_demand=per_arm["on-demand"], + ) + report: Dict[str, Any] = { + "benchmark": "BM3", + "models": list(models), + "priorities": {band: _ARM_PRIORITY[band] for band in PRIORITY_BANDS}, + "medians": { + band: { + "exposed_fetch_seconds": per_arm[band].exposed_fetch_seconds, + "tokens_per_second": per_arm[band].tokens_per_second, + } + for band in PRIORITY_BANDS + }, + "repetitions": { + band: len(arms[band]["exposed_fetch_seconds"]) + for band in PRIORITY_BANDS + }, + "exposed_fetch_improved": decision.exposed_fetch_improved, + "throughput_preserved": decision.throughput_preserved, + "on_demand_fastest": decision.on_demand_fastest, + "ship_priority_band": decision.ship_priority_band, + } + if extra: + report.update(dict(extra)) + return report + + +# --------------------------------------------------------------------------- +# GPU ablation runner (lazily imports torch / moe_infinity) +# --------------------------------------------------------------------------- + + +def run_priority_ablation( + *, + model_repo: str, + draft_repo: str, + offload_path: str, + device_memory_ratio: float, + repetitions: int, + block_size: int, + requests: int, + warmup_rounds: int, + seed: int, +) -> Dict[str, Any]: + """Measure the three priority arms for one offloaded target. + + Requires a genuinely offloaded target and its DFlash draft. For each arm we + force route-ahead issuance onto that native band (via + ``ExpertPrefetcher.route_ahead_priority``), reset cache state, run identical + seeded requests, and record exposed-fetch seconds and tokens/s. Route-ahead + issuance must expose an explicit priority for this ablation to be + meaningful; a build without the priority-band plumbing raises. + """ + import time + + import torch + + import moe_infinity._v4_fp4 # noqa: F401 (assert native FP4 path present) + from moe_infinity import MoE + from moe_infinity.memory.expert_prefetcher import ExpertPrefetcher + from moe_infinity.spec_decode import DFlashSpeculator + + if not hasattr(ExpertPrefetcher, "route_ahead_priority") and not any( + hasattr(ExpertPrefetcher, attr) + for attr in ("route_ahead_priority", "_route_ahead_priority") + ): + raise RuntimeError( + "ExpertPrefetcher exposes no route_ahead_priority knob; the " + "priority-band candidate (plan Task 9 Step 4) is not built in, so " + "the ablation cannot distinguish the route-ahead band" + ) + + model = MoE( + model_repo, + { + "offload_path": offload_path, + "device_memory_ratio": device_memory_ratio, + }, + ) + engine = model.engine + prefetcher = engine.expert_prefetcher + if prefetcher is None or prefetcher.archer_engine is None: + raise RuntimeError( + "loaded target has no offloaded ExpertPrefetcher/archer_engine; " + "lower --device-memory-ratio below 0.9 so experts offload" + ) + speculator = DFlashSpeculator(model, draft_repo) + enable = getattr(speculator, "enable_route_ahead_stats", None) + if callable(enable): + enable() + + from benchmarks.dflash._serving_measure import ( + _deterministic_prompt_ids, + _greedy_generate, + ) + + prompt_ids = _deterministic_prompt_ids(model, model_repo) + tokens_per_request = max(block_size * 4, 32) + + def _reset_cache() -> None: + reset = getattr(prefetcher.archer_engine, "clean_up_resources", None) + if callable(reset): + try: + reset() + except Exception: + pass + torch.cuda.synchronize() + + def _measure_once() -> Dict[str, float]: + torch.cuda.synchronize() + started = time.perf_counter() + generated = 0 + for _ in range(max(1, requests)): + out = _greedy_generate( + model, prompt_ids, speculator, tokens_per_request + ) + generated += len(out) + torch.cuda.synchronize() + elapsed = max(time.perf_counter() - started, 1e-9) + exposed = _exposed_fetch_seconds(engine, speculator) + return { + "exposed_fetch_seconds": exposed, + "tokens_per_second": generated / elapsed, + } + + arms: Dict[str, Dict[str, List[float]]] = {} + for band in PRIORITY_BANDS: + torch.manual_seed(seed) + prefetcher.route_ahead_priority = _ARM_PRIORITY[band] + # Warm up this arm. + for _ in range(max(0, warmup_rounds)): + _greedy_generate(model, prompt_ids, speculator, max(1, block_size)) + exposed_samples: List[float] = [] + tps_samples: List[float] = [] + for _ in range(max(1, repetitions)): + _reset_cache() + sample = _measure_once() + exposed_samples.append(sample["exposed_fetch_seconds"]) + tps_samples.append(sample["tokens_per_second"]) + arms[band] = { + "exposed_fetch_seconds": exposed_samples, + "tokens_per_second": tps_samples, + } + + return build_bm3_report( + models=[model_repo], + arms=arms, + extra={ + "offload_path": offload_path, + "device_memory_ratio": device_memory_ratio, + "draft": draft_repo, + "block_size": block_size, + "requests": requests, + "seed": seed, + }, + ) + + +def _exposed_fetch_seconds(engine: Any, speculator: Any) -> float: + """Best-effort exposed on-demand fetch seconds for the last run. + + Prefers a native/instrumented accessor; falls back to the route-ahead + stats' exposed-fetch term. Returns 0.0 only when nothing is instrumented, + in which case the report's ``warnings`` should be consulted. + """ + for source in (engine, getattr(engine, "expert_prefetcher", None)): + if source is None: + continue + for name in ( + "exposed_fetch_seconds", + "get_exposed_fetch_seconds", + "on_demand_fetch_seconds", + ): + value = getattr(source, 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) + stats = getattr(speculator, "route_ahead_stats", None) + if stats is not None: + snapshot = stats.as_dict() + exposed = snapshot.get("exposed_fetch_seconds") + if isinstance(exposed, (int, float)) and not isinstance(exposed, bool): + return float(exposed) + return 0.0 + + +def _run_all_models(args: argparse.Namespace) -> Dict[str, Any]: + if not (len(args.models) == len(args.drafts) == len(args.offload_dirs)): + raise SystemExit( + "--models, --drafts, and --offload-dirs must have equal length" + ) + per_model: List[Dict[str, Any]] = [] + ship_all = True + for model_repo, draft_repo, offload in zip( + args.models, args.drafts, args.offload_dirs + ): + report = run_priority_ablation( + model_repo=model_repo, + draft_repo=draft_repo, + offload_path=offload, + device_memory_ratio=args.device_memory_ratio, + repetitions=args.repetitions, + block_size=args.block_size, + requests=args.requests, + warmup_rounds=args.warmup_rounds, + seed=args.seed, + ) + ship_all = ship_all and bool(report["ship_priority_band"]) + per_model.append(report) + return { + "benchmark": "BM3", + "per_model": per_model, + # The candidate ships only when *every* required target passes. + "ship_priority_band": ship_all and len(per_model) > 0, + } + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m benchmarks.dflash.bench_prefetch_priority", + description="BM3 route-ahead prefetch priority-band ablation.", + ) + parser.add_argument("--models", nargs="+", required=True) + parser.add_argument("--drafts", nargs="+", required=True) + parser.add_argument("--offload-dirs", nargs="+", required=True) + parser.add_argument( + "--priorities", + nargs="+", + default=list(PRIORITY_BANDS), + choices=list(PRIORITY_BANDS), + help="informational; the runner always sweeps all three bands", + ) + parser.add_argument("--repetitions", type=int, default=10) + parser.add_argument("--block-size", type=int, default=16) + parser.add_argument("--requests", type=int, default=16) + parser.add_argument("--warmup-rounds", type=int, default=5) + parser.add_argument("--seed", type=int, default=1408) + parser.add_argument("--device-memory-ratio", type=float, default=0.85) + parser.add_argument("--output", required=True) + args = parser.parse_args(argv) + + if not os.environ.get("MOE_DFLASH_SERVING_GPU"): + parser.error( + "MOE_DFLASH_SERVING_GPU must be set (opt-in GPU priority ablation)" + ) + + report = _run_all_models(args) + os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True) + with open(args.output, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, sort_keys=True) + handle.write("\n") + json.dump(report, sys.stdout, indent=2, sort_keys=True) + sys.stdout.write("\n") + return 0 + + +__all__ = [ + "ON_DEMAND", + "ROUTE_AHEAD", + "BACKGROUND", + "PRIORITY_BANDS", + "PriorityArm", + "Bm3Decision", + "priority_arm", + "bm3_decision", + "median", + "build_bm3_report", + "run_priority_ablation", + "main", +] + + +if __name__ == "__main__": # pragma: no cover - CLI entry + raise SystemExit(main()) diff --git a/benchmarks/dflash/parse_overlap.py b/benchmarks/dflash/parse_overlap.py new file mode 100644 index 00000000..76e354ed --- /dev/null +++ b/benchmarks/dflash/parse_overlap.py @@ -0,0 +1,325 @@ +# Copyright (c) EfficientMoE. +# SPDX-License-Identifier: Apache-2.0 + +# EfficientMoE Team + +"""BM4 -- expert-H2D / compute overlap from an nsys trace (design §10). + +Task 10 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md``. +This is the ground truth for whether offloaded-expert fetch is *hidden*: the +fraction of expert host->device (H2D) memcpy **bytes** that overlap the DFlash +draft/router/verify compute NVTX ranges the serving runner emits. + +Two layers: + +* pure interval arithmetic (``compute_overlap``) apportions each memcpy's bytes + by its overlapped-duration fraction against the *union* of the compute ranges + -- a partially hidden copy contributes a proportional slice of its bytes, and + overlapping compute ranges are unioned so overlap can never exceed 100%; and +* an nsys CSV reader (``parse_nsys_rep``) that runs + ``nsys stats --report cuda_gpu_trace,nvtx_pushpop_trace --format csv`` on a + ``.nsys-rep`` and feeds the parsed H2D memcpys and NVTX ranges into + ``compute_overlap``. + +The pure layer is import-safe and unit-tested off-hardware; only +``parse_nsys_rep`` shells out to ``nsys``. +""" + +from __future__ import annotations + +import argparse +import csv +import io +import json +import subprocess +import sys +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Sequence, Tuple + +# NVTX ranges that count as compute able to hide an expert fetch (design §10); +# these are three of the five ranges the serving runner emits. ``expert_h2d`` +# and ``route_ahead_issue`` are issuance/transfer ranges, not hiding compute. +DEFAULT_COMPUTE_RANGES: Tuple[str, ...] = ( + "dflash_draft", + "route_ahead_router", + "target_verify", +) + + +@dataclass(frozen=True) +class Memcpy: + """One H2D memcpy interval with its transferred byte count.""" + + start: float + end: float + bytes: int + + +@dataclass(frozen=True) +class NvtxRange: + """One NVTX push/pop range interval.""" + + name: str + start: float + end: float + + +@dataclass(frozen=True) +class OverlapResult: + """Aggregate BM4 overlap over a set of memcpys and compute ranges.""" + + total_h2d_bytes: int + overlapped_h2d_bytes: float + overlap_fraction: float + per_memcpy_fraction: Tuple[float, ...] + + +def _merge_intervals( + intervals: Sequence[Tuple[float, float]], +) -> List[Tuple[float, float]]: + """Union a set of ``(start, end)`` intervals into disjoint spans.""" + spans = sorted( + (float(a), float(b)) for a, b in intervals if float(b) > float(a) + ) + merged: List[Tuple[float, float]] = [] + for start, end in spans: + if merged and start <= merged[-1][1]: + prev_start, prev_end = merged[-1] + merged[-1] = (prev_start, max(prev_end, end)) + else: + merged.append((start, end)) + return merged + + +def _overlap_duration( + interval: Tuple[float, float], spans: Sequence[Tuple[float, float]] +) -> float: + """Total length of ``interval`` covered by the disjoint ``spans``.""" + start, end = interval + covered = 0.0 + for span_start, span_end in spans: + lo = max(start, span_start) + hi = min(end, span_end) + if hi > lo: + covered += hi - lo + return covered + + +def compute_overlap( + memcpys: Sequence[Memcpy], + ranges: Sequence[NvtxRange], + compute_ranges: Sequence[str] = DEFAULT_COMPUTE_RANGES, +) -> OverlapResult: + """Apportion H2D bytes hidden behind the union of compute ranges. + + Each memcpy contributes ``bytes * (overlapped_duration / duration)`` hidden + bytes; a zero-duration memcpy contributes its full bytes only when its + start instant lies inside a compute span, and nothing otherwise. The + aggregate fraction is hidden bytes over total bytes (``0`` when no bytes). + """ + compute_names = set(compute_ranges) + spans = _merge_intervals( + [(r.start, r.end) for r in ranges if r.name in compute_names] + ) + + total_bytes = 0 + overlapped_bytes = 0.0 + fractions: List[float] = [] + for memcpy in memcpys: + nbytes = int(memcpy.bytes) + total_bytes += nbytes + duration = float(memcpy.end) - float(memcpy.start) + if duration > 0.0: + fraction = ( + _overlap_duration( + (float(memcpy.start), float(memcpy.end)), spans + ) + / duration + ) + else: + instant = float(memcpy.start) + fraction = ( + 1.0 if any(lo <= instant <= hi for lo, hi in spans) else 0.0 + ) + fraction = min(max(fraction, 0.0), 1.0) + fractions.append(fraction) + overlapped_bytes += nbytes * fraction + + overlap_fraction = ( + overlapped_bytes / total_bytes if total_bytes > 0 else 0.0 + ) + return OverlapResult( + total_h2d_bytes=total_bytes, + overlapped_h2d_bytes=overlapped_bytes, + overlap_fraction=overlap_fraction, + per_memcpy_fraction=tuple(fractions), + ) + + +# --------------------------------------------------------------------------- +# nsys CSV reader (shells out to the ``nsys`` CLI) +# --------------------------------------------------------------------------- + + +def _nsys_stats_csv(rep_path: str, report: str) -> str: + proc = subprocess.run( + [ + "nsys", + "stats", + "--report", + report, + "--format", + "csv", + "--force-export=true", + rep_path, + ], + check=True, + capture_output=True, + text=True, + ) + return proc.stdout + + +def _read_csv_rows(text: str) -> List[Dict[str, str]]: + rows: List[Dict[str, str]] = [] + reader = csv.reader(io.StringIO(text)) + header: Optional[List[str]] = None + for record in reader: + if not record: + continue + if header is None: + # nsys prefixes stats blocks with a title line before the header; + # the header is the first row that contains a recognised column. + lowered = [c.strip().lower() for c in record] + if any("duration" in c or "name" in c for c in lowered): + header = [c.strip() for c in record] + continue + if len(record) < len(header): + continue + rows.append(dict(zip(header, record))) + return rows + + +def _to_ns(value: str) -> float: + return float(str(value).replace(",", "").strip()) + + +def _column(row: Dict[str, str], *candidates: str) -> Optional[str]: + lowered = {k.lower(): k for k in row} + for candidate in candidates: + key = lowered.get(candidate.lower()) + if key is not None and row[key] != "": + return row[key] + return None + + +def parse_memcpys(csv_text: str) -> List[Memcpy]: + """Parse H2D memcpy intervals from an nsys ``cuda_gpu_trace`` CSV.""" + memcpys: List[Memcpy] = [] + for row in _read_csv_rows(csv_text): + name = (_column(row, "Name") or "").strip() + if "memcpy" not in name.lower() and "htod" not in name.lower(): + continue + if "hto" not in name.lower().replace("-", "") and "htod" not in ( + name.lower() + ): + if "host-to-device" not in name.lower(): + continue + start = _column(row, "Start (ns)", "Start") + duration = _column(row, "Duration (ns)", "Duration") + nbytes = _column(row, "Bytes (MB)", "Bytes", "Size (MB)", "Size") + if start is None or duration is None or nbytes is None: + continue + start_ns = _to_ns(start) + dur_ns = _to_ns(duration) + raw_bytes = _to_ns(nbytes) + byte_count = ( + int(raw_bytes * 1_000_000) + if "MB" in " ".join(row.keys()) + else int(raw_bytes) + ) + memcpys.append( + Memcpy(start=start_ns, end=start_ns + dur_ns, bytes=byte_count) + ) + return memcpys + + +def parse_nvtx_ranges(csv_text: str) -> List[NvtxRange]: + """Parse NVTX push/pop ranges from an nsys ``nvtx_pushpop_trace`` CSV.""" + ranges: List[NvtxRange] = [] + for row in _read_csv_rows(csv_text): + name = (_column(row, "Name", "Text") or "").strip() + if not name: + continue + start = _column(row, "Start (ns)", "Start") + duration = _column(row, "Duration (ns)", "Duration") + if start is None or duration is None: + continue + start_ns = _to_ns(start) + dur_ns = _to_ns(duration) + ranges.append( + NvtxRange(name=name, start=start_ns, end=start_ns + dur_ns) + ) + return ranges + + +def parse_nsys_rep( + rep_path: str, + compute_ranges: Sequence[str] = DEFAULT_COMPUTE_RANGES, +) -> Dict[str, Any]: + """Run ``nsys stats`` on ``rep_path`` and return the BM4 overlap report.""" + memcpys = parse_memcpys(_nsys_stats_csv(rep_path, "cuda_gpu_trace")) + ranges = parse_nvtx_ranges(_nsys_stats_csv(rep_path, "nvtx_pushpop_trace")) + result = compute_overlap(memcpys, ranges, compute_ranges) + return { + "benchmark": "BM4", + "rep": rep_path, + "total_h2d_bytes": result.total_h2d_bytes, + "overlapped_h2d_bytes": result.overlapped_h2d_bytes, + "overlap_fraction": result.overlap_fraction, + "exposed_fetch_bytes": result.total_h2d_bytes + - result.overlapped_h2d_bytes, + "num_memcpys": len(memcpys), + "compute_ranges": list(compute_ranges), + } + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m benchmarks.dflash.parse_overlap", + description="BM4 expert-H2D / compute overlap from an nsys trace.", + ) + parser.add_argument("--rep", required=True, help="path to a .nsys-rep") + parser.add_argument("--output", required=True) + parser.add_argument( + "--compute-range", + nargs="+", + default=list(DEFAULT_COMPUTE_RANGES), + help="NVTX range names that count as fetch-hiding compute", + ) + args = parser.parse_args(argv) + + report = parse_nsys_rep(args.rep, args.compute_range) + with open(args.output, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, sort_keys=True) + handle.write("\n") + json.dump(report, sys.stdout, indent=2, sort_keys=True) + sys.stdout.write("\n") + return 0 + + +__all__ = [ + "DEFAULT_COMPUTE_RANGES", + "Memcpy", + "NvtxRange", + "OverlapResult", + "compute_overlap", + "parse_memcpys", + "parse_nvtx_ranges", + "parse_nsys_rep", + "main", +] + + +if __name__ == "__main__": # pragma: no cover - CLI entry + raise SystemExit(main()) diff --git a/benchmarks/dflash/run_phase_c.sh b/benchmarks/dflash/run_phase_c.sh new file mode 100755 index 00000000..b8554820 --- /dev/null +++ b/benchmarks/dflash/run_phase_c.sh @@ -0,0 +1,144 @@ +#!/usr/bin/env bash +# =========================================================================== +# run_phase_c.sh -- one-command PD-DFlash Phase-C benchmark-gated C++ hops. +# +# Runs BM3 (route-ahead priority-band ablation), BM4 (expert-H2D / compute +# overlap trace via nsys), and BM5 (end-to-end Python-issue vs shipped +# C++-issue) on ONE RTX PRO 6000 (sm_120, capability 12.0) for both required +# MoE targets with FP4-offloaded experts, then aggregates the final keep/remove +# verdict per C++ hop. This is the hardware harness for plan Tasks 9-10 +# (docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md). +# +# The plan's hard rule: no C++ change ships without its paired BM passing on +# BOTH targets. BM3 gates the route-ahead priority band; BM4 supplies the +# overlap ground truth; BM5 proves C++ issue does not change correctness. +# +# USAGE +# benchmarks/dflash/run_phase_c.sh +# +# All inputs are environment variables (defaults assume this project's layout); +# override inline, e.g.: +# QWEN_OFFLOAD=/data/qwen-fp4 benchmarks/dflash/run_phase_c.sh +# +# REQUIRED on the GPU box: +# HF_HOME cached checkpoints (default /mnt/raid0nvme0/public/huggingface) +# CUDA_VISIBLE_DEVICES the single RTX PRO 6000 to use (default 0) +# QWEN_OFFLOAD dir of FP4-offloaded Qwen3-Coder-30B-A3B experts +# GPTOSS_OFFLOAD dir of FP4-offloaded gpt-oss-20b experts (needs #137) +# +# KEY KNOBS +# DEVICE_MEMORY_RATIO weight-resident fraction; MUST be < 0.9 to offload +# REPETITIONS BM3 ablation reps per arm (default 10) +# BLOCK_SIZE draft block size for BM3/BM4/BM5 (default 16) +# REQUESTS requests per measurement (default 16; short by design) +# PD_DFLASH_BUILD=1 rebuild the native sm_120 extensions first +# +# OUTPUTS (under $OUTPUT_DIR, default /tmp/pd-dflash-results) +# bm3.json three-way priority ablation + ship verdict +# nsys/qwen-final.nsys-rep, nsys/gptoss-final.nsys-rep +# bm4-qwen.json, bm4-gptoss.json overlap fraction ground truth +# bm5-python.json, bm5-cpp.json end-to-end Python vs C++ issue +# final.csv, final.md aggregated §8 + BM1-BM5 + hop verdicts +# =========================================================================== +set -euo pipefail + +export HF_HOME="${HF_HOME:-/mnt/raid0nvme0/public/huggingface}" +export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" +export MOE_ENABLE_SM120="${MOE_ENABLE_SM120:-1}" +export MOE_DFLASH_SERVING_GPU="${MOE_DFLASH_SERVING_GPU:-1}" + +OUTPUT_DIR="${OUTPUT_DIR:-/tmp/pd-dflash-results}" +REPETITIONS="${REPETITIONS:-10}" +BLOCK_SIZE="${BLOCK_SIZE:-16}" +REQUESTS="${REQUESTS:-16}" +WARMUP="${WARMUP:-5}" +SEED="${SEED:-1408}" +DEVICE_MEMORY_RATIO="${DEVICE_MEMORY_RATIO:-0.85}" + +QWEN_MODEL="${QWEN_MODEL:-Qwen/Qwen3-Coder-30B-A3B}" +QWEN_DRAFT="${QWEN_DRAFT:-z-lab/Qwen3-Coder-30B-A3B-DFlash}" +QWEN_OFFLOAD="${QWEN_OFFLOAD:-/mnt/raid0nvme0/offload/qwen3-coder-30b-a3b-fp4}" + +GPTOSS_MODEL="${GPTOSS_MODEL:-openai/gpt-oss-20b}" +GPTOSS_DRAFT="${GPTOSS_DRAFT:-z-lab/gpt-oss-20b-DFlash}" +GPTOSS_OFFLOAD="${GPTOSS_OFFLOAD:-/mnt/raid0nvme0/offload/gpt-oss-20b-fp4}" + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "$REPO_ROOT" +mkdir -p "$OUTPUT_DIR/raw" "$OUTPUT_DIR/nsys" + +echo "[run_phase_c] repo=$REPO_ROOT out=$OUTPUT_DIR device=$CUDA_VISIBLE_DEVICES" + +if awk "BEGIN{exit !($DEVICE_MEMORY_RATIO >= 0.9)}"; then + echo "[run_phase_c] ERROR: DEVICE_MEMORY_RATIO=$DEVICE_MEMORY_RATIO >= 0.9 will not offload" >&2 + exit 2 +fi + +if [[ "${PD_DFLASH_BUILD:-0}" == "1" ]]; then + echo "[run_phase_c] building native sm_120 extensions" + MOE_ENABLE_SM120=1 MOE_ENABLE_SM90=0 CUTLASS_DIR="${CUTLASS_DIR:-$HOME/cutlass}" \ + pip install --no-build-isolation -e . +fi + +# --- BM3: three-way priority-band ablation on both targets ----------------- +echo "[run_phase_c] === BM3 priority ablation ===" +python -m benchmarks.dflash.bench_prefetch_priority \ + --models "$QWEN_MODEL" "$GPTOSS_MODEL" \ + --drafts "$QWEN_DRAFT" "$GPTOSS_DRAFT" \ + --offload-dirs "$QWEN_OFFLOAD" "$GPTOSS_OFFLOAD" \ + --repetitions "$REPETITIONS" --block-size "$BLOCK_SIZE" \ + --requests "$REQUESTS" --warmup-rounds "$WARMUP" --seed "$SEED" \ + --device-memory-ratio "$DEVICE_MEMORY_RATIO" \ + --output "$OUTPUT_DIR/bm3.json" + +# --- BM4: nsys overlap capture + parse for the OURS config ----------------- +capture_bm4() { + local model="$1" draft="$2" offload="$3" tag="$4" + echo "[run_phase_c] === BM4 overlap capture: $model ===" + nsys profile --trace=cuda,nvtx --sample=none --force-overwrite=true \ + --output="$OUTPUT_DIR/nsys/$tag-final" \ + python -m benchmarks.dflash.pd_dflash_serving \ + --model "$model" --draft "$draft" --offload-dir "$offload" \ + --baseline OURS --block-size "$BLOCK_SIZE" --concurrency 8 \ + --requests 32 --warmup-rounds "$WARMUP" --seed "$SEED" \ + --device-memory-ratio "$DEVICE_MEMORY_RATIO" \ + --output "$OUTPUT_DIR/raw/$tag-final.json" + python -m benchmarks.dflash.parse_overlap \ + --rep "$OUTPUT_DIR/nsys/$tag-final.nsys-rep" \ + --output "$OUTPUT_DIR/bm4-$tag.json" +} +capture_bm4 "$QWEN_MODEL" "$QWEN_DRAFT" "$QWEN_OFFLOAD" "qwen" +capture_bm4 "$GPTOSS_MODEL" "$GPTOSS_DRAFT" "$GPTOSS_OFFLOAD" "gptoss" + +# --- BM5: end-to-end Python-issue vs shipped C++-issue --------------------- +run_bm5() { + local issue_mode="$1" output="$2" + for pair in "$QWEN_MODEL|$QWEN_DRAFT|$QWEN_OFFLOAD" \ + "$GPTOSS_MODEL|$GPTOSS_DRAFT|$GPTOSS_OFFLOAD"; do + IFS='|' read -r model draft offload <<<"$pair" + python -m benchmarks.dflash.pd_dflash_serving \ + --model "$model" --draft "$draft" --offload-dir "$offload" \ + --baseline OURS --block-size "$BLOCK_SIZE" --concurrency 1 8 32 \ + --requests "$REQUESTS" --warmup-rounds "$WARMUP" --seed "$SEED" \ + --device-memory-ratio "$DEVICE_MEMORY_RATIO" \ + --output "$output" || true + done +} +echo "[run_phase_c] === BM5 python-issue ===" +MOE_PREFETCH_ISSUE_MODE=python-per-expert run_bm5 python-per-expert \ + "$OUTPUT_DIR/bm5-python.json" +echo "[run_phase_c] === BM5 cpp-issue ===" +MOE_PREFETCH_ISSUE_MODE=cpp-batched run_bm5 cpp-batched \ + "$OUTPUT_DIR/bm5-cpp.json" + +# --- final aggregation ----------------------------------------------------- +echo "[run_phase_c] aggregating final gate report" +python -m benchmarks.dflash.report \ + --input "$OUTPUT_DIR/raw/qwen.json" "$OUTPUT_DIR/raw/gpt-oss.json" \ + --csv "$OUTPUT_DIR/final.csv" --markdown "$OUTPUT_DIR/final.md" || true + +echo "[run_phase_c] done:" +echo " BM3: $OUTPUT_DIR/bm3.json" +echo " BM4: $OUTPUT_DIR/bm4-{qwen,gptoss}.json" +echo " BM5: $OUTPUT_DIR/bm5-{python,cpp}.json" +echo " final: $OUTPUT_DIR/final.{csv,md}" diff --git a/tests/python/dflash/test_prefetch_perf_reports.py b/tests/python/dflash/test_prefetch_perf_reports.py index 46a43631..74c4adbd 100644 --- a/tests/python/dflash/test_prefetch_perf_reports.py +++ b/tests/python/dflash/test_prefetch_perf_reports.py @@ -37,6 +37,18 @@ build_bm2_report, percentiles_us, ) +from benchmarks.dflash.bench_prefetch_priority import ( + BACKGROUND, + ON_DEMAND, + PRIORITY_BANDS, + ROUTE_AHEAD, + Bm3Decision, + PriorityArm, + bm3_decision, + build_bm3_report, + median, + priority_arm, +) def test_bm2_candidate_required_when_per_expert_exceeds_window(): @@ -159,3 +171,257 @@ def test_bm2_report_ships_when_batched_mode_present_and_fast(): assert report["medians_us"][PYTHON_PER_EXPERT] == pytest.approx(900.0) assert report["candidate_required"] is True assert report["ship_batched"] is True + + +# =========================================================================== +# BM3 -- route-ahead priority-band ablation decision rule (design §10, Task 9). +# +# A three-way ablation over median *exposed-fetch seconds* and *tokens/s*: +# default background (prio 2), dedicated route-ahead band (prio 1), and +# on-demand (prio 0). The dedicated route-ahead band ships iff all three +# hold: +# (1) route-ahead has *lower* exposed fetch than default background; +# (2) route-ahead does not *reduce* tokens/s vs default background; and +# (3) on-demand remains the fastest service class (no priority inversion). +# All arms are pure medians, so the ship gate is verifiable off-hardware. +# =========================================================================== + + +def _arm(exposed_fetch_seconds: float, tokens_per_second: float) -> PriorityArm: + return priority_arm( + exposed_fetch_seconds=exposed_fetch_seconds, + tokens_per_second=tokens_per_second, + ) + + +def test_bm3_ships_when_route_ahead_helps_without_inversion_or_regression(): + decision = bm3_decision( + default=_arm(0.020, 100.0), + route_ahead=_arm(0.012, 105.0), + on_demand=_arm(0.008, 106.0), + ) + assert isinstance(decision, Bm3Decision) + assert decision.exposed_fetch_improved is True + assert decision.throughput_preserved is True + assert decision.on_demand_fastest is True + assert decision.ship_priority_band is True + + +def test_bm3_no_improvement_does_not_ship(): + # Route-ahead exposed fetch is not lower than default background. + decision = bm3_decision( + default=_arm(0.012, 100.0), + route_ahead=_arm(0.012, 100.0), + on_demand=_arm(0.008, 101.0), + ) + assert decision.exposed_fetch_improved is False + assert decision.ship_priority_band is False + + +def test_bm3_throughput_regression_does_not_ship(): + # Route-ahead lowers exposed fetch but *reduces* tokens/s vs background. + decision = bm3_decision( + default=_arm(0.020, 100.0), + route_ahead=_arm(0.012, 97.0), + on_demand=_arm(0.008, 101.0), + ) + assert decision.exposed_fetch_improved is True + assert decision.throughput_preserved is False + assert decision.ship_priority_band is False + + +def test_bm3_priority_inversion_does_not_ship(): + # On-demand is no longer the fastest class (route-ahead starves it). + decision = bm3_decision( + default=_arm(0.020, 100.0), + route_ahead=_arm(0.010, 105.0), + on_demand=_arm(0.014, 103.0), + ) + assert decision.exposed_fetch_improved is True + assert decision.throughput_preserved is True + assert decision.on_demand_fastest is False + assert decision.ship_priority_band is False + + +def test_bm3_throughput_preserved_is_inclusive_at_equality(): + decision = bm3_decision( + default=_arm(0.020, 100.0), + route_ahead=_arm(0.012, 100.0), + on_demand=_arm(0.008, 100.0), + ) + assert decision.throughput_preserved is True + assert decision.on_demand_fastest is True + assert decision.ship_priority_band is True + + +def test_bm3_decision_is_frozen(): + decision = bm3_decision( + default=_arm(0.020, 100.0), + route_ahead=_arm(0.012, 105.0), + on_demand=_arm(0.008, 106.0), + ) + with pytest.raises(Exception): + decision.ship_priority_band = False # type: ignore[misc] + + +@pytest.mark.parametrize("bad", [-1.0, float("nan"), float("inf")]) +def test_bm3_arm_rejects_nonfinite_or_negative_exposed_fetch(bad): + with pytest.raises(ValueError): + priority_arm(exposed_fetch_seconds=bad, tokens_per_second=100.0) + + +@pytest.mark.parametrize("bad", [-1.0, float("nan"), float("inf")]) +def test_bm3_arm_rejects_nonfinite_or_negative_tokens_per_second(bad): + with pytest.raises(ValueError): + priority_arm(exposed_fetch_seconds=0.01, tokens_per_second=bad) + + +def test_bm3_median_nearest_rank_odd_and_even(): + assert median([0.01, 0.03, 0.02]) == pytest.approx(0.02) + # even count -> mean of the two central order statistics + assert median([0.01, 0.02, 0.03, 0.06]) == pytest.approx(0.025) + + +def test_bm3_median_requires_samples(): + with pytest.raises(ValueError): + median([]) + + +def test_bm3_priority_bands_are_named_and_ordered_high_to_low_service(): + # On-demand (0) is serviced first, then route-ahead (1), then background (2). + assert ON_DEMAND == 0 + assert ROUTE_AHEAD == 1 + assert BACKGROUND == 2 + assert PRIORITY_BANDS == ("background", "route-ahead", "on-demand") + + +def test_bm3_report_assembles_medians_and_ship_verdict(): + report = build_bm3_report( + models=["tiny/fixture"], + arms={ + "background": { + "exposed_fetch_seconds": [0.021, 0.020, 0.019], + "tokens_per_second": [99.0, 100.0, 101.0], + }, + "route-ahead": { + "exposed_fetch_seconds": [0.013, 0.012, 0.011], + "tokens_per_second": [104.0, 105.0, 106.0], + }, + "on-demand": { + "exposed_fetch_seconds": [0.009, 0.008, 0.007], + "tokens_per_second": [105.0, 106.0, 107.0], + }, + }, + ) + assert report["benchmark"] == "BM3" + assert report["medians"]["route-ahead"]["exposed_fetch_seconds"] == ( + pytest.approx(0.012) + ) + assert report["exposed_fetch_improved"] is True + assert report["throughput_preserved"] is True + assert report["on_demand_fastest"] is True + assert report["ship_priority_band"] is True + + +def test_bm3_report_requires_all_three_arms(): + with pytest.raises(ValueError): + build_bm3_report( + models=["tiny/fixture"], + arms={ + "background": { + "exposed_fetch_seconds": [0.02], + "tokens_per_second": [100.0], + }, + "route-ahead": { + "exposed_fetch_seconds": [0.01], + "tokens_per_second": [100.0], + }, + }, + ) + + +# =========================================================================== +# BM4 -- H2D/compute overlap from nsys/CUPTI intervals (design §10, Task 10). +# +# Ground truth for whether expert fetch is hidden: the fraction of expert-H2D +# *bytes* overlapped with the draft/router/verify compute NVTX ranges. Bytes on +# a partially overlapped memcpy are apportioned by overlapped-duration fraction +# (never a whole-copy all-or-nothing). Pure interval arithmetic -> off-hardware. +# =========================================================================== + +from benchmarks.dflash.parse_overlap import ( # noqa: E402 + Memcpy, + NvtxRange, + compute_overlap, +) + +COMPUTE_RANGES = ("dflash_draft", "route_ahead_router", "target_verify") + + +def test_bm4_partial_overlap_apportions_bytes_by_duration(): + # One 100-byte memcpy over [0,10]; a compute range covers [0,8] -> 80% of + # the duration overlaps -> 80 of 100 bytes are hidden. + memcpys = [Memcpy(start=0.0, end=10.0, bytes=100)] + ranges = [NvtxRange(name="dflash_draft", start=0.0, end=8.0)] + result = compute_overlap(memcpys, ranges, COMPUTE_RANGES) + assert result.total_h2d_bytes == 100 + assert result.overlapped_h2d_bytes == pytest.approx(80.0) + assert result.overlap_fraction == pytest.approx(0.8) + + +def test_bm4_disjoint_ranges_hide_nothing(): + memcpys = [Memcpy(start=0.0, end=10.0, bytes=100)] + ranges = [NvtxRange(name="target_verify", start=20.0, end=30.0)] + result = compute_overlap(memcpys, ranges, COMPUTE_RANGES) + assert result.total_h2d_bytes == 100 + assert result.overlapped_h2d_bytes == pytest.approx(0.0) + assert result.overlap_fraction == pytest.approx(0.0) + + +def test_bm4_fully_overlapped_hides_everything(): + memcpys = [Memcpy(start=2.0, end=6.0, bytes=64)] + ranges = [NvtxRange(name="route_ahead_router", start=0.0, end=10.0)] + result = compute_overlap(memcpys, ranges, COMPUTE_RANGES) + assert result.overlapped_h2d_bytes == pytest.approx(64.0) + assert result.overlap_fraction == pytest.approx(1.0) + + +def test_bm4_zero_bytes_is_defined_and_not_a_divide_by_zero(): + result = compute_overlap([], [], COMPUTE_RANGES) + assert result.total_h2d_bytes == 0 + assert result.overlapped_h2d_bytes == pytest.approx(0.0) + assert result.overlap_fraction == pytest.approx(0.0) + + +def test_bm4_multiple_compute_ranges_union_not_double_counted(): + # Two overlapping compute ranges cover [0,4] and [2,10] -> union [0,10] + # fully covers a 50-byte copy over [0,10]; overlap must not exceed 100%. + memcpys = [Memcpy(start=0.0, end=10.0, bytes=50)] + ranges = [ + NvtxRange(name="dflash_draft", start=0.0, end=4.0), + NvtxRange(name="target_verify", start=2.0, end=10.0), + ] + result = compute_overlap(memcpys, ranges, COMPUTE_RANGES) + assert result.overlapped_h2d_bytes == pytest.approx(50.0) + assert result.overlap_fraction == pytest.approx(1.0) + + +def test_bm4_only_named_compute_ranges_count(): + # A range whose name is not a compute range (e.g. expert_h2d itself) does + # not count as compute that hides the fetch. + memcpys = [Memcpy(start=0.0, end=10.0, bytes=100)] + ranges = [NvtxRange(name="expert_h2d", start=0.0, end=10.0)] + result = compute_overlap(memcpys, ranges, COMPUTE_RANGES) + assert result.overlapped_h2d_bytes == pytest.approx(0.0) + + +def test_bm4_per_memcpy_fraction_reported(): + memcpys = [ + Memcpy(start=0.0, end=10.0, bytes=100), + Memcpy(start=0.0, end=10.0, bytes=100), + ] + ranges = [NvtxRange(name="dflash_draft", start=0.0, end=5.0)] + result = compute_overlap(memcpys, ranges, COMPUTE_RANGES) + assert result.total_h2d_bytes == 200 + assert result.overlapped_h2d_bytes == pytest.approx(100.0) + assert result.overlap_fraction == pytest.approx(0.5) From 5391673f1511bc7a758f87ca16728bf13b6a968d Mon Sep 17 00:00:00 2001 From: drunkcoding Date: Sat, 15 Aug 2026 12:31:34 +0000 Subject: [PATCH 9/9] bench(dflash): verify route-ahead overlap (BM5 equivalence + hop verdicts) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Task 10 final-gate aggregation (design §10): * summarise_bm5_equivalence -- switching Python per-expert issue to shipped C++ batched issue may only move tokens/s; acceptance, route-ahead coverage, and wasted bytes must match within tolerance. * cpp_hop_verdicts -- keep/remove each benchmark-gated C++ hop from its paired BM: batched issuance needs BM2 ship_batched, the priority band needs BM3 ship_priority_band; a missing BM removes the hop (no C++ ships without its BM). * test_pd_dflash_report.py -- CPU-only BM5 equivalence + hop-verdict tests. --- benchmarks/dflash/report.py | 82 +++++++++++++++++++- tests/python/dflash/test_pd_dflash_report.py | 65 ++++++++++++++++ 2 files changed, 146 insertions(+), 1 deletion(-) diff --git a/benchmarks/dflash/report.py b/benchmarks/dflash/report.py index c465dbf1..367a1fcf 100644 --- a/benchmarks/dflash/report.py +++ b/benchmarks/dflash/report.py @@ -19,7 +19,7 @@ import math import sys from dataclasses import dataclass -from typing import Any, Dict, List, Mapping, Sequence, Tuple +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple REQUIRED_METRICS = ( "output_tokens_per_second", @@ -181,6 +181,84 @@ def summarise_row(row: Mapping[str, object]) -> Mapping[str, object]: } +BM5_INVARIANTS = ( + "acceptance_length_a", + "route_ahead_prefetch_coverage", + "wasted_prefetch_bytes", +) + + +def summarise_bm5_equivalence( + *, + python_row: Mapping[str, Any], + cpp_row: Mapping[str, Any], + rel_tol: float = 1e-6, +) -> Dict[str, Any]: + """BM5: shipped C++ issue must not change correctness-visible terms. + + Switching Python per-expert issuance to the shipped C++ batched issuance may + only move tokens/s: acceptance length, route-ahead coverage, wasted bytes, + and output tokens must match within ``rel_tol``. Returns the equivalence + verdict, the mismatched invariants, and both tokens/s for the plot. + """ + mismatched: List[str] = [] + for key in BM5_INVARIANTS: + if key not in python_row or key not in cpp_row: + mismatched.append(key) + continue + a = float(python_row[key]) + b = float(cpp_row[key]) + scale = max(abs(a), abs(b), 1.0) + if abs(a - b) > rel_tol * scale: + mismatched.append(key) + return { + "benchmark": "BM5", + "equivalent": not mismatched, + "mismatched": mismatched, + "python_tokens_per_second": float( + python_row.get("output_tokens_per_second", float("nan")) + ), + "cpp_tokens_per_second": float( + cpp_row.get("output_tokens_per_second", float("nan")) + ), + } + + +def cpp_hop_verdicts( + *, + bm2: Optional[Mapping[str, Any]], + bm3: Optional[Mapping[str, Any]], +) -> Dict[str, Dict[str, Any]]: + """Keep/remove each benchmark-gated C++ hop from its paired BM verdict. + + A hop is kept only when its paired benchmark is present and passes: the + batched-issuance hop needs BM2 ``ship_batched``; the route-ahead + priority-band hop needs BM3 ``ship_priority_band``. A missing BM removes the + hop (design §10: no C++ change ships without its BM). + """ + + def verdict(report: Optional[Mapping[str, Any]], flag: str, hop: str): + if report is None: + return { + "keep": False, + "reason": f"{hop} removed: paired benchmark absent", + } + keep = bool(report.get(flag, False)) + return { + "keep": keep, + "reason": ( + f"{hop} kept: {flag} is true" + if keep + else f"{hop} removed: {flag} is false" + ), + } + + return { + "batched_issuance": verdict(bm2, "ship_batched", "batched_issuance"), + "priority_band": verdict(bm3, "ship_priority_band", "priority_band"), + } + + def aggregate_result_matrices( rows: Sequence[Mapping[str, Any]], ) -> Dict[Tuple[str, int, int], Dict[str, Mapping[str, Any]]]: @@ -353,9 +431,11 @@ def main(argv: Sequence[str] | None = None) -> int: "UNAVAILABLE_CAPACITY", "HideInequality", "aggregate_result_matrices", + "cpp_hop_verdicts", "evaluate_hide_inequality", "evaluate_matrix", "main", + "summarise_bm5_equivalence", "summarise_row", "validate_result_matrix", ] diff --git a/tests/python/dflash/test_pd_dflash_report.py b/tests/python/dflash/test_pd_dflash_report.py index 8743b412..e2ca7b04 100644 --- a/tests/python/dflash/test_pd_dflash_report.py +++ b/tests/python/dflash/test_pd_dflash_report.py @@ -268,3 +268,68 @@ def test_required_metrics_are_frozen_and_complete(): ) assert len(set(REQUIRED_METRICS)) == len(REQUIRED_METRICS) assert math.isfinite(1.0) + + +# =========================================================================== +# BM5 equivalence + final C++-hop keep/remove verdicts (design §10, Task 10). +# +# BM5: switching from Python per-expert issue to shipped C++ batched issue must +# not alter acceptance, route-ahead coverage, wasted bytes, or output tokens -- +# only tokens/s may move. The final gate reports keep/remove per C++ hop from +# each hop's paired BM ship flag; a hop lacking a passing BM is removed. +# =========================================================================== + +from benchmarks.dflash.report import ( # noqa: E402 + cpp_hop_verdicts, + summarise_bm5_equivalence, +) + + +def _bm5_row(tokens_per_second): + return { + "acceptance_length_a": 3.0, + "route_ahead_prefetch_coverage": 0.75, + "wasted_prefetch_bytes": 12_582_912.0, + "output_tokens_per_second": tokens_per_second, + } + + +def test_bm5_equivalence_holds_when_only_throughput_moves(): + summary = summarise_bm5_equivalence( + python_row=_bm5_row(100.0), cpp_row=_bm5_row(140.0) + ) + assert summary["equivalent"] is True + assert summary["python_tokens_per_second"] == 100.0 + assert summary["cpp_tokens_per_second"] == 140.0 + + +def test_bm5_equivalence_fails_when_coverage_changes(): + cpp = _bm5_row(140.0) + cpp["route_ahead_prefetch_coverage"] = 0.50 + summary = summarise_bm5_equivalence(python_row=_bm5_row(100.0), cpp_row=cpp) + assert summary["equivalent"] is False + assert "route_ahead_prefetch_coverage" in summary["mismatched"] + + +def test_bm5_equivalence_fails_when_waste_bytes_change(): + cpp = _bm5_row(140.0) + cpp["wasted_prefetch_bytes"] = 0.0 + summary = summarise_bm5_equivalence(python_row=_bm5_row(100.0), cpp_row=cpp) + assert summary["equivalent"] is False + assert "wasted_prefetch_bytes" in summary["mismatched"] + + +def test_cpp_hop_verdicts_keep_only_hops_with_passing_bm(): + verdicts = cpp_hop_verdicts( + bm2={"ship_batched": True}, + bm3={"ship_priority_band": False}, + ) + assert verdicts["batched_issuance"]["keep"] is True + assert verdicts["priority_band"]["keep"] is False + + +def test_cpp_hop_verdicts_missing_bm_removes_hop(): + verdicts = cpp_hop_verdicts(bm2=None, bm3=None) + assert verdicts["batched_issuance"]["keep"] is False + assert verdicts["priority_band"]["keep"] is False + assert verdicts["batched_issuance"]["reason"]