From a6e21ff480f464aad75616399bd01a2f504a0956 Mon Sep 17 00:00:00 2001 From: DiFanrui <2829107824@qq.com> Date: Sun, 12 Jul 2026 11:02:19 +0800 Subject: [PATCH 1/3] =?UTF-8?q?[T2-1-2]=E4=BC=8A=E9=87=8C=E6=96=AF=5Fcommi?= =?UTF-8?q?t?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...33\351\242\230\346\212\245\345\221\212.md" | 237 ++++++++++++++++++ python/infinilm/llm/scheduler.py | 67 ++++- 2 files changed, 301 insertions(+), 3 deletions(-) create mode 100644 "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" 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..e5932974 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,12 @@ 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 +131,42 @@ 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 +174,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 +191,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 +362,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 +372,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 +405,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." ) From 612571f408943b20e2a9fbd0efec57c67fa58d76 Mon Sep 17 00:00:00 2001 From: DiFanrui <2829107824@qq.com> Date: Mon, 13 Jul 2026 01:15:19 +0800 Subject: [PATCH 2/3] format scheduler --- python/infinilm/llm/scheduler.py | 1277 +++++++++++++++--------------- 1 file changed, 642 insertions(+), 635 deletions(-) diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index e5932974..ba2cf6dd 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -1,642 +1,649 @@ -""" -Scheduler - Request scheduling and batch management with Paged Attention KV Cache. -""" - -import logging -import os -import queue -from typing import List, Optional - -import janus - -from infinilm.llm.cache_manager import BlockManager, MambaCacheManager -from infinilm.llm.request import InferenceRequest, RequestStatus - -logger = logging.getLogger(__name__) - - -class SpeculativeCacheOps: - """Limited cache operations needed by speculative verification.""" - - def __init__(self, cache_manager: BlockManager): - self._cache_manager = cache_manager - - def append_verify_slots( - self, - block_table: List[int], - start_length: int, - num_slots: int, - token_ids: List[int], - ): - return self._cache_manager.append_slots( - block_table, - start_length, - num_slots, - token_ids, - update_hash=False, - ) - - def rollback_to_length(self, block_table: List[int], keep_tokens: int): - return self._cache_manager.truncate_blocks(block_table, keep_tokens) - - def commit_accepted_tokens( - self, block_table: List[int], token_ids: List[int], num_tokens: int - ) -> None: - self._cache_manager.commit_blocks_hash(block_table, token_ids, num_tokens) - - -class SchedulerOutput: - """Scheduler output containing scheduled requests and execution phase info.""" - - def __init__( - self, - scheduled_requests: List[InferenceRequest], - is_prefill: bool = False, - speculative_cache_ops: Optional[SpeculativeCacheOps] = None, - ): - self.scheduled_requests = scheduled_requests - self.num_requests = len(scheduled_requests) - self.is_prefill = is_prefill - self.speculative_cache_ops = speculative_cache_ops - self.kv_connector_metadata = None - - -class Scheduler: - """Request scheduler with integrated BlockManager for KV cache management. - - Scheduling logic: - 1. Running queue: Check for new blocks needed, update slot_mapping - 2. Waiting queue: Try block reuse (prefix caching), allocate new blocks - 3. Reference counting: Free blocks when requests complete - """ - - def __init__( - self, - max_batch_size: int = 16, - num_blocks: int = 512, - block_size: int = 256, - max_num_batched_tokens: int = 1024, - connector=None, - has_mamba_cache: bool = False, - num_mamba_cache_blocks: int | None = None, - ): - 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() - self.pending_free_blocks: dict[str, list[int]] = {} - self.pending_kv_decode_blocks: int = 0 - self.remote_kv_requests: dict[str, InferenceRequest] = {} - - self.cache_manager = BlockManager(num_blocks=num_blocks, block_size=block_size) - self.has_mamba_cache = has_mamba_cache - self.mamba_cache_manager = ( - MambaCacheManager(num_mamba_cache_blocks or max(2, num_blocks // 4)) - if has_mamba_cache - else None - ) - self.speculative_cache_ops = SpeculativeCacheOps(self.cache_manager) - self.block_size = block_size - self.max_num_batched_tokens = max_num_batched_tokens - self.connector = connector - - def add_request(self, request: InferenceRequest): - if request is not None: - request.status = RequestStatus.WAITING - self.waiting_queue.sync_q.put(request) - - def _exceeds_token_budget( - self, - current_num_batched_tokens: int, - num_tokens_this_step: int, - num_scheduled_requests: int, - ) -> bool: - """Return True when adding this request should be deferred for token budget. - - A single request is always allowed to make progress, even if it is larger - than max_num_batched_tokens. - """ - if num_scheduled_requests == 0: - return False +""" +Scheduler - Request scheduling and batch management with Paged Attention KV Cache. +""" + +import logging +import os +import queue +from typing import List, Optional + +import janus + +from infinilm.llm.cache_manager import BlockManager, MambaCacheManager +from infinilm.llm.request import InferenceRequest, RequestStatus + +logger = logging.getLogger(__name__) + + +class SpeculativeCacheOps: + """Limited cache operations needed by speculative verification.""" + + def __init__(self, cache_manager: BlockManager): + self._cache_manager = cache_manager + + def append_verify_slots( + self, + block_table: List[int], + start_length: int, + num_slots: int, + token_ids: List[int], + ): + return self._cache_manager.append_slots( + block_table, + start_length, + num_slots, + token_ids, + update_hash=False, + ) + + def rollback_to_length(self, block_table: List[int], keep_tokens: int): + return self._cache_manager.truncate_blocks(block_table, keep_tokens) + + def commit_accepted_tokens( + self, block_table: List[int], token_ids: List[int], num_tokens: int + ) -> None: + self._cache_manager.commit_blocks_hash(block_table, token_ids, num_tokens) + + +class SchedulerOutput: + """Scheduler output containing scheduled requests and execution phase info.""" + + def __init__( + self, + scheduled_requests: List[InferenceRequest], + is_prefill: bool = False, + speculative_cache_ops: Optional[SpeculativeCacheOps] = None, + ): + self.scheduled_requests = scheduled_requests + self.num_requests = len(scheduled_requests) + self.is_prefill = is_prefill + self.speculative_cache_ops = speculative_cache_ops + self.kv_connector_metadata = None + + +class Scheduler: + """Request scheduler with integrated BlockManager for KV cache management. + + Scheduling logic: + 1. Running queue: Check for new blocks needed, update slot_mapping + 2. Waiting queue: Try block reuse (prefix caching), allocate new blocks + 3. Reference counting: Free blocks when requests complete + """ + + def __init__( + self, + max_batch_size: int = 16, + num_blocks: int = 512, + block_size: int = 256, + max_num_batched_tokens: int = 1024, + connector=None, + has_mamba_cache: bool = False, + num_mamba_cache_blocks: int | None = None, + ): + 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() + self.pending_free_blocks: dict[str, list[int]] = {} + self.pending_kv_decode_blocks: int = 0 + self.remote_kv_requests: dict[str, InferenceRequest] = {} + + self.cache_manager = BlockManager(num_blocks=num_blocks, block_size=block_size) + self.has_mamba_cache = has_mamba_cache + self.mamba_cache_manager = ( + MambaCacheManager(num_mamba_cache_blocks or max(2, num_blocks // 4)) + if has_mamba_cache + else None + ) + self.speculative_cache_ops = SpeculativeCacheOps(self.cache_manager) + self.block_size = block_size + self.max_num_batched_tokens = max_num_batched_tokens + self.connector = connector + + def add_request(self, request: InferenceRequest): + if request is not None: + request.status = RequestStatus.WAITING + self.waiting_queue.sync_q.put(request) + + def _exceeds_token_budget( + self, + current_num_batched_tokens: int, + num_tokens_this_step: int, + num_scheduled_requests: int, + ) -> bool: + """Return True when adding this request should be deferred for token budget. + + A single request is always allowed to make progress, even if it is larger + than max_num_batched_tokens. + """ + if num_scheduled_requests == 0: + return False + return ( + current_num_batched_tokens + num_tokens_this_step + > 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 ( - current_num_batched_tokens + num_tokens_this_step - > self.max_num_batched_tokens + total_len >= self.long_context_threshold + or prompt_len + max_tokens >= self.long_context_threshold ) - - 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 - + + 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 + 1 + for scheduled_req in scheduled_requests if self._is_long_context_request(scheduled_req) ) - req_is_long = self._is_long_context_request(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 = [] - scheduled_requests = [] - 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) < prefill_batch_limit - and current_num_batched_tokens < self.max_num_batched_tokens - ): - try: - req = self.waiting_queue.sync_q.get_nowait() - except queue.Empty: - break - # Skip requests that were already finished (e.g., timed out/canceled while waiting) - if req.is_finished(): - 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: - if self.has_mamba_cache: - cached_block_table = [] - num_local_computed_tokens = 0 - blocks_blueprint = [] - load_kv_async = False - num_external_computed_tokens = 0 - else: - ( - cached_block_table, - num_local_computed_tokens, - blocks_blueprint, - ) = self.cache_manager.get_computed_blocks( - req_tokens, req.get_mm_token_index_mappings() - ) - if self.connector is not None: - ext_tokens, load_kv_async = ( - self.connector.get_num_new_matched_tokens( - req, num_local_computed_tokens - ) - ) - num_external_computed_tokens = ext_tokens - else: - load_kv_async = False - num_external_computed_tokens = 0 - - num_computed_tokens = ( - num_local_computed_tokens + num_external_computed_tokens - ) - if load_kv_async: - num_computed_tokens -= 1 - num_new_tokens = req.get_prompt_length() - num_computed_tokens - - # Early token budget check: skip can_accept_request and allocate_slots - # for requests that would exceed the per-schedule token budget. - if not load_kv_async: - num_tokens_this_step = ( - req.get_prompt_length() - num_local_computed_tokens - ) - if self._exceeds_token_budget( - current_num_batched_tokens, - num_tokens_this_step, - len(scheduled_requests), - ): - if num_local_computed_tokens > 0: - self.cache_manager.free_blocks(cached_block_table) - deferred_requests.append(req) - break - - if not self.can_accept_request( - req, - num_local_computed_tokens, - current_prefill_extra_blocks, - ): - logger.warning( - "Insufficient KV cache blocks for request %s, deferring.", - req.request_id, - ) - - if num_local_computed_tokens > 0: - self.cache_manager.free_blocks(cached_block_table) - deferred_requests.append(req) - break - - req_blocks, slot_mapping = self.cache_manager.allocate_slots( - req_tokens, - num_new_tokens, - num_computed_tokens=num_computed_tokens, - cached_block_table=cached_block_table, - blocks_blueprint=blocks_blueprint, - delay_cache_blocks=load_kv_async, - ) - - if req_blocks is None: - logger.warning( - "Failed to allocate KV cache blocks for request: %s", - req.request_id, - ) - if num_local_computed_tokens > 0: - self.cache_manager.free_blocks(cached_block_table) - deferred_requests.append(req) - break - - if self.has_mamba_cache and req.mamba_cache_index is None: - req.mamba_cache_index = self.mamba_cache_manager.allocate() - if req.mamba_cache_index is None: - self.cache_manager.free_blocks(req_blocks) - logger.warning( - "Insufficient mamba cache rows for request %s, deferring.", - req.request_id, - ) - deferred_requests.append(req) - break - - req.block_table = req_blocks - req.slot_mapping = slot_mapping - req.num_blocks = len(req_blocks) - req.num_local_cached_tokens = num_local_computed_tokens - req.num_computed_tokens = num_computed_tokens - - if self.connector is not None: - self.connector.update_state_after_alloc( - req, - req.block_table, - num_external_computed_tokens, - self.block_size, - ) - else: - load_kv_async = False - num_tokens_this_step = ( - req.get_prompt_length() - req.num_local_cached_tokens - ) - if self._exceeds_token_budget( - current_num_batched_tokens, - num_tokens_this_step, - len(scheduled_requests), - ): - deferred_requests.append(req) - break - self.cache_manager.update_blocks_hash( - req.block_table, req.num_local_cached_tokens - ) - - if load_kv_async: - req.status = RequestStatus.WAITING_FOR_REMOTE_KVS - self.remote_kv_requests[req.request_id] = req - self.pending_kv_decode_blocks += ( - req.sampling_params.max_tokens + self.block_size - 1 - ) // self.block_size - continue - - current_prefill_extra_blocks += self._get_prefill_extra_blocks(req) - scheduled_requests.append(req) - - num_tokens_this_step = req.get_prompt_length() - req.num_local_cached_tokens - current_num_batched_tokens += num_tokens_this_step - - req.status = RequestStatus.RUNNING - - if deferred_requests: - for req in deferred_requests: - self.waiting_queue.sync_q.put(req) - - # Return prefill batch if any waiting requests were scheduled - if scheduled_requests: - is_prefill = True - scheduler_output = SchedulerOutput( - scheduled_requests=scheduled_requests, - is_prefill=is_prefill, - speculative_cache_ops=self.speculative_cache_ops, - ) - if self.connector is not None: - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output - - # Process Running queue (decode phase) - while len(scheduled_requests) < decode_batch_limit: - try: - req = self.running_queue.sync_q.get_nowait() - except queue.Empty: - break - # Skip requests that were already finished (e.g., timed out/canceled while running) - if req.is_finished(): - 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( - req.block_table, req.get_total_length(), req.get_all_token_ids() - ) - req.slot_mapping = [new_slot] - req.num_blocks = len(req.block_table) - req.num_local_cached_tokens = req.get_total_length() - 1 - scheduled_requests.append(req) - - except RuntimeError as e: - raise RuntimeError("No available cache blocks for new token") from e - - # Promote completed remote KV transfers (lower priority than running queue). - # Cleanup (is_finished, failed re-queue) runs unconditionally; batch append only if slots remain. - if self.connector is not None and self.remote_kv_requests: - for req_id in list(self.remote_kv_requests.keys()): - req = self.remote_kv_requests[req_id] - if req.is_finished(): - self.complete_requests([req]) - continue - if req_id in self.failed_receiving_kv_req_ids: - logger.warning( - f"Request {req_id[:8]}... failed receiving KV, re-queuing for prefill." - ) - self.update_waiting_for_remote_kv(req) - 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) < decode_batch_limit: - logger.info( - f"Request {req_id[:8]}... finished receiving KV, scheduling for decode." - ) - self.update_waiting_for_remote_kv(req) - req.status = RequestStatus.RUNNING - scheduled_requests.append(req) - else: - break # Defer promotion to next schedule() if batch is full - - # Return decode batch if any running requests were scheduled - if scheduled_requests: - is_prefill = False - scheduler_output = SchedulerOutput( - scheduled_requests=scheduled_requests, - is_prefill=is_prefill, - speculative_cache_ops=self.speculative_cache_ops, - ) - - if self.connector is not None: - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output - - if self.connector is not None: - scheduler_output = SchedulerOutput( - scheduled_requests=[], - speculative_cache_ops=self.speculative_cache_ops, - ) - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output - - return None - - def update_waiting_for_remote_kv(self, request: InferenceRequest): - self.remote_kv_requests.pop(request.request_id, None) - self.pending_kv_decode_blocks -= ( - request.sampling_params.max_tokens + self.block_size - 1 - ) // self.block_size - if request.request_id in self.failed_receiving_kv_req_ids: - if request.num_computed_tokens: - valid_block_count = request.num_computed_tokens // self.block_size - self.cache_manager.update_blocks_hash( - request.block_table[:valid_block_count], - request.num_local_cached_tokens, - ) - request.slot_mapping = self.cache_manager.update_blocks_slot( - request.block_table, - request.num_computed_tokens, - request.get_prompt_length(), - ) - request.num_local_cached_tokens = request.num_computed_tokens - else: - self.cache_manager.free_blocks(request.block_table) - request.block_table = [] - request.slot_mapping = [] - request.num_local_cached_tokens = 0 - self.failed_receiving_kv_req_ids.discard(request.request_id) - else: - self.cache_manager.update_blocks_hash( - request.block_table, request.num_local_cached_tokens - ) - request.num_local_cached_tokens = request.num_computed_tokens - self.finished_receiving_kv_req_ids.discard(request.request_id) - - def complete_requests(self, requests: List[InferenceRequest]): - """Handle completed requests and free their blocks.""" - for req in requests: - if req.status in [ - RequestStatus.FINISHED, - RequestStatus.CANCELED, - RequestStatus.FAILED, - RequestStatus.TIMEOUT, - ]: - delay_free_blocks = False - if self.connector is not None: - delay_free_blocks, _ = self.connector.request_finished( - req, req.block_table, self.block_size - ) - - if req.request_id in self.remote_kv_requests: - self.pending_kv_decode_blocks -= ( - req.sampling_params.max_tokens + self.block_size - 1 - ) // self.block_size - self.remote_kv_requests.pop(req.request_id, None) - if req.request_id in self.finished_receiving_kv_req_ids: - self.finished_receiving_kv_req_ids.discard(req.request_id) - self.failed_receiving_kv_req_ids.discard(req.request_id) - else: - delay_free_blocks = True - if req.block_table and not delay_free_blocks: - self.cache_manager.free_blocks(req.block_table) - elif req.block_table and delay_free_blocks: - self.pending_free_blocks[req.request_id] = list(req.block_table) - if self.mamba_cache_manager is not None: - self.mamba_cache_manager.free(req.mamba_cache_index) - req.mamba_cache_index = None - - if req.status == RequestStatus.CANCELED: - logger.info( - f"Request {req.request_id[:8]}... canceled: {req.finish_reason}" - ) - elif req.status == RequestStatus.FAILED: - logger.error( - f"Request {req.request_id[:8]}... failed: {req.finish_reason}" - ) - elif req.status == RequestStatus.TIMEOUT: - logger.error( - f"Request {req.request_id[:8]}... timed out: {req.finish_reason}" - ) - else: - # Still running, put back in running queue - self.running_queue.sync_q.put(req) - - def can_accept_request( - self, - request: InferenceRequest, - num_local_computed_tokens: int, - current_prefill_extra_blocks: int = 0, - ) -> bool: - if ( - self.mamba_cache_manager is not None - and request.mamba_cache_index is None - and not self.mamba_cache_manager.can_allocate() - ): - return False - - total_required_blocks = 0 - - # Calculate blocks needed for running requests - running_queue_size = self.running_queue.sync_q.qsize() - for _ in range(running_queue_size): - req = self.running_queue.sync_q.get() - remaining_tokens = ( - req.sampling_params.max_tokens - req.get_num_generated_tokens() - ) - num_blocks_needed = ( - remaining_tokens + self.block_size - 1 - ) // self.block_size - total_required_blocks += num_blocks_needed - self.running_queue.sync_q.put(req) - - # Calculate blocks needed for the new request - total_length = request.get_prompt_length() - num_local_computed_tokens - total_length += request.sampling_params.max_tokens - num_blocks_needed = (total_length + self.block_size - 1) // self.block_size - total_required_blocks += num_blocks_needed - - # Include decode headroom for WAITING_FOR_REMOTE_KVS requests, which - # hold prompt blocks but will also need decode blocks once promoted. - total_required_blocks += self.pending_kv_decode_blocks - - # Include decode headroom for requests accepted earlier in this batch. - total_required_blocks += current_prefill_extra_blocks - - # Compare with total usable blocks in cache manager - return total_required_blocks <= self.cache_manager.get_total_usable_blocks() - - def _get_prefill_extra_blocks(self, request: InferenceRequest) -> int: - total_length = request.get_prompt_length() - total_length += request.sampling_params.max_tokens - total_required_blocks = (total_length + self.block_size - 1) // self.block_size - return max(total_required_blocks - len(request.block_table), 0) - - def update_from_output(self, model_output): - if self.connector is None or model_output.kv_connector_output is None: - return - - finished_recving_req_ids = ( - getattr(model_output.kv_connector_output, "finished_recving", None) or [] - ) - finished_sending_req_ids = ( - getattr(model_output.kv_connector_output, "finished_sending", None) or [] - ) - failed_recving_req_ids = ( - getattr(model_output.kv_connector_output, "failed_recving", None) or [] - ) - invalid_block_ids = ( - getattr(model_output.kv_connector_output, "invalid_block_ids", None) or [] - ) - - for req_id in finished_recving_req_ids: - if req_id in self.pending_free_blocks: - # Aborted request: transfer complete, now safe to free blocks. - self.cache_manager.free_blocks(self.pending_free_blocks.pop(req_id)) - elif req_id in self.remote_kv_requests: - # Active request: mark ready for promotion in schedule(). - self.finished_receiving_kv_req_ids.add(req_id) - # else: already processed or unknown, discard to avoid stale entries. - for req_id in finished_sending_req_ids: - self.cache_manager.free_blocks(self.pending_free_blocks.pop(req_id, [])) - for req_id in failed_recving_req_ids: - # Only track failures for active (non-aborted) requests; aborted - # requests are handled via pending_free_blocks in finished_recving. - if req_id in self.remote_kv_requests: - self.failed_receiving_kv_req_ids.add(req_id) - - if invalid_block_ids: - invalid_set = set(invalid_block_ids) - - for req in self.remote_kv_requests.values(): - start_block_idx = req.num_local_cached_tokens // self.block_size - for i, block_id in enumerate( - req.block_table[start_block_idx:], start=start_block_idx - ): - if block_id in invalid_set: - req.num_computed_tokens = i * self.block_size - break - elif self.failed_receiving_kv_req_ids: - for req_id in self.failed_receiving_kv_req_ids: - req = self.remote_kv_requests[req_id] - req.num_computed_tokens = req.num_local_cached_tokens - - def get_cache_stats(self) -> dict: - """Get cache statistics.""" - stats = { - "num_blocks": self.cache_manager.num_blocks, - "block_size": self.cache_manager.block_size, - "num_free_blocks": self.cache_manager.get_num_free_blocks(), - "usable_blocks": self.cache_manager.get_total_usable_blocks(), - "num_pending_blocks": len(self.cache_manager.pending_block_ids), - "num_used_blocks": len(self.cache_manager.used_block_ids), - } - if self.mamba_cache_manager is not None: - stats.update( - { - "num_mamba_cache_blocks": self.mamba_cache_manager.num_blocks, - "num_free_mamba_cache_blocks": self.mamba_cache_manager.get_num_free_blocks(), - "num_used_mamba_cache_blocks": len( - self.mamba_cache_manager.used_block_ids - ), - } + return ( + req_is_long + and len(scheduled_requests) >= self.long_context_max_batch ) - return stats + return long_in_batch >= self.long_context_max_batch + + def schedule(self) -> Optional[SchedulerOutput]: + """Schedule and return batch of requests to execute.""" + deferred_requests = [] + scheduled_requests = [] + 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) < prefill_batch_limit + and current_num_batched_tokens < self.max_num_batched_tokens + ): + try: + req = self.waiting_queue.sync_q.get_nowait() + except queue.Empty: + break + # Skip requests that were already finished (e.g., timed out/canceled while waiting) + if req.is_finished(): + 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: + if self.has_mamba_cache: + cached_block_table = [] + num_local_computed_tokens = 0 + blocks_blueprint = [] + load_kv_async = False + num_external_computed_tokens = 0 + else: + ( + cached_block_table, + num_local_computed_tokens, + blocks_blueprint, + ) = self.cache_manager.get_computed_blocks( + req_tokens, req.get_mm_token_index_mappings() + ) + if self.connector is not None: + ext_tokens, load_kv_async = ( + self.connector.get_num_new_matched_tokens( + req, num_local_computed_tokens + ) + ) + num_external_computed_tokens = ext_tokens + else: + load_kv_async = False + num_external_computed_tokens = 0 + + num_computed_tokens = ( + num_local_computed_tokens + num_external_computed_tokens + ) + if load_kv_async: + num_computed_tokens -= 1 + num_new_tokens = req.get_prompt_length() - num_computed_tokens + + # Early token budget check: skip can_accept_request and allocate_slots + # for requests that would exceed the per-schedule token budget. + if not load_kv_async: + num_tokens_this_step = ( + req.get_prompt_length() - num_local_computed_tokens + ) + if self._exceeds_token_budget( + current_num_batched_tokens, + num_tokens_this_step, + len(scheduled_requests), + ): + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + + if not self.can_accept_request( + req, + num_local_computed_tokens, + current_prefill_extra_blocks, + ): + logger.warning( + "Insufficient KV cache blocks for request %s, deferring.", + req.request_id, + ) + + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + + req_blocks, slot_mapping = self.cache_manager.allocate_slots( + req_tokens, + num_new_tokens, + num_computed_tokens=num_computed_tokens, + cached_block_table=cached_block_table, + blocks_blueprint=blocks_blueprint, + delay_cache_blocks=load_kv_async, + ) + + if req_blocks is None: + logger.warning( + "Failed to allocate KV cache blocks for request: %s", + req.request_id, + ) + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + + if self.has_mamba_cache and req.mamba_cache_index is None: + req.mamba_cache_index = self.mamba_cache_manager.allocate() + if req.mamba_cache_index is None: + self.cache_manager.free_blocks(req_blocks) + logger.warning( + "Insufficient mamba cache rows for request %s, deferring.", + req.request_id, + ) + deferred_requests.append(req) + break + + req.block_table = req_blocks + req.slot_mapping = slot_mapping + req.num_blocks = len(req_blocks) + req.num_local_cached_tokens = num_local_computed_tokens + req.num_computed_tokens = num_computed_tokens + + if self.connector is not None: + self.connector.update_state_after_alloc( + req, + req.block_table, + num_external_computed_tokens, + self.block_size, + ) + else: + load_kv_async = False + num_tokens_this_step = ( + req.get_prompt_length() - req.num_local_cached_tokens + ) + if self._exceeds_token_budget( + current_num_batched_tokens, + num_tokens_this_step, + len(scheduled_requests), + ): + deferred_requests.append(req) + break + self.cache_manager.update_blocks_hash( + req.block_table, req.num_local_cached_tokens + ) + + if load_kv_async: + req.status = RequestStatus.WAITING_FOR_REMOTE_KVS + self.remote_kv_requests[req.request_id] = req + self.pending_kv_decode_blocks += ( + req.sampling_params.max_tokens + self.block_size - 1 + ) // self.block_size + continue + + current_prefill_extra_blocks += self._get_prefill_extra_blocks(req) + scheduled_requests.append(req) + + num_tokens_this_step = req.get_prompt_length() - req.num_local_cached_tokens + current_num_batched_tokens += num_tokens_this_step + + req.status = RequestStatus.RUNNING + + if deferred_requests: + for req in deferred_requests: + self.waiting_queue.sync_q.put(req) + + # Return prefill batch if any waiting requests were scheduled + if scheduled_requests: + is_prefill = True + scheduler_output = SchedulerOutput( + scheduled_requests=scheduled_requests, + is_prefill=is_prefill, + speculative_cache_ops=self.speculative_cache_ops, + ) + if self.connector is not None: + meta = self.connector.build_connector_meta() + scheduler_output.kv_connector_metadata = meta + return scheduler_output + + # Process Running queue (decode phase) + while len(scheduled_requests) < decode_batch_limit: + try: + req = self.running_queue.sync_q.get_nowait() + except queue.Empty: + break + # Skip requests that were already finished (e.g., timed out/canceled while running) + if req.is_finished(): + 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( + req.block_table, req.get_total_length(), req.get_all_token_ids() + ) + req.slot_mapping = [new_slot] + req.num_blocks = len(req.block_table) + req.num_local_cached_tokens = req.get_total_length() - 1 + scheduled_requests.append(req) + + except RuntimeError as e: + raise RuntimeError("No available cache blocks for new token") from e + + # Promote completed remote KV transfers (lower priority than running queue). + # Cleanup (is_finished, failed re-queue) runs unconditionally; batch append only if slots remain. + if self.connector is not None and self.remote_kv_requests: + for req_id in list(self.remote_kv_requests.keys()): + req = self.remote_kv_requests[req_id] + if req.is_finished(): + self.complete_requests([req]) + continue + if req_id in self.failed_receiving_kv_req_ids: + logger.warning( + f"Request {req_id[:8]}... failed receiving KV, re-queuing for prefill." + ) + self.update_waiting_for_remote_kv(req) + 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) < decode_batch_limit: + logger.info( + f"Request {req_id[:8]}... finished receiving KV, scheduling for decode." + ) + self.update_waiting_for_remote_kv(req) + req.status = RequestStatus.RUNNING + scheduled_requests.append(req) + else: + break # Defer promotion to next schedule() if batch is full + + # Return decode batch if any running requests were scheduled + if scheduled_requests: + is_prefill = False + scheduler_output = SchedulerOutput( + scheduled_requests=scheduled_requests, + is_prefill=is_prefill, + speculative_cache_ops=self.speculative_cache_ops, + ) + + if self.connector is not None: + meta = self.connector.build_connector_meta() + scheduler_output.kv_connector_metadata = meta + return scheduler_output + + if self.connector is not None: + scheduler_output = SchedulerOutput( + scheduled_requests=[], + speculative_cache_ops=self.speculative_cache_ops, + ) + meta = self.connector.build_connector_meta() + scheduler_output.kv_connector_metadata = meta + return scheduler_output + + return None + + def update_waiting_for_remote_kv(self, request: InferenceRequest): + self.remote_kv_requests.pop(request.request_id, None) + self.pending_kv_decode_blocks -= ( + request.sampling_params.max_tokens + self.block_size - 1 + ) // self.block_size + if request.request_id in self.failed_receiving_kv_req_ids: + if request.num_computed_tokens: + valid_block_count = request.num_computed_tokens // self.block_size + self.cache_manager.update_blocks_hash( + request.block_table[:valid_block_count], + request.num_local_cached_tokens, + ) + request.slot_mapping = self.cache_manager.update_blocks_slot( + request.block_table, + request.num_computed_tokens, + request.get_prompt_length(), + ) + request.num_local_cached_tokens = request.num_computed_tokens + else: + self.cache_manager.free_blocks(request.block_table) + request.block_table = [] + request.slot_mapping = [] + request.num_local_cached_tokens = 0 + self.failed_receiving_kv_req_ids.discard(request.request_id) + else: + self.cache_manager.update_blocks_hash( + request.block_table, request.num_local_cached_tokens + ) + request.num_local_cached_tokens = request.num_computed_tokens + self.finished_receiving_kv_req_ids.discard(request.request_id) + + def complete_requests(self, requests: List[InferenceRequest]): + """Handle completed requests and free their blocks.""" + for req in requests: + if req.status in [ + RequestStatus.FINISHED, + RequestStatus.CANCELED, + RequestStatus.FAILED, + RequestStatus.TIMEOUT, + ]: + delay_free_blocks = False + if self.connector is not None: + delay_free_blocks, _ = self.connector.request_finished( + req, req.block_table, self.block_size + ) + + if req.request_id in self.remote_kv_requests: + self.pending_kv_decode_blocks -= ( + req.sampling_params.max_tokens + self.block_size - 1 + ) // self.block_size + self.remote_kv_requests.pop(req.request_id, None) + if req.request_id in self.finished_receiving_kv_req_ids: + self.finished_receiving_kv_req_ids.discard(req.request_id) + self.failed_receiving_kv_req_ids.discard(req.request_id) + else: + delay_free_blocks = True + if req.block_table and not delay_free_blocks: + self.cache_manager.free_blocks(req.block_table) + elif req.block_table and delay_free_blocks: + self.pending_free_blocks[req.request_id] = list(req.block_table) + if self.mamba_cache_manager is not None: + self.mamba_cache_manager.free(req.mamba_cache_index) + req.mamba_cache_index = None + + if req.status == RequestStatus.CANCELED: + logger.info( + f"Request {req.request_id[:8]}... canceled: {req.finish_reason}" + ) + elif req.status == RequestStatus.FAILED: + logger.error( + f"Request {req.request_id[:8]}... failed: {req.finish_reason}" + ) + elif req.status == RequestStatus.TIMEOUT: + logger.error( + f"Request {req.request_id[:8]}... timed out: {req.finish_reason}" + ) + else: + # Still running, put back in running queue + self.running_queue.sync_q.put(req) + + def can_accept_request( + self, + request: InferenceRequest, + num_local_computed_tokens: int, + current_prefill_extra_blocks: int = 0, + ) -> bool: + if ( + self.mamba_cache_manager is not None + and request.mamba_cache_index is None + and not self.mamba_cache_manager.can_allocate() + ): + return False + + total_required_blocks = 0 + + # Calculate blocks needed for running requests + running_queue_size = self.running_queue.sync_q.qsize() + for _ in range(running_queue_size): + req = self.running_queue.sync_q.get() + remaining_tokens = ( + req.sampling_params.max_tokens - req.get_num_generated_tokens() + ) + num_blocks_needed = ( + remaining_tokens + self.block_size - 1 + ) // self.block_size + total_required_blocks += num_blocks_needed + self.running_queue.sync_q.put(req) + + # Calculate blocks needed for the new request + total_length = request.get_prompt_length() - num_local_computed_tokens + total_length += request.sampling_params.max_tokens + num_blocks_needed = (total_length + self.block_size - 1) // self.block_size + total_required_blocks += num_blocks_needed + + # Include decode headroom for WAITING_FOR_REMOTE_KVS requests, which + # hold prompt blocks but will also need decode blocks once promoted. + total_required_blocks += self.pending_kv_decode_blocks + + # Include decode headroom for requests accepted earlier in this batch. + total_required_blocks += current_prefill_extra_blocks + + # Compare with total usable blocks in cache manager + return total_required_blocks <= self.cache_manager.get_total_usable_blocks() + + def _get_prefill_extra_blocks(self, request: InferenceRequest) -> int: + total_length = request.get_prompt_length() + total_length += request.sampling_params.max_tokens + total_required_blocks = (total_length + self.block_size - 1) // self.block_size + return max(total_required_blocks - len(request.block_table), 0) + + def update_from_output(self, model_output): + if self.connector is None or model_output.kv_connector_output is None: + return + + finished_recving_req_ids = ( + getattr(model_output.kv_connector_output, "finished_recving", None) or [] + ) + finished_sending_req_ids = ( + getattr(model_output.kv_connector_output, "finished_sending", None) or [] + ) + failed_recving_req_ids = ( + getattr(model_output.kv_connector_output, "failed_recving", None) or [] + ) + invalid_block_ids = ( + getattr(model_output.kv_connector_output, "invalid_block_ids", None) or [] + ) + + for req_id in finished_recving_req_ids: + if req_id in self.pending_free_blocks: + # Aborted request: transfer complete, now safe to free blocks. + self.cache_manager.free_blocks(self.pending_free_blocks.pop(req_id)) + elif req_id in self.remote_kv_requests: + # Active request: mark ready for promotion in schedule(). + self.finished_receiving_kv_req_ids.add(req_id) + # else: already processed or unknown, discard to avoid stale entries. + for req_id in finished_sending_req_ids: + self.cache_manager.free_blocks(self.pending_free_blocks.pop(req_id, [])) + for req_id in failed_recving_req_ids: + # Only track failures for active (non-aborted) requests; aborted + # requests are handled via pending_free_blocks in finished_recving. + if req_id in self.remote_kv_requests: + self.failed_receiving_kv_req_ids.add(req_id) + + if invalid_block_ids: + invalid_set = set(invalid_block_ids) + + for req in self.remote_kv_requests.values(): + start_block_idx = req.num_local_cached_tokens // self.block_size + for i, block_id in enumerate( + req.block_table[start_block_idx:], start=start_block_idx + ): + if block_id in invalid_set: + req.num_computed_tokens = i * self.block_size + break + elif self.failed_receiving_kv_req_ids: + for req_id in self.failed_receiving_kv_req_ids: + req = self.remote_kv_requests[req_id] + req.num_computed_tokens = req.num_local_cached_tokens + + def get_cache_stats(self) -> dict: + """Get cache statistics.""" + stats = { + "num_blocks": self.cache_manager.num_blocks, + "block_size": self.cache_manager.block_size, + "num_free_blocks": self.cache_manager.get_num_free_blocks(), + "usable_blocks": self.cache_manager.get_total_usable_blocks(), + "num_pending_blocks": len(self.cache_manager.pending_block_ids), + "num_used_blocks": len(self.cache_manager.used_block_ids), + } + if self.mamba_cache_manager is not None: + stats.update( + { + "num_mamba_cache_blocks": self.mamba_cache_manager.num_blocks, + "num_free_mamba_cache_blocks": self.mamba_cache_manager.get_num_free_blocks(), + "num_used_mamba_cache_blocks": len( + self.mamba_cache_manager.used_block_ids + ), + } + ) + return stats From 072bafa073778acae64fafc9c19e59593eb3a62b Mon Sep 17 00:00:00 2001 From: DiFanrui <2829107824@qq.com> Date: Mon, 13 Jul 2026 01:16:35 +0800 Subject: [PATCH 3/3] run ruff format on scheduler --- python/infinilm/llm/scheduler.py | 1281 +++++++++++++++--------------- 1 file changed, 644 insertions(+), 637 deletions(-) diff --git a/python/infinilm/llm/scheduler.py b/python/infinilm/llm/scheduler.py index ba2cf6dd..0619cedf 100644 --- a/python/infinilm/llm/scheduler.py +++ b/python/infinilm/llm/scheduler.py @@ -1,649 +1,656 @@ -""" -Scheduler - Request scheduling and batch management with Paged Attention KV Cache. -""" - -import logging -import os -import queue -from typing import List, Optional - -import janus - -from infinilm.llm.cache_manager import BlockManager, MambaCacheManager -from infinilm.llm.request import InferenceRequest, RequestStatus - -logger = logging.getLogger(__name__) - - -class SpeculativeCacheOps: - """Limited cache operations needed by speculative verification.""" - - def __init__(self, cache_manager: BlockManager): - self._cache_manager = cache_manager - - def append_verify_slots( - self, - block_table: List[int], - start_length: int, - num_slots: int, - token_ids: List[int], - ): - return self._cache_manager.append_slots( - block_table, - start_length, - num_slots, - token_ids, - update_hash=False, - ) - - def rollback_to_length(self, block_table: List[int], keep_tokens: int): - return self._cache_manager.truncate_blocks(block_table, keep_tokens) - - def commit_accepted_tokens( - self, block_table: List[int], token_ids: List[int], num_tokens: int - ) -> None: - self._cache_manager.commit_blocks_hash(block_table, token_ids, num_tokens) - - -class SchedulerOutput: - """Scheduler output containing scheduled requests and execution phase info.""" - - def __init__( - self, - scheduled_requests: List[InferenceRequest], - is_prefill: bool = False, - speculative_cache_ops: Optional[SpeculativeCacheOps] = None, - ): - self.scheduled_requests = scheduled_requests - self.num_requests = len(scheduled_requests) - self.is_prefill = is_prefill - self.speculative_cache_ops = speculative_cache_ops - self.kv_connector_metadata = None - - -class Scheduler: - """Request scheduler with integrated BlockManager for KV cache management. - - Scheduling logic: - 1. Running queue: Check for new blocks needed, update slot_mapping - 2. Waiting queue: Try block reuse (prefix caching), allocate new blocks - 3. Reference counting: Free blocks when requests complete - """ - - def __init__( - self, - max_batch_size: int = 16, - num_blocks: int = 512, - block_size: int = 256, - max_num_batched_tokens: int = 1024, - connector=None, - has_mamba_cache: bool = False, - num_mamba_cache_blocks: int | None = None, - ): - 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() - self.pending_free_blocks: dict[str, list[int]] = {} - self.pending_kv_decode_blocks: int = 0 - self.remote_kv_requests: dict[str, InferenceRequest] = {} - - self.cache_manager = BlockManager(num_blocks=num_blocks, block_size=block_size) - self.has_mamba_cache = has_mamba_cache - self.mamba_cache_manager = ( - MambaCacheManager(num_mamba_cache_blocks or max(2, num_blocks // 4)) - if has_mamba_cache - else None - ) - self.speculative_cache_ops = SpeculativeCacheOps(self.cache_manager) - self.block_size = block_size - self.max_num_batched_tokens = max_num_batched_tokens - self.connector = connector - - def add_request(self, request: InferenceRequest): - if request is not None: - request.status = RequestStatus.WAITING - self.waiting_queue.sync_q.put(request) - - def _exceeds_token_budget( - self, - current_num_batched_tokens: int, - num_tokens_this_step: int, - num_scheduled_requests: int, - ) -> bool: - """Return True when adding this request should be deferred for token budget. - - A single request is always allowed to make progress, even if it is larger - than max_num_batched_tokens. - """ - if num_scheduled_requests == 0: - return False - return ( - current_num_batched_tokens + num_tokens_this_step - > 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 +""" +Scheduler - Request scheduling and batch management with Paged Attention KV Cache. +""" + +import logging +import os +import queue +from typing import List, Optional + +import janus + +from infinilm.llm.cache_manager import BlockManager, MambaCacheManager +from infinilm.llm.request import InferenceRequest, RequestStatus + +logger = logging.getLogger(__name__) + + +class SpeculativeCacheOps: + """Limited cache operations needed by speculative verification.""" + + def __init__(self, cache_manager: BlockManager): + self._cache_manager = cache_manager + + def append_verify_slots( + self, + block_table: List[int], + start_length: int, + num_slots: int, + token_ids: List[int], + ): + return self._cache_manager.append_slots( + block_table, + start_length, + num_slots, + token_ids, + update_hash=False, + ) + + def rollback_to_length(self, block_table: List[int], keep_tokens: int): + return self._cache_manager.truncate_blocks(block_table, keep_tokens) + + def commit_accepted_tokens( + self, block_table: List[int], token_ids: List[int], num_tokens: int + ) -> None: + self._cache_manager.commit_blocks_hash(block_table, token_ids, num_tokens) + + +class SchedulerOutput: + """Scheduler output containing scheduled requests and execution phase info.""" + + def __init__( + self, + scheduled_requests: List[InferenceRequest], + is_prefill: bool = False, + speculative_cache_ops: Optional[SpeculativeCacheOps] = None, + ): + self.scheduled_requests = scheduled_requests + self.num_requests = len(scheduled_requests) + self.is_prefill = is_prefill + self.speculative_cache_ops = speculative_cache_ops + self.kv_connector_metadata = None + + +class Scheduler: + """Request scheduler with integrated BlockManager for KV cache management. + + Scheduling logic: + 1. Running queue: Check for new blocks needed, update slot_mapping + 2. Waiting queue: Try block reuse (prefix caching), allocate new blocks + 3. Reference counting: Free blocks when requests complete + """ + + def __init__( + self, + max_batch_size: int = 16, + num_blocks: int = 512, + block_size: int = 256, + max_num_batched_tokens: int = 1024, + connector=None, + has_mamba_cache: bool = False, + num_mamba_cache_blocks: int | None = None, + ): + 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() + self.pending_free_blocks: dict[str, list[int]] = {} + self.pending_kv_decode_blocks: int = 0 + self.remote_kv_requests: dict[str, InferenceRequest] = {} + + self.cache_manager = BlockManager(num_blocks=num_blocks, block_size=block_size) + self.has_mamba_cache = has_mamba_cache + self.mamba_cache_manager = ( + MambaCacheManager(num_mamba_cache_blocks or max(2, num_blocks // 4)) + if has_mamba_cache + else None + ) + self.speculative_cache_ops = SpeculativeCacheOps(self.cache_manager) + self.block_size = block_size + self.max_num_batched_tokens = max_num_batched_tokens + self.connector = connector + + def add_request(self, request: InferenceRequest): + if request is not None: + request.status = RequestStatus.WAITING + self.waiting_queue.sync_q.put(request) + + def _exceeds_token_budget( + self, + current_num_batched_tokens: int, + num_tokens_this_step: int, + num_scheduled_requests: int, + ) -> bool: + """Return True when adding this request should be deferred for token budget. + + A single request is always allowed to make progress, even if it is larger + than max_num_batched_tokens. + """ + if num_scheduled_requests == 0: + return False + return ( + current_num_batched_tokens + num_tokens_this_step + > 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 - + + 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) - + 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 + 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 = [] + scheduled_requests = [] + 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) < prefill_batch_limit + and current_num_batched_tokens < self.max_num_batched_tokens + ): + try: + req = self.waiting_queue.sync_q.get_nowait() + except queue.Empty: + break + # Skip requests that were already finished (e.g., timed out/canceled while waiting) + if req.is_finished(): + 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: + if self.has_mamba_cache: + cached_block_table = [] + num_local_computed_tokens = 0 + blocks_blueprint = [] + load_kv_async = False + num_external_computed_tokens = 0 + else: + ( + cached_block_table, + num_local_computed_tokens, + blocks_blueprint, + ) = self.cache_manager.get_computed_blocks( + req_tokens, req.get_mm_token_index_mappings() + ) + if self.connector is not None: + ext_tokens, load_kv_async = ( + self.connector.get_num_new_matched_tokens( + req, num_local_computed_tokens + ) + ) + num_external_computed_tokens = ext_tokens + else: + load_kv_async = False + num_external_computed_tokens = 0 + + num_computed_tokens = ( + num_local_computed_tokens + num_external_computed_tokens + ) + if load_kv_async: + num_computed_tokens -= 1 + num_new_tokens = req.get_prompt_length() - num_computed_tokens + + # Early token budget check: skip can_accept_request and allocate_slots + # for requests that would exceed the per-schedule token budget. + if not load_kv_async: + num_tokens_this_step = ( + req.get_prompt_length() - num_local_computed_tokens + ) + if self._exceeds_token_budget( + current_num_batched_tokens, + num_tokens_this_step, + len(scheduled_requests), + ): + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + + if not self.can_accept_request( + req, + num_local_computed_tokens, + current_prefill_extra_blocks, + ): + logger.warning( + "Insufficient KV cache blocks for request %s, deferring.", + req.request_id, + ) + + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + + req_blocks, slot_mapping = self.cache_manager.allocate_slots( + req_tokens, + num_new_tokens, + num_computed_tokens=num_computed_tokens, + cached_block_table=cached_block_table, + blocks_blueprint=blocks_blueprint, + delay_cache_blocks=load_kv_async, + ) + + if req_blocks is None: + logger.warning( + "Failed to allocate KV cache blocks for request: %s", + req.request_id, + ) + if num_local_computed_tokens > 0: + self.cache_manager.free_blocks(cached_block_table) + deferred_requests.append(req) + break + + if self.has_mamba_cache and req.mamba_cache_index is None: + req.mamba_cache_index = self.mamba_cache_manager.allocate() + if req.mamba_cache_index is None: + self.cache_manager.free_blocks(req_blocks) + logger.warning( + "Insufficient mamba cache rows for request %s, deferring.", + req.request_id, + ) + deferred_requests.append(req) + break + + req.block_table = req_blocks + req.slot_mapping = slot_mapping + req.num_blocks = len(req_blocks) + req.num_local_cached_tokens = num_local_computed_tokens + req.num_computed_tokens = num_computed_tokens + + if self.connector is not None: + self.connector.update_state_after_alloc( + req, + req.block_table, + num_external_computed_tokens, + self.block_size, + ) + else: + load_kv_async = False + num_tokens_this_step = ( + req.get_prompt_length() - req.num_local_cached_tokens + ) + if self._exceeds_token_budget( + current_num_batched_tokens, + num_tokens_this_step, + len(scheduled_requests), + ): + deferred_requests.append(req) + break + self.cache_manager.update_blocks_hash( + req.block_table, req.num_local_cached_tokens + ) + + if load_kv_async: + req.status = RequestStatus.WAITING_FOR_REMOTE_KVS + self.remote_kv_requests[req.request_id] = req + self.pending_kv_decode_blocks += ( + req.sampling_params.max_tokens + self.block_size - 1 + ) // self.block_size + continue + + current_prefill_extra_blocks += self._get_prefill_extra_blocks(req) + scheduled_requests.append(req) + + num_tokens_this_step = req.get_prompt_length() - req.num_local_cached_tokens + current_num_batched_tokens += num_tokens_this_step + + req.status = RequestStatus.RUNNING + + if deferred_requests: + for req in deferred_requests: + self.waiting_queue.sync_q.put(req) + + # Return prefill batch if any waiting requests were scheduled + if scheduled_requests: + is_prefill = True + scheduler_output = SchedulerOutput( + scheduled_requests=scheduled_requests, + is_prefill=is_prefill, + speculative_cache_ops=self.speculative_cache_ops, + ) + if self.connector is not None: + meta = self.connector.build_connector_meta() + scheduler_output.kv_connector_metadata = meta + return scheduler_output + + # Process Running queue (decode phase) + while len(scheduled_requests) < decode_batch_limit: + try: + req = self.running_queue.sync_q.get_nowait() + except queue.Empty: + break + # Skip requests that were already finished (e.g., timed out/canceled while running) + if req.is_finished(): + 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( + req.block_table, req.get_total_length(), req.get_all_token_ids() + ) + req.slot_mapping = [new_slot] + req.num_blocks = len(req.block_table) + req.num_local_cached_tokens = req.get_total_length() - 1 + scheduled_requests.append(req) + + except RuntimeError as e: + raise RuntimeError("No available cache blocks for new token") from e + + # Promote completed remote KV transfers (lower priority than running queue). + # Cleanup (is_finished, failed re-queue) runs unconditionally; batch append only if slots remain. + if self.connector is not None and self.remote_kv_requests: + for req_id in list(self.remote_kv_requests.keys()): + req = self.remote_kv_requests[req_id] + if req.is_finished(): + self.complete_requests([req]) + continue + if req_id in self.failed_receiving_kv_req_ids: + logger.warning( + f"Request {req_id[:8]}... failed receiving KV, re-queuing for prefill." + ) + self.update_waiting_for_remote_kv(req) + 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) < decode_batch_limit: + logger.info( + f"Request {req_id[:8]}... finished receiving KV, scheduling for decode." + ) + self.update_waiting_for_remote_kv(req) + req.status = RequestStatus.RUNNING + scheduled_requests.append(req) + else: + break # Defer promotion to next schedule() if batch is full + + # Return decode batch if any running requests were scheduled + if scheduled_requests: + is_prefill = False + scheduler_output = SchedulerOutput( + scheduled_requests=scheduled_requests, + is_prefill=is_prefill, + speculative_cache_ops=self.speculative_cache_ops, + ) + + if self.connector is not None: + meta = self.connector.build_connector_meta() + scheduler_output.kv_connector_metadata = meta + return scheduler_output + + if self.connector is not None: + scheduler_output = SchedulerOutput( + scheduled_requests=[], + speculative_cache_ops=self.speculative_cache_ops, + ) + meta = self.connector.build_connector_meta() + scheduler_output.kv_connector_metadata = meta + return scheduler_output + + return None + + def update_waiting_for_remote_kv(self, request: InferenceRequest): + self.remote_kv_requests.pop(request.request_id, None) + self.pending_kv_decode_blocks -= ( + request.sampling_params.max_tokens + self.block_size - 1 + ) // self.block_size + if request.request_id in self.failed_receiving_kv_req_ids: + if request.num_computed_tokens: + valid_block_count = request.num_computed_tokens // self.block_size + self.cache_manager.update_blocks_hash( + request.block_table[:valid_block_count], + request.num_local_cached_tokens, + ) + request.slot_mapping = self.cache_manager.update_blocks_slot( + request.block_table, + request.num_computed_tokens, + request.get_prompt_length(), + ) + request.num_local_cached_tokens = request.num_computed_tokens + else: + self.cache_manager.free_blocks(request.block_table) + request.block_table = [] + request.slot_mapping = [] + request.num_local_cached_tokens = 0 + self.failed_receiving_kv_req_ids.discard(request.request_id) + else: + self.cache_manager.update_blocks_hash( + request.block_table, request.num_local_cached_tokens + ) + request.num_local_cached_tokens = request.num_computed_tokens + self.finished_receiving_kv_req_ids.discard(request.request_id) + + def complete_requests(self, requests: List[InferenceRequest]): + """Handle completed requests and free their blocks.""" + for req in requests: + if req.status in [ + RequestStatus.FINISHED, + RequestStatus.CANCELED, + RequestStatus.FAILED, + RequestStatus.TIMEOUT, + ]: + delay_free_blocks = False + if self.connector is not None: + delay_free_blocks, _ = self.connector.request_finished( + req, req.block_table, self.block_size + ) + + if req.request_id in self.remote_kv_requests: + self.pending_kv_decode_blocks -= ( + req.sampling_params.max_tokens + self.block_size - 1 + ) // self.block_size + self.remote_kv_requests.pop(req.request_id, None) + if req.request_id in self.finished_receiving_kv_req_ids: + self.finished_receiving_kv_req_ids.discard(req.request_id) + self.failed_receiving_kv_req_ids.discard(req.request_id) + else: + delay_free_blocks = True + if req.block_table and not delay_free_blocks: + self.cache_manager.free_blocks(req.block_table) + elif req.block_table and delay_free_blocks: + self.pending_free_blocks[req.request_id] = list(req.block_table) + if self.mamba_cache_manager is not None: + self.mamba_cache_manager.free(req.mamba_cache_index) + req.mamba_cache_index = None + + if req.status == RequestStatus.CANCELED: + logger.info( + f"Request {req.request_id[:8]}... canceled: {req.finish_reason}" + ) + elif req.status == RequestStatus.FAILED: + logger.error( + f"Request {req.request_id[:8]}... failed: {req.finish_reason}" + ) + elif req.status == RequestStatus.TIMEOUT: + logger.error( + f"Request {req.request_id[:8]}... timed out: {req.finish_reason}" + ) + else: + # Still running, put back in running queue + self.running_queue.sync_q.put(req) + + def can_accept_request( + self, + request: InferenceRequest, + num_local_computed_tokens: int, + current_prefill_extra_blocks: int = 0, + ) -> bool: + if ( + self.mamba_cache_manager is not None + and request.mamba_cache_index is None + and not self.mamba_cache_manager.can_allocate() + ): + return False + + total_required_blocks = 0 + + # Calculate blocks needed for running requests + running_queue_size = self.running_queue.sync_q.qsize() + for _ in range(running_queue_size): + req = self.running_queue.sync_q.get() + remaining_tokens = ( + req.sampling_params.max_tokens - req.get_num_generated_tokens() + ) + num_blocks_needed = ( + remaining_tokens + self.block_size - 1 + ) // self.block_size + total_required_blocks += num_blocks_needed + self.running_queue.sync_q.put(req) + + # Calculate blocks needed for the new request + total_length = request.get_prompt_length() - num_local_computed_tokens + total_length += request.sampling_params.max_tokens + num_blocks_needed = (total_length + self.block_size - 1) // self.block_size + total_required_blocks += num_blocks_needed + + # Include decode headroom for WAITING_FOR_REMOTE_KVS requests, which + # hold prompt blocks but will also need decode blocks once promoted. + total_required_blocks += self.pending_kv_decode_blocks + + # Include decode headroom for requests accepted earlier in this batch. + total_required_blocks += current_prefill_extra_blocks + + # Compare with total usable blocks in cache manager + return total_required_blocks <= self.cache_manager.get_total_usable_blocks() + + def _get_prefill_extra_blocks(self, request: InferenceRequest) -> int: + total_length = request.get_prompt_length() + total_length += request.sampling_params.max_tokens + total_required_blocks = (total_length + self.block_size - 1) // self.block_size + return max(total_required_blocks - len(request.block_table), 0) + + def update_from_output(self, model_output): + if self.connector is None or model_output.kv_connector_output is None: + return + + finished_recving_req_ids = ( + getattr(model_output.kv_connector_output, "finished_recving", None) or [] + ) + finished_sending_req_ids = ( + getattr(model_output.kv_connector_output, "finished_sending", None) or [] + ) + failed_recving_req_ids = ( + getattr(model_output.kv_connector_output, "failed_recving", None) or [] + ) + invalid_block_ids = ( + getattr(model_output.kv_connector_output, "invalid_block_ids", None) or [] + ) + + for req_id in finished_recving_req_ids: + if req_id in self.pending_free_blocks: + # Aborted request: transfer complete, now safe to free blocks. + self.cache_manager.free_blocks(self.pending_free_blocks.pop(req_id)) + elif req_id in self.remote_kv_requests: + # Active request: mark ready for promotion in schedule(). + self.finished_receiving_kv_req_ids.add(req_id) + # else: already processed or unknown, discard to avoid stale entries. + for req_id in finished_sending_req_ids: + self.cache_manager.free_blocks(self.pending_free_blocks.pop(req_id, [])) + for req_id in failed_recving_req_ids: + # Only track failures for active (non-aborted) requests; aborted + # requests are handled via pending_free_blocks in finished_recving. + if req_id in self.remote_kv_requests: + self.failed_receiving_kv_req_ids.add(req_id) + + if invalid_block_ids: + invalid_set = set(invalid_block_ids) + + for req in self.remote_kv_requests.values(): + start_block_idx = req.num_local_cached_tokens // self.block_size + for i, block_id in enumerate( + req.block_table[start_block_idx:], start=start_block_idx + ): + if block_id in invalid_set: + req.num_computed_tokens = i * self.block_size + break + elif self.failed_receiving_kv_req_ids: + for req_id in self.failed_receiving_kv_req_ids: + req = self.remote_kv_requests[req_id] + req.num_computed_tokens = req.num_local_cached_tokens + + def get_cache_stats(self) -> dict: + """Get cache statistics.""" + stats = { + "num_blocks": self.cache_manager.num_blocks, + "block_size": self.cache_manager.block_size, + "num_free_blocks": self.cache_manager.get_num_free_blocks(), + "usable_blocks": self.cache_manager.get_total_usable_blocks(), + "num_pending_blocks": len(self.cache_manager.pending_block_ids), + "num_used_blocks": len(self.cache_manager.used_block_ids), + } + if self.mamba_cache_manager is not None: + stats.update( + { + "num_mamba_cache_blocks": self.mamba_cache_manager.num_blocks, + "num_free_mamba_cache_blocks": self.mamba_cache_manager.get_num_free_blocks(), + "num_used_mamba_cache_blocks": len( + self.mamba_cache_manager.used_block_ids + ), + } ) - return long_in_batch >= self.long_context_max_batch - - def schedule(self) -> Optional[SchedulerOutput]: - """Schedule and return batch of requests to execute.""" - deferred_requests = [] - scheduled_requests = [] - 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) < prefill_batch_limit - and current_num_batched_tokens < self.max_num_batched_tokens - ): - try: - req = self.waiting_queue.sync_q.get_nowait() - except queue.Empty: - break - # Skip requests that were already finished (e.g., timed out/canceled while waiting) - if req.is_finished(): - 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: - if self.has_mamba_cache: - cached_block_table = [] - num_local_computed_tokens = 0 - blocks_blueprint = [] - load_kv_async = False - num_external_computed_tokens = 0 - else: - ( - cached_block_table, - num_local_computed_tokens, - blocks_blueprint, - ) = self.cache_manager.get_computed_blocks( - req_tokens, req.get_mm_token_index_mappings() - ) - if self.connector is not None: - ext_tokens, load_kv_async = ( - self.connector.get_num_new_matched_tokens( - req, num_local_computed_tokens - ) - ) - num_external_computed_tokens = ext_tokens - else: - load_kv_async = False - num_external_computed_tokens = 0 - - num_computed_tokens = ( - num_local_computed_tokens + num_external_computed_tokens - ) - if load_kv_async: - num_computed_tokens -= 1 - num_new_tokens = req.get_prompt_length() - num_computed_tokens - - # Early token budget check: skip can_accept_request and allocate_slots - # for requests that would exceed the per-schedule token budget. - if not load_kv_async: - num_tokens_this_step = ( - req.get_prompt_length() - num_local_computed_tokens - ) - if self._exceeds_token_budget( - current_num_batched_tokens, - num_tokens_this_step, - len(scheduled_requests), - ): - if num_local_computed_tokens > 0: - self.cache_manager.free_blocks(cached_block_table) - deferred_requests.append(req) - break - - if not self.can_accept_request( - req, - num_local_computed_tokens, - current_prefill_extra_blocks, - ): - logger.warning( - "Insufficient KV cache blocks for request %s, deferring.", - req.request_id, - ) - - if num_local_computed_tokens > 0: - self.cache_manager.free_blocks(cached_block_table) - deferred_requests.append(req) - break - - req_blocks, slot_mapping = self.cache_manager.allocate_slots( - req_tokens, - num_new_tokens, - num_computed_tokens=num_computed_tokens, - cached_block_table=cached_block_table, - blocks_blueprint=blocks_blueprint, - delay_cache_blocks=load_kv_async, - ) - - if req_blocks is None: - logger.warning( - "Failed to allocate KV cache blocks for request: %s", - req.request_id, - ) - if num_local_computed_tokens > 0: - self.cache_manager.free_blocks(cached_block_table) - deferred_requests.append(req) - break - - if self.has_mamba_cache and req.mamba_cache_index is None: - req.mamba_cache_index = self.mamba_cache_manager.allocate() - if req.mamba_cache_index is None: - self.cache_manager.free_blocks(req_blocks) - logger.warning( - "Insufficient mamba cache rows for request %s, deferring.", - req.request_id, - ) - deferred_requests.append(req) - break - - req.block_table = req_blocks - req.slot_mapping = slot_mapping - req.num_blocks = len(req_blocks) - req.num_local_cached_tokens = num_local_computed_tokens - req.num_computed_tokens = num_computed_tokens - - if self.connector is not None: - self.connector.update_state_after_alloc( - req, - req.block_table, - num_external_computed_tokens, - self.block_size, - ) - else: - load_kv_async = False - num_tokens_this_step = ( - req.get_prompt_length() - req.num_local_cached_tokens - ) - if self._exceeds_token_budget( - current_num_batched_tokens, - num_tokens_this_step, - len(scheduled_requests), - ): - deferred_requests.append(req) - break - self.cache_manager.update_blocks_hash( - req.block_table, req.num_local_cached_tokens - ) - - if load_kv_async: - req.status = RequestStatus.WAITING_FOR_REMOTE_KVS - self.remote_kv_requests[req.request_id] = req - self.pending_kv_decode_blocks += ( - req.sampling_params.max_tokens + self.block_size - 1 - ) // self.block_size - continue - - current_prefill_extra_blocks += self._get_prefill_extra_blocks(req) - scheduled_requests.append(req) - - num_tokens_this_step = req.get_prompt_length() - req.num_local_cached_tokens - current_num_batched_tokens += num_tokens_this_step - - req.status = RequestStatus.RUNNING - - if deferred_requests: - for req in deferred_requests: - self.waiting_queue.sync_q.put(req) - - # Return prefill batch if any waiting requests were scheduled - if scheduled_requests: - is_prefill = True - scheduler_output = SchedulerOutput( - scheduled_requests=scheduled_requests, - is_prefill=is_prefill, - speculative_cache_ops=self.speculative_cache_ops, - ) - if self.connector is not None: - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output - - # Process Running queue (decode phase) - while len(scheduled_requests) < decode_batch_limit: - try: - req = self.running_queue.sync_q.get_nowait() - except queue.Empty: - break - # Skip requests that were already finished (e.g., timed out/canceled while running) - if req.is_finished(): - 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( - req.block_table, req.get_total_length(), req.get_all_token_ids() - ) - req.slot_mapping = [new_slot] - req.num_blocks = len(req.block_table) - req.num_local_cached_tokens = req.get_total_length() - 1 - scheduled_requests.append(req) - - except RuntimeError as e: - raise RuntimeError("No available cache blocks for new token") from e - - # Promote completed remote KV transfers (lower priority than running queue). - # Cleanup (is_finished, failed re-queue) runs unconditionally; batch append only if slots remain. - if self.connector is not None and self.remote_kv_requests: - for req_id in list(self.remote_kv_requests.keys()): - req = self.remote_kv_requests[req_id] - if req.is_finished(): - self.complete_requests([req]) - continue - if req_id in self.failed_receiving_kv_req_ids: - logger.warning( - f"Request {req_id[:8]}... failed receiving KV, re-queuing for prefill." - ) - self.update_waiting_for_remote_kv(req) - 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) < decode_batch_limit: - logger.info( - f"Request {req_id[:8]}... finished receiving KV, scheduling for decode." - ) - self.update_waiting_for_remote_kv(req) - req.status = RequestStatus.RUNNING - scheduled_requests.append(req) - else: - break # Defer promotion to next schedule() if batch is full - - # Return decode batch if any running requests were scheduled - if scheduled_requests: - is_prefill = False - scheduler_output = SchedulerOutput( - scheduled_requests=scheduled_requests, - is_prefill=is_prefill, - speculative_cache_ops=self.speculative_cache_ops, - ) - - if self.connector is not None: - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output - - if self.connector is not None: - scheduler_output = SchedulerOutput( - scheduled_requests=[], - speculative_cache_ops=self.speculative_cache_ops, - ) - meta = self.connector.build_connector_meta() - scheduler_output.kv_connector_metadata = meta - return scheduler_output - - return None - - def update_waiting_for_remote_kv(self, request: InferenceRequest): - self.remote_kv_requests.pop(request.request_id, None) - self.pending_kv_decode_blocks -= ( - request.sampling_params.max_tokens + self.block_size - 1 - ) // self.block_size - if request.request_id in self.failed_receiving_kv_req_ids: - if request.num_computed_tokens: - valid_block_count = request.num_computed_tokens // self.block_size - self.cache_manager.update_blocks_hash( - request.block_table[:valid_block_count], - request.num_local_cached_tokens, - ) - request.slot_mapping = self.cache_manager.update_blocks_slot( - request.block_table, - request.num_computed_tokens, - request.get_prompt_length(), - ) - request.num_local_cached_tokens = request.num_computed_tokens - else: - self.cache_manager.free_blocks(request.block_table) - request.block_table = [] - request.slot_mapping = [] - request.num_local_cached_tokens = 0 - self.failed_receiving_kv_req_ids.discard(request.request_id) - else: - self.cache_manager.update_blocks_hash( - request.block_table, request.num_local_cached_tokens - ) - request.num_local_cached_tokens = request.num_computed_tokens - self.finished_receiving_kv_req_ids.discard(request.request_id) - - def complete_requests(self, requests: List[InferenceRequest]): - """Handle completed requests and free their blocks.""" - for req in requests: - if req.status in [ - RequestStatus.FINISHED, - RequestStatus.CANCELED, - RequestStatus.FAILED, - RequestStatus.TIMEOUT, - ]: - delay_free_blocks = False - if self.connector is not None: - delay_free_blocks, _ = self.connector.request_finished( - req, req.block_table, self.block_size - ) - - if req.request_id in self.remote_kv_requests: - self.pending_kv_decode_blocks -= ( - req.sampling_params.max_tokens + self.block_size - 1 - ) // self.block_size - self.remote_kv_requests.pop(req.request_id, None) - if req.request_id in self.finished_receiving_kv_req_ids: - self.finished_receiving_kv_req_ids.discard(req.request_id) - self.failed_receiving_kv_req_ids.discard(req.request_id) - else: - delay_free_blocks = True - if req.block_table and not delay_free_blocks: - self.cache_manager.free_blocks(req.block_table) - elif req.block_table and delay_free_blocks: - self.pending_free_blocks[req.request_id] = list(req.block_table) - if self.mamba_cache_manager is not None: - self.mamba_cache_manager.free(req.mamba_cache_index) - req.mamba_cache_index = None - - if req.status == RequestStatus.CANCELED: - logger.info( - f"Request {req.request_id[:8]}... canceled: {req.finish_reason}" - ) - elif req.status == RequestStatus.FAILED: - logger.error( - f"Request {req.request_id[:8]}... failed: {req.finish_reason}" - ) - elif req.status == RequestStatus.TIMEOUT: - logger.error( - f"Request {req.request_id[:8]}... timed out: {req.finish_reason}" - ) - else: - # Still running, put back in running queue - self.running_queue.sync_q.put(req) - - def can_accept_request( - self, - request: InferenceRequest, - num_local_computed_tokens: int, - current_prefill_extra_blocks: int = 0, - ) -> bool: - if ( - self.mamba_cache_manager is not None - and request.mamba_cache_index is None - and not self.mamba_cache_manager.can_allocate() - ): - return False - - total_required_blocks = 0 - - # Calculate blocks needed for running requests - running_queue_size = self.running_queue.sync_q.qsize() - for _ in range(running_queue_size): - req = self.running_queue.sync_q.get() - remaining_tokens = ( - req.sampling_params.max_tokens - req.get_num_generated_tokens() - ) - num_blocks_needed = ( - remaining_tokens + self.block_size - 1 - ) // self.block_size - total_required_blocks += num_blocks_needed - self.running_queue.sync_q.put(req) - - # Calculate blocks needed for the new request - total_length = request.get_prompt_length() - num_local_computed_tokens - total_length += request.sampling_params.max_tokens - num_blocks_needed = (total_length + self.block_size - 1) // self.block_size - total_required_blocks += num_blocks_needed - - # Include decode headroom for WAITING_FOR_REMOTE_KVS requests, which - # hold prompt blocks but will also need decode blocks once promoted. - total_required_blocks += self.pending_kv_decode_blocks - - # Include decode headroom for requests accepted earlier in this batch. - total_required_blocks += current_prefill_extra_blocks - - # Compare with total usable blocks in cache manager - return total_required_blocks <= self.cache_manager.get_total_usable_blocks() - - def _get_prefill_extra_blocks(self, request: InferenceRequest) -> int: - total_length = request.get_prompt_length() - total_length += request.sampling_params.max_tokens - total_required_blocks = (total_length + self.block_size - 1) // self.block_size - return max(total_required_blocks - len(request.block_table), 0) - - def update_from_output(self, model_output): - if self.connector is None or model_output.kv_connector_output is None: - return - - finished_recving_req_ids = ( - getattr(model_output.kv_connector_output, "finished_recving", None) or [] - ) - finished_sending_req_ids = ( - getattr(model_output.kv_connector_output, "finished_sending", None) or [] - ) - failed_recving_req_ids = ( - getattr(model_output.kv_connector_output, "failed_recving", None) or [] - ) - invalid_block_ids = ( - getattr(model_output.kv_connector_output, "invalid_block_ids", None) or [] - ) - - for req_id in finished_recving_req_ids: - if req_id in self.pending_free_blocks: - # Aborted request: transfer complete, now safe to free blocks. - self.cache_manager.free_blocks(self.pending_free_blocks.pop(req_id)) - elif req_id in self.remote_kv_requests: - # Active request: mark ready for promotion in schedule(). - self.finished_receiving_kv_req_ids.add(req_id) - # else: already processed or unknown, discard to avoid stale entries. - for req_id in finished_sending_req_ids: - self.cache_manager.free_blocks(self.pending_free_blocks.pop(req_id, [])) - for req_id in failed_recving_req_ids: - # Only track failures for active (non-aborted) requests; aborted - # requests are handled via pending_free_blocks in finished_recving. - if req_id in self.remote_kv_requests: - self.failed_receiving_kv_req_ids.add(req_id) - - if invalid_block_ids: - invalid_set = set(invalid_block_ids) - - for req in self.remote_kv_requests.values(): - start_block_idx = req.num_local_cached_tokens // self.block_size - for i, block_id in enumerate( - req.block_table[start_block_idx:], start=start_block_idx - ): - if block_id in invalid_set: - req.num_computed_tokens = i * self.block_size - break - elif self.failed_receiving_kv_req_ids: - for req_id in self.failed_receiving_kv_req_ids: - req = self.remote_kv_requests[req_id] - req.num_computed_tokens = req.num_local_cached_tokens - - def get_cache_stats(self) -> dict: - """Get cache statistics.""" - stats = { - "num_blocks": self.cache_manager.num_blocks, - "block_size": self.cache_manager.block_size, - "num_free_blocks": self.cache_manager.get_num_free_blocks(), - "usable_blocks": self.cache_manager.get_total_usable_blocks(), - "num_pending_blocks": len(self.cache_manager.pending_block_ids), - "num_used_blocks": len(self.cache_manager.used_block_ids), - } - if self.mamba_cache_manager is not None: - stats.update( - { - "num_mamba_cache_blocks": self.mamba_cache_manager.num_blocks, - "num_free_mamba_cache_blocks": self.mamba_cache_manager.get_num_free_blocks(), - "num_used_mamba_cache_blocks": len( - self.mamba_cache_manager.used_block_ids - ), - } - ) - return stats + return stats