opt: refactor some code for acc

Author: niushengxiao

lightllm/models/deepseek2/infer_struct.py

@@ -0,0 +1,14 @@
+import torch
+import numpy as np
+from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+
+
+class Deepseek2InferStateInfo(LlamaInferStateInfo):
+    def __init__(self):
+        super().__init__()
+        self.kv_starts = None
+
+    def init_some_extra_state(self, model, input_ids: torch.Tensor):
+        super().init_some_extra_state(model, input_ids)
+        self.kv_starts = torch.cat([self.b_start_loc, self.b_start_loc[-1:] + self.b_seq_len[-1:]], dim=0)
+        return

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -19,7 +19,7 @@
 from lightllm.models.llama.layer_infer.transformer_layer_infer import LlamaTransformerLayerInfer
 from lightllm.models.llama.triton_kernel.rmsnorm import rmsnorm_forward
 from lightllm.models.chatglm2.triton_kernel.rotary_emb import rotary_emb_fwd
-from lightllm.models.llama.infer_struct import LlamaInferStateInfo
+from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
 from functools import partial
 from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
 import os
@@ -83,7 +83,7 @@ def _get_qkv(
         self,
         input: torch.Tensor,
         cache_kv,
-        infer_state: LlamaInferStateInfo,
+        infer_state: Deepseek2InferStateInfo,
         layer_weight: Deepseek2TransformerLayerWeight,
     ) -> torch.Tensor:
         input = input.view(-1, self.embed_dim_)
@@ -133,7 +133,7 @@ def _get_qkv(
         return (q_nope, q_rope), cache_kv
 
     def _get_o(
-        self, input, infer_state: LlamaInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
+        self, input, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
     ) -> torch.Tensor:
         if not self.disable_vo_absorb:
             input = input.view(-1, self.tp_q_head_num_ * self.kv_lora_rank)
@@ -145,7 +145,7 @@ def _get_o(
         return o_tensor
 
     def _CC_method(
-        self, q, compressed_kv, infer_state: LlamaInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
+        self, q, compressed_kv, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
     ):
         num_local_heads = self.num_heads
         num_local_kv_heads = self.num_kv_heads
@@ -176,7 +176,7 @@ def _CC_method(
         return self._context_attention_kernel_with_v(q, [k_nope, k_pe], v, infer_state, layer_weight)
 
     def _ACC_method(
-        self, q, compressed_kv, infer_state: LlamaInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
+        self, q, compressed_kv, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
     ):
         q_nope, q_rope = q
         num_local_heads = self.num_heads
@@ -185,22 +185,21 @@ def _ACC_method(
             num_local_heads //= self.world_size_
             num_local_kv_heads //= self.world_size_
         # ACC
-        q_nope = layer_weight.k_b_proj_.weight.bmm(
+        q_nope = layer_weight.k_b_proj_.bmm(
             q_nope.transpose(0, 1),
         ).transpose(0, 1)
         if self.enable_opt_decoding_mha:
             import lightllm_ppl_mla
 
             o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype)
-            kvstarts = torch.cat(
-                [infer_state.b_start_loc, infer_state.b_start_loc[-1:] + infer_state.b_seq_len[-1:]], dim=0
-            )
+            q = torch.cat([q_nope, q_rope], dim=-1)
             lightllm_ppl_mla.decode_mla(
                 o_tensor,
                 q,
-                compressed_kv[: infer_state.mem_end, :, :],
-                infer_state.b_start_loc,
-                kvstarts,
+                compressed_kv,
+                infer_state.req_manager.req_to_token_indexs,
+                infer_state.kv_starts,
+                infer_state.b_req_idx,
                 self.softmax_scale,
                 q.shape[-1],
                 q_nope.shape[-1],
@@ -214,20 +213,20 @@ def _ACC_method(
         return vo
 
     def _context_attention_kernel(
-        self, q, kv, infer_state: LlamaInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
+        self, q, kv, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ) -> torch.Tensor:
         if self.mla_type == "MIX":
             return self._context_attention_kernel_with_CC(q, kv, infer_state, layer_weight, out)
         else:
             return self._context_attention_kernel_origin(q, kv, infer_state, layer_weight, out)
 
     def _context_attention_kernel_with_CC(
-        self, q, kv, infer_state: LlamaInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
+        self, q, kv, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ) -> torch.Tensor:
         return self._CC_method(q, kv, infer_state, layer_weight)
 
     def _context_attention_kernel_with_v(
-        self, q: Tuple[torch.Tensor, torch.Tensor], k, v, infer_state: LlamaInferStateInfo, layer_weight, out=None
+        self, q: Tuple[torch.Tensor, torch.Tensor], k, v, infer_state: Deepseek2InferStateInfo, layer_weight, out=None
     ) -> torch.Tensor:
         q_nope, q_rope = q
         k_nope, k_rope = k
@@ -267,7 +266,7 @@ def _context_attention_kernel_with_v(
         return o_tensor
 
     def _context_attention_kernel_origin(
-        self, q: Tuple[torch.Tensor, torch.Tensor], kv, infer_state: LlamaInferStateInfo, layer_weight, out=None
+        self, q: Tuple[torch.Tensor, torch.Tensor], kv, infer_state: Deepseek2InferStateInfo, layer_weight, out=None
     ) -> torch.Tensor:
         q_nope, q_rope = q
         o_tensor = self.alloc_tensor(q_nope.shape, dtype=q_nope.dtype) if out is None else out
@@ -304,20 +303,22 @@ def _context_attention_kernel_origin(
         q_rope = None
         return o_tensor
 
-    def _token_gqa_decode_attention_flashdecoding(self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None):
+    def _token_gqa_decode_attention_flashdecoding(
+        self, q, infer_state: Deepseek2InferStateInfo, layer_weight, out=None
+    ):
         if self.mla_type == "MIX":
             return self._token_gqa_decode_attention_flashdecoding_with_ACC(q, infer_state, layer_weight, out)
         else:
             return self._token_gqa_decode_attention_flashdecoding_origin(q, infer_state, layer_weight, out)
 
     def _token_gqa_decode_attention_flashdecoding_with_ACC(
-        self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None
+        self, q, infer_state: Deepseek2InferStateInfo, layer_weight, out=None
     ):
-        # compressed_kv = infer_state.mem_manager.kv_buffer[self.layer_num_][: infer_state.mem_end, :, :]
-        return self._ACC_method(q, None, infer_state, layer_weight)
+        compressed_kv = infer_state.mem_manager.kv_buffer[self.layer_num_]
+        return self._ACC_method(q, compressed_kv, infer_state, layer_weight)
 
     def _token_gqa_decode_attention_flashdecoding_origin(
-        self, q, infer_state: LlamaInferStateInfo, layer_weight, out=None
+        self, q, infer_state: Deepseek2InferStateInfo, layer_weight, out=None
     ):
         q_nope, q_rope = q
         kv = infer_state.mem_manager.kv_buffer[self.layer_num_][:, :, : self.kv_lora_rank]
@@ -347,7 +348,7 @@ def _copy_kv_to_mem_cache_normal(self, buffer, mem_index, mem_manager):
         return
 
     def _moe_ffn(
-        self, input, infer_state: LlamaInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
+        self, input, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight
     ) -> torch.Tensor:
         hidden_states = input.view(-1, self.embed_dim_)
         num_tokens, hidden_dim = hidden_states.shape

lightllm/models/deepseek2/model.py

@@ -3,6 +3,7 @@
 import torch
 from lightllm.models.deepseek2.layer_infer.transformer_layer_infer import Deepseek2TransformerLayerInfer
 from lightllm.models.deepseek2.layer_weights.transformer_layer_weight import Deepseek2TransformerLayerWeight
+from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
 from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights
 
 from lightllm.models.llama.model import LlamaTpPartModel
@@ -20,6 +21,9 @@ class Deepseek2TpPartModel(LlamaTpPartModel):
     # infer class
     transformer_layer_infer_class = Deepseek2TransformerLayerInfer
 
+    # infer state class
+    infer_state_class = Deepseek2InferStateInfo
+
     def __init__(self, kvargs):
         self.disable_qk_absorb = os.getenv("DISABLE_QK_ABSORB", "False").upper() in ["ON", "TRUE", "1"]
         self.disable_vo_absorb = os.getenv("DISABLE_VO_ABSORB", "False").upper() in ["ON", "TRUE", "1"]