Shared Code

import torch
import platform
import os
import time

print("=" * 60)
print("LANIE GPU CAPABILITY PROBE")
print("=" * 60)

print("Python:", platform.python_version())
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())

if not torch.cuda.is_available():
    raise SystemExit("NO CUDA GPU")

device = torch.device("cuda")

props = torch.cuda.get_device_properties(0)

print("GPU:", torch.cuda.get_device_name(0))
print("CUDA runtime:", torch.version.cuda)
print("VRAM GB:", round(props.total_memory / 1024**3, 2))
print("Compute capability:", f"{props.major}.{props.minor}")
print("Multiprocessors:", props.multi_processor_count)

free, total = torch.cuda.mem_get_info()

print("Free VRAM GB:", round(free / 1024**3, 2))
print("Total VRAM GB:", round(total / 1024**3, 2))

# Real GPU matrix multiplication
n = 4096

a = torch.randn(n, n, device=device, dtype=torch.float16)
b = torch.randn(n, n, device=device, dtype=torch.float16)

torch.cuda.synchronize()
start = time.time()

c = a @ b

torch.cuda.synchronize()
elapsed = time.time() - start

print("4096x4096 FP16 matmul:", round(elapsed, 4), "sec")

allocated = torch.cuda.memory_allocated() / 1024**3
reserved = torch.cuda.memory_reserved() / 1024**3

print("Allocated GB:", round(allocated, 3))
print("Reserved GB:", round(reserved, 3))
print("Checksum:", float(c[0, 0]))

print("=" * 60)
print("REAL CUDA TEST COMPLETE")