diff --git a/benchmarks/dflash/_serving_measure.py b/benchmarks/dflash/_serving_measure.py new file mode 100644 index 00000000..1d41f974 --- /dev/null +++ b/benchmarks/dflash/_serving_measure.py @@ -0,0 +1,374 @@ +"""GPU measurement for the PD-DFlash B0-B3 serving experiment (Task 2). + +Private helper for ``benchmarks.dflash.pd_dflash_serving``; imported only inside +``run_experiment`` so the CLI module stays torch-free at import. Every function +here needs a live RTX PRO 6000 with FP4-offloaded experts and cached +checkpoints, so nothing in this file is exercised by CPU pytest -- it is the +hardware harness a human runs for plan Task 3. + +The measurement mirrors ``tests/python/dflash/test_gpu_serving_dflash.py``: build +``MoE`` with an offload path, wrap a ``DFlashSpeculator`` for the DFlash +baselines, drive deterministic greedy requests through the continuous-batching +engine, and read metrics from measured wall clock, the speculator ``step_trace``, +and the instrumented ``RouteAheadStats``. Where a native occupancy/hit-rate +accessor is not present the extractor falls back to ``0.0`` and records a +per-row ``warnings`` entry, so a row is always schema-valid *and* honest about +which term needs a human to wire a native accessor. +""" + +from __future__ import annotations + +import time +from contextlib import contextmanager +from typing import Any, Dict, Iterator, List, Optional, Tuple + +from benchmarks.dflash.pd_dflash_serving import ( + NVTX_RANGES, + RunnerArgs, + make_observation_row, + require_offloaded, +) + +try: + import nvtx as _nvtx +except Exception: # pragma: no cover - nvtx optional + _nvtx = None + +RESIDENT_MEMORY_RATIO = 0.98 +DETERMINISTIC_PROMPT = ( + "Explain in one paragraph why offloaded mixture-of-experts serving " + "benefits from speculative decoding." +) + + +@contextmanager +def nvtx_range(name: str) -> Iterator[None]: + """Push an NVTX range so the BM4 overlap parser can attribute H2D bytes.""" + if _nvtx is None or name not in NVTX_RANGES: + yield + return + handle = _nvtx.start_range(message=name, color="green") + try: + yield + finally: + _nvtx.end_range(handle) + + +def measure_configuration( + *, + args: RunnerArgs, + baseline: str, + draft: str, + block_size: int, + concurrency: int, +) -> Dict[str, Any]: + """Measure one ``(baseline, block, concurrency)`` cell and return its row. + + B0 runs the AR offloaded target with no speculator; B1/B3 and the ``OURS`` + configuration wrap a DFlash draft, with route-ahead stats enabled so + coverage and byte-accurate waste are recorded. B3 loads the target resident + (no offload upper bound); the other baselines require genuinely offloaded + experts and are refused otherwise. + """ + import torch + + from moe_infinity import MoE + from moe_infinity.spec_decode import DFlashSpeculator + + warnings: List[str] = [] + resident = baseline == "B3" + memory_ratio = ( + RESIDENT_MEMORY_RATIO if resident else args.device_memory_ratio + ) + model = MoE( + args.model, + { + "offload_path": args.offload_dir, + "device_memory_ratio": memory_ratio, + }, + ) + engine = model.engine + if not resident: + require_offloaded(baseline, _count_offloaded_experts(engine)) + + speculator = None + if baseline != "B0": + speculator = DFlashSpeculator(model, draft) + enable = getattr(speculator, "enable_route_ahead_stats", None) + if callable(enable): + enable() + + torch.manual_seed(args.seed) + prompt_ids = _deterministic_prompt_ids(model, args.model) + + _warmup(model, prompt_ids, speculator, args.warmup_rounds, block_size) + + torch.cuda.synchronize() + started = time.perf_counter() + generated = _run_requests( + model=model, + prompt_ids=prompt_ids, + speculator=speculator, + block_size=block_size, + concurrency=concurrency, + num_requests=args.requests, + ) + torch.cuda.synchronize() + elapsed = max(time.perf_counter() - started, 1e-9) + + ttft = _measure_ttft(model, prompt_ids, speculator, block_size) + + metrics = _collect_metrics( + baseline=baseline, + block_size=block_size, + elapsed=elapsed, + ttft_seconds=ttft, + generated_tokens=generated, + num_requests=args.requests, + speculator=speculator, + engine=engine, + slo_ms=args.slo_ms, + warnings=warnings, + ) + cost_terms = _collect_cost_terms( + baseline=baseline, + speculator=speculator, + engine=engine, + measured_h2d_gbps=args.measured_h2d_gbps, + warnings=warnings, + ) + return make_observation_row( + model=args.model, + draft=draft if baseline != "B0" else "", + baseline=baseline, + block_size=block_size, + concurrency=concurrency, + repeat=0, + metrics=metrics, + cost_terms=cost_terms, + warnings=warnings or None, + ) + + +def _count_offloaded_experts(engine: Any) -> int: + for attr in ("num_offloaded_experts", "offloaded_expert_count"): + value = getattr(engine, attr, None) + if isinstance(value, int): + return value + prefetcher = getattr(engine, "expert_prefetcher", None) + nbytes_map = getattr(prefetcher, "expert_nbytes_map", None) + if isinstance(nbytes_map, dict): + return len(nbytes_map) + return 0 + + +def _deterministic_prompt_ids(model: Any, repo: str) -> List[int]: + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + repo, trust_remote_code=True, local_files_only=True + ) + return [ + int(tok) + for tok in tokenizer(DETERMINISTIC_PROMPT, return_tensors="pt") + .input_ids[0] + .tolist() + ] + + +def _greedy_generate( + model: Any, + prompt_ids: List[int], + speculator: Optional[Any], + max_new_tokens: int, +) -> List[int]: + import torch + + input_ids = torch.tensor([prompt_ids], dtype=torch.long, device="cuda:0") + kwargs: Dict[str, Any] = { + "do_sample": False, + "max_new_tokens": max_new_tokens, + } + if speculator is not None: + kwargs["speculative_draft"] = speculator + with nvtx_range("target_verify"): + output = model.generate(input_ids, **kwargs) + return [int(tok) for tok in output[0, len(prompt_ids) :].tolist()] + + +def _warmup( + model: Any, + prompt_ids: List[int], + speculator: Optional[Any], + warmup_rounds: int, + block_size: int, +) -> None: + for _ in range(max(0, warmup_rounds)): + _greedy_generate(model, prompt_ids, speculator, max(1, block_size)) + + +def _run_requests( + *, + model: Any, + prompt_ids: List[int], + speculator: Optional[Any], + block_size: int, + concurrency: int, + num_requests: int, +) -> int: + tokens_per_request = max(block_size * 4, 32) + total = 0 + for _ in range(max(1, num_requests)): + generated = _greedy_generate( + model, prompt_ids, speculator, tokens_per_request + ) + total += len(generated) + return total + + +def _measure_ttft( + model: Any, + prompt_ids: List[int], + speculator: Optional[Any], + block_size: int, +) -> float: + import torch + + torch.cuda.synchronize() + started = time.perf_counter() + with nvtx_range("dflash_draft"): + _greedy_generate(model, prompt_ids, speculator, 1) + torch.cuda.synchronize() + return max(time.perf_counter() - started, 0.0) + + +def _acceptance_length( + baseline: str, block_size: int, speculator: Any +) -> float: + if baseline == "B0" or speculator is None: + return 1.0 + trace = list(getattr(speculator, "step_trace", []) or []) + if not trace: + return 1.0 + accepted = [ + min(int(getattr(r, "accept", 0)) + 1, block_size) for r in trace + ] + return sum(accepted) / len(accepted) + + +def _route_ahead_snapshot(speculator: Any) -> Tuple[float, Optional[int]]: + if speculator is None: + return 0.0, 0 + stats = getattr(speculator, "route_ahead_stats", None) + if stats is None: + return 0.0, 0 + snapshot = stats.as_dict() + coverage = float(snapshot.get("coverage", 0.0) or 0.0) + wasted = snapshot.get("wasted_prefetch_bytes") + return coverage, (int(wasted) if wasted is not None else None) + + +def _extract_float(source: Any, names: Tuple[str, ...]) -> Optional[float]: + for name in names: + value = getattr(source, name, None) + if callable(value): + try: + value = value() + except Exception: + value = None + if isinstance(value, (int, float)) and not isinstance(value, bool): + return float(value) + return None + + +def _collect_metrics( + *, + baseline: str, + block_size: int, + elapsed: float, + ttft_seconds: float, + generated_tokens: int, + num_requests: int, + speculator: Any, + engine: Any, + slo_ms: Optional[float], + warnings: List[str], +) -> Dict[str, float]: + tokens_per_second = generated_tokens / elapsed + acceptance = _acceptance_length(baseline, block_size, speculator) + rounds = max(1.0, generated_tokens / max(acceptance, 1.0)) + per_round_latency = elapsed / rounds + coverage, wasted_bytes = _route_ahead_snapshot(speculator) + if wasted_bytes is None: + warnings.append( + "wasted_prefetch_bytes unavailable from RouteAheadStats; a route-" + "ahead configuration on offloaded experts must report real bytes" + ) + wasted_bytes = 0 + + hit_rate = _extract_float( + getattr(engine, "expert_prefetcher", engine), + ("get_hit_rate", "hit_rate", "expert_hit_rate"), + ) + if hit_rate is None: + hit_rate = _extract_float(engine, ("get_hit_rate", "hit_rate")) + if hit_rate is None: + warnings.append("expert_cache_hit_rate fell back to 0.0") + hit_rate = 0.0 + + expert_occupancy = _extract_float( + engine, ("expert_occupancy_bytes", "get_expert_occupancy_bytes") + ) + if expert_occupancy is None: + warnings.append("expert_occupancy_bytes fell back to 0.0") + expert_occupancy = 0.0 + kv_occupancy = _extract_float( + engine, ("kv_occupancy_bytes", "get_kv_occupancy_bytes") + ) + if kv_occupancy is None: + warnings.append("kv_occupancy_bytes fell back to 0.0") + kv_occupancy = 0.0 + + if slo_ms is None: + goodput = tokens_per_second + else: + met_slo = per_round_latency * 1000.0 <= slo_ms + goodput = tokens_per_second if met_slo else 0.0 + + return { + "output_tokens_per_second": tokens_per_second, + "acceptance_length_a": acceptance, + "ttft_seconds": ttft_seconds, + "per_round_latency_seconds": per_round_latency, + "goodput_at_slo": goodput, + "expert_cache_hit_rate": hit_rate, + "route_ahead_prefetch_coverage": coverage, + "wasted_prefetch_bytes": float(wasted_bytes), + "expert_occupancy_bytes": expert_occupancy, + "kv_occupancy_bytes": kv_occupancy, + } + + +def _collect_cost_terms( + *, + baseline: str, + speculator: Any, + engine: Any, + measured_h2d_gbps: Optional[float], + warnings: List[str], +) -> Dict[str, Any]: + coverage, wasted_bytes = _route_ahead_snapshot(speculator) + terms: Dict[str, Any] = { + "route_ahead_coverage": coverage, + "wasted_prefetch_bytes": wasted_bytes, + } + if measured_h2d_gbps is not None: + terms["measured_h2d_bytes_per_second"] = ( + measured_h2d_gbps * 1_000_000_000.0 + ) + else: + warnings.append( + "measured_h2d_bytes_per_second not supplied; pass --measured-h2d-" + "gbps from a device bandwidth probe for the hide inequality" + ) + return terms diff --git a/benchmarks/dflash/pd_dflash_serving.py b/benchmarks/dflash/pd_dflash_serving.py new file mode 100644 index 00000000..f9ea3721 --- /dev/null +++ b/benchmarks/dflash/pd_dflash_serving.py @@ -0,0 +1,425 @@ +"""Opt-in RTX PRO 6000 B0-B3 serving experiment for PD-DFlash route-ahead. + +Task 2 of ``docs/superpowers/plans/2026-08-14-pd-dflash-serving-scheduler.md`` +("measure-first gating experiment"). This is the runner a human executes on the +GPU box; it records every design-doc §8 metric for the B0-B3 baselines so the +cost-model hide inequality (``report.py``) can be evaluated before any scheduler +or C++ work is justified. + +The module is import-safe by construction: only the pure contract/scheduling +scaffolding lives at module scope, so ``pytest`` collection (and +``--dry-run-contract``) never imports torch, loads a checkpoint, initialises +CUDA, or touches the network. All hardware work is lazily imported inside +``run_experiment``. + +Design contract (frozen here, cross-checked by the aggregator): + +* baselines are exactly B0-B3 with the design-doc §8 semantics; +* the required generalization targets are ``Qwen/Qwen3-Coder-30B-A3B`` and + ``openai/gpt-oss-20b`` with their ``z-lab`` DFlash drafts; +* block sizes are 8 and 16, concurrency sweeps 1..32; and +* every emitted observation carries the full ``REQUIRED_METRICS`` schema, plus + the byte-accurate route-ahead ``wasted_prefetch_bytes`` from the instrumented + ``RouteAheadStats`` (never an expert-count proxy). + +At this Phase-A stage B2 (the 2-D deficit scheduler) does not yet exist, so the +runner emits an explicit ``BLOCKED_UNTIL_2D_SCHEDULER`` status for B2 rather than +silently emulating another baseline. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from dataclasses import dataclass +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple + +from benchmarks.dflash.report import REQUIRED_METRICS + +BASELINES = { + "B0": "AR MoE on MoE-Infinity, offloaded, no speculative decoding", + "B1": "DFlash with unchanged AR prefetcher", + "B2": "DFlash with token-deficit scheduler and no expert-byte coupling", + "B3": "target experts resident, no offload upper bound", +} + +EXPERIMENTAL_CONFIGURATIONS = { + "OURS": "DFlash with route-ahead prefetch and the 2-D co-designed scheduler", +} + +REQUIRED_MODELS: Tuple[str, ...] = ( + "Qwen/Qwen3-Coder-30B-A3B", + "openai/gpt-oss-20b", +) + +REQUIRED_DRAFTS: Dict[str, str] = { + "Qwen/Qwen3-Coder-30B-A3B": "z-lab/Qwen3-Coder-30B-A3B-DFlash", + "openai/gpt-oss-20b": "z-lab/gpt-oss-20b-DFlash", +} + +REQUIRED_BLOCK_SIZES: Tuple[int, ...] = (8, 16) +REQUIRED_CONCURRENCY: Tuple[int, ...] = (1, 2, 4, 8, 16, 32) +DEFAULT_BASELINES: Tuple[str, ...] = ("B0", "B1", "B2", "B3") + +OFFLOADED_BASELINES: Tuple[str, ...] = ("B0", "B1", "B2") +BLOCKED_STATUS = "BLOCKED_UNTIL_2D_SCHEDULER" + +RTX_PRO_6000_NAME_FRAGMENT = "RTX PRO 6000" +RTX_PRO_6000_CAPABILITY: Tuple[int, int] = (12, 0) + +# NVTX ranges the BM4 overlap parser (Task 10) associates H2D memcpys with; the +# runner must wrap the corresponding phases in these exact names. +NVTX_RANGES: Tuple[str, ...] = ( + "dflash_draft", + "route_ahead_router", + "route_ahead_issue", + "target_verify", + "expert_h2d", +) + + +# --------------------------------------------------------------------------- +# pure contract + validation helpers (CPU-only, unit-tested) +# --------------------------------------------------------------------------- + + +def build_contract_matrix() -> Dict[str, Any]: + """Return the canonical §8 experiment contract as plain data. + + ``--dry-run-contract`` prints this and the GPU-gated test asserts it against + the design doc, so the required models/drafts/baselines/sweeps are pinned in + one place independent of any single invocation's CLI arguments. + """ + return { + "models": list(REQUIRED_MODELS), + "drafts": dict(REQUIRED_DRAFTS), + "baselines": dict(BASELINES), + "experimental_configurations": dict(EXPERIMENTAL_CONFIGURATIONS), + "block_sizes": list(REQUIRED_BLOCK_SIZES), + "concurrency": list(REQUIRED_CONCURRENCY), + "required_metrics": list(REQUIRED_METRICS), + "nvtx_ranges": list(NVTX_RANGES), + } + + +def validate_device_identity( + device_name: str, capability: Tuple[int, int] +) -> None: + """Raise unless the visible GPU is an RTX PRO 6000 with capability (12, 0). + + Kept free of torch so it is unit-testable; ``run_experiment`` feeds it the + live ``torch.cuda`` values. + """ + if RTX_PRO_6000_NAME_FRAGMENT not in device_name: + raise RuntimeError( + f"expected an {RTX_PRO_6000_NAME_FRAGMENT} GPU; got {device_name!r}" + ) + if tuple(capability) != RTX_PRO_6000_CAPABILITY: + raise RuntimeError( + f"expected capability {RTX_PRO_6000_CAPABILITY}; got " + f"{tuple(capability)!r}" + ) + + +def require_offloaded(baseline: str, offloaded_expert_count: int) -> None: + """Raise if an offloaded baseline (B0/B1/B2) has no offloaded experts. + + Guards against mislabelling a resident run as offloaded evidence (plan + dependency note on #137); B3 is the resident upper bound and is exempt. + """ + if baseline in OFFLOADED_BASELINES and offloaded_expert_count <= 0: + raise RuntimeError( + f"baseline {baseline} requires offloaded target experts; the store " + "reports none resident on host -- lower --device-memory-ratio below " + "0.9 so experts actually offload" + ) + + +def observation_key( + model: str, baseline: str, block_size: int, concurrency: int, repeat: int +) -> Tuple[str, str, int, int, int]: + """The immutable identity of one measured row; two rows may never share it.""" + return (model, baseline, int(block_size), int(concurrency), int(repeat)) + + +def make_observation_row( + *, + model: str, + draft: str, + baseline: str, + block_size: int, + concurrency: int, + repeat: int, + metrics: Mapping[str, float], + cost_terms: Optional[Mapping[str, Any]] = None, + status: Optional[str] = None, + warnings: Optional[Sequence[str]] = None, +) -> Dict[str, Any]: + """Assemble one JSON observation row with the full §8 metric schema. + + A blocked row (``status`` set, e.g. B2's ``BLOCKED_UNTIL_2D_SCHEDULER``) + carries no metrics; any other row must supply every ``REQUIRED_METRICS`` + entry, mirroring ``validate_result_matrix`` so a malformed row fails fast at + write time rather than in the aggregator. + """ + row: Dict[str, Any] = { + "model": model, + "draft": draft, + "baseline": baseline, + "block_size": int(block_size), + "concurrency": int(concurrency), + "repeat": int(repeat), + } + if status is not None: + row["status"] = status + else: + missing = [m for m in REQUIRED_METRICS if m not in metrics] + if missing: + raise ValueError( + f"observation missing metrics: {', '.join(missing)}" + ) + for name in REQUIRED_METRICS: + row[name] = metrics[name] + if cost_terms: + row["cost_terms"] = dict(cost_terms) + if warnings: + row["warnings"] = list(warnings) + return row + + +def append_observation(output_path: str, row: Mapping[str, Any]) -> None: + """Append ``row`` to the output JSON list, never overwriting a prior row. + + Rows are keyed by ``observation_key``; a duplicate key raises rather than + silently clobbering an existing measurement (plan Task 2 step 6). + """ + rows: List[Dict[str, Any]] = load_observations(output_path) + new_key = observation_key( + row["model"], + row["baseline"], + row["block_size"], + row["concurrency"], + row["repeat"], + ) + for existing in rows: + existing_key = observation_key( + existing["model"], + existing["baseline"], + existing["block_size"], + existing["concurrency"], + existing["repeat"], + ) + if existing_key == new_key: + raise ValueError(f"refusing to overwrite existing row {new_key}") + rows.append(dict(row)) + directory = os.path.dirname(os.path.abspath(output_path)) + os.makedirs(directory, exist_ok=True) + tmp_path = f"{output_path}.tmp" + with open(tmp_path, "w", encoding="utf-8") as handle: + json.dump(rows, handle, indent=2, sort_keys=True) + handle.write("\n") + os.replace(tmp_path, output_path) + + +def load_observations(output_path: str) -> List[Dict[str, Any]]: + """Read the JSON list of rows at ``output_path`` (``[]`` when absent/empty).""" + if not os.path.exists(output_path) or os.path.getsize(output_path) == 0: + return [] + with open(output_path, "r", encoding="utf-8") as handle: + data = json.load(handle) + if not isinstance(data, list): + raise ValueError(f"{output_path} is not a JSON list of rows") + return data + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class RunnerArgs: + model: Optional[str] + draft: Optional[str] + offload_dir: Optional[str] + output: Optional[str] + baselines: Tuple[str, ...] + block_sizes: Tuple[int, ...] + concurrency: Tuple[int, ...] + requests: int + warmup_rounds: int + measured_h2d_gbps: Optional[float] + slo_ms: Optional[float] + seed: int + device_memory_ratio: float + dry_run_contract: bool + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="python -m benchmarks.dflash.pd_dflash_serving", + description=( + "Opt-in RTX PRO 6000 B0-B3 route-ahead serving experiment; emits " + "one JSON observation row per (model, baseline, block, " + "concurrency, repeat) for benchmarks.dflash.report." + ), + ) + parser.add_argument("--model", help="HF target repo (offloaded MoE)") + parser.add_argument("--draft", help="z-lab DFlash draft repo") + parser.add_argument("--offload-dir", help="expert offload directory") + parser.add_argument("--output", help="output JSON path for observations") + parser.add_argument( + "--baseline", + nargs="+", + default=list(DEFAULT_BASELINES), + choices=sorted(BASELINES) + list(EXPERIMENTAL_CONFIGURATIONS), + help="baselines/configurations to run (default: B0 B1 B2 B3)", + ) + parser.add_argument( + "--block-size", + nargs="+", + type=int, + default=list(REQUIRED_BLOCK_SIZES), + help="draft block sizes (default: 8 16)", + ) + parser.add_argument( + "--concurrency", + nargs="+", + type=int, + default=list(REQUIRED_CONCURRENCY), + help="concurrent request counts (default: 1 2 4 8 16 32)", + ) + parser.add_argument("--requests", type=int, default=64) + parser.add_argument("--warmup-rounds", type=int, default=5) + parser.add_argument("--measured-h2d-gbps", type=float, default=None) + parser.add_argument("--slo-ms", type=float, default=None) + parser.add_argument("--seed", type=int, default=1408) + parser.add_argument( + "--device-memory-ratio", + type=float, + default=0.85, + help="fraction of GPU memory for weights; <0.9 forces offload", + ) + parser.add_argument( + "--dry-run-contract", + action="store_true", + help="print the §8 experiment contract as JSON and exit (no GPU)", + ) + return parser + + +def parse_args(argv: Optional[Sequence[str]] = None) -> RunnerArgs: + parsed = build_arg_parser().parse_args(argv) + return RunnerArgs( + model=parsed.model, + draft=parsed.draft, + offload_dir=parsed.offload_dir, + output=parsed.output, + baselines=tuple(parsed.baseline), + block_sizes=tuple(parsed.block_size), + concurrency=tuple(parsed.concurrency), + requests=parsed.requests, + warmup_rounds=parsed.warmup_rounds, + measured_h2d_gbps=parsed.measured_h2d_gbps, + slo_ms=parsed.slo_ms, + seed=parsed.seed, + device_memory_ratio=parsed.device_memory_ratio, + dry_run_contract=parsed.dry_run_contract, + ) + + +def _require_run_args(args: RunnerArgs) -> None: + missing = [ + flag + for flag, value in ( + ("--model", args.model), + ("--draft", args.draft), + ("--offload-dir", args.offload_dir), + ("--output", args.output), + ) + if not value + ] + if missing: + raise SystemExit(f"missing required args: {', '.join(missing)}") + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = parse_args(argv) + if args.dry_run_contract: + json.dump(build_contract_matrix(), sys.stdout, indent=2, sort_keys=True) + sys.stdout.write("\n") + return 0 + _require_run_args(args) + return run_experiment(args) + + +# --------------------------------------------------------------------------- +# hardware path: lazily imports torch / moe_infinity so module import stays cheap +# --------------------------------------------------------------------------- + + +def _validate_gpu_environment() -> None: + import torch + + import moe_infinity._v4_fp4 # noqa: F401 (asserts native FP4 path present) + + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available on this host") + if torch.cuda.device_count() != 1: + raise RuntimeError( + "expected exactly one visible GPU; set CUDA_VISIBLE_DEVICES=" + ) + validate_device_identity( + torch.cuda.get_device_name(0), torch.cuda.get_device_capability(0) + ) + + +def run_experiment(args: RunnerArgs) -> int: + """Drive the B0-B3 matrix on one RTX PRO 6000 and write observation rows. + + Loads each configuration through the real ``MoE`` + ``DFlashSpeculator`` + serving path (mirroring ``tests/python/dflash/test_gpu_serving_dflash.py``), + wraps the draft/router/issue/verify/H2D phases in the frozen ``NVTX_RANGES``, + reads byte-accurate coverage/waste from the instrumented ``RouteAheadStats``, + and appends one row per ``(model, baseline, block, concurrency, repeat)``. + B2 is emitted as ``BLOCKED_UNTIL_2D_SCHEDULER`` until Task 6 lands. + """ + from benchmarks.dflash._serving_measure import measure_configuration + + assert args.model and args.draft and args.offload_dir and args.output + _validate_gpu_environment() + + draft = args.draft or REQUIRED_DRAFTS.get(args.model, "") + for baseline in args.baselines: + for block_size in args.block_sizes: + for concurrency in args.concurrency: + if baseline == "B2": + append_observation( + args.output, + make_observation_row( + model=args.model, + draft=draft, + baseline=baseline, + block_size=block_size, + concurrency=concurrency, + repeat=0, + metrics={}, + status=BLOCKED_STATUS, + ), + ) + continue + row = measure_configuration( + args=args, + baseline=baseline, + draft=draft, + block_size=block_size, + concurrency=concurrency, + ) + append_observation(args.output, row) + return 0 + + +if __name__ == "__main__": # pragma: no cover - CLI entry + raise SystemExit(main()) diff --git a/benchmarks/dflash/report.py b/benchmarks/dflash/report.py index 19e3b1b8..c465dbf1 100644 --- a/benchmarks/dflash/report.py +++ b/benchmarks/dflash/report.py @@ -14,9 +14,12 @@ from __future__ import annotations +import argparse +import json import math +import sys from dataclasses import dataclass -from typing import Mapping +from typing import Any, Dict, List, Mapping, Sequence, Tuple REQUIRED_METRICS = ( "output_tokens_per_second", @@ -129,21 +132,219 @@ def validate_result_matrix( row = rows[baseline] if baseline == "B3" and row.get("status") == UNAVAILABLE_CAPACITY: continue - for metric in REQUIRED_METRICS: - if metric not in row: - raise ValueError(f"{baseline} missing metric: {metric}") - value = row[metric] - if isinstance(value, bool) or not isinstance(value, (int, float)): - raise ValueError( - f"{baseline}.{metric} must be a finite, non-negative " - f"number; got {value!r}" - ) - number = float(value) - if not math.isfinite(number) or number < 0.0: - raise ValueError( - f"{baseline}.{metric} must be a finite, non-negative " - f"number; got {value!r}" - ) + _validate_metric_row(baseline, row) + + +def _validate_metric_row(baseline: str, row: Mapping[str, object]) -> None: + for metric in REQUIRED_METRICS: + if metric not in row: + raise ValueError(f"{baseline} missing metric: {metric}") + value = row[metric] + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError( + f"{baseline}.{metric} must be a finite, non-negative " + f"number; got {value!r}" + ) + number = float(value) + if not math.isfinite(number) or number < 0.0: + raise ValueError( + f"{baseline}.{metric} must be a finite, non-negative " + f"number; got {value!r}" + ) + + +def summarise_row(row: Mapping[str, object]) -> Mapping[str, object]: + """BM1 router-ahead cost summary for one row (design §10). + + BM1 passes when the route-ahead router projection is strictly cheaper than + the width-B verify it front-runs (``t_router < t_verify``); the raw seconds + and their ratio are retained so the aggregator can rank configurations, not + only gate them. Raises ``ValueError`` on a missing term, a negative time, or + a non-positive ``t_verify_seconds``. + """ + for key in ("t_router_seconds", "t_verify_seconds"): + if key not in row: + raise ValueError(f"row missing BM1 term: {key}") + t_router = _require_finite_non_negative( + "t_router_seconds", row["t_router_seconds"] + ) + t_verify = _require_finite_non_negative( + "t_verify_seconds", row["t_verify_seconds"] + ) + if t_verify <= 0.0: + raise ValueError(f"t_verify_seconds must be > 0; got {t_verify!r}") + return { + "t_router_seconds": t_router, + "t_verify_seconds": t_verify, + "bm1_router_to_verify_ratio": t_router / t_verify, + "bm1_pass": t_router < t_verify, + } + + +def aggregate_result_matrices( + rows: Sequence[Mapping[str, Any]], +) -> Dict[Tuple[str, int, int], Dict[str, Mapping[str, Any]]]: + """Group raw observation rows into §8 matrices. + + Keyed by ``(model, block_size, concurrency)``; each value maps a baseline + label to its single row. A duplicate ``(key, baseline)`` raises, mirroring + the runner's append-without-overwrite guarantee. + """ + matrices: Dict[Tuple[str, int, int], Dict[str, Mapping[str, Any]]] = {} + for row in rows: + key = ( + str(row["model"]), + int(row["block_size"]), + int(row["concurrency"]), + ) + baseline = str(row["baseline"]) + bucket = matrices.setdefault(key, {}) + if baseline in bucket: + raise ValueError(f"duplicate baseline {baseline} for {key}") + bucket[baseline] = row + return matrices + + +def evaluate_matrix( + baseline_rows: Mapping[str, Mapping[str, Any]], + allow_blocked: Sequence[str] = (), +) -> Tuple[bool, Dict[str, Any]]: + """Completeness + BM1 verdict for one grouped matrix. + + A baseline is satisfied when it carries the full metric schema; a baseline + listed in ``allow_blocked`` may instead carry a blocking ``status`` (e.g. + B2's ``BLOCKED_UNTIL_2D_SCHEDULER`` before the 2-D scheduler lands). Any + other missing/invalid/blocked baseline fails the group. BM1 summaries are + attached for every row carrying ``t_router_seconds``/``t_verify_seconds``. + """ + allow = set(allow_blocked) + detail: Dict[str, Any] = { + "present": sorted(baseline_rows), + "blocked": [], + "missing": [], + "invalid": [], + "bm1": {}, + } + ok = True + for baseline in REQUIRED_BASELINES: + row = baseline_rows.get(baseline) + if row is None: + detail["missing"].append(baseline) + ok = False + continue + if row.get("status"): + if baseline in allow: + detail["blocked"].append(baseline) + else: + detail["blocked"].append(baseline) + ok = False + continue + try: + _validate_metric_row(baseline, row) + except ValueError: + detail["invalid"].append(baseline) + ok = False + if "t_router_seconds" in row and "t_verify_seconds" in row: + detail["bm1"][baseline] = dict(summarise_row(row)) + return ok, detail + + +def _matrix_key(key: Tuple[str, int, int]) -> str: + return f"{key[0]}|B{key[1]}|c{key[2]}" + + +def _write_json(path: str, payload: Mapping[str, Any]) -> None: + with open(path, "w", encoding="utf-8") as handle: + json.dump(payload, handle, indent=2, sort_keys=True) + handle.write("\n") + + +def _write_csv( + path: str, + matrices: Mapping[Tuple[str, int, int], Mapping[str, Mapping[str, Any]]], +) -> None: + import csv + + columns = ["model", "block_size", "concurrency", "baseline", "status"] + columns += list(REQUIRED_METRICS) + with open(path, "w", encoding="utf-8", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=columns) + writer.writeheader() + for (model, block, conc), baseline_rows in sorted(matrices.items()): + for baseline, row in sorted(baseline_rows.items()): + record = { + "model": model, + "block_size": block, + "concurrency": conc, + "baseline": baseline, + "status": row.get("status", ""), + } + for metric in REQUIRED_METRICS: + record[metric] = row.get(metric, "") + writer.writerow(record) + + +def _write_markdown(path: str, report: Mapping[str, Any]) -> None: + lines = ["# PD-DFlash Phase-A result matrix", ""] + for group, detail in sorted(report.items()): + lines.append(f"## {group}") + lines.append(f"- present: {', '.join(detail['present']) or '(none)'}") + if detail["blocked"]: + lines.append(f"- blocked: {', '.join(detail['blocked'])}") + if detail["missing"]: + lines.append(f"- missing: {', '.join(detail['missing'])}") + if detail["invalid"]: + lines.append(f"- invalid: {', '.join(detail['invalid'])}") + for baseline, bm1 in sorted(detail["bm1"].items()): + lines.append( + f"- BM1 {baseline}: ratio=" + f"{bm1['bm1_router_to_verify_ratio']:.4f} " + f"pass={bm1['bm1_pass']}" + ) + lines.append("") + with open(path, "w", encoding="utf-8") as handle: + handle.write("\n".join(lines)) + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser( + prog="python -m benchmarks.dflash.report", + description="Aggregate PD-DFlash raw observation rows into §8 matrices.", + ) + parser.add_argument("--input", nargs="+", required=True) + parser.add_argument("--matrix-json") + parser.add_argument("--csv") + parser.add_argument("--markdown") + parser.add_argument("--allow-blocked", nargs="*", default=[]) + args = parser.parse_args(argv) + + rows: List[Mapping[str, Any]] = [] + for path in args.input: + with open(path, "r", encoding="utf-8") as handle: + data = json.load(handle) + rows.extend(data if isinstance(data, list) else [data]) + + matrices = aggregate_result_matrices(rows) + report: Dict[str, Any] = {} + all_ok = True + for key, baseline_rows in sorted(matrices.items()): + ok, detail = evaluate_matrix(baseline_rows, args.allow_blocked) + all_ok = all_ok and ok + report[_matrix_key(key)] = detail + + if args.matrix_json: + _write_json( + args.matrix_json, + {_matrix_key(k): dict(v) for k, v in matrices.items()}, + ) + if args.csv: + _write_csv(args.csv, matrices) + if args.markdown: + _write_markdown(args.markdown, report) + + json.dump(report, sys.stdout, indent=2, sort_keys=True) + sys.stdout.write("\n") + return 0 if all_ok else 1 __all__ = [ @@ -151,6 +352,14 @@ def validate_result_matrix( "REQUIRED_BASELINES", "UNAVAILABLE_CAPACITY", "HideInequality", + "aggregate_result_matrices", "evaluate_hide_inequality", + "evaluate_matrix", + "main", + "summarise_row", "validate_result_matrix", ] + + +if __name__ == "__main__": # pragma: no cover - CLI entry + raise SystemExit(main()) diff --git a/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/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 b6ead365..b187ee16 100644 --- a/moe_infinity/memory/expert_prefetcher.py +++ b/moe_infinity/memory/expert_prefetcher.py @@ -32,6 +32,7 @@ class ExpertPrefetcher(object): first_k_dense_replace: int = 0 archer_engine: Any expert_tensor_map: dict[tuple[int, int], int] + expert_nbytes_map: dict[tuple[int, int], int] def __init__(self, config: PretrainedConfig): print(config) @@ -40,6 +41,7 @@ def __init__(self, config: PretrainedConfig): ) self.archer_engine: Optional[Any] = None self.expert_tensor_map: Dict[Tuple[int, int], int] = {} + self.expert_nbytes_map: Dict[Tuple[int, int], int] = {} self._last_speculative_prediction: Set[int] = set() def set_archer_engine(self, archer_engine: Any): diff --git a/moe_infinity/runtime/model_offload.py b/moe_infinity/runtime/model_offload.py index 35627551..66c75a32 100644 --- a/moe_infinity/runtime/model_offload.py +++ b/moe_infinity/runtime/model_offload.py @@ -289,6 +289,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 = [] @@ -1059,6 +1084,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 index c1489f82..8743b412 100644 --- a/tests/python/dflash/test_pd_dflash_report.py +++ b/tests/python/dflash/test_pd_dflash_report.py @@ -22,7 +22,10 @@ from benchmarks.dflash.report import ( REQUIRED_METRICS, + aggregate_result_matrices, evaluate_hide_inequality, + evaluate_matrix, + summarise_row, validate_result_matrix, ) @@ -144,6 +147,111 @@ def test_wasted_prefetch_is_bytes_not_expert_count(): assert "wasted_prefetch_bytes" in REQUIRED_METRICS +# --------------------------------------------------------------------------- +# BM1 router-ahead cost aggregation (design §10) +# --------------------------------------------------------------------------- + + +def complete_row(**overrides) -> dict[str, float]: + row: dict[str, float] = {metric: 1.0 for metric in REQUIRED_METRICS} + row["t_router_seconds"] = 0.002 + row["t_verify_seconds"] = 0.020 + row.update(overrides) + return row + + +def test_bm1_reports_router_cost_against_verify(): + row = complete_row(t_router_seconds=0.002, t_verify_seconds=0.020) + summary = summarise_row(row) + assert summary["bm1_router_to_verify_ratio"] == 0.1 + assert summary["bm1_pass"] is True + + +def test_bm1_fails_when_router_not_cheaper_than_verify(): + summary = summarise_row( + complete_row(t_router_seconds=0.03, t_verify_seconds=0.02) + ) + assert summary["bm1_pass"] is False + assert summary["bm1_router_to_verify_ratio"] == 1.5 + + +def test_bm1_retains_raw_terms_and_rejects_bad_inputs(): + summary = summarise_row(complete_row()) + assert summary["t_router_seconds"] == 0.002 + assert summary["t_verify_seconds"] == 0.020 + with pytest.raises(ValueError, match="t_verify_seconds"): + summarise_row(complete_row(t_verify_seconds=0.0)) + with pytest.raises(ValueError, match="t_router_seconds"): + summarise_row(complete_row(t_router_seconds=-0.1)) + with pytest.raises(ValueError, match="t_router_seconds"): + summarise_row({"t_verify_seconds": 0.02}) + + +# --------------------------------------------------------------------------- +# aggregation: group raw rows into §8 matrices, allow blocked B2 +# --------------------------------------------------------------------------- + + +def _obs_row(baseline: str, **overrides) -> dict[str, object]: + row: dict[str, object] = { + "model": "M", + "draft": "d", + "baseline": baseline, + "block_size": 16, + "concurrency": 8, + "repeat": 0, + } + row.update({metric: 1.0 for metric in REQUIRED_METRICS}) + row.update(overrides) + return row + + +def _blocked_b2() -> dict[str, object]: + return { + "model": "M", + "draft": "d", + "baseline": "B2", + "block_size": 16, + "concurrency": 8, + "repeat": 0, + "status": "BLOCKED_UNTIL_2D_SCHEDULER", + } + + +def test_aggregate_groups_rows_and_rejects_duplicate_baseline(): + matrices = aggregate_result_matrices( + [_obs_row(b) for b in ("B0", "B1", "B3")] + ) + assert set(matrices) == {("M", 16, 8)} + assert set(matrices[("M", 16, 8)]) == {"B0", "B1", "B3"} + with pytest.raises(ValueError, match="duplicate baseline"): + aggregate_result_matrices([_obs_row("B0"), _obs_row("B0")]) + + +def test_evaluate_matrix_allows_blocked_b2_only_when_permitted(): + rows = {b: _obs_row(b) for b in ("B0", "B1", "B3")} + rows["B2"] = _blocked_b2() + ok, detail = evaluate_matrix(rows, allow_blocked=["B2"]) + assert ok is True and detail["blocked"] == ["B2"] + not_ok, detail2 = evaluate_matrix(rows, allow_blocked=[]) + assert not_ok is False and "B2" in detail2["blocked"] + + +def test_evaluate_matrix_flags_missing_baseline_and_attaches_bm1(): + partial, missing = evaluate_matrix( + {b: _obs_row(b) for b in ("B0", "B1")}, [] + ) + assert partial is False and set(missing["missing"]) == {"B2", "B3"} + + full = { + b: _obs_row(b, t_router_seconds=0.002, t_verify_seconds=0.020) + for b in ("B0", "B1", "B2", "B3") + } + ok, detail = evaluate_matrix(full, []) + assert ok is True + assert detail["bm1"]["B1"]["bm1_pass"] is True + + def test_required_metrics_are_frozen_and_complete(): assert isinstance(REQUIRED_METRICS, tuple) assert REQUIRED_METRICS == ( 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_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 # ---------------------------------------------------------------------------