Eliminate the redundant handling of index_select in get_splitfuse_attn_mask

When max_position_embeddings are large and max_seq_len is small, index_select will move a large amount of unused data, thereby increasing the utilization rate of the cpu, and the index_select operator executed on the cpu takes a very long time. this change can eliminate the redundant handling of index_select in get_splitfuse_attn_mask.

Signed-off-by: tt545571022 <tjl545571022@hotmail.com>

Author: tt545571022

vllm_ascend/attention/attention_mask.py

@@ -62,7 +62,7 @@ def get_attn_mask(self, max_seq_len: int, dtype: torch.dtype,
                       device: torch.device):
         self._update_attn_cache(max_seq_len, dtype)
         return self.attn_mask_cache[:max_seq_len, :max_seq_len].contiguous(
-        ).to(device)
+        ).to(device, non_blocking=True)
 
     def get_splitfuse_attn_mask(
         self,
@@ -79,11 +79,10 @@ def get_splitfuse_attn_mask(
         # FIXME: Currently the mask value of chunked-prefill situation and Prefill-Only situation
         # is not the same. Fix this in the future when kernel is ready.
         mask_scale_factor = AttentionMaskBuilder.get_mask_scale_factor(dtype)
-        attn_mask = torch.index_select(self.attn_mask_cache,
-                                       dim=0,
-                                       index=position)[:, :max_seq_len]
+        attn_mask = self.attn_mask_cache[position, :max_seq_len].to(
+            device, non_blocking=True)
         attn_mask *= mask_scale_factor
-        return attn_mask.contiguous().to(device, non_blocking=True)
+        return attn_mask
 
     def _update_attn_cache(self, seqlen: int, dtype: torch.dtype):
         if seqlen > self._seq_len_cached: