Fix datatype issue with sparse attention softmax (#363)
Fixes a dataype issue with softmax where the number of blocks being sent to the Triton kernel source was a torch.Tensor but should have been a python integer. On some environments (e.g., conda) this resulted in triton not knowing how to serialize the input (and crashing in our tests). Once switching to the correct datatype that triton expects this seems to have solved the issue.
Co-authored-by: Shaden Smith <Shaden.Smith@microsoft.com>
