Fabian
132 lines
4.7 KiB
Python
132 lines
4.7 KiB
Python
"""
|
|
BC-250 (gfx1010) Softmax Fix Test
|
|
Tests if manual softmax works where fused kernel crashes with VGPR overflow.
|
|
"""
|
|
import torch
|
|
import time
|
|
import os
|
|
import gc
|
|
|
|
print(f'Device: {torch.cuda.get_device_name(0)}')
|
|
print()
|
|
|
|
def safe_softmax(x, dim=-1):
|
|
"""Manual softmax using elementwise ops - avoids fused kernel VGPR overflow."""
|
|
x_max = x.max(dim=dim, keepdim=True).values
|
|
exp_x = torch.exp(x - x_max)
|
|
return exp_x / exp_x.sum(dim=dim, keepdim=True)
|
|
|
|
# Test 1: Does safe_softmax work where torch.softmax crashes?
|
|
print('=== Test 1: safe_softmax on [1, 24, 512, 4096] ===', flush=True)
|
|
try:
|
|
x = torch.randn(1, 24, 512, 4096, device='cuda', dtype=torch.float32)
|
|
torch.cuda.synchronize()
|
|
t = time.time()
|
|
y = safe_softmax(x, dim=-1)
|
|
torch.cuda.synchronize()
|
|
print(f' OK: {time.time()-t:.3f}s', flush=True)
|
|
# Verify correctness
|
|
assert y.min() >= 0, "Negative values in softmax"
|
|
assert abs(y.sum(dim=-1).mean().item() - 1.0) < 0.001, "Softmax doesn't sum to 1"
|
|
print(f' Correctness: OK (sum={y.sum(dim=-1).mean().item():.6f})', flush=True)
|
|
del x, y
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
print(f' FAIL: {e}', flush=True)
|
|
|
|
# Test 2: Full attention pipeline with safe_softmax
|
|
print('=== Test 2: Full attention with safe_softmax (24h, 4096 seq) ===', flush=True)
|
|
try:
|
|
heads, seq, hd = 24, 4096, 128
|
|
chunk = 512
|
|
|
|
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)
|
|
|
|
scale = hd ** -0.5
|
|
|
|
torch.cuda.synchronize()
|
|
t = time.time()
|
|
for i in range(0, seq, chunk):
|
|
q_chunk = q[:, :, i:i+chunk, :] * scale
|
|
scores = torch.matmul(q_chunk, k.transpose(-2, -1))
|
|
attn = safe_softmax(scores, dim=-1)
|
|
out[:, :, i:i+chunk, :] = torch.matmul(attn, v)
|
|
del scores, attn
|
|
torch.cuda.synchronize()
|
|
elapsed = time.time() - t
|
|
print(f' OK: {elapsed:.3f}s', flush=True)
|
|
del q, k, v, out
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
print(f' FAIL: {e}', flush=True)
|
|
|
|
# Test 3: Also verify torch.softmax on SMALL tensors still works
|
|
print('=== Test 3: torch.softmax on small tensor [1, 8, 256, 256] ===', flush=True)
|
|
try:
|
|
x = torch.randn(1, 8, 256, 256, device='cuda', dtype=torch.float32)
|
|
torch.cuda.synchronize()
|
|
y = torch.softmax(x, dim=-1)
|
|
torch.cuda.synchronize()
|
|
print(f' OK (native softmax works on small tensors)', flush=True)
|
|
del x, y
|
|
except Exception as e:
|
|
print(f' FAIL: {e}', flush=True)
|
|
|
|
# Test 4: Find the threshold where torch.softmax breaks
|
|
print('=== Test 4: Find native softmax failure threshold ===', flush=True)
|
|
for last_dim in [256, 512, 1024, 2048, 4096]:
|
|
try:
|
|
x = torch.randn(1, 24, 512, last_dim, device='cuda', dtype=torch.float32)
|
|
torch.cuda.synchronize()
|
|
y = torch.softmax(x, dim=-1)
|
|
torch.cuda.synchronize()
|
|
print(f' dim={last_dim}: OK', flush=True)
|
|
del x, y
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
except Exception as e:
|
|
print(f' dim={last_dim}: FAIL ({e})', flush=True)
|
|
gc.collect()
|
|
torch.cuda.empty_cache()
|
|
break
|
|
|
|
# Test 5: Test SDPA (scaled_dot_product_attention) - what ComfyUI actually uses
|
|
print('=== Test 5: F.scaled_dot_product_attention ===', flush=True)
|
|
try:
|
|
q = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
|
|
k = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
|
|
v = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
|
|
torch.cuda.synchronize()
|
|
t = time.time()
|
|
out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
|
torch.cuda.synchronize()
|
|
print(f' SDPA OK: {time.time()-t:.3f}s', flush=True)
|
|
del q, k, v, out
|
|
except Exception as e:
|
|
print(f' SDPA FAIL: {e}', flush=True)
|
|
|
|
# Test 5b: SDPA with smaller dims
|
|
print('=== Test 5b: SDPA smaller (8h, 1024 seq) ===', flush=True)
|
|
try:
|
|
q = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
|
|
k = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
|
|
v = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
|
|
torch.cuda.synchronize()
|
|
t = time.time()
|
|
out = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
|
torch.cuda.synchronize()
|
|
print(f' SDPA OK: {time.time()-t:.3f}s', flush=True)
|
|
del q, k, v, out
|
|
except Exception as e:
|
|
print(f' SDPA FAIL: {e}', flush=True)
|
|
|
|
print()
|
|
print('=== ALL DONE ===', flush=True)
|
|
torch.cuda.synchronize()
|
|
os._exit(0)
|