243 lines
8.9 KiB
Python
243 lines
8.9 KiB
Python
#!/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)
|
|
parser.add_argument(
|
|
"--manifest-schema",
|
|
default="frontier-profile-v5-kvgrowth-v1",
|
|
)
|
|
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": args.manifest_schema,
|
|
"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()
|