Fabian
96 lines
3.2 KiB
Python
96 lines
3.2 KiB
Python
import torch
|
|
import time
|
|
import os
|
|
import sys
|
|
import gc
|
|
|
|
print(f'Device: {torch.cuda.get_device_name(0)}')
|
|
print(f'Total Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')
|
|
print()
|
|
|
|
def mem_info():
|
|
alloc = torch.cuda.memory_allocated() / 1e6
|
|
reserved = torch.cuda.memory_reserved() / 1e6
|
|
return f'alloc={alloc:.0f}MB, reserved={reserved:.0f}MB'
|
|
|
|
def test_attention(heads, seq_len, head_dim, dtype=torch.float32):
|
|
label = f'Attn h={heads} s={seq_len} d={head_dim} {"fp32" if dtype==torch.float32 else "fp16"}'
|
|
print(f'=== {label} ===', flush=True)
|
|
|
|
# Calculate memory needed
|
|
attn_size = heads * seq_len * seq_len * (4 if dtype==torch.float32 else 2)
|
|
qk_size = 2 * heads * seq_len * head_dim * (4 if dtype==torch.float32 else 2)
|
|
total_est = (attn_size + qk_size) / 1e6
|
|
print(f' Est memory: {total_est:.0f}MB ({mem_info()})', flush=True)
|
|
|
|
try:
|
|
q = torch.randn(1, heads, seq_len, head_dim, device='cuda', dtype=dtype)
|
|
k = torch.randn(1, heads, seq_len, head_dim, device='cuda', dtype=dtype)
|
|
print(f' Q/K allocated ({mem_info()})', flush=True)
|
|
|
|
torch.cuda.synchronize()
|
|
t = time.time()
|
|
scores = torch.matmul(q, k.transpose(-2, -1))
|
|
torch.cuda.synchronize()
|
|
elapsed = time.time() - t
|
|
|
|
print(f' OK: {elapsed:.3f}s, scores shape={list(scores.shape)} ({mem_info()})', flush=True)
|
|
del q, k, scores
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return True
|
|
except Exception as e:
|
|
print(f' FAIL: {e}', flush=True)
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
return False
|
|
|
|
# Progressive attention scaling
|
|
test_attention(4, 512, 64)
|
|
test_attention(8, 1024, 64)
|
|
test_attention(8, 1024, 128)
|
|
test_attention(16, 2048, 128)
|
|
test_attention(24, 2048, 128)
|
|
test_attention(24, 4096, 128) # full Lumina2 scale!
|
|
|
|
# If full scale fails in fp32, try fp16
|
|
print()
|
|
print('=== fp16 ATTENTION TESTS ===', flush=True)
|
|
test_attention(24, 4096, 128, torch.float16)
|
|
|
|
# Test split attention approach (process in chunks)
|
|
print()
|
|
print('=== SPLIT ATTENTION (simulate ComfyUI split attn) ===', flush=True)
|
|
try:
|
|
heads = 24
|
|
seq = 4096
|
|
hd = 128
|
|
chunk = 512 # process 512 tokens at a time
|
|
|
|
q = torch.randn(1, heads, seq, hd, device='cuda', dtype=torch.float32)
|
|
k = torch.randn(1, heads, seq, hd, device='cuda', dtype=torch.float32)
|
|
v = torch.randn(1, heads, seq, hd, device='cuda', dtype=torch.float32)
|
|
out = torch.zeros(1, heads, seq, hd, device='cuda', dtype=torch.float32)
|
|
|
|
torch.cuda.synchronize()
|
|
t = time.time()
|
|
for i in range(0, seq, chunk):
|
|
q_chunk = q[:, :, i:i+chunk, :]
|
|
scores = torch.matmul(q_chunk, k.transpose(-2, -1))
|
|
attn = torch.softmax(scores, dim=-1)
|
|
out[:, :, i:i+chunk, :] = torch.matmul(attn, v)
|
|
del scores, attn
|
|
torch.cuda.synchronize()
|
|
elapsed = time.time() - t
|
|
print(f' Split attention OK: {elapsed:.3f}s', flush=True)
|
|
del q, k, v, out
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
print(f' Split attention FAIL: {e}', flush=True)
|
|
|
|
print()
|
|
print('=== DONE ===', flush=True)
|
|
torch.cuda.synchronize()
|
|
os._exit(0)
|