From cad00fae2625fb10865cd5e94f699a6f7a6b84ac Mon Sep 17 00:00:00 2001 From: Jeff Rasley Date: Fri, 4 Sep 2020 10:57:57 -0700 Subject: [PATCH] Fix datatype issue with sparse attention softmax 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. --- deepspeed/ops/sparse_attention/softmax.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/ops/sparse_attention/softmax.py b/deepspeed/ops/sparse_attention/softmax.py index 814e9cc50e19..41267298a0a4 100644 --- a/deepspeed/ops/sparse_attention/softmax.py +++ b/deepspeed/ops/sparse_attention/softmax.py @@ -234,7 +234,7 @@ def __init__(self, layout, block, bench=False): bench: optional: set if you want to do benchmarking """ - self.num_blocks = layout.sum() + self.num_blocks = layout.sum().item() self.spdims = layout.shape self.layout = layout self.block = block