ViperEkura
01eacbde51
- parameterize warp tile (WarpM/WarpN) in Fp8GemmTraits; MMA loops, fragment arrays and epilogue scale with kMt/kNt instead of the fixed 64x32/4x4, enabling cuBLAS-style 64x64 CTAs of 32x32 warps - dispatch by output tiling (grid-searched via csrc/tests/fp8_sweep.cu): fewer than 48 output tiles take 64x64/32x32 with a lean ring (4 CTAs/SM fill the wave-quantization gap: 512^3 goes 16 -> 64 CTAs); larger shapes keep 128x128 with the kStages+1 ring - kStages+1 canonic ring rotation drops the post-compute barrier on the congruous path (one __syncthreads per k-tile); LeanRing keeps the kStages ring for the small CTA; direct-crosswise operands always rotate kStages+1 (their prefetch issues right after barrier 1 and would race a lean ring - caught by the pure C layout suite) - stage the bf16 epilogue through the reclaimed operand smem: swizzled scatter + barrier + coalesced 16B copy-out replaces 8 disjoint 16B per-warp segments (~50% write efficiency before) - hoist per-lane ldmatrix swizzle offsets out of the mainloop (stage-relative table + ring-base add) so the innermost loop stops recomputing IMAD/LOP3 address chains - bypass the torch.library dispatch for real CUDA tensors in quantize/mm_fp8 wrappers (~5us/call, ~40% of a 512-wide call's wall time); fake/subclass tensors keep the custom_op route vs the previous kernel + python path, wall clock on NT squares: 512^3 52 -> 13us (4.0x, 5.2 -> 20.5 TF, now 1.36x cuBLAS _scaled_mm), 1024^3 1.05x, 2048^3 1.02x (46.9 -> 48.2 TF kernel-only); correctness: 4 layouts x 6 shapes pure C suite PASS, 588 pytest PASS
89 lines
2.7 KiB
Python
89 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
"""FP8 GEMM micro-benchmark: AstrAI kernel vs cuBLAS (torch._scaled_mm).
|
|
|
|
Sizes 512-8192, forward (NT) layout: x8[M,K] @ w8[N,K]^T -> bf16.
|
|
cuBLAS reference uses the same pre-quantized fp8 operands and the same
|
|
combined scale, so the comparison isolates the GEMM loop itself.
|
|
|
|
python scripts/tools/bench_fp8_gemm.py --sizes 512 1024 2048
|
|
"""
|
|
|
|
import argparse
|
|
import sys
|
|
|
|
import torch
|
|
|
|
sys.path.insert(0, ".")
|
|
|
|
from astrai.extension import loader
|
|
|
|
|
|
def bench(fn, iters=50, warmup=10):
|
|
for _ in range(warmup):
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
start = torch.cuda.Event(enable_timing=True)
|
|
end = torch.cuda.Event(enable_timing=True)
|
|
start.record()
|
|
for _ in range(iters):
|
|
fn()
|
|
end.record()
|
|
torch.cuda.synchronize()
|
|
return start.elapsed_time(end) / iters # ms
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--sizes", type=int, nargs="+", default=[512, 1024, 2048, 4096, 8192]
|
|
)
|
|
parser.add_argument(
|
|
"--rect", action="store_true", help="also bench M=4096xN=1024 style rectangles"
|
|
)
|
|
parser.add_argument("--iters", type=int, default=50)
|
|
args = parser.parse_args()
|
|
|
|
assert loader.is_available("fp8_ops"), "fp8_ops extension not built"
|
|
from astrai.extension.ops.fp8 import mm_fp8
|
|
|
|
dev = torch.device("cuda")
|
|
torch.manual_seed(0)
|
|
|
|
shapes = [(s, s, s) for s in args.sizes]
|
|
if args.rect:
|
|
shapes += [(4096, 1024, 4096), (8192, 4096, 8192), (2048, 8192, 2048)]
|
|
|
|
print(
|
|
f"{'M':>6} {'N':>6} {'K':>6} | {'ours(ms)':>9} {'TFLOPs':>7} | "
|
|
f"{'cublas(ms)':>10} {'TFLOPs':>7} | {'ratio':>6}"
|
|
)
|
|
print("-" * 72)
|
|
for m, n, k in shapes:
|
|
x8 = (torch.randn(m, k, device=dev) * 0.05).to(torch.float8_e4m3fn)
|
|
w8 = (torch.randn(n, k, device=dev) * 0.05).to(torch.float8_e4m3fn)
|
|
scale = torch.ones(1, device=dev, dtype=torch.float32)
|
|
|
|
# ours: NT (x8 @ w8^T, LayoutB=ColMajor = weight layout)
|
|
t_ours = bench(lambda: mm_fp8(x8, w8, scale, trans_b=True), iters=args.iters)
|
|
# cuBLAS: _scaled_mm needs A row-major, B column-major (= w8.t())
|
|
wt = w8.t()
|
|
sa = torch.ones(1, device=dev)
|
|
sb = torch.ones(1, device=dev)
|
|
|
|
def cublas():
|
|
return torch._scaled_mm(x8, wt, sa, sb, out_dtype=torch.bfloat16)
|
|
|
|
t_cublas = bench(cublas, iters=args.iters)
|
|
flops = 2.0 * m * n * k
|
|
tf_ours = flops / (t_ours * 1e-3) / 1e12
|
|
tf_cublas = flops / (t_cublas * 1e-3) / 1e12
|
|
print(
|
|
f"{m:>6} {n:>6} {k:>6} | {t_ours:>9.3f} {tf_ours:>7.1f} | "
|
|
f"{t_cublas:>10.3f} {tf_cublas:>7.1f} | "
|
|
f"{tf_ours / tf_cublas:>5.0%}"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|