[NPU] update Mixed chunk op to FIA (#15518)

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -87,9 +87,8 @@ class AscendAttnMaskBuilder:
         self.mtp_mask = self.generate_mask_flag(mtp_mask_len).to(self.device)
 
         # Initialize mixed chunk mask cache
-        mixed_chunk_cache_len = 8192
+        mixed_mask_len = 2048
-        self.mix_mask_cache = self.generate_attn_mask(mixed_chunk_cache_len, "mix")
+        self.mixed_chunk_attn_mask = self.get_splitfuse_attn_mask(mixed_mask_len)
-        self.mix_seq_len_cached = self.mix_mask_cache.shape[0]
 
         if use_mla:
             # Initialize RingMla mask
@@ -183,48 +182,19 @@ class AscendAttnMaskBuilder:
     def get_splitfuse_attn_mask(
         self,
         seq_lens: torch.Tensor = None,
-        position: torch.Tensor = None,
-        dtype: torch.dtype = None,
-        device: torch.device = None,
     ) -> torch.Tensor:
         """
         Generate a splitfuse attention mask.
 
         :param seq_lens: Sequence lengths.
-        :param position: Position indices for the mask.
-        :param dtype: Data type of the mask tensor.
-        :param device: Device to run the model on.
         :return: A tensor representing the splitfuse attention mask.
         """
-        if dtype not in [torch.float16, torch.bfloat16]:
+        attn_mask = (
-            raise ValueError("splitfuse_attn_mask now only supports bf16 and fp16")
+            torch.triu(torch.ones(seq_lens, seq_lens), diagonal=1)
-        max_seq_len = max(seq_lens, default=0)
+            .to(torch.int8)
-        self.mix_mask_cache, self.mix_seq_len_cached = self.update_attn_cache(
+            .to(self.device)
-            max_seq_len, self.mix_mask_cache, self.mix_seq_len_cached, dtype, mode="mix"
         )
-        attn_mask = torch.index_select(self.mix_mask_cache, dim=0, index=position)[
+        return attn_mask
-            :, :max_seq_len
-        ]
-        return attn_mask.contiguous().to(device, non_blocking=True)
-
-    def update_mask(self, forward_metadata):
-        """
-        Update the splitfuse attention mask based on forward metadata.
-
-        :param forward_metadata: Forward metadata containing sequence lengths and extended lengths.
-        :return: Updated splitfuse attention mask.
-        """
-        attn_mask_id = self.get_attention_mask_id(
-            forward_metadata.seq_lens_cpu_int,
-            forward_metadata.extend_seq_lens_cpu_int,
-        )
-        mix_mask = self.get_splitfuse_attn_mask(
-            seq_lens=forward_metadata.seq_lens_cpu_int,
-            position=attn_mask_id,
-            dtype=torch.float16,
-            device=self.device,
-        ).to(torch.bfloat16)
-        return mix_mask
 
 
 class AscendAttnBackend(AttentionBackend):
@@ -259,7 +229,7 @@ class AscendAttnBackend(AttentionBackend):
             self.ascend_attn_mask_builder.mask,
             self.ascend_attn_mask_builder.fia_mask,
             self.ascend_attn_mask_builder.mtp_mask,
-            self.ascend_attn_mask_builder.mix_mask_cache,
+            self.ascend_attn_mask_builder.mixed_chunk_attn_mask,
         )
         if self.use_mla:
             self.ringmla_mask = self.ascend_attn_mask_builder.ringmla_mask
@@ -333,10 +303,6 @@ class AscendAttnBackend(AttentionBackend):
                     )
                 )
 
-        if forward_batch.forward_mode.is_mixed():
-            self.mix_mask = self.ascend_attn_mask_builder.update_mask(
-                self.forward_metadata
-            )
         self.graph_mode = False
 
     def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int):
@@ -1450,29 +1416,28 @@ class AscendAttnBackend(AttentionBackend):
             )
         k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
         v_cache = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id)
+        num_block, block_size, _, _ = k_cache.shape
+        key = k_cache.view(num_block, block_size, -1)
+        value = v_cache.view(num_block, block_size, -1)
 
         query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
 
-        # Initialize the output tensor for attention results
+        attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score(
-        attn_output = torch.empty(
+            query,
-            (query.shape[0], layer.tp_q_head_num, layer.v_head_dim),
+            key,
-            dtype=query.dtype,
+            value,
-            device=query.device,
+            num_heads=layer.tp_q_head_num,
+            num_key_value_heads=layer.tp_k_head_num,
+            input_layout="TND",
+            block_size=block_size,
+            block_table=self.forward_metadata.block_tables,
+            atten_mask=self.mix_mask,
+            sparse_mode=3,
+            actual_seq_lengths=self.forward_metadata.seq_lens_list_cumsum,
+            actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_int,
+            scale=layer.scaling,
         )
 
-        torch_npu._npu_paged_attention_splitfuse(
-            query=query,
-            key_cache=k_cache,
-            value_cache=v_cache,
-            block_table=self.forward_metadata.block_tables,
-            context_lens=self.forward_metadata.seq_lens_cpu_int,
-            mask=self.mix_mask,
-            seq_len=self.forward_metadata.extend_seq_lens_cpu_int,
-            scale_value=layer.scaling,
-            num_heads=layer.tp_q_head_num,
-            num_kv_heads=layer.tp_k_head_num,
-            out=attn_output,
-        )
         return attn_output.view(
             attn_output.shape[0], layer.tp_q_head_num * layer.v_head_dim
         )