From fcca259475c309f149ffc75eae859da8db851f81 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Thu, 16 Jul 2026 22:43:55 +0800 Subject: [PATCH] Profile Qwen MoE router on vLLM 0.20 --- .../jobs_frontier_linear_full.toml | 6 +- .../jobs_router_full.toml | 18 ++ .../profile_vllm020_router.py | 195 ++++++++++++++++++ .../run_frontier_linear_full.sh | 1 + .../run_router_full.sh | 49 +++++ 5 files changed, 266 insertions(+), 3 deletions(-) create mode 100644 runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml create mode 100644 runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py create mode 100644 runs/frontier-qwen30-vllm020-profile-v1/run_router_full.sh diff --git a/runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_full.toml b/runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_full.toml index 374470e..c24a55c 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_full.toml +++ b/runs/frontier-qwen30-vllm020-profile-v1/jobs_frontier_linear_full.toml @@ -1,18 +1,18 @@ version = 1 [[jobs]] -name = "qwen30-vllm020-frontier-linear-full-20260716-v1" +name = "qwen30-vllm020-frontier-linear-full-20260716-v2-max-tokens" gpus = 1 gpu_model = "H20" hosts = ["dash0"] command = "cd /home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-qwen30-vllm020-profile-v1 && timeout --signal=TERM --kill-after=30s 1620 bash run_frontier_linear_full.sh" -artifacts = ["artifacts/frontier-linear-full-v1"] +artifacts = ["artifacts/frontier-linear-full-v2"] [jobs.env] HOME = "/tmp/wjh" XDG_CACHE_HOME = "/tmp/wjh/.cache" VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm" -OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/frontier-linear-full-v1" +OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/frontier-linear-full-v2" VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1" FRONTIER_ROOT = "/home/admin/cpfs/wjh/frontier-qwen30-vllm020-profile-v1/Frontier" VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build" diff --git a/runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml b/runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml new file mode 100644 index 0000000..95cbc71 --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/jobs_router_full.toml @@ -0,0 +1,18 @@ +version = 1 + +[[jobs]] +name = "qwen30-vllm020-router-full-20260716-v1" +gpus = 1 +gpu_model = "H20" +hosts = ["dash0"] +command = "cd /home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-v1/runs/frontier-qwen30-vllm020-profile-v1 && timeout --signal=TERM --kill-after=30s 1020 bash run_router_full.sh" +artifacts = ["artifacts/router-full-v1"] + +[jobs.env] +HOME = "/tmp/wjh" +XDG_CACHE_HOME = "/tmp/wjh/.cache" +VLLM_CACHE_ROOT = "/tmp/wjh/.cache/vllm" +OUTPUT_ROOT = "/home/admin/cpfs/wjh/aituner/aituner-qwen30-vllm020-profile-fleet/artifacts/router-full-v1" +VENV_ROOT = "/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1" +VLLM_SOURCE = "/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build" +MODEL = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B" diff --git a/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py new file mode 100644 index 0000000..b1feca8 --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_router.py @@ -0,0 +1,195 @@ +#!/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() diff --git a/runs/frontier-qwen30-vllm020-profile-v1/run_frontier_linear_full.sh b/runs/frontier-qwen30-vllm020-profile-v1/run_frontier_linear_full.sh index 7371f3a..e68cf88 100644 --- a/runs/frontier-qwen30-vllm020-profile-v1/run_frontier_linear_full.sh +++ b/runs/frontier-qwen30-vllm020-profile-v1/run_frontier_linear_full.sh @@ -47,6 +47,7 @@ timeout --signal=TERM --kill-after=30s 1380 \ --device h20 \ --models qwen3-a3b-30b-moe \ --num_tensor_parallel_workers 1 2 4 \ + --max_tokens 8192 \ --num_tokens_list "${TOKENS[@]}" \ --profile_method cuda_event \ --precision BF16 \ diff --git a/runs/frontier-qwen30-vllm020-profile-v1/run_router_full.sh b/runs/frontier-qwen30-vllm020-profile-v1/run_router_full.sh new file mode 100644 index 0000000..b4fd7a5 --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/run_router_full.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash + +set -euo pipefail + +OUTPUT_ROOT="${OUTPUT_ROOT:?OUTPUT_ROOT must be set}" +VENV_ROOT="${VENV_ROOT:-/tmp/wjh/venvs/vllm-0.20.0-cu129-profiler-v1}" +VLLM_SOURCE="${VLLM_SOURCE:-/home/admin/cpfs/wjh/agentic-kv/third_party/vllm_v20_build}" +MODEL="${MODEL:-/home/admin/cpfs/wjh/models/Qwen/Qwen3-30B-A3B}" +TOKENS=(1 8 16 32 64 128 256 512 1024 2048 4096 8192) +mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" "${OUTPUT_ROOT}/raw" +exec > >(tee -a "${OUTPUT_ROOT}/logs/full.log") 2>&1 + +IFS=',' read -r -a GPU_IDS <<< "${CUDA_VISIBLE_DEVICES:?fleet GPU is required}" +if [[ "${#GPU_IDS[@]}" -ne 1 ]]; then + echo "ERROR: expected exactly one GPU, got ${CUDA_VISIBLE_DEVICES}" >&2 + exit 1 +fi + +echo "PROFILE_LAUNCH_ECHO host=$(hostname) gpu=${CUDA_VISIBLE_DEVICES} model=${MODEL} runtime=vLLM-0.20.0+cu129 operators=ReplicatedLinear,fused_topk tokens=${TOKENS[*]} dtype=BF16 output=${OUTPUT_ROOT} expected_wall=2-6m hard_wall=900s hard_gpu_cap=0.25_H20h" +date -u +"START_UTC=%Y-%m-%dT%H:%M:%SZ" +nvidia-smi --query-gpu=index,name,driver_version,memory.used,utilization.gpu --format=csv,noheader + +test "$(git -C "${VLLM_SOURCE}" rev-parse HEAD)" = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1" +test -s "${MODEL}/config.json" +git rev-parse HEAD > "${OUTPUT_ROOT}/provenance/aituner.commit" +git -C "${VLLM_SOURCE}" rev-parse HEAD > "${OUTPUT_ROOT}/provenance/vllm-source.commit" +sha256sum profile_vllm020_router.py run_router_full.sh \ + > "${OUTPUT_ROOT}/provenance/source.sha256" +sha256sum "${MODEL}/config.json" > "${OUTPUT_ROOT}/provenance/model-config.sha256" +uv pip freeze --python "${VENV_ROOT}/bin/python" \ + > "${OUTPUT_ROOT}/provenance/pip-freeze.txt" +nvidia-smi --query-gpu=index,uuid,name,driver_version,memory.total \ + --format=csv,noheader > "${OUTPUT_ROOT}/provenance/gpus.csv" +printf '%s\n' "${TOKENS[@]}" > "${OUTPUT_ROOT}/provenance/tokens.txt" + +timeout --signal=TERM --kill-after=30s 780 \ + "${VENV_ROOT}/bin/python" profile_vllm020_router.py \ + --vllm-source "${VLLM_SOURCE}" \ + --model "${MODEL}" \ + --output "${OUTPUT_ROOT}/raw/router.json" \ + --num-tokens "${TOKENS[@]}" \ + --warmup-iters 5 \ + --repeats 20 + +test -s "${OUTPUT_ROOT}/raw/router.json" +sha256sum "${OUTPUT_ROOT}/raw/router.json" "${OUTPUT_ROOT}/provenance"/* \ + > "${OUTPUT_ROOT}/artifacts.sha256" +date -u +"END_UTC=%Y-%m-%dT%H:%M:%SZ" +echo "ROUTER_FULL_COMPLETE cases=${#TOKENS[@]}"