273 lines
9.7 KiB
Python
273 lines
9.7 KiB
Python
#!/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()
|