Freeze vLLM 0.20 profiles and capture trace routing
This commit is contained in:
@@ -0,0 +1,108 @@
|
||||
{
|
||||
"contains_prompt_text": false,
|
||||
"fixture_path": "/tmp/qwen30-routing-fixture-20260716.jsonl",
|
||||
"fixture_sha256": "e9e7f5b4e0d3a788dcd99d432f939d9e36bff2a64e412649b407b0609f0e39bb",
|
||||
"prefix_pairs": [
|
||||
{
|
||||
"child_row_id": "17366",
|
||||
"child_turn": 13,
|
||||
"parent_row_id": "16914",
|
||||
"parent_turn": 12,
|
||||
"trace_hash_common_prefix_blocks": 59
|
||||
}
|
||||
],
|
||||
"request_count": 8,
|
||||
"rows": [
|
||||
{
|
||||
"chat_id": "184516",
|
||||
"fixture_index": 0,
|
||||
"input_length": 3791,
|
||||
"output_length": 73,
|
||||
"parent_chat_id": "183921",
|
||||
"prompt_sha256": "6b59c3ecb9a7d8bd3178f10b847a7c16c221fef113df3e80ede55543bc66ea6b",
|
||||
"row_id": "16914",
|
||||
"trace_hash_blocks": 60,
|
||||
"turn": 12
|
||||
},
|
||||
{
|
||||
"chat_id": "183512",
|
||||
"fixture_index": 1,
|
||||
"input_length": 264,
|
||||
"output_length": 975,
|
||||
"parent_chat_id": "-1",
|
||||
"prompt_sha256": "0af4fe018ad7afc152951995e5cb18f389bacf4076239f5a72165a9c84dd566d",
|
||||
"row_id": "15910",
|
||||
"trace_hash_blocks": 5,
|
||||
"turn": 1
|
||||
},
|
||||
{
|
||||
"chat_id": "189951",
|
||||
"fixture_index": 2,
|
||||
"input_length": 488,
|
||||
"output_length": 863,
|
||||
"parent_chat_id": "-1",
|
||||
"prompt_sha256": "b84c0580e9cd80e8dd388f176bd51119f06ee4686838d957b4018b8a1feccb15",
|
||||
"row_id": "22349",
|
||||
"trace_hash_blocks": 8,
|
||||
"turn": 1
|
||||
},
|
||||
{
|
||||
"chat_id": "177472",
|
||||
"fixture_index": 3,
|
||||
"input_length": 1037,
|
||||
"output_length": 895,
|
||||
"parent_chat_id": "-1",
|
||||
"prompt_sha256": "146dc187af0eeb342dd7bd6ebd9453973074209ab4da5b122718b5d9e06d46d1",
|
||||
"row_id": "9870",
|
||||
"trace_hash_blocks": 17,
|
||||
"turn": 1
|
||||
},
|
||||
{
|
||||
"chat_id": "177528",
|
||||
"fixture_index": 4,
|
||||
"input_length": 1993,
|
||||
"output_length": 654,
|
||||
"parent_chat_id": "-1",
|
||||
"prompt_sha256": "c92ba8acd0d637b796d60b5b01e3ac54bde70a986e83809b434c46f50e5242cf",
|
||||
"row_id": "9926",
|
||||
"trace_hash_blocks": 32,
|
||||
"turn": 1
|
||||
},
|
||||
{
|
||||
"chat_id": "193539",
|
||||
"fixture_index": 5,
|
||||
"input_length": 4088,
|
||||
"output_length": 1842,
|
||||
"parent_chat_id": "-1",
|
||||
"prompt_sha256": "38eecbf8766bd8432ac41d6a061a74244c4b3f8b20e8dfee7237b5e6c0e9e13e",
|
||||
"row_id": "25937",
|
||||
"trace_hash_blocks": 64,
|
||||
"turn": 1
|
||||
},
|
||||
{
|
||||
"chat_id": "177590",
|
||||
"fixture_index": 6,
|
||||
"input_length": 7995,
|
||||
"output_length": 1128,
|
||||
"parent_chat_id": "-1",
|
||||
"prompt_sha256": "4bee5b9d1aceec7010a80b8585fb5071b6d468c751e5903abbcd656a8285fcd8",
|
||||
"row_id": "9988",
|
||||
"trace_hash_blocks": 125,
|
||||
"turn": 1
|
||||
},
|
||||
{
|
||||
"chat_id": "184968",
|
||||
"fixture_index": 7,
|
||||
"input_length": 4017,
|
||||
"output_length": 72,
|
||||
"parent_chat_id": "184516",
|
||||
"prompt_sha256": "6885aff07780bff5906669057d46935d78148fb7f4edd0bf58d44c2ffec76952",
|
||||
"row_id": "17366",
|
||||
"trace_hash_blocks": 63,
|
||||
"turn": 13
|
||||
}
|
||||
],
|
||||
"schema_version": "qwen30_trace_routing_fixture.v1",
|
||||
"source_trace": "/home/gahow/phd/aituner/trace_windows/traces/chat_w20260311_1000.jsonl",
|
||||
"source_trace_sha256": "f539f38eb0ee0f750e3c23ff47df6eed3faf723a25f1444d55665a85871750b9"
|
||||
}
|
||||
272
runs/frontier-qwen30-vllm020-profile-v1/capture_trace_routing.py
Normal file
272
runs/frontier-qwen30-vllm020-profile-v1/capture_trace_routing.py
Normal file
@@ -0,0 +1,272 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Capture exact Qwen3 routed-expert IDs from vLLM 0.20 on trace prompts."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import vllm
|
||||
|
||||
|
||||
VLLM_VERSION = "0.20.0"
|
||||
VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
|
||||
NUM_EXPERTS = 128
|
||||
TOP_K = 8
|
||||
NUM_LAYERS = 48
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--vllm-source", type=Path, required=True)
|
||||
parser.add_argument("--model", type=Path, required=True)
|
||||
parser.add_argument("--fixture", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--routes", type=Path, required=True)
|
||||
parser.add_argument("--decode-override", type=int)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def git_head(repo: Path) -> str:
|
||||
return subprocess.check_output(
|
||||
["git", "-C", str(repo), "rev-parse", "HEAD"], text=True
|
||||
).strip()
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
return hashlib.sha256(path.read_bytes()).hexdigest()
|
||||
|
||||
|
||||
def common_prefix(left: list[int], right: list[int]) -> int:
|
||||
count = 0
|
||||
for lhs, rhs in zip(left, right):
|
||||
if lhs != rhs:
|
||||
break
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def distribution(counts: np.ndarray) -> dict[str, Any]:
|
||||
values = counts.astype(np.float64)
|
||||
total = float(values.sum())
|
||||
mean = float(values.mean())
|
||||
probabilities = values[values > 0] / total
|
||||
entropy = float(-(probabilities * np.log2(probabilities)).sum())
|
||||
variance = float(((values - mean) ** 2).mean())
|
||||
ordered = np.sort(values)
|
||||
gini = float(
|
||||
2.0 * np.dot(np.arange(1, len(values) + 1), ordered)
|
||||
/ (len(values) * total)
|
||||
- (len(values) + 1) / len(values)
|
||||
)
|
||||
hottest = np.argsort(values)[-8:][::-1]
|
||||
return {
|
||||
"total_routed_tokens": int(total),
|
||||
"tokens_per_expert_mean": mean,
|
||||
"load_cv": math.sqrt(variance) / mean,
|
||||
"load_gini": gini,
|
||||
"load_entropy_bits": entropy,
|
||||
"min_load_ratio": float(values.min() / mean),
|
||||
"max_load_ratio": float(values.max() / mean),
|
||||
"expert_utilization": float(np.count_nonzero(values) / len(values)),
|
||||
"hottest_experts": [int(value) for value in hottest],
|
||||
"hottest_counts": [int(values[value]) for value in hottest],
|
||||
"counts": counts.astype(int).tolist(),
|
||||
}
|
||||
|
||||
|
||||
def phase_summary(routes: list[np.ndarray]) -> dict[str, Any]:
|
||||
counts = np.zeros(NUM_EXPERTS, dtype=np.int64)
|
||||
per_layer = np.zeros((NUM_LAYERS, NUM_EXPERTS), dtype=np.int64)
|
||||
token_count = 0
|
||||
for route in routes:
|
||||
token_count += route.shape[0]
|
||||
counts += np.bincount(route.reshape(-1), minlength=NUM_EXPERTS)
|
||||
for layer in range(NUM_LAYERS):
|
||||
per_layer[layer] += np.bincount(
|
||||
route[:, layer, :].reshape(-1), minlength=NUM_EXPERTS
|
||||
)
|
||||
return {
|
||||
"token_count": token_count,
|
||||
"all_layers": distribution(counts),
|
||||
"per_layer": [distribution(row) for row in per_layer],
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if vllm.__version__ != VLLM_VERSION:
|
||||
raise SystemExit(f"expected vLLM {VLLM_VERSION}, got {vllm.__version__}")
|
||||
source_head = git_head(args.vllm_source)
|
||||
if source_head != VLLM_COMMIT:
|
||||
raise SystemExit(f"expected vLLM source {VLLM_COMMIT}, got {source_head}")
|
||||
rows = [json.loads(line) for line in args.fixture.read_text().splitlines() if line]
|
||||
if not rows:
|
||||
raise SystemExit("empty routing fixture")
|
||||
requested_decode = [
|
||||
args.decode_override
|
||||
if args.decode_override is not None
|
||||
else int(row["output_length"])
|
||||
for row in rows
|
||||
]
|
||||
if any(value <= 0 for value in requested_decode):
|
||||
raise SystemExit("all requested decode lengths must be positive")
|
||||
|
||||
from vllm import LLM, SamplingParams
|
||||
|
||||
llm = LLM(
|
||||
model=str(args.model),
|
||||
dtype="bfloat16",
|
||||
tensor_parallel_size=1,
|
||||
max_model_len=16384,
|
||||
max_num_batched_tokens=8192,
|
||||
max_num_seqs=64,
|
||||
gpu_memory_utilization=0.90,
|
||||
enable_chunked_prefill=True,
|
||||
enable_prefix_caching=True,
|
||||
enable_return_routed_experts=True,
|
||||
attention_backend="FLASH_ATTN",
|
||||
disable_log_stats=False,
|
||||
)
|
||||
sampling = [
|
||||
SamplingParams(temperature=0, min_tokens=value, max_tokens=value)
|
||||
for value in requested_decode
|
||||
]
|
||||
conversations = [
|
||||
[{"role": "user", "content": row["prompt"]}] for row in rows
|
||||
]
|
||||
outputs = llm.chat(conversations, sampling_params=sampling, use_tqdm=False)
|
||||
if len(outputs) != len(rows):
|
||||
raise SystemExit(f"expected {len(rows)} outputs, got {len(outputs)}")
|
||||
|
||||
prompt_tokens_by_chat: dict[str, list[int]] = {}
|
||||
prefill_routes: list[np.ndarray] = []
|
||||
decode_routes: list[np.ndarray] = []
|
||||
raw_routes: dict[str, np.ndarray] = {}
|
||||
request_summaries = []
|
||||
for row, output, decode_tokens in zip(rows, outputs, requested_decode):
|
||||
completion = output.outputs[0]
|
||||
routed = completion.routed_experts
|
||||
if routed is None:
|
||||
raise SystemExit(f"row {row['row_id']} returned no routed experts")
|
||||
routed = np.asarray(routed)
|
||||
prompt_tokens = list(output.prompt_token_ids)
|
||||
generated_tokens = list(completion.token_ids)
|
||||
expected = len(prompt_tokens) + len(generated_tokens) - 1
|
||||
if routed.shape != (expected, NUM_LAYERS, TOP_K):
|
||||
raise SystemExit(
|
||||
f"row {row['row_id']} routes {routed.shape}, expected "
|
||||
f"{(expected, NUM_LAYERS, TOP_K)}"
|
||||
)
|
||||
if routed.min() < 0 or routed.max() >= NUM_EXPERTS:
|
||||
raise SystemExit(f"row {row['row_id']} returned invalid expert IDs")
|
||||
prefill = routed[: len(prompt_tokens)]
|
||||
decode = routed[len(prompt_tokens) :]
|
||||
if decode.shape[0] != decode_tokens - 1:
|
||||
raise SystemExit(f"row {row['row_id']} decode route length mismatch")
|
||||
prefill_routes.append(prefill)
|
||||
decode_routes.append(decode)
|
||||
raw_routes[f"row_{row['row_id']}"] = routed.astype(np.int16)
|
||||
prompt_tokens_by_chat[str(row["chat_id"])] = prompt_tokens
|
||||
request_summaries.append(
|
||||
{
|
||||
"fixture_index": row["fixture_index"],
|
||||
"row_id": row["row_id"],
|
||||
"turn": row["turn"],
|
||||
"input_length_trace": row["input_length"],
|
||||
"prompt_tokens_vllm": len(prompt_tokens),
|
||||
"chat_wrapper_delta": len(prompt_tokens) - int(row["input_length"]),
|
||||
"generated_tokens": len(generated_tokens),
|
||||
"requested_decode_tokens": decode_tokens,
|
||||
"routed_shape": list(routed.shape),
|
||||
"prompt_sha256": row["prompt_sha256"],
|
||||
"trace_hash_blocks": len(row["hash_ids"]),
|
||||
}
|
||||
)
|
||||
|
||||
prefix_pairs = []
|
||||
by_chat = {str(row["chat_id"]): row for row in rows}
|
||||
for child in rows:
|
||||
parent = by_chat.get(str(child["parent_chat_id"]))
|
||||
if parent is None:
|
||||
continue
|
||||
parent_tokens = prompt_tokens_by_chat[str(parent["chat_id"])]
|
||||
child_tokens = prompt_tokens_by_chat[str(child["chat_id"])]
|
||||
prefix_pairs.append(
|
||||
{
|
||||
"parent_row_id": parent["row_id"],
|
||||
"child_row_id": child["row_id"],
|
||||
"trace_hash_common_prefix_blocks": common_prefix(
|
||||
parent["hash_ids"], child["hash_ids"]
|
||||
),
|
||||
"vllm_token_common_prefix": common_prefix(parent_tokens, child_tokens),
|
||||
"vllm_full_common_blocks_16": common_prefix(
|
||||
parent_tokens, child_tokens
|
||||
)
|
||||
// 16,
|
||||
}
|
||||
)
|
||||
|
||||
args.routes.parent.mkdir(parents=True, exist_ok=True)
|
||||
np.savez_compressed(args.routes, **raw_routes)
|
||||
payload = {
|
||||
"schema_version": "qwen30_vllm020_trace_routing.v1",
|
||||
"environment": {
|
||||
"vllm_version": vllm.__version__,
|
||||
"vllm_source_commit": source_head,
|
||||
"torch_version": torch.__version__,
|
||||
"torch_cuda": torch.version.cuda,
|
||||
"gpu": torch.cuda.get_device_name(0),
|
||||
"model": str(args.model),
|
||||
"dtype": "bfloat16",
|
||||
"tensor_parallel_size": 1,
|
||||
"max_num_batched_tokens": 8192,
|
||||
"max_num_seqs": 64,
|
||||
"prefix_caching": True,
|
||||
"chunked_prefill": True,
|
||||
"attention_backend": "FLASH_ATTN",
|
||||
},
|
||||
"capture_contract": {
|
||||
"api": "LLM.chat",
|
||||
"enable_return_routed_experts": True,
|
||||
"route_shape": "[prompt_tokens + generated_tokens - 1, layers, topk]",
|
||||
"decode_policy": (
|
||||
f"fixed_override_{args.decode_override}"
|
||||
if args.decode_override is not None
|
||||
else "exact_trace_output_length"
|
||||
),
|
||||
"contains_prompt_text": False,
|
||||
"fixture_sha256": sha256(args.fixture),
|
||||
"routes_npz": str(args.routes),
|
||||
},
|
||||
"requests": request_summaries,
|
||||
"prefix_pairs": prefix_pairs,
|
||||
"phases": {
|
||||
"prefill": phase_summary(prefill_routes),
|
||||
"decode": phase_summary(decode_routes),
|
||||
},
|
||||
}
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"requests": len(rows),
|
||||
"prefill_tokens": payload["phases"]["prefill"]["token_count"],
|
||||
"decode_tokens": payload["phases"]["decode"]["token_count"],
|
||||
"prefix_pairs": prefix_pairs,
|
||||
},
|
||||
sort_keys=True,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,162 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Extract a prompt-bearing routing fixture while emitting a prompt-free manifest."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--trace", type=Path, required=True)
|
||||
parser.add_argument("--row-ids", type=int, nargs="+", required=True)
|
||||
parser.add_argument("--parent-of", type=int, nargs="*", default=[])
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--manifest", type=Path, required=True)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def common_prefix(left: list[Any], right: list[Any]) -> int:
|
||||
count = 0
|
||||
for lhs, rhs in zip(left, right):
|
||||
if lhs != rhs:
|
||||
break
|
||||
count += 1
|
||||
return count
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
target_ids = set(args.row_ids)
|
||||
parent_targets = set(args.parent_of)
|
||||
if not parent_targets.issubset(target_ids):
|
||||
raise SystemExit("--parent-of must be a subset of --row-ids")
|
||||
|
||||
trace_digest = hashlib.sha256()
|
||||
by_row: dict[int, dict[str, Any]] = {}
|
||||
by_chat_id: dict[str, dict[str, Any]] = {}
|
||||
with args.trace.open("rb") as handle:
|
||||
row_id = 0
|
||||
while True:
|
||||
offset = handle.tell()
|
||||
line = handle.readline()
|
||||
if not line:
|
||||
break
|
||||
trace_digest.update(line)
|
||||
row = json.loads(line)
|
||||
meta = {
|
||||
"row_id": row_id,
|
||||
"offset": offset,
|
||||
"chat_id": str(row.get("chat_id")),
|
||||
"parent_chat_id": str(row.get("parent_chat_id")),
|
||||
"turn": int(row.get("turn", 1)),
|
||||
"input_length": int(row["input_length"]),
|
||||
"output_length": int(row["output_length"]),
|
||||
"hash_ids": row.get("hash_ids") or [],
|
||||
}
|
||||
by_chat_id[meta["chat_id"]] = meta
|
||||
if row_id in target_ids:
|
||||
by_row[row_id] = meta
|
||||
row_id += 1
|
||||
|
||||
missing = target_ids - set(by_row)
|
||||
if missing:
|
||||
raise SystemExit(f"missing target row IDs: {sorted(missing)}")
|
||||
parent_rows: list[dict[str, Any]] = []
|
||||
for row_id in args.parent_of:
|
||||
parent_id = by_row[row_id]["parent_chat_id"]
|
||||
parent = by_chat_id.get(parent_id)
|
||||
if parent is None:
|
||||
raise SystemExit(f"row {row_id} parent chat {parent_id} is absent")
|
||||
parent_rows.append(parent)
|
||||
|
||||
# Put parents first so online prefix caching can materialize their shared
|
||||
# blocks before descendants are admitted.
|
||||
ordered_meta = parent_rows + [by_row[row_id] for row_id in args.row_ids]
|
||||
if len({row["row_id"] for row in ordered_meta}) != len(ordered_meta):
|
||||
raise SystemExit("fixture rows must be unique")
|
||||
|
||||
output_rows: list[dict[str, Any]] = []
|
||||
with args.trace.open("rb") as handle:
|
||||
for fixture_index, meta in enumerate(ordered_meta):
|
||||
handle.seek(meta["offset"])
|
||||
source = json.loads(handle.readline())
|
||||
prompt = source.get("prompt")
|
||||
if not isinstance(prompt, str) or not prompt:
|
||||
raise SystemExit(f"row {meta['row_id']} has no prompt")
|
||||
output_rows.append(
|
||||
{
|
||||
"fixture_index": fixture_index,
|
||||
"row_id": str(meta["row_id"]),
|
||||
"prompt": prompt,
|
||||
"prompt_sha256": hashlib.sha256(prompt.encode()).hexdigest(),
|
||||
"input_length": meta["input_length"],
|
||||
"output_length": meta["output_length"],
|
||||
"turn": meta["turn"],
|
||||
"chat_id": meta["chat_id"],
|
||||
"parent_chat_id": meta["parent_chat_id"],
|
||||
"hash_ids": meta["hash_ids"],
|
||||
}
|
||||
)
|
||||
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
with args.output.open("w") as handle:
|
||||
for row in output_rows:
|
||||
handle.write(json.dumps(row, sort_keys=True) + "\n")
|
||||
|
||||
pair_coverage = []
|
||||
by_chat = {row["chat_id"]: row for row in output_rows}
|
||||
for child in output_rows:
|
||||
parent = by_chat.get(child["parent_chat_id"])
|
||||
if parent is None:
|
||||
continue
|
||||
pair_coverage.append(
|
||||
{
|
||||
"parent_row_id": parent["row_id"],
|
||||
"child_row_id": child["row_id"],
|
||||
"parent_turn": parent["turn"],
|
||||
"child_turn": child["turn"],
|
||||
"trace_hash_common_prefix_blocks": common_prefix(
|
||||
parent["hash_ids"], child["hash_ids"]
|
||||
),
|
||||
}
|
||||
)
|
||||
fixture_digest = hashlib.sha256(args.output.read_bytes()).hexdigest()
|
||||
manifest = {
|
||||
"schema_version": "qwen30_trace_routing_fixture.v1",
|
||||
"source_trace": str(args.trace.resolve()),
|
||||
"source_trace_sha256": trace_digest.hexdigest(),
|
||||
"fixture_path": str(args.output.resolve()),
|
||||
"fixture_sha256": fixture_digest,
|
||||
"contains_prompt_text": False,
|
||||
"request_count": len(output_rows),
|
||||
"rows": [
|
||||
{
|
||||
key: row[key]
|
||||
for key in (
|
||||
"fixture_index",
|
||||
"row_id",
|
||||
"prompt_sha256",
|
||||
"input_length",
|
||||
"output_length",
|
||||
"turn",
|
||||
"chat_id",
|
||||
"parent_chat_id",
|
||||
)
|
||||
}
|
||||
| {"trace_hash_blocks": len(row["hash_ids"])}
|
||||
for row in output_rows
|
||||
],
|
||||
"prefix_pairs": pair_coverage,
|
||||
}
|
||||
args.manifest.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.manifest.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
|
||||
print(json.dumps({"requests": len(output_rows), "prefix_pairs": pair_coverage}))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,695 @@
|
||||
#!/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()
|
||||
@@ -0,0 +1,19 @@
|
||||
version = 1
|
||||
|
||||
[[jobs]]
|
||||
name = "qwen30-vllm020-trace-routing-20260716-v1"
|
||||
gpus = 1
|
||||
gpu_model = "H20"
|
||||
hosts = ["dash0"]
|
||||
command = "cd /home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-qwen30-vllm020-profile-v1 && timeout --signal=TERM --kill-after=30s 3720 bash run_trace_routing.sh"
|
||||
artifacts = ["artifacts/trace-routing-v1"]
|
||||
|
||||
[jobs.env]
|
||||
HOME = "/tmp/wjh"
|
||||
XDG_CACHE_HOME = "/tmp/wjh/.cache"
|
||||
VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm"
|
||||
OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/trace-routing-v1"
|
||||
FIXTURE = "/tmp/wjh/qwen30-routing-fixture-20260716.jsonl"
|
||||
VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1"
|
||||
VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build"
|
||||
MODEL = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B"
|
||||
@@ -0,0 +1,104 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Prepare old/new profile-only Frontier manifests from the frozen P1 probes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
PROFILE_KEYS = {
|
||||
"linear_op_input_file": "linear_op.csv",
|
||||
"atten_input_file": "attention.csv",
|
||||
"moe_input_file": "moe.csv",
|
||||
}
|
||||
|
||||
|
||||
def sha256(path: Path) -> str:
|
||||
digest = hashlib.sha256()
|
||||
with path.open("rb") as handle:
|
||||
for chunk in iter(lambda: handle.read(1 << 20), b""):
|
||||
digest.update(chunk)
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def write_json(path: Path, payload: Any) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--source", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--mode", choices=("old-profile-only", "new-profile-only"), required=True)
|
||||
parser.add_argument("--profile-root", type=Path)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
source = json.loads(args.source.read_text())
|
||||
if source.get("status") != "PASS" or len(source.get("entries", [])) != 12:
|
||||
raise SystemExit("source manifest must contain 12 passing P1 probes")
|
||||
if args.mode == "new-profile-only" and args.profile_root is None:
|
||||
raise SystemExit("--profile-root is required for new-profile-only")
|
||||
|
||||
output = args.output.resolve()
|
||||
config_root = output / "configs"
|
||||
cache_root = output / "prediction-cache"
|
||||
entries: list[dict[str, Any]] = []
|
||||
profile_hashes: dict[str, str] = {}
|
||||
if args.profile_root is not None:
|
||||
profile_root = args.profile_root.resolve()
|
||||
for filename in PROFILE_KEYS.values():
|
||||
path = profile_root / filename
|
||||
if not path.is_file():
|
||||
raise SystemExit(f"missing frozen profile: {path}")
|
||||
profile_hashes[str(path)] = sha256(path)
|
||||
|
||||
for entry in source["entries"]:
|
||||
config_path = Path(entry["config"])
|
||||
config = json.loads(config_path.read_text())
|
||||
config["mode"] = args.mode
|
||||
config["config_id"] = f"{config['cell_id']}__{args.mode}"
|
||||
config["calibration"]["a_tp"] = 1.0
|
||||
knobs = config["frontier"]["knobs"]
|
||||
knobs["cache_dir"] = str(cache_root)
|
||||
knobs["no_cache"] = False
|
||||
if args.mode == "new-profile-only":
|
||||
for key, filename in PROFILE_KEYS.items():
|
||||
knobs[key] = str((args.profile_root.resolve() / filename))
|
||||
|
||||
target_config = config_root / f"{entry['fixture_id']}.json"
|
||||
write_json(target_config, config)
|
||||
updated_entry = dict(entry)
|
||||
updated_entry["config"] = str(target_config)
|
||||
updated_entry["calibration_scale"] = 1.0
|
||||
entries.append(updated_entry)
|
||||
|
||||
prepared = {
|
||||
"schema": "frontier-qwen30-profile-comparison-prepared.v1",
|
||||
"status": "PASS",
|
||||
"mode": args.mode,
|
||||
"source": {
|
||||
"manifest": str(args.source.resolve()),
|
||||
"sha256": sha256(args.source),
|
||||
},
|
||||
"profile_hashes": profile_hashes,
|
||||
"isolation": {
|
||||
"calibration_a_tp": 1.0,
|
||||
"prediction_cache": str(cache_root),
|
||||
"all_non_profile_knobs_inherited": True,
|
||||
},
|
||||
"entries": entries,
|
||||
}
|
||||
write_json(output / "prepared-manifest.json", prepared)
|
||||
print(output / "prepared-manifest.json")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
50
runs/frontier-qwen30-vllm020-profile-v1/run_trace_routing.sh
Normal file
50
runs/frontier-qwen30-vllm020-profile-v1/run_trace_routing.sh
Normal file
@@ -0,0 +1,50 @@
|
||||
#!/usr/bin/env bash
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
OUTPUT_ROOT="${OUTPUT_ROOT:?OUTPUT_ROOT must be set}"
|
||||
FIXTURE="${FIXTURE:?FIXTURE must be set}"
|
||||
VENV_ROOT="${VENV_ROOT:-/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1}"
|
||||
VLLM_SOURCE="${VLLM_SOURCE:-/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build}"
|
||||
MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}"
|
||||
mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" "${OUTPUT_ROOT}/raw"
|
||||
exec > >(tee -a "${OUTPUT_ROOT}/logs/full.log") 2>&1
|
||||
|
||||
IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES:?fleet GPU is required}"
|
||||
if [[ "${#GPU_IDS[@]}" -ne 1 ]]; then
|
||||
echo "ERROR: expected exactly one GPU, got ${CUDA_VISIBLE_DEVICES}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "ROUTING_LAUNCH_ECHO host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES} model=${MODEL} runtime=vLLM-0.20.0+cu129 trace_fixture=${FIXTURE} fixture_sha256=e9e7f5b4e0d3a788dcd99d432f939d9e36bff2a64e412649b407b0609f0e39bb requests=8 input_tokens_trace=23673 output_tokens_trace=6502 TP=1 MBT=8192 MNS=64 prefix_cache=true chunked_prefill=true dtype=BF16 output=${OUTPUT_ROOT} expected_wall=15-40m hard_wall=3600s hard_gpu_cap=1.0_H20h"
|
||||
date -u +"START_UTC=%Y-%m-%dT%H:%M:%SZ"
|
||||
nvidia-smi --query-gpu=index,name,driver_version,memory.used,utilization.gpu --format=csv,noheader
|
||||
|
||||
test "$(git -C "${VLLM_SOURCE}" rev-parse HEAD)" = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1"
|
||||
test -s "${MODEL}/config.json"
|
||||
echo "e9e7f5b4e0d3a788dcd99d432f939d9e36bff2a64e412649b407b0609f0e39bb ${FIXTURE}" | sha256sum -c -
|
||||
git rev-parse HEAD > "${OUTPUT_ROOT}/provenance/aituner.commit"
|
||||
git -C "${VLLM_SOURCE}" rev-parse HEAD > "${OUTPUT_ROOT}/provenance/vllm-source.commit"
|
||||
sha256sum capture_trace_routing.py run_trace_routing.sh \
|
||||
> "${OUTPUT_ROOT}/provenance/source.sha256"
|
||||
sha256sum "${MODEL}/config.json" > "${OUTPUT_ROOT}/provenance/model-config.sha256"
|
||||
sha256sum "${FIXTURE}" > "${OUTPUT_ROOT}/provenance/fixture.sha256"
|
||||
uv pip freeze --python "${VENV_ROOT}/bin/python" \
|
||||
> "${OUTPUT_ROOT}/provenance/pip-freeze.txt"
|
||||
nvidia-smi --query-gpu=index,uuid,name,driver_version,memory.total \
|
||||
--format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.csv"
|
||||
|
||||
timeout --signal=TERM --kill-after=30s 3300 \
|
||||
"${VENV_ROOT}/bin/python" capture_trace_routing.py \
|
||||
--vllm-source "${VLLM_SOURCE}" \
|
||||
--model "${MODEL}" \
|
||||
--fixture "${FIXTURE}" \
|
||||
--output "${OUTPUT_ROOT}/raw/routing.json" \
|
||||
--routes "${OUTPUT_ROOT}/raw/routes.npz"
|
||||
|
||||
test -s "${OUTPUT_ROOT}/raw/routing.json"
|
||||
test -s "${OUTPUT_ROOT}/raw/routes.npz"
|
||||
sha256sum "${OUTPUT_ROOT}/raw/routing.json" "${OUTPUT_ROOT}/raw/routes.npz" \
|
||||
"${OUTPUT_ROOT}/provenance"/* > "${OUTPUT_ROOT}/artifacts.sha256"
|
||||
date -u +"END_UTC=%Y-%m-%dT%H:%M:%SZ"
|
||||
echo "TRACE_ROUTING_COMPLETE requests=8"
|
||||
Reference in New Issue
Block a user