221 lines
6.6 KiB
Python
221 lines
6.6 KiB
Python
#!/usr/bin/env python3
|
|
"""Compare patched Frontier MoE decomposition with vLLM's serving path."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from frontier.profiling.moe.moe_vllm_kernel import (
|
|
profile_fused_moe_kernel,
|
|
quantize_weights_to_fp8,
|
|
)
|
|
from vllm.model_executor.layers.fused_moe.fused_moe import (
|
|
fused_experts,
|
|
get_config_dtype_str,
|
|
moe_align_block_size,
|
|
try_get_optimal_moe_config,
|
|
)
|
|
|
|
|
|
def _measure(step, warmup: int, active: int) -> dict[str, float]:
|
|
for _ in range(warmup):
|
|
step()
|
|
torch.cuda.synchronize()
|
|
|
|
samples = []
|
|
for _ in range(active):
|
|
start = torch.cuda.Event(enable_timing=True)
|
|
end = torch.cuda.Event(enable_timing=True)
|
|
start.record()
|
|
step()
|
|
end.record()
|
|
torch.cuda.synchronize()
|
|
samples.append(start.elapsed_time(end))
|
|
values = torch.tensor(samples)
|
|
return {
|
|
"min": float(values.min()),
|
|
"median": float(values.median()),
|
|
"mean": float(values.mean()),
|
|
"max": float(values.max()),
|
|
"std": float(values.std()),
|
|
}
|
|
|
|
|
|
def _routing(num_tokens: int, top_k: int, num_experts: int, seed: int):
|
|
generator = torch.Generator(device="cuda")
|
|
generator.manual_seed(seed)
|
|
topk_ids = torch.randint(
|
|
num_experts,
|
|
(num_tokens, top_k),
|
|
generator=generator,
|
|
device="cuda",
|
|
dtype=torch.int64,
|
|
)
|
|
topk_weights = torch.rand(
|
|
(num_tokens, top_k),
|
|
generator=generator,
|
|
device="cuda",
|
|
dtype=torch.float32,
|
|
)
|
|
topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
|
|
return topk_weights, topk_ids
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--tokens", nargs="+", type=int, default=[16, 256, 1024])
|
|
parser.add_argument("--tp", type=int, default=4)
|
|
parser.add_argument("--ep", type=int, default=1)
|
|
parser.add_argument("--warmup", type=int, default=2)
|
|
parser.add_argument("--active", type=int, default=20)
|
|
parser.add_argument("--seed", type=int, default=20260715)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
args = parser.parse_args()
|
|
|
|
hidden_dim = 4096
|
|
expert_hidden_dim = 1536
|
|
global_num_experts = 128
|
|
top_k = 8
|
|
block_shape = [128, 128]
|
|
if global_num_experts % args.ep:
|
|
raise ValueError("EP must divide 128 experts")
|
|
if expert_hidden_dim % args.tp:
|
|
raise ValueError("TP must divide the expert intermediate dimension")
|
|
|
|
local_num_experts = global_num_experts // args.ep
|
|
local_intermediate = expert_hidden_dim // args.tp
|
|
device = torch.device("cuda")
|
|
|
|
torch.manual_seed(args.seed)
|
|
w1_bf16 = torch.randn(
|
|
local_num_experts,
|
|
2 * local_intermediate,
|
|
hidden_dim,
|
|
dtype=torch.bfloat16,
|
|
device=device,
|
|
)
|
|
w2_bf16 = torch.randn(
|
|
local_num_experts,
|
|
hidden_dim,
|
|
local_intermediate,
|
|
dtype=torch.bfloat16,
|
|
device=device,
|
|
)
|
|
w1, w1_scale = quantize_weights_to_fp8(w1_bf16, block_shape=block_shape)
|
|
w2, w2_scale = quantize_weights_to_fp8(w2_bf16, block_shape=block_shape)
|
|
del w1_bf16, w2_bf16
|
|
torch.cuda.empty_cache()
|
|
|
|
rows = []
|
|
for index, num_tokens in enumerate(args.tokens):
|
|
topk_weights, topk_ids = _routing(
|
|
num_tokens,
|
|
top_k,
|
|
global_num_experts,
|
|
args.seed + index,
|
|
)
|
|
hidden_states = torch.randn(
|
|
num_tokens,
|
|
hidden_dim,
|
|
dtype=torch.bfloat16,
|
|
device=device,
|
|
)
|
|
|
|
frontier_grouped = profile_fused_moe_kernel(
|
|
num_tokens=num_tokens,
|
|
num_experts=local_num_experts,
|
|
hidden_dim=hidden_dim,
|
|
expert_hidden_dim=expert_hidden_dim,
|
|
top_k=top_k,
|
|
topk_weights=topk_weights,
|
|
topk_ids=topk_ids,
|
|
tensor_parallel_size=args.tp,
|
|
dtype=torch.bfloat16,
|
|
warmup_steps=args.warmup,
|
|
active_steps=args.active,
|
|
use_fp8=True,
|
|
per_channel_quant=False,
|
|
block_shape=block_shape,
|
|
global_num_experts=global_num_experts,
|
|
)
|
|
|
|
config = try_get_optimal_moe_config(
|
|
w1_shape=w1.shape,
|
|
w2_shape=w2.shape,
|
|
top_k=top_k,
|
|
dtype=get_config_dtype_str(
|
|
torch.bfloat16,
|
|
use_fp8_w8a8=True,
|
|
),
|
|
M=num_tokens,
|
|
block_shape=block_shape,
|
|
)
|
|
|
|
def align_step() -> None:
|
|
moe_align_block_size(
|
|
topk_ids,
|
|
config["BLOCK_SIZE_M"],
|
|
global_num_experts,
|
|
)
|
|
|
|
alignment = _measure(align_step, args.warmup, args.active)
|
|
|
|
def serving_step() -> None:
|
|
fused_experts(
|
|
hidden_states=hidden_states,
|
|
w1=w1,
|
|
w2=w2,
|
|
topk_weights=topk_weights,
|
|
topk_ids=topk_ids,
|
|
inplace=True,
|
|
use_fp8_w8a8=True,
|
|
per_channel_quant=False,
|
|
global_num_experts=global_num_experts,
|
|
w1_scale=w1_scale,
|
|
w2_scale=w2_scale,
|
|
block_shape=block_shape,
|
|
)
|
|
|
|
serving = _measure(serving_step, args.warmup, args.active)
|
|
decomposed_ms = frontier_grouped["median"] + alignment["median"]
|
|
rows.append(
|
|
{
|
|
"num_tokens": num_tokens,
|
|
"block_size_m": config["BLOCK_SIZE_M"],
|
|
"frontier_grouped_ms": frontier_grouped,
|
|
"frontier_alignment_ms": alignment,
|
|
"frontier_decomposed_median_ms": decomposed_ms,
|
|
"vllm_fused_experts_ms": serving,
|
|
"decomposed_over_serving": decomposed_ms / serving["median"],
|
|
}
|
|
)
|
|
|
|
result = {
|
|
"contract": "Frontier grouped_gemm + shuffling alignment vs vLLM fused_experts",
|
|
"model_shape": {
|
|
"hidden_dim": hidden_dim,
|
|
"expert_hidden_dim": expert_hidden_dim,
|
|
"global_num_experts": global_num_experts,
|
|
"local_num_experts": local_num_experts,
|
|
"top_k": top_k,
|
|
"tp": args.tp,
|
|
"ep": args.ep,
|
|
"dtype": "block_fp8_w8a8_bf16_output",
|
|
"block_shape": block_shape,
|
|
},
|
|
"warmup": args.warmup,
|
|
"active": args.active,
|
|
"rows": rows,
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
|
|
print(json.dumps(result, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|