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