[Refactor] add fia_v3 attention & remove other attention operator.

Signed-off-by: weijinqian_v1 <[email protected]>

Author: weijinqian_v1

vllm_ascend/attention/attention_v1.py

@@ -281,18 +281,18 @@ def build_for_graph_capture(
 class AscendAttentionBackendImpl(AttentionImpl):
 
     def __init__(
-            self,
-            num_heads: int,
-            head_size: int,
-            scale: float,
-            num_kv_heads: int,
-            alibi_slopes: Optional[List[float]],
-            sliding_window: Optional[int],
-            kv_cache_dtype: str,
-            logits_soft_cap: Optional[float],
-            attn_type: str,
-            kv_sharing_target_layer_name: Optional[str],
-            **kwargs,
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: Optional[List[float]],
+        sliding_window: Optional[int],
+        kv_cache_dtype: str,
+        logits_soft_cap: Optional[float],
+        attn_type: str,
+        kv_sharing_target_layer_name: Optional[str],
+        **kwargs,
     ) -> None:
         self.num_heads = num_heads
         self.head_size = head_size
@@ -313,11 +313,8 @@ def __init__(
         self.key_cache = None
         self.value_cache = None
 
-    def _forward_prefill(self,
-                         query: torch.Tensor,
-                         key: torch.Tensor,
-                         value: torch.Tensor,
-                         attn_metadata: AscendMetadata,
+    def _forward_prefill(self, query: torch.Tensor, key: torch.Tensor,
+                         value: torch.Tensor, attn_metadata: AscendMetadata,
                          output: torch.Tensor):
         if attn_metadata.attn_state == AscendAttentionState.PrefillNoCache:
             block_size = 128
@@ -365,18 +362,19 @@ def _forward_prefill(self,
             sparse_mode=3,
         )
 
-        attn_output = attn_output.view(num_tokens, self.num_heads, self.head_size)
+        attn_output = attn_output.view(num_tokens, self.num_heads,
+                                       self.head_size)
         output[:num_tokens] = attn_output[:num_tokens]
         return output
 
     def _forward_decode_only(
-            self,
-            query: torch.Tensor,
-            attn_metadata: AscendMetadata,
-            output: torch.Tensor,
+        self,
+        query: torch.Tensor,
+        attn_metadata: AscendMetadata,
+        output: torch.Tensor,
     ) -> torch.Tensor:
         if self.sliding_window is not None and attn_metadata.seq_lens.shape[
-            0] == query.size(0):
+                0] == query.size(0):
             batch_size = attn_metadata.seq_lens.shape[0]
             block_size = 128
             query = query.view(batch_size, 1, self.num_heads * self.head_size)
@@ -470,12 +468,12 @@ def _forward_decode_only(
         return output
 
     def _forward_encode(
-            self,
-            query: torch.Tensor,
-            key: torch.Tensor,
-            value: torch.Tensor,
-            attn_metadata: AscendMetadata,
-            output: torch.Tensor,
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: AscendMetadata,
+        output: torch.Tensor,
     ) -> torch.Tensor:
         cum_seq_len = attn_metadata.query_start_loc[1:].tolist()
         output = torch_npu.npu_fusion_attention(
@@ -495,15 +493,15 @@ def _forward_encode(
         return output
 
     def forward(
-            self,
-            layer: AttentionLayer,
-            query: torch.Tensor,
-            key: torch.Tensor,
-            value: torch.Tensor,
-            kv_cache: Tuple[torch.Tensor],
-            attn_metadata: AscendMetadata,
-            output: Optional[torch.Tensor] = None,
-            trace_flag: bool = True,
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: Tuple[torch.Tensor],
+        attn_metadata: AscendMetadata,
+        output: Optional[torch.Tensor] = None,
+        trace_flag: bool = True,
     ) -> torch.Tensor:
         """Forward pass with Ascend attention.
         Args:
@@ -546,7 +544,7 @@ def forward(
 
         if attn_metadata is None:
             return output.view(num_tokens, self.hidden_size).fill_(0)
-        # ori_output = output
+
         if hasattr(layer, 'quant_method') and use_kv_cache_int8:
             output = layer.quant_method.apply(layer, query, key, value,
                                               kv_cache, attn_metadata,
@@ -563,7 +561,8 @@ def forward(
 
         if self.attn_type == AttentionType.ENCODER_ONLY:
             ori_output = output
-            output = self._forward_encode(query, key, value, attn_metadata, output)
+            output = self._forward_encode(query, key, value, attn_metadata,
+                                          output)
             ori_output[:num_tokens, :, :] = output[:num_tokens, :, :]
             return ori_output.view(num_tokens, self.hidden_size)
 
@@ -582,7 +581,8 @@ def forward(
         if attn_metadata.attn_state == AscendAttentionState.DecodeOnly:
             output = self._forward_decode_only(query, attn_metadata, output)
         else:
-            output = self._forward_prefill(query, key, value, attn_metadata, output)
+            output = self._forward_prefill(query, key, value, attn_metadata,
+                                           output)
 
         return output.view(num_tokens, self.hidden_size)