Chunked prefill (#717)

Co-authored-by: baishihao <baishihao@sensetime.com>
Co-authored-by: hiworldwzj <30762946+hiworldwzj@users.noreply.github.com>

Author: shihaobai

docs/CN/source/models/test.rst

@@ -135,12 +135,12 @@ internlm2-1_8b
 
     $ python -m lightllm.server.api_server 
     $           --model_dir ~/models/internlm2-1_8b  \
-    $           --splitfuse_mode                     \
+    $           --enable_chunked_prefill                     \
     $           --trust_remote_code               
 
 .. tip::
 
-    ``--splitfuse_mode`` 表示使用splitfuse进行加速。
+    ``--enable_chunked_prefill`` 表示使用chunkedprefill进行长文本推理。
 
 
 **测试服务**

docs/EN/source/models/test.rst

@@ -213,12 +213,12 @@ internlm2-1_8b
 .. code-block:: console
 
     $ python -m lightllm.server.api_server --model_dir ~/models/internlm2-1_8b  \
-    $                                       --splitfuse_mode                     \
+    $                                       --enable_chunked_prefill                     \
     $                                       --trust_remote_code               
 
 .. tip::
 
-    ``--splitfuse_mode`` Indicates the use of splitfuse for acceleration.
+    ``--enable_chunked_prefill`` Indicates the use of chunkedprefill for long context.
 
 
 **Test Server**

lightllm/common/basemodel/__init__.py

@@ -9,7 +9,6 @@
 from .layer_infer.template.pre_layer_infer_template import PreLayerInferTpl
 from .layer_infer.template.post_layer_infer_template import PostLayerInferTpl
 from .infer_struct import InferStateInfo
-from .splitfuse_infer_struct import SplitFuseInferStateInfo
 from .basemodel import TpPartBaseModel
 
 
@@ -26,5 +25,4 @@
     "TpPartBaseModel",
     "PreLayerInferTpl",
     "PostLayerInferTpl",
-    "SplitFuseInferStateInfo",
 ]

lightllm/common/basemodel/basemodel.py

@@ -7,13 +7,11 @@
 
 from lightllm.common.basemodel.layer_weights.hf_load_utils import load_hf_weights
 from lightllm.common.basemodel.infer_struct import InferStateInfo
-from lightllm.common.basemodel.splitfuse_infer_struct import SplitFuseInferStateInfo
 from lightllm.common.mem_manager import MemoryManager
 from lightllm.common.req_manager import ReqManager
 from lightllm.common.infer_utils import init_req_to_token_indexes
 from lightllm.common.build_utils import repair_config
 from lightllm.common.basemodel.triton_kernel.copy_kv_index_to_req import copy_kv_index_to_req
-from lightllm.common.basemodel.triton_kernel.splitfuse_copy_kv_index_to_req import splitfuse_copy_kv_index_to_req
 from lightllm.common.basemodel.layer_infer.cache_tensor_manager import g_cache_manager
 from lightllm.common.basemodel.cuda_graph import CudaGraph
 from lightllm.common.quantization import Quantcfg
@@ -36,7 +34,6 @@ class TpPartBaseModel:
 
     # infer state class
     infer_state_class = InferStateInfo
-    splitfuse_infer_state_class = SplitFuseInferStateInfo
 
     def __init__(self, kvargs):
         self.run_mode = kvargs["run_mode"]
@@ -57,6 +54,8 @@ def __init__(self, kvargs):
         self.return_all_prompt_logics = kvargs.get("return_all_prompt_logics", False)
         assert not (self.is_token_healing and self.return_all_prompt_logics), "can not be true in same time"
         self.use_dynamic_prompt_cache = kvargs.get("use_dynamic_prompt_cache", False)
+        enable_chunked_prefill = kvargs.get("enable_chunked_prefill", False)  # chunked prefill is default on.
+        self.use_dynamic_prompt_cache = self.use_dynamic_prompt_cache or enable_chunked_prefill
         self.data_type = kvargs.get("data_type", "float16")
         self.graph_max_batch_size = kvargs.get("graph_max_batch_size", 16)
         self.graph_max_len_in_batch = kvargs.get("graph_max_len_in_batch", 8192)
