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,108 @@
+"""Test F.linear with exact shapes from Z-Image-Turbo model on GPU.
+This isolates whether the GPU matmul hangs for model-specific shapes."""
+import torch
+import torch.nn.functional as F
+import time
+import os
+
+os.environ['HSA_OVERRIDE_GFX_VERSION'] = '10.1.0'
+os.environ['HSA_ENABLE_SDMA'] = '0'
+os.environ['TORCHDYNAMO_DISABLE'] = '1'
+
+print(f"Device: {torch.cuda.get_device_name(0)}")
+print(f"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1024**3:.1f} GB" if hasattr(torch.cuda.get_device_properties(0), 'total_mem') else "")
+
+# Weight shapes from the model (logged as CUDA#505-510)
+shapes = [
+    ("CBW505", (1024, 256)),
+    ("CBW506", (256, 1024)),
+    ("CBW507", (3840, 2560)),
+    ("CBW508", (3840, 64)),
+    ("CBW509", (11520, 3840)),
+    ("CBW510", (3840, 3840)),
+]
+
+# Batch size for Lumina2 at 512x512: typically (1, seq_len, dim)
+# Sequence length ~ 1024 patches
+batch_seq = 1024
+
+print(f"\n=== Test 1: Individual F.linear with model shapes ===")
+for name, (out_feat, in_feat) in shapes:
+    print(f"\n{name}: Linear({in_feat} -> {out_feat})")
+    try:
+        # Create tensors
+        x = torch.randn(1, batch_seq, in_feat, dtype=torch.float16, device='cuda:0')
+        w = torch.randn(out_feat, in_feat, dtype=torch.float16, device='cuda:0')
+        b = torch.randn(out_feat, dtype=torch.float16, device='cuda:0')
+        torch.cuda.synchronize()
+        print(f"  Tensors created on GPU", flush=True)
+        
+        # F.linear = x @ w.T + b
+        t0 = time.time()
+        out = F.linear(x, w, b)
+        torch.cuda.synchronize()
+        dt = time.time() - t0
+        print(f"  F.linear done: {dt*1000:.1f}ms, out={list(out.shape)}", flush=True)
+        
+        del x, w, b, out
+    except Exception as e:
+        print(f"  ERROR: {e}", flush=True)
+
+print(f"\n=== Test 2: Sequential F.linear chain (simulating forward pass) ===")
+# Simulate the model's sequence: transfer from CPU, compute, repeat
+x = torch.randn(1, batch_seq, 256, dtype=torch.float16, device='cuda:0')
+torch.cuda.synchronize()
+t0 = time.time()
+
+for i, (name, (out_feat, in_feat)) in enumerate(shapes):
+    # Simulate: CPU dequant → transfer → F.linear
+    w_cpu = torch.randn(out_feat, in_feat, dtype=torch.float16)
+    w_gpu = w_cpu.to('cuda:0', non_blocking=False)
+    
+    # Adjust input dimension if needed
+    if x.shape[-1] != in_feat:
+        x = torch.randn(1, batch_seq, in_feat, dtype=torch.float16, device='cuda:0')
+        torch.cuda.synchronize()
+    
+    out = F.linear(x, w_gpu)
+    torch.cuda.synchronize()
+    dt = time.time() - t0
+    print(f"  {name}: transfer+linear done, dt_total={dt:.2f}s, out={list(out.shape)}", flush=True)
+    x = out
+    del w_cpu, w_gpu
+
+print(f"\nTotal chain time: {time.time()-t0:.2f}s")
+
+print(f"\n=== Test 3: Manual softmax on GPU ===")
+x = torch.randn(1, 16, 1024, 1024, dtype=torch.float16, device='cuda:0')
+torch.cuda.synchronize()
+t0 = time.time()
+x_max = x.max(dim=-1, keepdim=True).values
+exp_x = torch.exp(x - x_max)
+result = exp_x / exp_x.sum(dim=-1, keepdim=True)
+torch.cuda.synchronize()
+print(f"  Manual softmax: {(time.time()-t0)*1000:.1f}ms, shape={list(result.shape)}")
+
+print(f"\n=== Test 4: Full manual SDPA on GPU ===")
+Q = torch.randn(1, 16, 1024, 64, dtype=torch.float16, device='cuda:0')
+K = torch.randn(1, 16, 1024, 64, dtype=torch.float16, device='cuda:0')
+V = torch.randn(1, 16, 1024, 64, dtype=torch.float16, device='cuda:0')
+torch.cuda.synchronize()
+t0 = time.time()
+scale = 64 ** -0.5
+attn = torch.matmul(Q, K.transpose(-2, -1)) * scale
+torch.cuda.synchronize()
+print(f"  Q@K^T: {(time.time()-t0)*1000:.1f}ms", flush=True)
+t1 = time.time()
+attn_max = attn.max(dim=-1, keepdim=True).values
+exp_attn = torch.exp(attn - attn_max)
+attn_weights = exp_attn / exp_attn.sum(dim=-1, keepdim=True)
+torch.cuda.synchronize()
+print(f"  Softmax: {(time.time()-t1)*1000:.1f}ms", flush=True)
+t2 = time.time()
+out = torch.matmul(attn_weights, V)
+torch.cuda.synchronize()
+print(f"  Attn@V: {(time.time()-t2)*1000:.1f}ms", flush=True)
+print(f"  Total SDPA: {(time.time()-t0)*1000:.1f}ms, out={list(out.shape)}")
+
+print("\n=== ALL TESTS PASSED ===")