混合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 计算，又不会被“强行凑前缀”拖垮整体吞吐。

从材料看，本文的关键词集中在 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 只需要计算一次 prefill，N 条请求共享同一段 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-caching

SGLang 风格则通常是：

# 启动 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 + f"Problem {i}: compute the value of expression {i}*{i+3}." 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, json=payload, timeout=180) 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_workers=8) as pool: for i, p in enumerate(prompts): pool.submit(sample_one, p, i)

这段脚本不是论文的复现代码，而是帮助你把“混合 rollout 请求”落到真实引擎上，观察开启前缀缓存前后的变化。

7. 功能测试与效果验证

部署好环境后，按照下面的实验矩阵逐项验证。建议每个实验跑 2 到 3 次，取稳定值，避免单次波动影响判断。

7.1 验证前缀缓存生效

测试目的：确认引擎确实复用了共享前缀，而不是每次都重新计算 prefill。

操作步骤：

1. 启动服务，确认日志中 prefix caching 已开启。
2. 连续发送两个请求，请求前缀完全相同。
3. 观察第二次请求的 TTFT 是否明显降低。

判断标准：第二次请求的首 token 时延显著低于第一次，或者引擎的缓存命中计数增加。如果两次 TTFT 几乎一样，说明前缀缓存没有生效，需要检查参数和引擎版本。

7.2 验证同 prompt 多路采样共享前缀

测试目的：模拟 RL rollout 最常见的请求结构，即同一 prompt 生成多条采样分支。

操作步骤：

1. 选择一个 prompt，设置n=8，一次请求 8 条采样结果。
2. 再用单条请求n=1的方式发 8 次同样 prompt，观察结果差异。
3. 对比两种方式的吞吐和显存占用。

预期结果：在支持前缀缓存的引擎中，设置n=8会比拆成 8 次请求节省大量重复 prefill，显存占用增长也更平缓。这也验证了“同一 prompt 多路采样”是该调度问题的天然受益场景。

7.3 验证混合前缀负载下的调度表现

测试目的：构造“不同 prompt + 不同长度 + 不同采样数”的混合负载，观察前缀缓存是否仍然稳定收益。

操作步骤：

1. 使用上文 6.2 节的负载构造思路，生成混合请求文件。
2. 分别在开启和关闭前缀缓存的设置下跑同一份负载。
3. 记录总耗时、缓存命中率、显存峰值。

判断标准：观察总耗时是否下降。如果下降不明显，说明前缀相似度太低或被其他瓶颈掩盖。此时可以进一步分析请求的 token 组成，看看公共前缀占比多大。这部分需要实测数据来判断。

7.4 验证调度等待成本

测试目的：验证“为了凑前缀而等待兄弟请求”是否会拖慢整体任务。

操作步骤：

1. 构造一批同一 prompt 的多路采样请求，但让它们分不同时间点到达。
2. 用两种策略对比：一种等到同一 prompt 的请求到齐后再统一执行；另一种是到齐多少就执行多少。
3. 对比整体完成时间。

预期结果：当请求到达时间分散时，“等齐再执行”通常完成时间更长。这正是 Beyond Prefix Locality 要解决的问题——前缀复用收益和排队等待成本必须做权衡。

7.5 功能测试结果记录

建议用下面的表格记录每轮实验结果：

这些实验做完，基本就能判断这个调度方向在你的负载下值不值得继续投入。

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-gpu=index,memory.used,utilization.gpu,power.draw \ --format=csv -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 调度在实现和验证阶段有一些高频率问题，整理成排查清单如下。

排查时优先看三层：第一层是引擎日志，启动参数有没有生效；第二层是 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 调度器。