Prhub

#32208 O(1) slot allocation in ReqToTokenPool.alloc()

原始 PR 作者 DmiSergeev 合并时间 2026-08-12 02:26 文件变更 3 提交数 2 评论 7 代码增减 +25 / -6

执行摘要

ReqToTokenPool 槽分配改尾部弹出,热路径由 O(pool) 降为 O(batch)

alloc() 每调度步准入新请求时执行一次,且运行在单线程调度进程上,直接门控 GPU 分发。原实现 self.free_slots = self.free_slots[need_size:] 会整体复制幸存指针,PR body 指出该行复杂度与剩余空闲槽数成正比而非与实际分配数成正比:“this line is O(len(free_slots)) — proportional to how many free slots remain, not to how many are being allocated (need_size)… a purely wasted allocation-dominated CPU cost on every scheduler iteration”。在 --max-running-requests(4096/16384)高并发稳态下 free_slots 常达数千、每步准入仅个位数到低百位,每次 alloc() 都是纯浪费的 CPU 开销;free() 已是 O(1) 均摊,本 PR 将 alloc 侧补齐。

值得精读。核心看点有两个:一是 memory_pool.py 中 alloc() 的尾部弹出与零长度切片防护,二是测试断言如何随分配策略变化而显式放宽。学习点包括 Python 列表切片复杂度、-0 切片陷阱,以及“不透明索引 + req_generation 代际计数解耦”的池设计思路。

讨论亮点

维护者 xiezhq-hermann 认可改动简洁,建议跳过配套 benchmark 与回归测试文件:“good call, the changes seem concise enough and I think we can skip the benchmark and the test file for the changes”,作者随后将测试与 benchmark 移入 gist。Review 中唯一代码级意见是精简内联注释:“please shorten the comments as well”(memory_pool.py:310),作者回复“Done.”后获得 APPROVED。另外,CI 分区 base-b-test-1-gpu-large 的 test_awq.py(Mixtral-8x7B AWQ INT4)失败被作者定位为 AWQ/Marlin MoE 的 JIT kernel align_single_token.cuh 在 CUDA graph capture 期间的编译错误,与本次内存池记账改动无关。

实现拆解

  1. 变更入口与核心替换python/sglang/srt/mem_cache/memory_pool.pyReqToTokenPool.alloc(),将 self.free_slots[:need_size] 头部切片加整体列表重建,替换为 select_index = self.free_slots[-need_size:] + del self.free_slots[-need_size:],只触碰被移除的 need_size 个元素,原地删除尾部不触发剩余元素拷贝。
  2. 边界防护need_size == 0free_slots[-0:]free_slots[:] 等价(-0 == 0),会取走整个空闲池;显式 select_index = [] 保证全批次复用槽位时(如 chunked prefill 续批)不消耗任何空闲槽。
  3. 测试契约同步:分配顺序从头部改为尾部后,test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py 的两个 extend 测试与 test/registered/unit/hardware_backend/mlx/test_attention_patching.py 的 MLX auxiliary state 测试放宽了原先“新请求固定拿到索引 2/1”的断言,改为只校验复用行索引不变、新行拿到差异化合法槽位,将“分配顺序无语义”显式化为测试契约。
  4. 配套调整:无配置、schema、部署变更;贡献者最初提交了独立回归测试与 benchmark 文件,经维护者确认改动足够简洁后移除并迁移至 gist 附录,仓库内回归测试与 CI 均通过。
文件 模块 状态 重要度
python/sglang/srt/mem_cache/memory_pool.py 内存池 modified 5.91
test/registered/unit/mem_cache/test_dllm_fdfo_kv_reuse.py 内存池 modified 4.42
test/registered/unit/hardware_backend/mlx/test_attention_patching.py 内存池 modified 3.49

关键符号

ReqToTokenPool.alloc

关键源码片段

python/sglang/srt/mem_cache/memory_pool.py core-logic

唯一源码改动点:ReqToTokenPool.alloc() 从头部切片改为尾部弹出,附带 need_size == 0 防护,是本次优化的核心。

