fix

Author: hiworldwzj

lightllm/common/basemodel/basemodel.py

@@ -36,7 +36,6 @@ class TpPartBaseModel:
     infer_state_class = InferStateInfo
 
     def __init__(self, kvargs):
-        self.infer_state = self.infer_state_class()
         self.run_mode = kvargs["run_mode"]
         self.tp_rank_ = kvargs["tp_rank"]
         self.world_size_ = kvargs["world_size"]
@@ -331,9 +330,7 @@ def _decode(
         b_seq_len,
         multimodal_params,
     ):
-        infer_state = self.infer_state
-        if self.graph is None or self.graph.need_capture(batch_size) or infer_state.is_prefill:
-            infer_state = self.infer_state_class()
+        infer_state = self.infer_state_class()
         infer_state.is_prefill = False
         infer_state.batch_size = batch_size
         infer_state.total_token_num = total_token_num

lightllm/models/deepseek2/flashinfer_struct.py

@@ -13,20 +13,14 @@ def __init__(self):
         super().__init__()
         self.prefill_wrapper = None
         self.decode_wrapper = None
+        self.flashinfer_extra_state = None
 
     def init_some_extra_state(self, model, input_ids: torch.Tensor):
         super().init_some_extra_state(model, input_ids)
+        self.flashinfer_extra_state = model.flashinfer_extra_state
 
         if not self.is_prefill:
             if enable_env_vars("ENABLE_FLASHINFER_DECODE_MLA"):
-                self.tp_q_head_num = model.flashinfer_state.tp_q_head_num
-                self.kv_lora_rank = model.flashinfer_state.kv_lora_rank
-                self.qk_rope_head_dim = model.flashinfer_state.qk_rope_head_dim
-                self.qk_nope_head_dim = model.flashinfer_state.qk_nope_head_dim
-                self.softmax_scale = model.flashinfer_state.softmax_scale
-                self.q_data_type = model.flashinfer_state.data_type
-                self.kv_data_type = model.flashinfer_state.data_type
-
                 self.q_indptr = torch.arange(self.batch_size + 1, dtype=torch.int32).to(input_ids.device)
                 self.kv_indices = torch.empty(
                     self.batch_size * model.flashinfer_state.max_seq_length, dtype=torch.int32
@@ -41,7 +35,7 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                 )
                 if self.decode_wrapper is None:
                     self.decode_wrapper = flashinfer.mla.BatchMLAPagedAttentionWrapper(
-                        model.flashinfer_state.workspace_buffer,
+                        self.flashinfer_extra_state.workspace_buffer,
                         use_cuda_graph=True,
                         qo_indptr=self.q_indptr,
                         kv_indices=self.kv_indices,
@@ -53,23 +47,17 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                         self.kv_starts,
                         self.kv_indices,
                         self.b_seq_len,
-                        self.tp_q_head_num,
-                        self.kv_lora_rank,
-                        self.qk_rope_head_dim,
+                        self.flashinfer_extra_state.tp_q_head_num,
+                        self.flashinfer_extra_state.kv_lora_rank,
+                        self.flashinfer_extra_state.qk_rope_head_dim,
                         1,
                         False,  # causal
-                        self.softmax_scale,
-                        self.q_data_type,
-                        self.kv_data_type,
+                        self.flashinfer_extra_state.softmax_scale,
+                        self.flashinfer_extra_state.q_data_type,
+                        self.flashinfer_extra_state.kv_data_type,
                     )
         else:
             if enable_env_vars("ENABLE_FLASHINFER_PREFILLED"):
-                self.tp_q_head_num = model.flashinfer_state.tp_q_head_num
-                self.qk_rope_head_dim = model.flashinfer_state.qk_rope_head_dim
-                self.qk_nope_head_dim = model.flashinfer_state.qk_nope_head_dim
-                self.softmax_scale = model.flashinfer_state.softmax_scale
-                self.q_data_type = model.flashinfer_state.data_type
-
                 q_starts = torch.cat(
                     [self.b_start_loc, self.b_start_loc[-1:] + (self.b_seq_len - self.b_ready_cache_len)[-1:]], dim=0
                 ).int()
