mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 09:14:35 +02:00
cleanup and update changelog
Signed-off-by: vladmandic <mandic00@live.com>
This commit is contained in:
vladmandic
+56
-61
@@ -1,18 +1,25 @@
|
||||
from typing import Dict, Any
|
||||
import time
|
||||
import warnings
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from typing import Dict, Any
|
||||
|
||||
warnings.filterwarnings("ignore", category=UserWarning)
|
||||
|
||||
warmup = 2
|
||||
repeats = 50
|
||||
dtypes = [torch.bfloat16] # , torch.float16]
|
||||
# if hasattr(torch, "float8_e4m3fn"):
|
||||
# dtypes.append(torch.float8_e4m3fn)
|
||||
|
||||
PROFILES = {
|
||||
backends = [
|
||||
"sdpa_math",
|
||||
"sdpa_mem_efficient",
|
||||
"sdpa_flash",
|
||||
"sdpa_all",
|
||||
"flex_attention",
|
||||
"xformers",
|
||||
"flash_attn",
|
||||
"sage_attn",
|
||||
]
|
||||
profiles = {
|
||||
"sdxl": {"l_q": 4096, "l_k": 4096, "h": 32, "d": 128},
|
||||
"flux.1": {"l_q": 16717, "l_k": 16717, "h": 24, "d": 128},
|
||||
"sd35": {"l_q": 16538, "l_k": 16538, "h": 24, "d": 128},
|
||||
@@ -30,19 +37,19 @@ def get_stats(reset: bool = False):
|
||||
m = torch.cuda.max_memory_allocated()
|
||||
t = time.perf_counter()
|
||||
return m / (1024 ** 2), t
|
||||
|
||||
|
||||
def print_gpu_info():
|
||||
if not torch.cuda.is_available():
|
||||
print("GPU: Not available")
|
||||
return
|
||||
|
||||
|
||||
device = torch.cuda.current_device()
|
||||
props = torch.cuda.get_device_properties(device)
|
||||
total_mem = props.total_memory / (1024**3)
|
||||
free_mem, _ = torch.cuda.mem_get_info(device)
|
||||
free_mem = free_mem / (1024**3)
|
||||
major, minor = torch.cuda.get_device_capability(device)
|
||||
|
||||
|
||||
print(f"gpu: {torch.cuda.get_device_name(device)}")
|
||||
print(f"vram: total={total_mem:.2f}GB free={free_mem:.2f}GB")
|
||||
print(f"cuda: capability={major}.{minor} version={torch.version.cuda}")
|
||||
@@ -56,11 +63,9 @@ def benchmark_attention(
|
||||
l_k: int = 4096,
|
||||
h: int = 32,
|
||||
d: int = 128,
|
||||
warmup: int = 10,
|
||||
repeats: int = 100
|
||||
) -> Dict[str, Any]:
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
|
||||
# Initialize tensors
|
||||
q = torch.randn(b, h, l_q, d, device=device, dtype=torch.float16 if dtype.is_floating_point and dtype.itemsize == 1 else dtype, requires_grad=False).to(dtype)
|
||||
k = torch.randn(b, h, l_k, d, device=device, dtype=torch.float16 if dtype.is_floating_point and dtype.itemsize == 1 else dtype, requires_grad=False).to(dtype)
|
||||
@@ -78,7 +83,7 @@ def benchmark_attention(
|
||||
try:
|
||||
if backend.startswith("sdpa_"):
|
||||
from torch.nn.attention import sdpa_kernel, SDPBackend
|
||||
sdp_type = backend[len("sdpa_"):]
|
||||
sdp_type = backend[len("sdpa_"):]
|
||||
# Map friendly names to new SDPA backends
|
||||
backend_map = {
|
||||
"math": [SDPBackend.MATH],
|
||||
@@ -88,21 +93,21 @@ def benchmark_attention(
|
||||
}
|
||||
if sdp_type not in backend_map:
|
||||
raise ValueError(f"Unknown SDPA type: {sdp_type}")
|
||||
|
||||
|
||||
results["version"] = torch.__version__
|
||||
|
||||
|
||||
with sdpa_kernel(backend_map[sdp_type]):
|
||||
# Warmup
|
||||
for _ in range(warmup):
|
||||
_ = F.scaled_dot_product_attention(q, k, v)
|
||||
|
||||
|
||||
start_mem, start_time = get_stats(True)
|
||||
|
||||
|
||||
for _ in range(repeats):
|
||||
_ = F.scaled_dot_product_attention(q, k, v)
|
||||
|
||||
|
||||
end_mem, end_time = get_stats()
|
||||
|
||||
|
||||
results["latency_ms"] = (end_time - start_time) / repeats * 1000
|
||||
results["memory_mb"] = end_mem - start_mem
|
||||
|
||||
@@ -113,17 +118,17 @@ def benchmark_attention(
|
||||
q_fa = q.transpose(1, 2)
|
||||
k_fa = k.transpose(1, 2)
|
||||
v_fa = v.transpose(1, 2)
|
||||
|
||||
|
||||
for _ in range(warmup):
|
||||
_ = flash_attn_func(q_fa, k_fa, v_fa)
|
||||
|
||||
|
||||
start_mem, start_time = get_stats(True)
|
||||
|
||||
|
||||
for _ in range(repeats):
|
||||
_ = flash_attn_func(q_fa, k_fa, v_fa)
|
||||
|
||||
|
||||
end_mem, end_time = get_stats()
|
||||
|
||||
|
||||
results["latency_ms"] = (end_time - start_time) / repeats * 1000
|
||||
results["memory_mb"] = end_mem - start_mem
|
||||
|
||||
@@ -135,17 +140,17 @@ def benchmark_attention(
|
||||
q_xf = q.transpose(1, 2)
|
||||
k_xf = k.transpose(1, 2)
|
||||
v_xf = v.transpose(1, 2)
|
||||
|
||||
|
||||
for _ in range(warmup):
|
||||
_ = memory_efficient_attention(q_xf, k_xf, v_xf)
|
||||
|
||||
|
||||
start_mem, start_time = get_stats(True)
|
||||
|
||||
|
||||
for _ in range(repeats):
|
||||
_ = memory_efficient_attention(q_xf, k_xf, v_xf)
|
||||
|
||||
|
||||
end_mem, end_time = get_stats()
|
||||
|
||||
|
||||
results["latency_ms"] = (end_time - start_time) / repeats * 1000
|
||||
results["memory_mb"] = end_mem - start_mem
|
||||
|
||||
@@ -158,18 +163,18 @@ def benchmark_attention(
|
||||
results["version"] = importlib.metadata.version("sageattention")
|
||||
except Exception:
|
||||
results["version"] = getattr(sageattention, "__version__", "N/A")
|
||||
|
||||
|
||||
# SageAttention expects (B, H, L, D) logic
|
||||
for _ in range(warmup):
|
||||
_ = sageattn(q, k, v)
|
||||
|
||||
|
||||
start_mem, start_time = get_stats(True)
|
||||
|
||||
|
||||
for _ in range(repeats):
|
||||
_ = sageattn(q, k, v)
|
||||
|
||||
|
||||
end_mem, end_time = get_stats()
|
||||
|
||||
|
||||
results["latency_ms"] = (end_time - start_time) / repeats * 1000
|
||||
results["memory_mb"] = end_mem - start_mem
|
||||
|
||||
@@ -179,62 +184,52 @@ def benchmark_attention(
|
||||
|
||||
# flex_attention requires torch.compile for performance
|
||||
flex_attention_compiled = torch.compile(flex_attention, dynamic=False)
|
||||
|
||||
|
||||
# Warmup (important to trigger compilation)
|
||||
for _ in range(warmup):
|
||||
_ = flex_attention_compiled(q, k, v)
|
||||
|
||||
|
||||
start_mem, start_time = get_stats(True)
|
||||
|
||||
|
||||
for _ in range(repeats):
|
||||
_ = flex_attention_compiled(q, k, v)
|
||||
|
||||
|
||||
end_mem, end_time = get_stats()
|
||||
|
||||
|
||||
results["latency_ms"] = (end_time - start_time) / repeats * 1000
|
||||
results["memory_mb"] = end_mem - start_mem
|
||||
except Exception as e:
|
||||
results["status"] = "fail"
|
||||
results["error"] = str(e)[:49]
|
||||
print(e)
|
||||
|
||||
return results
|
||||
|
||||
def main():
|
||||
backends = [
|
||||
"sdpa_math",
|
||||
"sdpa_mem_efficient",
|
||||
"sdpa_flash",
|
||||
"flex_attention",
|
||||
"xformers",
|
||||
"flash_attn",
|
||||
"sage_attn",
|
||||
]
|
||||
|
||||
def main():
|
||||
|
||||
all_results = []
|
||||
|
||||
print_gpu_info()
|
||||
print(f'config: warmup={warmup} repeats={repeats} dtypes={dtypes}')
|
||||
for name, config in PROFILES.items():
|
||||
print(f'config: warmup={warmup} repeats={repeats} dtypes={dtypes}')
|
||||
for name, config in profiles.items():
|
||||
print(f"profile: {name} (L_q={config['l_q']}, L_k={config['l_k']}, H={config['h']}, D={config['d']})")
|
||||
for dtype in dtypes:
|
||||
print(f" dtype: {dtype}")
|
||||
print(f" {'backend':<20} | {'version':<12} | {'status':<8} | {'latency':<10} | {'memory':<12} | ")
|
||||
for backend in backends:
|
||||
res = benchmark_attention(
|
||||
backend,
|
||||
dtype,
|
||||
l_q=config["l_q"],
|
||||
l_k=config["l_k"],
|
||||
h=config["h"],
|
||||
d=config["d"],
|
||||
warmup=warmup,
|
||||
repeats=repeats
|
||||
backend,
|
||||
dtype,
|
||||
l_q=config["l_q"],
|
||||
l_k=config["l_k"],
|
||||
h=config["h"],
|
||||
d=config["d"],
|
||||
)
|
||||
all_results.append(res)
|
||||
|
||||
|
||||
latency = f"{res['latency_ms']:.4f} ms"
|
||||
memory = f"{res['memory_mb']:.2f} MB"
|
||||
|
||||
|
||||
print(f" {res['backend']:<20} | {res['version']:<12} | {res['status']:<8} | {latency:<10} | {memory:<12} | {res['error']}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user