Freeze vLLM 0.20 profiles and capture trace routing
This commit is contained in:
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()
|
||||
Reference in New Issue
Block a user