#!/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, ) # This process profiles one TP-local weight shard. Keep the global # runtime context single-rank so vLLM does not initialize a collective; # the action-conditioned shard size remains explicit in moe_config and # the real TP2/TP4 all-reduce is profiled in a separate multi-GPU run. vllm_config = VllmConfig( parallel_config=ParallelConfig(tensor_parallel_size=1) ) 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": ( "one TP-local weight shard: 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()