Uploaded sanitized BC250/ROCm Repository.

2026-08-20 00:45:43 +02:00
parent 7d2184f1e8
commit d7d22e93b3
678 changed files with 65963 additions and 1 deletions
@@ -0,0 +1,131 @@
+"""
+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)