#!/usr/bin/env python3 """Build profile-v5-kvgrowth: profile-v4-trace-final + chunked-prefill KV-context grid. Appends attention rows built from the new raw flashattn-kvgrowth-tp{1,2,4}.json artifacts to the frozen profile-v4 attention.csv. moe/linear are copied unchanged (the fix targets only the attention KV-growth term; one variable at a time). Row schema replicates freeze_frontier_profiles.py's single-segment prefill path (mode=vllm020_batch_spec) so the Frontier predictor ingests the new rows exactly like the existing chunked-prefill samples. Anchor check: re-measured q1ks8k/q512s4k must agree with profile-v4 within ANCHOR_TOLERANCE, otherwise the profiler stack drifted and the merge aborts. """ from __future__ import annotations import argparse import csv import hashlib import json import math import re import shutil from pathlib import Path from typing import Any STAT_NAMES = ("min", "max", "mean", "median", "std") ATTENTION_OPS = ( "attn_input_reshape", "attn_kv_cache_save", "attn_prefill", "attn_decode", "attn_output_reshape", ) ANCHOR_SPECS = {"q1ks8k", "q512s4k"} ANCHOR_TOLERANCE = 0.25 # relative diff on attn core mean vs profile-v4 def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--base", type=Path, required=True, help="profile-v4-trace-final directory") parser.add_argument("--raw", type=Path, nargs="+", required=True, help="flashattn-kvgrowth-tp*.json raw artifacts") parser.add_argument("--output", type=Path, required=True, help="profile-v5-kvgrowth output directory") parser.add_argument("--max-model-len", type=int, default=40960) return parser.parse_args() def parse_size(value: str, suffix: str) -> int: return int(value) * (1024 if suffix == "k" else 1) def parse_batch_spec(spec: str) -> list[tuple[int, int]]: requests: list[tuple[int, int]] = [] pattern = re.compile(r"^(?:(\d+))?q(\d+)(k?)(?:s(\d+)(k?))?$") for segment in spec.split("_"): match = pattern.match(segment) if match is None: raise ValueError(f"invalid vLLM batch spec: {spec}") count = int(match.group(1) or 1) query = parse_size(match.group(2), match.group(3)) seq = parse_size(match.group(4), match.group(5)) if match.group(4) else query requests.extend([(query, seq)] * count) return requests def stat_columns(prefix: str, stats: dict[str, float]) -> dict[str, float]: return {f"time_stats.{prefix}.{name}": stats[name] for name in STAT_NAMES} def zero_stat_columns(prefix: str) -> dict[str, float]: return {f"time_stats.{prefix}.{name}": 0.0 for name in STAT_NAMES} def attention_core_stats(raw: dict[str, Any]) -> dict[str, float]: # Same policy as freeze_frontier_profiles.py: benchmark exports aggregate # mean, not samples; mean stands in for median. return { "min": 1000.0 * float(raw["min_time"]), "max": 1000.0 * float(raw["max_time"]), "mean": 1000.0 * float(raw["mean_time"]), "median": 1000.0 * float(raw["mean_time"]), "std": 1000.0 * float(raw["std_time"]), } def kv_update_stats(raw: dict[str, Any]) -> dict[str, float]: stats = raw["kv_cache_update_time"] return {name: float(stats[f"{name}_ms"]) for name in STAT_NAMES} def build_row(raw: dict[str, Any], max_model_len: int = 40960) -> dict[str, Any]: spec = raw["config"]["batch_spec"] tp = int(raw["tensor_parallel_size"]) requests = parse_batch_spec(spec) if any(q == 1 for q, _ in requests): raise ValueError(f"decode segment in prefill grid spec: {spec}") queries = [q for q, _ in requests] contexts = [kv - q for q, kv in requests] avg_query = sum(queries) / len(queries) variance = sum((q - avg_query) ** 2 for q in queries) / len(queries) std = math.sqrt(variance) avg_context = sum(contexts) / len(contexts) row: dict[str, Any] = {} for op in ATTENTION_OPS: row.update(zero_stat_columns(op)) row.update(stat_columns("attn_kv_cache_save", kv_update_stats(raw))) row.update(stat_columns("attn_prefill", attention_core_stats(raw))) row.update( { "n_embd": 2048, "n_q_head": 32, "n_kv_head": 4, "block_size": 16, "num_tensor_parallel_workers": tp, "max_model_len": max_model_len, "batch_size": len(requests), "prefill_chunk_size": sum(queries), "kv_cache_size": avg_context, "is_prefill": True, "attention_backend": "FLASH_ATTN", "is_mixed_batch": False, "mode": "vllm020_batch_spec", "seq_lens": json.dumps(queries), "total_tokens": sum(queries), "max_seq_len": max(queries), "min_seq_len": min(queries), "avg_seq_len": avg_query, "equal_seq_len": len(set(queries)) == 1, "seq_len_variance": variance, "seq_len_std": std, "seq_len_cv": std / avg_query if avg_query else 0.0, "is_chunked_prefill_sample": avg_context > 0, "chunk_start_token": avg_context, "chunk_end_token": avg_context + sum(queries), "total_prefill_tokens": sum(queries), "profiling_precision": "BF16", "model_arch": "generic", "quant_signature": "none", "measurement_type": "CUDA_EVENT", "is_true_mixed_batch": False, "prefill_seq_lens": "", "prefill_kv_cache_sizes": "", "decode_kv_cache_sizes": "", "num_prefill_seqs": "", "num_decode_seqs": "", "decode_batch_size": "", "total_batch_size": "", "total_decode_tokens": "", "decode_avg_kv_cache_size": "", "batch_spec": spec, "projection_policy": ( "measured_FA3_core_plus_measured_KV;reshape_assumed_zero;" "mean_as_median" ), } ) return row def check_anchor(row: dict[str, Any], base_rows: list[dict[str, str]]) -> str: tp = str(row["num_tensor_parallel_workers"]) spec = row["batch_spec"] matches = [ r for r in base_rows if r["batch_spec"] == spec and r["num_tensor_parallel_workers"] == tp ] if not matches: return f"anchor {spec}/TP{tp}: no profile-v4 row (skipped)" old = float(matches[0]["time_stats.attn_prefill.mean"]) new = float(row["time_stats.attn_prefill.mean"]) rel = abs(new - old) / old if old else float("inf") verdict = f"anchor {spec}/TP{tp}: v4={old:.4f}ms new={new:.4f}ms rel_diff={rel:.1%}" if rel > ANCHOR_TOLERANCE: raise SystemExit(f"PROFILER DRIFT: {verdict} exceeds {ANCHOR_TOLERANCE:.0%}") return verdict def sha256(path: Path) -> str: return hashlib.sha256(path.read_bytes()).hexdigest() def main() -> None: args = parse_args() base_csv = args.base / "attention.csv" with base_csv.open(newline="") as stream: reader = csv.DictReader(stream) fields = list(reader.fieldnames or []) base_rows = list(reader) new_rows: list[dict[str, Any]] = [] anchors: list[str] = [] seen: set[tuple[str, str]] = set() for raw_path in args.raw: raw_payload = json.loads(raw_path.read_text()) for raw in raw_payload["rows"]: if raw.get("error"): raise SystemExit(f"{raw_path}: failed sample {raw['config']['batch_spec']}") row = build_row(raw, args.max_model_len) key = (row["batch_spec"], str(row["num_tensor_parallel_workers"])) if key in seen: raise SystemExit(f"duplicate sample {key} across raw inputs") seen.add(key) if row["batch_spec"] in ANCHOR_SPECS: anchors.append(check_anchor(row, base_rows)) else: new_rows.append(row) args.output.mkdir(parents=True, exist_ok=True) out_csv = args.output / "attention.csv" with out_csv.open("w", newline="") as stream: writer = csv.DictWriter(stream, fieldnames=fields) writer.writeheader() writer.writerows(base_rows) for row in new_rows: writer.writerow({field: row.get(field, "") for field in fields}) for name in ("moe.csv", "linear_op.csv"): shutil.copy2(args.base / name, args.output / name) manifest = { "schema": "frontier-profile-v5-kvgrowth-v1", "base": str(base_csv), "base_sha256": sha256(base_csv), "raw_inputs": {str(p): sha256(p) for p in args.raw}, "appended_rows": len(new_rows), "max_model_len": args.max_model_len, "anchor_checks": anchors, "output_sha256": sha256(out_csv), } (args.output / "manifest.json").write_text(json.dumps(manifest, indent=2)) print(json.dumps(manifest, indent=2)) if __name__ == "__main__": main()