fix: expand shape of attn_mask (#10)

Author: JackWeiw

lmdeploy/pytorch/engine/devices/ascend.py

@@ -17,7 +17,7 @@ def update_step_context(cls, step_context):
             single_attention_mask = torch.logical_not(
                 torch.tril(
                     torch.ones(step_context.q_seq_length[i],
-                               (step_context.kv_seq_length[i] + 31) & (~31),
+                               step_context.block_offsets.shape[1] * block_size,
                                dtype=torch.bool).cuda(),
                     diagonal=step_context.kv_seq_length[i] -
                     step_context.q_seq_length[i],