From 5122d21e14c247958f824176ec8614fe561f01d7 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 23 Sep 2026 08:42:03 +0000 Subject: [PATCH 1/2] feat: add configurable prefill queue scheduling --- docs/CN/source/index.rst | 1 + docs/CN/source/tutorial/api_server_args.rst | 34 ++- .../source/tutorial/prefill_queue_strategy.md | 200 ++++++++++++++++++ docs/EN/source/tutorial/api_server_args.rst | 43 +++- lightllm/server/api_cli.py | 13 +- lightllm/server/core/objs/sampling_params.py | 15 +- lightllm/server/core/objs/start_args_type.py | 2 +- lightllm/server/httpserver/manager.py | 8 +- .../httpserver_for_pd_master/manager.py | 14 +- .../model_infer/mode_backend/base_backend.py | 66 +++--- .../prefill_queue_strategy/__init__.py | 24 +++ .../prefill_queue_strategy/base.py | 40 ++++ .../prefill_queue_strategy/strategies.py | 43 ++++ .../server/router/req_queue/base_queue.py | 6 +- .../server/router/req_queue/dp_base_queue.py | 4 +- lightllm/utils/envs_utils.py | 4 +- .../test_pd_master_multi_choice.py | 9 +- .../test_pd_node_request_limit.py | 43 ++-- unit_tests/common/test_req_manager_page.py | 9 +- .../server/core/objs/test_sampling_params.py | 6 + .../test_running_request_lifecycle.py | 6 +- .../mode_backend/test_pd_dynamic_split.py | 3 +- .../test_prefill_queue_strategy.py | 161 ++++++++++++++ unit_tests/utils/test_envs_utils.py | 4 +- 24 files changed, 647 insertions(+), 111 deletions(-) create mode 100644 docs/CN/source/tutorial/prefill_queue_strategy.md create mode 100644 lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/__init__.py create mode 100644 lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/base.py create mode 100644 lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/strategies.py create mode 100644 unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py diff --git a/docs/CN/source/index.rst b/docs/CN/source/index.rst index 8f79e5126f..a594e31066 100755 --- a/docs/CN/source/index.rst +++ b/docs/CN/source/index.rst @@ -56,6 +56,7 @@ Lightllm 整合了众多的开源方案的优点,包括但不限于 FasterTran OpenAI 接口使用 工具调用(Function Calling) 思考解析(Reasoning Parser) + Prefill 排队策略 APIServer 参数详解 lightllm api介绍 diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index 1015e7b1d1..d7aebc4f5e 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -145,16 +145,17 @@ PD 分离模式参数 请求最终完成的成功率。 设置为非负数时,超时会导致 ``Server is busy``; 其中已进入 Router 但仍未进入推理系统的请求会主动标记为 aborted,由 PD Master 转换为 HTTP 429。 - 本功能启用时,PD Master 收到 ``Server is busy`` 会重新选择 P/D 节点并重试;最长探测周期由 - ``LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS`` 控制,默认 120 秒。若请求已经向客户端输出 token, - 则不再从头重试,以免产生重复内容。设置 ``--disable_pd_node_self_request_limit`` 后,PD Master 不再下发 + 本功能启用时,PD Master 收到 ``Server is busy`` 默认直接返回,不进行重试。 + ``LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS`` 控制最长重试探测周期,默认值为 ``0``;设置为正数后, + PD Master 才会在该周期内重新选择 P/D 节点并重试。若请求已经向客户端输出 token,则即使配置了正数也不再 + 从头重试,以免产生重复内容。设置 ``--disable_pd_node_self_request_limit`` 后,PD Master 不再下发 有限的资源等待时间;P/D 节点永久等待,其他原因产生的 ``Server is busy`` 也会直接返回,不触发重试。 多机 TP 场景仅由 master 节点执行超时判断,slave 节点永久等待。cache 命中记录允许提升优先级的最大年龄由 ``LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MAX_AGE_SECONDS`` 控制,默认 36 秒。cache 命中提权还要求输入 token 数达到 ``LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MIN_PROMPT_TOKENS`` 配置的门槛(默认 4096),避免短请求仅因 cache 命中率高而提升优先级。 - 启动示例: + 以下示例显式开启最长 120 秒的节点繁忙重试: .. code-block:: bash @@ -349,6 +350,31 @@ PD 分离模式参数 是否禁用分块预填充 +.. option:: --prefill_queue_strategy + + 推理后端在资源分配前使用的 prefill 排序策略,默认值为 ``default``。 + 策略会先识别 decode 请求,将其稳定地放在队列最左侧且不参与后续排序; + ``infer_high_priority`` 和所选策略只作用于右侧的 prefill 请求。 + + prefill 请求带有内部字段 ``infer_high_priority``。该字段默认为 ``0``,不接受外部请求设置; + PD 分离模式会为续跑分段设置负值,使其在推理进程中优先于普通请求。 + + * ``default``:按 ``infer_high_priority`` 从小到大稳定排序,相同值内保持 FCFS 顺序。 + * ``promote_shortest``:负优先级 prefill 请求按到达顺序排在前面;在非负优先级的 + 普通请求中选择剩余 prefill token 最少的一个,将其移动到普通请求队头,其余普通 + 请求保持相对顺序。已完成 prompt 计算的请求按零计算。该模式主要适合 PD 分离的 + Prefill 节点,用于改善部分短请求的 TTFT 和首字体验。 + + 例如,在原启动命令中添加 ``--prefill_queue_strategy promote_shortest``。 + 应结合实际流量的 TTFT、TPOT 和尾延迟压测选择策略,策略本身不保证 SLA 改善。 + 完整的调度流程、算法细节和排序示例见 :doc:`Prefill 排队策略 `。 + + 扩展策略时,在 ``lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy`` 中继承 + ``PrefillQueueStrategy`` 并实现 ``reorder_prefill``,然后加入 ``PREFILL_QUEUE_STRATEGIES`` + 映射即可通过启动参数选择。实现必须保持请求集合、请求状态和输入列表不变, + 并保证各 TP rank 的排序结果一致。策略构造函数会接收当前 ``ModeBackend`` 实例, + 可通过 ``self.backend`` 读取调度所需的全局状态。 + .. option:: --diverse_mode 多结果输出模式 diff --git a/docs/CN/source/tutorial/prefill_queue_strategy.md b/docs/CN/source/tutorial/prefill_queue_strategy.md new file mode 100644 index 0000000000..f83f8cb0cb --- /dev/null +++ b/docs/CN/source/tutorial/prefill_queue_strategy.md @@ -0,0 +1,200 @@ +# Prefill 排队策略 + +LightLLM 支持在推理进程中选择 prefill 请求的排队策略。该功能用于调整请求进入本轮资源分配的顺序,便于根据业务流量优化 TTFT、尾延迟和长短请求之间的公平性。 + +启动参数为: + +```bash +--prefill_queue_strategy default +``` + +当前支持两种策略: + +| 策略 | 行为 | +| --- | --- | +| `default` | 负优先级 prefill 请求在前;相同优先级内保持 FCFS 顺序。 | +| `promote_shortest` | 负优先级请求按到达顺序排在前面,只把一个最短的普通 prefill 请求提升到普通请求队头。 | + +## 调度阶段 + +排队策略在 `ModeBackend._get_classed_reqs()` 中执行。此时后端已经取得当前可处理的 `ready_reqs`,但还没有为本轮请求分配计算 token 和 KV cache。 + +整体流程如下: + +```text +ready_reqs + │ + ├─ 按本轮运行条件识别 decode 和 prefill 请求 + │ + ├─ decode 请求稳定放到最左侧,不参与 prefill 策略排序 + │ + ├─ 对右侧 prefill 请求应用选定策略 + │ + └─ 按新顺序检查请求状态并分配计算与 KV cache 预算 +``` + +最终队列始终具有以下结构: + +```text +[decode 请求,保持原始相对顺序] + [经过策略排序的 prefill 请求] +``` + +排序发生在资源检查之前,因此队列靠前的请求会更早尝试获得本轮计算 token 和 KV cache。该机制只改变尚未执行请求的检查顺序,不会抢占正在 GPU 上运行的请求。 + +## Decode 与 prefill 的识别 + +策略通过 `ModeBackend._is_decode_req()` 使用与后续请求分类完全相同的规则: + +1. `no_decode=True` 时,所有请求都按 prefill 处理。 +2. 通常情况下,当 `cur_kv_len + 1 == get_cur_total_len()` 时,请求进入 decode 阶段。 +3. `strict_prefill=True` 时,如果请求位于 prompt 边界,即 `cur_kv_len + 1 == input_len`,它仍按 prefill 处理。 + +`no_decode` 主要用于只执行 prefill 的后端。`strict_prefill` 用于需要严格区分 prompt 阶段的运行模式。 + +识别出的 decode 请求会稳定地移动到队列最左侧。decode 请求之间保持到达顺序,并且不会读取 `infer_high_priority`,也不会按剩余 prefill token 数排序。 + +## 内部推理优先级 + +`infer_high_priority` 是共享内存 `SamplingParams` 中的内部整数字段: + +```text +0 表示普通请求 +负数表示高优先级请求 +``` + +该字段不属于公开 API,外部请求不能设置它。普通请求初始化为 `0`。在 PD 分离模式中,PD Master 仅为第二段及后续续跑分段设置 `-1`,使这些请求在推理进程的 prefill 队列中优先处理。首段即使满足高 cache 命中条件,也只会通过 `high_priority_request` 提升 HTTP 资源申请和 Router 等待队列的优先级,不会改变推理进程中的 `infer_high_priority`。 + +`infer_high_priority` 与 `high_priority_request` 的职责不同: + +| 字段 | 类型 | 使用位置 | 作用 | +| --- | --- | --- | --- | +| `high_priority_request` | `bool` | HTTP server 与 Router 等待队列 | 调整共享内存资源重试和进入 Router 的等待顺序。 | +| `infer_high_priority` | `int` | 推理进程的 prefill 排队策略 | 调整 prefill 请求参与计算和 KV cache 资源检查的顺序。 | + +两者都是内部字段。`high_priority_request` 不会被推理侧的 prefill 策略读取。 + +## `default` 策略 + +`default` 策略对 prefill 请求执行以下稳定排序: + +```python +sorted(prefill_reqs, key=infer_priority) +``` + +Python 的 `sorted()` 是稳定排序。因此: + +- 负优先级请求排在普通请求之前。 +- 如果以后出现多个负优先级等级,数值更小的请求排在前面。 +- `infer_high_priority` 相同时保持原始相对顺序,也就是 FCFS。 +- 当所有请求的优先级都是默认值 `0` 时,队列顺序完全不变。 + +该策略适合希望保持简单、公平排队,同时确保 PD 内部高优先级分段能够尽快继续执行的场景。 + +## `promote_shortest` 策略 + +`promote_shortest` 不会把所有普通请求按长度排序。它只提升一个请求,具体步骤如下: + +1. 按 `infer_high_priority < 0` 和 `infer_high_priority >= 0` 将请求分为高优先级组和普通组,各组保持到达顺序。 +2. 高优先级组保持在普通组之前。 +3. 只在非负优先级的普通请求中查找最短请求。 +4. 按下式计算每个普通请求的剩余 prefill token 数: + + ```python + max(0, req.shm_req.input_len - req.cur_kv_len) + ``` + +5. 选择剩余 token 数最少的一个普通请求。 +6. 将它移动到普通请求队头。 +7. 其他普通请求保持原始相对顺序。 + +如果多个普通请求的剩余 token 数相同,会选择其中原本最靠前的请求。如果不存在普通请求,则高优先级请求的到达顺序保持不变。 + +这种设计只允许一个短请求越过其他普通请求,可以改善短请求的 TTFT,同时避免完整 SJF 排序持续改变整个队列并显著增加长请求的等待时间。 + +该模式主要适合在 PD 分离架构的 Prefill 节点上使用。Prefill 节点同时排队处理多个长短不一的 prompt 时,优先调度一个剩余 prefill token 最少的普通请求,可以让部分短请求更早完成 Prefill 并进入 Decode,从而改善其 TTFT 和首字体验。续跑分段仍位于普通请求之前,不会被短请求越过。 + +## 排序示例 + +假设原始队列如下: + +| 到达顺序 | 请求 | 阶段 | `infer_high_priority` | 剩余 prefill token | +| --- | --- | --- | ---: | ---: | +| 1 | `normal-a` | prefill | 0 | 100 | +| 2 | `decode-a` | decode | 0 | 不参与 | +| 3 | `pd-high` | prefill | -1 | 800 | +| 4 | `normal-b` | prefill | 0 | 50 | +| 5 | `decode-b` | decode | -1 | 不参与 | +| 6 | `normal-short` | prefill | 0 | 20 | + +`default` 的结果为: + +```text +decode-a, decode-b, pd-high, normal-a, normal-b, normal-short +``` + +`promote_shortest` 的结果为: + +```text +decode-a, decode-b, pd-high, normal-short, normal-a, normal-b +``` + +可以看到: + +- decode 请求始终在最左侧,并保持 `decode-a, decode-b` 的相对顺序。 +- `pd-high` 不会被普通短请求越过。 +- 只有 `normal-short` 被提升。 +- `normal-a` 和 `normal-b` 的相对顺序没有改变。 + +## 选择策略 + +使用默认策略: + +```bash +python -m lightllm.server.api_server \ + --model_dir /path/to/model \ + --prefill_queue_strategy default +``` + +提升一个最短普通请求: + +```bash +python -m lightllm.server.api_server \ + --model_dir /path/to/model \ + --prefill_queue_strategy promote_shortest +``` + +建议使用与生产环境一致的请求长度分布、并发度和 prefix cache 命中率进行压测,重点观察: + +- TTFT 及其 P95/P99。 +- TPOT 和端到端延迟。 +- 长 prompt 请求的等待时间。 +- 计算 token 与 KV cache 容量紧张时的尾延迟。 +- PD 续跑分段能否及时获得推理资源。 + +`promote_shortest` 主要用于改善 PD 分离模式下 Prefill 节点中部分短请求的首字体验。它只提升一个普通请求,通常比完整 SJF 更温和,但具体 TTFT 和 SLA 收益仍取决于实际流量。 + +## 扩展新策略 + +策略代码位于: + +```text +lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy +``` + +新增策略时,需要: + +1. 继承 `PrefillQueueStrategy`。 +2. 实现 `reorder_prefill(prefill_reqs)`。 +3. 在 `PREFILL_QUEUE_STRATEGIES` 中注册名称与实现类。 +4. 同步更新 CLI 和 `StartArgs` 的可选值。 +5. 补充策略顺序、稳定性、空队列、单请求、decode 前置和输入列表不变等测试。 + +基类已经负责识别并前置 decode 请求。子类只能重排传入的 prefill 请求,并应遵守以下约束: + +- 不丢弃或复制请求。 +- 不修改请求状态。 +- 不原地修改输入列表。 +- 相同排序条件下保持稳定顺序。 +- 在各 TP rank 上产生一致、确定的结果。 + +策略对象持有完整的 `ModeBackend`,可以通过 `self.backend` 读取模型、缓存、rank 和其他调度状态。读取这些状态时仍需保证所有 TP rank 得到一致的排序结果。 diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index a58c607971..a1f66dcbc4 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -153,10 +153,11 @@ PD disaggregation Mode Parameters that requests which have already produced partial results complete successfully. When set to a non-negative value, a timeout reports ``Server is busy``; a request that has entered the Router but not inference is proactively marked aborted, and PD Master converts this to HTTP 429. - While this feature is enabled, PD Master selects P/D nodes again and retries after receiving ``Server is busy``. - The maximum probing period is controlled by ``LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS`` and defaults to - 120 seconds. Once response tokens have been streamed to the client, the request is not restarted because doing so - would duplicate output. With ``--disable_pd_node_self_request_limit``, PD Master no longer supplies a finite + While this feature is enabled, PD Master returns ``Server is busy`` without retrying by default. + ``LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS`` controls the maximum retry probing period and defaults to ``0``; + set it to a positive value to make PD Master select P/D nodes again and retry within that period. Once response + tokens have been streamed to the client, the request is not restarted even when a positive timeout is configured, + because doing so would duplicate output. With ``--disable_pd_node_self_request_limit``, PD Master no longer supplies a finite resource wait timeout; all P/D nodes wait indefinitely, and a ``Server is busy`` raised for another reason is returned immediately without retrying. In multi-node TP deployments, only the master node evaluates the timeout; slave nodes wait indefinitely. @@ -166,7 +167,7 @@ PD disaggregation Mode Parameters ``LIGHTLLM_PD_CACHE_HIGH_PRIORITY_MIN_PROMPT_TOKENS`` (4096 by default), so short requests do not gain priority solely from a high cache-hit rate. - Startup example: + The following example explicitly enables busy retries for up to 120 seconds: .. code-block:: bash @@ -364,6 +365,38 @@ Scheduling Parameters Whether to disable chunked prefill +.. option:: --prefill_queue_strategy + + Prefill ordering policy in the inference backend before resource allocation. Defaults + to ``default``. The policy first identifies decode requests, keeps them at the left in + arrival order, and excludes them from subsequent sorting. ``infer_high_priority`` and + the selected policy apply only to the prefill requests on the right. + + Prefill requests carry an internal ``infer_high_priority`` field. It defaults to ``0``, + cannot be set by external requests, and is assigned a negative value to continuation + segments through internal PD-disaggregation communication. + + * ``default``: stably sort by ``infer_high_priority`` from smallest to largest and preserve + FCFS order within the same value. + * ``promote_shortest``: keep negative-priority prefill requests first in arrival order. + Among non-negative-priority normal requests, move only the request with the fewest + remaining prefill tokens to the front of the normal queue. Preserve the relative order + of all other normal requests. Fully processed prompts count as zero. This mode is mainly + intended for Prefill nodes in PD-disaggregated deployments to improve TTFT for some short + requests. + + For example, append ``--prefill_queue_strategy promote_shortest`` to your launch command. + Compare TTFT, TPOT, and tail latency under representative traffic before choosing a + policy; changing policies does not guarantee an SLA improvement. + + To extend, subclass ``PrefillQueueStrategy`` in + ``lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy``, implement + ``reorder_prefill``, and + add the class to ``PREFILL_QUEUE_STRATEGIES`` to expose its name through the CLI. + Implementations must preserve the request set, request state, and input list, and + produce deterministic ordering across TP ranks. Each strategy receives the current + ``ModeBackend`` instance and can read scheduling state through ``self.backend``. + .. option:: --diverse_mode Multi-result output mode diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index 926f9f2030..e7473e919c 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -371,11 +371,14 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument("--chunked_prefill_size", type=int, default=None, help="chunked prefill size") parser.add_argument("--disable_chunked_prefill", action="store_true", help="whether to disable chunked prefill") parser.add_argument( - "--short_prefill_token_threshold", - type=int, - default=None, - help="""Enable short prefill request priority scheduling. - The remaining tokens are calculated after prefix-cache matching. Disabled by default.""", + "--prefill_queue_strategy", + choices=("default", "promote_shortest"), + default="default", + help="""Order ready requests before resource allocation. + Decode requests are identified first, remain in arrival order, and are not reordered by these policies. + default: preserve FCFS order for prefill requests with the same infer_high_priority value. + promote_shortest: keep negative-priority requests first in arrival order, then move one + shortest non-negative-priority prefill request to the front of the normal requests.""", ) parser.add_argument("--diverse_mode", action="store_true", help="diversity generation mode") vocab_topk_choices = [2, 8, 16, 32, 64, 128, 256, 512] diff --git a/lightllm/server/core/objs/sampling_params.py b/lightllm/server/core/objs/sampling_params.py index f0dda1dd80..dc36259b72 100644 --- a/lightllm/server/core/objs/sampling_params.py +++ b/lightllm/server/core/objs/sampling_params.py @@ -286,9 +286,10 @@ class SamplingParams(ctypes.Structure): ("stop_sequences", StopSequenceGroups), ("exponential_decay_length_penalty", ExponentialDecayLengthPenalty), ("group_request_id", ctypes.c_int64), # p d mode used params - # 由 PD Master 为分段续跑或预计 cache 命中率较高的请求设置,表示请求需 - # 以高优先级插入 Router 调度队列。 - ("pd_high_priority_request", ctypes.c_bool), + # 由 PD Master 设置,仅用于 HTTP 资源申请和 Router 等待队列调度。 + ("high_priority_request", ctypes.c_bool), + # PD 内部通信使用的 prefill 排队优先级;0 表示普通请求,负值表示高优先级请求。 + ("infer_high_priority", ctypes.c_int), # P/D 节点的资源等待超时,由 PD Master 下发。非负值用于控制 shm_req 申请和 # Router 等待进入推理系统的时限;负数表示永久等待。 ("pd_node_resource_wait_timeout_seconds", ctypes.c_int), @@ -337,8 +338,9 @@ def init(self, tokenizer, **kwargs): self.min_new_tokens = kwargs.get("min_new_tokens", 1) self.input_penalty = kwargs.get("input_penalty", DEFAULT_INPUT_PENALTY) self.group_request_id = kwargs.get("group_request_id", -1) - # 这两个字段是 PD Master 的内部调度信息,不能由外部请求参数开启或修改。 - self.pd_high_priority_request = False + # PD Master 的内部调度信息,不能由外部请求参数开启或修改。 + self.high_priority_request = False + self.infer_high_priority = 0 self.pd_node_resource_wait_timeout_seconds = -1 self.suggested_dp_index = kwargs.get("suggested_dp_index", -1) @@ -508,7 +510,8 @@ def to_dict(self): "allowed_token_ids": self.allowed_token_ids.to_list(), "invalid_token_ids": self.invalid_token_ids.to_list(), "group_request_id": self.group_request_id, - "pd_high_priority_request": self.pd_high_priority_request, + "high_priority_request": self.high_priority_request, + "infer_high_priority": self.infer_high_priority, "pd_node_resource_wait_timeout_seconds": self.pd_node_resource_wait_timeout_seconds, "skip_special_tokens": self.skip_special_tokens, "add_special_tokens": self.add_special_tokens, diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 13ad6367a5..8abdf9fdb3 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -101,7 +101,7 @@ class StartArgs: disable_dynamic_prompt_cache: bool = field(default=False) chunked_prefill_size: int = field(default=None) disable_chunked_prefill: bool = field(default=False) - short_prefill_token_threshold: Optional[int] = field(default=None) + prefill_queue_strategy: str = field(default="default", metadata={"choices": ["default", "promote_shortest"]}) diverse_mode: bool = field(default=False) target_vocab_topk_sampling: Optional[int] = field( default=None, metadata={"choices": [2, 8, 16, 32, 64, 128, 256, 512]} diff --git a/lightllm/server/httpserver/manager.py b/lightllm/server/httpserver/manager.py index 420c55eece..7a5e2557be 100644 --- a/lightllm/server/httpserver/manager.py +++ b/lightllm/server/httpserver/manager.py @@ -430,11 +430,11 @@ async def generate( await self._register_running_request() running_request_registered = True - # 申请资源并存储。PD 高优先级请求仍以更短的间隔重试;资源等待上限 + # 申请资源并存储。高优先级请求仍以更短的间隔重试;资源等待上限 # 完全由 PD Master 下发,与请求优先级无关。 alloced_req_indexes = await self._alloc_shm_req_indexes( sampling_params.n, - pd_high_priority_request=sampling_params.pd_high_priority_request, + high_priority_request=sampling_params.high_priority_request, pd_node_resource_wait_timeout_seconds=sampling_params.pd_node_resource_wait_timeout_seconds, ) req_objs: List[Req] = [] @@ -545,7 +545,7 @@ def _count_multimodal_tokens(self, multimodal_params: MultimodalParams) -> Tuple async def _alloc_shm_req_indexes( self, req_num: int, - pd_high_priority_request: bool = False, + high_priority_request: bool = False, pd_node_resource_wait_timeout_seconds: int = -1, ) -> List[int]: """为一个请求申请全部 shm_req 索引,申请失败时回滚已分配的索引。 @@ -565,7 +565,7 @@ async def _alloc_shm_req_indexes( while len(alloced_req_indexes) < req_num: alloc_req_index = await self.shm_req_manager.async_alloc_req_index() # 保持相同的退避起点,仅通过系数让高优先级请求更快地重新尝试获取 shm_req。 - sleep_time_factor = 0.2 if pd_high_priority_request else 1 + sleep_time_factor = 0.2 if high_priority_request else 1 sleep_time = 0.1 while alloc_req_index is None: if alloc_deadline is not None and time.monotonic() >= alloc_deadline: diff --git a/lightllm/server/httpserver_for_pd_master/manager.py b/lightllm/server/httpserver_for_pd_master/manager.py index 96f0d7203e..77c86c7bc6 100644 --- a/lightllm/server/httpserver_for_pd_master/manager.py +++ b/lightllm/server/httpserver_for_pd_master/manager.py @@ -230,7 +230,7 @@ async def _generate_one( origin_request_id: int, input_token_num: int, ): - """节点繁忙时重新选择 P/D 节点,并在配置的探测周期内重试。""" + """处理节点繁忙错误;配置正数探测周期时重新选择 P/D 节点并重试。""" retry_start_time = time.monotonic() has_yielded_result = False @@ -338,11 +338,13 @@ async def _generate_one_attempt( 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 = remaining_max_new_tokens - # 首段仅在输入达到长度门槛、预计 cache 命中率高于 0.8 且命中记录仍在 - # 有效时间窗内时提升优先级,避免短请求或可能已被 P 节点淘汰的陈旧 - # KV cache 插队。第二段及后续分段仍统一使用高优先级,避免因临时资源 - # 紧张导致分段续跑失败。 - sampling_params.pd_high_priority_request = segment_index > 0 or has_fresh_high_cache_hit + # high_priority_request 只影响 HTTP 资源申请和 Router 等待队列:首段仅在 + # 输入达到长度门槛、预计 cache 命中率高于 0.8 且命中记录仍有效时提升; + # 第二段及后续分段统一提升,避免因临时资源紧张导致分段续跑失败。 + sampling_params.high_priority_request = segment_index > 0 or has_fresh_high_cache_hit + # infer_high_priority 是 PD 内部通信字段,不接受外部请求设置。 + # 仅第二段及后续分段在推理进程中优先排队。 + sampling_params.infer_high_priority = -1 if segment_index > 0 else 0 # 仅在 Master 开启限流时下发资源等待超时。续跑分段已经产生了部分结果, # 使用独立配置的等待时间,提高请求最终完成的成功率。 if self.enable_pd_node_self_request_limit: diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 650cb8493c..27e3ce3e00 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -48,6 +48,7 @@ CacheTier, create_cache_placement_controller, ) +from .prefill_queue_strategy import create_prefill_queue_strategy from .multi_level_kv_cache import MultiLevelKvCacheModule from lightllm.utils.profiler import ProcessProfiler, ProfilerCmd @@ -259,6 +260,9 @@ def init_model(self, kvargs): prof_mode = self.args.enable_profiling self.profiler = ProcessProfiler(mode=prof_mode, name=prof_name, use_multi_thread=True) if prof_mode else None + # 策略可能读取模型、缓存管理器和 rank 等全局调度状态,因此在 backend 完成初始化后再创建。 + self.prefill_queue_strategy = create_prefill_queue_strategy(self) + # 启动infer_loop_thread, 启动两个线程进行推理,对于具备双batch推理折叠得场景 # 可以降低 cpu overhead,大幅提升gpu得使用率。 self.infer_loop_thread = threading.Thread(target=self.infer_loop, daemon=True) @@ -614,35 +618,6 @@ def _timer_merge_radix_tree(self): ) return - def _reorder_pd_high_priority_reqs(self, ready_reqs: List[InferReq]) -> List[InferReq]: - """将 PD 分段续跑的高优先级请求前置,普通请求保持在其后。""" - # PD 分段续跑请求已经完成前一段推理,需要优先进入本轮调度;将请求拆分后再拼接, - # 保持各自原有顺序,并确保高优先级请求位于普通请求之前。 - high_priority_reqs = [req for req in ready_reqs if req.shm_req.sample_params.pd_high_priority_request] - normal_reqs = [req for req in ready_reqs if not req.shm_req.sample_params.pd_high_priority_request] - return high_priority_reqs + normal_reqs - - def _reorder_long_prefill_reqs(self, ready_reqs: List[InferReq]) -> List[InferReq]: - """ - 提升一个短 prefill 请求的优先级。 - """ - short_token_threshold = self.args.short_prefill_token_threshold - if short_token_threshold is None: - return ready_reqs - - def remaining_prefill_tokens(req: InferReq) -> int: - return max(0, req.shm_req.input_len - req.cur_kv_len) - - sorted_reqs = sorted( - ready_reqs, - key=lambda req: (remaining_prefill_tokens(req), req.shm_req.group_req_id), - ) - if sorted_reqs and remaining_prefill_tokens(sorted_reqs[0]) <= short_token_threshold: - target_req = sorted_reqs[0] - ready_reqs.remove(target_req) - ready_reqs.insert(0, target_req) - return ready_reqs - # 一些可以复用的通用功能函数 def _alloc_req_kv_mem( self, @@ -667,6 +642,23 @@ def _alloc_req_kv_mem( req_obj.hold_kv_len = new_hold_kv_len return mem_indexes + @staticmethod + def _is_decode_req(req: InferReq, no_decode: bool, strict_prefill: bool) -> bool: + """判断请求在当前调度轮次中是否按 decode 请求处理。""" + if no_decode: + return False + + is_decode = req.cur_kv_len + 1 == req.get_cur_total_len() + if not is_decode: + return False + + if strict_prefill: + is_at_prompt_boundary = req.cur_kv_len + 1 == req.shm_req.input_len + if is_at_prompt_boundary: + return False + + return True + def _get_classed_reqs( self, req_ids: List[int] = None, @@ -706,8 +698,11 @@ def _get_classed_reqs( ready_reqs = self._filter_not_ready_reqs(req_ids) support_overlap = self.support_overlap - ready_reqs = self._reorder_pd_high_priority_reqs(ready_reqs) - ready_reqs = self._reorder_long_prefill_reqs(ready_reqs) + ready_reqs = self.prefill_queue_strategy.reorder( + ready_reqs, + no_decode=no_decode, + strict_prefill=strict_prefill, + ) wait_pause_reqs = [] paused_reqs = [] @@ -746,14 +741,7 @@ def _get_classed_reqs( finished_reqs.append(req_obj) continue - if no_decode: - is_decode = False - else: - is_decode = req_obj.cur_kv_len + 1 == req_obj.get_cur_total_len() - if is_decode and strict_prefill and req_obj.cur_kv_len + 1 == req_obj.shm_req.input_len: - is_decode = False - - if is_decode: + if self._is_decode_req(req_obj, no_decode=no_decode, strict_prefill=strict_prefill): # KV 容量检查使用额外分配量,已有页的剩余容量可以覆盖部分或全部 decode 需求。 _, alloc_token_num = req_obj.decode_need_token_num() # page_size 较小时,decode 会频繁触发 KV 内存分配。此处额外预申请不超过 8 个 token, diff --git a/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/__init__.py b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/__init__.py new file mode 100644 index 0000000000..4e621d5541 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/__init__.py @@ -0,0 +1,24 @@ +from typing import TYPE_CHECKING + +from .base import PrefillQueueStrategy +from .strategies import FCFSStrategy, PromoteShortestPrefillStrategy + +if TYPE_CHECKING: + from ..base_backend import ModeBackend + + +# 在这里注册新的 PrefillQueueStrategy 子类,即可通过启动参数选择。 +PREFILL_QUEUE_STRATEGIES = { + "default": FCFSStrategy, + "promote_shortest": PromoteShortestPrefillStrategy, +} + + +def create_prefill_queue_strategy(backend: "ModeBackend") -> PrefillQueueStrategy: + name = backend.args.prefill_queue_strategy + if name not in PREFILL_QUEUE_STRATEGIES: + raise ValueError(f"Unknown prefill queue strategy {name!r}; choose from {', '.join(PREFILL_QUEUE_STRATEGIES)}") + return PREFILL_QUEUE_STRATEGIES[name](backend) + + +__all__ = ["PrefillQueueStrategy", "PREFILL_QUEUE_STRATEGIES", "create_prefill_queue_strategy"] diff --git a/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/base.py b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/base.py new file mode 100644 index 0000000000..c8033b0905 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/base.py @@ -0,0 +1,40 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, List + +if TYPE_CHECKING: + from lightllm.server.router.model_infer.infer_batch import InferReq + from ..base_backend import ModeBackend + + +class PrefillQueueStrategy(ABC): + """在请求分类和 token 预算分配前调整就绪请求的排队顺序。 + + 输入中可能包含 decode、暂停或已结束的请求。基类先把 decode 请求稳定地放到最左侧, + 子类只负责排序右侧的 prefill 请求。实现必须返回原请求集合的一个排列,不能修改请求状态、 + 丢弃请求或原地修改输入列表。各 TP rank 的排序结果必须一致;除非策略明确提供次级排序条件, + 否则应保持相同排序键请求的原始顺序。 + 策略可通过 ``self.backend`` 读取已完成初始化的模型、缓存、rank 等调度状态。 + """ + + def __init__(self, backend: "ModeBackend") -> None: + self.backend = backend + + def reorder( + self, + ready_reqs: List["InferReq"], + no_decode: bool = False, + strict_prefill: bool = False, + ) -> List["InferReq"]: + decode_reqs = [] + prefill_reqs = [] + for req in ready_reqs: + if self.backend._is_decode_req(req, no_decode=no_decode, strict_prefill=strict_prefill): + decode_reqs.append(req) + else: + prefill_reqs.append(req) + + return decode_reqs + self.reorder_prefill(prefill_reqs) + + @abstractmethod + def reorder_prefill(self, prefill_reqs: List["InferReq"]) -> List["InferReq"]: + raise NotImplementedError diff --git a/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/strategies.py b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/strategies.py new file mode 100644 index 0000000000..76efb6f3d1 --- /dev/null +++ b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/strategies.py @@ -0,0 +1,43 @@ +from typing import TYPE_CHECKING, List + +from .base import PrefillQueueStrategy + +if TYPE_CHECKING: + from lightllm.server.router.model_infer.infer_batch import InferReq + + +def remaining_prefill_tokens(req: "InferReq") -> int: + return max(0, req.shm_req.input_len - req.cur_kv_len) + + +def infer_priority(req: "InferReq") -> int: + return req.shm_req.sample_params.infer_high_priority + + +class FCFSStrategy(PrefillQueueStrategy): + """按 prefill 推理优先级排队;优先级相同时保持 FCFS 顺序。""" + + def reorder_prefill(self, prefill_reqs: List["InferReq"]) -> List["InferReq"]: + # 这里依赖 Python 的 sorted 为稳定排序,确保相同优先级的请求保持 FCFS 顺序。 + return sorted(prefill_reqs, key=infer_priority) + + +class PromoteShortestPrefillStrategy(PrefillQueueStrategy): + """在普通 prefill 请求中,只将剩余 token 最少的一个请求提升到普通队列头部。 + + 负优先级请求按原始相对顺序排在前面,非负优先级请求视为普通请求。 + 除被提升的最短请求外,其他普通请求保持原始相对顺序。 + prompt 已处理完成的请求按零个剩余 token 计算。 + """ + + def reorder_prefill(self, prefill_reqs: List["InferReq"]) -> List["InferReq"]: + high_priority_reqs = [req for req in prefill_reqs if infer_priority(req) < 0] + normal_reqs = [req for req in prefill_reqs if infer_priority(req) >= 0] + if not normal_reqs: + return high_priority_reqs + + shortest_req = min(normal_reqs, key=remaining_prefill_tokens) + normal_reqs.remove(shortest_req) + high_priority_reqs.append(shortest_req) + high_priority_reqs.extend(normal_reqs) + return high_priority_reqs diff --git a/lightllm/server/router/req_queue/base_queue.py b/lightllm/server/router/req_queue/base_queue.py index 0c25b8949a..8d4ac38085 100644 --- a/lightllm/server/router/req_queue/base_queue.py +++ b/lightllm/server/router/req_queue/base_queue.py @@ -81,12 +81,12 @@ def filter_aborted_reqs(self): def extend(self, req_group: List[Req]): for req in req_group: req.sample_params.suggested_dp_index = self.dp_index - # PD 高优先级请求应排在普通请求之前,但高优先级请求之间仍按到达顺序排队, + # 高优先级请求应排在普通请求之前,但高优先级请求之间仍按到达顺序排队, # 避免后到请求反复插到队头而阻塞先到的高优先级请求。 - if req_group and req_group[0].sample_params.pd_high_priority_request: + if req_group and req_group[0].sample_params.high_priority_request: first_normal_req_index = len(self.waiting_req_list) for index, waiting_req in enumerate(self.waiting_req_list): - if not waiting_req.sample_params.pd_high_priority_request: + if not waiting_req.sample_params.high_priority_request: first_normal_req_index = index break # req_group 可能包含同一请求组的多个 Req,整体插入可以保持组内顺序。 diff --git a/lightllm/server/router/req_queue/dp_base_queue.py b/lightllm/server/router/req_queue/dp_base_queue.py index f7f7173ebb..2d5c5b039b 100644 --- a/lightllm/server/router/req_queue/dp_base_queue.py +++ b/lightllm/server/router/req_queue/dp_base_queue.py @@ -56,11 +56,11 @@ def extend(self, req_group: List[Req]): suggested_dp_index = req_group[0].sample_params.suggested_dp_index if suggested_dp_index >= self.dp_size_in_node or suggested_dp_index < 0: # 同一个组的,要分配在同一个 dp 上 - if req_group[0].sample_params.pd_high_priority_request: + if req_group[0].sample_params.high_priority_request: # 高优先级请求组插在第一个普通请求组之前,同时保持高优先级组之间的 FIFO 顺序。 first_normal_group_index = len(self.reqs_waiting_for_dp_index) for index, waiting_group in enumerate(self.reqs_waiting_for_dp_index): - if not waiting_group[0].sample_params.pd_high_priority_request: + if not waiting_group[0].sample_params.high_priority_request: first_normal_group_index = index break self.reqs_waiting_for_dp_index.insert(first_normal_group_index, req_group) diff --git a/lightllm/utils/envs_utils.py b/lightllm/utils/envs_utils.py index 66bed05122..66bdd8f39d 100644 --- a/lightllm/utils/envs_utils.py +++ b/lightllm/utils/envs_utils.py @@ -327,8 +327,8 @@ def get_pd_node_continuation_resource_wait_timeout_seconds() -> int: @lru_cache(maxsize=None) def get_pd_node_busy_retry_timeout_seconds() -> int: - """PD Master 收到节点繁忙错误后的最长重试时间,单位为秒。""" - return max(0, int(os.getenv("LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS", 120))) + """PD Master 收到节点繁忙错误后的最长重试时间,单位为秒;0 表示不重试。""" + return max(0, int(os.getenv("LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS", 0))) @lru_cache(maxsize=None) diff --git a/test/test_pd_selector/test_pd_master_multi_choice.py b/test/test_pd_selector/test_pd_master_multi_choice.py index 38f43232fa..b822ecc6f0 100644 --- a/test/test_pd_selector/test_pd_master_multi_choice.py +++ b/test/test_pd_selector/test_pd_master_multi_choice.py @@ -394,6 +394,7 @@ async def run(): dispatched_loads = [] dispatched_req_counts = [] high_priority_request_flags = [] + infer_high_priorities = [] dispatched_max_new_tokens = [] async def wait_to_token_package( @@ -404,7 +405,8 @@ async def wait_to_token_package( dispatched_prompts.append(block_prompt) dispatched_loads.append(selected_p_node.dispatched_prompt_chars) dispatched_req_counts.append(selected_p_node.dispatched_req_num) - high_priority_request_flags.append(sampling_params.pd_high_priority_request) + high_priority_request_flags.append(sampling_params.high_priority_request) + infer_high_priorities.append(sampling_params.infer_high_priority) dispatched_max_new_tokens.append(sampling_params.max_new_tokens) yield ( sampling_params.group_request_id, @@ -447,6 +449,7 @@ async def wait_to_token_package( assert dispatched_loads == [other_request_load + len("prompt"), other_request_load + len("promptx")] assert dispatched_req_counts == [other_request_count + 1, other_request_count + 1] assert high_priority_request_flags == [False, True] + assert infer_high_priorities == [0, -1] assert dispatched_max_new_tokens == [2, 1] assert p_node.dispatched_prompt_chars == other_request_load assert p_node.dispatched_req_num == other_request_count @@ -561,7 +564,7 @@ async def run(): high_priority_request_flags = [] async def wait_to_token_package(_p_node, _d_node, _start_time, _prompt, sampling_params, *_args): - high_priority_request_flags.append(sampling_params.pd_high_priority_request) + high_priority_request_flags.append(sampling_params.high_priority_request) yield ( sampling_params.group_request_id, "x", @@ -619,7 +622,7 @@ async def run(): async def wait_to_token_package(_p_node, _d_node, _start_time, _prompt, sampling_params, *_args): captured_request_settings.append( ( - sampling_params.pd_high_priority_request, + sampling_params.high_priority_request, sampling_params.pd_node_resource_wait_timeout_seconds, ) ) diff --git a/test/test_pd_selector/test_pd_node_request_limit.py b/test/test_pd_selector/test_pd_node_request_limit.py index 3cf20573a9..a6f9b394b0 100644 --- a/test/test_pd_selector/test_pd_node_request_limit.py +++ b/test/test_pd_selector/test_pd_node_request_limit.py @@ -38,14 +38,21 @@ def test_pd_node_resource_wait_timeout_is_internal_and_defaults_to_waiting_forev sampling_params = SamplingParams() sampling_params.init( None, - pd_high_priority_request=True, + high_priority_request=True, pd_node_resource_wait_timeout_seconds=99, ) - assert sampling_params.pd_high_priority_request is False + assert sampling_params.high_priority_request is False assert sampling_params.pd_node_resource_wait_timeout_seconds == -1 +def test_infer_high_priority_is_internal_and_cannot_be_set_by_request(): + sampling_params = SamplingParams() + sampling_params.init(None, infer_high_priority=-3) + + assert sampling_params.infer_high_priority == 0 + + def test_shm_req_partial_allocations_are_released_on_failure(): async def run(): manager = _manager() @@ -95,7 +102,7 @@ async def run(): assert ( await manager._alloc_shm_req_indexes( 1, - pd_high_priority_request=True, + high_priority_request=True, pd_node_resource_wait_timeout_seconds=60, ) == [3] @@ -106,8 +113,8 @@ async def run(): asyncio.run(run()) -@pytest.mark.parametrize("pd_high_priority_request", [False, True]) -def test_shm_req_allocation_uses_master_timeout_independently_of_priority(pd_high_priority_request): +@pytest.mark.parametrize("high_priority_request", [False, True]) +def test_shm_req_allocation_uses_master_timeout_independently_of_priority(high_priority_request): async def run(): manager = _manager() manager.shm_req_manager.async_alloc_req_index = AsyncMock(return_value=None) @@ -118,7 +125,7 @@ async def run(): ): await manager._alloc_shm_req_indexes( 1, - pd_high_priority_request=pd_high_priority_request, + high_priority_request=high_priority_request, pd_node_resource_wait_timeout_seconds=60, ) @@ -126,7 +133,7 @@ async def run(): @pytest.mark.parametrize( - ("infer_start_time", "pd_high_priority_request", "resource_wait_timeout_seconds", "expected"), + ("infer_start_time", "high_priority_request", "resource_wait_timeout_seconds", "expected"), [ (0, False, -1, False), (0, False, 60, True), @@ -137,7 +144,7 @@ async def run(): ) def test_router_wait_uses_master_timeout_independently_of_priority( infer_start_time, - pd_high_priority_request, + high_priority_request, resource_wait_timeout_seconds, expected, ): @@ -148,7 +155,7 @@ def test_router_wait_uses_master_timeout_independently_of_priority( infer_start_time=infer_start_time, router_arrival_time=100, sample_params=SimpleNamespace( - pd_high_priority_request=pd_high_priority_request, + high_priority_request=high_priority_request, pd_node_resource_wait_timeout_seconds=resource_wait_timeout_seconds, ), ) @@ -159,13 +166,13 @@ def test_router_wait_uses_master_timeout_independently_of_priority( assert req_status.has_timed_out_waiting_for_inference() is expected -def test_pd_high_priority_request_is_inserted_before_first_normal_request(): +def test_high_priority_request_is_inserted_before_first_normal_request(): queue = BaseQueue.__new__(BaseQueue) queue.dp_index = 0 - earlier_high_req = SimpleNamespace(sample_params=SimpleNamespace(pd_high_priority_request=True)) - normal_req = SimpleNamespace(sample_params=SimpleNamespace(pd_high_priority_request=False)) - high_req_1 = SimpleNamespace(sample_params=SimpleNamespace(pd_high_priority_request=True)) - high_req_2 = SimpleNamespace(sample_params=SimpleNamespace(pd_high_priority_request=True)) + earlier_high_req = SimpleNamespace(sample_params=SimpleNamespace(high_priority_request=True)) + normal_req = SimpleNamespace(sample_params=SimpleNamespace(high_priority_request=False)) + high_req_1 = SimpleNamespace(sample_params=SimpleNamespace(high_priority_request=True)) + high_req_2 = SimpleNamespace(sample_params=SimpleNamespace(high_priority_request=True)) queue.waiting_req_list = [earlier_high_req, normal_req] queue.extend([high_req_1, high_req_2]) @@ -175,15 +182,15 @@ def test_pd_high_priority_request_is_inserted_before_first_normal_request(): assert high_req_2.sample_params.suggested_dp_index == 0 -def test_pd_high_priority_request_group_keeps_fifo_order_while_waiting_for_dp_index(): +def test_high_priority_request_group_keeps_fifo_order_while_waiting_for_dp_index(): queue = DpQueue.__new__(DpQueue) queue.dp_size_in_node = 2 - earlier_high_group = [SimpleNamespace(sample_params=SimpleNamespace(pd_high_priority_request=True))] - normal_group = [SimpleNamespace(sample_params=SimpleNamespace(pd_high_priority_request=False))] + earlier_high_group = [SimpleNamespace(sample_params=SimpleNamespace(high_priority_request=True))] + normal_group = [SimpleNamespace(sample_params=SimpleNamespace(high_priority_request=False))] new_high_group = [ SimpleNamespace( sample_params=SimpleNamespace( - pd_high_priority_request=True, + high_priority_request=True, suggested_dp_index=-1, ) ) diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index f45ac26ceb..f1d7a5ec67 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -185,8 +185,7 @@ def test_prefill_scheduler_checks_compute_and_kv_budgets_separately( backend.batch_max_tokens = batch_max_tokens backend.is_master_in_dp = True backend._timer_merge_radix_tree = lambda: None - backend._reorder_pd_high_priority_reqs = lambda reqs: reqs - backend._reorder_long_prefill_reqs = lambda reqs: reqs + backend.prefill_queue_strategy = SimpleNamespace(reorder=lambda reqs, **kwargs: reqs) context.get_can_alloc_token_num = lambda: can_alloc_token_num context.cache_placement_controller = SimpleNamespace(set_req_cache_way=lambda reqs: None) context.filter_reqs = lambda finished_reqs: None @@ -246,8 +245,7 @@ def test_decode_scheduler_uses_non_blocking_request_table_copy(monkeypatch): backend.batch_max_tokens = 8 backend.is_master_in_dp = True backend._timer_merge_radix_tree = lambda: None - backend._reorder_pd_high_priority_reqs = lambda reqs: reqs - backend._reorder_long_prefill_reqs = lambda reqs: reqs + backend.prefill_queue_strategy = SimpleNamespace(reorder=lambda reqs, **kwargs: reqs) context.get_can_alloc_token_num = lambda: 12 context.cache_placement_controller = SimpleNamespace(set_req_cache_way=lambda reqs: None) context.filter_reqs = lambda finished_reqs: None @@ -308,8 +306,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step backend.batch_max_tokens = 8 backend.is_master_in_dp = True backend._timer_merge_radix_tree = lambda: None - backend._reorder_pd_high_priority_reqs = lambda reqs: reqs - backend._reorder_long_prefill_reqs = lambda reqs: reqs + backend.prefill_queue_strategy = SimpleNamespace(reorder=lambda reqs, **kwargs: reqs) context.get_can_alloc_token_num = lambda: 32 context.cache_placement_controller = SimpleNamespace(set_req_cache_way=lambda reqs: None) context.filter_reqs = lambda finished_reqs: None diff --git a/unit_tests/server/core/objs/test_sampling_params.py b/unit_tests/server/core/objs/test_sampling_params.py index d0e7fde893..8a82f37141 100644 --- a/unit_tests/server/core/objs/test_sampling_params.py +++ b/unit_tests/server/core/objs/test_sampling_params.py @@ -170,6 +170,12 @@ def test_sampling_params_initialization(): assert params.pd_kv_trans_params.get() == pd_kv_trans_params +def test_infer_high_priority_is_internal_and_ignores_external_value(): + params = SamplingParams() + params.init(None, infer_high_priority=-3) + assert params.infer_high_priority == 0 + + # Mock tokenizer for testing class MockTokenizer: def encode(self, text, add_special_tokens=False): diff --git a/unit_tests/server/httpserver/test_running_request_lifecycle.py b/unit_tests/server/httpserver/test_running_request_lifecycle.py index 6fe231a0a6..2ffc10d952 100644 --- a/unit_tests/server/httpserver/test_running_request_lifecycle.py +++ b/unit_tests/server/httpserver/test_running_request_lifecycle.py @@ -243,7 +243,7 @@ async def run(): router_arrival_time=0, infer_start_time=0, sample_params=SimpleNamespace( - pd_high_priority_request=False, + high_priority_request=False, pd_node_resource_wait_timeout_seconds=20, ), ) @@ -302,14 +302,14 @@ def test_httpserver_keeps_started_requests_and_requests_with_remaining_master_ti router_arrival_time=1.0, infer_start_time=0.0, sample_params=SimpleNamespace( - pd_high_priority_request=False, + high_priority_request=False, pd_node_resource_wait_timeout_seconds=60, ), ) started_req = SimpleNamespace( router_arrival_time=1.0, infer_start_time=2.0, - sample_params=SimpleNamespace(pd_high_priority_request=False), + sample_params=SimpleNamespace(high_priority_request=False), ) req_status = _req_status([waiting_req, started_req]) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py b/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py index 73e133cc18..e8d4f65826 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_pd_dynamic_split.py @@ -49,8 +49,7 @@ def _classify_without_token_capacity(monkeypatch, req, support_overlap=True): backend.logger = logger backend._timer_merge_radix_tree = MagicMock() backend._filter_not_ready_reqs = MagicMock(return_value=reqs) - backend._reorder_pd_high_priority_reqs = MagicMock(side_effect=lambda reqs: reqs) - backend._reorder_long_prefill_reqs = MagicMock(side_effect=lambda reqs: reqs) + backend.prefill_queue_strategy = SimpleNamespace(reorder=MagicMock(side_effect=lambda reqs, **kwargs: reqs)) infer_context = base_backend.g_infer_context monkeypatch.setattr(infer_context, "get_can_alloc_token_num", MagicMock(return_value=0)) diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py b/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py new file mode 100644 index 0000000000..be5ce32844 --- /dev/null +++ b/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py @@ -0,0 +1,161 @@ +from types import SimpleNamespace + +import pytest + +from lightllm.server.api_cli import make_argument_parser +from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend +from lightllm.server.router.model_infer.mode_backend.prefill_queue_strategy import ( + PREFILL_QUEUE_STRATEGIES, + PrefillQueueStrategy, + create_prefill_queue_strategy, +) + + +def req(group_id, input_len, kv_len=0, priority=0, router_priority=False, total_len=None): + return SimpleNamespace( + cur_kv_len=kv_len, + shm_req=SimpleNamespace( + input_len=input_len, + group_req_id=group_id, + sample_params=SimpleNamespace( + high_priority_request=router_priority, + infer_high_priority=priority, + ), + ), + get_cur_total_len=lambda: input_len if total_len is None else total_len, + ) + + +def strategy(name="default"): + backend = ModeBackend.__new__(ModeBackend) + backend.args = SimpleNamespace(prefill_queue_strategy=name) + return create_prefill_queue_strategy(backend) + + +@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) +def test_empty_and_singleton(name): + policy = strategy(name) + assert isinstance(policy, PrefillQueueStrategy) + assert policy.reorder([]) == [] + request = req(1, 10) + assert policy.reorder([request]) == [request] + + +@pytest.mark.parametrize( + "name, expected", + [ + ("default", [2, 4, 1, 3]), + ("promote_shortest", [2, 4, 3, 1]), + ], +) +def test_ordering_and_input_preservation(name, expected): + requests = [req(1, 20), req(2, 50, priority=-2), req(3, 102, 100), req(4, 5, priority=-1)] + original = list(requests) + result = strategy(name).reorder(requests) + assert [r.shm_req.group_req_id for r in result] == expected + assert all(a is b for a, b in zip(requests, original)) + assert sorted(map(id, result)) == sorted(map(id, original)) + + +def test_inference_priority_is_independent_from_router_priority(): + router_high_priority = req(1, 10, priority=0, router_priority=True) + infer_high_priority = req(2, 10, priority=-1, router_priority=False) + + assert strategy().reorder([router_high_priority, infer_high_priority]) == [ + infer_high_priority, + router_high_priority, + ] + + +def test_promote_shortest_clamps_remaining_tokens_and_preserves_ties(): + requests = [req(9, 5, 6), req(8, 5, 10), req(7, 5, 5), req(6, 10, 8), req(1, 2)] + assert strategy("promote_shortest").reorder(requests) == requests + + +def test_promote_shortest_keeps_high_priority_requests_first(): + long_high_priority = req(1, 1000, priority=-1) + short_normal_priority = req(2, 2, priority=0) + + assert strategy("promote_shortest").reorder([short_normal_priority, long_high_priority]) == [ + long_high_priority, + short_normal_priority, + ] + + +def test_promote_shortest_preserves_high_priority_arrival_order(): + first_high_priority = req(1, 100, priority=-1) + second_high_priority = req(2, 100, priority=-2) + + assert strategy("promote_shortest").reorder([first_high_priority, second_high_priority]) == [ + first_high_priority, + second_high_priority, + ] + + +def test_promote_shortest_moves_only_one_normal_request_and_preserves_other_order(): + first_normal = req(1, 20) + shortest_normal = req(2, 2) + third_normal = req(3, 10) + fourth_normal = req(4, 5) + + assert strategy("promote_shortest").reorder([first_normal, shortest_normal, third_normal, fourth_normal]) == [ + shortest_normal, + first_normal, + third_normal, + fourth_normal, + ] + + +@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) +def test_decode_requests_stay_left_and_are_not_reordered(name): + decode_normal = req(1, 5, kv_len=4, priority=0) + prefill_high = req(2, 100, priority=-2) + decode_high = req(3, 5, kv_len=4, priority=-10) + prefill_normal = req(4, 2, priority=0) + + assert strategy(name).reorder([decode_normal, prefill_normal, decode_high, prefill_high]) == [ + decode_normal, + decode_high, + prefill_high, + prefill_normal, + ] + + +def test_no_decode_and_strict_prefill_match_backend_classification_rules(): + boundary_req = req(1, 5, kv_len=4, priority=0) + high_priority_prefill = req(2, 10, priority=-1) + + policy = strategy() + assert policy.reorder([boundary_req, high_priority_prefill]) == [boundary_req, high_priority_prefill] + assert policy.reorder([boundary_req, high_priority_prefill], no_decode=True) == [ + high_priority_prefill, + boundary_req, + ] + assert policy.reorder([boundary_req, high_priority_prefill], strict_prefill=True) == [ + high_priority_prefill, + boundary_req, + ] + + +@pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) +def test_cli_strategy_selection(name): + args = make_argument_parser().parse_args(["--prefill_queue_strategy", name]) + backend = SimpleNamespace(args=args) + policy = create_prefill_queue_strategy(backend) + assert isinstance(policy, PREFILL_QUEUE_STRATEGIES[name]) + assert policy.backend is backend + + +def test_cli_default_and_invalid_strategy(): + parser = make_argument_parser() + assert parser.parse_args([]).prefill_queue_strategy == "default" + assert set(PREFILL_QUEUE_STRATEGIES) == {"default", "promote_shortest"} + with pytest.raises(SystemExit): + parser.parse_args(["--prefill_queue_strategy", "unknown"]) + with pytest.raises(ValueError, match="Unknown prefill queue strategy"): + strategy("unknown") + + +def test_strategy_base_is_abstract(): + with pytest.raises(TypeError): + PrefillQueueStrategy(SimpleNamespace()) diff --git a/unit_tests/utils/test_envs_utils.py b/unit_tests/utils/test_envs_utils.py index 0cac7caa9e..4ecfb10b56 100644 --- a/unit_tests/utils/test_envs_utils.py +++ b/unit_tests/utils/test_envs_utils.py @@ -79,11 +79,11 @@ def test_pd_node_continuation_resource_wait_timeout_reads_environment_variable(m get_pd_node_continuation_resource_wait_timeout_seconds.cache_clear() -def test_pd_node_busy_retry_timeout_defaults_to_120_seconds(monkeypatch): +def test_pd_node_busy_retry_timeout_defaults_to_zero(monkeypatch): monkeypatch.delenv("LIGHTLLM_PD_NODE_BUSY_RETRY_TIMEOUT_SECONDS", raising=False) get_pd_node_busy_retry_timeout_seconds.cache_clear() - assert get_pd_node_busy_retry_timeout_seconds() == 120 + assert get_pd_node_busy_retry_timeout_seconds() == 0 get_pd_node_busy_retry_timeout_seconds.cache_clear() From 5574809dbc3c094f924a6872775fe9dd75fb5226 Mon Sep 17 00:00:00 2001 From: wangzaijun Date: Wed, 23 Sep 2026 09:44:27 +0000 Subject: [PATCH 2/2] feat: add HRRN prefill queue scheduling --- docs/CN/source/tutorial/api_server_args.rst | 10 ++- .../source/tutorial/prefill_queue_strategy.md | 49 ++++++++++- docs/EN/source/tutorial/api_server_args.rst | 13 ++- lightllm/server/api_cli.py | 6 +- lightllm/server/core/objs/start_args_type.py | 4 +- .../server/router/model_infer/infer_batch.py | 4 + .../model_infer/mode_backend/base_backend.py | 8 +- .../mode_backend/chunked_prefill/impl.py | 3 +- .../mode_backend/dp_backend/impl.py | 3 +- .../prefill_queue_strategy/__init__.py | 3 +- .../prefill_queue_strategy/base.py | 6 +- .../prefill_queue_strategy/strategies.py | 28 +++++++ unit_tests/common/test_req_manager_page.py | 8 +- .../test_prefill_queue_strategy.py | 84 ++++++++++++++++--- 14 files changed, 194 insertions(+), 35 deletions(-) diff --git a/docs/CN/source/tutorial/api_server_args.rst b/docs/CN/source/tutorial/api_server_args.rst index d7aebc4f5e..d452097cb2 100644 --- a/docs/CN/source/tutorial/api_server_args.rst +++ b/docs/CN/source/tutorial/api_server_args.rst @@ -364,8 +364,14 @@ PD 分离模式参数 普通请求中选择剩余 prefill token 最少的一个,将其移动到普通请求队头,其余普通 请求保持相对顺序。已完成 prompt 计算的请求按零计算。该模式主要适合 PD 分离的 Prefill 节点,用于改善部分短请求的 TTFT 和首字体验。 - - 例如,在原启动命令中添加 ``--prefill_queue_strategy promote_shortest``。 + * ``hrrn``:负优先级 prefill 请求仍排在最前面;普通请求按基于 token aging 的最高 + 响应比优先(Highest Response Ratio Next)策略排序。响应比使用请求等待期间已处理的 + prefill token 数与该请求首次参与 HRRN 排序时未缓存的 prefill token 数计算,既倾向于较短请求,也会随等待量 + 增加逐步提升长请求优先级,避免纯最短任务优先造成饥饿。该模式适合 PD 分离的 Prefill + 节点中请求长度差异较大的流量。该策略参考了 + `SGLang PR #32911 `_,感谢原作者及 SGLang 社区的贡献。 + + 例如,在原启动命令中添加 ``--prefill_queue_strategy hrrn``。 应结合实际流量的 TTFT、TPOT 和尾延迟压测选择策略,策略本身不保证 SLA 改善。 完整的调度流程、算法细节和排序示例见 :doc:`Prefill 排队策略 `。 diff --git a/docs/CN/source/tutorial/prefill_queue_strategy.md b/docs/CN/source/tutorial/prefill_queue_strategy.md index f83f8cb0cb..cebe84c850 100644 --- a/docs/CN/source/tutorial/prefill_queue_strategy.md +++ b/docs/CN/source/tutorial/prefill_queue_strategy.md @@ -8,12 +8,13 @@ LightLLM 支持在推理进程中选择 prefill 请求的排队策略。该功 --prefill_queue_strategy default ``` -当前支持两种策略: +当前支持三种策略: | 策略 | 行为 | | --- | --- | | `default` | 负优先级 prefill 请求在前;相同优先级内保持 FCFS 顺序。 | | `promote_shortest` | 负优先级请求按到达顺序排在前面,只把一个最短的普通 prefill 请求提升到普通请求队头。 | +| `hrrn` | 负优先级请求保持在前,普通请求按基于 token aging 的最高响应比排序。 | ## 调度阶段 @@ -113,11 +114,34 @@ Python 的 `sorted()` 是稳定排序。因此: 该模式主要适合在 PD 分离架构的 Prefill 节点上使用。Prefill 节点同时排队处理多个长短不一的 prompt 时,优先调度一个剩余 prefill token 最少的普通请求,可以让部分短请求更早完成 Prefill 并进入 Decode,从而改善其 TTFT 和首字体验。续跑分段仍位于普通请求之前,不会被短请求越过。 +## `hrrn` 策略 + +`hrrn`(Highest Response Ratio Next,最高响应比优先)在偏向短请求的同时引入 aging,避免持续到来的短请求使长请求一直无法获得调度。LightLLM 使用与墙上时间无关的 token 形式: + +该策略的设计与实现参考了 [SGLang PR #32911](https://github.com/sgl-project/sglang/pull/32911),感谢原作者及 SGLang 社区的工作与分享。 + +```text +priority_score = waited_prefill_tokens / uncached_prefill_tokens +waited_prefill_tokens = processed_prefill_tokens - arrival_processed_prefill_tokens +uncached_prefill_tokens = max(1, input_len - cur_kv_len_at_first_hrrn_sort) +``` + +- `processed_prefill_tokens` 是当前推理后端累计提交给 prefill forward 执行的 token 数。 +- 请求进入推理等待队列时,会将当前计数保存为 `arrival_processed_prefill_tokens`。 +- 请求第一次参与 HRRN 排序时,策略通过动态属性保存 `uncached_prefill_tokens`。因此此前完成的 Prefix/CPU Cache 初始化会反映在服务成本中;属性绑定后保持不变,后续 chunked prefill 不会持续缩小分母。 +- 分数越高越先调度;服务成本最小按一个 token 计算,避免零分母分支。 +- 分数相同时使用请求 ID 作为确定性的次级排序键,确保各 TP rank 得到相同顺序。 +- `infer_high_priority < 0` 的内部高优先级请求始终位于 HRRN 普通请求之前。 + +如果吞吐量在一段时间内近似稳定,处理过的 prefill token 数与等待时间成正比,首次 HRRN 排序时的未缓存 prefill token 数与预计服务时间成正比。因此该公式与经典 HRRN 的 `1 + wait_time / service_time` 具有相同的排序效果,同时不依赖各 TP rank 可能略有差异的墙上时钟,也不需要配置预计 TPS。 + +该模式主要适合 PD 分离架构的 Prefill 节点,尤其是 prompt 长度和 Prefix Cache 命中长度差异较大的流量。短请求通常能更早完成 Prefill,从而改善 TTFT 和首字体验;长请求的分数会随等待期间处理的 token 数持续增长,从而降低饥饿和尾延迟风险。 + ## 排序示例 假设原始队列如下: -| 到达顺序 | 请求 | 阶段 | `infer_high_priority` | 剩余 prefill token | +| 到达顺序 | 请求 | 阶段 | `infer_high_priority` | 首次 HRRN 排序时的 prefill 工作量 | | --- | --- | --- | ---: | ---: | | 1 | `normal-a` | prefill | 0 | 100 | | 2 | `decode-a` | decode | 0 | 不参与 | @@ -138,12 +162,18 @@ decode-a, decode-b, pd-high, normal-a, normal-b, normal-short decode-a, decode-b, pd-high, normal-short, normal-a, normal-b ``` +假设三个普通请求进入队列时保存的计数相同,且之后已经处理了一批 prefill token,`hrrn` 的结果为: + +```text +decode-a, decode-b, pd-high, normal-short, normal-b, normal-a +``` + 可以看到: - decode 请求始终在最左侧,并保持 `decode-a, decode-b` 的相对顺序。 - `pd-high` 不会被普通短请求越过。 -- 只有 `normal-short` 被提升。 -- `normal-a` 和 `normal-b` 的相对顺序没有改变。 +- `promote_shortest` 只提升 `normal-short`,`normal-a` 和 `normal-b` 的相对顺序不变。 +- `hrrn` 会按响应比排列全部普通请求;随着等待量变化,长请求也能逐步前移。 ## 选择策略 @@ -163,6 +193,14 @@ python -m lightllm.server.api_server \ --prefill_queue_strategy promote_shortest ``` +使用 HRRN 调度普通 Prefill 请求: + +```bash +python -m lightllm.server.api_server \ + --model_dir /path/to/model \ + --prefill_queue_strategy hrrn +``` + 建议使用与生产环境一致的请求长度分布、并发度和 prefix cache 命中率进行压测,重点观察: - TTFT 及其 P95/P99。 @@ -170,9 +208,12 @@ python -m lightllm.server.api_server \ - 长 prompt 请求的等待时间。 - 计算 token 与 KV cache 容量紧张时的尾延迟。 - PD 续跑分段能否及时获得推理资源。 +- 长请求是否能通过 aging 在合理时间内得到调度。 `promote_shortest` 主要用于改善 PD 分离模式下 Prefill 节点中部分短请求的首字体验。它只提升一个普通请求,通常比完整 SJF 更温和,但具体 TTFT 和 SLA 收益仍取决于实际流量。 +`hrrn` 更适合长短请求差异明显且需要兼顾平均 TTFT 与长请求公平性的场景。其服务成本使用请求首次参与 HRRN 排序时的未缓存 token 数近似;对于每 token 成本随上下文长度明显增长的模型,这个估算并不完全等价于真实 Prefill 时间,仍应通过生产流量压测确认收益。 + ## 扩展新策略 策略代码位于: diff --git a/docs/EN/source/tutorial/api_server_args.rst b/docs/EN/source/tutorial/api_server_args.rst index a1f66dcbc4..d9a1d3eb4d 100644 --- a/docs/EN/source/tutorial/api_server_args.rst +++ b/docs/EN/source/tutorial/api_server_args.rst @@ -384,8 +384,17 @@ Scheduling Parameters of all other normal requests. Fully processed prompts count as zero. This mode is mainly intended for Prefill nodes in PD-disaggregated deployments to improve TTFT for some short requests. - - For example, append ``--prefill_queue_strategy promote_shortest`` to your launch command. + * ``hrrn``: keep negative-priority prefill requests first, then order normal requests with + token-based Highest Response Ratio Next. The response ratio uses the number of prefill + tokens processed while a request has waited and the uncached prefill tokens captured when it first + participates in HRRN ordering. + It favors short requests while gradually aging long-waiting requests to prevent starvation. + This mode is intended for Prefill nodes in PD-disaggregated deployments with highly varied + request lengths. This strategy is based on + `SGLang PR #32911 `_; thanks to its authors + and the SGLang community for sharing their work. + + For example, append ``--prefill_queue_strategy hrrn`` to your launch command. Compare TTFT, TPOT, and tail latency under representative traffic before choosing a policy; changing policies does not guarantee an SLA improvement. diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index e7473e919c..2d0b264d99 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -372,13 +372,15 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser.add_argument("--disable_chunked_prefill", action="store_true", help="whether to disable chunked prefill") parser.add_argument( "--prefill_queue_strategy", - choices=("default", "promote_shortest"), + choices=("default", "promote_shortest", "hrrn"), default="default", help="""Order ready requests before resource allocation. Decode requests are identified first, remain in arrival order, and are not reordered by these policies. default: preserve FCFS order for prefill requests with the same infer_high_priority value. promote_shortest: keep negative-priority requests first in arrival order, then move one - shortest non-negative-priority prefill request to the front of the normal requests.""", + shortest non-negative-priority prefill request to the front of the normal requests. + hrrn: order normal prefill requests by token-based highest response ratio next to favor + short requests while aging long-waiting requests; negative-priority requests remain first.""", ) parser.add_argument("--diverse_mode", action="store_true", help="diversity generation mode") vocab_topk_choices = [2, 8, 16, 32, 64, 128, 256, 512] diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 8abdf9fdb3..f66eae79af 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -101,7 +101,9 @@ class StartArgs: disable_dynamic_prompt_cache: bool = field(default=False) chunked_prefill_size: int = field(default=None) disable_chunked_prefill: bool = field(default=False) - prefill_queue_strategy: str = field(default="default", metadata={"choices": ["default", "promote_shortest"]}) + prefill_queue_strategy: str = field( + default="default", metadata={"choices": ["default", "promote_shortest", "hrrn"]} + ) diverse_mode: bool = field(default=False) target_vocab_topk_sampling: Optional[int] = field( default=None, metadata={"choices": [2, 8, 16, 32, 64, 128, 256, 512]} diff --git a/lightllm/server/router/model_infer/infer_batch.py b/lightllm/server/router/model_infer/infer_batch.py index 99f58ad67b..69c5ebe49b 100644 --- a/lightllm/server/router/model_infer/infer_batch.py +++ b/lightllm/server/router/model_infer/infer_batch.py @@ -34,6 +34,7 @@ @dataclass class InferenceContext: + backend: Optional["ModeBackend"] = None req_manager: ReqManager = None # gpu 请求管理 radix_cache: Union[HybridAttPagedRadixCache, RadixCache] = None shm_req_manager: ShmReqManager = None # 共享内存请求对象管理 @@ -554,6 +555,9 @@ def __init__( self.need_out_token_id_statistics = True self.out_token_id_count: Dict[int, int] = None + # InferReq 创建时直接快照后端累计的 prefill token,作为 HRRN aging 的起点。 + self.arrival_processed_prefill_tokens: int = g_infer_context.backend.processed_prefill_tokens + # diverse mode 下,用于标记请求组之间的依赖关系 self.slave_reqs: List[InferReq] = [] self.related_master_req: InferReq = None diff --git a/lightllm/server/router/model_infer/mode_backend/base_backend.py b/lightllm/server/router/model_infer/mode_backend/base_backend.py index 27e3ce3e00..7e284f1fa0 100644 --- a/lightllm/server/router/model_infer/mode_backend/base_backend.py +++ b/lightllm/server/router/model_infer/mode_backend/base_backend.py @@ -83,6 +83,8 @@ def __init__(self) -> None: self._radix_tree_merge_counter: int = 0 self._enable_radix_tree_timer_merge: bool = enable_radix_tree_timer_merge() self._radix_tree_merge_update_delta: int = get_radix_tree_merge_update_delta() + # 已提交执行的 prefill token 总数,作为 HRRN 与墙上时间无关的 aging 轴。 + self.processed_prefill_tokens: int = 0 pass def init_model(self, kvargs): @@ -683,6 +685,7 @@ def _get_classed_reqs( 3. finished_reqs 需要释放的请求, 包含正常结束和aborted退出的请求。 4. prefill_reqs 需要进行prefill操作的请求 5. decode_reqs 需要进行decode操作的请求 + 6. prefill_tokens 本轮选出的 prefill batch 需要计算的 token 数 """ # 定期对 radix cache 进行 merge,防止查询插入的操作效率下降 self._timer_merge_radix_tree() @@ -694,7 +697,7 @@ def _get_classed_reqs( req_ids = g_infer_context.infer_req_ids if len(req_ids) == 0: - return [], [] + return [], [], 0 ready_reqs = self._filter_not_ready_reqs(req_ids) support_overlap = self.support_overlap @@ -830,6 +833,7 @@ def _get_classed_reqs( # 则将 decode 请求合并到 prefill 请求中。 if self.args.enable_prefill_decode_mixed and len(prefill_reqs) > 0 and len(decode_reqs) > 0: if prefill_tokens + len(decode_reqs) <= self.batch_max_tokens: + prefill_tokens += len(decode_reqs) for decode_req in decode_reqs: # 给 decode req 添加一个属性标签,标识其为混合prefill的请求。 # 在 prefill 阶段,会根据这个属性标签, 对这些请求的处理进行一些 @@ -838,7 +842,7 @@ def _get_classed_reqs( prefill_reqs.append(decode_req) decode_reqs = [] - return prefill_reqs, decode_reqs + return prefill_reqs, decode_reqs, prefill_tokens # 一些可以复用的通用功能函数 def _pre_post_handle(self, run_reqs: List[InferReq], is_chuncked_mode: bool) -> List[InferReqUpdatePack]: diff --git a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py index d0026da9bd..6af071238f 100644 --- a/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/chunked_prefill/impl.py @@ -62,7 +62,7 @@ def infer_loop(self): self._try_read_new_reqs() - prefill_reqs, decode_reqs = self._get_classed_reqs( + prefill_reqs, decode_reqs, prefill_tokens = self._get_classed_reqs( no_decode=self.classed_req_no_decode, strict_prefill=self.classed_req_strict_prefill, recover_paused=self.control_state_machine.try_recover_paused_reqs(), @@ -74,6 +74,7 @@ def infer_loop(self): # 进行一次流同步,保证 _try_read_new_reqs 中的一些算子操作,必然已经完成。 # 防止后续的推理流程读取到显存中可能存在错误的数据。 g_infer_context.get_overlap_stream().wait_stream(torch.cuda.current_stream()) + self.processed_prefill_tokens += prefill_tokens self.prefill( event_pack=event_pack, prefill_reqs=prefill_reqs, diff --git a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py index 2cb4e89e5b..a14610f5ba 100644 --- a/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py +++ b/lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py @@ -123,7 +123,7 @@ def infer_loop(self): self._try_read_new_reqs() - prefill_reqs, decode_reqs = self._get_classed_reqs( + prefill_reqs, decode_reqs, prefill_tokens = self._get_classed_reqs( no_decode=self.classed_req_no_decode, strict_prefill=self.classed_req_strict_prefill, recover_paused=self.control_state_machine.try_recover_paused_reqs(), @@ -144,6 +144,7 @@ def infer_loop(self): # 进行一次流同步,保证 _try_read_new_reqs 中的一些算子操作,必然已经完成。 # 防止后续的推理流程读取到显存中可能存在错误的数据。 g_infer_context.get_overlap_stream().wait_stream(torch.cuda.current_stream()) + self.processed_prefill_tokens += prefill_tokens self.prefill( event_pack=event_pack, prefill_reqs=prefill_reqs, diff --git a/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/__init__.py b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/__init__.py index 4e621d5541..3b7a57a3e7 100644 --- a/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/__init__.py +++ b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/__init__.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING from .base import PrefillQueueStrategy -from .strategies import FCFSStrategy, PromoteShortestPrefillStrategy +from .strategies import FCFSStrategy, HRRNPrefillStrategy, PromoteShortestPrefillStrategy if TYPE_CHECKING: from ..base_backend import ModeBackend @@ -11,6 +11,7 @@ PREFILL_QUEUE_STRATEGIES = { "default": FCFSStrategy, "promote_shortest": PromoteShortestPrefillStrategy, + "hrrn": HRRNPrefillStrategy, } diff --git a/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/base.py b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/base.py index c8033b0905..88042676d5 100644 --- a/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/base.py +++ b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/base.py @@ -10,9 +10,9 @@ class PrefillQueueStrategy(ABC): """在请求分类和 token 预算分配前调整就绪请求的排队顺序。 输入中可能包含 decode、暂停或已结束的请求。基类先把 decode 请求稳定地放到最左侧, - 子类只负责排序右侧的 prefill 请求。实现必须返回原请求集合的一个排列,不能修改请求状态、 - 丢弃请求或原地修改输入列表。各 TP rank 的排序结果必须一致;除非策略明确提供次级排序条件, - 否则应保持相同排序键请求的原始顺序。 + 子类只负责排序右侧的 prefill 请求。实现必须返回原请求集合的一个排列;除策略自身的惰性 + 调度元数据外,不能修改请求状态、丢弃请求或原地修改输入列表。各 TP rank 的排序结果必须 + 一致;除非策略明确提供次级排序条件,否则应保持相同排序键请求的原始顺序。 策略可通过 ``self.backend`` 读取已完成初始化的模型、缓存、rank 等调度状态。 """ diff --git a/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/strategies.py b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/strategies.py index 76efb6f3d1..fa796e13e0 100644 --- a/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/strategies.py +++ b/lightllm/server/router/model_infer/mode_backend/prefill_queue_strategy/strategies.py @@ -41,3 +41,31 @@ def reorder_prefill(self, prefill_reqs: List["InferReq"]) -> List["InferReq"]: high_priority_reqs.append(shortest_req) high_priority_reqs.extend(normal_reqs) return high_priority_reqs + + +class HRRNPrefillStrategy(PrefillQueueStrategy): + """使用基于 token aging 的最高响应比优先策略排列普通 prefill 请求。 + + 负优先级请求仍按原始相对顺序排在最前面。普通请求按 + ``waited_prefill_tokens / uncached_prefill_tokens`` 从大到小排列;等待越久或 + 首次参与 HRRN 排序时的未缓存工作量越小,越早获得调度。请求 ID 用作相同比率下的 + 确定性次级排序键。 + """ + + def reorder_prefill(self, prefill_reqs: List["InferReq"]) -> List["InferReq"]: + high_priority_reqs = [req for req in prefill_reqs if infer_priority(req) < 0] + normal_reqs = [req for req in prefill_reqs if infer_priority(req) >= 0] + processed_prefill_tokens = self.backend.processed_prefill_tokens + + def hrrn_key(req: "InferReq"): + # 请求第一次进入 HRRN 排序时才绑定服务成本,使此前完成的 Cache 初始化 + # 能反映在分母中;后续 chunked prefill 不再改变该请求的服务成本。 + if not hasattr(req, "uncached_prefill_tokens"): + req.uncached_prefill_tokens = max(1, remaining_prefill_tokens(req)) + uncached_prefill_tokens = req.uncached_prefill_tokens + waited_tokens = max(0, processed_prefill_tokens - req.arrival_processed_prefill_tokens) + return (-waited_tokens / uncached_prefill_tokens, req.req_id) + + normal_reqs.sort(key=hrrn_key) + high_priority_reqs.extend(normal_reqs) + return high_priority_reqs diff --git a/unit_tests/common/test_req_manager_page.py b/unit_tests/common/test_req_manager_page.py index f1d7a5ec67..224c308f28 100644 --- a/unit_tests/common/test_req_manager_page.py +++ b/unit_tests/common/test_req_manager_page.py @@ -224,10 +224,11 @@ def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): backend._alloc_req_kv_mem = _record_alloc - prefill_reqs, decode_reqs = backend._get_classed_reqs(req_ids=[0, 1]) + prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0, 1]) assert prefill_reqs == [reqs[0]] assert decode_reqs == [] + assert prefill_tokens == 3 assert context.req_manager.mem_manager.alloc_sizes == expected_alloc_sizes assert reqs[0].hold_kv_len == 8 assert not reqs[0].wait_pause @@ -280,10 +281,11 @@ def _record_alloc(req_obj, alloc_token_num, no_blcoking_copy=False): backend._alloc_req_kv_mem = _record_alloc - prefill_reqs, decode_reqs = backend._get_classed_reqs(req_ids=[0]) + prefill_reqs, decode_reqs, prefill_tokens = backend._get_classed_reqs(req_ids=[0]) assert prefill_reqs == [] assert decode_reqs == [req] + assert prefill_tokens == 0 assert req.hold_kv_len == 16 assert context.req_manager.mem_manager.alloc_sizes == [12] assert copy_modes == [True] @@ -329,7 +331,7 @@ def test_decode_small_page_preallocation_reuses_one_allocation_for_multiple_step backend._filter_not_ready_reqs = lambda req_ids: [req] for _ in range(expected_alloc_token_num): - _, decode_reqs = backend._get_classed_reqs(req_ids=[0]) + _, decode_reqs, _ = backend._get_classed_reqs(req_ids=[0]) assert decode_reqs == [req] req.cur_kv_len += 1 diff --git a/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py b/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py index be5ce32844..7ce0d3f9c5 100644 --- a/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py +++ b/unit_tests/server/router/model_infer/mode_backend/test_prefill_queue_strategy.py @@ -3,6 +3,7 @@ import pytest from lightllm.server.api_cli import make_argument_parser +from lightllm.server.router.model_infer.infer_batch import InferReq from lightllm.server.router.model_infer.mode_backend.base_backend import ModeBackend from lightllm.server.router.model_infer.mode_backend.prefill_queue_strategy import ( PREFILL_QUEUE_STRATEGIES, @@ -11,24 +12,35 @@ ) -def req(group_id, input_len, kv_len=0, priority=0, router_priority=False, total_len=None): - return SimpleNamespace( - cur_kv_len=kv_len, - shm_req=SimpleNamespace( - input_len=input_len, - group_req_id=group_id, - sample_params=SimpleNamespace( - high_priority_request=router_priority, - infer_high_priority=priority, - ), +def req( + group_id, + input_len, + kv_len=0, + priority=0, + router_priority=False, + total_len=None, + arrival_processed_prefill_tokens=0, +): + request = InferReq.__new__(InferReq) + request.req_id = group_id + request.cur_kv_len = kv_len + request.arrival_processed_prefill_tokens = arrival_processed_prefill_tokens + request.shm_req = SimpleNamespace( + input_len=input_len, + group_req_id=group_id, + sample_params=SimpleNamespace( + high_priority_request=router_priority, + infer_high_priority=priority, ), - get_cur_total_len=lambda: input_len if total_len is None else total_len, ) + request.get_cur_total_len = lambda: input_len if total_len is None else total_len + return request -def strategy(name="default"): +def strategy(name="default", processed_prefill_tokens=0): backend = ModeBackend.__new__(ModeBackend) backend.args = SimpleNamespace(prefill_queue_strategy=name) + backend.processed_prefill_tokens = processed_prefill_tokens return create_prefill_queue_strategy(backend) @@ -46,6 +58,7 @@ def test_empty_and_singleton(name): [ ("default", [2, 4, 1, 3]), ("promote_shortest", [2, 4, 3, 1]), + ("hrrn", [2, 4, 1, 3]), ], ) def test_ordering_and_input_preservation(name, expected): @@ -106,6 +119,51 @@ def test_promote_shortest_moves_only_one_normal_request_and_preserves_other_orde ] +def test_hrrn_prioritizes_short_work_for_equally_aged_requests(): + requests = [req(1, 100), req(2, 10), req(3, 2)] + + assert strategy("hrrn", processed_prefill_tokens=1000).reorder(requests) == [ + requests[2], + requests[1], + requests[0], + ] + + +def test_hrrn_aging_allows_long_request_to_overtake_new_short_request(): + long_old = req(1, 100, arrival_processed_prefill_tokens=0) + short_new = req(2, 2, arrival_processed_prefill_tokens=100_000) + + assert strategy("hrrn", processed_prefill_tokens=100_000).reorder([short_new, long_old]) == [ + long_old, + short_new, + ] + + +def test_hrrn_binds_uncached_tokens_once_and_uses_request_id_tiebreaker(): + less_cached = req(2, 100, kv_len=0) + equally_cached_later_id = req(3, 100, kv_len=90) + more_cached = req(1, 100, kv_len=90) + policy = strategy("hrrn", processed_prefill_tokens=1000) + + requests = [less_cached, equally_cached_later_id, more_cached] + assert policy.reorder(requests) == [more_cached, equally_cached_later_id, less_cached] + assert [request.uncached_prefill_tokens for request in requests] == [100, 10, 10] + + less_cached.cur_kv_len = 95 + more_cached.cur_kv_len = 0 + assert policy.reorder(requests) == [more_cached, equally_cached_later_id, less_cached] + + +def test_hrrn_keeps_internal_high_priority_requests_first(): + shortest_normal = req(1, 2, priority=0) + long_high_priority = req(2, 1000, priority=-1) + + assert strategy("hrrn", processed_prefill_tokens=1000).reorder([shortest_normal, long_high_priority]) == [ + long_high_priority, + shortest_normal, + ] + + @pytest.mark.parametrize("name", PREFILL_QUEUE_STRATEGIES) def test_decode_requests_stay_left_and_are_not_reordered(name): decode_normal = req(1, 5, kv_len=4, priority=0) @@ -149,7 +207,7 @@ def test_cli_strategy_selection(name): def test_cli_default_and_invalid_strategy(): parser = make_argument_parser() assert parser.parse_args([]).prefill_queue_strategy == "default" - assert set(PREFILL_QUEUE_STRATEGIES) == {"default", "promote_shortest"} + assert set(PREFILL_QUEUE_STRATEGIES) == {"default", "promote_shortest", "hrrn"} with pytest.raises(SystemExit): parser.parse_args(["--prefill_queue_strategy", "unknown"]) with pytest.raises(ValueError, match="Unknown prefill queue strategy"):