| """Sweep tile configs for the GEMM / GEGLU / attention kernels at the shapes laya actually hits.""" |
| import torch, sys, time, itertools, json |
| sys.path.insert(0, "kernels"); import tl_kernels as K |
| dev = "cuda" |
| def bench(f, n=20): |
| for _ in range(3): f() |
| torch.cuda.synchronize(); t = time.perf_counter() |
| for _ in range(n): f() |
| torch.cuda.synchronize(); return (time.perf_counter() - t) / n * 1e3 |
| res = {"gemm": {}, "geglu": {}, "attn": {}} |
| Ms = [128, 1024, 8192, 28672] |
| cfgs = [(bm, bn, th) for bm in (64, 128) for bn in (64, 128, 256) for th in (128, 256)] |
| for (N, Kd) in [(2304, 768), (768, 768), (768, 1152)]: |
| W = torch.randn(N, Kd, device=dev, dtype=torch.bfloat16) * 0.02; b = torch.zeros(N, device=dev) |
| for M in Ms: |
| A = torch.randn(M, Kd, device=dev, dtype=torch.bfloat16); C = torch.empty(M, N, device=dev, dtype=torch.bfloat16) |
| tref = bench(lambda: torch.nn.functional.linear(A, W)) |
| best = None |
| for bm, bn, th in cfgs: |
| try: |
| k = K.gemm_kernel(N, Kd, bm=bm, bn=bn, bk=64, stages=3, threads=th); ms = bench(lambda: k(A, W, b, C)) |
| if best is None or ms < best[0]: best = (ms, (bm, bn, th)) |
| except Exception as e: print("fail", N, Kd, M, (bm, bn, th), str(e)[:60], flush=True) |
| tf = 2 * M * N * Kd / best[0] / 1e9 |
| print(f"gemm N={N} K={Kd} M={M}: best {best[1]} {best[0]:.3f}ms ({tf:.1f} TFLOPS) cublas {tref:.3f}ms (default 64x128x128thr)", flush=True) |
| res["gemm"][f"{N},{Kd},{M}"] = best |
| for M in Ms: |
| A = torch.randn(M, 768, device=dev, dtype=torch.bfloat16); Wi = torch.randn(2304, 768, device=dev, dtype=torch.bfloat16) * 0.02; C = torch.empty(M, 1152, device=dev, dtype=torch.bfloat16) |
| best = None |
| for bm, bn, th in cfgs: |
| if bn > 128: continue |
| try: |
| k = K.gemm_geglu_kernel(1152, 768, bm=bm, bn=bn, bk=64, stages=3, threads=th); ms = bench(lambda: k(A, Wi, C)) |
| if best is None or ms < best[0]: best = (ms, (bm, bn, th)) |
| except Exception as e: print("fail geglu", M, (bm, bn, th), str(e)[:60], flush=True) |
| print(f"geglu M={M}: best {best[1]} {best[0]:.3f}ms ({2*M*2304*768/best[0]/1e9:.1f} TFLOPS)", flush=True) |
| res["geglu"][str(M)] = best |
| H, Dh = 12, 64 |
| for (B, L) in [(32, 1024), (32, 128), (4, 128)]: |
| qkv = torch.randn(B, L, 3, H, Dh, device=dev, dtype=torch.bfloat16); lens = torch.full((B,), L - 3, device=dev, dtype=torch.int32); O = torch.empty(B, L, H * Dh, device=dev, dtype=torch.bfloat16) |
| for window in (0, 65): |
| best = None |
| for bm, bn, st, th in itertools.product((64, 128), (64, 128), (1, 2), (128, 256)): |
| try: |
| k = K.attn_kernel(B, L, H, Dh, window=window, bm=bm, bn=bn, stages=st, threads=th); ms = bench(lambda: k(qkv, lens, O)) |
| if best is None or ms < best[0]: best = (ms, (bm, bn, st, th)) |
| except Exception as e: print("fail attn", (B, L, window), (bm, bn, st, th), str(e)[:60], flush=True) |
| fl = 4 * B * H * L * L * Dh if window == 0 else 4 * B * H * L * (2 * 65 + 1) * Dh |
| print(f"attn B={B} L={L} window={window}: best {best[1]} {best[0]:.3f}ms ({fl/best[0]/1e9:.1f} TFLOPS)", flush=True) |
| res["attn"][f"{B},{L},{window}"] = best |
| json.dump(res, open("kernels/tune_results.json", "w"), indent=1) |
| print("DONE") |
|
|
