CUDA errors on kernels after block sparse Triton ops

I'm experimenting with block sparse Linear layers and have been getting various CUDA errors when adding triton ops.

On triton==2.0.0.dev20221105, the following causes RuntimeError: CUDA error: CUBLAS_STATUS_NOT_INITIALIZED when calling cublasCreate(handle)`` when running F.linear on the last Pytorch Linear layer.

import torch
import torch.nn as nn
import triton

def sparsify_tensor(x, mask, block):
    ret = torch.empty((x.size(0), mask.sum(), block, block), dtype=x.dtype, device=x.device)
    for idx, (h, i, j) in enumerate(zip(*mask.nonzero(as_tuple=True))):
        ret[:, idx, :, :] = x[:, h, i * block:(i + 1) * block, j * block:(j + 1) * block]
    return ret

class BlockSparseLinear(nn.Module):
    def __init__(self, in_features: int, out_features: int, sparsity: float = 0.8):
        assert in_features % 32 == 0
        assert out_features % 32 == 0
        super().__init__()
        mask = (torch.rand((1, out_features //32, in_features //32)).uniform_() > sparsity).long()
        self.op = triton.ops.blocksparse.matmul(mask, 32, 'dsd', trans_a=False, trans_b=True, device='cuda')
        self.weight = torch.nn.Parameter(sparsify_tensor(torch.empty((1, 1, out_features , in_features)), mask, 32))

    def forward(self, input):
        input.unsqueeze_(2).unsqueeze_(3)
        return self.op(self.weight, input).squeeze_(3)

if __name__ == '__main__':
    model = nn.Sequential(
        BlockSparseLinear(512, 2048),
        nn.GELU(),
        BlockSparseLinear(2048, 512),
        nn.Linear(512, 10),
    )
    model.half().cuda()
    batch = torch.rand((512, 256, 512)).half().cuda()
    model(batch)

If I add a Linear layer before the triton ops, CuBLAS gets properly initialized.

Thanks for the help and let me know if you need more details!