feat: dynamically split PD decode requests

deploy_dspark_1p1d.sh

@@ -28,7 +28,6 @@ export LIGHTLLM_TRITON_AUTOTUNE_LEVEL="${LIGHTLLM_TRITON_AUTOTUNE_LEVEL:-1}"
 export LIGHTLLM_ANTHROPIC_ENABLE_PDF_PARSING="${LIGHTLLM_ANTHROPIC_ENABLE_PDF_PARSING:-1}"
 export LIGHTLLM_LOG_LEVEL="${LIGHTLLM_LOG_LEVEL:-debug}"
 export PYTHONUNBUFFERED=1
-export LIGHTLLM_PD_SPLIT_MAX_NEW_TOKENS="${LIGHTLLM_PD_SPLIT_MAX_NEW_TOKENS:-4096}"
 
 # Keep same-host PD WebSocket traffic away from HTTP/WebSocket proxies.
 export NO_PROXY="${NO_PROXY:+${NO_PROXY},}${PD_MASTER_IP},localhost"

lightllm/server/api_start.py

@@ -84,10 +84,10 @@ def _launch_subprocesses(args: StartArgs):
 
     # 调度参数的自动设置, 人工设置则听人工的
     if args.router_token_ratio is None:
-        if args.run_mode in ["normal"]:
+        if args.run_mode in ["normal", "decode"]:
             args.router_token_ratio = 0.85
         else:
-            # pd 分离模式下，不开启高级调度
+            # PD 分离模式下，prefill 节点不开启高级调度
             args.router_token_ratio = 0.0
     # 部分模式还不能支持与高级动态调度算法协同，to do.
     if args.diverse_mode:

lightllm/server/core/objs/req.py

@@ -27,6 +27,8 @@ class FinishStatus(ctypes.Structure):
     - ``NO_FINISH``: 未结束
     - ``FINISHED_STOP``: 正常停止（EOS / stop 序列等），finish_reason=``stop``
     - ``FINISHED_LENGTH``: 达到 max_new_tokens 等长度上限，finish_reason=``length``
+    - ``FINISHED_PD_DECODE_CAPACITY``: PD Decode 因 token 容量不足提前结束当前分段。
+      该状态仅用于 PD 内部分段续跑，不应直接暴露给 API 用户。
     - ``FINISHED_ABORTED``: 客户端/调度主动 abort，finish_reason=``abort``
     - ``FINISHED_ERROR``: 服务端内部错误导致无法继续生成，finish_reason=``error``。
       典型场景：PD 分离 decode 节点 KV 传输失败。与 abort 区分：非用户取消，
@@ -39,6 +41,9 @@ class FinishStatus(ctypes.Structure):
     NO_FINISH = 0
     FINISHED_STOP = 1
     FINISHED_LENGTH = 2
+    # PD Decode 因 token 容量不足时，用此状态结束当前 segment。PD Master 会吞掉
+    # 模拟结束 token，并携带剩余 max_new_tokens 继续下一个 segment。值 5 用于保持旧状态编号兼容。
+    FINISHED_PD_DECODE_CAPACITY = 5
     FINISHED_ABORTED = 3
     # 内部错误结束（如 PD KV 传输失败）；见类文档。
     FINISHED_ERROR = 4
@@ -47,14 +52,14 @@ def __init__(self, init_state=NO_FINISH):
         self.status = init_state
 
     def set_status(self, new_status):
-        assert 0 <= new_status <= 4
+        assert 0 <= new_status <= self.FINISHED_PD_DECODE_CAPACITY
         self.status = new_status
 
     def get_status(self):
         return self.status
 
     def is_finished(self):
-        return self.FINISHED_STOP <= self.status <= self.FINISHED_ERROR
+        return self.FINISHED_STOP <= self.status <= self.FINISHED_PD_DECODE_CAPACITY
 
     def is_stopped(self):
         return self.status == self.FINISHED_STOP
@@ -65,6 +70,9 @@ def is_finished_length(self):
     def is_finished_error(self):
         return self.status == self.FINISHED_ERROR
 
+    def is_finished_pd_decode_capacity(self):
+        return self.status == self.FINISHED_PD_DECODE_CAPACITY
+
     def is_error_finished(self):
         return self.status in (self.FINISHED_ABORTED, self.FINISHED_ERROR)
 
@@ -77,6 +85,9 @@ def get_finish_reason(self):
             return "abort"
         elif self.status == self.FINISHED_ERROR:
             return "error"
+        elif self.status == self.FINISHED_PD_DECODE_CAPACITY:
+            # 该内部状态正常会被 PD Master 吞掉；泄漏时按 length 降级处理。
+            return "length"
         return None
 
 

lightllm/server/httpserver_for_pd_master/manager.py

