Align Frontier FP8 profiling with vLLM runtime
This commit is contained in:
@@ -0,0 +1,220 @@
|
||||
#!/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=False,
|
||||
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()
|
||||
Reference in New Issue
Block a user