@@ -368,81 +367,6 @@ def _decode(
             predict_logics = self._token_forward(input_ids, infer_state)
         return predict_logics
 
-    @torch.no_grad()
-    def splitfuse_forward(
-        self,
-        input_ids,
-        mem_indexes,
-        decode_req_num,
-        decode_total_token_num,
-        decode_b_req_idx: torch.Tensor,
-        decode_b_start_loc: torch.Tensor,
-        decode_b_seq_len: torch.Tensor,
-        decode_max_len_in_batch,
-        prefill_req_num,
-        prefill_b_req_idx: torch.Tensor,
-        prefill_b_split_start_loc: torch.Tensor,
-        prefill_b_split_ready_cache_len: torch.Tensor,
-        prefill_max_split_seq_len_in_batch,
-        prefill_b_seq_len: torch.Tensor,
-    ):
-
-        infer_state = self.splitfuse_infer_state_class()
-        infer_state.use_dynamic_prompt_cache = self.use_dynamic_prompt_cache
-        infer_state.batch_size = decode_req_num + prefill_req_num
-
-        infer_state.decode_req_num = decode_req_num
-        infer_state.decode_total_token_num = decode_total_token_num
-        infer_state.decode_b_req_idx = decode_b_req_idx
-        infer_state.decode_b_start_loc = decode_b_start_loc
-        infer_state.decode_b_seq_len = decode_b_seq_len
-        infer_state.decode_max_len_in_batch = decode_max_len_in_batch
-
-        infer_state.prefill_req_num = prefill_req_num
-        infer_state.prefill_b_req_idx = prefill_b_req_idx
-        infer_state.prefill_b_split_start_loc = prefill_b_split_start_loc
-        infer_state.prefill_b_split_ready_cache_len = prefill_b_split_ready_cache_len
-        infer_state.prefill_max_split_seq_len_in_batch = prefill_max_split_seq_len_in_batch
-        infer_state.prefill_b_seq_len = prefill_b_seq_len
-        # infer_state.event = [torch.cuda.Event() for _ in range(self.layers_num)]
-
-        infer_state.mem_manager = self.mem_manager
-        infer_state.req_manager = self.req_manager
-
-        alloc_size = len(input_ids)
-        infer_state.mem_is_contiguous = False
-        infer_state.mem_index = mem_indexes
-        infer_state.kv_buffer = torch.empty(
-            (alloc_size, self.tp_k_head_num_ + self.tp_v_head_num_, self.head_dim_),
-            dtype=self.data_type,
-            device="cuda",
-        )
-
-        # decode 部分
-        if decode_req_num != 0:
-            copy_kv_index_to_req(
-                self.req_manager.req_to_token_indexs,
-                decode_b_req_idx,
-                decode_b_seq_len,
-                infer_state.mem_index[0:decode_req_num],
-            )
-
-        # split prefill 部分
-        if prefill_req_num != 0:
-            splitfuse_copy_kv_index_to_req(
-                self.req_manager.req_to_token_indexs,
-                prefill_b_req_idx,
-                prefill_b_split_ready_cache_len,
-                prefill_b_seq_len,
-                infer_state.mem_index[decode_req_num:],
-            )
-
-        infer_state.init_some_extra_state(self, input_ids)
-        infer_state.create_inner_decode_infer_status()
-        infer_state.create_inner_prefill_infer_status()
-        predict_logics = self._splitfuse_forward(input_ids, infer_state)
-        return predict_logics
-
     @final
     def _context_forward(self, input_ids, infer_state: InferStateInfo):
         g_cache_manager.cache_env_in()
@@ -469,17 +393,6 @@ def _token_forward(self, input_ids, infer_state: InferStateInfo):
         g_cache_manager.cache_env_out()
         return predict_logics
 
-    @final
-    def _splitfuse_forward(self, input_ids, infer_state: SplitFuseInferStateInfo):
-        g_cache_manager.cache_env_in()
-        cuda_input_ids = input_ids
-        input_embs = self.pre_infer.splitfuse_forward(cuda_input_ids, infer_state, self.pre_post_weight)
-        for i in range(0, self.layers_num):
-            input_embs = self.layers_infer[i].splitfuse_forward(input_embs, infer_state, self.trans_layers_weight[i])
-        predict_logics = self.post_infer.splitfuse_forward(input_embs, infer_state, self.pre_post_weight)
-        g_cache_manager.cache_env_out()
-        return predict_logics
-
     @final
     @torch.no_grad()
     def _check_max_len_infer(self):

lightllm/common/basemodel/infer_struct.py

@@ -30,7 +30,6 @@ def __init__(self):
         self.mem_end = None
         self.kv_buffer = None
 
-        self.is_splitfuse = False
         self.is_token_healing = False
         self.return_all_prompt_logics = False
         self.use_dynamic_prompt_cache = False

lightllm/common/basemodel/layer_infer/base_layer_infer.py

@@ -1,7 +1,6 @@
 import torch
 from typing import Dict, Iterable, Literal, Tuple, Union, List
 from lightllm.common.basemodel.infer_struct import InferStateInfo
-from lightllm.common.basemodel.splitfuse_infer_struct import SplitFuseInferStateInfo
 from lightllm.common.basemodel.layer_weights.base_layer_weight import BaseLayerWeight
 from .cache_tensor_manager import g_cache_manager
 
@@ -16,9 +15,6 @@ def context_forward(self, input_ids, infer_state: InferStateInfo, layer_weight:
     def token_forward(self, input_ids, infer_state: InferStateInfo, layer_weight: BaseLayerWeight):
         raise Exception("need to impl")
 
-    def splitfuse_forward(self, input_ids, infer_state: SplitFuseInferStateInfo, layer_weight: BaseLayerWeight):
-        raise Exception("need to impl")
-
     def alloc_tensor(
         self,
         shape: Union[torch.Size, Iterable[int]],

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

@@ -8,7 +8,6 @@
 from lightllm.utils.infer_utils import mark_cost_time
 
 from ...infer_struct import InferStateInfo
-from ...splitfuse_infer_struct import SplitFuseInferStateInfo
 from ..transformer_layer_infer import TransformerLayerInfer
 
 
@@ -69,11 +68,6 @@ def _context_attention_kernel(self, q, kv, infer_state: InferStateInfo, layer_we
     def _token_attention_kernel(self, q, infer_state: InferStateInfo, layer_weight, out=None) -> torch.Tensor:
         raise Exception("need to impl")
 
-    def _splitfuse_attention_kernel(
-        self, q, infer_state: SplitFuseInferStateInfo, layer_weight, out=None
-    ) -> torch.Tensor:
-        raise Exception("need to impl")
-
     def _get_o(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
         raise Exception("need to impl")
 
@@ -118,25 +112,6 @@ def _token_ffn(self, input_embdings, infer_state: InferStateInfo, layer_weight):
         infer_state._ffn_out = ffn_out
         return
 
-    def _splitfuse_attention(self, input_embding, infer_state: SplitFuseInferStateInfo, layer_weight):
-        cache_kv = self._pre_cache_kv(infer_state, layer_weight)
-        q, cache_kv = self._get_qkv(input_embding, cache_kv, infer_state, layer_weight)
-        self._post_cache_kv(cache_kv, infer_state, layer_weight)
-        o = self._splitfuse_attention_kernel(q, infer_state, layer_weight)
-        q = None
-        o = self._get_o(o, infer_state, layer_weight)
-        if self.world_size_ > 1:
-            dist.all_reduce(o, op=dist.ReduceOp.SUM, async_op=False)
-        infer_state._attn_out = o
-        return
-
-    def _splitfuse_ffn(self, input_embdings, infer_state: SplitFuseInferStateInfo, layer_weight):
-        ffn_out = self._ffn(input_embdings, infer_state, layer_weight)
-        if self.world_size_ > 1:
-            dist.all_reduce(ffn_out, op=dist.ReduceOp.SUM, async_op=False)
-        infer_state._ffn_out = ffn_out
-        return
-
     def _cohere_residual(self, input_embdings, infer_state: InferStateInfo):
         # emb_addr = input_embdings.data_ptr()
         # attn_out_addr = infer_state._attn_out.data_ptr()
@@ -161,10 +136,3 @@ def token_forward(self, input_embdings, infer_state: InferStateInfo, layer_weigh
         self._token_ffn(input1, infer_state, layer_weight)
         self._cohere_residual(input_embdings, infer_state)
         return input_embdings
-
-    def splitfuse_forward(self, input_embdings, infer_state: SplitFuseInferStateInfo, layer_weight):
-        input1 = self._att_norm(input_embdings, infer_state, layer_weight)
-        self._splitfuse_attention(input1, infer_state, layer_weight=layer_weight)
-        self._splitfuse_ffn(input1, infer_state, layer_weight)
-        self._cohere_residual(input_embdings, infer_state)
-        return input_embdings

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

@@ -3,7 +3,6 @@
 import torch.distributed as dist
 from ..transformer_layer_infer import TransformerLayerInfer
 from ...infer_struct import InferStateInfo
-from ...splitfuse_infer_struct import SplitFuseInferStateInfo
 from lightllm.utils.infer_utils import mark_cost_time
 from lightllm.common.basemodel.triton_kernel.destindex_copy_kv import destindex_copy_kv
 from typing import Tuple
@@ -61,11 +60,6 @@ def _context_attention_kernel(self, q, kv, infer_state: InferStateInfo, layer_we
     def _token_attention_kernel(self, q, infer_state: InferStateInfo, layer_weight, out=None) -> torch.Tensor:
         raise Exception("need to impl")
 
-    def _splitfuse_attention_kernel(
-        self, q, infer_state: SplitFuseInferStateInfo, layer_weight, out=None
-    ) -> torch.Tensor:
-        raise Exception("need to impl")
-
     def _get_o(self, input, infer_state: InferStateInfo, layer_weight) -> torch.Tensor:
         raise Exception("need to impl")
 
@@ -118,29 +112,6 @@ def _token_ffn(self, input_embdings, infer_state: InferStateInfo, layer_weight):
         input_embdings.add_(ffn_out.view(-1, self.embed_dim_))
         return
 
-    def _splitfuse_attention(self, input_embding, infer_state: SplitFuseInferStateInfo, layer_weight):
-        input1 = self._att_norm(input_embding, infer_state, layer_weight)
-        cache_kv = self._pre_cache_kv(infer_state, layer_weight)
-        q, cache_kv = self._get_qkv(input1, cache_kv, infer_state, layer_weight)
-        input1 = None
-        self._post_cache_kv(cache_kv, infer_state, layer_weight)
-        o = self._splitfuse_attention_kernel(q, infer_state, layer_weight)
-        q = None
-        o = self._get_o(o, infer_state, layer_weight)
-        if self.world_size_ > 1:
-            dist.all_reduce(o, op=dist.ReduceOp.SUM, async_op=False)
-        input_embding.add_(o.view(-1, self.embed_dim_))
-        return
-
-    def _splitfuse_ffn(self, input_embdings, infer_state: SplitFuseInferStateInfo, layer_weight):
-        input1 = self._ffn_norm(input_embdings, infer_state, layer_weight)
-        ffn_out = self._ffn(input1, infer_state, layer_weight)
-        input1 = None
-        if self.world_size_ > 1:
-            dist.all_reduce(ffn_out, op=dist.ReduceOp.SUM, async_op=False)
-        input_embdings.add_(ffn_out.view(-1, self.embed_dim_))
-        return
-
     def context_forward(self, input_embdings, infer_state: InferStateInfo, layer_weight):
         self._context_attention(input_embdings, infer_state, layer_weight=layer_weight)
         self._context_ffn(input_embdings, infer_state, layer_weight)
@@ -150,8 +121,3 @@ def token_forward(self, input_embdings, infer_state: InferStateInfo, layer_weigh
         self._token_attention(input_embdings, infer_state, layer_weight=layer_weight)
         self._token_ffn(input_embdings, infer_state, layer_weight)
         return input_embdings
-
-    def splitfuse_forward(self, input_embdings, infer_state: SplitFuseInferStateInfo, layer_weight):
-        self._splitfuse_attention(input_embdings, infer_state, layer_weight=layer_weight)
-        self._splitfuse_ffn(input_embdings, infer_state, layer_weight)
-        return input_embdings