[WIP][Prefill Performance] Parallel Strategy Optimizations (VRAM-for-Speed Tradeoff)

vllm_ascend/attention/mla_v1.py

@@ -28,6 +28,12 @@
     from vllm.v1.core.sched.output import SchedulerOutput
     from vllm.v1.worker.gpu_input_batch import InputBatch
 
+try:
+    from torch_npu.atb import npu_mla_prefill  # noqa: F401
+    ATB_MLA_PREFILL_ENABLED = True
+except ImportError:
+    ATB_MLA_PREFILL_ENABLED = False
+
 
 class AscendMLABackend(AttentionBackend):
 
@@ -623,6 +629,7 @@ def __init__(
         ascend_config = get_ascend_config()
         self.torchair_graph_enabled = ascend_config.torchair_graph_config.enabled
         self.enable_kv_nz = ascend_config.torchair_graph_config.enable_kv_nz
+        self.enable_prefill_optimizations = ascend_config.enable_prefill_optimizations
 
         # Adapt torch air graph mode with spec decoding.
         speculative_config = get_current_vllm_config().speculative_config
@@ -882,17 +889,41 @@ def _forward_prefill(
                 query, kv_c_and_k_pe_cache, self.qk_rope_head_dim, attn_metadata, attn_output, attn_lse)
 
         elif attn_metadata.attn_state == AscendAttentionState.PrefillNoCache:
-            key = torch.cat((k_nope, k_pe), dim=-1)
-            torch_npu._npu_flash_attention(
-                query=query,
-                key=key,
-                value=value,
-                mask=attn_metadata.attn_mask,
-                seq_len=attn_metadata.prefill.context_lens,
-                scale_value=self.scale,
-                num_heads=self.num_heads,
-                num_kv_heads=self.num_heads,
-                out=attn_output)
+            if not self.enable_prefill_optimizations or not ATB_MLA_PREFILL_ENABLED:
+                key = torch.cat((k_nope, k_pe), dim=-1)
+                torch_npu._npu_flash_attention(
+                    query=query,
+                    key=key,
+                    value=value,
+                    mask=attn_metadata.attn_mask,
+                    seq_len=attn_metadata.prefill.context_lens,
+                    scale_value=self.scale,
+                    num_heads=self.num_heads,
+                    num_kv_heads=self.num_heads,
+                    out=attn_output)
+            else:
+                q_pe = query[..., self.qk_nope_head_dim:]
+                q_nope = query[..., :self.qk_nope_head_dim]
+                mask = torch.triu(
+                    torch.ones(512,
+                               512,
+                               device=query.device,
+                               dtype=query.dtype),
+                    1)  # 512: mask only support 512
+                torch_npu.atb.npu_mla_prefill(
+                    q=q_nope,
+                    q_rope=q_pe,
+                    k=k_nope,
+                    k_rope=k_pe,
+                    v=value,
+                    q_seqlen=attn_metadata.prefill.context_lens,
+                    kv_seqlen=attn_metadata.prefill.context_lens,
+                    q_headnum=self.num_heads,
+                    qk_scale=self.scale,
+                    kv_headnum=self.num_heads,
+                    mask=mask,
+                    mask_type="mask_type_free",
+                    output=attn_output)
             attn_output = attn_output.view(-1, self.num_heads, self.v_head_dim)
         attn_output = attn_output.reshape(
             [num_tokens, self.num_heads * self.v_head_dim])

vllm_ascend/quantization/w8a8_dynamic.py

@@ -520,6 +520,18 @@ def fused_experts_with_all2all(hidden_states: torch.Tensor,
                 expert_tokens_before_capacity_flag=False,
                 quant_mode=1,
             )
+        elif hasattr(torch_npu, "npu_moe_init_routing_quantv2"):  # TODO: Remove it
+            quantized_tokens, expanded_row_idx, global_expert_tokens, _, token_scales = torch_npu.npu_moe_init_routing_quantv2(
+                hidden_states,
+                expert_idx=topk_ids.to(torch.int32),
+                active_num=0,
+                expert_capacity=0,
+                expert_num=global_num_experts,
+                drop_pad_mode=0,
+                expert_tokens_count_or_cumsum_flag=2,
+                expert_tokens_before_capacity_flag=False,
+                quant_mode=1,
+            )
         else:
             quantized_tokens, expanded_row_idx, global_expert_tokens, token_scales = init_routing_quant(
                 hidden_states, top_k, topk_ids, global_num_experts)

vllm_ascend/worker/model_runner_v1.py

@@ -712,7 +712,7 @@ def _make_attention_mask(self, seq_lens, query_lens, position,
                 seq_lens, query_lens, position, self.dtype, self.device)
         # Prefill without cache situation.
         elif attn_state == AscendAttentionState.PrefillNoCache:
-            max_seq_len = max(seq_lens, default=0)
+            max_seq_len = 128
             return self.attn_mask_builder.get_attn_mask(
                 max_seq_len, self.dtype, self.device)
         # Prefill with cache hit.