@@ -23,7 +23,7 @@
 from lightllm.utils.statics_utils import MovingAverage
 from lightllm.server.httpserver.manager import AsyncQueue
 from lightllm.utils.error_utils import ClientDisconnected, ServerBusyError
-from lightllm.utils.envs_utils import get_pd_high_priority_request_timeout_seconds, get_pd_split_max_new_tokens
+from lightllm.utils.envs_utils import get_pd_high_priority_request_timeout_seconds
 from lightllm.utils.shm_port_args import get_shm_port_args
 from .pd_selector import create_selector
 
@@ -215,13 +215,6 @@ async def _generate_one(
         start_time: float,
         origin_request_id: int,
     ):
-        # 先将请求根据max_new_tokens 参数进行分块操作，主要是 pd 分离场景中，
-        # 只能使用保守调度，但是如果用户都设置一个很大的 max_new_tokens 值，会
-        # 导致极大显存预留，照成系统的吞吐能力下降，所以我们将请求分割成几段进行
-        # 推理，只要保证分块合理，实际分段推理是极少发生的情况，系统吞吐就不会受
-        # 到影响。
-        max_new_tokens_list = self._split_max_new_tokens(max_new_tokens=origin_sampling_params.max_new_tokens)
-
         block_group_request_id = origin_request_id
         p_node = None
         d_node = None
