99 lines
2.7 KiB
Python
99 lines
2.7 KiB
Python
#!/usr/bin/env python3
|
|
import csv
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
from itertools import product
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parent.parent
|
|
BUILD = ROOT / "kernels_op_7" / "build"
|
|
CSV_PATH = ROOT / "kernels_op_7" / "sweep.csv"
|
|
|
|
# Just BR x BC. 9 configs. BR=128 will OOM on smem (48KB cap), expected NOT_FOUND.
|
|
GRID = {
|
|
"br": [32, 64, 128],
|
|
"bc": [32, 64, 128],
|
|
}
|
|
|
|
# from Makefile
|
|
NVCC = "nvcc"
|
|
CFLAGS = "-O3 -arch=sm_86 -std=c++17 -rdc=true -g -lineinfo"
|
|
INCLUDES = "-I/usr/local/cuda/include -I/work/hdd/bche/hshah6/cutlass/include"
|
|
LIBS = "-L/usr/local/cuda/lib64 -lcublas -lcublasLt"
|
|
|
|
KERNEL_NAME = "flash_attention_forward_kernel"
|
|
|
|
|
|
def run(cmd):
|
|
return subprocess.run(cmd, shell=True, cwd=ROOT, capture_output=True, text=True)
|
|
|
|
|
|
def compile_one(tag, defs):
|
|
"""Two-step compile matching the Makefile."""
|
|
obj = BUILD / f"{tag}.o"
|
|
binary = BUILD / tag
|
|
r = run(f"{NVCC} {CFLAGS} {INCLUDES} {defs} -c test_gpt2.cu -o {obj}")
|
|
if r.returncode != 0:
|
|
return None, f"COMPILE FAILED:\n{r.stderr}"
|
|
r = run(f"{NVCC} {CFLAGS} -o {binary} {obj} {LIBS}")
|
|
if r.returncode != 0:
|
|
return None, f"LINK FAILED:\n{r.stderr}"
|
|
return binary, ""
|
|
|
|
|
|
def profile_one(binary, tag):
|
|
report = BUILD / tag
|
|
log = BUILD / f"{tag}.log"
|
|
cmd = (
|
|
f"nsys profile --stats=true --force-overwrite=true "
|
|
f"--cuda-event-trace=false -o {report} {binary}"
|
|
)
|
|
r = run(cmd)
|
|
log.write_text((r.stdout or "") + "\n" + (r.stderr or ""))
|
|
return log
|
|
|
|
|
|
def parse_avg_ns(log_path, kernel_name):
|
|
pattern = re.compile(rf"^\s*\d.*\b{re.escape(kernel_name)}\b")
|
|
for line in log_path.read_text().splitlines():
|
|
if pattern.search(line):
|
|
cols = line.split()
|
|
if len(cols) >= 4:
|
|
return cols[3]
|
|
return "NOT_FOUND"
|
|
|
|
|
|
def main():
|
|
BUILD.mkdir(parents=True, exist_ok=True)
|
|
configs = list(product(GRID["br"], GRID["bc"]))
|
|
|
|
first_fail_printed = False
|
|
with CSV_PATH.open("w", newline="") as fp:
|
|
out = csv.writer(fp)
|
|
out.writerow(["BR", "BC", "fa2_avg_ns"])
|
|
|
|
for i, (br, bc) in enumerate(configs, 1):
|
|
tag = f"br{br}_bc{bc}"
|
|
|
|
defs = f"-DSWEEP_BR={br} -DSWEEP_BC={bc}"
|
|
binary, err = compile_one(tag, defs)
|
|
|
|
if binary is None:
|
|
if not first_fail_printed:
|
|
print(err)
|
|
first_fail_printed = True
|
|
out.writerow([br, bc, "BUILD_FAIL"])
|
|
fp.flush()
|
|
continue
|
|
|
|
log = profile_one(binary, tag)
|
|
out.writerow([br, bc, parse_avg_ns(log, KERNEL_NAME)])
|
|
fp.flush()
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|