fix cann version barrier for npu_fused_infer_attention_score

Signed-off-by: MengqingCao <cmq0113@163.com>

Author: MengqingCao

vllm_ascend/attention/attention_mask.py

@@ -50,7 +50,7 @@ def __init__(
         self._seq_len_cached = attn_mask.shape[0]
         self.attn_mask_cache = attn_mask
         self.device = device
-        if torch.version.cann.startswith("8.3"):
+        if torch.version.cann > "8.3.RC1.alpha002":
             assigned_mask_dim = 2048
             self.chunked_prefill_attn_mask = torch.triu(
                 torch.ones(assigned_mask_dim, assigned_mask_dim),
@@ -82,7 +82,7 @@ def get_splitfuse_attn_mask(
         dtype: torch.dtype = None,
         device: torch.device = None,
     ) -> torch.Tensor:
-        if torch.version.cann.startswith("8.3"):
+        if torch.version.cann > "8.3.RC1.alpha002":
             return self.chunked_prefill_attn_mask
         else:
             if dtype not in [torch.float16, torch.bfloat16]:

vllm_ascend/attention/attention_v1.py

@@ -483,7 +483,7 @@ def _forward_v1_style(
             attn_metadata.seq_lens = \
                 attn_metadata.seq_lens.to(device=query.device)
 
-        if torch.version.cann.startswith("8.3"):
+        if torch.version.cann > "8.3.RC1.alpha002":
             # TODO:The npu_fused_infer_attention_score op is planned to
             # be utilized in a wider range in upcoming versions.
             num_block, block_size, _, _ = self.key_cache.shape  # type: ignore
@@ -613,7 +613,7 @@ def forward(
                                                    output)
             # Normal V1 situation.
             else:
-                if torch.version.cann.startswith("8.3"):
+                if torch.version.cann > "8.3.RC1.alpha002":
                     # npu_fused_infer_attention_score does not support cases
                     # where query.shape[0] != attn_metadata.query_start_loc[-1].
                     # Thus we need unpad it here.

vllm_ascend/worker/model_runner_v1.py

@@ -314,7 +314,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
             use_mla=self.model_config.use_mla,
         )
 
-        if torch.version.cann.startswith("8.3"):
+        if torch.version.cann > "8.3.RC1.alpha002":
             self.attn_mask_builder = AttentionMaskBuilder(
                 self.scheduler_config.max_num_batched_tokens, self.dtype,
                 self.device)
@@ -881,7 +881,7 @@ def _make_attention_mask(self, seq_lens, position,
                              attn_state) -> torch.Tensor:
         # Chunk Prefill situation.
         if attn_state == AscendAttentionState.ChunkedPrefill and not self.vllm_config.model_config.use_mla:
-            if torch.version.cann.startswith("8.3"):
+            if torch.version.cann > "8.3.RC1.alpha002":
                 return self.attn_mask_builder.get_splitfuse_attn_mask()
             else:
                 return self.attn_mask_builder.get_splitfuse_attn_mask(