import torch
print("  device      ", torch.cuda.get_device_name(0))
print("  capability  ", torch.cuda.get_device_capability(0))
print("  arch_list   ", torch.cuda.get_arch_list())
a = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
b = torch.randn(4096, 4096, device="cuda", dtype=torch.bfloat16)
print("  matmul ok   ", torch.matmul(a, b).float().abs().mean().item() > 0)
import transformer_engine.pytorch as te
lin = te.Linear(1024, 1024, params_dtype=torch.bfloat16).cuda()
print("  TE Linear   ", tuple(lin(torch.randn(8, 1024, device="cuda", dtype=torch.bfloat16)).shape))
from mamba_ssm.ops.selective_scan_interface import selective_scan_fn
import causal_conv1d
from causal_conv1d import causal_conv1d_fn
x = torch.randn(2, 64, 128, device="cuda", dtype=torch.bfloat16)
w = torch.randn(64, 4, device="cuda", dtype=torch.bfloat16)
print("  causal_conv1d CUDA op", tuple(causal_conv1d_fn(x, w).shape))
print("GPU_GATE_PASS")