File size: 1,515 Bytes
1ef8270 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | import argparse
import torch
import int4_blackwell
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--iterations", type=int, default=8192)
parser.add_argument("--repeats", type=int, default=10)
parser.add_argument("--launches", type=int, default=20)
args = parser.parse_args()
props = torch.cuda.get_device_properties(0)
blocks = props.multi_processor_count * 4
warps = blocks * 8
flops = warps * 4 * args.iterations * 2 * 16 * 8 * 64
out = torch.empty((blocks, 256), device="cuda", dtype=torch.float32)
for mode in ("e2m1", "a", "b", "ab"):
int4_blackwell.mma_probe(
mode, iterations=args.iterations, blocks=blocks, out=out
)
torch.cuda.synchronize()
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
best_ms = float("inf")
for _ in range(args.repeats):
start.record()
int4_blackwell.mma_probe(
mode,
iterations=args.iterations,
blocks=blocks,
launches=args.launches,
out=out,
)
end.record()
end.synchronize()
best_ms = min(best_ms, start.elapsed_time(end))
per_launch_ms = best_ms / args.launches
tflops = flops / (per_launch_ms * 1e-3) / 1e12
print(f"{mode:5s} {per_launch_ms * 1e3:9.3f} us {tflops:8.1f} TFLOPS")
if __name__ == "__main__":
main()
|