跳转至

批处理作业感知调度器

批处理作业感知调度器是一种专为离线批处理推理场景设计的专用调度器,在这些场景中,吞吐量和硬件利用率是主要目标。当并发处理多个批处理作业且每个作业具有不同的一组请求时,它尤其有效。

注意: 该调度器未实现饥饿预防。它仅适用于批处理场景或不关注请求等待时间的在线推理场景。如果请求有严格的延迟或公平性要求,则此调度器可能不适用。

概述

该调度器实现了三种关键策略以提升吞吐量:

  1. LPT(最长处理时间优先)调度:优先调度较长的任务,并调度较短的任务以填补间隙,尤其是在解码步骤中。这提高了每个调度轮次平均计算的 token 数。

  2. KV 缓存预留:为正在运行的请求提前估计并预留 KV 缓存预算,从而减少抢占开销。

  3. 作业感知请求分组:按作业名称(从请求 ID 中提取)对请求进行分组,为每个作业分配独立的请求桶,并根据 KV 缓存可用情况动态调整作业调度顺序:

  4. 当可用 token 数 > 阈值(默认 4096)时:优先调度长解码作业
  5. 当可用 token 数 ≤ 阈值时:优先调度短解码作业

该调度器通过扩展等待队列并加入动态优先级调度来实现,同时继承原始调度器类的大多数功能,包括分块预填充、异步调度等。

工作原理

作业名称提取

调度器通过嵌入在请求 ID 中的 #job_name[${JOB_NAME}]# 标签来识别请求属于哪个作业。例如:

request_id = "req_001#job_name[my_batch_job]#"

没有作业名称标签的请求将被归入 __default__ 作业。

解码长度估计

每个作业的解码长度使用纯 EWMA(指数加权移动平均) 估计器进行预测:

  • 尚无样本:返回冷启动默认值(128 个 token)。
  • 首次观测:EWMA 初始化为观测到的解码长度。
  • 后续观测:EWMA 增量更新,对近期数据赋予更高权重,同时平滑波动。

该方法简单、响应迅速,避免了多阶段估计的复杂性,即使在数据有限的情况下也能产生稳定的预测。

这使得调度器能够区分“长解码”作业(在资源充足时优先调度更有益)和“短解码”作业(在资源稀缺时优先调度)。

快速开始

前提条件

  • 需要 vLLM v1 引擎(批处理作业感知调度器基于 v1 调度框架构建)。
  • 具有足够内存以支持目标模型的 Ascend NPU。

启用该功能

批处理作业感知调度器通过 additional_config 参数启用。目前仅支持离线批处理模式。

python -m vllm.entrypoints.openai.run_batch \
    --model /path/to/model \
    -i /path/to/input.jsonl \
    -o /path/to/output.jsonl \
    --additional-config '{"scheduler_config": {"batch_job_sched_config": {"enabled": true}}}'

请求 ID 格式

要利用作业感知调度,请在批处理输入文件中的请求 ID 内嵌入 #job_name[...]# 标签:

# In your batch input file (e.g., input.jsonl):
{"custom_id": "#job_name[job_A]#req_001", "method": "POST", "url": "/v1/chat/completions", "body": {"request_id": "#job_name[job_A]#req_001", "messages": [{"role": "user", "content": "Hello"}], "n": 1}}
{"custom_id": "#job_name[job_A]#req_002", "method": "POST", "url": "/v1/chat/completions", "body": {"request_id": "#job_name[job_A]#req_002", "messages": [{"role": "user", "content": "What is AI?"}], "n": 1}}
{"custom_id": "#job_name[job_B]#req_003", "method": "POST", "url": "/v1/chat/completions", "body": {"request_id": "#job_name[job_B]#req_003", "messages": [{"role": "user", "content": "Explain quantum computing"}], "n": 1}}

没有作业名称标签的请求将被归入默认作业,并且仍然可以从调度器的 KV 缓存预留和 LPT 调度中受益。

配置参数

所有参数都嵌套在 additional_config: {"scheduler_config":: {}} 中的 batch_job_sched_config 下。

参数 类型 默认值 描述
enabled bool false 启用批处理作业感知调度器
max_jobs int 20 跟踪的最大作业数
reserve_margin_blocks int 2 作为安全缓冲添加到 KV 缓存预留中的额外块余量
reserve_max_blocks int 8 可预留的最大块数
low_available_tokens_threshold int 4096 用于优先调度长解码作业与短解码作业的阈值
short_decode_token_threshold int 32 将作业分类为“短解码”的阈值

使用

基础离线批处理

python -m vllm.entrypoints.openai.run_batch \
    --model /path/to/model \
    -i /path/to/input.jsonl \
    -o /path/to/output.jsonl \
    --max-model-len 4096 \
    --gpu-memory-utilization 0.9 \
    --additional-config '{"scheduler_config": {"batch_job_sched_config": {"enabled": true}}}'

使用自定义配置的离线批处理

python -m vllm.entrypoints.openai.run_batch \
    --model /path/to/model \
    -i /path/to/input.jsonl \
    -o /path/to/output.jsonl \
    --max-model-len 4096 \
    --gpu-memory-utilization 0.9 \
    --additional-config '{
        "scheduler_config": {
            "batch_job_sched_config": {
                "enabled": true,
                "max_jobs": 10,
                "reserve_margin_blocks": 4,
                "reserve_max_blocks": 12,
                "low_available_tokens_threshold": 2048,
                "short_decode_token_threshold": 32
            }
        }
    }'

通过 Python API 使用

from vllm import LLM

llm = LLM(
    model="/path/to/model",
    max_model_len=4096,
    gpu_memory_utilization=0.9,
    additional_config={
        "scheduler_config": {
            "batch_job_sched_config": {
                "enabled": True,
            },
        },
    },
)

最佳实践

  1. 在请求 ID 中编码作业名称:在请求 ID 中使用 #job_name[${JOB_NAME}]# 前缀,以帮助调度器有效地分组和优先处理请求。

  2. 调整 low_available_tokens_threshold:如果您的工作负载持续以长解码为主,请考虑降低此阈值,以保持对长作业的优先处理。对于混合工作负载,请保留默认值。

本页原文 Markdown:在 AtomGit 查看·内容源自开源项目 vl/vllm-ascend