[Perf] Optimize torch.where and vectorize PCP/DCP loops in mla_v1.py

Signed-off-by: F.Liu <[email protected]>

Author: F.Liu

vllm_ascend/attention/mla_v1.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
-from typing import (TYPE_CHECKING, ClassVar, List, NamedTuple, Optional, Tuple,
-                    Type, TypeVar)
+from typing import (TYPE_CHECKING, ClassVar, NamedTuple, Optional, Tuple, Type,
+                    TypeVar)
 
 import numpy as np
 import torch
@@ -376,7 +376,7 @@ def pad_actual_seq_len_q_mtp_disable_pad(self, num_reqs_pad_size, num_reqs,
         Only use for acl full graph mode.
         Pad the last element of the actual_seq_lengths_q equal to the TND(T) and
         the num of dimensions equal to the batch_size of main model.
-        
+
         For example:
         batch_size = 8, num_reqs = 4, num_speculative_tokens = 1
         input actual_seq_lengths_q = [1, 2, 4, 5]  (the 3rd req was accept a token)
@@ -661,7 +661,7 @@ def build(
                     batch_seq_mask, non_blocking=True)
                 batch_seq_mask = self.batch_seq_mask_buf[:batch_seq_mask.
                                                          shape[0]]
-                cp_seq_len = torch.where(cp_seq_len == 0, 1, cp_seq_len)
+                cp_seq_len.masked_fill_(cp_seq_len == 0, 1)
             else:
                 cp_seq_len, batch_seq_mask = None, None
 
@@ -1971,20 +1971,32 @@ def _forward_decode_pcp_dcp(
         attn_output = self._npu_attention_update(attn_out_lse_list)
         return self._v_up_proj(attn_output)
 
-    def _npu_attention_update(
-            self, attn_out_lse_list: List[torch.Tensor]) -> torch.Tensor:
-        attn_out_split_cp = []
-        attn_lse_split_cp = []
-
-        for attn_out_lse in attn_out_lse_list:
-            attn_out_allgather, attn_lse_allgather = self._out_lse_reshape(
-                *torch.split(attn_out_lse, [self.kv_lora_rank, 1], dim=-1))
-            attn_out_split_cp.append(attn_out_allgather)
-            attn_lse_split_cp.append(attn_lse_allgather)
-        attn_out, _ = torch_npu.npu_attention_update(attn_lse_split_cp,
-                                                     attn_out_split_cp, 0)
-        attn_out = attn_out.view(-1, attn_out_lse_list[0].shape[1],
-                                 self.kv_lora_rank)
+    def _npu_attention_update(self,
+                              attn_out_lse: torch.Tensor) -> torch.Tensor:
+        B_total, H_total, D_plus_1 = attn_out_lse.shape
+        S = B_total // self.pcp_size
+        H = H_total // self.dcp_size
+        D = self.kv_lora_rank
+        assert D_plus_1 == D + 1
+        # [PCP, S, DCP, H, D+1]
+        x = attn_out_lse.view(self.pcp_size, S, self.dcp_size, H, D_plus_1)
+        # [PCP, DCP, S, H, D+1]
+        x = x.permute(0, 2, 1, 3, 4).contiguous()
+        # Flatten [N, S, H, D+1], N = pcp_size * dcp_size
+        x = x.view(-1, S, H, D_plus_1)
+        # Split out lse
+        out_flat, lse_flat = torch.split(x, [D, 1],
+                                         dim=-1)  # [N, S, H, D], [N, S, H, 1]
+        #    out: [N, S, H, D] -> [N, S*H, D]
+        #    lse: [N, S, H, 1] -> [N, S*H]
+        out_flat = out_flat.flatten(1, 2)  # [N, S*H, D]
+        lse_flat = lse_flat.squeeze(-1).flatten(1)  # [N, S*H]
+        #  unbind to list
+        out_list = out_flat.unbind(0)  # [S*H, D]
+        lse_list = lse_flat.unbind(0)  # [S*H]
+
+        attn_out, _ = torch_npu.npu_attention_update(lse_list, out_list, 0)
+        attn_out = attn_out.view(-1, H, self.kv_lora_rank)
         return attn_out
 
     def _out_lse_reshape(self, attn_out: torch.Tensor,
@@ -2000,14 +2012,10 @@ def _process_attn_out_lse(
         attn_output: torch.Tensor,
         softmax_lse: torch.Tensor,
         decode_meta: AscendMLADecodeMetadata,
-    ) -> List[torch.Tensor]:
-        attn_out_lse_list = []
-        out_mask = decode_meta.batch_seq_mask[:, None,
-                                              None].expand_as(attn_output)
-        attn_output = torch.where(out_mask, 0, attn_output)
-        lse_mask = decode_meta.batch_seq_mask[:, None,
-                                              None].expand_as(softmax_lse)
-        softmax_lse = torch.where(lse_mask, -torch.inf, softmax_lse)
+    ) -> torch.Tensor:
+        out_lse_mask = decode_meta.batch_seq_mask[:, None, None].bool()
+        attn_output.masked_fill_(out_lse_mask, 0)
+        softmax_lse.masked_fill_(out_lse_mask, -torch.inf)
 
         softmax_lse = softmax_lse.to(torch.float32)
         attn_output = attn_output.to(torch.float32)
@@ -2020,30 +2028,15 @@ def _process_attn_out_lse(
             dist.all_to_all_single(attn_out_lse_all2all,
                                    attn_out_lse,
                                    group=self.dcp_group)
-            # permute: [num_heads, v_head_dim+1, bs] -> [bs, num_heads, v_head_dim+1]
             attn_out_lse_all2all = attn_out_lse_all2all.permute([2, 0, 1])
             if self.pcp_size > 1:
                 attn_out_lse = attn_out_lse_all2all.contiguous()
-            attn_out_lse_list = list(
-                torch.chunk(attn_out_lse_all2all, self.dcp_size, dim=1))
 
         if self.pcp_size > 1:
-            # AllGather out&lse within PCP group
-            attn_out_lse_list = [
-                torch.empty_like(attn_out_lse) for _ in range(self.pcp_size)
-            ]
-            dist.all_gather(attn_out_lse_list,
-                            attn_out_lse,
-                            group=self.pcp_group)
-        if self.dcp_size > 1 and self.pcp_size > 1:
-            attn_out_lse_list_pcp_dcp = []
-            for s in attn_out_lse_list:
-                attn_out_lse_list_split = list(
-                    torch.chunk(s, self.dcp_size, dim=1))
-                attn_out_lse_list_pcp_dcp += attn_out_lse_list_split
-            attn_out_lse_list = attn_out_lse_list_pcp_dcp
-
-        return attn_out_lse_list
+            # AllGather out&lse within CP group
+            attn_out_lse = get_pcp_group().all_gather(attn_out_lse, dim=0)
+
+        return attn_out_lse
 
     def _reorg_kvcache(
         self,
@@ -2116,4 +2109,4 @@ def _reorg_kvcache(
         assert reorganized_kv_c_normed.shape[0] == sum_seq_len
         assert reorganized_k_pe.shape[0] == sum_seq_len
         assert max_seq_len_check == max_seq_len
-        return reorganized_kv_c_normed, reorganized_k_pe
+        return reorganized_kv_c_normed, reorganized_k_pe