diff --git "a/T2-1-2_\344\274\212\351\207\214\346\226\257_\350\265\233\351\242\230\346\212\245\345\221\212.md" "b/T2-1-2_\344\274\212\351\207\214\346\226\257_\350\265\233\351\242\230\346\212\245\345\221\212.md" new file mode 100644 index 00000000..5d75bfa7 --- /dev/null +++ "b/T2-1-2_\344\274\212\351\207\214\346\226\257_\350\265\233\351\242\230\346\212\245\345\221\212.md" @@ -0,0 +1,237 @@ +# T2-1-2 服务能力优化报告 + +## 摘要 + +本提交针对 InfiniLM 在 RTX 4090 24GB 上运行 Llama-3.1-8B-Instruct 的高并发与长文本服务场景进行调度优化。本组成员在调测过程中发现，原始调度在部分长输入或高并发 prefill/decode 交叠场景下会触发 `cudaMalloc failed`，使 RankWorker 停止并导致服务 unhealthy。改动将 prefill 与 decode 的 batch 上限拆分，并对真正的长输入实施保守保护，在保持小规模请求性能无明显下降的同时，提升高并发短输出场景吞吐，并消除已复现的长输入服务失效。 + +受当前可用硬件与模型存储条件限制，本次提交仅完成 8B 模型的完整测试矩阵与优化验证，未对 70B 模型给出实测结果。本文所有性能数据、截图和结论均以 8B 模型为准；调度策略本身不依赖 8B 特有结构，理论上可迁移到 70B 服务，但仍需要在具备足够显存、模型存储和多卡通信条件的环境中单独验证。 + +## 环境与复现 + +- InfiniLM commit: `7c1efb5` +- 模型: `Llama-3.1-8B-Instruct` +- GPU: NVIDIA RTX 4090 24GB, TP=1 +- KV cache: `--num-blocks 64`, `--block-size 256` +- 测试工具：`vllm bench serve` +- 测试范围：官方 8B 测试矩阵，共 30 个点 + + | 并发数 | 输入长度（tokens） | 输出长度（tokens） | + | -------- | ------------- | --------------- | + | 1, 4 | 32, 256, 4096 | 256, 1024, 4096 | + | 16, 64 | 32, 256 | 256, 1024, 4096 | + +- 说明：本节参数为本组在 RTX 4090 24GB 单卡环境下针对 8B 测试矩阵调测得到的稳定配置，并不宣称为所有硬件和模型上的全局最优配置。 + +关键参数选择： + +- `--num-blocks 64`：在 4090 24GB 上可稳定运行完整 8B 矩阵。 +- `--max-batch-size 64`：保留高并发请求的 admission 上限，实际 prefill/decode 批次由调度器环境变量进一步约束。 +- `INFINILM_PREFILL_BATCH_CAP=8`：限制 prefill 阶段瞬时显存峰值。 +- `INFINILM_DECODE_BATCH_CAP=24`：允许 decode 阶段使用更大的安全批次，提高输出吞吐。 +- `INFINILM_LONG_CONTEXT_THRESHOLD=4096` 与 `INFINILM_LONG_CONTEXT_MAX_BATCH=1`：对真正长输入请求进行保守保护。 +- `INFINILM_DECODE_FIRST_WHEN_RUNNING=1`：已有 decode 请求运行时延后新的长输入 prefill，避免长 prefill 与 decode 叠加触发 OOM。 + +反例参数： + +- `num-blocks=80/96`：服务可启动，但在长输入 forward 时触发 OOM。 +- `block-size=128`：触发 paged-attention/graph 编译失败。 +- 仅降低全局 `max-batch-size` 到 2 或 3：仍不能稳定解决长输入卡点，说明问题不只是请求数，而是 prefill/decode 阶段显存压力不同。 + +复现启动命令： + +```bash +export INFINILM_LONG_CONTEXT_THRESHOLD=4096 +export INFINILM_LONG_CONTEXT_MAX_BATCH=1 +unset INFINILM_LONG_CONTEXT_USE_TOTAL_LEN +export INFINILM_DECODE_FIRST_WHEN_RUNNING=1 +export INFINILM_PREFILL_BATCH_CAP=8 +export INFINILM_DECODE_BATCH_CAP=24 + +python python/infinilm/server/inference_server.py \ + --device nvidia --model=/root/autodl-tmp/models/Llama-3.1-8B-Instruct \ + --port 8102 --tp 1 --max-new-tokens 4096 --num-blocks 64 \ + --max-batch-size 64 --enable-graph --enable-paged-attn \ + --attn flash-attn --ignore-eos +``` + +## 优化内容 + +修改文件：`python/infinilm/llm/scheduler.py`。 + +1. 增加 `INFINILM_PREFILL_BATCH_CAP` 与 `INFINILM_DECODE_BATCH_CAP`：允许 prefill 保持保守批次以控制瞬时显存，同时让 decode 使用更大的安全批次以提升输出吞吐。 +2. 增加长输入判定与 `INFINILM_LONG_CONTEXT_MAX_BATCH`：默认只将 `prompt_len >= 4096` 的请求视为长上下文，避免把短输入、长输出误串行化。 +3. 增加 `INFINILM_DECODE_FIRST_WHEN_RUNNING`：当已有 decode 在运行时，延后新的长输入 prefill，避免长 prefill 与长 decode 重叠导致 OOM。 +4. 所有策略均由环境变量控制；默认值保持原始 batch 行为，便于回退和复现。 + +## 问题与验证 + +本组在原始配置调测过程中，于以下点复现 `cudaMalloc failed`、`RankWorker stopped` 与 `/health` 异常： + +```text +con=4, input=4096, output=1024 +con=16, input=256, output=256 +``` + +启用本文调度策略后，长输入卡点稳定完成： + +| 场景 | 成功请求 | 输出吞吐 | Mean TTFT | Mean/P99 TPOT | Mean/P99 ITL | 服务状态 | +| ------------------------ | ----: | ----------: | -----------: | ---------------: | ---------------: | ------- | +| con=4, in=4096, out=1024 | 40/40 | 53.89 tok/s | 54604.96 ms | 18.09 / 18.24 ms | 18.07 / 18.90 ms | healthy | +| con=4, in=4096, out=4096 | 40/40 | 54.46 tok/s | 203503.93 ms | 17.33 / 18.61 ms | 18.24 / 19.64 ms | healthy | +| con=16, in=4096, out=256 | 8/8 | 50.75 tok/s | 17291.99 ms | 22.60 / 36.95 ms | 22.51 / 36.27 ms | healthy | + +高并发分相位限批的消融结果： + +| 场景 | 策略 | 成功请求 | 输出 tok/s | 总 tok/s | Mean/P99 TTFT ms | Mean/P99 TPOT ms | Mean/P99 ITL ms | 结果 | +|---|---|---:|---:|---:|---:|---:|---:|---| +| con=16, in=256, out=256 | 静态 max-batch=8 | 80/80 | 407.30 | 814.65 | 437.84 / 619.32 | 37.70 / 39.20 | 37.54 / 39.61 | healthy | +| con=16, in=256, out=256 | prefill=8, decode=24 | 80/80 | 726.03 | 1445.77 | 468.25 / 648.21 | 20.26 / 21.79 | 20.20 / 22.70 | healthy | +| con=64, in=256, out=256 | 静态 max-batch=8 | 128/128 | 376.64 | 753.37 | 21498.23 / 69212.98 | 51.69 / 53.70 | 51.50 / 61.84 | healthy | +| con=64, in=256, out=256 | prefill=8, decode=24 | 128/128 | 725.92 | 1449.00 | 10944.62 / 25269.96 | 24.19 / 25.64 | 24.09 / 27.34 | healthy | +| con=64, in=32, out=256 | prefill=8, decode=24 | 64/64 | 1052.16 | 1180.10 | 4299.80 / 8402.76 | 29.31 / 30.72 | 29.20 / 47.44 | healthy | + +## 边界与解释 + +`output=4096` 的高并发受 KV 容量而非 decode kernel 限制。以 `input=256, output=4096` 为例，每个请求约需要 `ceil((256+4096)/256)=17` 个 block；64 个 block 只能容纳约 3 个完整请求。因此 `con=16` 测得 137.23 tok/s、Mean TPOT 19.61ms、Mean TTFT 176.76s。该现象是安全 admission 导致的排队，服务保持 healthy。 + +不采用“只少预留几个 block”的激进策略，因为没有 KV swap/preemption 时，请求会在生成中途耗尽 block，风险从排队变成失败。 + +## 8B 测试矩阵运行结果 + +按照赛题给出的 8B 测试矩阵，本组共完成 30 个测试点，所有测试点均成功返回，服务保持 healthy。其中 `con=16/64, in=256, out=256/1024` 采用参数扫描阶段中与官方矩阵完全相同配置的成功结果，其余高并发点来自最终续跑目录。 + +| con | input | output | success | 输出 tok/s | 总 tok/s | Mean/P99 TTFT ms | Mean/P99 TPOT ms | Mean/P99 ITL ms | +|---:|---:|---:|---:|---:|---:|---:|---:|---:| +| 1 | 32 | 256 | 20/20 | 56.86 | 63.77 | 32.04 / 35.86 | 17.53 / 17.59 | 17.46 / 18.12 | +| 1 | 32 | 1024 | 20/20 | 56.42 | 58.13 | 33.57 / 37.35 | 17.71 / 17.78 | 17.70 / 18.34 | +| 1 | 32 | 4096 | 20/20 | 56.19 | 56.62 | 36.27 / 39.52 | 17.79 / 17.89 | 17.79 / 18.52 | +| 1 | 256 | 256 | 20/20 | 55.84 | 111.68 | 59.00 / 65.34 | 17.75 / 17.83 | 17.68 / 18.29 | +| 1 | 256 | 1024 | 20/20 | 56.34 | 70.40 | 62.17 / 70.27 | 17.71 / 18.07 | 17.68 / 18.26 | +| 1 | 256 | 4096 | 20/20 | 56.07 | 59.56 | 72.66 / 95.72 | 17.82 / 17.86 | 17.84 / 19.03 | +| 1 | 4096 | 256 | 20/20 | 49.95 | 852.13 | 530.22 / 571.01 | 18.01 / 18.08 | 17.94 / 18.94 | +| 1 | 4096 | 1024 | 20/20 | 53.96 | 269.84 | 540.27 / 581.55 | 18.02 / 18.11 | 18.01 / 19.12 | +| 1 | 4096 | 4096 | 20/20 | 54.40 | 108.78 | 552.19 / 587.24 | 18.25 / 18.30 | 18.25 / 19.30 | +| 4 | 32 | 256 | 40/40 | 219.68 | 246.37 | 68.45 / 83.37 | 18.01 / 18.30 | 17.94 / 19.76 | +| 4 | 32 | 1024 | 40/40 | 218.34 | 224.96 | 71.65 / 85.85 | 18.27 / 18.37 | 18.31 / 20.60 | +| 4 | 32 | 4096 | 40/40 | 147.13 | 148.25 | 26035.76 / 81451.53 | 19.49 / 19.92 | 19.49 / 21.63 | +| 4 | 256 | 256 | 40/40 | 211.18 | 422.36 | 149.97 / 188.61 | 18.42 / 18.83 | 18.35 / 19.56 | +| 4 | 256 | 1024 | 40/40 | 212.85 | 265.91 | 155.94 / 196.74 | 18.66 / 18.79 | 18.65 / 19.79 | +| 4 | 256 | 4096 | 40/40 | 142.62 | 151.50 | 26983.21 / 83198.15 | 20.12 / 20.30 | 20.11 / 21.85 | +| 4 | 4096 | 256 | 40/40 | 106.39 | 1814.25 | 2939.61 / 7078.04 | 24.76 / 25.73 | 24.66 / 24.14 | +| 4 | 4096 | 1024 | 40/40 | 53.89 | 269.49 | 54604.96 / 462816.24 | 18.09 / 18.24 | 18.07 / 18.90 | +| 4 | 4096 | 4096 | 40/40 | 54.46 | 111.78 | 203503.93 / 1000116.28 | 17.33 / 18.61 | 18.24 / 19.64 | +| 16 | 32 | 256 | 16/16 | 798.81 | 895.94 | 174.00 / 237.02 | 19.40 / 19.79 | 19.32 / 23.38 | +| 16 | 32 | 1024 | 16/16 | 414.56 | 427.12 | 5312.69 / 20810.77 | 19.69 / 20.24 | 19.68 / 23.08 | +| 16 | 32 | 4096 | 16/16 | 138.96 | 140.02 | 174643.55 / 386805.00 | 19.36 / 19.54 | 19.36 / 21.08 | +| 16 | 256 | 256 | 80/80 | 726.03 | 1445.77 | 468.25 / 648.21 | 20.26 / 21.79 | 20.20 / 22.70 | +| 16 | 256 | 1024 | 80/80 | 483.11 | 602.73 | 11469.24 / 21660.05 | 20.43 / 20.83 | 20.42 / 23.17 | +| 16 | 256 | 4096 | 16/16 | 137.23 | 145.77 | 176758.32 / 392478.99 | 19.61 / 19.88 | 19.61 / 21.67 | +| 64 | 32 | 256 | 64/64 | 1052.16 | 1180.10 | 4299.80 / 8402.76 | 29.31 / 30.72 | 29.20 / 47.44 | +| 64 | 32 | 1024 | 64/64 | 530.96 | 547.03 | 45873.64 / 104608.76 | 20.20 / 20.45 | 20.18 / 23.19 | +| 64 | 32 | 4096 | 64/64 | 153.98 | 155.97 | 682827.66 / 1000204.93 | 11.86 / 19.67 | 19.48 / 21.49 | +| 64 | 256 | 256 | 128/128 | 725.92 | 1449.00 | 10944.62 / 25269.96 | 24.19 / 25.64 | 24.09 / 27.34 | +| 64 | 256 | 1024 | 128/128 | 475.28 | 593.76 | 86502.94 / 207693.01 | 20.44 / 20.78 | 20.43 / 23.28 | +| 64 | 256 | 4096 | 64/64 | 152.06 | 168.37 | 686751.99 / 1000352.92 | 12.07 / 19.81 | 19.72 / 21.50 | + +结果上，短输出场景受益最明显：`con=64, in=32, out=256` 输出吞吐达到 1052.16 tok/s，`con=64, in=256, out=256` 达到 725.92 tok/s。长输出场景的 TPOT/ITL 仍保持在约 19-21ms，但 TTFT 会随 KV 容量排队显著升高，这是 24GB 单卡与 64 个 KV block 下的容量边界。 + +## 单并发基线与 decode 稳定性 + +为确认问题并非来自单请求 decode 性能，本组首先记录了 8B 模型在单并发下的基线表现。以下结果均使用 Flash Attention 与 Paged Attention，每个测试点 20 个请求。 + +| con=1 | input | output | 输出 tok/s | 总 tok/s | Mean TTFT ms | Mean/P99 TPOT ms | Mean/P99 ITL ms | +| ----: | ----: | -----: | -------: | ------: | -----------: | ----------------: | ---------------: | +| 1 | 32 | 256 | 56.86 | 63.77 | 32.04 | 17.53 / 17.59 | 17.46 / 18.12 | +| 1 | 32 | 1024 | 56.42 | 58.13 | 33.57 | 17.71 / 17.78 | 17.70 / 18.34 | +| 1 | 32 | 4096 | 56.19 | 56.62 | 36.27 | 17.79 / 17.89 | 17.79 / 18.52 | +| 1 | 256 | 256 | 55.84 | 111.68 | 59.00 | 17.75 / 17.83 | 17.68 / 18.29 | +| 1 | 256 | 1024 | 56.34 | 70.40 | 62.17 | 17.71 / 18.07 | 17.68 / 18.26 | +| 1 | 256 | 4096 | 56.07 | 59.56 | 72.66 | 17.82 / 17.86 | 17.84 / 19.03 | +| 1 | 4096 | 256 | 49.95 | 852.13 | 530.22 | 18.01 / 18.08 | 17.94 / 18.94 | +| 1 | 4096 | 1024 | 53.96 | 269.84 | 540.27 | 18.02 / 18.11 | 18.01 / 19.12 | +| 1 | 4096 | 4096 | 54.40 | 108.78 | 552.19 | 18.25 / 18.30 | 18.25 / 19.30 | + +结论：单流 decode 始终约 56 tok/s、TPOT 约 18ms；4096 输入主要增加 prefill/TTFT，而未破坏 decode 路径。这为后续将问题定位到“并发调度时的瞬时显存峰值”提供了对照。 + +## 并发基线与故障复现 + +| con=4 | input | output | 输出 tok/s | Mean TTFT ms | Mean/P99 TPOT ms | Mean/P99 ITL ms | 现象 | +|---:|---:|---:|---:|---:|---:|---:|---| +| 4 | 32 | 256 | 219.68 | 68.45 | 18.01 / 18.30 | 17.94 / 19.76 | 正常 | +| 4 | 32 | 1024 | 218.34 | 71.65 | 18.27 / 18.37 | 18.31 / 20.60 | 正常 | +| 4 | 32 | 4096 | 147.13 | 26035.76 | 19.49 / 19.92 | 19.49 / 21.63 | KV 排队 | +| 4 | 256 | 256 | 211.18 | 149.97 | 18.42 / 18.83 | 18.35 / 19.56 | 正常 | +| 4 | 256 | 1024 | 212.85 | 155.94 | 18.66 / 18.79 | 18.65 / 19.79 | 正常 | +| 4 | 256 | 4096 | 142.62 | 26983.21 | 20.12 / 20.30 | 20.11 / 21.85 | KV 排队 | +| 4 | 4096 | 256 | 106.39 | 2939.61 | 24.76 / 25.73 | 24.66 / 24.14 | 长输入干扰 decode | + +在进一步调测中，`con=4, in=4096, out=1024` 于原始实现下复现 `cudaMalloc failed`、`RankWorker stopped` 和服务 unhealthy；`con=16, in=256, out=256` 也复现相同类型的 failure。这说明仅按请求数准入并不足以覆盖长 prefill 或高并发 prefill/decode 交叠时的临时显存开销。 + +## 参数扫描：排除不可靠路径 + +| 尝试 | 结果 | 决策 | +|---|---|---| +| `num-blocks=96` / block 256 | 空载约 22587 MiB，首个长输入 forward OOM | 超出 24GB 余量，不采用 | +| `num-blocks=80` / block 256 | 长输入 forward OOM | 不采用 | +| `num-blocks=128` / block 128 | paged attention/graph 编译失败 | 不采用 | +| max batch=3 或 2 | 长输入仍 OOM/unhealthy | 仅全局降 batch 不足 | + +这组反例说明，问题不能通过单纯扩大 KV cache 或全局降低 batch size 稳定解决。最终方案需要区分 prefill 与 decode 阶段的显存压力，为两类阶段设置不同的安全上限。 + +## 优化迭代与消融 + +### 长输入保护的两次迭代 + +第一版策略仅限制同一轮 schedule 中的长输入请求数量，但新的长输入 prefill 仍可能与已经在运行队列中的 decode 请求交叠。在 `con=4, in=4096, out=1024` 场景下，该版本仍会触发服务 unhealthy。 + +第二版加入 decode-first 保护：当 running queue 非空时，新的长输入 prefill 会被延后调度，优先让已有 decode 请求继续推进。该策略避免了长输入 prefill 与 decode 阶段叠加造成的瞬时显存峰值，`con=4, in=4096, out=1024` 最终稳定完成 40/40 请求。 + +随后，本组将 long-context 判定从“总长度达到 4096”收窄为默认仅 `prompt_len >= 4096`。这样可以只保护真正的长输入请求，避免把短输入、长输出请求误判为长上下文并强制串行化。以 `con=4, in=32, out=4096` 为例，该调整使输出吞吐从 108.46 提升到 137.01 tok/s，Median TTFT 从 73100ms 降至 66ms，同时长输入卡点仍保持稳定。 + +长输入保护解决了服务稳定性问题，但高并发短输出场景仍会受到全局 batch 上限限制。为此，本组进一步将 prefill 与 decode 的 batch cap 拆分，使 prefill 保持保守批次以控制显存峰值，同时让 decode 使用更大的安全批次提升输出吞吐。 + +### 分相位 batch cap + +| 场景 | 策略 | 输出 tok/s | 总 tok/s | Mean/P99 TTFT ms | Mean/P99 TPOT ms | Mean/P99 ITL ms | +|---|---|---:|---:|---:|---:|---:| +| con=16, in=256, out=256 | max batch=8 | 407.30 | 814.65 | 437.84 / 619.32 | 37.70 / 39.20 | 37.54 / 39.61 | +| con=16, in=256, out=256 | prefill=8/decode=16 | 724.84 | 1449.68 | 468.86 / 648.64 | 20.30 / 21.82 | 20.23 / 22.91 | +| con=16, in=256, out=256 | prefill=8/decode=24 | 726.03 | 1445.77 | 468.25 / 648.21 | 20.26 / 21.79 | 20.20 / 22.70 | +| con=16, in=256, out=1024 | max batch=8 | 399.02 | 497.97 | 13727.04 / 51497.49 | 24.91 / 25.38 | 24.90 / 41.97 | +| con=16, in=256, out=1024 | prefill=8/decode=16 | 483.19 | 603.04 | 11460.90 / 21678.92 | 20.43 / 20.78 | 20.41 / 23.36 | +| con=16, in=256, out=1024 | prefill=8/decode=24 | 483.11 | 602.73 | 11469.24 / 21660.05 | 20.43 / 20.83 | 20.42 / 23.17 | +| con=64, in=256, out=256 | max batch=8 | 376.64 | 753.37 | 21498.23 / 69212.98 | 51.69 / 53.70 | 51.50 / 61.84 | +| con=64, in=256, out=256 | prefill=8/decode=16 | 631.06 | 1261.17 | 12516.38 / 31741.32 | 29.07 / 30.60 | 28.96 / 46.19 | +| con=64, in=256, out=256 | prefill=8/decode=24 | 725.92 | 1449.00 | 10944.62 / 25269.96 | 24.19 / 25.64 | 24.09 / 27.34 | + +从消融结果看，`decode=24` 对 con=16 场景收益有限，但在 con=64 短输出场景中，将输出吞吐从 631.06 提升到 725.92 tok/s，并将 Mean TPOT 从 29.07ms 降至 24.19ms、Mean ITL 从 28.96ms 降至 24.09ms。因此，本次 8B 提交采用 `prefill=8/decode=24` 作为最终配置。 + +## 调度策略说明 + +waiting queue: + - prefill 阶段最多调度 prefill_batch_cap 个请求 + - 若请求 prompt_len >= 4096，且已有 decode 请求正在运行，则延后该 prefill + - 同一 batch 中 long-context 请求数不超过 long_context_max_batch + +running queue: + - decode 阶段最多调度 decode_batch_cap 个请求 + - 同样遵守 long-context batch 上限 + +该策略将稳定性保护集中在真正的长输入 prefill 上，同时避免短输入、长输出请求被不必要地串行化。对于普通短/中输入请求，decode 阶段可以使用更大的 batch cap 提升吞吐；对于真正长输入请求，则避免其 prefill 与正在执行的 decode 阶段叠加到不安全的瞬时显存峰值。 + + +## 结果文件与审计 + +- 原始详细日志：`/root/autodl-tmp/bench-results-official-8b/*.log` +- vLLM 原始 JSON：`/root/autodl-tmp/bench-results-official-8b/*.json` +- 参数扫描结果：`/root/autodl-tmp/bench-results-param-scan/` +- PR 测试截图：`doc_extern/pr_screenshot.png` +- 实验全过程记录：`doc_extern/实验过程记录与卡点分析.md` + +最后 6 个官方矩阵点已由 `doc_extern/run_remaining_official_8b.sh` 顺序完成；所有点均有日志或 JSON 可追溯。 + +PR内容包括： + +- 修改文件：`python/infinilm/llm/scheduler.py` +- 赛题报告：`T2-1-2_伊里斯_赛题报告.md` + diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index a1153e91..0619cedf 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -3,6 +3,7 @@ """ import logging +import os import queue from typing import List, Optional @@ -82,6 +83,20 @@ def __init__( self.waiting_queue = janus.Queue() self.running_queue = janus.Queue() self.max_batch_size = max_batch_size + self.long_context_threshold = int( + os.getenv("INFINILM_LONG_CONTEXT_THRESHOLD", "4096") + ) + self.long_context_max_batch = int( + os.getenv("INFINILM_LONG_CONTEXT_MAX_BATCH", "1") + ) + self.long_context_use_total_len = ( + os.getenv("INFINILM_LONG_CONTEXT_USE_TOTAL_LEN", "0") == "1" + ) + self.decode_first_when_running = ( + os.getenv("INFINILM_DECODE_FIRST_WHEN_RUNNING", "0") == "1" + ) + self.prefill_batch_cap = int(os.getenv("INFINILM_PREFILL_BATCH_CAP", "0")) + self.decode_batch_cap = int(os.getenv("INFINILM_DECODE_BATCH_CAP", "0")) self.finished_receiving_kv_req_ids: set[str] = set() self.failed_receiving_kv_req_ids: set[str] = set() @@ -124,6 +139,48 @@ def _exceeds_token_budget( > self.max_num_batched_tokens ) + def _phase_batch_limit(self, cap: int) -> int: + if cap <= 0: + return self.max_batch_size + return min(self.max_batch_size, cap) + + def _is_long_context_request(self, req: InferenceRequest) -> bool: + if self.long_context_max_batch <= 0: + return False + prompt_len = req.get_prompt_length() + if prompt_len >= self.long_context_threshold: + return True + if not self.long_context_use_total_len: + return False + max_tokens = req.sampling_params.max_tokens or 0 + try: + total_len = req.get_total_length() + except Exception: + total_len = prompt_len + max_tokens + return ( + total_len >= self.long_context_threshold + or prompt_len + max_tokens >= self.long_context_threshold + ) + + def _exceeds_long_context_batch( + self, req: InferenceRequest, scheduled_requests: List[InferenceRequest] + ) -> bool: + if self.long_context_max_batch <= 0 or not scheduled_requests: + return False + + long_in_batch = sum( + 1 + for scheduled_req in scheduled_requests + if self._is_long_context_request(scheduled_req) + ) + req_is_long = self._is_long_context_request(req) + + if long_in_batch == 0: + return ( + req_is_long and len(scheduled_requests) >= self.long_context_max_batch + ) + return long_in_batch >= self.long_context_max_batch + def schedule(self) -> Optional[SchedulerOutput]: """Schedule and return batch of requests to execute.""" deferred_requests = [] @@ -131,10 +188,12 @@ def schedule(self) -> Optional[SchedulerOutput]: is_prefill = False current_num_batched_tokens = 0 current_prefill_extra_blocks = 0 + prefill_batch_limit = self._phase_batch_limit(self.prefill_batch_cap) + decode_batch_limit = self._phase_batch_limit(self.decode_batch_cap) # Process Waiting queue (prefill phase) while ( - len(scheduled_requests) < self.max_batch_size + len(scheduled_requests) < prefill_batch_limit and current_num_batched_tokens < self.max_num_batched_tokens ): try: @@ -146,6 +205,18 @@ def schedule(self) -> Optional[SchedulerOutput]: self.complete_requests([req]) continue + if ( + self.decode_first_when_running + and self._is_long_context_request(req) + and self.running_queue.sync_q.qsize() > 0 + ): + deferred_requests.append(req) + break + + if self._exceeds_long_context_batch(req, scheduled_requests): + deferred_requests.append(req) + break + req_tokens = req.get_input_tokens() if req.num_computed_tokens == 0: @@ -305,7 +376,7 @@ def schedule(self) -> Optional[SchedulerOutput]: return scheduler_output # Process Running queue (decode phase) - while len(scheduled_requests) < self.max_batch_size: + while len(scheduled_requests) < decode_batch_limit: try: req = self.running_queue.sync_q.get_nowait() except queue.Empty: @@ -315,6 +386,10 @@ def schedule(self) -> Optional[SchedulerOutput]: self.complete_requests([req]) continue + if self._exceeds_long_context_batch(req, scheduled_requests): + self.running_queue.sync_q.put(req) + break + # Decode phase: allocate slot for newly generated token try: req.block_table, new_slot = self.cache_manager.append_slot( @@ -344,7 +419,7 @@ def schedule(self) -> Optional[SchedulerOutput]: req.status = RequestStatus.WAITING self.waiting_queue.sync_q.put(req) elif req_id in self.finished_receiving_kv_req_ids: - if len(scheduled_requests) < self.max_batch_size: + if len(scheduled_requests) < decode_batch_limit: logger.info( f"Request {req_id[:8]}... finished receiving KV, scheduling for decode." )