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