@@ -237,17 +230,24 @@ async def _generate_one(
 
             history_gen_token_strs = []
             origin_prompt_cache_len = None
-
-            for iter_index, block_max_new_tokens in enumerate(max_new_tokens_list):
+            remaining_max_new_tokens = origin_sampling_params.max_new_tokens
+            segment_index = 0
+            # 后续分段的 prompt 会追加已生成内容；始终保留所有分段中最小的 prompt token 数，
+            # 对外 usage 才能反映用户的原始输入长度，而不是最后一次续跑的 block_prompt 长度。
+            prompt_tokens = sys.maxsize
+
+            # Decode 节点容量不足时会用专用状态结束当前分段。
+            # PD Master 吞掉该内部分段 marker，并用剩余 token 限额在同一组 P/D 节点上继续。
+            while remaining_max_new_tokens > 0:
                 sampling_params = SamplingParams.from_buffer_copy(origin_sampling_params)
                 block_group_request_id = self.id_gen.generate_id()
                 sampling_params.group_request_id = block_group_request_id
                 logger.info(f"pd log gen sub req id {block_group_request_id} for main req id {origin_request_id}")
-                sampling_params.max_new_tokens = block_max_new_tokens
+                sampling_params.max_new_tokens = remaining_max_new_tokens
                 # 预计输入 cache 命中率高于 0.8 时，将首段也标记为 PD 高优先级请求，
                 # 使其优先进入 Router 调度队列，尽快复用已命中的 KV cache。第二段及后续
                 # 分段仍统一使用高优先级，避免因临时资源紧张导致分段续跑失败。
-                sampling_params.pd_high_priority_request = iter_index > 0 or estimated_cache_hit_rate > 0.8
+                sampling_params.pd_high_priority_request = segment_index > 0 or estimated_cache_hit_rate > 0.8
                 # 为高优先级请求下发较长的有限等待时间；P/D 节点仅在自身开启
                 # 本地限流时使用该值，未开启限流时仍保持无限等待。
                 if sampling_params.pd_high_priority_request:
@@ -270,35 +270,44 @@ async def _generate_one(
                     multimodal_params,
                     request,
                 )
-                is_last_block = iter_index == len(max_new_tokens_list) - 1
-                prompt_tokens = sys.maxsize  # 因为分段的原因
-                async for sub_req_id, request_output, metadata, finish_status in results_generator:
-                    # pd 分离模式下，返回的 metadata 可能序号信息可能存在不准确性。
+                raw_finish_status = FinishStatus()
+                async for sub_req_id, request_output, metadata, raw_finish_status in results_generator:
+                    # PD 分离模式下 metadata 中的 token 序号可能不准确，按实际产出计数。
                     assert sub_req_id == block_group_request_id
-                    if finish_status.is_finished_length() and not is_last_block:
-                        finish_status = FinishStatus()  # 转换为NoFinished
+
+                    # 收到当前分段的任意输出，说明该请求已经完成 P 节点的 prefill 派发阶段。
+                    # 立即归还 selector 中记录的在途 prompt 字符数和请求数，并通过置空确保每段只更新一次。
+                    if pending_prefill_load_chars is not None:
+                        p_node.dispatched_prompt_chars = max(
+                            0, p_node.dispatched_prompt_chars - pending_prefill_load_chars
+                        )
+                        p_node.dispatched_req_num = max(0, p_node.dispatched_req_num - 1)
+                        pending_prefill_load_chars = None
+
+                    if raw_finish_status.is_finished_pd_decode_capacity():
+                        # 容量不足状态是 PD 内部分段边界：吞掉模拟结束 token，继续生成剩余 token。
+                        break
+
+                    # 容量 marker 已在上方过滤，能走到这里的每个 token 都立即扣减全局剩余输出额度。
+                    remaining_max_new_tokens -= 1
                     history_gen_token_strs.append(request_output)
                     prompt_tokens = min(prompt_tokens, metadata["prompt_tokens"])
                     metadata["prompt_tokens"] = prompt_tokens
-                    if iter_index == 0 and origin_prompt_cache_len is None:
+                    if origin_prompt_cache_len is None:
                         origin_prompt_cache_len = metadata.get("prompt_cache_len", 0)
                         prompt_cache_hit_rate = origin_prompt_cache_len / max(prompt_tokens, 1)
                         self.pd_manager.selector.record_prompt_cache_hit_rate(prompt_cache_hit_rate)
-                        if not finish_status.is_error_finished():
+                        if not raw_finish_status.is_error_finished():
                             # 只有收到成功的推理结果后才将 prompt 写入前缀树，避免尚未进入
                             # 推理或已失败的请求被后续请求误判为可复用 cache。
                             self.pd_manager.selector.insert_prompt_cache(prompt, p_node)
                     metadata["prompt_cache_len"] = origin_prompt_cache_len or 0
-                    if pending_prefill_load_chars is not None:
-                        p_node.dispatched_prompt_chars = max(
-                            0, p_node.dispatched_prompt_chars - pending_prefill_load_chars
-                        )
-                        p_node.dispatched_req_num = max(0, p_node.dispatched_req_num - 1)
-                        pending_prefill_load_chars = None
-                    yield origin_request_id, request_output, metadata, finish_status
+                    yield origin_request_id, request_output, metadata, raw_finish_status
 
                 await self.remove_req(group_request_id=block_group_request_id)
-                if finish_status.is_finished():
+                segment_index += 1
+                # 只有 PD Decode 容量不足产生的内部分段需要续跑；其他状态都结束整个请求。
+                if not raw_finish_status.is_finished_pd_decode_capacity():
                     break
 
         except (ClientDisconnected, BaseException) as e:
@@ -718,14 +727,6 @@ async def handle_loop(self):
                 logger.exception(str(e))
         return
 
-    def _split_max_new_tokens(self, max_new_tokens: int) -> List[int]:
-        block_max_new_tokens = get_pd_split_max_new_tokens()
-        ans_list = [block_max_new_tokens for _ in range(max_new_tokens // block_max_new_tokens)]
-        left_token = max_new_tokens - (max_new_tokens // block_max_new_tokens) * block_max_new_tokens
-        if left_token > 0:
-            ans_list.append(left_token)
-        return ans_list
-
 
 class ReqStatus:
     def __init__(self, req_id, p_node, d_node) -> None:

lightllm/server/router/model_infer/infer_batch.py

@@ -907,11 +907,20 @@ def mark_shm_aborted_finished(self):
         ``Req.mark_simulated_finished``（在已有输出末尾追加 EOS）。不回写本地
         ``cur_output_len`` / ``finish_status``（本 InferReq 即将释放）。
         """
-        # 本地已由 stop / eos / length 等正确结束，保留原 shm finish_status。
+        # 请求本身已由 stop / eos / length / error 等状态结束时，
+        # 请求自身的结束原因优先，不能被后续的容量不足标记覆盖。
         if self.finish_status.is_finished():
             return
+
+        # 仅在请求本身尚未结束时，才将 finished_by_pd_decode_capacity
+        # 转换为 PD 内部分段状态，补模拟结束 token 并交给 PD Master 续跑。
+        if getattr(self, "finished_by_pd_decode_capacity", False):
+            finish_status = FinishStatus.FINISHED_PD_DECODE_CAPACITY
+        else:
+            finish_status = FinishStatus.FINISHED_ABORTED
+
         self.shm_req.mark_simulated_finished(
-            FinishStatus.FINISHED_ABORTED,
+            finish_status,
             output_len=self.cur_output_len,
         )
         return

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -694,10 +694,8 @@ def _get_classed_reqs(
         prefill_reqs = []
         decode_reqs = []
 
-        # 一次性最多暂停请求的数量, 防止盲目暂停大量请求
-        # 因为部分请求释放占用的token容量后，就会使推理可以正常进行。
-        # 如果因为一次推理容量不足，就以当前token容量的判断暂停了大量
-        # 请求，其逻辑是不适合的。
+        # 单轮最多处理少量因 token 容量不足而无法继续的请求，避免一次性影响大量请求。
+        # 普通 Decode 请求进入暂停队列等待恢复；PD Decode 请求则强制提前结束并进入清理流程。
         pause_max_req_num = 2
         wait_pause_count = 0
         prefill_tokens = 0
@@ -741,8 +739,24 @@ def _get_classed_reqs(
                     can_alloc_token_num -= token_num
                 else:
                     if wait_pause_count < pause_max_req_num:
-                        req_obj.wait_pause = True
-                        wait_pause_count += 1
+                        if self.args.run_mode == "decode":
+                            # PD Decode 节点的 token 容量不足时，强制当前请求提前结束以释放资源。
+                            # 单轮只处理 pause_max_req_num 个请求，避免所有资源不足的请求同时退出。
+                            wait_pause_count += 1
+                            setattr(req_obj, "finished_by_pd_decode_capacity", True)
+                            if support_overlap:
+                                # overlap 模式可能仍有异步计算在访问请求，先标记，下一轮再安全清理。
+                                req_obj.filter_mark = True
+                            else:
+                                # 非 overlap 模式没有在途的异步计算，可以在本轮直接清理。
+                                finished_reqs.append(req_obj)
+                            self.logger.info(
+                                f"force early finish for PD decode req_id={req_obj.req_id} "
+                                f"because token capacity is insufficient"
+                            )
+                        else:
+                            req_obj.wait_pause = True
+                            wait_pause_count += 1
             else:
                 # 在 diverse mode 模式下，prefill 只会使用 master 状态的请求，slave 请求依靠后续
                 # 的推理代码中将master请求的状态复制到slave请求中去， 所以这里 slave 状态的请求，不

lightllm/server/router/req_queue/__init__.py

@@ -5,15 +5,15 @@
 
 
 def _get_req_queue_class(args, router, dp_size_in_node: int):
+    if args.run_mode in ["prefill", "decode"]:
+        return PDQueue
+
     if args.diverse_mode:
         return ChunkedBeamContinuesBatchQueue
     if args.output_constraint_mode != "none":
         return ChunkedPrefillQueue
     if args.first_token_constraint_mode:
         return ChunkedPrefillQueue
-    if args.run_mode in ["prefill", "decode"]:
-        return PDQueue
-
     if args.disable_chunked_prefill:
         # 虽然也使用chuncked prefill queue 但是由于 args.chunked_prefill_size = args.max_req_total_len
         # 所以调度的实际行为类似过去的 continues batch 调度，所以将两种调度的实现统一为一种实现，减少代码重复。

lightllm/server/router/req_queue/chunked_prefill/impl_for_pd.py

@@ -11,7 +11,14 @@ def __init__(self, args, router, dp_index, dp_size_in_node) -> None:
 
     # @calculate_time(show=True, min_cost_ms=0.1)
     def _can_add_new_req(self, req: Req, estimated_peak_token_num: int, batch_req_num: int) -> Tuple[bool, int, int]:
-        estimated_peak_token_num += req.input_len + req.sample_params.max_new_tokens
+        if self.args.run_mode == "decode":
+            estimated_output_len = min(
+                self.router.router_statics.ema_req_out_len,
+                req.sample_params.max_new_tokens,
+            )
+        else:
+            estimated_output_len = req.sample_params.max_new_tokens
+        estimated_peak_token_num += req.input_len + estimated_output_len
         ok_token_num = estimated_peak_token_num < self.max_total_tokens
         batch_req_num += 1
         ok_req_num = batch_req_num <= self.running_max_req_size
@@ -38,7 +45,14 @@ def _caclu_batch_estimated_peak_token_num(self, batch: Batch):
                             req.get_tuple_tokens(is_busy, self.router.router_statics.ema_req_out_len)
                         )
                     else:
-                        estimated_peak_token_num += req.input_len + req.sample_params.max_new_tokens
+                        if self.args.run_mode == "decode":
+                            estimated_output_len = min(
+                                self.router.router_statics.ema_req_out_len,
+                                req.sample_params.max_new_tokens,
+                            )
+                        else:
+                            estimated_output_len = req.sample_params.max_new_tokens
+                        estimated_peak_token_num += req.input_len + estimated_output_len
 
         if decoding_req_list:
             decoding_req_list.sort(key=lambda x: -x[1])

lightllm/utils/envs_utils.py

@@ -306,11 +306,6 @@ def _get_mtp_draft_backbone_layer_num(draft_model_dir: str) -> int:
     return int(layer_num)
 
 
-@lru_cache(maxsize=None)
-def get_pd_split_max_new_tokens() -> int:
-    return int(os.getenv("LIGHTLLM_PD_SPLIT_MAX_NEW_TOKENS", 2048))
-
-
 @lru_cache(maxsize=None)
 def get_pd_node_shm_req_alloc_timeout_seconds() -> int:
     """PD 节点申请 ``shm_req`` 对象的最长等待时间，单位为秒。"""