86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
|
"""Microbenchmark of candidate CPU 4-bit weight kernels at Spark-X2.5 shapes.
|
||
|
|
|
||
|
|
Decode is one token, so every linear is a GEMV and the only thing that matters
|
||
|
|
is how fast the weight bytes can be streamed. llama.cpp's Q4_K_M moves ~4.5
|
||
|
|
bits per weight; anything we pick has to get close to that at a comparable
|
||
|
|
fraction of peak bandwidth.
|
||
|
|
"""
|
||
|
|
import argparse, json, time
|
||
|
|
import torch
|
||
|
|
|
||
|
|
SHAPES = [
|
||
|
|
("qkv", 3072, 2048),
|
||
|
|
("o", 2048, 2048),
|
||
|
|
("gate_up", 13312, 2048),
|
||
|
|
("down", 2048, 6656),
|
||
|
|
("lm_head", 131072, 2048),
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def timeit(fn, iters, warmup=5):
|
||
|
|
for _ in range(warmup):
|
||
|
|
fn()
|
||
|
|
best = float("inf")
|
||
|
|
for _ in range(iters):
|
||
|
|
t = time.perf_counter()
|
||
|
|
fn()
|
||
|
|
best = min(best, time.perf_counter() - t)
|
||
|
|
return best
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
ap = argparse.ArgumentParser()
|
||
|
|
ap.add_argument("--m", type=int, default=1)
|
||
|
|
ap.add_argument("--iters", type=int, default=30)
|
||
|
|
ap.add_argument("--group", type=int, default=128)
|
||
|
|
ap.add_argument("--out", default="int4_kernels.json")
|
||
|
|
a = ap.parse_args()
|
||
|
|
torch.set_num_threads(torch.get_num_threads())
|
||
|
|
print("threads", torch.get_num_threads())
|
||
|
|
|
||
|
|
rows = []
|
||
|
|
for name, N, K in SHAPES:
|
||
|
|
x16 = torch.randn(a.m, K, dtype=torch.bfloat16)
|
||
|
|
x32 = x16.float()
|
||
|
|
w16 = torch.randn(N, K, dtype=torch.bfloat16)
|
||
|
|
rec = {"name": name, "N": N, "K": K, "M": a.m}
|
||
|
|
|
||
|
|
# bf16 reference (oneDNN through aten)
|
||
|
|
t = timeit(lambda: torch.nn.functional.linear(x16, w16), a.iters)
|
||
|
|
rec["bf16_ms"] = t * 1e3
|
||
|
|
rec["bf16_GBps"] = N * K * 2 / t / 1e9
|
||
|
|
|
||
|
|
# int8 weight-only (aten _weight_int8pack_mm)
|
||
|
|
try:
|
||
|
|
w8 = torch.randint(-127, 127, (N, K), dtype=torch.int8)
|
||
|
|
s8 = torch.rand(N, dtype=torch.bfloat16) * 0.01
|
||
|
|
t = timeit(lambda: torch.ops.aten._weight_int8pack_mm(x16, w8, s8), a.iters)
|
||
|
|
rec["int8_ms"] = t * 1e3
|
||
|
|
rec["int8_GBps"] = N * K / t / 1e9
|
||
|
|
except Exception as e:
|
||
|
|
rec["int8_err"] = f"{type(e).__name__}: {e}"
|
||
|
|
|
||
|
|
# int4 group-wise (aten _weight_int4pack_mm_for_cpu)
|
||
|
|
for dt, xx in (("bf16", x16), ("fp32", x32)):
|
||
|
|
try:
|
||
|
|
wint = torch.randint(0, 15, (N, K), dtype=torch.int32)
|
||
|
|
packed = torch.ops.aten._convert_weight_to_int4pack_for_cpu(wint, 1)
|
||
|
|
sz = torch.rand(K // a.group, N, 2, dtype=xx.dtype)
|
||
|
|
f = lambda: torch.ops.aten._weight_int4pack_mm_for_cpu(xx, packed, a.group, sz)
|
||
|
|
out = f()
|
||
|
|
t = timeit(f, a.iters)
|
||
|
|
rec[f"int4_{dt}_ms"] = t * 1e3
|
||
|
|
rec[f"int4_{dt}_GBps"] = N * K * 0.5 / t / 1e9
|
||
|
|
rec[f"int4_{dt}_out"] = str(tuple(out.shape)) + " " + str(out.dtype)
|
||
|
|
except Exception as e:
|
||
|
|
rec[f"int4_{dt}_err"] = f"{type(e).__name__}: {e}"
|
||
|
|
|
||
|
|
rows.append(rec)
|
||
|
|
print(json.dumps(rec))
|
||
|
|
|
||
|
|
with open(a.out, "w") as f:
|
||
|
|
json.dump(rows, f, indent=2)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|