add vanilla_chunk_prefill_mla back

Signed-off-by: whx-sjtu <2952154980@qq.com>

Author: whx-sjtu

vllm_ascend/ascend_config.py

@@ -45,6 +45,8 @@ def __init__(self, vllm_config):
             ascend_scheduler_config)
 
         self.expert_map_path = additional_config.get("expert_map_path", None)
+        self.chunked_prefill_for_mla = additional_config.get(
+            "chunked_prefill_for_mla", False)
         self.enable_shared_expert_dp = additional_config.get(
             "enable_shared_expert_dp", False
         ) and not self.torchair_graph_config.enabled and vllm_config.parallel_config.enable_expert_parallel

vllm_ascend/attention/mla_v1.py

@@ -20,6 +20,7 @@
 from vllm_ascend.multistream.base import MSAttentionMetadataSplitConfig
 from vllm_ascend.multistream.context import get_multistream_comm_context
 from vllm_ascend.multistream.ms_split import model_input_split_v1_mla_attn
+from vllm_ascend.ops.attention import vanilla_chunked_prefill_mla
 from vllm_ascend.utils import npu_prefetch
 from vllm_ascend.worker.npu_input_batch import InputBatch
 
@@ -184,10 +185,7 @@ def __init__(self,
                            self.block_size - 1) // self.block_size
         self.chunked_prefill_enabled = scheduler_config.chunked_prefill_enabled
 
-        if vllm_config.speculative_config is not None:
-            self.decode_threshold = vllm_config.speculative_config.num_speculative_tokens + 1
-        else:
-            self.decode_threshold = 1
+        self.decode_threshold = 1
 
         if self.chunked_prefill_enabled:
             self.chunked_prefill_workspace_size = min(
@@ -483,6 +481,9 @@ def __init__(
         self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp
         self.enable_mla_prefetch = ascend_config.enable_mla_prefetch
         self.enable_kv_nz = ascend_config.torchair_graph_config.enable_kv_nz
+        self.chunked_prefill_for_mla = ascend_config.chunked_prefill_for_mla
+
+        self.prefill_mask = None
 
         # Adapt torch air graph mode with spec decoding.
         speculative_config = get_current_vllm_config().speculative_config
@@ -673,14 +674,18 @@ def _forward_prefill(
                 num_heads=self.num_heads,
                 num_kv_heads=self.num_heads,
                 out=attn_output)
-        else:
+        elif self.chunked_prefill_for_mla:
             attn_lse = torch.empty(self.num_heads,
                                    num_tokens,
                                    dtype=torch.float32,
                                    device=q_nope.device)
-            self.prefill_mask = torch.triu(
-                torch.ones(512, 512, device=q_nope.device, dtype=q_nope.dtype),
-                1)  # 512: mask only support 512
+            if self.prefill_mask is None:
+                self.prefill_mask = torch.triu(
+                    torch.ones(512,
+                               512,
+                               device=q_nope.device,
+                               dtype=q_nope.dtype),
+                    1)  # 512: mask only support 512
             if attn_metadata.num_prefills > 1:
                 self.prefill_mask = self.prefill_mask.unsqueeze(0).repeat(
                     attn_metadata.num_prefills, 1, 1)
@@ -706,9 +711,38 @@ def _forward_prefill(
                 softmax_lse=attn_lse)
             attn_output, attn_lse = self._compute_prefill_context( \
                 q_nope, q_pe, kv_c_and_k_pe_cache, self.qk_rope_head_dim, attn_metadata, attn_output, attn_lse)
+        else:
+            query = torch.cat((q_nope, q_pe), dim=-1)
+            attn_output_torch = torch.empty(num_tokens,
+                                            self.num_heads * self.v_head_dim,
+                                            dtype=query.dtype,
+                                            device=query.device)
+            # current requests is chunked in prefill, disable flash attention with chunked prefill
+            vanilla_chunked_prefill_mla(
+                output=attn_output_torch,
+                query=query,
+                kv_cache=kv_c_and_k_pe_cache,
+                block_tables=attn_metadata.prefill.block_table,
+                query_lens=attn_metadata.prefill.query_lens,
+                context_lens=attn_metadata.prefill.context_lens,
+                kv_b_proj=self.kv_b_proj,
+                max_query_len=attn_metadata.prefill.max_query_len,
+                max_context_len=attn_metadata.prefill.max_seq_lens,
+                nope_dim=self.qk_nope_head_dim,
+                rope_dim=self.qk_rope_head_dim,
+                v_head_dim=self.v_head_dim,
+                scale=self.scale,
+                alibi_slopes=None,
+                causal=True)
 
         attn_output = attn_output.reshape(
             [num_tokens, self.num_heads * self.v_head_dim])
+        if attn_metadata.attn_state in [
+                AscendAttentionState.ChunkedPrefill,
+                AscendAttentionState.SpecDecoding,
+                AscendAttentionState.PrefillCacheHit
+        ] and not self.chunked_prefill_for_mla:
+            attn_output = attn_output_torch
         return attn_output
 
     def exec_kv_decode(