大规模RL Rollout：超越Prefix Locality的调度策略设计

大规模强化学习训练阶段模型不是只跑训练迭代还要同时做大量 Rollout 生成。每个训练 step 都可能触发成百上千条推理请求这些请求混合了短回复、长生成、多轮对话和奖励模型打分调度起来和传统在线推理服务完全不是一回事。Prefix Locality 是这几年 KV Cache 优化里一个很有效的思路但在混合 RL Rollout 场景下只靠前缀复用并不能解决所有问题。这篇文章我们不聊概念直接拆解“Scheduling Mixed RL Rollouts Beyond Prefix Locality”背后的调度问题讲清楚为什么前缀局部性不够用、混合 Rollout 到底混合了什么、调度器应该怎么重新设计以及落地时该看哪些指标。如果你正在做 RL 训练框架、推理引擎选型或者要给大模型训练集群设计调度模块这篇文章可以直接收藏。全文会覆盖核心机制拆解、调度器设计路径、配置示例、评测方法和常见坑位避免你在实际搭建时走弯路。1. 核心概念速览能力项说明主题类型大规模 RL 训练中的 Rollout 推理调度策略设计核心问题打破 Prefix Locality 的单一优化视角处理混合请求、动态前缀和训练-推理资源竞争关键技术动态前缀索引、优先级调度、两阶段资源预算、冷启动调度策略、训练感知调度适用场景LLM RL 训练管道、分布式推理引擎、在线 Rollout 服务、多任务混合推理队列主要收益提升样本吞吐、降低训练空转、提高 KV Cache 复用率、平滑训练早期冷启动波动不适用场景单模型单请求的简单推理、无 RL 训练的纯在线对话服务落地难度中高涉及调度器、缓存索引、训练器多模块协同先回到最基础的问题Prefix Locality 是什么它利用了“多个请求共享同一段前缀”的特征在推理时只计算一次共享前缀的 KV Cache后续请求直接复用。这在共享系统提示词、Few-shot 示例、多轮对话历史等场景下非常有效。vLLM 的 prefix caching、SGLang 的 RadixAttention 都是这个思路的代表实现。但 RL Rollout 和传统在线推理有一个本质差别请求分布不是静态的。策略网络每轮都在更新生成结果越来越偏离初始前缀分布同时 Rollout 请求还混着奖励模型打分、短回答、长生成等不同类型。单纯按前缀组织调度容易出现前缀树碎片化、缓存命中率忽高忽低、训练器饿死等问题。下面我们把 Prefix Locality 的适用边界和失效场景拆开来看。2. 为什么 Prefix Locality 在 RL Rollout 场景中不够用2.1 在线推理与 RL Rollout 的差异在线推理服务里请求来源是用户前缀结构相对稳定。比如同一产品的所有用户都共享系统提示词或者同一个会话的后续轮次复用之前的历史。这种情况下前缀复用收益稳定且可预测调度器可以放心地把缓存命中率当作核心优化项。RL Rollout 的场景完全不同第一产生请求的主体是训练器。每个训练 step 之后策略网络的参数变了导致后续请求的内容分布发生变化。虽然 prompt 本身可能来自同一个任务族但策略更新后生成的响应、采样的动作、奖励模型打分的对象都会漂移。第二请求类型是混合的。同一个调度周期内可能同时存在需要生成 512 个 token 的长回答只需要输出“0/1”或一个分数的奖励查询需要多轮交互的状态查询从回放缓冲区里采样的旧策略数据重放。这些请求对前缀的敏感程度、对延迟的要求、对资源的消耗完全不同。如果一个调度器只盯前缀复用很容易把资源全部喂给了长序列复用请求导致短请求和关键路径上的训练请求饿死。第三训练与推理共享显存和算力。RL 训练环节GPU 既要跑训练的前向反向又要跑 Rollout 的推理生成还要跑奖励模型的打分。调度器如果不感知训练器的当前状态就会出现推理请求把算力占满、训练 step 迟迟等不到 Rollout 数据的情况。2.2 Prefix Locality 失效的具体场景场景一请求前缀高度动态。当任务本身没有稳定共享前缀时比如每个 prompt 都是独立采样的数学题、代码题前缀复用率天然很低。此时 Prefix Locality 的收益很小继续按前缀聚合请求反而增加调度器的组织成本。场景二前缀树碎片化。RL 训练里同一个任务族往往基于一个基础 prompt 做少量扰动。这些 prompt 之间共享很长一段前缀但结尾略有差异。如果调度器把每条请求都当作独立前缀插入Radix Tree 会越来越碎索引本身占用内存缓存命中率却没有提升。需要定期合并、剪枝和过期清理。场景三训练冷启动阶段。训练初期策略网络接近随机初始化生成的输出五花八门前缀复用率极低。此时如果调度器执着于“找到可复用的前缀块”反而会延迟请求执行。材料里提到的“rl 冷启动”就是这个问题冷启动时模型输出熵高、分布不稳定调度策略应该偏重快速收集多样样本而不是追求缓存复用等训练中后期策略收敛、输出分布稳定后再逐步提高前缀复用权重。场景四延迟敏感的训练关键路径。RL 训练中Rollout 数据是训练 step 的“原料”。如果调度器把所有请求都按最大吞吐模式排布某个关键 batch 的 Rollout 结果迟迟不返回训练器只能空转。调度器需要区分“关键路径请求”和“后台请求”前者优先后者可以等待。所以超越 Prefix Locality 的核心不是抛弃前缀复用而是把它从“唯一目标”降级为“一个加权因素”同时引入训练感知、多样性保障和混合请求优先级。3. Mixed RL Rollouts 到底混合了什么设计调度策略之前先把“混合”这个词拆清楚。混合 RL Rollout 至少在三个维度上混合。3.1 任务类型混合请求类型输出长度计算特征典型用途生成式 Rollout中到长128~1024 token高计算量KV Cache 增长快策略采样、轨迹生成判别式/奖励查询极短1~10 token低计算量延迟敏感奖励模型打分、价值估计状态/动作查询短到中中等计算量多轮决策、环境交互回放数据重放任意无在线生成只做训练输入旧策略样本训练不同任务类型的资源消耗差异极大。一个生成式请求的 token 消耗可能是奖励查询的百倍。调度器必须按类型拆分队列否则一个长生成请求会阻塞后面大量短请求。3.2 数据分布混合RL 训练过程中Rollout 数据来自不同阶段当前策略采样旧策略的 replay buffer混有探索噪声的样本来自不同任务群体的数据。这些数据的前缀分布不同冗余程度也不同。调度器如果只按请求到达顺序处理可能出现某个任务族的样本严重过采样另一个任务族长期饥饿导致训练数据分布偏移。3.3 资源需求混合训练机器上GPU 资源不是全部分给推理的。训练 step 和 Rollout 生成存在时间上的交错训练前向/反向阶段计算资源被训练占用Rollout 阶段训练器等待数据资源释放给推理。调度器需要做资源预算控制明确“当前时间片里推理最多占多少显存和算力”。否则训练和推理会出现互相争抢吞吐反而低于串行执行。3.4 混合的意义“Mixed”是问题同时也是优化机会。不同类型的请求可以互相填补调度空隙。例如短请求可以插在长请求的 Prefill 阶段中间执行缓存复用率高的请求可以优先填充等待窗口。调度器的核心设计目标就是在这些混合请求之间找到一组调度顺序和资源分配策略使得训练样本吞吐最大化、训练空转时间最小化。4. 调度器核心设计思路超越 Prefix Locality4.1 设计目标调度器不再追求单一指标例如缓存命中率而是维护一个综合目标调度收益 前缀复用收益 训练关键路径紧迫度 数据多样性贡献 - 执行维护成本。每一项都可以配置权重权重随训练阶段动态调整。4.2 动态前缀索引与老化机制Prefix Locality 仍然是重要工具但需要做三件事第一前缀索引支持动态过期。RL 训练中策略更新后旧策略生成的前缀不再有复用价值应该从索引中淘汰。可以记录每个前缀块的“最近命中时间”和“策略版本号”当策略版本落后超过阈值时直接标记失效。第二前缀合并与剪枝。对于差异极小的 prompt可以对前缀树做自动合并减少碎片化。相似度计算可以用公共前缀长度、token 编辑距离等策略。第三分阶段调整复用权重。训练早期前缀复用率低调度器降低复用权重优先保证请求快速执行训练中后期策略收敛输出分布稳定逐步提高复用权重。4.3 混合请求分类与优先级调度器把请求分为三类Critical Request当前训练 step 依赖的 Rollout 请求必须尽快完成Reusable Request前缀命中率高、可以延迟执行换取缓存复用的请求Background Request回放数据、评估数据、日志生成等后台任务可以在资源空闲时执行。调度优先级按 Critical Reusable Background 处理同时配合时间片限制防止高优请求无限占用资源。4.4 两阶段调度资源预算分配 执行顺序编排一个大致的实现思路是两阶段阶段一预算分配。根据训练器当前状态是否在 waiting、下一个 step 需要多少样本决定当前调度窗口内推理最多消耗多少显存和算力哪些任务配额多少。阶段二执行顺序。在预算约束内对就绪请求排队按优先级、前缀复用率、预计执行时间进行排序选择当前批次的最优组合。这种两阶段解耦的好处是资源控制与缓存优化互不影响训练器侧只需要关心预算接口调度器内部可以自由调整执行策略。4.5 冷启动调度策略对应“rl 冷启动”问题调度器需要在训练初期采用不同参数调度参数冷启动阶段稳定训练阶段前缀复用权重低高批量请求大小小快速出结果大追求吞吐多样性采样权重高覆盖任务分布中按需调整缓存过期时限短旧前缀尽快清理长保留稳定前缀关键路径放松度高优先喂数据给训练器中平衡吞吐冷启动阶段的目标是“快速产生多样且有效的 Rollout 数据”让训练器尽快收敛到一个分布更稳定的策略上。等策略稳定后再切换到高复用、高吞吐模式。5. 调度器核心模块与实现路径这里给出一个简化的调度器模块划分与实现思路。注意下面代码是设计示例不是某个特定开源项目的完整实现落地时需要结合你的推理框架和训练框架接口调整。5.1 模块划分整体可以拆成五个核心模块请求接收器接收训练器、回放缓冲区、评估任务发来的请求统一封装为RLRequest。前缀索引器维护动态前缀树支持插入、命中、过期、剪枝。预算控制器从训练器获取资源状态计算调度窗口的 GPU 资源预算。优先级队列按分类和优先级组织就绪请求。执行器与推理引擎交互提交 batch收集结果。5.2 请求封装与优先级打分代码示例# 请求封装示例 import time from dataclasses import dataclass, field dataclass class RLRequest: req_id: str prompt_tokens: list task_type: str # generation / reward / state_query / replay max_tokens: int arrival_ts: float field(default_factorytime.time) is_critical: bool False # 当前训练 step 是否依赖 policy_version: int 0 # 生成该请求的策略版本 prefix_hit_len: int 0 # 调度时动态填充 python # 优先级打分示例 def score_request(req: RLRequest, weights: dict) - float: prefix_score req.prefix_hit_len / max(len(req.prompt_tokens), 1) critical_score 1.0 if req.is_critical else 0.0 waiting_score min((time.time() - req.arrival_ts) / 60.0, 1.0) score ( weights[prefix] * prefix_score weights[critical] * critical_score weights[waiting] * waiting_score ) return score5.3 前缀索引示例实际前缀树实现比较复杂这里给一个基于字典的简化版本便于理解核心逻辑# 简化前缀索引示例 class PrefixIndex: def __init__(self, expiry_version_threshold: int 3): self.trie {} self.expiry_threshold expiry_version_threshold def insert(self, token_ids: list, policy_version: int): node self.trie for tid in token_ids: if tid not in node: node[tid] {children: {}, version: policy_version, hits: 0} node node[tid][children] node[_end] {version: policy_version, hits: 0} def match(self, token_ids: list, current_version: int) - int: node self.trie hit_len 0 for tid in token_ids: if tid not in node: break node node[tid][children] if node.get(_end) is None: # 非完整前缀按需判断是否记录 hit_len 1 elif current_version - node[_end][version] self.expiry_threshold: node[_end][hits] 1 hit_len 1 else: break return hit_len生产环境中建议使用 Radix Tree 或类似 SGLang RadixAttention 的 Prefix Cache并加上并发锁和异步清理策略。5.4 调度主循环示例# 调度主循环伪代码 def scheduling_loop(ready_queue, prefix_index, budget_controller, executor): while True: requests ready_queue.get_ready_batch(max_batchbudget_controller.current_batch_size()) for req in requests: req.prefix_hit_len prefix_index.match(req.prompt_tokens, req.policy_version) requests.sort(keylambda r: score_request(r, current_weights()), reverseTrue) selected budget_controller.filter_by_resource_budget(requests) if selected: executor.submit(selected) for req in selected: prefix_index.insert(req.prompt_tokens, req.policy_version) budget_controller.release_waiting_resources() time.sleep(0.01)核心要点是排序前先做前缀匹配排序混合多因子提交前再做资源预算过滤。这样可以在不牺牲资源控制的前提下吸收 Prefix Locality 的部分收益。6. 接口与配置设计调度器通常通过 gRPC 或 HTTP 与训练器交互。下面给出一个通用的配置和请求示例落地时按项目调整。6.1 调度器配置示例# scheduler_config.yaml 示例 scheduler: max_batch_size: 32 scheduling_interval_ms: 10 queue_capacity: 4096 weights: prefix: 0.4 critical: 0.4 waiting: 0.2 policy: cold_start_steps: 200 # 训练前 200 步视为冷启动 cold_start_prefix_weight: 0.1 # 冷启动时降低前缀复用权重 stable_prefix_weight: 0.4 # 稳定阶段恢复 prefix_ttl_steps: 5 # 前缀块策略版本过期阈值 budget: max_gpu_memory_mb: 40960 training_reserved_mb: 20480 # 为训练保留的显存 max_inference_mb: 20480 # 推理可用显存上限 executor: engine: vllm # 也可以是 sglang / tensorrt-llm timeout_seconds: 1206.2 训练器提交 Rollout 请求示例import requests payload { task_type: generation, prompt: Solve the following math problem step by step., max_tokens: 512, is_critical: True, policy_version: 12, } resp requests.post(http://127.0.0.1:8900/submit_rollout, jsonpayload, timeout10) print(resp.json())6.3 调度器返回结果示例{ req_id: rl_20250101_00001, status: accepted, estimated_batch: batch_20250101_0003, queue_position: 2, prefix_match_length: 128 }实际生产环境建议用 gRPC 流式接口支持长连接和批量结果返回避免 HTTP 短连接在大规模 RL 训练下成为瓶颈。7. 评测方法与性能观察超越 Prefix Locality 的效果不能只看缓存命中率需要一套针对 RL 训练场景的评测指标。7.1 核心指标指标含义目标Sample Throughput每秒产出有效 Rollout 样本数越高越好GPU Idle Ratio训练器等待 Rollout 数据的时间占比越低越好Prefix Hit RatioKV Cache 前缀命中比例辅助指标不是唯一目标Training Step Time单次训练 step 总耗时越低越好Diversity Score生成样本的分布覆盖度冷启动阶段重点观察Cache Memory Overhead前缀索引和缓存占用额外内存控制合理范围7.2 评测方法推荐做 A/B 对比Baseline 1纯 Prefix Locality 调度按前缀命中优先Baseline 2纯 FIFO 调度不做特殊优化Test混合调度前缀复用 关键路径 冷启动策略。每个方案固定相同的训练数据、模型和训练超参数跑相同的训练步数对比总耗时、样本吞吐、GPU 空闲时间和最终模型效果。评测脚本示例# 评测调度策略对比示例 import time import statistics def run_training_steps(scheduler_type, steps100): start time.time() sample_throughputs [] for step in range(steps): t0 time.time() rollout_data scheduler_collect(scheduler_type, step) train_one_step(rollout_data) cost time.time() - t0 sample_throughputs.append(len(rollout_data) / cost) total_time time.time() - start return { total_time: total_time, avg_throughput: statistics.mean(sample_throughputs), gpu_idle_ratio: compute_gpu_idle_ratio(), prefix_hit_ratio: compute_prefix_hit_ratio(), }7.3 性能观察重点实际运行时要重点观察三个位置第一训练器侧等待日志。如果训练 step 的日志里频繁出现waiting for rollout data说明调度器没有优先保障关键路径请求。第二GPU 显存分配趋势。训练和推理的显存分配是否在交替攀升还是长期互相挤压。可以通过nvidia-smi定时采样观察显存水位。第三前缀索引的内存开销。当请求量达到百万级时前缀索引本身可能占用数 GB 内存。如果发现索引内存异常增长需要检查过期清理策略是否生效。8. 资源占用与性能权衡8.1 显存开销前缀 Cache 是提升复用率的重要手段但会占用显存。前缀索引本身存放在 CPU 内存时还需要考虑 CPU 内存和 GPU 显存之间的搬运开销。经验上生产环境会限制前缀 Cache 的最大显存容量。比如设置一个max_cache_gpu_mb超过后触发 LRU 淘汰。调度器应该把缓存淘汰策略暴露为配置项而不是写死。8.2 CPU 调度开销调度器本身是一个高频决策模块。每 10ms 做一次批次决策时如果请求队列非常大排序成本不可忽视。优化方向使用堆队列代替全量排序对相同前缀的请求预先分组减少排序规模调度决策异步化不阻塞请求接收。8.3 长请求与短请求的互相影响一个长生成请求如果占满了执行器的 batch 窗口后台短请求就会堆积。一个常见做法是时间片轮转执行器每执行 N 个长请求后强制插入一批短请求。这比纯粹按优先级排序更稳定。8.4 权衡总结权衡点偏向 Prefix 复用偏向训练吞吐长处节省重复计算提升推理效率保证训练器持续有数据风险训练器等待、多样性下降计算重叠增加、KV 缓存利用率下降适用阶段策略稳定、前缀稳定冷启动、策略频繁更新、任务多样调度器的参数不应该是一次性固定的而应该随训练状态动态调整。例如用训练步数、策略更新频率、最近 N 步 Rollout 数据的平均熵作为状态信号。9. 常见问题与排查方法问题现象可能原因排查方式解决方案训练 step 频繁等待 Rollout 数据关键路径请求优先级不足查看调度日志中 critical 请求的排队时间提高 critical 权重或为关键路径单独开队列前缀命中率很高但训练速度反而下降过度追求复用牺牲了样本多样性对比命中率和样本多样性评分降低 prefix 权重增加多样性采样冷启动阶段 Rollout 输出无效样本多策略未收敛调度器未切换冷启动模式观察策略熵变化和输出分布冷启动阶段降低复用权重、缩短缓存过期时间前缀索引占内存过大过期前缀未及时清理检查策略版本淘汰日志缩短 prefix_ttl_steps定期合并前缀树训练与推理显存互相挤压资源预算没有生效查看 budget 配置和显存采样曲线设置训练保留显存控制推理最大显存长请求阻塞短请求批次调度没有时间片轮转查看短请求的排队时延增加时间片轮转策略定期强制插入短请求GPU 利用率波动大调度窗口和训练步调不匹配对比调度间隔与训练 step 时间调整 scheduling_interval_ms对齐训练步调回放数据与在线数据比例失衡后台请求优先级过高查看不同任务类型样本占比降低 Background 请求权重限制回放请求配额10. 最佳实践与落地建议第一先跑通最小闭环再上规模。建议先用一个小模型、小显存环境把“训练器 - 调度器 - 推理引擎 - 结果回填”这条链路跑通验证调度器不会成为新的瓶颈再逐步扩大请求量和 batch size。第二调度参数不要拍脑袋。每次调整 Prefix 权重、Critical 权重后至少跑 50~100 个训练 step对比样本吞吐和训练总耗时。第三指标日志要打全。训练器侧至少记录每个 step 的等待时间调度器侧记录请求排队时长、每个 batch 的执行时间、前缀命中分布执行器侧记录 KV Cache 命中率、显存占用峰值。三者拼起来才能定位问题。第四冷启动和稳定阶段分开设计。不要在训练一开始就开满所有优化容易让调度器在异常分布上自我消耗。建议阶段切换使用训练步数和策略熵双阈值触发。第五合规与数据安全。RL 训练中涉及的用户数据、生成内容、奖励模型标注数据要确保来源合法、授权清晰。涉及人脸、声音、版权素材的生成任务必须确认授权边界模型进入生产或商用前要对逻辑漏洞、偏见和有害内容做额外评估。调度器相关实验建议在隔离的开发环境进行并保留完整的训练与推理日志方便追溯。第六注意开源许可。引用或集成 vLLM、SGLang、Ray、训练框架等三方组件时检查各自的 License 是否满足你的分发要求。不要因为调度器的代码只是薄薄一层就忽略底层框架的许可约束。11. 总结与下一步这个方向最值得尝试的点在于把调度器从“KV Cache 的工具人”升级为 RL 训练管道的核心组件。具体来说就是先定义清楚训练目标再把前缀复用、关键路径、数据多样性、冷启动策略加权融合到一个可配置的调度器里。它解决的问题是真实存在的大规模 RL 训练中Rollout 请求混杂、前缀漂移、训练推理争资源这些问题只靠传统 Prefix Locality 无法根除。最先应该验证的功能是两阶段调度中的预算控制。因为如果训练和推理的资源边界划不清楚后面所有优化都会建立在沙地上。最容易踩的坑是照着在线推理服务的思路做调度把缓存命中率当作 KPI忽略训练器的实时状态。结果是缓存看起来很漂亮训练吞吐却上不去。后续可以继续扩展的方向包括基于强化学习的自适应调度参数调整、多机多卡环境下的跨节点前缀共享、以及调度器与训练器联合编排的端到端 pipeline。这里的关键是“先设计训练感知的调度接口再逐步加缓存优化”不要把顺序搞反了。