#!/usr/bin/env python3 """Profile Qwen3's replicated MoE gate and fused top-k in vLLM 0.20.""" from __future__ import annotations import argparse import json import statistics import subprocess from pathlib import Path from typing import Any, Callable import torch import vllm VLLM_VERSION = "0.20.0" VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1" HIDDEN_DIM = 2048 NUM_EXPERTS = 128 TOP_K = 8 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("--output", type=Path, required=True) parser.add_argument("--num-tokens", type=int, nargs="+", required=True) parser.add_argument("--warmup-iters", type=int, default=5) parser.add_argument("--repeats", type=int, default=20) parser.add_argument("--device", default="cuda:0") 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 stats_ms(samples: list[float]) -> dict[str, float]: return { "min": min(samples), "max": max(samples), "mean": statistics.fmean(samples), "median": statistics.median(samples), "std": statistics.pstdev(samples), } def measure_ms( fn: Callable[[], Any], warmup_iters: int, repeats: int ) -> tuple[Any, dict[str, float]]: result = None for _ in range(warmup_iters): result = fn() torch.accelerator.synchronize() samples: list[float] = [] for _ in range(repeats): start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() result = fn() end.record() torch.accelerator.synchronize() samples.append(float(start.elapsed_time(end))) return result, stats_ms(samples) 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}") raw_model_config = json.loads(args.model.joinpath("config.json").read_text()) observed = { "hidden_size": raw_model_config.get("hidden_size"), "num_experts": raw_model_config.get("num_experts"), "num_experts_per_tok": raw_model_config.get("num_experts_per_tok"), "norm_topk_prob": raw_model_config.get("norm_topk_prob"), } expected = { "hidden_size": HIDDEN_DIM, "num_experts": NUM_EXPERTS, "num_experts_per_tok": TOP_K, "norm_topk_prob": True, } if observed != expected: raise SystemExit(f"model contract mismatch: expected {expected}, got {observed}") from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.model_executor.layers.fused_moe import fused_topk from vllm.model_executor.layers.linear import ReplicatedLinear device = torch.device(args.device) torch.accelerator.set_device_index(device) torch.manual_seed(20260716) model_config = ModelConfig( model=str(args.model), dtype="bfloat16", max_model_len=8192, skip_tokenizer_init=True, generation_config="vllm", ) rows: list[dict[str, Any]] = [] with set_current_vllm_config(VllmConfig(model_config=model_config)): gate = ReplicatedLinear( HIDDEN_DIM, NUM_EXPERTS, bias=False, quant_config=None, prefix="model.layers.0.mlp.gate", ).to(device=device, dtype=torch.bfloat16) gate.weight.data.uniform_(-0.01, 0.01) for num_tokens in args.num_tokens: hidden = torch.empty( (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16 ).uniform_(-0.1, 0.1) logits, gate_time = measure_ms( lambda: gate(hidden)[0], args.warmup_iters, args.repeats ) topk_result, topk_time = measure_ms( lambda: fused_topk(hidden, logits, TOP_K, renormalize=True), args.warmup_iters, args.repeats, ) def gate_and_topk() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: current_logits, _ = gate(hidden) return fused_topk(hidden, current_logits, TOP_K, renormalize=True) combined_result, combined_time = measure_ms( gate_and_topk, args.warmup_iters, args.repeats ) topk_weights, topk_ids, _ = topk_result combined_weights, combined_ids, _ = combined_result if logits.shape != (num_tokens, NUM_EXPERTS): raise SystemExit(f"invalid gate output shape: {tuple(logits.shape)}") if topk_ids.shape != (num_tokens, TOP_K): raise SystemExit(f"invalid top-k shape: {tuple(topk_ids.shape)}") torch.testing.assert_close( topk_weights.sum(dim=-1), torch.ones(num_tokens, device=device), atol=1e-5, rtol=1e-5, ) torch.testing.assert_close(combined_weights, topk_weights) torch.testing.assert_close(combined_ids, topk_ids) additive_median = gate_time["median"] + topk_time["median"] row = { "num_tokens": num_tokens, "gate_linear_time_ms": gate_time, "routing_topk_time_ms": topk_time, "gate_plus_topk_time_ms": combined_time, "median_nonadditivity_ratio": ( combined_time["median"] / additive_median if additive_median > 0 else 1.0 ), } rows.append(row) print(json.dumps(row, sort_keys=True), flush=True) payload = { "schema_version": "qwen30_vllm020_router_raw.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(device), "model": str(args.model), "dtype": "bfloat16", "gate_replication": "replicated_across_tp", "top_k": TOP_K, "norm_topk_prob": True, }, "measurement_scope": ( "vLLM ReplicatedLinear gate and fused_topk; measured separately and " "as the actual sequential router path" ), "rows": rows, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") if __name__ == "__main__": main()