Add exact vLLM 0.20 MoE profile smoke
This commit is contained in:
361
runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py
Normal file
361
runs/frontier-qwen30-vllm020-profile-v1/profile_vllm020_moe.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user