Files
aituner/runs/frontier-qwen30-vllm020-profile-v1/freeze_frontier_profiles.py

696 lines
28 KiB
Python

#!/usr/bin/env python3
"""Freeze vLLM 0.20 microprofiles into Frontier-compatible CSV inputs."""
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",
)
ATTENTION_METADATA = (
"n_embd",
"n_q_head",
"n_kv_head",
"block_size",
"num_tensor_parallel_workers",
"max_model_len",
"batch_size",
"prefill_chunk_size",
"kv_cache_size",
"is_prefill",
"attention_backend",
"is_mixed_batch",
"mode",
"seq_lens",
"total_tokens",
"max_seq_len",
"min_seq_len",
"avg_seq_len",
"equal_seq_len",
"seq_len_variance",
"seq_len_std",
"seq_len_cv",
"is_chunked_prefill_sample",
"chunk_start_token",
"chunk_end_token",
"total_prefill_tokens",
"profiling_precision",
"model_arch",
"quant_signature",
"measurement_type",
"is_true_mixed_batch",
"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_composition_ratio",
"batch_spec",
"projection_policy",
)
MOE_OPS = (
"moe_gating_linear",
"moe_gating_routing_topk",
"moe_shuffling",
"moe_grouped_gemm",
)
MOE_METADATA = (
"num_tokens",
"num_experts",
"num_experts_per_device",
"expert_parallel_size",
"routing_runtime_path",
"routing_assignment_policy",
"routing_weight_policy",
"routing_uses_router_logits",
"gating_runtime_context",
"gating_runtime_context_impl",
"router_topk",
"hidden_dim",
"expert_hidden_dim",
"use_gated",
"num_tensor_parallel_workers",
"total_routed_tokens",
"model_expansion_ratio",
"tokens_per_expert_avg",
"tokens_to_experts_ratio",
"expert_utilization",
"min_load_ratio",
"load_imbalance_cv",
"max_load_ratio",
"load_entropy",
"load_gini_coefficient",
"load_distribution",
"seed",
"moe_grouped_gemm_backend",
"measurement_type",
"profiling_precision",
"model_arch",
"quant_signature",
"router_median_nonadditivity_ratio",
"projection_policy",
)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--linear", type=Path, required=True)
parser.add_argument("--attention", type=Path, nargs=3, required=True)
parser.add_argument("--moe", type=Path, required=True)
parser.add_argument("--router", type=Path, required=True)
parser.add_argument("--allreduce", type=Path, nargs=2, required=True)
parser.add_argument("--output", type=Path, required=True)
return parser.parse_args()
def sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def load_json(path: Path) -> dict[str, Any]:
return json.loads(path.read_text())
def stat_columns(prefix: str, stats: dict[str, float]) -> dict[str, float]:
return {f"time_stats.{prefix}.{name}": float(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]:
# vLLM's benchmark result exports aggregate mean but not the raw samples.
# Preserve that mean as Frontier's training target and record the proxy in
# the manifest rather than inventing an unobserved 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 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))
kv = (
parse_size(match.group(4), match.group(5))
if match.group(4)
else query
)
requests.extend([(query, kv)] * count)
return requests
def write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, Any]]) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
with path.open("w", newline="") as handle:
writer = csv.DictWriter(handle, fieldnames=fieldnames, extrasaction="raise")
writer.writeheader()
writer.writerows(rows)
def freeze_attention(
inputs: list[Path], output: Path
) -> tuple[int, int, list[str]]:
rows: list[dict[str, Any]] = []
mixed_rows: list[dict[str, Any]] = []
seen_tps: set[int] = set()
raw_by_tp: dict[int, list[dict[str, Any]]] = {}
for path in inputs:
payload = load_json(path)
if payload.get("schema_version") != "qwen30_vllm020_flashattn_raw.v1":
raise ValueError(f"unexpected attention schema in {path}")
if payload["environment"].get("vllm_version") != "0.20.0":
raise ValueError(f"unexpected vLLM version in {path}")
for raw in payload["rows"]:
if raw.get("error") is not None:
raise ValueError(f"failed attention row in {path}: {raw['error']}")
tp = int(raw["tensor_parallel_size"])
seen_tps.add(tp)
raw_by_tp.setdefault(tp, []).append(raw)
def pure_reference_ms(
tp: int, requests: list[tuple[int, int]], *, decode_phase: bool
) -> float:
candidates: list[tuple[list[tuple[int, int]], float]] = []
for candidate in raw_by_tp[tp]:
parsed = parse_batch_spec(candidate["config"]["batch_spec"])
is_decode = all(query == 1 for query, _ in parsed)
if is_decode != decode_phase:
continue
if decode_phase and not is_decode:
continue
if not decode_phase and any(query == 1 for query, _ in parsed):
continue
candidates.append((parsed, 1000.0 * float(candidate["mean_time"])))
for parsed, mean_ms in candidates:
if parsed == requests:
return mean_ms
if not decode_phase:
raise ValueError(f"no exact pure prefill reference for TP{tp}: {requests}")
if len({kv for _, kv in requests}) != 1:
raise ValueError(f"decode interpolation requires one KV length: {requests}")
target_batch = len(requests)
target_kv = requests[0][1]
points = sorted(
(parsed[0][1], mean_ms)
for parsed, mean_ms in candidates
if len(parsed) == target_batch
and len({kv for _, kv in parsed}) == 1
)
if not points:
raise ValueError(f"no pure decode reference for TP{tp}: {requests}")
if target_kv <= points[0][0]:
return points[0][1]
if target_kv >= points[-1][0]:
return points[-1][1]
for (left_kv, left_ms), (right_kv, right_ms) in zip(points, points[1:]):
if left_kv <= target_kv <= right_kv:
fraction = (target_kv - left_kv) / (right_kv - left_kv)
return left_ms + fraction * (right_ms - left_ms)
raise AssertionError("unreachable decode interpolation")
for tp in sorted(raw_by_tp):
for raw in raw_by_tp[tp]:
spec = raw["config"]["batch_spec"]
requests = parse_batch_spec(spec)
prefill = [(q, kv) for q, kv in requests if q > 1]
decode = [(q, kv) for q, kv in requests if q == 1]
core = attention_core_stats(raw)
kv_stats = kv_update_stats(raw)
if prefill and decode:
prefill_reference_ms = pure_reference_ms(
tp, prefill, decode_phase=False
)
decode_reference_ms = pure_reference_ms(
tp, decode, decode_phase=True
)
reference_total_ms = prefill_reference_ms + decode_reference_ms
prefill_share = prefill_reference_ms / reference_total_ms
decode_share = decode_reference_ms / reference_total_ms
projected_prefill = {
name: value * prefill_share for name, value in core.items()
}
projected_decode = {
name: value * decode_share for name, value in core.items()
}
row = {}
for op in ATTENTION_OPS:
row.update(zero_stat_columns(op))
row.update(stat_columns("attn_kv_cache_save", kv_stats))
row.update(stat_columns("attn_prefill", projected_prefill))
row.update(stat_columns("attn_decode", projected_decode))
prefill_queries = [q for q, _ in prefill]
prefill_contexts = [kv - q for q, kv in prefill]
decode_kv_lengths = [kv for _, kv in decode]
total_batch = len(requests)
row.update(
{
"n_embd": 2048,
"n_q_head": 32,
"n_kv_head": 4,
"block_size": 16,
"num_tensor_parallel_workers": tp,
"max_model_len": 40960,
"batch_size": total_batch,
"prefill_chunk_size": 0,
"kv_cache_size": 0,
"is_prefill": True,
"attention_backend": "FLASH_ATTN",
"is_mixed_batch": False,
"mode": "true_mixed_fused_projected",
"seq_lens": "",
"total_tokens": sum(prefill_queries) + len(decode),
"max_seq_len": "",
"min_seq_len": "",
"avg_seq_len": "",
"equal_seq_len": "",
"seq_len_variance": "",
"seq_len_std": "",
"seq_len_cv": "",
"is_chunked_prefill_sample": False,
"chunk_start_token": "",
"chunk_end_token": "",
"total_prefill_tokens": sum(prefill_queries),
"profiling_precision": "BF16",
"model_arch": "generic",
"quant_signature": "none",
"measurement_type": "CUDA_EVENT",
"is_true_mixed_batch": True,
"prefill_seq_lens": json.dumps(prefill_queries),
"prefill_kv_cache_sizes": json.dumps(prefill_contexts),
"decode_kv_cache_sizes": json.dumps(decode_kv_lengths),
"num_prefill_seqs": len(prefill),
"num_decode_seqs": len(decode),
"decode_batch_size": len(decode),
"total_batch_size": total_batch,
"total_decode_tokens": len(decode),
"decode_avg_kv_cache_size": (
sum(decode_kv_lengths) / len(decode_kv_lengths)
),
"batch_composition_ratio": len(prefill) / total_batch,
"batch_spec": spec,
"projection_policy": (
"fused_total_conserving_projection_by_same_tp_pure_"
"prefill_decode_reference_ratio"
),
}
)
rows.append(row)
mixed_rows.append(
{
"num_tensor_parallel_workers": tp,
"batch_spec": spec,
"num_prefill_seqs": len(prefill),
"num_decode_seqs": len(decode),
"total_prefill_tokens": sum(q for q, _ in prefill),
"total_decode_tokens": len(decode),
"decode_avg_kv_cache_size": sum(kv for _, kv in decode)
/ len(decode),
"attention_core_mean_ms": core["mean"],
"attention_core_mean_as_median_ms": core["median"],
"kv_cache_update_median_ms": kv_stats["median"],
"pure_prefill_reference_mean_ms": prefill_reference_ms,
"pure_decode_reference_mean_ms": decode_reference_ms,
"projected_prefill_mean_ms": projected_prefill["mean"],
"projected_decode_mean_ms": projected_decode["mean"],
"projection_sum_error_ms": (
projected_prefill["mean"]
+ projected_decode["mean"]
- core["mean"]
),
"representation": (
"one_fused_FA3_call_projected_for_Frontier_with_"
"total_conservation"
),
}
)
continue
is_decode = bool(decode)
queries = [q for q, _ in requests]
contexts = [kv if is_decode else kv - q for q, kv in requests]
avg_query = sum(queries) / len(queries)
variance = sum((query - avg_query) ** 2 for query 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_stats))
row.update(
stat_columns("attn_decode" if is_decode else "attn_prefill", core)
)
row.update(
{
"n_embd": 2048,
"n_q_head": 32,
"n_kv_head": 4,
"block_size": 16,
"num_tensor_parallel_workers": tp,
"max_model_len": 40960,
"batch_size": len(requests),
"prefill_chunk_size": 0 if is_decode else sum(queries),
"kv_cache_size": avg_context,
"is_prefill": not is_decode,
"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": (not is_decode and avg_context > 0),
"chunk_start_token": avg_context if not is_decode else 0,
"chunk_end_token": avg_context + sum(queries) if not is_decode else 0,
"total_prefill_tokens": 0 if is_decode else 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_composition_ratio": "",
"batch_spec": spec,
"projection_policy": (
"measured_FA3_core_plus_measured_KV;reshape_assumed_zero;"
"mean_as_median"
),
}
)
rows.append(row)
if seen_tps != {1, 2, 4}:
raise ValueError(f"attention TP coverage mismatch: {seen_tps}")
attention_fields = [
f"time_stats.{op}.{stat}" for op in ATTENTION_OPS for stat in STAT_NAMES
] + list(ATTENTION_METADATA)
write_csv(output / "attention.csv", attention_fields, rows)
mixed_fields = [
"num_tensor_parallel_workers",
"batch_spec",
"num_prefill_seqs",
"num_decode_seqs",
"total_prefill_tokens",
"total_decode_tokens",
"decode_avg_kv_cache_size",
"attention_core_mean_ms",
"attention_core_mean_as_median_ms",
"kv_cache_update_median_ms",
"pure_prefill_reference_mean_ms",
"pure_decode_reference_mean_ms",
"projected_prefill_mean_ms",
"projected_decode_mean_ms",
"projection_sum_error_ms",
"representation",
]
write_csv(output / "attention_true_mixed_fused.csv", mixed_fields, mixed_rows)
return len(rows), len(mixed_rows), sorted(seen_tps)
def load_features(counts: list[int]) -> dict[str, float]:
total = sum(counts)
count = len(counts)
mean = total / count
variance = sum((value - mean) ** 2 for value in counts) / count
probabilities = [value / total for value in counts if value > 0]
entropy = -sum(probability * math.log2(probability) for probability in probabilities)
sorted_counts = sorted(counts)
gini = (
2 * sum((index + 1) * value for index, value in enumerate(sorted_counts))
/ (count * total)
- (count + 1) / count
)
return {
"total_routed_tokens": total,
"num_experts_per_device": count,
"hidden_dim": 2048,
"expert_hidden_dim": 768,
"router_topk": 8,
"model_expansion_ratio": 768 / 2048,
"tokens_per_expert_avg": mean,
"tokens_to_experts_ratio": mean,
"expert_utilization": sum(value > 0 for value in counts) / count,
"min_load_ratio": min(counts) / mean,
"load_imbalance_cv": math.sqrt(variance) / mean,
"max_load_ratio": max(counts) / mean,
"load_entropy": entropy,
"load_gini_coefficient": gini,
}
def freeze_moe(moe_path: Path, router_path: Path, output: Path) -> int:
moe = load_json(moe_path)
router = load_json(router_path)
if moe.get("schema_version") != "qwen30_vllm020_moe_raw.v1":
raise ValueError(f"unexpected MoE schema in {moe_path}")
if router.get("schema_version") != "qwen30_vllm020_router_raw.v1":
raise ValueError(f"unexpected router schema in {router_path}")
router_by_tokens = {int(row["num_tokens"]): row for row in router["rows"]}
rows: list[dict[str, Any]] = []
seen_pairs: set[tuple[int, int, str]] = set()
for raw in moe["rows"]:
tp = int(raw["tensor_parallel_size"])
num_tokens = int(raw["num_tokens"])
routing_mode = str(raw["routing_mode"])
key = (tp, num_tokens, routing_mode)
if key in seen_pairs:
raise ValueError(f"duplicate MoE row: {key}")
seen_pairs.add(key)
router_row = router_by_tokens[num_tokens]
counts = [int(value) for value in raw["routing_load"]["counts"]]
if sum(counts) != num_tokens * 8 or len(counts) != 128:
raise ValueError(f"invalid routing counts for {key}")
row: dict[str, Any] = {}
row.update(stat_columns("moe_gating_linear", router_row["gate_linear_time_ms"]))
row.update(
stat_columns(
"moe_gating_routing_topk", router_row["routing_topk_time_ms"]
)
)
row.update(zero_stat_columns("moe_shuffling"))
row.update(stat_columns("moe_grouped_gemm", raw["time_ms"]))
row.update(load_features(counts))
row.update(
{
"num_tokens": num_tokens,
"num_experts": 128,
"expert_parallel_size": 1,
"routing_runtime_path": "standard_fused_topk",
"routing_assignment_policy": (
"logit_topk"
if routing_mode == "uniform_random_logits"
else "fixed_hotset8"
),
"routing_weight_policy": "softmax_renorm",
"routing_uses_router_logits": routing_mode == "uniform_random_logits",
"gating_runtime_context": "standalone_legacy",
"gating_runtime_context_impl": "vllm020_replicated_linear",
"use_gated": True,
"num_tensor_parallel_workers": tp,
"load_distribution": routing_mode,
"seed": 20260716,
"moe_grouped_gemm_backend": raw["backend"],
"measurement_type": "CUDA_EVENT",
"profiling_precision": "BF16",
"model_arch": "generic",
"quant_signature": "none",
"router_median_nonadditivity_ratio": router_row[
"median_nonadditivity_ratio"
],
"projection_policy": (
"measured_gate+topk+modular_expert;shuffling_zero_because_"
"expert_measurement_includes_prepare_finalize"
),
}
)
rows.append(row)
expected = 3 * 12 * 2
if len(rows) != expected:
raise ValueError(f"expected {expected} MoE rows, got {len(rows)}")
moe_fields = [
f"time_stats.{op}.{stat}" for op in MOE_OPS for stat in STAT_NAMES
] + list(MOE_METADATA)
write_csv(output / "moe.csv", moe_fields, rows)
return len(rows)
def freeze_allreduce(inputs: list[Path], output: Path) -> int:
rows: list[dict[str, Any]] = []
environments: list[dict[str, Any]] = []
for path in inputs:
payload = load_json(path)
if payload.get("schema_version") != "qwen30_vllm020_allreduce_raw.v1":
raise ValueError(f"unexpected all-reduce schema in {path}")
rows.extend(payload["rows"])
environments.append(payload["environment"])
if {(row["tensor_parallel_size"], row["num_tokens"]) for row in rows} != {
(tp, tokens)
for tp in (2, 4)
for tokens in (1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192)
}:
raise ValueError("all-reduce TP/token coverage mismatch")
(output / "allreduce.json").write_text(
json.dumps(
{
"schema_version": "qwen30_vllm020_allreduce_frozen.v1",
"environment": environments,
"rows": sorted(
rows,
key=lambda row: (
row["tensor_parallel_size"],
row["num_tokens"],
),
),
"frontier_consumption": (
"diagnostic_only_in_base_profile_only_run; measured lookup "
"requires a separate CC-backend injection ablation"
),
},
indent=2,
sort_keys=True,
)
+ "\n"
)
return len(rows)
def main() -> None:
args = parse_args()
all_inputs = [args.linear, *args.attention, args.moe, args.router, *args.allreduce]
for path in all_inputs:
if not path.is_file():
raise SystemExit(f"missing input: {path}")
args.output.mkdir(parents=True, exist_ok=False)
linear_output = args.output / "linear_op.csv"
shutil.copyfile(args.linear, linear_output)
with linear_output.open(newline="") as handle:
linear_rows = list(csv.DictReader(handle))
if len(linear_rows) != 36:
raise ValueError(f"expected 36 linear rows, got {len(linear_rows)}")
attention_rows, mixed_rows, attention_tps = freeze_attention(
list(args.attention), args.output
)
moe_rows = freeze_moe(args.moe, args.router, args.output)
allreduce_rows = freeze_allreduce(list(args.allreduce), args.output)
output_files = [
linear_output,
args.output / "attention.csv",
args.output / "attention_true_mixed_fused.csv",
args.output / "moe.csv",
args.output / "allreduce.json",
]
manifest = {
"schema_version": "frontier_qwen30_vllm020_frozen_profile.v2",
"profile_id": (
"qwen3-30b-a3b-bf16-vllm020-h20-tp1-2-4-"
"fused-mixed-total-conserving"
),
"environment_contract": {
"hardware": "NVIDIA H20",
"model": "Qwen3-30B-A3B",
"dtype": "bfloat16",
"vllm_version": "0.20.0",
"vllm_source_commit": "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1",
"frontier_commit": "d9cfeb6d8791fbf2f295dd9744c56a666171776e",
"tensor_parallel_sizes": [1, 2, 4],
},
"row_counts": {
"linear": len(linear_rows),
"attention_frontier_compatible": attention_rows,
"attention_true_mixed_fused_diagnostic": mixed_rows,
"moe": moe_rows,
"allreduce": allreduce_rows,
},
"attention_tp_coverage": attention_tps,
"projection_contract": {
"linear": "Frontier profiler using vLLM 0.20 CUDA operators",
"attention": (
"Pure prefill/extend/decode FA3 core plus separately measured KV update; "
"input/output reshape assumed zero; exported mean is used as median target; "
"true mixed rows use a total-conserving compatibility projection"
),
"attention_true_mixed": (
"The directly measured fused total is preserved in diagnostics. Frontier's "
"two targets are projected by the same-TP pure prefill/decode reference "
"ratio, with projected prefill + decode exactly equal to the fused total; "
"the split is a schema compatibility attribution, not an observation"
),
"moe": (
"Replicated gate and fused top-k plus TP-local modular expert kernel; "
"expert measurement already includes prepare/finalize so shuffling is zero"
),
"allreduce": (
"Frozen exact runtime measurements; base profile-only comparison keeps the "
"historical Frontier CC backend fixed to isolate compute profile fidelity"
),
},
"inputs": {str(path.resolve()): sha256(path) for path in all_inputs},
"outputs": {path.name: sha256(path) for path in output_files},
}
manifest_path = args.output / "manifest.json"
manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
print(json.dumps(manifest["row_counts"], sort_keys=True))
if __name__ == "__main__":
main()