ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

混合RL Rollout调度:超越Prefix Locality的推理优化实践

混合RL Rollout调度:超越Prefix Locality的推理优化实践 关于“Scheduling Mixed RL Rollouts Beyond Prefix Locality”这个方向很多人第一次看到会把它当成一篇纯推理优化论文实际上它卡在 RL 训练和 LLM 推理的交叉点上核心矛盾非常具体RL 训练每轮都要用当前策略模型生成大量 rollout也就是采样轨迹这些请求既有一大段完全相同的共享前缀又有大量采样分支、长度、优先级完全不同的内容。如果调度器只会盯着“前缀相同就复用一个 prefix”很快会发现训练 step 在傻等、KV Cache 被挤爆、GPU 利用率忽高忽低。这篇文章不打算给你一个虚构的“一键启动包”而是把这套调度思路拆开讲清楚Prefix Locality 到底复用什么为什么只做前缀复用不够混合 RL Rollout 调度需要权衡哪些指标以及怎么用现有推理引擎搭一个最小实验去验证效果。适合做 RL 训练框架、推理引擎优化、以及大模型训练推理一体化平台的工程师阅读。先交代本文覆盖的内容第一RL Rollout 调度的问题定位第二Prefix Locality 的原理和收益边界第三Beyond Prefix Locality 的调度目标和关键指标第四一套可复现的验证环境和测试流程第五资源占用观察、问题排查和最佳实践。全程不给编造数据遇到需要以实际环境为准的地方会明确标出。1. 核心问题速览这个方向不是某个可以直接 clone 的开源仓库而是一个调度机制设计方向通常以系统论文或推理引擎特性形式出现。它解决的核心问题是在 RL 训练过程中如何高效调度推理侧的大量 rollout 请求让采样生成既不浪费共享前缀的 KV Cache 计算又不会被“强行凑前缀”拖垮整体吞吐。能力项说明项目类型LLM 推理调度与 RL 训练系统的交叉方向核心机制Prefix Locality、KV Cache 复用、Rollout 批量调度主要场景RLHF / PPO / GRPO / 代码 RL / 数学推理 RL 的在线采样输入特征少量共享前缀 大量采样分支 变长输出 不同优先级调度目标提高 token 吞吐降低训练 step 等待时间控制显存占用硬件门槛GPU 集群为主单卡可做小规模模拟验证显存占用取决于模型规模、KV Cache 策略和并发数需实测可复用组件vLLM 的 prefix caching、SGLang 的 RadixAttention、continuous batchingAPI 能力通常复用推理引擎的 schedule 接口或自定义采样服务批量任务是核心场景RL rollout 天然是高并发批量采样适合读者RL 训练框架开发、推理引擎优化、MLOps 工程师从材料看本文的关键词集中在 Scheduling、RL、Rollouts、Prefix Locality因此下面的分析都围绕这四个词展开不涉及无关特性。2. RL Rollout 调度为什么是训练效率的瓶颈先把概念对齐。RL 训练的主流做法是在每个训练迭代中做两步第一步是 rollout用当前策略模型对一批 prompt 做多次采样生成若干条完整轨迹第二步是 update用这些轨迹计算策略梯度并更新模型参数。PPO、GRPO、REINFORCE 等算法虽然在更新公式上不同但 rollout 阶段的结构高度相似同一个 prompt 往往要求生成 4、8、16 条甚至更多不同采样结果。Rollout 为什么是瓶颈从三个角度看就很清楚。第一个角度是算力比例。Rollout 是纯推理过程需要跑大量 decode step而 update 阶段虽然要跑前向和反向但 token 总量通常远小于 rollout 生成量。常见经验是 rollout 的计算量在训练总计算量中占比很高尤其是输出越长、采样数越多推理侧压力越明显。所以 rollout 调度做得好不好直接决定训练整体吞吐。第二个角度是对实时性的要求。RL 训练是一个在线过程策略模型每更新一次旧的 rollout 数据就不能再用于下一轮更新。这意味着 rollout 任务有明确的截止时间意识——训练 step 在等 rollout 结果rollout 迟迟不返回GPU 上训练进程可能空转。推理侧如果为了提升某个局部指标把请求排很久训练整体效率反而下降。第三个角度是请求结构的复杂性。RL 场景的 rollout 请求不是传统在线推理那种“用户来一个请求我尽快回一个请求”的模式。它通常是一个 batch 里有几十上百个不同 prompt每个 prompt 又要生成多条采样结果每条结果长度不同有的 prompt 还带着很长的系统提示词或对话历史。这种请求结构天然混合调度器不能假设所有请求共享一个前缀也不能假设每条请求完全独立。这个方向的适用者很明确正在做 RLHF 或 RLAIF 训练发现 rollout 阶段 GPU 利用率不高的人。使用 PPO / GRPO 训练代码模型、数学模型需要大量在线采样的人。自己维护推理服务想把 vLLM、SGLang 等引擎接入训练链路的人。做推理引擎内核优化想理解 prefill、decode、前缀复用如何影响调度的人。同时也要说清边界如果你的场景是纯离线生成数据、对时延不敏感前缀复用的收益虽然还在但调度复杂度可以大幅降低如果你的场景是普通在线对话服务请求之间几乎没有共享前缀那这个方向的红利就不明显。它真正的适用场景是“共享前缀 多分支采样 在线持续迭代”的 RL 训练链路。3. 前置知识Prefix Locality 到底复用了什么Prefix Locality直译是前缀局部性。要理解它先理解 Transformer 推理的 prefill 和 decode 两阶段。在生成第一个 token 之前模型需要把输入序列中的所有 token 都做一次前向计算这一步叫 prefill对每个输入 token 计算 Key 和 Value并写入 KV Cache。之后每生成一个新 token只需要做一次 decode用一个 query 和 KV Cache 里所有 Key 做 attention。所以 KV Cache 是推理性能的关键。输入前缀越长prefill 的计算量越大KV Cache 占用的显存也越大。如果两个请求的输入前缀完全相同例如系统提示词都是“Solve the following math problem step by step:”那么第一个请求在 prefill 阶段算出来的这些 token 的 KV Cache第二个请求可以原样复用不需要重新计算。这就是 prefix locality 的收益来源——按 token 数量看能省掉一大块重复的 prefill 计算也能显著降低首 token 时延。在 RL rollout 场景里这个特性被放大了。训练集里同一个 prompt 要生成多条采样结果这 N 个请求共享同一个 prompt 前缀只有采样分支不同。如果调度器把同一 prompt 的 N 个请求分到同一批并且引擎支持前缀感知的 KV Cache 复用那么整段 prompt 只需要计算一次 prefillN 条请求共享同一段 KV Cache各自从不同位置开始 decode。除了同一个 prompt 的多路采样RL 训练里还有另一种前缀复用机会多个不同训练 prompt 可能包含相同的系统提示词、相同的 few-shot 示例、相同的数据集描述这些内容作为公共前缀时也能被批量复用。在工程实现层面前缀复用并不是简单“缓存一份计算结果”就行它需要满足几个前提模型推理支持前缀缓存通常要求使用支持 PagedAttention 或类似分页 KV Cache 机制的引擎。调度器能识别请求之间的共享前缀而不是把每个请求当成独立序列。引擎的显存管理器支持对 KV Cache 块的引用计数和淘汰避免缓存占用过多导致 decode 空间不足。采样参数不会影响 prefill 阶段的 KV Cache 计算也就是说同一前缀下的不同采样温度、top-p 等参数可以安全共享前缀。这里有一个容易混淆的点prefix locality 说的是“共享前缀的 KV Cache 可以被复用”不等于“只要前缀相同调度就一定优先把它们放一起”。后者是调度策略前者是硬件和引擎能力。很多实现把这两件事当成同一件事就会在后续混合负载上吃大亏。4. 为什么不能只依赖 Prefix Locality题目里的 “Beyond Prefix Locality” 是整篇文章最值得琢磨的词。它暗示了一个结论单纯依赖前缀局部性做调度在真实 RL 负载下不够甚至会有反效果。原因要从负载结构说起。RL 训练的 rollout 请求通常不是单一前缀树而是混合负载。一个典型训练 step 的采样 batch 长这样有 50 个不同 prompt每个 prompt 不共享前缀每个 prompt 要采样 8 条结果这 8 条共享同一个 prompt 前缀部分 prompt 带很长的历史上下文部分 prompt 是短问题每条采样的 max_tokens 不同有的要求 128有的要求 1024训练进度约束下整个 batch 必须在某个时间窗口内返回否则训练侧空转。如果把“最大化前缀复用”作为唯一调度目标会出现三种问题。第一种是排队阻塞。调度器为了把共享同一前缀的请求凑到一起会让早到的请求在队列里等待后来的兄弟请求。等待时间一长训练 step 那边已经等不及了空转的 GPU 代价远超省下的 prefill 计算。前缀复用的收益是确定的但等待成本是累加的trade-off 必须显式建模。第二种是显存挤占。前缀缓存命中越多需要驻留在显存里的 KV Cache 块就越多。如果同时缓存大量不同 prompt 的前缀显存被 cache 吃满活跃 decode 请求的可用 KV Cache 变少可能触发频繁驱逐驱逐后又要重新 prefill反而形成“缓存抖动”。系统整体吞吐不升反降。第三种是优先级和保序需求被忽略。RL 训练里 rollout 请求不是同等重要的有的 prompt 是当前迭代的关键数据有的 batch 必须在 update 前完成。只按前缀相似度分组会打乱任务的完成顺序导致训练 step 等一个不重要的长尾请求。所以 Beyond Prefix Locality 的含义是调度器要把前缀复用当成一个优化维度而不是唯一维度。它需要同时考虑前缀复用收益省多少 prefill 计算省多少 TTFT等待代价为了凑前缀请求多等了多久训练 step 空转多久显存代价缓存占了多少空间驱逐概率多大优先级约束哪些 rollout 必须在截止时间前完成设备利用率每个 GPU 上的 prefill/decode 负载是否均衡。只有把这些目标放到同一个调度决策里才是真正意义上的混合 RL Rollout 调度。5. 调度建模与关键指标要把调度问题讲清楚可以先建一个简化模型再逐步加约束。假设一个训练 step 产生一组 rollout 请求每个请求可以表示为dataclass class RolloutRequest: request_id: int prompt_id: int # 同一个 prompt_id 表示共享 prompt 前缀 input_tokens: int # prompt 长度 max_tokens: int # 本次采样最大输出长度 num_samples: int # 采样分支数可拆成多个独立请求 priority: float # 训练侧给定的优先级/截止时间 arrival_time: float # 到达调度器的时间调度器的任务是把这些请求切分成若干个 batch分配到不同 GPU 上去执行同时决定哪些请求可以共享已经缓存的 prefix。目标函数不是单一吞吐而是类似下面的联合目标maximize rollout_token_throughput minimize max(step_wait_time) subject to KV_cache_memory GPU_memory_budget per_request_deadline_slo 某些高优请求的截止时间在这个模型下纯 prefix locality 的调度策略等价于“优先按 prompt_id 分组组内共享 prefill”而 beyond prefix locality 的调度策略则会考虑如果同一 prompt 的采样分支还没到齐是否先用当前已到请求启动 decode不同 prompt 之间能否混排以填满 prefill 引擎的空隙当缓存命中收益不足以抵消等待成本时是否主动放弃复用。关键指标可以分成三类。第一类是推理性能指标Prefill 计算量 / 总计算量衡量前缀复用省了多少重复 prefill。前缀缓存命中率衡量有多少输入 token 直接命中缓存。Decode 吞吐单位时间生成的采样 token 数。TTFT首 token 时延和 TPOT每输出 token 时延。第二类是训练侧指标Step 等待时间即训练进程等待 rollout 返回的时间。每 train step 的墙钟时间这是最终收益的衡量标准。Rollout 数据到达率单位时间能产出的有效训练样本数。第三类是资源指标KV Cache 占用和命中缓存块数量。缓存驱逐次数。GPU 显存峰值。每个 GPU 的利用率曲线。一个值得强调的点是单独看任何一项指标都可能误判。缓存命中率提高 20%如果 step 等待时间翻倍整体训练效率反而下降。所以验证这个方向时一定要把“训练 step 墙钟时间”作为第一结果指标而不是只盯推理侧吞吐。6. 验证环境用推理引擎搭最小可复现实验这个方向没有统一的开箱即用项目但可以用现有推理引擎复现它的核心问题。下面给出一套通用的最小验证环境方案不绑定具体版本实际使用时需要根据本机环境调整。6.1 基础组件建议准备以下组件一个支持前缀缓存的推理引擎例如 vLLM 或 SGLang。vLLM 有开启动态前缀缓存的参数SGLang 的 RadixAttention 天然支持前缀树复用。一个支持 RL 训练的框架作为 rollout 客户端例如直接写 Python 脚本调用引擎的采样接口不必引入完整 RL 框架。一套混合负载构造脚本模拟“多个不同 prompt 每 prompt 多路采样 不同输出长度”的请求分布。监控工具包括nvidia-smi、引擎自带 metrics 端口、Python 侧的日志统计。6.2 负载构造思路先构造一个模拟 RL rollout batch 的请求文件关键点是让请求同时包含共享前缀和独立前缀。例如取 20 个不同 prompt每个 prompt 前都拼接同一段系统提示词然后每个 prompt 请求 8 条采样。这样系统提示词部分是公共前缀prompt 部分是独立前缀采样分支是共享 prompt 前缀的后缀。6.3 启动推理引擎以 vLLM 风格为例启动参数可以这样组织实际参数名以你安装的版本为准# 启动一个兼容 OpenAI 采样接口的本地推理服务 python -m vllm.entrypoints.openai.api_server \ --model /path/to/your/policy/model \ --tensor-parallel-size 1 \ --gpu-memory-utilization 0.85 \ --max-num-seqs 256 \ --enable-prefix-cachingSGLang 风格则通常是# 启动 SGLang 推理服务 python -m sglang.launch_server \ --model-path /path/to/your/policy/model \ --host 127.0.0.1 \ --port 30000 \ --mem-fraction-static 0.8上面命令是通用模板具体模型路径、端口、并行度必须按实际环境修改。多数引擎会在启动日志里打印是否启用 prefix caching这一步建议先确认“已经开启”。6.4 最小客户端脚本下面是一个通用 Python 客户端模板用 OpenAI 兼容接口向推理服务发送多路采样请求。接口地址和参数需要按实际部署调整import time from concurrent.futures import ThreadPoolExecutor import requests API_URL http://127.0.0.1:8000/v1/completions MODEL /path/to/your/policy/model # 构造一个模拟 RL rollout 的混合负载 base_system_prompt Solve the following problem step by step. prompts [ base_system_prompt fProblem {i}: compute the value of expression {i}*{i3}. for i in range(20) ] def sample_one(prompt: str, idx: int) - None: payload { model: MODEL, prompt: prompt, max_tokens: 256, temperature: 0.8, n: 4, # 每个 prompt 采样 4 条 } t0 time.time() resp requests.post(API_URL, jsonpayload, timeout180) dt time.time() - t0 print(f[{idx}] prompt_len{len(prompt.split())} tokens, time{dt:.2f}s, status{resp.status_code}) # 并发发送模拟训练侧一次性提交整个 rollout batch with ThreadPoolExecutor(max_workers8) as pool: for i, p in enumerate(prompts): pool.submit(sample_one, p, i)这段脚本不是论文的复现代码而是帮助你把“混合 rollout 请求”落到真实引擎上观察开启前缀缓存前后的变化。7. 功能测试与效果验证部署好环境后按照下面的实验矩阵逐项验证。建议每个实验跑 2 到 3 次取稳定值避免单次波动影响判断。7.1 验证前缀缓存生效测试目的确认引擎确实复用了共享前缀而不是每次都重新计算 prefill。操作步骤启动服务确认日志中 prefix caching 已开启。连续发送两个请求请求前缀完全相同。观察第二次请求的 TTFT 是否明显降低。判断标准第二次请求的首 token 时延显著低于第一次或者引擎的缓存命中计数增加。如果两次 TTFT 几乎一样说明前缀缓存没有生效需要检查参数和引擎版本。7.2 验证同 prompt 多路采样共享前缀测试目的模拟 RL rollout 最常见的请求结构即同一 prompt 生成多条采样分支。操作步骤选择一个 prompt设置n8一次请求 8 条采样结果。再用单条请求n1的方式发 8 次同样 prompt观察结果差异。对比两种方式的吞吐和显存占用。预期结果在支持前缀缓存的引擎中设置n8会比拆成 8 次请求节省大量重复 prefill显存占用增长也更平缓。这也验证了“同一 prompt 多路采样”是该调度问题的天然受益场景。7.3 验证混合前缀负载下的调度表现测试目的构造“不同 prompt 不同长度 不同采样数”的混合负载观察前缀缓存是否仍然稳定收益。操作步骤使用上文 6.2 节的负载构造思路生成混合请求文件。分别在开启和关闭前缀缓存的设置下跑同一份负载。记录总耗时、缓存命中率、显存峰值。判断标准观察总耗时是否下降。如果下降不明显说明前缀相似度太低或被其他瓶颈掩盖。此时可以进一步分析请求的 token 组成看看公共前缀占比多大。这部分需要实测数据来判断。7.4 验证调度等待成本测试目的验证“为了凑前缀而等待兄弟请求”是否会拖慢整体任务。操作步骤构造一批同一 prompt 的多路采样请求但让它们分不同时间点到达。用两种策略对比一种等到同一 prompt 的请求到齐后再统一执行另一种是到齐多少就执行多少。对比整体完成时间。预期结果当请求到达时间分散时“等齐再执行”通常完成时间更长。这正是 Beyond Prefix Locality 要解决的问题——前缀复用收益和排队等待成本必须做权衡。7.5 功能测试结果记录建议用下面的表格记录每轮实验结果实验编号开关状态请求数量总耗时缓存命中率显存峰值结论7.1开2---确认缓存生效7.2开/关1x8 vs 8x1---对比多路采样优势7.3开/关混合 80 条---对比混合负载收益7.4策略 A/B分批到达---对比等待成本这些实验做完基本就能判断这个调度方向在你的负载下值不值得继续投入。8. 资源占用与性能观察混合 RL Rollout 调度里资源占用主要看三块KV Cache、显存总量、GPU 利用率。8.1 显存占用观察显存占用最直接的影响因素有三个模型权重本身。前缀缓存驻留的 KV Cache 块。活跃请求 decode 过程中新增的 KV Cache。在 RL 训练场景中KV Cache 往往比模型权重更吃显存。长 prompt、多路采样、高并发这三个条件叠加KV Cache 可能占到显存的大半。观察方式是在请求运行期间持续采样nvidia-smi# 每 2 秒记录一次显存和利用率 nvidia-smi --query-gpuindex,memory.used,utilization.gpu,power.draw \ --formatcsv -l 2 gpu_monitor.log如果过程中出现显存溢出或缓存驱逐日志里会出现大量重新 prefill 的迹象表现为 TTFT 突增。8.2 吞吐和利用率观察推理侧要看两个阶段的比例prefill 阶段计算密度高但 token 输入是一次性的decode 阶段单步计算量小但步数多。理想的调度是 prefill 和 decode 尽量重叠不让 GPU 在 prefill 结束后的间隙空转。混合 rollout 调度的难点就在于如果把相同前缀的请求聚在一起做一次大 prefill虽然省了重复计算但 prefill 阶段结束后所有请求同时进入 decode如果并发数过高KV Cache 会快速膨胀。反之如果把请求分散到不同时间点前缀复用率下降但显存曲线更平滑。所以性能观察不能只看某一时刻的利用率要看整条时间线。建议记录以下序列每批请求的 prefill 耗时、decode 总耗时、KV Cache 增量、GPU 利用率曲线然后画在同一个坐标系里分析瓶颈是 prefill、decode 还是显存。8.3 如何降低显存压力给几个通用方向具体效果以实际测试为准降低max-num-seqs限制并发 decode 请求数减少 KV Cache 峰值。调整gpu-memory-utilization给前缀缓存留出明确空间。使用 GQA 或 MQA 架构模型KV Cache 本身就会更小。关闭不需要的采样分支缓存缓存只保留真正会复用的前缀。对长 prompt 做分段只缓存高频复用的系统提示词部分。这些手段的目的不是“一劳永逸”而是让调度器在显存约束下有更多可操作空间。9. 常见问题与排查方法混合 RL Rollout 调度在实现和验证阶段有一些高频率问题整理成排查清单如下。问题现象可能原因排查方式解决方案前缀缓存命中率很低请求前缀本身不共享或引擎未开启 prefix caching检查启动日志统计请求 token 前缀分布构造公共系统提示词确认参数开启缓存命中但 TTFT 没有下降显存不足导致缓存被驱逐或引擎不支持跨请求复用观察缓存驱逐计数和显存曲线调整 gpu-memory-utilization降低并发同一 prompt 多路采样时吞吐不升反降采样分支同时 decode 导致 KV Cache 膨胀GPU 受限对比 n1 和 n8 的显存占用减少 n 值或拆分到多卡为了凑前缀训练 step 空转时间变长调度策略过度偏好前缀复用记录请求到达时间与完成时间改用“到多少做多少”的混合策略批量任务卡住不返回某个长 max_tokens 请求拖尾或并发线程耗尽检查服务端日志和请求超时配置设置请求级超时限制最大输出长度显存溢出 OOMKV Cache 峰值过高看启动参数和监控日志降低 max-num-seqs清理缓存CUDA / 引擎版本不匹配引擎与显卡驱动版本不兼容查看引擎启动报错日志按引擎官方文档匹配 CUDA 版本多卡场景负载不均前缀分布集中在某一台设备检查每卡利用率调整请求路由策略或手动切分前缀排查时优先看三层第一层是引擎日志启动参数有没有生效第二层是 GPU 监控显存和利用率在哪个阶段异常第三层是请求时间线是 prefill 慢、decode 慢还是排队慢。大部分问题都能在这三层里定位。10. 最佳实践与使用建议基于当前分析把这个方向落地到 RL 训练链路时有几个实操建议值得先落实。第一先做负载画像再优化。不要一上来就改调度器。先用现有引擎跑几轮真实训练 batch统计共享前缀占比、prompt 长度分布、采样分支数、max_tokens 分布。这部分数据决定了 Prefix Locality 是不是你场景的主要矛盾。第二把“训练 step 墙钟时间”作为唯一最终指标。调度器设计很容易陷入局部指标优化比如缓存命中率从 60% 提到 80%但如果 step 等待时间变长就没有意义。每次改动都要回归到训练侧总耗时。第三调度策略要加超时和放弃机制。如果同一前缀的兄弟请求迟迟不到就先用已到请求启动放弃这次复用机会。这个机制能避免“为了省一次 prefill让训练空转数十秒”的极端情况。第四显存预算要显式建模。调度器应该知道当前 KV Cache 缓存了多少块、还能分配多少块。当缓存空间不足时优先保留复用频率高的公共前缀淘汰低价值前缀。第五目录和日志管理要规范。RL 训练需要区分 prompt 输入、采样输出、模型权重、评估结果。建议按训练 step 组织目录每个 rollout batch 输出带时间戳的日志方便回溯是哪一步调度问题导致训练变慢。第六合规边界要提前确认。RL 训练通常涉及大量 prompt 数据和模型权重使用前需要确认数据来源合法、模型权重遵循对应开源协议、生成内容不包含违规信息。如果用于代码或数学推理训练还要留意训练数据的版权和敏感性。涉及内部业务数据时建议在受控环境部署推理服务不要暴露到公网。11. 总结与下一步这个方向最值得关注的点是把“前缀复用”从一个局部优化手段提升为整个 RL 训练链路的一个调度维度并且明确承认“复用前缀不是免费的”——等待、显存、优先级都是成本。只要你负责的 RL 训练任务存在“同 prompt 多路采样 在线迭代 训练步等结果”的组合就应该认真评估这套调度思路。建议最先验证三件事第一当前推理引擎是否启用前缀缓存同前缀请求的 TTFT 是否真的下降第二真实训练 batch 的请求结构里共享前缀占比到底有多高第三开启前缀缓存后训练 step 总耗时是升还是降。最容易踩的坑有两个一个是把缓存命中率当成终极指标忽略了训练步等待另一个是盲目增加缓存空间导致活跃请求的 KV Cache 不够用引发频繁驱逐。这两个坑在真实负载里几乎一定会出现。下一步可以从三个方向继续深入一是把调度策略改成可配置的“前缀复用 等待预算 显存预算”联合决策二是接入完整 RL 框架把 rollout 服务的接口和训练侧 step 同步起来三是在多机多卡环境下验证请求路由是否均衡。等你有了一组稳定的实验数据后就能判断是否值得在内部训练框架里实现一个专门的 rollout 调度器。
返回列表