remove _pre_cache_kv

Author: hiworldwzj

lightllm/common/basemodel/basemodel.py

@@ -290,10 +290,6 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0)
         infer_state.req_manager = self.req_manager
 
         infer_state.mem_index = model_input.mem_indexes
-        infer_state.kv_buffer_shapedtype = (
-            (model_input.input_ids.shape[0], self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
-            self.data_type,
-        )
         infer_state.microbatch_index = microbatch_index
         infer_state.dist_group = dist_group_manager.get_group(microbatch_index)
 

lightllm/common/basemodel/infer_struct.py

@@ -40,7 +40,6 @@ def __init__(self):
         self.req_manager: ReqManager = None
 
         self.mem_index: torch.Tensor = None
-        self.kv_buffer_shapedtype: Tuple[Any, Any] = None
 
         self.is_token_healing: bool = False
         self.return_all_prompt_logics: bool = False

lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_cohere_template.py

@@ -44,13 +44,12 @@ def _bind_rotary_emb_fwd(self):
     def _get_qkv(
         self, input, infer_state: InferStateInfo, layer_weight
     ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
         q = torch.mm(input.view(-1, self.embed_dim_), layer_weight.q_weight_)
-        torch.mm(
+        cache_kv = torch.mm(
             input.view(-1, self.embed_dim_),
             layer_weight.kv_weight_,
-            out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_),
-        )
+        ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
+
         if self.use_qk_norm_:
             q = q.view(-1, self.tp_q_head_num_, self.head_dim_)
             k = cache_kv[:, 0 : self.tp_k_head_num_, :]

lightllm/common/basemodel/layer_infer/template/transformer_layer_infer_template.py

@@ -31,22 +31,6 @@ def _att_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.T
     def _ffn_norm(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
         raise Exception("need to impl")
 
-    def _pre_cache_kv(self, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
-        if infer_state.is_prefill and get_env_start_args().enable_dp_prefill_balance:
-            shape = infer_state.kv_buffer_shapedtype[0]
-            shape = (len(infer_state.position_ids), *shape[1:])
-        else:
-            shape = infer_state.kv_buffer_shapedtype[0]
-
-        cache_kv = self.alloc_tensor(
-            shape=shape,
-            dtype=infer_state.kv_buffer_shapedtype[1],
-            device="cuda",
-            is_graph_out=False,
-            microbatch_index=infer_state.microbatch_index,
-        )
-        return cache_kv
-
     def _get_qkv(self, input, infer_state: InferStateInfo, layer_weight) -> Tuple[torch.Tensor, torch.Tensor]:
         raise Exception("need to impl")
 

lightllm/models/bloom/layer_infer/transformer_layer_infer.py

@@ -47,10 +47,7 @@ def _get_qkv(
         self, input, infer_state: InferStateInfo, layer_weight: BloomTransformerLayerWeight
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         q = layer_weight.q_proj.mm(input.view(-1, self.embed_dim_))
-        cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
-        cache_kv = layer_weight.kv_proj.mm(
-            input, out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_)
-        ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
+        cache_kv = layer_weight.kv_proj.mm(input).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
         return q, cache_kv
 
     def _context_attention_kernel(

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -143,14 +143,6 @@ 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,
@@ -161,8 +153,7 @@ def _get_qkv(
 
         if self.q_lora_rank is None:
             q = layer_weight.q_weight_.mm(input)
-            cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
-            layer_weight.kv_a_proj_with_mqa_.mm(input, out=cache_kv.view(-1, self.kv_lora_rank + self.qk_rope_head_dim))
+            cache_kv = layer_weight.kv_a_proj_with_mqa_.mm(input).view(-1, 1, self.kv_lora_rank + self.qk_rope_head_dim)
         else:
             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
@@ -203,8 +194,7 @@ def _tpsp_get_qkv(
 
             input = input.view(-1, self.embed_dim_)
             q = layer_weight.q_weight_.mm(input)
-            cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
-            layer_weight.kv_a_proj_with_mqa_.mm(input, out=cache_kv.view(-1, self.kv_lora_rank + self.qk_rope_head_dim))
+            cache_kv = layer_weight.kv_a_proj_with_mqa_.mm(input).view(-1, 1, self.kv_lora_rank + self.qk_rope_head_dim)
         else:
             input = input.view(-1, self.embed_dim_)
             qkv = layer_weight.qkv_a_proj_with_mqa_.mm(input)

lightllm/models/gemma3/layer_infer/transformer_layer_infer.py

@@ -87,9 +87,9 @@ def _get_qkv(
         # kv = kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
         k = layer_weight.k_proj.mm(input)
         v = layer_weight.v_proj.mm(input)
-        cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
-        cache_kv[:, 0 : self.tp_k_head_num_, :] = k.view(-1, self.tp_k_head_num_, self.head_dim_)
-        cache_kv[:, self.tp_k_head_num_ :, :] = v.view(-1, self.tp_v_head_num_, self.head_dim_)
+        cache_kv = torch.cat(
+            [k.view(-1, self.tp_k_head_num_, self.head_dim_), v.view(-1, self.tp_v_head_num_, self.head_dim_)], dim=1
+        )
 
         # gemma3 use qk norm
         q = q.view(-1, self.tp_q_head_num_, self.head_dim_)

lightllm/models/llama/layer_infer/transformer_layer_infer.py

@@ -197,10 +197,7 @@ def _get_qkv(
         self, input, infer_state: LlamaInferStateInfo, layer_weight: LlamaTransformerLayerWeight
     ) -> torch.Tensor:
         q = layer_weight.q_proj.mm(input)
-        cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
-        cache_kv = layer_weight.kv_proj.mm(
-            input, out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_)
-        ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
+        cache_kv = layer_weight.kv_proj.mm(input).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
 
         rotary_emb_fwd(
             q.view(-1, self.tp_q_head_num_, self.head_dim_),
@@ -222,10 +219,7 @@ def _tpsp_get_qkv(
             input = gather_input[0 : len(infer_state.position_cos), :]
 
         q = layer_weight.q_proj.mm(input)
-        cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
-        cache_kv = layer_weight.kv_proj.mm(
-            input, out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_)
-        ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
+        cache_kv = layer_weight.kv_proj.mm(input).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
 
         rotary_emb_fwd(
             q.view(-1, self.tp_q_head_num_, self.head_dim_),

lightllm/models/phi3/layer_infer/transformer_layer_infer.py

@@ -29,10 +29,8 @@ def _bind_attention(self):
 
     def _get_qkv(self, input_emb, infer_state: LlamaInferStateInfo, layer_weight: Phi3TransformerLayerWeight):
         q = layer_weight.q_proj.mm(input_emb.view(-1, self.embed_dim_))
-        cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
         cache_kv = layer_weight.kv_proj.mm(
             input_emb.view(-1, self.embed_dim_),
-            out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_),
         ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
         rotary_emb_fwd(
             q.view(-1, self.tp_q_head_num_, self.head_dim_),

lightllm/models/qwen/layer_infer/transformer_layer_infer.py

@@ -18,10 +18,9 @@ def __init__(self, layer_num, network_config, mode=[]):
 
     def _get_qkv(self, input_emb, infer_state: QwenInferStateInfo, layer_weight: QwenTransformerLayerWeight):
         q = layer_weight.q_proj.mm(input_emb)
-        cache_kv = self._pre_cache_kv(infer_state=infer_state, layer_weight=layer_weight)
-        cache_kv = layer_weight.kv_proj.mm(
-            input_emb, out=cache_kv.view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_) * self.head_dim_)
-        ).view(-1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_)
+        cache_kv = layer_weight.kv_proj.mm(input_emb).view(
+            -1, (self.tp_k_head_num_ + self.tp_v_head_num_), self.head_dim_
+        )
 
         rotary_emb_fwd(
             q.view(-1, self.tp_q_head_num_, self.head_dim_),