def alloc(self, reqs: list[Req]) -> Optional[List[int]]:
    # 收集本批次中已有 req_pool_idx、将复用旧槽位的请求(例如跨 chunk 的
    # chunked prefill 继续请求),它们不需要新分配槽位。
    reusing = [i for i, r in enumerate(reqs) if r.req_pool_idx is not None]
    # 复用请求必须是 chunked 或已有 committed KV;该限制已由 #20476
    # 暂时放宽,故此处不再断言同批次只能有一个 chunked 复用请求。
    assert all(
        reqs[i].inflight_middle_chunks > 0 or reqs[i].kv_committed_len > 0
        for i in reusing
    ), "reusing request must be chunked or have committed KV"
​
    need_size = len(reqs) - len(reusing)
    if need_size > len(self.free_slots):
        return None
    if need_size > 0:
        # 从尾部弹出 N 个槽位:只搬运被移除的元素,成本 O(need_size);
        # 从头部弹出并重建列表会整体复制剩余指针,成本 O(len(free_slots))。
        select_index = self.free_slots[-need_size:]
        del self.free_slots[-need_size:]
    else:
        # 单独处理 need_size == 0:free_slots[-0:] 等价于 free_slots[:],
        # 会取到整个列表而不是空列表,所以必须显式返回空列表。
        select_index = []
    offset = 0
    for r in reqs:
        if r.req_pool_idx is None:
            r.req_pool_idx = select_index[offset]
            self.req_generation[r.req_pool_idx] += 1
            offset += 1
    return [r.req_pool_idx for r in reqs]

评论区精华

是否保留配套 benchmark 与回归测试文件 测试

作者最初提交了独立回归测试与 benchmark 文件;维护者 xiezhq-hermann 认为改动足够简洁,建议跳过:“good call, the changes seem concise enough and I think we can skip the benchmark and the test file for the changes。”

结论:作者将测试与 benchmark 迁移至 gist 并从 PR 移除,回归测试改为依赖既有单元测试文件。 · 已解决

alloc() 内联注释冗长需精简 style

reviewer xiezhq-hermann 在 memory_pool.py:310 要求缩短注释:“please shorten the comments as well”,作者回复“Done.”。

结论:作者按要求精简注释,随后获得 APPROVED。 · 已解决

CI 失败(test_awq.py JIT 编译错误)与本 PR 无关 other

base-b-test-1-gpu-large 分区中 test_awq.py(Mixtral-8x7B AWQ INT4)失败,作者排查为 AWQ/Marlin MoE 的 JIT kernel align_single_token.cuh 在 CUDA graph capture 期间报 namespace 错误,并贴出完整编译日志。

结论:失败根因在 JIT kernel 编译环境,与本次内存池记账改动无关。 · 已解决

风险与影响

  1. 分配顺序语义变化:尾部弹出与 free() 的尾部追加构成 LIFO 栈,刚释放的槽位最可能被立即复用。PR 明确 req_pool_idx 是不透明行索引、req_generation 已单独跟踪复用,但若存在隐式依赖“编号小的槽位先分配”的监控或外部断言,可能受影响;测试断言放宽是该变化的直接信号。
  2. 负零切片回归风险:need_size == 0 的 guard 一旦在后续重构中丢失,会出现“整池空闲槽被一次性分配”的严重错误,建议保留该分支及注释。
  3. 影响面收敛:free()、available_size()、clear() 与 free_slots 的类型和语义均未变,invariant_checker.py、streaming_session.py 等外部读取方无需修改。
  4. 性能影响:尾弹出为原地删除,无列表整体复制;基准显示池 16384、空闲 14745、单步准入 1-4 时均值由约 32 us 降至约 3 us(约 10.5x),病理场景约 3.9x,小池大批次时新旧实现持平。属调度器 CPU 分发开销优化,端到端吞吐仅在极高并发/短 step 场景体现。

对用户与系统:不改变模型输出与 KV 缓存内容,端到端正确性无影响;调度器 CPU 账本开销下降,对超大规模部署(--max-running-requests 4096/16384 级别、高并发短 step)的步间开销改善明显,PR body 估算病理场景每百万调度步约节省 28 CPU 秒。对团队与工程:将“分配顺序无语义”上升为测试契约,后续修改分配策略时测试兼容性更好,同时示范了一种低成本消除热路径列表拷贝的模式,可推广到同类空闲池实现。

调度器热路径变更 槽位分配顺序语义变化 Python 负零切片陷阱防护 测试断言显式放宽

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论