Function get_cublas

src/fused_mm_sampling/tl_matmul.py:190–195

()

def get_cublas():
    from triton._C.libtriton import nvidia

    device_workspace = torch.empty(32 * 1024 * 1024, device="cuda", dtype=torch.uint8)
    cublas = nvidia.cublas.CublasLt(device_workspace)
    return cublas