merge q_a_proj and kv_a_proj to reduce the kernel launch overhead (#1055)

Author: shihaobai

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -140,6 +140,14 @@ def _bind_attention(self):
                     Deepseek2TransformerLayerInfer._context_attention_kernel_origin, self
                 )
 
+    def _pre_cache_kv(
+        self, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
+    ) -> torch.Tensor:
+        # q_lora_rank 不是None的时候，融合 q_a_proj 和 kv_a_proj_with_mqa
+        if self.q_lora_rank is None:
+            return super()._pre_cache_kv(infer_state, layer_weight)
+        return None
+
     def _get_qkv(
         self,
         input: torch.Tensor,
@@ -151,13 +159,16 @@ def _get_qkv(
 
         if self.q_lora_rank is None:
             q = layer_weight.q_weight_.mm(input)
+            layer_weight.kv_a_proj_with_mqa_.mm(input, out=cache_kv.view(-1, self.kv_lora_rank + self.qk_rope_head_dim))
         else:
-            q = layer_weight.q_a_proj_.mm(input)
-            rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_, out=q)
+            q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split(
+                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
+            )
+            q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_)
             q = layer_weight.q_b_proj_.mm(q)
+            cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank + self.qk_rope_head_dim)
         q = q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim)
         q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
-        layer_weight.kv_a_proj_with_mqa_.mm(input, out=cache_kv.view(-1, self.kv_lora_rank + self.qk_rope_head_dim))
         rmsnorm_forward(
             cache_kv[:, :, : self.kv_lora_rank],
             weight=layer_weight.kv_a_layernorm_.weight,
@@ -185,16 +196,18 @@ def _tpsp_get_qkv(
             input = gather_input[0 : len(infer_state.position_cos), :]
 
         input = input.view(-1, self.embed_dim_)
-
         if self.q_lora_rank is None:
             q = layer_weight.q_weight_.mm(input)
+            layer_weight.kv_a_proj_with_mqa_.mm(input, out=cache_kv.view(-1, self.kv_lora_rank + self.qk_rope_head_dim))
         else:
-            q = layer_weight.q_a_proj_.mm(input)
-            rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_, out=q)
+            q, cache_kv = layer_weight.qkv_a_proj_with_mqa_.mm(input).split(
+                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
+            )
+            q = rmsnorm_forward(q, weight=layer_weight.q_a_layernorm_.weight, eps=self.eps_)
             q = layer_weight.q_b_proj_.mm(q)
+            cache_kv = cache_kv.view(-1, 1, self.kv_lora_rank + self.qk_rope_head_dim)
         q = q.view(-1, self.tp_q_head_num_, self.qk_nope_head_dim + self.qk_rope_head_dim)
         q_nope, q_rope = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
-        layer_weight.kv_a_proj_with_mqa_.mm(input, out=cache_kv.view(-1, self.kv_lora_rank + self.qk_rope_head_dim))
         rmsnorm_forward(
             cache_kv[:, :, : self.kv_lora_rank],
             weight=layer_weight.kv_a_layernorm_.weight,

lightllm/models/deepseek2/layer_weights/transformer_layer_weight.py

@@ -148,13 +148,25 @@ def _init_qkvo(self):
                 layer_num=self.layer_num_,
                 name="q_weight",
             )
+            self.kv_a_proj_with_mqa_ = ROWMMWeight(
+                weight_name=f"model.layers.{self.layer_num_}.self_attn.kv_a_proj_with_mqa.weight",
+                data_type=self.data_type_,
+                quant_cfg=self.quant_cfg,
+                layer_num=self.layer_num_,
+                name="kv_a_proj_with_mqa",
+                tp_rank=0,
+                tp_world_size=1,
+            )
         else:
-            self.q_a_proj_ = ROWMMWeight(
-                weight_name=f"model.layers.{self.layer_num_}.self_attn.q_a_proj.weight",
+            self.qkv_a_proj_with_mqa_ = MultiROWMMWeight(
+                weight_names=[
+                    f"model.layers.{self.layer_num_}.self_attn.q_a_proj.weight",
+                    f"model.layers.{self.layer_num_}.self_attn.kv_a_proj_with_mqa.weight",
+                ],
                 data_type=self.data_type_,
                 quant_cfg=self.quant_cfg,
                 layer_num=self.layer_num_,
-                name="q_a_proj",
+                name="qkv_a_proj_with_mqa",
                 tp_rank=0,
                 tp_world_size=1,
             )
@@ -165,16 +177,6 @@ def _init_qkvo(self):
                 layer_num=self.layer_num_,
                 name="q_b_proj",
             )
-
-        self.kv_a_proj_with_mqa_ = ROWMMWeight(
-            weight_name=f"model.layers.{self.layer_num_}.self_attn.kv_a_proj_with_mqa.weight",
-            data_type=self.data_type_,
-            quant_cfg=self.quant_cfg,
-            layer_num=self.layer_num_,
-            name="kv_a_proj_with_mqa",
-            tp_rank=0,
-            tp_world_size=1,
-        )
         self.k_b_proj_ = ROWBMMWeight(
             weight_name=f"model.layers.{self.layer_num_}.self_attn.k_b_proj.weight",
             data_type=self.data_type_,