#!/usr/bin/env python3 """Freeze a profile-only Frontier ranking for the Qwen235B prefill grid.""" from __future__ import annotations import argparse import csv import hashlib import json import os import subprocess import sys import time from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Iterable SCHEMA = "frontier-qwen235b-prefill-grid-v1" MODEL = "Qwen3-235B-A22B-FP8" TRACE_SHA256 = "f878e9af18f94dcfaced94a8e1e6b20a2f7d97d64aa862448025660dbbd965b2" PROFILE_RELATIVE = Path("compute/h20") / MODEL NETWORK_RELATIVE = Path("network/h20_nccl/all_reduce.csv") SEARCH_LOW = 0.0 SEARCH_HIGH = 0.125 SEARCH_PROBES = 6 WINDOW_DURATION_S = 600.0 MAX_INPUT_TOKENS = 32768 MAX_MODEL_TOKENS = 40960 BLOCK_SIZE_TOKENS = 16 MOE_ROUTING_SEED = 42 @dataclass(frozen=True) class GridConfig: tp: int mns: int mbt: int moe_tp: int moe_ep: int num_gpu_blocks: int @property def name(self) -> str: return f"tp{self.tp}_mns{self.mns}_mbt{self.mbt}" @property def gpu_count(self) -> int: return self.tp GRID = tuple( GridConfig( tp=tp, mns=mns, mbt=mbt, moe_tp=4 if tp == 4 else 1, moe_ep=1 if tp == 4 else 8, # Measured by community vLLM on the same checkpoint/runtime. The real # grid will use --num-gpu-blocks-override with these same values. num_gpu_blocks=26101 if tp == 4 else 62351, ) for tp in (4, 8) for mns in (64, 128) for mbt in (8192, 16384) ) def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as source: for chunk in iter(lambda: source.read(1 << 20), b""): digest.update(chunk) return digest.hexdigest() def order_hash(values: Iterable[object]) -> str: payload = "\n".join(str(value) for value in values).encode() return hashlib.sha256(payload).hexdigest() def write_json(path: Path, payload: Any) -> None: path.parent.mkdir(parents=True, exist_ok=True) temporary = path.with_suffix(path.suffix + ".tmp") temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") os.replace(temporary, path) def sha256_bytes(payload: bytes) -> str: return hashlib.sha256(payload).hexdigest() def frontier_source_fingerprint( source: Path, declared_commit: str ) -> dict[str, Any]: model_config = source / "data" / "config" / "models" / f"{MODEL}.json" if not model_config.is_file(): raise FileNotFoundError(f"Frontier model config is missing: {model_config}") device_config = source / "data" / "config" / "device" / "h20.json" files = sorted((source / "frontier").glob("**/*.py")) + [ model_config, device_config, ] missing = [str(path) for path in files if not path.is_file()] if missing: raise FileNotFoundError(f"Frontier fingerprint inputs are missing: {missing}") digest = hashlib.sha256() for path in files: relative = path.relative_to(source).as_posix() digest.update(relative.encode() + b"\0") digest.update(bytes.fromhex(sha256(path))) fingerprint = { "declared_upstream_commit": declared_commit, "python_and_config_tree_sha256": digest.hexdigest(), "fingerprinted_file_count": len(files), "model_config": { "path": str(model_config.resolve()), "sha256": sha256(model_config), }, } git_dir = source / ".git" if git_dir.exists(): commit = subprocess.run( ["git", "rev-parse", "HEAD"], cwd=source, check=True, stdout=subprocess.PIPE, text=True, ).stdout.strip() status = subprocess.run( ["git", "status", "--porcelain=v1"], cwd=source, check=True, stdout=subprocess.PIPE, text=True, ).stdout diff = subprocess.run( ["git", "diff", "--binary", "HEAD"], cwd=source, check=True, stdout=subprocess.PIPE, ).stdout fingerprint.update( { "git_commit": commit, "status_porcelain": status.splitlines(), "tracked_diff_sha256": sha256_bytes(diff), } ) else: fingerprint["git_metadata"] = "absent_source_snapshot" return fingerprint def anchor_key(anchor: float) -> str: return f"{anchor:.12f}" def anchor_filename(anchor: float) -> str: return f"u_{anchor_key(anchor).replace('.', 'p')}.csv" def binary_search_lattice() -> list[float]: intervals = [(SEARCH_LOW, SEARCH_HIGH)] anchors: list[float] = [] for _ in range(SEARCH_PROBES): next_intervals = [] for low, high in intervals: midpoint = (low + high) / 2.0 anchors.append(midpoint) next_intervals.extend(((low, midpoint), (midpoint, high))) intervals = next_intervals return sorted(set(anchors)) def ttft_slo_ms(input_tokens: int) -> float: return 1000.0 if input_tokens <= 8191 else 2000.0 def load_trace_rows(path: Path) -> list[dict[str, Any]]: rows = [] with path.open(encoding="utf-8") as source: for source_index, line in enumerate(source): if not line.strip(): continue raw = json.loads(line) input_tokens = int(raw["input_length"]) if not 0 <= input_tokens <= MAX_INPUT_TOKENS: continue timestamp = float(raw["timestamp"]) sampling_u = float(raw["sampling_u"]) rows.append( { "arrived_at": timestamp, "num_prefill_tokens": input_tokens, "num_decode_tokens": 1, "source_row_index": source_index, "source_request_id": str( raw.get("request_id") or raw.get("id") or source_index ), "sampling_u": sampling_u, "slo_ttft_ms": ttft_slo_ms(input_tokens), } ) # This reproduces AITuner's stable arrival-only sort. Tied timestamps retain # source order; sampling_u must not become a secondary key. rows.sort(key=lambda row: row["arrived_at"]) return rows def prepare_traces(trace: Path, output_root: Path, expected_sha256: str) -> Path: actual_sha256 = sha256(trace) if expected_sha256 and actual_sha256 != expected_sha256: raise ValueError( f"trace SHA256 mismatch: expected={expected_sha256}, actual={actual_sha256}" ) rows = load_trace_rows(trace) if not rows: raise ValueError("no trace rows remain after input-length filtering") if any( rows[index]["arrived_at"] > rows[index + 1]["arrived_at"] for index in range(len(rows) - 1) ): raise ValueError("materialized trace is not monotonic by arrival") trace_dir = output_root / "traces" trace_dir.mkdir(parents=True, exist_ok=True) fields = [ "arrived_at", "num_prefill_tokens", "num_decode_tokens", "source_row_index", "source_request_id", "sampling_u", "slo_ttft_ms", ] anchors: dict[str, Any] = {} for anchor in binary_search_lattice(): selected = [row for row in rows if row["sampling_u"] <= anchor] target = trace_dir / anchor_filename(anchor) with target.open("w", encoding="utf-8", newline="") as output: writer = csv.DictWriter(output, fieldnames=fields) writer.writeheader() writer.writerows(selected) key = anchor_key(anchor) anchors[key] = { "anchor": anchor, "path": str(target.resolve()), "sha256": sha256(target), "request_count": len(selected), "request_rate": len(selected) / WINDOW_DURATION_S, "source_row_order_sha256": order_hash( row["source_row_index"] for row in selected ), "arrival_order_sha256": order_hash( f"{row['arrived_at']:.12f}" for row in selected ), "input_length_order_sha256": order_hash( row["num_prefill_tokens"] for row in selected ), } manifest = { "schema": SCHEMA, "source": { "path": str(trace.resolve()), "sha256": actual_sha256, "filtered_request_count": len(rows), }, "selection_contract": { "input_tokens": [0, MAX_INPUT_TOKENS], "completion_tokens_override": 1, "stable_sort_key": "arrived_at_only", "selection": "sampling_u <= anchor", "window_duration_s": WINDOW_DURATION_S, "search_low": SEARCH_LOW, "search_high": SEARCH_HIGH, "search_probes": SEARCH_PROBES, "lattice_anchor_count": len(anchors), "ttft_slo_ms": {"input_le_8191": 1000, "otherwise": 2000}, "target_pass_rate": 0.95, }, "anchors": anchors, } manifest_path = output_root / "trace_manifest.json" write_json(manifest_path, manifest) return manifest_path def resolve_profile_paths(profile_root: Path) -> dict[str, Path]: paths = { "linear": profile_root / PROFILE_RELATIVE / "linear_op.csv", "attention": profile_root / PROFILE_RELATIVE / "attention.csv", "moe": profile_root / PROFILE_RELATIVE / "moe.csv", "all_reduce": profile_root / NETWORK_RELATIVE, "manifest": profile_root / "profile_manifest.json", } missing = [f"{name}={path}" for name, path in paths.items() if not path.is_file()] if missing: raise FileNotFoundError("missing profile inputs: " + ", ".join(missing)) return paths def build_command( *, python: Path, frontier_source: Path, profile_root: Path, profile_paths: dict[str, Path], trace: Path, config: GridConfig, probe_dir: Path, run_id: str, cache_root: Path, ) -> list[str]: return [ str(python), "-m", "frontier.main", "--simulation_mode", "offline", # Frontier's offline default rewrites every request arrival to time 0. # Preserve the trace arrival process so the simulator and serving run # exercise the same queueing workload. "--offline_use_generated_request_arrivals", "--sys_arch", "co-location", "--cluster_config_num_replicas", "1", "--replica_config_model_name", MODEL, "--replica_config_attn_tensor_parallel_size", str(config.tp), "--replica_config_attn_data_parallel_size", "1", "--replica_config_moe_tensor_parallel_size", str(config.moe_tp), "--replica_config_moe_expert_parallel_size", str(config.moe_ep), "--replica_config_total_expert_num", "128", "--replica_config_router_topk", "8", "--replica_config_moe_routing_mode", "simulation", "--replica_config_moe_routing_seed", str(MOE_ROUTING_SEED), "--replica_config_num_pipeline_stages", "1", "--replica_config_device", "h20", "--replica_config_network_device", "h20_dgx", "--cc_backend_config_type", "vidur", "--vidur_cc_backend_config_profiling_data_dir", str(profile_root), "--vidur_cc_backend_config_cache_dir", str(cache_root / "collectives"), "--vidur_cc_backend_config_all_reduce_input_file", str(profile_paths["all_reduce"]), "--replica_scheduler_config_type", "vllm_v1", "--decode_cuda_graph_mode", "none", "--vllm_v1_scheduler_config_batch_size_cap", str(config.mns), "--vllm_v1_scheduler_config_block_size", str(BLOCK_SIZE_TOKENS), "--vllm_v1_scheduler_config_num_blocks", str(config.num_gpu_blocks), "--vllm_v1_scheduler_config_num_blocks_mode", "explicit", "--vllm_v1_scheduler_config_max_tokens_in_batch", str(config.mbt), "--vllm_v1_scheduler_config_enable_chunked_prefill", "--no-vllm_v1_scheduler_config_enable_prefix_caching", "--request_generator_config_type", "trace_replay", "--trace_request_generator_config_trace_file", str(trace), "--trace_request_generator_config_time_scale_factor", "1", "--trace_request_generator_config_prefill_scale_factor", "1", "--trace_request_generator_config_decode_scale_factor", "1", "--trace_request_generator_config_max_tokens", str(MAX_MODEL_TOKENS), "--no-random_forrest_execution_time_predictor_config_enable_dummy_mode", "--random_forrest_execution_time_predictor_config_linear_op_input_file", str(profile_paths["linear"]), "--random_forrest_execution_time_predictor_config_atten_input_file", str(profile_paths["attention"]), "--random_forrest_execution_time_predictor_config_moe_input_file", str(profile_paths["moe"]), "--random_forrest_execution_time_predictor_config_all_reduce_input_file", str(profile_paths["all_reduce"]), "--random_forrest_execution_time_predictor_config_prediction_max_prefill_chunk_size", "16384", "--random_forrest_execution_time_predictor_config_prediction_max_tokens_per_request", str(MAX_INPUT_TOKENS + 1), "--random_forrest_execution_time_predictor_config_prediction_max_batch_size", "128", "--random_forrest_execution_time_predictor_config_skip_cpu_overhead_modeling", "--metrics_config_cache_dir", str(cache_root / "execution"), "--metrics_config_output_dir", str(probe_dir / "metrics"), "--metrics_config_run_id", run_id, "--metrics_config_write_metrics", "--metrics_config_store_request_metrics", "--no-metrics_config_store_plots", "--no-metrics_config_enable_chrome_trace", "--no-metrics_config_write_json_trace", ] def score_request_metrics(trace: Path, request_metrics: Path) -> dict[str, Any]: with trace.open(encoding="utf-8", newline="") as source: trace_rows = list(csv.DictReader(source)) with request_metrics.open(encoding="utf-8", newline="") as source: metric_rows = list(csv.DictReader(source)) if len(metric_rows) != len(trace_rows): raise ValueError( f"request count mismatch: trace={len(trace_rows)}, metrics={len(metric_rows)}" ) metrics_by_id = {int(row["Request Id"]): row for row in metric_rows} expected_ids = set(range(len(trace_rows))) if set(metrics_by_id) != expected_ids: raise ValueError("Frontier Request Ids do not match trace row positions") passed = 0 ttfts = [] for request_id, trace_row in enumerate(trace_rows): ttft = float(metrics_by_id[request_id]["ttft"]) threshold = float(trace_row["slo_ttft_ms"]) ttfts.append(ttft) passed += int(ttft <= threshold) count = len(trace_rows) pass_rate = passed / count if count else 0.0 ordered = sorted(ttfts) def percentile(fraction: float) -> float | None: if not ordered: return None index = round((len(ordered) - 1) * fraction) return ordered[index] return { "request_count": count, "passed_request_count": passed, "slo_pass_rate": pass_rate, "feasible": pass_rate >= 0.95, "ttft_ms": { "min": min(ordered) if ordered else None, "p50": percentile(0.50), "p95": percentile(0.95), "p99": percentile(0.99), "max": max(ordered) if ordered else None, }, } def find_request_metrics(probe_dir: Path) -> Path: candidates = list((probe_dir / "metrics").glob("**/request_metrics.csv")) if len(candidates) != 1: raise ValueError( f"expected one request_metrics.csv under {probe_dir}, got {candidates}" ) return candidates[0] def run_probe( *, python: Path, frontier_source: Path, profile_root: Path, profile_paths: dict[str, Path], trace_record: dict[str, Any], config: GridConfig, probe_index: int, output_root: Path, cache_root: Path, ) -> dict[str, Any]: anchor = float(trace_record["anchor"]) probe_dir = output_root / "runs" / config.name / f"probe_{probe_index}_{anchor_filename(anchor)[:-4]}" result_path = probe_dir / "result.json" if result_path.is_file(): result = json.loads(result_path.read_text()) if result.get("status") == "completed": return result probe_dir.mkdir(parents=True, exist_ok=True) trace = Path(trace_record["path"]) run_id = f"{config.name}_probe{probe_index}_{anchor_filename(anchor)[:-4]}" command = build_command( python=python, frontier_source=frontier_source, profile_root=profile_root, profile_paths=profile_paths, trace=trace, config=config, probe_dir=probe_dir, run_id=run_id, cache_root=cache_root, ) write_json(probe_dir / "command.json", command) environment = os.environ.copy() environment.update( { "PYTHONPATH": str(frontier_source), "WANDB_DISABLED": "true", "VIDUR_DISABLE_WANDB": "1", } ) started = time.time() with (probe_dir / "stdout.log").open("w", encoding="utf-8") as output: completed = subprocess.run( command, cwd=frontier_source, env=environment, stdout=output, stderr=subprocess.STDOUT, check=False, ) elapsed = time.time() - started if completed.returncode != 0: result = { "status": "failed", "returncode": completed.returncode, "elapsed_seconds": elapsed, "config": asdict(config), "anchor": anchor, "trace_sha256": trace_record["sha256"], } write_json(result_path, result) raise RuntimeError(f"Frontier probe failed: {config.name}, anchor={anchor}") request_metrics = find_request_metrics(probe_dir) score = score_request_metrics(trace, request_metrics) result = { "status": "completed", "elapsed_seconds": elapsed, "config": asdict(config), "anchor": anchor, "trace_path": str(trace), "trace_sha256": trace_record["sha256"], "request_metrics_path": str(request_metrics), "request_metrics_sha256": sha256(request_metrics), "request_rate": score["request_count"] / WINDOW_DURATION_S, "request_rate_per_gpu": score["request_count"] / WINDOW_DURATION_S / config.gpu_count, **score, } write_json(result_path, result) return result def selected_configs(names: list[str] | None) -> list[GridConfig]: if not names: return list(GRID) by_name = {config.name: config for config in GRID} unknown = sorted(set(names) - set(by_name)) if unknown: raise ValueError(f"unknown configs: {unknown}; available={sorted(by_name)}") return [by_name[name] for name in names] def run_grid(args: argparse.Namespace) -> None: trace_manifest = json.loads(args.trace_manifest.read_text()) if trace_manifest.get("schema") != SCHEMA: raise ValueError(f"unexpected trace manifest schema: {trace_manifest.get('schema')}") profile_paths = resolve_profile_paths(args.profile_root) args.output_root.mkdir(parents=True, exist_ok=True) cache_root = args.output_root / "cache" configs = selected_configs(args.config) run_manifest = { "schema": SCHEMA, "frontier": { "source": str(args.frontier_source.resolve()), "python": str(args.python.resolve()), "fingerprint": frontier_source_fingerprint( args.frontier_source, args.frontier_commit ), }, "trace_manifest": { "path": str(args.trace_manifest.resolve()), "sha256": sha256(args.trace_manifest), }, "profiles": { name: {"path": str(path.resolve()), "sha256": sha256(path)} for name, path in profile_paths.items() }, "kv_capacity_evidence": { "tp4": { "path": str(args.tp4_capacity_artifact.resolve()), "sha256": sha256(args.tp4_capacity_artifact), "num_gpu_blocks": 26101, }, "tp8": { "path": str(args.tp8_capacity_artifact.resolve()), "sha256": sha256(args.tp8_capacity_artifact), "num_gpu_blocks": 62351, }, }, "contract": { "metric": "maximum SLO-feasible request_rate_per_gpu", "target_pass_rate": 0.95, "cpu_overhead_modeling": "skipped_no_community_vllm_native_records", "end_to_end_calibration": False, "moe_routing_mode": "simulation", "moe_routing_seed": MOE_ROUTING_SEED, "kv_blocks_source": "community_vllm_measured_and_fixed_per_topology", "network_device": "h20_dgx", }, "configs": [asdict(config) | {"name": config.name} for config in configs], } write_json(args.output_root / "run_manifest.json", run_manifest) all_results = {} for config in configs: low = SEARCH_LOW high = SEARCH_HIGH probes = [] best = None for probe_index in range(SEARCH_PROBES): anchor = (low + high) / 2.0 record = trace_manifest["anchors"].get(anchor_key(anchor)) if record is None: raise ValueError(f"trace manifest lacks anchor {anchor_key(anchor)}") result = run_probe( python=args.python, frontier_source=args.frontier_source, profile_root=args.profile_root, profile_paths=profile_paths, trace_record=record, config=config, probe_index=probe_index, output_root=args.output_root, cache_root=cache_root, ) probes.append(result) print( json.dumps( { "config": config.name, "probe": probe_index, "anchor": anchor, "pass_rate": result["slo_pass_rate"], "feasible": result["feasible"], "elapsed_seconds": result["elapsed_seconds"], }, sort_keys=True, ), flush=True, ) if result["feasible"]: low = anchor best = result else: high = anchor summary = { "config": asdict(config) | {"name": config.name}, "probes": probes, "capacity_interval": [low, high], "best_feasible": best, } write_json(args.output_root / "results" / f"{config.name}.json", summary) all_results[config.name] = summary ranked = sorted( all_results.values(), key=lambda item: ( -( item["best_feasible"]["request_rate_per_gpu"] if item["best_feasible"] is not None else -1.0 ), item["config"]["name"], ), ) freeze = { "schema": SCHEMA, "run_manifest_sha256": sha256(args.output_root / "run_manifest.json"), "ranking": [ { "rank": index + 1, "config": item["config"], "capacity_interval_sampling_u": item["capacity_interval"], "best_feasible": item["best_feasible"], } for index, item in enumerate(ranked) ], } write_json(args.output_root / "frontier_ranking_frozen.json", freeze) print(json.dumps(freeze, indent=2), flush=True) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() subparsers = parser.add_subparsers(dest="command", required=True) prepare = subparsers.add_parser("prepare") prepare.add_argument("--trace", type=Path, required=True) prepare.add_argument("--output-root", type=Path, required=True) prepare.add_argument("--expected-trace-sha256", default=TRACE_SHA256) run = subparsers.add_parser("run") run.add_argument("--frontier-source", type=Path, required=True) run.add_argument( "--frontier-commit", default="d9cfeb6d8791fbf2f295dd9744c56a666171776e", ) run.add_argument("--python", type=Path, required=True) run.add_argument("--profile-root", type=Path, required=True) run.add_argument("--trace-manifest", type=Path, required=True) run.add_argument("--output-root", type=Path, required=True) run.add_argument("--tp4-capacity-artifact", type=Path, required=True) run.add_argument("--tp8-capacity-artifact", type=Path, required=True) run.add_argument("--config", action="append") return parser.parse_args() def main() -> None: args = parse_args() if args.command == "prepare": manifest = prepare_traces( args.trace, args.output_root, args.expected_trace_sha256 ) print(manifest) return if args.command == "run": run_grid(args) return raise AssertionError(args.command) if __name__ == "__main__": main()