wehub-resource-sync
ec436095dd
Book-CI / test (macos-latest) (push) Has been cancelled
Book-CI / test (ubuntu-latest) (push) Has been cancelled
Book-CI / test (windows-latest) (push) Has been cancelled
Release Fake Tag / publish (push) Has been cancelled
Deploy / deploy (macos-latest) (push) Has been cancelled
Deploy / deploy (ubuntu-latest) (push) Has been cancelled
Deploy / deploy (windows-latest) (push) Has been cancelled
Release to PyPI / Build & publish sglang-kt (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.11) (push) Has been cancelled
Release to PyPI / Build kt-kernel (Python 3.12) (push) Has been cancelled
Release to PyPI / Publish kt-kernel to PyPI (push) Has been cancelled
394 lines
15 KiB
Python
394 lines
15 KiB
Python
#!/usr/bin/env python
|
|
# coding=utf-8
|
|
"""
|
|
Torch MoE benchmark with multiple execution paths:
|
|
1) expert: Python loop over experts
|
|
2) batched_bmm: batched matmul path (selected experts only)
|
|
3) batched_einsum: einsum path (selected experts only)
|
|
"""
|
|
import argparse
|
|
import os
|
|
import time
|
|
|
|
import torch
|
|
import torch.nn.quantized as nnq
|
|
|
|
scale, zero_point = 0.1, 0
|
|
|
|
# Keep defaults aligned with bench_moe_amx.py.
|
|
expert_num = 256
|
|
hidden_size = 7168
|
|
intermediate_size = 2048
|
|
num_experts_per_tok = 8
|
|
layer_num = 5
|
|
qlen = 1
|
|
warm_up_iter = 1000
|
|
test_iter = 10000
|
|
gen_iter = 3000
|
|
|
|
num_threads = 64
|
|
interop_threads = 1
|
|
exclude_input_quant_time = True
|
|
|
|
|
|
def parse_csv(value: str):
|
|
return [item.strip() for item in value.split(",") if item.strip()]
|
|
|
|
|
|
def configure_torch_threads(threads: int, interop: int):
|
|
os.environ["OMP_NUM_THREADS"] = str(threads)
|
|
os.environ["MKL_NUM_THREADS"] = str(threads)
|
|
torch.set_num_threads(threads)
|
|
torch.set_num_interop_threads(interop)
|
|
|
|
|
|
def act_fn(x):
|
|
return x / (1.0 + torch.exp(-x))
|
|
|
|
|
|
def build_common_inputs():
|
|
expert_ids = (
|
|
torch.rand(gen_iter * qlen, expert_num, device="cpu")
|
|
.argsort(dim=-1)[:, :num_experts_per_tok]
|
|
.reshape(gen_iter, qlen, num_experts_per_tok)
|
|
.contiguous()
|
|
)
|
|
weights = torch.rand((gen_iter, qlen, num_experts_per_tok), dtype=torch.float32, device="cpu").contiguous()
|
|
inputs = torch.randn((layer_num, qlen, hidden_size), dtype=torch.bfloat16, device="cpu").contiguous()
|
|
return expert_ids, weights, inputs
|
|
|
|
|
|
def build_float_projections(proj_dtype: torch.dtype):
|
|
gate_projs, up_projs, down_projs = [], [], []
|
|
for _ in range(layer_num):
|
|
gate = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device="cpu").contiguous()
|
|
up = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device="cpu").contiguous()
|
|
down = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device="cpu").contiguous()
|
|
gate_projs.append(gate.to(proj_dtype))
|
|
up_projs.append(up.to(proj_dtype))
|
|
down_projs.append(down.to(proj_dtype))
|
|
return gate_projs, up_projs, down_projs
|
|
|
|
|
|
def build_int8_projections():
|
|
gate_projs, up_projs, down_projs = [], [], []
|
|
for _ in range(layer_num):
|
|
gate = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device="cpu").contiguous()
|
|
up = torch.randn((expert_num, intermediate_size, hidden_size), dtype=torch.float32, device="cpu").contiguous()
|
|
down = torch.randn((expert_num, hidden_size, intermediate_size), dtype=torch.float32, device="cpu").contiguous()
|
|
|
|
q_gate_layer, q_up_layer, q_down_layer = [], [], []
|
|
for i in range(expert_num):
|
|
gate_q = torch.quantize_per_tensor(gate[i], scale, zero_point, torch.qint8)
|
|
up_q = torch.quantize_per_tensor(up[i], scale, zero_point, torch.qint8)
|
|
down_q = torch.quantize_per_tensor(down[i], scale, zero_point, torch.qint8)
|
|
|
|
q_gate = nnq.Linear(hidden_size, intermediate_size)
|
|
q_up = nnq.Linear(hidden_size, intermediate_size)
|
|
q_down = nnq.Linear(intermediate_size, hidden_size)
|
|
q_gate.set_weight_bias(gate_q, None)
|
|
q_up.set_weight_bias(up_q, None)
|
|
q_down.set_weight_bias(down_q, None)
|
|
|
|
q_gate_layer.append(q_gate)
|
|
q_up_layer.append(q_up)
|
|
q_down_layer.append(q_down)
|
|
|
|
gate_projs.append(q_gate_layer)
|
|
up_projs.append(q_up_layer)
|
|
down_projs.append(q_down_layer)
|
|
|
|
return gate_projs, up_projs, down_projs
|
|
|
|
|
|
def moe_expert_float(input_fp, expert_ids_one, weights_one, gate_proj, up_proj, down_proj):
|
|
counts = expert_ids_one.new_zeros((expert_ids_one.shape[0], expert_num))
|
|
counts.scatter_(1, expert_ids_one, 1)
|
|
tokens_per_expert = counts.sum(dim=0)
|
|
|
|
idxs = expert_ids_one.reshape(-1).argsort()
|
|
sorted_tokens = input_fp[idxs // expert_ids_one.shape[1]]
|
|
|
|
outputs = []
|
|
start_idx = 0
|
|
for expert_idx, num_tokens in enumerate(tokens_per_expert):
|
|
end_idx = start_idx + num_tokens
|
|
if num_tokens == 0:
|
|
continue
|
|
token_block = sorted_tokens[start_idx:end_idx]
|
|
gate_buf = torch.mm(token_block.to(gate_proj.dtype), gate_proj[expert_idx].t())
|
|
up_buf = torch.mm(token_block.to(up_proj.dtype), up_proj[expert_idx].t())
|
|
inter = act_fn(gate_buf) * up_buf
|
|
out = torch.mm(inter.to(down_proj.dtype), down_proj[expert_idx].t())
|
|
outputs.append(out)
|
|
start_idx = end_idx
|
|
|
|
concat_out = torch.cat(outputs, dim=0) if outputs else sorted_tokens.new_empty(0)
|
|
reordered = torch.empty_like(concat_out)
|
|
reordered[idxs] = concat_out
|
|
return (
|
|
reordered.view(*expert_ids_one.shape, -1)
|
|
.type(weights_one.dtype)
|
|
.mul_(weights_one.unsqueeze(dim=-1))
|
|
.sum(dim=1)
|
|
.type(reordered.dtype)
|
|
)
|
|
|
|
|
|
def moe_expert_int8(input_fp, expert_ids_one, weights_one, gate_proj, up_proj, down_proj, input_q=None):
|
|
counts = expert_ids_one.new_zeros((expert_ids_one.shape[0], expert_num))
|
|
counts.scatter_(1, expert_ids_one, 1)
|
|
tokens_per_expert = counts.sum(dim=0)
|
|
|
|
idxs = expert_ids_one.reshape(-1).argsort()
|
|
if input_q is None:
|
|
input_q = torch.quantize_per_tensor(input_fp.to(torch.float32), scale, zero_point, torch.quint8)
|
|
sorted_tokens_q = input_q[idxs // expert_ids_one.shape[1]]
|
|
|
|
outputs = []
|
|
start_idx = 0
|
|
for expert_idx, num_tokens in enumerate(tokens_per_expert):
|
|
end_idx = start_idx + num_tokens
|
|
if num_tokens == 0:
|
|
continue
|
|
token_block_q = sorted_tokens_q[start_idx:end_idx]
|
|
gate_buf = gate_proj[expert_idx](token_block_q).dequantize()
|
|
up_buf = up_proj[expert_idx](token_block_q).dequantize()
|
|
inter = act_fn(gate_buf) * up_buf
|
|
inter_q = torch.quantize_per_tensor(inter, scale, zero_point, torch.quint8)
|
|
out = down_proj[expert_idx](inter_q).dequantize()
|
|
outputs.append(out)
|
|
start_idx = end_idx
|
|
|
|
concat_out = torch.cat(outputs, dim=0) if outputs else torch.empty((0, hidden_size), dtype=torch.float32)
|
|
reordered = torch.empty_like(concat_out)
|
|
reordered[idxs] = concat_out
|
|
return (
|
|
reordered.view(*expert_ids_one.shape, -1)
|
|
.type(weights_one.dtype)
|
|
.mul_(weights_one.unsqueeze(dim=-1))
|
|
.sum(dim=1)
|
|
.type(reordered.dtype)
|
|
)
|
|
|
|
|
|
def moe_batched_bmm(input_fp, expert_ids_one, weights_one, gate_proj, up_proj, down_proj):
|
|
q, k = expert_ids_one.shape
|
|
x = input_fp.to(gate_proj.dtype)
|
|
flat_ids = expert_ids_one.reshape(-1)
|
|
|
|
gate_sel = gate_proj.index_select(0, flat_ids).view(q, k, intermediate_size, hidden_size)
|
|
up_sel = up_proj.index_select(0, flat_ids).view(q, k, intermediate_size, hidden_size)
|
|
down_sel = down_proj.index_select(0, flat_ids).view(q, k, hidden_size, intermediate_size)
|
|
|
|
x_rep = x.unsqueeze(1).expand(-1, k, -1).reshape(-1, hidden_size).unsqueeze(-1)
|
|
gate_buf = (
|
|
torch.bmm(gate_sel.reshape(-1, intermediate_size, hidden_size), x_rep).squeeze(-1).view(q, k, intermediate_size)
|
|
)
|
|
up_buf = (
|
|
torch.bmm(up_sel.reshape(-1, intermediate_size, hidden_size), x_rep).squeeze(-1).view(q, k, intermediate_size)
|
|
)
|
|
|
|
inter = act_fn(gate_buf) * up_buf
|
|
out = (
|
|
torch.bmm(
|
|
down_sel.reshape(-1, hidden_size, intermediate_size),
|
|
inter.reshape(-1, intermediate_size).unsqueeze(-1),
|
|
)
|
|
.squeeze(-1)
|
|
.view(q, k, hidden_size)
|
|
)
|
|
|
|
return (out.type(weights_one.dtype) * weights_one.unsqueeze(-1)).sum(dim=1).type(out.dtype)
|
|
|
|
|
|
def moe_batched_einsum(input_fp, expert_ids_one, weights_one, gate_proj, up_proj, down_proj):
|
|
q, k = expert_ids_one.shape
|
|
x = input_fp.to(gate_proj.dtype)
|
|
flat_ids = expert_ids_one.reshape(-1)
|
|
|
|
gate_sel = gate_proj.index_select(0, flat_ids).view(q, k, intermediate_size, hidden_size)
|
|
up_sel = up_proj.index_select(0, flat_ids).view(q, k, intermediate_size, hidden_size)
|
|
down_sel = down_proj.index_select(0, flat_ids).view(q, k, hidden_size, intermediate_size)
|
|
|
|
gate_buf = torch.einsum("qh,qkih->qki", x, gate_sel)
|
|
up_buf = torch.einsum("qh,qkih->qki", x, up_sel)
|
|
inter = act_fn(gate_buf) * up_buf
|
|
out = torch.einsum("qki,qkhi->qkh", inter, down_sel)
|
|
|
|
return (out.type(weights_one.dtype) * weights_one.unsqueeze(-1)).sum(dim=1).type(out.dtype)
|
|
|
|
|
|
def run_one_iter(
|
|
exec_path, quant_mode, input_tensor, expert_ids_one, weights_one, gate_proj, up_proj, down_proj, input_q=None
|
|
):
|
|
if quant_mode == "qint8":
|
|
if exec_path != "expert":
|
|
raise ValueError("qint8 only supports expert path in this benchmark")
|
|
return moe_expert_int8(input_tensor, expert_ids_one, weights_one, gate_proj, up_proj, down_proj, input_q)
|
|
|
|
if exec_path == "expert":
|
|
return moe_expert_float(input_tensor, expert_ids_one, weights_one, gate_proj, up_proj, down_proj)
|
|
if exec_path == "batched_bmm":
|
|
return moe_batched_bmm(input_tensor, expert_ids_one, weights_one, gate_proj, up_proj, down_proj)
|
|
if exec_path == "batched_einsum":
|
|
return moe_batched_einsum(input_tensor, expert_ids_one, weights_one, gate_proj, up_proj, down_proj)
|
|
|
|
raise ValueError(f"Unknown exec_path={exec_path}")
|
|
|
|
|
|
def bench_moe(quant_mode: str, exec_path: str = "expert"):
|
|
with torch.inference_mode(mode=True):
|
|
if quant_mode == "fp32":
|
|
proj_type = torch.float32
|
|
bytes_per_elem = 4.0
|
|
elif quant_mode == "fp16":
|
|
proj_type = torch.float16
|
|
bytes_per_elem = 2.0
|
|
elif quant_mode == "bf16":
|
|
proj_type = torch.bfloat16
|
|
bytes_per_elem = 2.0
|
|
elif quant_mode == "qint8":
|
|
proj_type = torch.qint8
|
|
bytes_per_elem = 1.0
|
|
else:
|
|
raise ValueError(f"Unsupported quant_mode={quant_mode}")
|
|
|
|
if quant_mode == "qint8":
|
|
gate_projs, up_projs, down_projs = build_int8_projections()
|
|
else:
|
|
gate_projs, up_projs, down_projs = build_float_projections(proj_type)
|
|
|
|
expert_ids, weights, inputs = build_common_inputs()
|
|
pre_quant_inputs = None
|
|
if quant_mode == "qint8" and exclude_input_quant_time:
|
|
pre_quant_inputs = [
|
|
torch.quantize_per_tensor(inputs[i].to(torch.float32), scale, zero_point, torch.quint8)
|
|
for i in range(layer_num)
|
|
]
|
|
|
|
for i in range(warm_up_iter):
|
|
layer_idx = i % layer_num
|
|
gen_idx = i % gen_iter
|
|
input_q = pre_quant_inputs[layer_idx] if pre_quant_inputs is not None else None
|
|
run_one_iter(
|
|
exec_path,
|
|
quant_mode,
|
|
inputs[layer_idx],
|
|
expert_ids[gen_idx],
|
|
weights[gen_idx],
|
|
gate_projs[layer_idx],
|
|
up_projs[layer_idx],
|
|
down_projs[layer_idx],
|
|
input_q,
|
|
)
|
|
|
|
start = time.perf_counter()
|
|
for i in range(test_iter):
|
|
layer_idx = i % layer_num
|
|
gen_idx = i % gen_iter
|
|
input_q = pre_quant_inputs[layer_idx] if pre_quant_inputs is not None else None
|
|
run_one_iter(
|
|
exec_path,
|
|
quant_mode,
|
|
inputs[layer_idx],
|
|
expert_ids[gen_idx],
|
|
weights[gen_idx],
|
|
gate_projs[layer_idx],
|
|
up_projs[layer_idx],
|
|
down_projs[layer_idx],
|
|
input_q,
|
|
)
|
|
end = time.perf_counter()
|
|
|
|
total_time = end - start
|
|
time_us = total_time / test_iter * 1e6
|
|
|
|
work_elems = hidden_size * intermediate_size * 3 * num_experts_per_tok * qlen
|
|
bandwidth = work_elems * bytes_per_elem * test_iter / total_time / 1e9
|
|
flops = work_elems * 2 * test_iter / total_time / 1e12
|
|
|
|
print("Quant mode:", quant_mode)
|
|
print("Exec path:", exec_path)
|
|
print("Time(s):", total_time)
|
|
print("Iteration:", test_iter)
|
|
print("Time(us) per iteration:", time_us)
|
|
print("Bandwidth:", bandwidth, "GB/s")
|
|
print("Flops:", flops, "TFLOPS")
|
|
if quant_mode == "qint8":
|
|
print("Exclude input quantization time:", exclude_input_quant_time)
|
|
print("Note: intermediate quant/dequant is still inside forward path.")
|
|
print("")
|
|
|
|
|
|
def main():
|
|
global expert_num
|
|
global hidden_size
|
|
global intermediate_size
|
|
global num_experts_per_tok
|
|
global layer_num
|
|
global qlen
|
|
global warm_up_iter
|
|
global test_iter
|
|
global gen_iter
|
|
global num_threads
|
|
global interop_threads
|
|
global exclude_input_quant_time
|
|
|
|
parser = argparse.ArgumentParser(description="Torch MoE benchmark")
|
|
parser.add_argument("--expert-num", type=int, default=expert_num)
|
|
parser.add_argument("--hidden-size", type=int, default=hidden_size)
|
|
parser.add_argument("--intermediate-size", type=int, default=intermediate_size)
|
|
parser.add_argument("--num-experts-per-tok", type=int, default=num_experts_per_tok)
|
|
parser.add_argument("--layer-num", type=int, default=layer_num)
|
|
parser.add_argument("--qlen", type=int, default=qlen)
|
|
parser.add_argument("--warm-up-iter", type=int, default=warm_up_iter)
|
|
parser.add_argument("--test-iter", type=int, default=test_iter)
|
|
parser.add_argument("--gen-iter", type=int, default=gen_iter)
|
|
parser.add_argument("--threads", type=int, default=num_threads)
|
|
parser.add_argument("--interop-threads", type=int, default=interop_threads)
|
|
parser.add_argument("--modes", type=str, default="bf16,qint8")
|
|
parser.add_argument("--exec-paths", type=str, default="expert,batched_bmm,batched_einsum")
|
|
parser.add_argument("--include-input-quant-time", action="store_true", default=False)
|
|
args = parser.parse_args()
|
|
|
|
expert_num = args.expert_num
|
|
hidden_size = args.hidden_size
|
|
intermediate_size = args.intermediate_size
|
|
num_experts_per_tok = args.num_experts_per_tok
|
|
layer_num = args.layer_num
|
|
qlen = args.qlen
|
|
warm_up_iter = args.warm_up_iter
|
|
test_iter = args.test_iter
|
|
gen_iter = args.gen_iter
|
|
num_threads = args.threads
|
|
interop_threads = args.interop_threads
|
|
exclude_input_quant_time = not args.include_input_quant_time
|
|
|
|
configure_torch_threads(num_threads, interop_threads)
|
|
|
|
modes = parse_csv(args.modes)
|
|
exec_paths = parse_csv(args.exec_paths)
|
|
|
|
print("[config] torch bench")
|
|
print(
|
|
f"[config] E={expert_num}, H={hidden_size}, I={intermediate_size}, topk={num_experts_per_tok}, "
|
|
f"layers={layer_num}, qlen={qlen}"
|
|
)
|
|
print(f"[config] warmup={warm_up_iter}, test={test_iter}, gen_iter={gen_iter}")
|
|
print(f"[config] threads={num_threads}, interop_threads={interop_threads}")
|
|
print(f"[config] modes={modes}, exec_paths={exec_paths}")
|
|
print(f"[config] exclude_input_quant_time={exclude_input_quant_time}")
|
|
|
|
for mode in modes:
|
|
for path in exec_paths:
|
|
if mode == "qint8" and path != "expert":
|
|
print(f"Skip mode={mode}, exec_path={path}: qint8 only supports expert path")
|
|
print("")
|
|
continue
|
|
bench_moe(mode, path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|