deepseek-MTP eagle, topk=1 (#1073)

Co-authored-by: hiworldwzj <[email protected]>
Co-authored-by: wangzaijun <[email protected]>

Author: 3 people

docs/CN/source/tutorial/api_server_args_zh.rst

@@ -445,9 +445,10 @@ MTP 多预测参数
 
 .. option:: --mtp_mode
 
-    支持的 mtp 模式，可选值：
+    支持的 mtp 模式，建议使用 deepseekv3_eagle获得更好的性能体验，可选值：
 
-    * ``deepseekv3``
+    * ``deepseekv3_vanilla``
+    * ``deepseekv3_eagle``
     * ``None``: 不启用 mtp（默认）
 
 .. option:: --mtp_draft_model_dir

docs/EN/source/tutorial/api_server_args_zh.rst

@@ -442,9 +442,10 @@ MTP Multi-Prediction Parameters
 
 .. option:: --mtp_mode
 
-    Supported mtp modes, optional values:
+    Supported mtp modes, it is recommended to use deepseekv3_eagle for better performance, optional values:
 
-    * ``deepseekv3``
+    * ``deepseekv3_vanilla``
+    * ``deepseekv3_eagle``
     * ``None``: Do not enable mtp (default)
 
 .. option:: --mtp_draft_model_dir

lightllm/common/basemodel/basemodel.py

@@ -67,12 +67,16 @@ 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.data_type = kvargs.get("data_type", "float16")
+        mtp_step = get_env_start_args().mtp_step
         self.graph_max_batch_size = kvargs.get("graph_max_batch_size", 16)
         self.graph_max_batch_size = (
             self.graph_max_batch_size // 2
             if get_env_start_args().enable_decode_microbatch_overlap
             else self.graph_max_batch_size
         )
+        # mtp 模式下需要修缮对应的最大batch size，为 （mtp_step + 1) 的倍数
+        self.graph_max_batch_size = self.graph_max_batch_size * (mtp_step + 1)
+
         self.graph_max_len_in_batch = kvargs.get("graph_max_len_in_batch", 8192)
         self.disable_cudagraph = kvargs.get("disable_cudagraph", False)
         self.quant_type = kvargs.get("quant_type", "none")
@@ -81,7 +85,7 @@ def __init__(self, kvargs):
         self.tp_world_size_ = get_dp_world_size()
         self.enable_tpsp_mix_mode = get_env_start_args().enable_tpsp_mix_mode
 
-        self.is_deepseekv3_mtp_mode = self.args.mtp_mode == "deepseekv3"
+        self.is_deepseekv3_mtp_mode = self.args.mtp_mode in ["deepseekv3_vanilla", "deepseekv3_eagle"]
 
         self._init_datatype()
         self._init_config()
@@ -258,6 +262,10 @@ def _create_inferstate(self, model_input: ModelInput, microbatch_index: int = 0)
         infer_state.batch_size = model_input.batch_size
         infer_state.total_token_num = model_input.total_token_num
         infer_state.max_len_in_batch = model_input.max_len_in_batch
+        infer_state.max_q_seq_len = model_input.max_q_seq_len
+        infer_state.max_kv_seq_len = model_input.max_kv_seq_len
+        infer_state.max_cache_len = model_input.max_cache_len
+        infer_state.prefix_total_token_num = model_input.prefix_total_token_num
         assert model_input.b_req_idx.shape[0] == model_input.b_seq_len.shape[0]
         infer_state.b_req_idx = model_input.b_req_idx
         infer_state.b_seq_len = model_input.b_seq_len
@@ -335,16 +343,16 @@ def _prefill(
         model_input: ModelInput,
     ):
         infer_state = self._create_inferstate(model_input)
+        infer_state.init_some_extra_state(self, model_input.input_ids)
         init_req_to_token_indexes(
-            self.req_manager.req_to_token_indexs,
-            model_input.b_req_idx,
-            model_input.b_seq_len,
-            infer_state.b_ready_cache_len,
-            model_input.max_len_in_batch,
-            infer_state.mem_index,
+            req_to_token_indexs=self.req_manager.req_to_token_indexs,
+            b_req_idx=infer_state.b_req_idx,
+            b_seq_len=infer_state.b_seq_len,
+            b_ready_cache_len=infer_state.b_ready_cache_len,
+            b_start_loc=infer_state.b_start_loc,
+            alloc_mem_index=infer_state.mem_index,
+            max_q_seq_len=infer_state.max_q_seq_len,
         )
-
-        infer_state.init_some_extra_state(self, model_input.input_ids)
         return self._context_forward(model_input.input_ids, infer_state)
 
     def _decode(
@@ -474,26 +482,28 @@ def microbatch_overlap_prefill(self, model_input0: ModelInput, model_input1: Mod
         input_ids0, input_ids1 = model_input0.input_ids, model_input1.input_ids
 
         infer_state0 = self._create_inferstate(model_input0, 0)
+        infer_state0.init_some_extra_state(self, input_ids0)
         init_req_to_token_indexes(
-            self.req_manager.req_to_token_indexs,
-            model_input0.b_req_idx,
-            model_input0.b_seq_len,
-            infer_state0.b_ready_cache_len,
-            model_input0.max_len_in_batch,
-            infer_state0.mem_index,
+            req_to_token_indexs=self.req_manager.req_to_token_indexs,
+            b_req_idx=infer_state0.b_req_idx,
+            b_seq_len=infer_state0.b_seq_len,
+            b_ready_cache_len=infer_state0.b_ready_cache_len,
+            b_start_loc=infer_state0.b_start_loc,
+            alloc_mem_index=infer_state0.mem_index,
+            max_q_seq_len=infer_state0.max_q_seq_len,
         )
-        infer_state0.init_some_extra_state(self, input_ids0)
 
         infer_state1 = self._create_inferstate(model_input1, 1)
+        infer_state1.init_some_extra_state(self, input_ids1)
         init_req_to_token_indexes(
-            self.req_manager.req_to_token_indexs,
-            model_input1.b_req_idx,
-            model_input1.b_seq_len,
-            infer_state1.b_ready_cache_len,
-            model_input1.max_len_in_batch,
-            infer_state1.mem_index,
+            req_to_token_indexs=self.req_manager.req_to_token_indexs,
+            b_req_idx=infer_state1.b_req_idx,
+            b_seq_len=infer_state1.b_seq_len,
+            b_ready_cache_len=infer_state1.b_ready_cache_len,
+            b_start_loc=infer_state1.b_start_loc,
+            alloc_mem_index=infer_state1.mem_index,
+            max_q_seq_len=infer_state1.max_q_seq_len,
         )
-        infer_state1.init_some_extra_state(self, input_ids1)
 
         model_output0, model_output1 = self._overlap_tpsp_context_forward(
             input_ids0, infer_state0, input_ids1=input_ids1, infer_state1=infer_state1
@@ -521,7 +531,7 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
                 model_input1.b_req_idx,
                 model_input1.b_mtp_index,
             )
-
+        # TODO 动态 mtp fix
         assert model_input0.batch_size == model_input1.batch_size
         assert model_input0.mem_indexes.is_cuda
         assert model_input1.mem_indexes.is_cuda
@@ -531,6 +541,8 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
 
         if self.graph is not None and self.graph.can_run(origin_batch_size, max_len_in_batch):
             find_graph_batch_size = self.graph.find_closest_graph_batch_size(origin_batch_size)
+            # TODO 如果支持动态步数的 mtp，在不同的mtp步上，model_input0 和 model_input1 的内部batch size可能不
+            # 一致，需要按照较高 batch size 进行graph的寻找，同时，进行有效的恢复。
             padded_model_input0 = self._create_padded_decode_model_input(model_input0, find_graph_batch_size)
             padded_model_input1 = self._create_padded_decode_model_input(model_input1, find_graph_batch_size)
             infer_state0 = self._create_inferstate(padded_model_input0, 0)
@@ -568,6 +580,8 @@ def microbatch_overlap_decode(self, model_input0: ModelInput, model_input1: Mode
                     input_ids1=padded_model_input1.input_ids,
                     infer_state1=infer_state1,
                 )
+
+            # TODO 动态 mtp fix
             model_output0 = self._create_unpad_decode_model_output(model_output0, origin_batch_size=origin_batch_size)
             model_output1 = self._create_unpad_decode_model_output(model_output1, origin_batch_size=origin_batch_size)
         else:
@@ -696,6 +710,10 @@ def _check_max_len_infer(self):
                 batch_size=1,
                 total_token_num=total_token_num,
                 max_len_in_batch=self.batch_max_tokens,
+                max_q_seq_len=self.batch_max_tokens,
+                max_kv_seq_len=self.batch_max_tokens,
+                max_cache_len=0,
+                prefix_total_token_num=0,
                 input_ids=dummy_input_ids,
                 mem_indexes=mem_indexes,
                 b_req_idx=b_req_idx,
@@ -766,6 +784,10 @@ def _autotune_warmup(self):
                     batch_size=1,
                     total_token_num=total_token_num,
                     max_len_in_batch=input_len,
+                    max_q_seq_len=input_len,
+                    max_kv_seq_len=input_len,
+                    max_cache_len=0,
+                    prefix_total_token_num=0,
                     input_ids=dummy_input_ids,
                     mem_indexes=mem_indexes,
                     b_req_idx=b_req_idx,
@@ -822,6 +844,10 @@ def _init_padded_req(self):
             batch_size=batch_size,
             total_token_num=total_token_num,
             max_len_in_batch=prefill_input_len,
+            max_q_seq_len=prefill_input_len,
+            max_kv_seq_len=prefill_input_len,
+            max_cache_len=0,
+            prefix_total_token_num=0,
             input_ids=dummy_input_ids,
             mem_indexes=mem_indexes,
             b_req_idx=b_req_idx,

lightllm/common/basemodel/batch_objs.py

@@ -10,10 +10,17 @@ class ModelInput:
     batch_size: int
     total_token_num: int
     max_len_in_batch: int
-    input_ids: torch.Tensor
-    b_req_idx: torch.Tensor
-    b_mtp_index: torch.Tensor
-    b_seq_len: torch.Tensor
+    # 在 decode 阶段， 常规模式下， max_q_seq_len 必定是 1，
+    # 在 mtp 模式下，max_q_seq_len 统计的是一个请求考虑了 mtp 步数的
+    # 最大长度，实际值是 max([(1 + req.mtp_step) for req in reqs])
+    max_q_seq_len: int
+    max_kv_seq_len: int
+    max_cache_len: int = None
+    prefix_total_token_num: int = None
+    input_ids: torch.Tensor = None
+    b_req_idx: torch.Tensor = None
+    b_mtp_index: torch.Tensor = None
+    b_seq_len: torch.Tensor = None
     mem_indexes: torch.Tensor = None
     is_prefill: bool = False
     b_ready_cache_len: torch.Tensor = None

lightllm/common/basemodel/cuda_graph.py

@@ -19,27 +19,27 @@ class CudaGraph:
     def __init__(self, max_batch_size=8, max_len_in_batch=8192):
         self.graph = {}
         self.mempool = torch.cuda.graph_pool_handle() if torch.cuda.is_available() else None
+        self.args = get_env_start_args()
+        self.mtp_step = self.args.mtp_step
         self.max_batch_size = max_batch_size
         self.graph_max_len_in_batch = max_len_in_batch
-        self.args = get_env_start_args()
         self.enable_decode_microbatch_overlap = self.args.enable_decode_microbatch_overlap
 
         # gen cuda graph batch_sizes
         # cuda graph gen for batch size = [1, 2, 3, ..., graph_split_batch_size]
         # and [graph_split_batch_size + graph_grow_step_size,
-        # graph_split_batch_size + 2 * graph_grow_step_size,  ...,  self.max_batch_size]
-        graph_split_batch_size = self.args.graph_split_batch_size
-        max_batch_size = self.max_batch_size
-        graph_grow_step_size = self.args.graph_grow_step_size
+        # if the mtp_step is not 0, then the batch_sizes will be multiply of (mtp_step + 1)
 
-        batch_sizes = [i for i in range(1, graph_split_batch_size + 1)]
+        graph_split_batch_size = self.args.graph_split_batch_size * (self.mtp_step + 1)
+        graph_grow_step_size = self.args.graph_grow_step_size * (self.mtp_step + 1)
+
+        batch_sizes = [i * (self.mtp_step + 1) for i in range(1, graph_split_batch_size + 1)]
         for _batch_size in range(graph_split_batch_size + graph_grow_step_size, max_batch_size, graph_grow_step_size):
             batch_sizes.append(_batch_size)
 
         batch_sizes = list(set([e for e in batch_sizes if e < max_batch_size]))
         batch_sizes.append(max_batch_size)
         batch_sizes.sort()
-
         self.cuda_graph_batch_sizes = batch_sizes
         assert batch_sizes[-1] == self.max_batch_size
         logger.info(f"cuda graph batch_sizes: {self.cuda_graph_batch_sizes}")
@@ -208,6 +208,8 @@ def warmup(self, model):
                 batch_size=batch_size,
                 total_token_num=total_token_num,
                 max_len_in_batch=max_len_in_batch,
+                max_q_seq_len=self.mtp_step + 1,
+                max_kv_seq_len=max_len_in_batch,
                 input_ids=input_ids,
                 mem_indexes=mem_indexes,
                 b_req_idx=b_req_idx,
@@ -265,6 +267,8 @@ def warmup_overlap(self, model):
                     batch_size=batch_size,
                     total_token_num=total_token_num,
                     max_len_in_batch=max_len_in_batch,
+                    max_q_seq_len=self.mtp_step + 1,
+                    max_kv_seq_len=max_len_in_batch,
                     input_ids=input_ids,
                     b_mtp_index=b_mtp_index,
                     mem_indexes=mem_indexes,

lightllm/common/basemodel/infer_struct.py

@@ -25,6 +25,11 @@ def __init__(self):
         # prefill 阶段指每个req 输入token的长度（不包括已经cache的部分）最大值
         # decode 阶段指的是每个req的总长 最大值
         self.max_len_in_batch: int = None
+        # max_cache_len 用于 prefill 阶段标识请求中最大 cache的kv 的长度
+        self.max_cache_len: int = None
+        # prefix_total_token_num 用于 prefill 阶段标识当前请求中所有已经ready的kv的长度
+        # 的sum值, 其值等于 sum(b_ready_cache_len)
+        self.prefix_total_token_num: int = None
         self.is_prefill: bool = None
 
         self.mem_manager: MemoryManager = None
@@ -72,8 +77,6 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                 self.b_kv_seq_len,
                 self.b1_cu_kv_seq_len,
                 self.position_ids,
-                self.max_q_seq_len,
-                self.max_kv_seq_len,
             ) = gen_prefill_params(
                 input_token_num=input_ids.shape[0],
                 b_ready_cache_len=self.b_ready_cache_len,
@@ -88,7 +91,6 @@ def init_some_extra_state(self, model, input_ids: torch.Tensor):
                 self.b1_cu_kv_seq_len,
                 self.position_ids,
             ) = gen_decode_params(self.b_seq_len)
-            self.max_q_seq_len = 1
             # TODO: check the correctness
             self.max_kv_seq_len = self.max_len_in_batch
             self.b_start_loc = self.b1_cu_kv_seq_len[0:-1]

lightllm/common/basemodel/triton_kernel/copy_kv_index_to_req.py

@@ -6,8 +6,7 @@
 
 @triton.jit
 def _fwd_kernel_copy_kv_index_to_req(
-    req_to_token_indexs, b_req_idx, b_seq_len, memindex,
-    stride_req_to_token_b, stride_req_to_token_s
+    req_to_token_indexs, b_req_idx, b_seq_len, memindex, stride_req_to_token_b, stride_req_to_token_s
 ):
     cur_index = tl.program_id(0)
     cur_req_idx = tl.load(b_req_idx + cur_index)
@@ -26,8 +25,86 @@ def copy_kv_index_to_req(req_to_token_indexs, b_req_idx, b_seq_len, memindex):
     num_warps = 1
 
     _fwd_kernel_copy_kv_index_to_req[grid](
-        req_to_token_indexs, b_req_idx, b_seq_len, memindex,
-        req_to_token_indexs.stride(0), req_to_token_indexs.stride(1),
+        req_to_token_indexs,
+        b_req_idx,
+        b_seq_len,
+        memindex,
+        req_to_token_indexs.stride(0),
+        req_to_token_indexs.stride(1),
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    return
+
+
+@triton.jit
+def _fwd_kernel_copy_kv_index_to_req_prefill(
+    req_to_token_indexs,
+    b_req_idx,
+    b_seq_len,
+    b_ready_cache_len,
+    b_start_loc,
+    memindex,
+    stride_req_to_token_b,
+    stride_req_to_token_s,
+    BLOCK: tl.constexpr,
+):
+
+    block_index = tl.program_id(0)
+    batch_index = tl.program_id(1)
+    cur_req_idx = tl.load(b_req_idx + batch_index)
+    cur_seq_len = tl.load(b_seq_len + batch_index)
+    cur_ready_cache_len = tl.load(b_ready_cache_len + batch_index)
+    cur_start_loc = tl.load(b_start_loc + batch_index)
+    copy_len = cur_seq_len - cur_ready_cache_len
+
+    block_range = block_index * BLOCK + tl.arange(0, BLOCK)
+    block_mask = block_range < copy_len
+    cur_token_index = tl.load(memindex + cur_start_loc + block_range, mask=block_mask)
+    dest_offset = (
+        req_to_token_indexs
+        + cur_req_idx * stride_req_to_token_b
+        + (cur_ready_cache_len + block_range) * stride_req_to_token_s
+    )
+    tl.store(dest_offset, cur_token_index, mask=block_mask)
+
+    return
+
+
+def get_triton_config(max_q_seq_len: int) -> tuple[int, int]:
+    if max_q_seq_len <= 512:
+        return 256, 2
+    elif max_q_seq_len <= 4096:
+        return 512, 4
+    else:
+        return 1024, 8
+
+
+@torch.no_grad()
+def copy_kv_index_to_req_prefill(
+    req_to_token_indexs: torch.Tensor,
+    b_req_idx: torch.Tensor,
+    b_seq_len: torch.Tensor,
+    b_ready_cache_len: torch.Tensor,
+    b_start_loc: torch.Tensor,
+    memindex: torch.Tensor,
+    max_q_seq_len: int,
+):
+    batch_size = b_req_idx.shape[0]
+    BLOCK, num_warps = get_triton_config(max_q_seq_len)
+    grid = (triton.cdiv(max_q_seq_len, BLOCK), batch_size)
+    num_warps = 1
+
+    _fwd_kernel_copy_kv_index_to_req_prefill[grid](
+        req_to_token_indexs,
+        b_req_idx,
+        b_seq_len,
+        b_ready_cache_len,
+        b_start_loc,
+        memindex,
+        req_to_token_indexs.stride(0),
+        req_to_token_indexs.stride(1),
+        BLOCK=BLOCK,
         num_warps=num_warps,
         num_stages=1,
     )