Sparse attention error: IndexError: index 0 is out of bounds for dimension 1 with size 0

I am trying to use sparse attention, but I get this error:

  File "/home/ubuntu/anaconda3/envs/meshattn/lib/python3.6/site-packages/deepspeed/ops/sparse_attention/sparse_self_attention.py", line 152, in forward
    attn_output_weights = sparse_dot_sdd_nt(query, key)
  File "/home/ubuntu/anaconda3/envs/meshattn/lib/python3.6/site-packages/deepspeed/ops/sparse_attention/matmul.py", line 712, in __call__
    db_lut, db_num_locks, db_width, db_packs = self.make_lut(a.dtype, a.device)
  File "/home/ubuntu/anaconda3/envs/meshattn/lib/python3.6/site-packages/deepspeed/ops/sparse_attention/matmul.py", line 641, in make_lut
    da_lut, da_num_locks, da_width, da_packs = _sparse_matmul.make_dxx_lut(layout, block, step, True, device)
  File "/home/ubuntu/anaconda3/envs/meshattn/lib/python3.6/site-packages/deepspeed/ops/sparse_attention/matmul.py", line 284, in make_dxx_lut
    xincs[:, 0] -= (div - 1) * step
IndexError: index 0 is out of bounds for dimension 1 with size 0

on this line - https://github.com/microsoft/DeepSpeed/blob/e4139a7dc14ba0a1d9263e7f11a16b168a180fe8/deepspeed/ops/sparse_attention/sparse_self_attention.py#L152.

My query and key tensors look fine when I set a breakpoint before this line, and I'm not sure where I am going wrong

(Pdb) query.shape
torch.Size([3, 4, 1304, 64])
(Pdb) key.shape
torch.Size([3, 4, 1304, 64])
(Pdb) sparse_dot_sdd_nt(query, key)
Using /home/ubuntu/.cache/torch_extensions as PyTorch extensions root...
Emitting ninja build file /home/ubuntu/.cache/torch_extensions/sparse_attn/build.ninja...
Building extension module sparse_attn...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module sparse_attn...
Time to load sparse_attn op: 0.35572361946105957 seconds
*** IndexError: index 0 is out of bounds for dimension 1 with size 0
(Pdb) SparseSelfAttention.ops
{1304: (<deepspeed.ops.sparse_attention.matmul.MatMul object at 0x7f46a5733470>, <deepspeed.ops.sparse_attention.matmul.MatMul object at 0x7f46a573e0b8>, <deepspeed.ops.sparse_attention.softmax.Softmax object at 0x7f46a573e860>)}
(Pdb) tgt_len
1304
(Pdb) tgt_len not in SparseSelfAttention.ops
False
(Pdb) SparseSelfAttention.ops[tgt_len]
(<deepspeed.ops.sparse_attention.matmul.MatMul object at 0x7f46a5733470>, <deepspeed.ops.sparse_attention.matmul.MatMul object at 0x7f46a573e0b8>, <deepspeed.ops.sparse_attention.softmax.Softmax object at 0x7f46a573e860>)

Start at deepspeed/ops/sparse_attention/sparse_self_attention.py line 152 and trace into deepspeed/ops/sparse_attention/matmul.py, especially make_dxx_lut at line 284. Inspect how the reported sequence length of 1304 produces the empty dimension before xincs[:, 0] is accessed. Done means the reported sparse-attention shapes no longer raise IndexError.