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/_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/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/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/benchmarks/dflash/report.py b/benchmarks/dflash/report.py new file mode 100644 index 00000000..c465dbf1 --- /dev/null +++ b/benchmarks/dflash/report.py @@ -0,0 +1,365 @@ +"""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 argparse +import json +import math +import sys +from dataclasses import dataclass +from typing import Any, Dict, List, Mapping, Sequence, Tuple + +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 + _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__ = [ + "REQUIRED_METRICS", + "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/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}" 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/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..637517b7 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): @@ -53,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/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_pd_dflash_report.py b/tests/python/dflash/test_pd_dflash_report.py new file mode 100644 index 00000000..8743b412 --- /dev/null +++ b/tests/python/dflash/test_pd_dflash_report.py @@ -0,0 +1,270 @@ +"""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, + aggregate_result_matrices, + evaluate_hide_inequality, + evaluate_matrix, + summarise_row, + 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 + + +# --------------------------------------------------------------------------- +# 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 == ( + "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) 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"} 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_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 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 # --------------------------------------------------------------------------- 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()