import torch, time
dev="cuda:1"
torch.cuda.set_device(dev)
from torch.nn.attention import SDPBackend, sdpa_kernel
S, H, D = 14100, 56, 128     # real H3 shape @ 480x832x124f
for dt in (torch.float16,):
    q=k=v=torch.randn(1,H,S,D,device=dev,dtype=dt)
    for be,nm in ((SDPBackend.EFFICIENT_ATTENTION,"EFFICIENT"),):
        try:
            with sdpa_kernel(be):
                torch.cuda.synchronize(); t0=time.time()
                o=torch.nn.functional.scaled_dot_product_attention(q,k,v)
                torch.cuda.synchronize()
                print(f"[OK ] SDPA {nm} fp16 S={S} H={H} D={D}: {o.shape} {time.time()-t0:.2f}s peak={torch.cuda.max_memory_allocated(dev)/2**30:.2f}G")
        except Exception as e:
            print(f"[ERR] SDPA {nm} fp16 D={D}: {type(e).__name__}: {str(e)[:200]}")
    del q,k,v
    torch.cuda.empty_cache(); torch.cuda.reset_peak_memory_stats(dev)

# fp16 GEMM throughput on this card (H3 block-ish shape)
a=torch.randn(14100,5376,device=dev,dtype=torch.float16); b=torch.randn(5376,28672,device=dev,dtype=torch.float16)
for _ in range(2): c=a@b
torch.cuda.synchronize(); t0=time.time()
for _ in range(5): c=a@b
torch.cuda.synchronize()
fl=2*14100*5376*28672*5
print(f"fp16 GEMM effective: {fl/(time.time()-t0)/1e12:.1f} TFLOPS")
# fp16 overflow sanity on a deep-ish accumulation
x=torch.randn(4096,4096,device=dev,dtype=torch.float16)*4
y=x
for i in range(20): y=torch.nn.functional.layer_norm(y@x/64,(4096,))
print("fp16 deep-chain finite:", bool(torch.isfinite(y).all()))