测试脚本已上传，镜像为vllm-metax-0.23.0-maca.ai3.8.0.103-torch2.10-py312-ubuntu22.04-amd64，宿主机的驱动版本
| MX-SMI 2.3.4 Kernel Mode Driver Version: 3.9.25 |
| MACA Version: 3.8.2.6 BIOS Version: 1.35.4.0 |
经过测试，除了一张GPU读写方向带宽基本一致以外，其余GPU读带宽均大于写带宽，读方向带宽为1600GB/s左右：
➜ ✗ python3 benchmark/probe_rw_direction.py --out /tmp/rw.json 2>&1 | tail -22
buffer 1024 MiB fp32, per-GPU allocate/free, min-of-20 cuda-event timing
gpu name w4 read w4 write w8 read w8 write rd/wr copy free GiB
0 MetaX C500 1603 412 1602 415 3.89 575 62.10
1 MetaX C500 1605 413 1595 416 3.88 574 62.10
2 MetaX C500 1605 411 1578 414 3.90 573 62.10
3 MetaX C500 1611 413 1608 415 3.90 574 62.10
4 MetaX C500 1457 1365 1214 1154 1.07 1189 31.67
5 MetaX C500 1605 412 1606 415 3.87 576 62.10
6 MetaX C500 1600 412 1610 415 3.88 574 62.09
执行脚本方法：
python3 probe_rw_direction.py --out /tmp/rw.json 2>&1
请问这是正常现象吗？
import argparse, json, os, statistics as st
import torch, triton, triton.language as tl
@triton.jit
def k_read(x, out, N, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
v = tl.load(x + offs, mask=offs < N, other=0.0)
tl.store(out + pid, tl.sum(v))
@triton.jit
def k_write(y, N, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offs = pid * BLOCK + tl.arange(0, BLOCK)
tl.store(y + offs, 1.0, mask=offs < N)
def t(fn, iters=20, warm=6):
for _ in range(warm):
fn()
torch.cuda.synchronize()
ts = []
for _ in range(iters):
a = torch.cuda.Event(True)
b = torch.cuda.Event(True)
a.record()
fn()
b.record()
b.synchronize()
ts.append(a.elapsed_time(b) / 1e3)
return min(ts)
def probe_one(g, mib, warps):
"""Everything for one GPU. Allocates and frees inside, so a card already
holding a server's weights is fine as long as the buffers fit."""
torch.cuda.set_device(g)
dev = f"cuda:{g}"
N = mib * 1024 * 1024 // 4 # fp32 elements
x = torch.randn(N, dtype=torch.float32, device=dev)
y = torch.empty_like(x)
BLOCK = 4096
G = triton.cdiv(N, BLOCK)
red = torch.empty(G, dtype=torch.float32, device=dev)
B = N * 4
r = {"triton": {}, "torch": {}}
for nw in warps:
tr = t(lambda: k_read[(G,)](x, red, N, BLOCK=BLOCK, num_warps=nw))
tw = t(lambda: k_write[(G,)](y, N, BLOCK=BLOCK, num_warps=nw))
r["triton"][str(nw)] = {"read": B / 1e9 / tr, "write": B / 1e9 / tw}
for nm, fn, by in (("sum_read", lambda: x.sum(), B),
("fill_write", lambda: y.fill_(1.0), B),
("copy_rdwr", lambda: y.copy_(x), 2 * B),
("mul_rmw", lambda: x.mul_(1.0001), 2 * B)):
r["torch"][nm] = by / 1e9 / t(fn)
del x, y, red
torch.cuda.empty_cache()
return r
def main():
p = argparse.ArgumentParser()
p.add_argument("--gpus", default="",
help="comma list of torch device indices; default every visible GPU")
p.add_argument("--mib", type=int, default=1024, help="buffer size, well past L2")
p.add_argument("--warps", default="4,8", help="comma list of num_warps to try")
p.add_argument("--out", default="", help="write the full result as JSON here")
a = p.parse_args()
warps = [int(w) for w in a.warps.split(",") if w != ""]
gpus = ([int(x) for x in a.gpus.split(",") if x != ""]
or list(range(torch.cuda.device_count())))
cvd = os.environ.get("CUDA_VISIBLE_DEVICES") or os.environ.get("MACA_VISIBLE_DEVICES")
if cvd:
print(f"!! CUDA_VISIBLE_DEVICES={cvd} is set -- the 'gpu' column below is the\n"
f"!! TORCH index, which is a position within that list, NOT the physical\n"
f"!! index mx-smi/nvidia-smi prints. Unset it before producing a report\n"
f"!! that names a card.")
print(f"buffer {a.mib} MiB fp32, per-GPU allocate/free, min-of-20 cuda-event timing\n")
out = {"mib": a.mib, "visible_devices": cvd, "gpus": {}}
hdr = "".join(f"{f'w{w} read':>11}{f'w{w} write':>11}" for w in warps)
print(f" {'gpu':<4}{'name':<24}{hdr}{'rd/wr':>8}{'copy':>9}{'free GiB':>10}")
for g in gpus:
torch.cuda.set_device(g)
free, _ = torch.cuda.mem_get_info(g)
name = torch.cuda.get_device_name(g)
if free < a.mib * 1024 * 1024 * 3:
print(f" {g:<4}{name[:23]:<24} skipped: needs 3x the buffer free "
f"({free / 1024 ** 3:.1f} GiB)")
continue
r = probe_one(g, a.mib, warps)
r["name"] = name
r["free_gib"] = round(free / 1024 ** 3, 2)
best = max(warps, key=lambda w: r["triton"][str(w)]["read"])
r["ratio_read_over_write"] = (r["triton"][str(best)]["read"]
/ r["triton"][str(best)]["write"])
out["gpus"][str(g)] = r
cells = "".join(f"{r['triton'][str(w)]['read']:>11.0f}"
f"{r['triton'][str(w)]['write']:>11.0f}" for w in warps)
print(f" {g:<4}{name[:23]:<24}{cells}{r['ratio_read_over_write']:>8.2f}"
f"{r['torch']['copy_rdwr']:>9.0f}{r['free_gib']:>10.2f}")
got = out["gpus"]
if len(got) > 1:
for k, get in (("read", lambda r: r["triton"][str(warps[0])]["read"]),
("write", lambda r: r["triton"][str(warps[0])]["write"]),
("copy", lambda r: r["torch"]["copy_rdwr"])):
v = {g: get(r) for g, r in got.items()}
lo, hi = min(v, key=v.get), max(v, key=v.get)
out[f"{k}_spread"] = round(v[hi] / v[lo], 3)
print(f"\n {k:<6} spread {v[hi] / v[lo]:.2f}x across GPUs "
f"(min {v[lo]:.0f} on gpu {lo}, max {v[hi]:.0f} on gpu {hi})")
print("\n A spread well above 1.1x on identical parts under no load is the\n"
" finding. Report the whole table, not one card.")
else:
print("\n ONE GPU ONLY. This number has no reference. Nothing here can\n"
" distinguish a slow card from a slow card model -- run every card.")
if a.out:
json.dump(out, open(a.out, "w"), indent=1)
print(f"\n -> {a.out}")
if __name__ == "__main__":
main()
