#!/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()