【Task.24】使用 SGLang 原生组采样复用多模态公共前缀

[overloadedHenry]

问题背景

GRPO 通常需要针对同一个 prompt 生成多条回复。在 Qwen3-VL 训练配置中，每个 prompt/image group 包含 8 条 sample。

Relax 当前会把这 8 条 sample 作为 8 个独立 /generate 请求发送。即使客户端侧已经共享了多模态编码结果，SGLang 仍会收到重复的 prompt 和媒体输入，可能重复执行公共前缀 prefill，或者依赖跨请求 prefix cache 与路由策略来弥补重复计算。

当前请求形态：

一个 prompt/image
  -> 8 个独立 HTTP 请求
  -> 最多执行 8 次重复公共前缀 prefill
  -> 8 条 decode

SGLang 支持通过 sampling_params.n 进行原生并行采样。对于一个输入和 n=8，SGLang 会先缓存公共前缀，再将请求展开为 8 条 decode 分支。

建议方案

增加一条默认关闭的 Relax rollout 路径，将一个满足条件的 prompt group 合并为一个非流式 SGLang 请求：

payload = {
    "input_ids": prompt_ids,
    "sampling_params": {**sampling_params, "n": len(group)},
    "return_logprob": True,
    # 存在媒体时，使用 batch-size-one 的多模态字段
}

SGLang 返回结果列表后，Relax 按稳定顺序回填到 group 内各个 sample，并复用现有逐 sample 路径的 token、logprob、文本与状态处理逻辑。

排查过程与路由策略选择

实验没有一开始就直接把 round-robin 写成默认值，而是分三步定位：

1. 先跑 baseline：每个 rollout 仍然发送约 512 个 /generate 请求。四轮合计 2048 个请求，每卡请求数为 218–298，CV 约 8.7%，整体比较均衡。
2. 只开启 native n=8：每个 rollout 的请求数从约 512 降到约 64，四轮合计从 2048 降到 256，说明请求合并机制已经真正生效；但 response-token throughput 反而从 9,678.52 降到 8,848.82 token/s，下降 8.6%。
3. 检查每个 engine 的访问日志：256 个 native group 请求的分布为 24, 24, 26, 27, 29, 29, 30, 67，CV 约 41.9%。单个 engine 承担了明显过多的 n=8 decode group，成为整轮完成时间的长尾。

这个现象说明问题不是 native prefix 复用没有发生，而是请求粒度发生变化后，原路由策略不再适合：

• Baseline 的一个 HTTP 请求只对应一条 decode，且请求数多，路由误差容易被后续请求摊平。
• Native 的一个 HTTP 请求代表 8 条 decode，请求数只有原来的 1/8。一次路由偏差会同时把整组 8 条 decode 压到同一 engine。
• 每个 group 的 prompt/image 不同，而 group 内共享前缀已经在单个 native 请求内部完成，因此跨请求 cache affinity 的价值下降。
• 从实验现象推断，cache-aware 的请求级 cache/load 信号没有充分反映 n=8 请求背后的 decode 工作量，导致少量、粗粒度 group 请求集中到同一 engine。

因此优化顺序是：

1. 保留 n=8，不放弃一次 prefill、8 条 decode 的核心收益。
2. 先用 round-robin 将完整 group 均匀分配给各 engine。
3. 如果 round-robin 后仍有明显的 group 内长度长尾，再考虑将一个 group 拆成两个 n=4 请求；本次实验中 round-robin 已经足够，因此没有继续拆分。

切换 round-robin 后，每卡请求数收敛到 31–33，CV 降到约 1.6%，native 的 prefix 复用收益才真正转化为墙钟和吞吐收益。

实现范围

• 新增 --sglang-native-group-sampling，默认关闭。
• 仅对同 prompt、同多模态对象、first-turn、训练阶段的 group-RM sample 启用 native 路径。
• 不满足条件时继续使用现有逐 sample 路径。
• 支持纯文本、单图和多图输入。
• 保留适用的 routing header。
• native 开启时自动启用 --group-rm。
• native 开启时默认使用 round_robin，允许通过 SGLANG_ROUTER_POLICY 显式覆盖。
• 严格验证返回类型与数量，不静默丢弃 sample。

非目标

• Evaluation 或确定性推理
• Partial/multi-turn rollout
• OPD
• Routing replay
• Slime router middleware
• LoRA adapter routing
• Custom generate function
• 修改 reward、advantage、actor train 或 update weights 逻辑

初步实验结果

配置：Qwen3-VL-4B-Instruct，64 prompts × 8 responses，8 张 H800。

方案	平均 response length	Response 吞吐	生成墙钟时间
Baseline	1449.26	9,678.52 token/s	76.7 s
Native n=8 + cache-aware/default	1428.72	8,848.82 token/s	82.7 s
Native n=8 + round-robin	1447.90	11,064.55 token/s	67.0 s

在平均 response length 只相差 0.094% 的情况下，response-token throughput 提升 14.3%。

Native 配合默认/cache-aware 路由时，生成时间比 baseline 增加 7.8%，response-token throughput 下降 8.6%；切换 round-robin 后，生成时间比 baseline 降低 12.6%，response-token throughput 提升 14.3%。因此均衡路由是方案的组成部分，而不是可选调优项。