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
314 lines
12 KiB
Python
314 lines
12 KiB
Python
#!/usr/bin/env python
|
||
# coding=utf-8
|
||
"""Benchmark MXFP4 MoE kernel — V4-Flash shape, mat-vec / mat-mat 双路径覆盖。
|
||
|
||
Synthesizes V4-Flash-shaped MXFP4 weights (random nibbles + bf16 group scales),
|
||
runs the chosen backend over a list of batch sizes M, prints a throughput table.
|
||
|
||
Routing modes (决定是否触发 mat-mat 路径):
|
||
balanced —— 每 token randperm(EXPERT_NUM)[:TOP_K]; 平均 per-expert m ≈
|
||
M*top_k/expert_num. V4 真实路由分布. 大 batch (M=1024) 才
|
||
平均触发 mat-mat (per-expert m ≥ 4).
|
||
concentrated —— 所有 M token 共用同一组 top_k expert; per-expert m = M.
|
||
M=4 立即触发 mat-mat —— 用来直观放大 mat-mat 性能优势.
|
||
|
||
Usage:
|
||
python bench/bench_fp4_moe.py --backend v1
|
||
python bench/bench_fp4_moe.py --backend v1 --routing concentrated
|
||
python bench/bench_fp4_moe.py --all --routing concentrated # 所有可用 backend 对比
|
||
|
||
`--backend` 是预留扩展点; 当前编译只绑定 v1 (AMXFP4_KGroup_MOE)。要选 v2/v3
|
||
需要 ext_bindings 里加新绑定。`--all` 会自动检测哪些 backend 可用。
|
||
"""
|
||
import argparse
|
||
import json
|
||
import os
|
||
import platform
|
||
import subprocess
|
||
import sys
|
||
import time
|
||
|
||
import torch
|
||
|
||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "build"))
|
||
from kt_kernel import kt_kernel_ext # noqa: E402
|
||
|
||
# ----- V4-Flash MoE shape -----
|
||
HIDDEN = 4096
|
||
INTER = 2048
|
||
EXPERT_NUM = 256
|
||
TOP_K = 6
|
||
K_GROUP_SIZE = 32
|
||
|
||
# ----- bench knobs -----
|
||
# M=1024 时平均 per-expert m ≈ M*6/256 = 24, balanced 路由也能触发 mat-mat。
|
||
DEFAULT_M_LIST = [1, 4, 16, 64, 256, 1024]
|
||
WARMUP_ITER = 200
|
||
TEST_ITER = 2000
|
||
|
||
# ----- WorkerPool: 2 NUMA × 40 thread (matches bench_k2_moe_amx.py) -----
|
||
WORKER_NUMA = 2
|
||
WORKER_THREADS_PER_NUMA = 40
|
||
|
||
# ----- Backend registry: name → kt_kernel_ext.moe class (None = not bound) -----
|
||
BACKENDS = {
|
||
"v1": getattr(kt_kernel_ext.moe, "AMXFP4_KGroup_MOE", None),
|
||
# 预留扩展点; 加新 backend 时在 ext_bindings 绑定后这里加一行即可。
|
||
"v2": getattr(kt_kernel_ext.moe, "AMXFP4_KGroup_MOE_V2", None),
|
||
}
|
||
|
||
# OCP MXFP4 (E2M1) codepoints — same order as the kernel's LUT.
|
||
E2M1_VALUES = torch.tensor(
|
||
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0,
|
||
-0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0],
|
||
dtype=torch.float32,
|
||
)
|
||
|
||
|
||
def quantize_mxfp4_tensor(weights: torch.Tensor, group_size: int):
|
||
"""[E, N, K] fp32/bf16 → packed nibbles uint8 [E, N, K/2] + bf16 scales [E, N, K/gs]."""
|
||
w = weights.to(torch.float32)
|
||
e, rows, cols = w.shape
|
||
assert cols % group_size == 0 and cols % 2 == 0
|
||
reshaped = w.view(e, rows, cols // group_size, group_size)
|
||
max_abs = reshaped.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8)
|
||
scales = (max_abs / 6.0).squeeze(-1)
|
||
normalized = reshaped / scales.unsqueeze(-1)
|
||
distances = torch.abs(normalized.unsqueeze(-1) - E2M1_VALUES.view(1, 1, 1, 1, 16))
|
||
nibbles = distances.argmin(dim=-1).to(torch.uint8).view(e, rows, cols // 2, 2)
|
||
lo = nibbles[..., 0]
|
||
hi = nibbles[..., 1]
|
||
packed = ((hi << 4) | lo).contiguous() # uint8 [E, N, K/2]
|
||
scales = scales.to(torch.bfloat16).contiguous() # bf16 [E, N, K/gs]
|
||
return packed, scales
|
||
|
||
|
||
def build_synth_weights():
|
||
torch.manual_seed(0)
|
||
gate = torch.randn((EXPERT_NUM, INTER, HIDDEN), dtype=torch.float32) / 100
|
||
up = torch.randn((EXPERT_NUM, INTER, HIDDEN), dtype=torch.float32) / 100
|
||
down = torch.randn((EXPERT_NUM, HIDDEN, INTER), dtype=torch.float32) / 100
|
||
gw, gs = quantize_mxfp4_tensor(gate, K_GROUP_SIZE)
|
||
uw, us = quantize_mxfp4_tensor(up, K_GROUP_SIZE)
|
||
dw, ds = quantize_mxfp4_tensor(down, K_GROUP_SIZE)
|
||
return {
|
||
"gate_w": gw, "up_w": uw, "down_w": dw,
|
||
"gate_s": gs, "up_s": us, "down_s": ds,
|
||
}
|
||
|
||
|
||
def build_moe(backend: str, weights, cpu_infer):
|
||
cls = BACKENDS.get(backend)
|
||
if cls is None:
|
||
raise RuntimeError(
|
||
f"backend={backend} not bound in this build. Available: "
|
||
f"{[k for k, v in BACKENDS.items() if v is not None]}"
|
||
)
|
||
cfg = kt_kernel_ext.moe.MOEConfig(EXPERT_NUM, TOP_K, HIDDEN, INTER, 0)
|
||
cfg.max_len = max(DEFAULT_M_LIST)
|
||
cfg.pool = cpu_infer.backend_
|
||
cfg.quant_config.bits = 4
|
||
cfg.quant_config.group_size = K_GROUP_SIZE
|
||
cfg.quant_config.zero_point = False
|
||
cfg.gate_projs = [[t.data_ptr() for t in weights["gate_w"]]]
|
||
cfg.up_projs = [[t.data_ptr() for t in weights["up_w"]]]
|
||
cfg.down_projs = [[t.data_ptr() for t in weights["down_w"]]]
|
||
cfg.gate_scales = [[t.data_ptr() for t in weights["gate_s"]]]
|
||
cfg.up_scales = [[t.data_ptr() for t in weights["up_s"]]]
|
||
cfg.down_scales = [[t.data_ptr() for t in weights["down_s"]]]
|
||
moe = cls(cfg)
|
||
p2l = torch.arange(EXPERT_NUM, dtype=torch.int64).contiguous()
|
||
cpu_infer.submit(moe.load_weights_task(p2l.data_ptr()))
|
||
cpu_infer.sync()
|
||
return moe
|
||
|
||
|
||
def make_expert_ids(M: int, routing: str) -> torch.Tensor:
|
||
"""[M, TOP_K] int64 (kernel forward_binding casts to const int64_t*)."""
|
||
if routing == "concentrated":
|
||
# 所有 M token 共用同组 top_k expert → per-expert m = M
|
||
hot = torch.randperm(EXPERT_NUM)[:TOP_K]
|
||
return hot.unsqueeze(0).expand(M, TOP_K).contiguous().to(torch.int64)
|
||
# balanced: 每 token 独立 randperm
|
||
return torch.stack(
|
||
[torch.randperm(EXPERT_NUM)[:TOP_K] for _ in range(M)]
|
||
).to(torch.int64).contiguous()
|
||
|
||
|
||
def bench_one_m(moe, cpu_infer, M: int, routing: str):
|
||
bsz = torch.tensor([M], dtype=torch.int32)
|
||
expert_ids = make_expert_ids(M, routing)
|
||
routing_w = torch.rand((M, TOP_K), dtype=torch.float32).contiguous()
|
||
x = (torch.randn((M, HIDDEN), dtype=torch.bfloat16) / 100).contiguous()
|
||
y = torch.empty((M, HIDDEN), dtype=torch.bfloat16).contiguous()
|
||
|
||
for _ in range(WARMUP_ITER):
|
||
cpu_infer.submit(moe.forward_task(
|
||
bsz.data_ptr(), TOP_K, expert_ids.data_ptr(),
|
||
routing_w.data_ptr(), x.data_ptr(), y.data_ptr(), False))
|
||
cpu_infer.sync()
|
||
|
||
start = time.perf_counter()
|
||
for _ in range(TEST_ITER):
|
||
cpu_infer.submit(moe.forward_task(
|
||
bsz.data_ptr(), TOP_K, expert_ids.data_ptr(),
|
||
routing_w.data_ptr(), x.data_ptr(), y.data_ptr(), False))
|
||
cpu_infer.sync()
|
||
total = time.perf_counter() - start
|
||
|
||
per_iter_us = total / TEST_ITER * 1e6
|
||
tok_per_s = M * TEST_ITER / total
|
||
unique_e = int(torch.unique(expert_ids).numel())
|
||
avg_m_per_expert = float(M * TOP_K) / max(unique_e, 1)
|
||
return {
|
||
"M": M, "iters": TEST_ITER, "total_s": total,
|
||
"per_iter_us": per_iter_us, "tokens_per_s": tok_per_s,
|
||
"unique_experts": unique_e, "avg_m_per_expert": avg_m_per_expert,
|
||
}
|
||
|
||
|
||
def run_backend(backend: str, weights, cpu_infer, m_list, routing):
|
||
print(f"\n[bench-fp4] backend={backend} routing={routing}")
|
||
moe = build_moe(backend, weights, cpu_infer)
|
||
rows = []
|
||
for M in m_list:
|
||
r = bench_one_m(moe, cpu_infer, M, routing)
|
||
rows.append(r)
|
||
print(f" M={M:>5} avg_m/expert={r['avg_m_per_expert']:>6.1f} "
|
||
f"per-iter={r['per_iter_us']:>9.1f} us tok/s={r['tokens_per_s']:>9.1f}")
|
||
return rows
|
||
|
||
|
||
def print_single_table(backend, rows, routing):
|
||
print(f"\n=== Summary ({backend}, routing={routing}) ===")
|
||
print(f"{'M':>5} {'avg_m':>6} {'per-iter us':>12} {'tok/s':>10}")
|
||
for r in rows:
|
||
print(f"{r['M']:>5} {r['avg_m_per_expert']:>6.1f} "
|
||
f"{r['per_iter_us']:>12.1f} {r['tokens_per_s']:>10.1f}")
|
||
|
||
|
||
def print_compare_table(all_rows: dict, routing: str):
|
||
backends = list(all_rows.keys())
|
||
if len(backends) < 2:
|
||
print_single_table(backends[0], all_rows[backends[0]], routing)
|
||
return
|
||
base = backends[0]
|
||
print(f"\n=== {' vs '.join(backends)} (routing={routing}, base={base}) ===")
|
||
header = f"{'M':>5} {'avg_m':>6}"
|
||
for be in backends:
|
||
header += f" {be + ' us':>10}"
|
||
for be in backends[1:]:
|
||
header += f" {be + '/' + base:>8}"
|
||
print(header)
|
||
n_rows = len(all_rows[base])
|
||
for i in range(n_rows):
|
||
line = f"{all_rows[base][i]['M']:>5} {all_rows[base][i]['avg_m_per_expert']:>6.1f}"
|
||
for be in backends:
|
||
line += f" {all_rows[be][i]['per_iter_us']:>10.1f}"
|
||
for be in backends[1:]:
|
||
ratio = all_rows[be][i]['per_iter_us'] / all_rows[base][i]['per_iter_us']
|
||
line += f" {ratio:>8.3f}"
|
||
print(line)
|
||
|
||
|
||
def get_git_commit():
|
||
try:
|
||
commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode().strip()
|
||
dirty = bool(subprocess.check_output(["git", "status", "--porcelain"]).decode().strip())
|
||
return {"commit": commit, "dirty": dirty}
|
||
except Exception as e:
|
||
return {"commit": None, "error": str(e)}
|
||
|
||
|
||
def get_system_info():
|
||
info = {"node": platform.node(), "system": platform.system()}
|
||
cpu_model = None
|
||
try:
|
||
with open("/proc/cpuinfo") as f:
|
||
for line in f:
|
||
if "model name" in line:
|
||
cpu_model = line.split(":", 1)[1].strip()
|
||
break
|
||
except Exception:
|
||
pass
|
||
info["cpu"] = cpu_model
|
||
info["cpu_cores"] = os.cpu_count()
|
||
return info
|
||
|
||
|
||
def main():
|
||
p = argparse.ArgumentParser()
|
||
p.add_argument("--backend", choices=list(BACKENDS.keys()), default="v1",
|
||
help="单 backend 模式（被 --all 覆盖）。当前可用: "
|
||
+ ",".join(k for k, v in BACKENDS.items() if v is not None))
|
||
p.add_argument("--all", action="store_true",
|
||
help="跑所有已绑定的 backend 并打印对比表（自动跳过未绑定的）")
|
||
p.add_argument("--routing", choices=["balanced", "concentrated"], default="balanced",
|
||
help="balanced=每 token randperm (V4 真实); "
|
||
"concentrated=所有 token 共用同组 top_k (per-expert m=M, 放大 mat-mat)")
|
||
p.add_argument("--m-list", type=str, default=None,
|
||
help=f"Comma-separated M values, default: {','.join(map(str, DEFAULT_M_LIST))}")
|
||
p.add_argument("--numa", type=int, default=WORKER_NUMA)
|
||
p.add_argument("--threads-per-numa", type=int, default=WORKER_THREADS_PER_NUMA)
|
||
args = p.parse_args()
|
||
|
||
m_list = [int(x) for x in args.m_list.split(",")] if args.m_list else DEFAULT_M_LIST
|
||
|
||
if args.all:
|
||
backends = [k for k, v in BACKENDS.items() if v is not None]
|
||
if not backends:
|
||
raise RuntimeError("No MXFP4 backend bound in this build.")
|
||
print(f"[bench-fp4] --all: detected backends = {backends}")
|
||
else:
|
||
if BACKENDS.get(args.backend) is None:
|
||
raise RuntimeError(
|
||
f"backend={args.backend} not bound. Available: "
|
||
f"{[k for k, v in BACKENDS.items() if v is not None]}"
|
||
)
|
||
backends = [args.backend]
|
||
|
||
print(f"[bench-fp4] shape=H{HIDDEN}/I{INTER}/E{EXPERT_NUM}/k{TOP_K}/gs{K_GROUP_SIZE} routing={args.routing}")
|
||
print(f"[bench-fp4] WorkerPool: numa={args.numa} threads_per_numa={args.threads_per_numa}")
|
||
print(f"[bench-fp4] m_list: {m_list}")
|
||
|
||
wp = kt_kernel_ext.WorkerPoolConfig()
|
||
wp.subpool_count = args.numa
|
||
wp.subpool_numa_map = list(range(args.numa))
|
||
wp.subpool_thread_count = [args.threads_per_numa] * args.numa
|
||
cpu_infer = kt_kernel_ext.CPUInfer(wp)
|
||
|
||
print("[bench-fp4] synthesizing MXFP4 weights …")
|
||
weights = build_synth_weights()
|
||
|
||
all_rows = {}
|
||
for be in backends:
|
||
all_rows[be] = run_backend(be, weights, cpu_infer, m_list, args.routing)
|
||
|
||
if len(backends) > 1:
|
||
print_compare_table(all_rows, args.routing)
|
||
else:
|
||
print_single_table(backends[0], all_rows[backends[0]], args.routing)
|
||
|
||
# JSONL log
|
||
out_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "bench_fp4_moe.jsonl")
|
||
ts = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
||
git = get_git_commit()
|
||
sys_info = get_system_info()
|
||
with open(out_path, "a") as f:
|
||
for be, rows in all_rows.items():
|
||
record = {
|
||
"backend": be,
|
||
"routing": args.routing,
|
||
"shape": {"hidden": HIDDEN, "inter": INTER, "expert_num": EXPERT_NUM,
|
||
"top_k": TOP_K, "k_group_size": K_GROUP_SIZE},
|
||
"worker_pool": {"numa": args.numa, "threads_per_numa": args.threads_per_numa},
|
||
"rows": rows,
|
||
"git": git, "system": sys_info, "timestamp": ts,
|
||
}
|
||
f.write(json.dumps(record) + "\n")
|
||
print(f"\n[bench-fp4] appended {len(backends)} record(s) → {out_path}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|