From 5a958077a79b63b562daa3fed9f2c01f30318341 Mon Sep 17 00:00:00 2001 From: Gahow Wang Date: Thu, 16 Jul 2026 21:27:35 +0800 Subject: [PATCH] Add exact vLLM 0.20 MoE profile smoke --- .../jobs_moe_smoke.toml | 18 + .../profile_vllm020_moe.py | 361 ++++++++++++++++++ .../run_moe_smoke.sh | 47 +++ 3 files changed, 426 insertions(+) create mode 100644 runs/frontier-qwen30-vllm020-profile-v1/jobs_moe_smoke.toml create mode 100644 runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py create mode 100644 runs/frontier-qwen30-vllm020-profile-v1/run_moe_smoke.sh diff --git a/runs/frontier-qwen30-vllm020-profile-v1/jobs_moe_smoke.toml b/runs/frontier-qwen30-vllm020-profile-v1/jobs_moe_smoke.toml new file mode 100644 index 0000000..48cd8db --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/jobs_moe_smoke.toml @@ -0,0 +1,18 @@ +version = 1 + +[[jobs]] +name = "qwen30-vllm020-moe-smoke-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_moe_smoke.sh" +artifacts = ["artifacts/moe-smoke"] + +[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/moe-smoke" +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_moe.py b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py new file mode 100644 index 0000000..927b316 --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py @@ -0,0 +1,361 @@ +#!/usr/bin/env python3 +"""Profile vLLM 0.20's Qwen3-30B unquantized MoE kernel at TP-local shapes.""" + +from __future__ import annotations + +import argparse +import json +import math +import statistics +import subprocess +from pathlib import Path +from typing import Any + +import torch +import vllm + + +VLLM_VERSION = "0.20.0" +VLLM_COMMIT = "88d34c6409e9fb3c7b8ca0c04756f061d2099eb1" +HIDDEN_DIM = 2048 +INTERMEDIATE_DIM = 768 +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("--tp", type=int, choices=(1, 2, 4), nargs="+", default=[1, 2, 4]) + parser.add_argument("--num-tokens", type=int, nargs="+", default=[8]) + parser.add_argument( + "--routing-modes", + choices=("uniform_random_logits", "hotset8"), + nargs="+", + default=["uniform_random_logits"], + ) + parser.add_argument("--warmup-iters", type=int, default=3) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--device", default="cuda:0") + parser.add_argument("--check-reference", action="store_true") + 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 routing_inputs( + mode: str, num_tokens: int, device: torch.device +) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + from vllm.model_executor.layers.fused_moe import fused_topk + + if mode == "uniform_random_logits": + logits = torch.randn( + (num_tokens, NUM_EXPERTS), device=device, dtype=torch.bfloat16 + ) + hidden_for_topk = torch.empty( + (num_tokens, HIDDEN_DIM), device=device, dtype=torch.bfloat16 + ) + weights, ids, _ = fused_topk( + hidden_for_topk, + logits, + TOP_K, + renormalize=True, + ) + elif mode == "hotset8": + ids = torch.arange(TOP_K, device=device, dtype=torch.int32).repeat( + num_tokens, 1 + ) + weights = torch.full( + (num_tokens, TOP_K), + 1.0 / TOP_K, + device=device, + dtype=torch.float32, + ) + else: + raise ValueError(mode) + + counts = torch.bincount(ids.flatten().to(torch.int64), minlength=NUM_EXPERTS) + counts_cpu = counts.cpu().tolist() + mean_load = num_tokens * TOP_K / NUM_EXPERTS + variance = sum((count - mean_load) ** 2 for count in counts_cpu) / NUM_EXPERTS + return weights, ids, { + "active_experts": sum(count > 0 for count in counts_cpu), + "min_tokens_per_expert": min(counts_cpu), + "max_tokens_per_expert": max(counts_cpu), + "load_cv": math.sqrt(variance) / mean_load if mean_load else 0.0, + "counts": counts_cpu, + } + + +def reference_partial_output( + hidden: torch.Tensor, + w13_original: torch.Tensor, + w2: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, +) -> torch.Tensor: + output = torch.zeros_like(hidden) + for token in range(hidden.shape[0]): + for route in range(TOP_K): + expert = int(topk_ids[token, route]) + gate_up = torch.mv(w13_original[expert], hidden[token]) + gate, up = gate_up.chunk(2) + activated = torch.nn.functional.silu(gate) * up + expert_output = torch.mv(w2[expert], activated) + output[token].add_( + expert_output * topk_weights[token, route].to(expert_output.dtype) + ) + return output + + +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}") + model_config = json.loads(args.model.joinpath("config.json").read_text()) + expected_model = { + "hidden_size": HIDDEN_DIM, + "moe_intermediate_size": INTERMEDIATE_DIM, + "num_experts": NUM_EXPERTS, + "num_experts_per_tok": TOP_K, + "norm_topk_prob": True, + "torch_dtype": "bfloat16", + } + observed_model = {key: model_config.get(key) for key in expected_model} + if observed_model != expected_model: + raise SystemExit( + f"model contract mismatch: expected {expected_model}, got {observed_model}" + ) + + from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config + from vllm.model_executor.layers.fused_moe.activation import MoEActivation + from vllm.model_executor.layers.fused_moe.config import ( + FUSED_MOE_UNQUANTIZED_CONFIG, + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, + ) + from vllm.model_executor.layers.fused_moe.oracle.unquantized import ( + UnquantizedMoeBackend, + convert_to_unquantized_kernel_format, + make_unquantized_moe_kernel, + select_unquantized_moe_backend, + ) + from vllm.utils.math_utils import next_power_of_2 + from vllm.v1.worker.workspace import init_workspace_manager + + device = torch.device(args.device) + torch.accelerator.set_device_index(device) + torch.manual_seed(20260716) + init_workspace_manager(args.device) + max_num_tokens = next_power_of_2(max(args.num_tokens)) + + rows: list[dict[str, Any]] = [] + for tp in args.tp: + parallel = FusedMoEParallelConfig( + tp_size=tp, + tp_rank=0, + pcp_size=1, + pcp_rank=0, + dp_size=1, + dp_rank=0, + ep_size=1, + ep_rank=0, + sp_size=1, + use_ep=False, + all2all_backend="allgather_reducescatter", + enable_eplb=False, + ) + moe_config = FusedMoEConfig( + num_experts=NUM_EXPERTS, + experts_per_token=TOP_K, + hidden_dim=HIDDEN_DIM, + intermediate_size_per_partition=INTERMEDIATE_DIM // tp, + num_local_experts=NUM_EXPERTS, + num_logical_experts=NUM_EXPERTS, + activation=MoEActivation.SILU, + device=device, + routing_method=RoutingMethodType.Renormalize, + moe_parallel_config=parallel, + in_dtype=torch.bfloat16, + max_num_tokens=max_num_tokens, + ) + vllm_config = VllmConfig( + parallel_config=ParallelConfig(tensor_parallel_size=tp) + ) + with set_current_vllm_config(vllm_config): + backend, experts_cls = select_unquantized_moe_backend(moe_config) + if backend != UnquantizedMoeBackend.FLASHINFER_CUTLASS: + raise SystemExit( + "runtime backend mismatch: expected FlashInfer CUTLASS, " + f"got {backend.value} at TP={tp}" + ) + if experts_cls is None: + raise SystemExit(f"missing experts class for {backend.value}") + + w13_original = torch.empty( + (NUM_EXPERTS, 2 * (INTERMEDIATE_DIM // tp), HIDDEN_DIM), + device=device, + dtype=torch.bfloat16, + ).uniform_(-0.01, 0.01) + w2 = torch.empty( + (NUM_EXPERTS, HIDDEN_DIM, INTERMEDIATE_DIM // tp), + device=device, + dtype=torch.bfloat16, + ).uniform_(-0.01, 0.01) + + class Layer: + pass + + layer = Layer() + layer.moe_config = moe_config + w13_kernel, w2_kernel = convert_to_unquantized_kernel_format( + backend, + layer=layer, + w13_weight=w13_original, + w2_weight=w2, + ) + kernel = make_unquantized_moe_kernel( + quant_config=FUSED_MOE_UNQUANTIZED_CONFIG, + moe_config=moe_config, + backend=backend, + experts_cls=experts_cls, + ) + + reference_checked = False + for routing_mode in args.routing_modes: + for num_tokens in args.num_tokens: + hidden = torch.empty( + (num_tokens, HIDDEN_DIM), + device=device, + dtype=torch.bfloat16, + ).uniform_(-0.1, 0.1) + topk_weights, topk_ids, load = routing_inputs( + routing_mode, num_tokens, device + ) + + for _ in range(args.warmup_iters): + output = kernel.apply( + hidden_states=hidden, + w1=w13_kernel, + w2=w2_kernel, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=NUM_EXPERTS, + expert_map=None, + apply_router_weight_on_input=False, + ) + torch.accelerator.synchronize() + + samples: list[float] = [] + for _ in range(args.repeats): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + output = kernel.apply( + hidden_states=hidden, + w1=w13_kernel, + w2=w2_kernel, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=MoEActivation.SILU, + global_num_experts=NUM_EXPERTS, + expert_map=None, + apply_router_weight_on_input=False, + ) + end.record() + torch.accelerator.synchronize() + samples.append(float(start.elapsed_time(end))) + + if output.shape != hidden.shape or not torch.isfinite(output).all(): + raise SystemExit( + f"invalid MoE output TP={tp} M={num_tokens} mode={routing_mode}" + ) + if args.check_reference and not reference_checked: + check_tokens = min(2, num_tokens) + reference = reference_partial_output( + hidden[:check_tokens], + w13_original, + w2, + topk_weights[:check_tokens], + topk_ids[:check_tokens], + ) + torch.testing.assert_close( + output[:check_tokens], reference, atol=0.03, rtol=0.03 + ) + reference_checked = True + + row = { + "tensor_parallel_size": tp, + "num_tokens": num_tokens, + "routing_mode": routing_mode, + "backend": backend.value, + "intermediate_size_per_partition": INTERMEDIATE_DIM // tp, + "output_is_reduced": kernel.output_is_reduced(), + "time_ms": stats_ms(samples), + "routing_load": load, + } + rows.append(row) + print( + json.dumps( + { + "tp": tp, + "num_tokens": num_tokens, + "routing_mode": routing_mode, + "backend": backend.value, + "median_ms": row["time_ms"]["median"], + }, + sort_keys=True, + ), + flush=True, + ) + + del kernel, w13_kernel, w2_kernel, w13_original, w2 + torch.accelerator.empty_cache() + + payload = { + "schema_version": "qwen30_vllm020_moe_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", + "weight_quantization": "none", + "top_k": TOP_K, + "norm_topk_prob": True, + }, + "measurement_scope": ( + "vLLM modular MoE prepare+FlashInfer CUTLASS experts+finalize; " + "router linear/top-k and TP all-reduce excluded" + ), + "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_moe_smoke.sh b/runs/frontier-qwen30-vllm020-profile-v1/run_moe_smoke.sh new file mode 100644 index 0000000..af7b784 --- /dev/null +++ b/runs/frontier-qwen30-vllm020-profile-v1/run_moe_smoke.sh @@ -0,0 +1,47 @@ +#!/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}" +mkdir -p "${OUTPUT_ROOT}/logs" "${OUTPUT_ROOT}/provenance" "${OUTPUT_ROOT}/raw" +exec > >(tee -a "${OUTPUT_ROOT}/logs/smoke.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 operator=FlashInfer-CUTLASS-unquantized-MoE tp_local_shapes=1,2,4 tokens=8 routing=uniform_random_logits reference_check=on dtype=BF16 output=${OUTPUT_ROOT} expected_wall=3-8m 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 + +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_moe.py run_moe_smoke.sh \ + > "${OUTPUT_ROOT}/provenance/source.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" + +timeout --signal=TERM --kill-after=30s 780 \ + "${VENV_ROOT}/bin/python" profile_vllm020_moe.py \ + --vllm-source "${VLLM_SOURCE}" \ + --model "${MODEL}" \ + --output "${OUTPUT_ROOT}/raw/moe-smoke.json" \ + --tp 1 2 4 \ + --num-tokens 8 \ + --routing-modes uniform_random_logits \ + --warmup-iters 3 \ + --repeats 5 \ + --check-reference + +test -s "${OUTPUT_ROOT}/raw/moe-smoke.json" +sha256sum "${OUTPUT_ROOT}/raw/moe-smoke.json" "${OUTPUT_ROOT}/provenance"/* \ + > "${OUTPUT_ROOT}/artifacts.sha256" +date -u +"END_UTC=%Y-%m-%dT%H:%M:%SZ" +echo "MOE_SMOKE_COMPLETE"