[long_seq_optim] support cp&sp

Author: LookAround0301

vllm_ascend/attention/mla_v1.py

@@ -1132,15 +1132,39 @@ def _mla_preprocess(self, hidden_states, kv_cache, attn_metadata,
             if self.cp_size * self.sp_size > 1:
                 decode_q_wo_k_up_pe = decode_q_wo_k_up[..., self.qk_nope_head_dim:]
                 decode_q_wo_k_up_pe = self.rope_single(decode_q_wo_k_up_pe, cos, sin)
+                decode_q_wo_k_up[..., self.qk_nope_head_dim:] = decode_q_wo_k_up_pe
+                decode_kv_no_split = kv_no_split[:num_decode_tokens]
+                kv_c, k_pe = decode_kv_no_split.split(
+                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+                k_pe = k_pe.unsqueeze(1) 
+                decode_k_pe = k_pe[:num_decode_tokens]
+                decode_k_pe = self.rope_single(decode_k_pe, cos, sin)
+                k_pe[:num_decode_tokens] = decode_k_pe
             else:
-            # TODO sp_cp的decode的适配这里还不完整，待思考位置编码的处理
                 decode_q_pe = self.rope_single(decode_q_pe, cos, sin)
             decode_slots = attn_metadata.slot_mapping[:num_decode_tokens]
             decode_kv_no_split = kv_no_split[:num_decode_tokens]
-            decode_k_pe, decode_k_nope = self.exec_kv_decode(
+            if self.cp_size * self.sp_size > 1:
+                kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
+                assert len(
+                    kv_cache
+                ) > 1, "the number of kv cache should be greater than 1, namely (nope_cache and rope_cache)"
+                kv_c_normed = kv_c_normed.view(
+                    [num_actual_tokens, self.num_kv_heads, -1])
+                torch_npu._npu_reshape_and_cache(
+                    key=kv_c_normed,
+                    value=k_pe,
+                    key_cache=kv_cache[0],
+                    value_cache=kv_cache[1],
+                    slot_indices=attn_metadata.slot_mapping)
+            else:
+                decode_k_pe, decode_k_nope = self.exec_kv_decode(
                 decode_kv_no_split, cos, sin, kv_cache, decode_slots)
-            decode_preprocess_res = DecodeMLAPreprocessResult(
-                decode_ql_nope, decode_q_pe, decode_k_nope, decode_k_pe, decode_q_wo_k_up)
+            if self.cp_size * self.sp_size > 1:
+                decode_preprocess_res = DecodeMLAPreprocessResult(decode_q_wo_k_up=decode_q_wo_k_up)
+            else:
+                decode_preprocess_res = DecodeMLAPreprocessResult(
+                    decode_ql_nope, decode_q_pe, decode_k_nope, decode_k_pe)
         # Preprocess for prefill tokens
         if has_prefill:
             prefill_kv_no_split = kv_no_split[
@@ -1154,22 +1178,46 @@ def _mla_preprocess(self, hidden_states, kv_cache, attn_metadata,
             sin = attn_metadata.prefill.sin
             prefill_slots = attn_metadata.slot_mapping[
                 num_decode_tokens:num_actual_tokens]
-            # TODO 待确认这里的位置编码是否会影响到kv cathe的存储
             prefill_q_pe = self.rope_single(prefill_q_pe, cos, sin)
             if self.cp_size > 1:
-                prefill_kv_no_split = get_cp_group().all_gather(prefill_kv_no_split, 0)
-                prefill_kv_no_split = torch.index_select(prefill_kv_no_split, 0, attn_metadata.prefill.cp_kv_recover_idx)
-            prefill_k_pe, prefill_k_c_normed = self.exec_kv_prefill(
-                prefill_kv_no_split, cos, sin, kv_cache, prefill_slots)
-            prefill_k_pe = prefill_k_pe.view(prefill_q_c.shape[0],
-                                             self.num_kv_heads, -1)
+                kv_c, k_pe = prefill_kv_no_split.split(
+                [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+                kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
+                assert len(
+                    kv_cache
+                ) > 1, "the number of kv cache should be greater than 1, namely (nope_cache and rope_cache)"
+                kv_c_normed = kv_c_normed.view(
+                    [num_actual_tokens, self.num_kv_heads, -1])
+                k_pe = k_pe.unsqueeze(1)    
+                prefill_k_pe = k_pe[num_decode_tokens:]
+                prefill_k_pe = self.rope_single(prefill_k_pe, cos, sin)
+                prefill_k_c_normed = kv_c_normed[num_decode_tokens:]
+                
+                prefill_kv_c_k_pe = torch.cat([prefill_k_c_normed, prefill_k_pe], dim=-1)
+                prefill_kv_c_k_pe = get_cp_group().all_gather(prefill_kv_c_k_pe, 0)
+                prefill_kv_c_k_pe = torch.index_select(prefill_kv_c_k_pe, 0, attn_metadata.prefill.cp_kv_recover_idx)
+                prefill_k_c_normed, prefill_k_pe = prefill_kv_c_k_pe.split([self.kv_lora_rank, self.qk_rope_head_dim],
+                                                                       dim=-1)
+                kv_c_normed, k_pe = prefill_k_c_normed, prefill_k_pe
+                prefill_k_c_normed = prefill_k_c_normed.squeeze()
+                torch_npu._npu_reshape_and_cache(
+                    key=kv_c_normed,
+                    value=k_pe,
+                    key_cache=kv_cache[0],
+                    value_cache=kv_cache[1],
+                    slot_indices=attn_metadata.slot_mapping)
+            else:
+                prefill_k_pe, prefill_k_c_normed = self.exec_kv_prefill(
+                    prefill_kv_no_split, cos, sin, kv_cache, prefill_slots)
             prefill_k_nope, prefill_value = self.kv_b_proj(
-                prefill_k_c_normed)[0].view(
-                    -1, self.num_heads,
+                    prefill_k_c_normed)[0].view(
+                    -1, self.num_heads, 
                     self.qk_nope_head_dim + self.v_head_dim).split(
                         [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
-            prefill_k_pe = prefill_k_pe.expand(
-                (*prefill_k_nope.shape[:-1], -1))
+            if not self.cp_size > 1:
+                prefill_k_pe = prefill_k_pe.view(prefill_q_c.shape[0],
+                                                self.num_kv_heads, -1)
+            prefill_k_pe = prefill_k_pe.expand((*prefill_k_nope.shape[:-1], -1))
             prefill_preprocess_res = PrefillMLAPreprocessResult(
                 prefill_q_nope, prefill_q_pe, prefill_k_nope, prefill_k_pe,
                 prefill_value)