gahow/aituner
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Align Frontier piecewise graph profiles

Browse Source
This commit is contained in:
Gahow Wang
2026-07-17 23:22:42 +08:00
parent 47355a9411
commit bdc357dc6c
10 changed files with 804 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

132
runs/frontier-qwen30-vllm020-profile-v1/freeze_frontier_profiles.py
View File

@@ -117,8 +117,18 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--attention", type=Path, nargs="+", 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("--allreduce", type=Path, nargs=2)
parser.add_argument("--allreduce-frozen", type=Path)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument(
"--measurement-type",
choices=("CUDA_EVENT", "KERNEL_ONLY"),
default="CUDA_EVENT",
)
parser.add_argument(
"--frontier-commit",
default="d9cfeb6d8791fbf2f295dd9744c56a666171776e",
)
return parser.parse_args()
@@ -196,7 +206,7 @@ def write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, Any]]) ->
def freeze_attention(
inputs: list[Path], output: Path
inputs: list[Path], output: Path, *, measurement_type: str
) -> tuple[int, int, list[str]]:
rows: list[dict[str, Any]] = []
mixed_rows: list[dict[str, Any]] = []
@@ -208,6 +218,13 @@ def freeze_attention(
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}")
expected_method = (
"record_function" if measurement_type == "KERNEL_ONLY" else "cuda_event"
)
if payload["environment"].get("profile_method", "cuda_event") != expected_method:
raise ValueError(
f"attention profile method mismatch in {path}: expected {expected_method}"
)
for raw in payload["rows"]:
if raw.get("error") is not None:
raise ValueError(f"failed attention row in {path}: {raw['error']}")
@@ -321,7 +338,7 @@ def freeze_attention(
"profiling_precision": "BF16",
"model_arch": "generic",
"quant_signature": "none",
"measurement_type": "CUDA_EVENT",
"measurement_type": measurement_type,
"is_true_mixed_batch": True,
"prefill_seq_lens": json.dumps(prefill_queries),
"prefill_kv_cache_sizes": json.dumps(prefill_contexts),
@@ -418,7 +435,7 @@ def freeze_attention(
"profiling_precision": "BF16",
"model_arch": "generic",
"quant_signature": "none",
"measurement_type": "CUDA_EVENT",
"measurement_type": measurement_type,
"is_true_mixed_batch": False,
"prefill_seq_lens": "",
"prefill_kv_cache_sizes": "",
@@ -498,13 +515,23 @@ def load_features(counts: list[int]) -> dict[str, float]:
}
def freeze_moe(moe_path: Path, router_path: Path, output: Path) -> int:
def freeze_moe(
moe_path: Path, router_path: Path, output: Path, *, measurement_type: str
) -> 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}")
expected_method = (
"record_function" if measurement_type == "KERNEL_ONLY" else "cuda_event"
)
for payload, label in ((moe, "moe"), (router, "router")):
if payload["environment"].get("profile_method", "cuda_event") != expected_method:
raise ValueError(
f"{label} profile method mismatch in {payload}: expected {expected_method}"
)
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()
@@ -550,7 +577,7 @@ def freeze_moe(moe_path: Path, router_path: Path, output: Path) -> int:
"load_distribution": routing_mode,
"seed": 20260716,
"moe_grouped_gemm_backend": raw["backend"],
"measurement_type": "CUDA_EVENT",
"measurement_type": measurement_type,
"profiling_precision": "BF16",
"model_arch": "generic",
"quant_signature": "none",
@@ -565,9 +592,17 @@ def freeze_moe(moe_path: Path, router_path: Path, output: Path) -> int:
)
rows.append(row)
expected = 3 * 12 * 2
if len(rows) != expected:
raise ValueError(f"expected {expected} MoE rows, got {len(rows)}")
tokens = {int(row["num_tokens"]) for row in router["rows"]}
modes = {str(row["routing_mode"]) for row in moe["rows"]}
expected = {(tp, tokens_value, mode) for tp in (1, 2, 4) for tokens_value in tokens for mode in modes}
actual = {
(int(row["num_tensor_parallel_workers"]), int(row["num_tokens"]), str(row["load_distribution"]))
for row in rows
}
if actual != expected:
raise ValueError(
f"MoE TP/token/routing coverage mismatch: missing={expected - actual}, extra={actual - expected}"
)
moe_fields = [
f"time_stats.{op}.{stat}" for op in MOE_OPS for stat in STAT_NAMES
] + list(MOE_METADATA)
@@ -617,7 +652,13 @@ def freeze_allreduce(inputs: list[Path], output: Path) -> int:
def main() -> None:
args = parse_args()
all_inputs = [args.linear, *args.attention, args.moe, args.router, *args.allreduce]
if args.allreduce is not None and args.allreduce_frozen is not None:
raise SystemExit("provide either --allreduce or --allreduce-frozen, not both")
all_inputs = [args.linear, *args.attention, args.moe, args.router]
if args.allreduce is not None:
all_inputs.extend(args.allreduce)
if args.allreduce_frozen is not None:
all_inputs.append(args.allreduce_frozen)
for path in all_inputs:
if not path.is_file():
raise SystemExit(f"missing input: {path}")
@@ -627,24 +668,39 @@ def main() -> None:
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)}")
if not linear_rows:
raise ValueError("linear profile has no rows")
if {row.get("measurement_type") for row in linear_rows} != {args.measurement_type}:
raise ValueError(
f"linear measurement family mismatch: expected {args.measurement_type}"
)
attention_rows, mixed_rows, attention_tps = freeze_attention(
list(args.attention), args.output
list(args.attention), args.output, measurement_type=args.measurement_type
)
moe_rows = freeze_moe(args.moe, args.router, args.output)
allreduce_rows = freeze_allreduce(list(args.allreduce), args.output)
moe_rows = freeze_moe(
args.moe, args.router, args.output, measurement_type=args.measurement_type
)
allreduce_rows = 0
allreduce_source = "not_included"
if args.allreduce is not None:
allreduce_rows = freeze_allreduce(list(args.allreduce), args.output)
allreduce_source = "raw_vllm020_measurements"
elif args.allreduce_frozen is not None:
shutil.copyfile(args.allreduce_frozen, args.output / "allreduce.json")
allreduce_rows = len(load_json(args.allreduce_frozen).get("rows", []))
allreduce_source = "carried_forward_frozen_measurements"
output_files = [
linear_output,
args.output / "attention.csv",
args.output / "attention_true_mixed_fused.csv",
args.output / "moe.csv",
args.output / "allreduce.json",
]
if (args.output / "allreduce.json").is_file():
output_files.append(args.output / "allreduce.json")
batch_composition_augmented = len(args.attention) > 3
long_context_augmented = any(
long_context_augmented = args.measurement_type == "KERNEL_ONLY" or any(
"long-context" in path.name for path in args.attention
)
long_context_coverage: dict[str, Any] = {"included": long_context_augmented}
@@ -659,34 +715,45 @@ def main() -> None:
if int(row["num_tensor_parallel_workers"]) == tp
and row["is_prefill"].lower() == "false"
}
mixed_kv = {
int(float(row["decode_avg_kv_cache_size"]))
for row in frozen_attention
if int(row["num_tensor_parallel_workers"]) == tp
and row.get("is_true_mixed_batch", "").lower() == "true"
}
if not {16384, 32768, 40960}.issubset(decode_kv):
raise ValueError(f"long-context decode coverage mismatch for TP{tp}")
if not {16384, 32768}.issubset(mixed_kv):
raise ValueError(f"long-context mixed coverage mismatch for TP{tp}")
required_decode = (
{128, 1024, 2048, 4096, 8192, 16384, 32768, 40960}
if args.measurement_type == "KERNEL_ONLY"
else {16384, 32768, 40960}
)
if not required_decode.issubset(decode_kv):
raise ValueError(f"decode KV coverage mismatch for TP{tp}")
by_tp[str(tp)] = {
"decode_kv_lengths": sorted(decode_kv),
"true_mixed_decode_avg_kv_lengths": sorted(mixed_kv),
}
if args.measurement_type != "KERNEL_ONLY":
mixed_kv = {
int(float(row["decode_avg_kv_cache_size"]))
for row in frozen_attention
if int(row["num_tensor_parallel_workers"]) == tp
and row.get("is_true_mixed_batch", "").lower() == "true"
}
if not {16384, 32768}.issubset(mixed_kv):
raise ValueError(f"long-context mixed coverage mismatch for TP{tp}")
by_tp[str(tp)]["true_mixed_decode_avg_kv_lengths"] = sorted(mixed_kv)
long_context_coverage["by_tp"] = by_tp
manifest = {
"schema_version": (
"frontier_qwen30_vllm020_frozen_profile.v4"
if long_context_augmented
"frontier_qwen30_vllm020_kernel_only_profile.v1"
if args.measurement_type == "KERNEL_ONLY"
else (
"frontier_qwen30_vllm020_frozen_profile.v4"
if long_context_augmented
else (
"frontier_qwen30_vllm020_frozen_profile.v3"
if batch_composition_augmented
else "frontier_qwen30_vllm020_frozen_profile.v2"
)
)
),
"profile_id": (
"qwen3-30b-a3b-bf16-vllm020-h20-tp1-2-4-"
"fused-mixed-total-conserving"
+ ("-kernel-only-record-function" if args.measurement_type == "KERNEL_ONLY" else "")
+ ("-pure-prefill-batch-composition" if batch_composition_augmented else "")
+ ("-long-context-decode-mixed" if long_context_augmented else "")
),
@@ -696,7 +763,7 @@ def main() -> None:
"dtype": "bfloat16",
"vllm_version": "0.20.0",
"vllm_source_commit": "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1",
"frontier_commit": "d9cfeb6d8791fbf2f295dd9744c56a666171776e",
"frontier_commit": args.frontier_commit,
"tensor_parallel_sizes": [1, 2, 4],
},
"row_counts": {
@@ -740,8 +807,7 @@ def main() -> None:
"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"
"Frozen exact runtime measurements; source=" + allreduce_source
),
},
"inputs": {str(path.resolve()): sha256(path) for path in all_inputs},