@@ -78,18 +66,19 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                 ).int()
                 if self.prefill_wrapper is None:
                     self.prefill_wrapper = flashinfer.prefill.BatchPrefillWithRaggedKVCacheWrapper(
-                        model.flashinfer_state.workspace_buffer, "NHD"
+                        self.flashinfer_extra_state.workspace_buffer, "NHD"
                     )
                 self.prefill_wrapper.plan(
                     qo_indptr=q_starts,
                     kv_indptr=kv_starts,
-                    num_qo_heads=self.tp_q_head_num,
-                    num_kv_heads=self.tp_q_head_num,
-                    head_dim_qk=self.qk_nope_head_dim + self.qk_rope_head_dim,
-                    head_dim_vo=self.qk_nope_head_dim,
-                    q_data_type=self.q_data_type,
+                    num_qo_heads=self.flashinfer_extra_state.tp_q_head_num,
+                    num_kv_heads=self.flashinfer_extra_state.tp_q_head_num,
+                    head_dim_qk=self.flashinfer_extra_state.qk_nope_head_dim
+                    + self.flashinfer_extra_state.qk_rope_head_dim,
+                    head_dim_vo=self.flashinfer_extra_state.qk_nope_head_dim,
+                    q_data_type=self.flashinfer_extra_state.q_data_type,
                     causal=True,
-                    sm_scale=self.softmax_scale,
+                    sm_scale=self.flashinfer_extra_state.softmax_scale,
                 )
         return
 
@@ -101,13 +90,13 @@ def copy_for_cuda_graph(self, new_infer_state):
                 new_infer_state.kv_starts,
                 new_infer_state.kv_indices,
                 new_infer_state.b_seq_len,
-                new_infer_state.tp_q_head_num,
-                new_infer_state.kv_lora_rank,
-                new_infer_state.qk_rope_head_dim,
+                new_infer_state.flashinfer_extra_state.tp_q_head_num,
+                new_infer_state.flashinfer_extra_state.kv_lora_rank,
+                new_infer_state.flashinfer_extra_state.qk_rope_head_dim,
                 1,
                 False,  # causal
-                new_infer_state.softmax_scale,
-                new_infer_state.q_data_type,
-                new_infer_state.kv_data_type,
+                new_infer_state.flashinfer_extra_state.softmax_scale,
+                new_infer_state.flashinfer_extra_state.q_data_type,
+                new_infer_state.flashinfer_extra_state.kv_data_type,
             )
         return

lightllm/models/deepseek2/layer_infer/transformer_layer_infer.py

@@ -21,6 +21,7 @@
 from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd
 from lightllm.models.deepseek2.triton_kernel.rotary_emb import rotary_emb_fwd
 from lightllm.models.deepseek2.infer_struct import Deepseek2InferStateInfo
+from lightllm.models.deepseek2.flashinfer_struct import Deepseek2FlashInferStateInfo
 from functools import partial
 from lightllm.models.llama.yarn_rotary_utils import get_deepseek_mscale
 import os
@@ -224,7 +225,7 @@ def _context_attention_flashinfer_kernel_with_CC(
         self,
         q: torch.Tensor,
         kv,
-        infer_state: Deepseek2InferStateInfo,
+        infer_state: Deepseek2FlashInferStateInfo,
         layer_weight: Deepseek2TransformerLayerWeight,
         out=None,
     ) -> torch.Tensor:
@@ -240,7 +241,7 @@ def _context_attention_flashinfer_kernel_with_CC_fp8(
         self,
         q: torch.Tensor,
         kv,
-        infer_state: Deepseek2InferStateInfo,
+        infer_state: Deepseek2FlashInferStateInfo,
         layer_weight: Deepseek2TransformerLayerWeight,
         out=None,
     ) -> torch.Tensor:
@@ -393,7 +394,7 @@ def _context_attention_kernel_origin_fp8(
         return o_tensor
 
     def _token_gqa_decode_attention_flashinfer(
-        self, q, infer_state: Deepseek2InferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
+        self, q, infer_state: Deepseek2FlashInferStateInfo, layer_weight: Deepseek2TransformerLayerWeight, out=None
     ):
         q_nope, q_rope = q[:, :, : -self.qk_rope_head_dim], q[:, :, -self.qk_rope_head_dim :]
         q_nope = layer_weight.k_b_proj_.bmm(q_nope.transpose(0, 1)).transpose(0, 1)

lightllm/models/deepseek2/model.py

@@ -23,7 +23,8 @@ def __init__(self, model):
         self.qk_nope_head_dim = model.qk_nope_head_dim
         self.qk_rope_head_dim = model.qk_rope_head_dim
         self.kv_lora_rank = model.kv_lora_rank
-        self.data_type = model.data_type
+        self.q_data_type = model.data_type
+        self.kv_data_type = model.data_type
         self.workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8).to(model.tp_rank_)
         self.max_seq_length = model.max_seq_length
         self.softmax_scale = (self.qk_nope_head_dim + self.qk_rope_head_dim) ** (-0.5)
@@ -65,7 +66,7 @@ def _init_some_value(self):
         self.kv_lora_rank = self.config["kv_lora_rank"]
         self.head_dim_ = self.kv_lora_rank + self.qk_rope_head_dim
         if self.enable_flashinfer:
-            self.flashinfer_state = FlashInferStateExtraInfo(self)
+            self.flashinfer_extra_state = FlashInferStateExtraInfo(self)
 
     def _init_custom(self):
         self._init_to_get_yarn_rotary()