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)