Batch SGLang rollout generation per prompt group

[supercharleszhu]

TLDR

Miles' SGLang rollout path currently sends one /generate request per sampled response in common GRPO-style workloads. For prompt groups with multiple samples, this creates avoidable HTTP/router/tokenizer/scheduler-dispatch overhead even though SGLang supports batched /generate inputs.

Proposed Fix

#1568

This PR proposed batching the default SGLang rollout path at the GRPO group boundary. Instead of issuing one /generate request per sample in a prompt group, Miles now sends one batched /generate request for the whole group when it is safe to do so.

Why this speeds things up

GRPO commonly uses multiple samples per prompt, e.g. n_samples_per_prompt=4. The old path created one async request per sample, so a rollout with rollout_batch_size=256 and n=4 submitted roughly 1024 /generate requests through the SGLang router. The new path submits one request per prompt group, reducing that to roughly 256 requests while preserving the same generated sample count and max output token cap.

In profiling on an 8xH200 Miles FSDP smoke with no-op reward and the same prompt/response shape (rollout_batch_size=256, n=4, max_new_tokens=256):

Path	Rollout time	Tokens/GPU/s
Per-sample requests	~28.1-30.7s	~995-1085
Batched per prompt group	13.9s	~2202

The output cap remained correct (response_len/max=256). Continuous batching was already enabled in SGLang; the speedup came from reducing Python/router/request overhead and letting each prompt group enter the engine as one batched request.

I also tried increasing SGLang memory, enabling CUDA graph, and raising max_prefill_tokens/chunked_prefill_size to 131072. Those did not materially improve the per-sample path, which points to request granularity rather than KV-cache capacity or max batched token limits as the main bottleneck.

OSS reproduction for upstream issue

The slowness can be reproduced without any private data or Miles-specific training stack by comparing two equivalent SGLang /generate workloads:

• sample-by-sample mode: 256 prompt groups x 4 samples = 1024 independent HTTP requests
• group-batched mode: 256 prompt groups x 1 batched request with 4 copies of the prompt = 256 HTTP requests

Start an OSS SGLang server with any public instruction model that fits your GPU, for example:

python -m sglang.launch_server \
  --model-path Qwen/Qwen2.5-7B-Instruct \
  --host 0.0.0.0 \
  --port 30000 \
  --tp-size 1

Then run this benchmark from another shell:

import argparse
import asyncio
import time

import aiohttp

PROMPT = "Write a concise summary of why request batching can improve LLM serving throughput."


def make_payload(prompt, max_new_tokens):
    return {
        "text": prompt,
        "sampling_params": {
            "max_new_tokens": max_new_tokens,
            "temperature": 1.0,
            "top_p": 1.0,
        },
    }


async def post_generate(session, url, payload):
    async with session.post(url, json=payload, timeout=aiohttp.ClientTimeout(total=600)) as resp:
        resp.raise_for_status()
        return await resp.json()


async def run_sample_by_sample(url, groups, samples_per_group, max_new_tokens):
    async with aiohttp.ClientSession() as session:
        tasks = []
        for group_idx in range(groups):
            prompt = f"{PROMPT}\nGroup id: {group_idx}"
            for _ in range(samples_per_group):
                tasks.append(post_generate(session, url, make_payload(prompt, max_new_tokens)))
        await asyncio.gather(*tasks)


async def run_group_batched(url, groups, samples_per_group, max_new_tokens):
    async with aiohttp.ClientSession() as session:
        tasks = []
        for group_idx in range(groups):
            prompt = f"{PROMPT}\nGroup id: {group_idx}"
            payload = make_payload([prompt] * samples_per_group, max_new_tokens)
            tasks.append(post_generate(session, url, payload))
        await asyncio.gather(*tasks)


async def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--url", default="http://127.0.0.1:30000/generate")
    parser.add_argument("--groups", type=int, default=256)
    parser.add_argument("--samples-per-group", type=int, default=4)
    parser.add_argument("--max-new-tokens", type=int, default=256)
    parser.add_argument("--mode", choices=["sample", "group"], required=True)
    args = parser.parse_args()

    start = time.perf_counter()
    if args.mode == "sample":
        await run_sample_by_sample(args.url, args.groups, args.samples_per_group, args.max_new_tokens)
    else:
        await run_group_batched(args.url, args.groups, args.samples_per_group, args.max_new_tokens)
    elapsed = time.perf_counter() - start

    total_requests = args.groups * args.samples_per_group if args.mode == "sample" else args.groups
    total_outputs = args.groups * args.samples_per_group
    print(
        f"mode={args.mode} groups={args.groups} samples_per_group={args.samples_per_group} "
        f"http_requests={total_requests} outputs={total_outputs} elapsed_sec={elapsed:.2f}"
    )


if __name__ == "__main__":
    asyncio.run(main())

Run both modes:

python repro_sglang_group_batch.py --mode sample
python repro_sglang_group_batch.py --mode group

Expected result: both modes request the same number of generated outputs with the same max_new_tokens, but sample-by-sample mode is slower because it sends 4x more HTTP/router requests. This mirrors the original Miles rollout path for GRPO-style sampling (n_samples_per_prompt=4) and the PR's group-batched path.

Devpod reproduction result on f08a506f0598d435ea6c-n0-0-master-0 using the public Qwen/Qwen2-0.5B-Instruct model with SGLang on one H200 GPU:

Workload	Sample-by-sample	Group-batched	Speedup
256 groups x 4 samples x 256 max_new_tokens	7.51s, 1024 HTTP requests	6.67s, 256 HTTP requests	1.13x
512 groups x 4 samples x 16 max_new_tokens	3.11s, 2048 HTTP requests	1.91s, 512 HTTP requests	1.63x

The short-output case isolates request/router overhead more clearly. The production Miles profile shows a larger gap because it combines this request granularity issue with the real rollout/training shape.

SGLang source code path

Pinned source links below use upstream SGLang tag v0.5.9 (bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2), matching the 0.5.9.x SGLang line used in the reproduction.

The group-batched client payload is batched because text or input_ids is a list of prompts/sequences, e.g. make_payload([prompt] * samples_per_group, max_new_tokens) sends one JSON object whose text field is List[str].

1. /generate receives one HTTP request and FastAPI parses it into one GenerateReqInput: https://github.com/sgl-project/sglang/blob/bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2/python/sglang/srt/entrypoints/http_server.py#L648-L680
2. GenerateReqInput explicitly supports both single input and batched input: text: Union[List[str], str] and input_ids: Union[List[List[int]], List[int]]: https://github.com/sgl-project/sglang/blob/bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2/python/sglang/srt/managers/io_struct.py#L171-L191
3. During normalization, SGLang determines batching from the input shape. For text, str means is_single=True, while List[str] means is_single=False and batch_size=len(text). For input_ids, List[int] is single, while List[List[int]] is batched: https://github.com/sgl-project/sglang/blob/bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2/python/sglang/srt/managers/io_struct.py#L289-L342
4. TokenizerManager.generate_request() branches on obj.is_single; batched objects go to _handle_batch_request, which tokenizes/processes the batch and sends one BatchTokenizedGenerateReqInput to the scheduler: https://github.com/sgl-project/sglang/blob/bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2/python/sglang/srt/managers/tokenizer_manager.py#L490-L524 and https://github.com/sgl-project/sglang/blob/bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2/python/sglang/srt/managers/tokenizer_manager.py#L1214-L1228

The scheduler has a dedicated BatchTokenizedGenerateReqInput handler. It expands the batch into normal internal requests, but this happens after the HTTP/tokenizer/scheduler-dispatch boundary has already been crossed once: https://github.com/sgl-project/sglang/blob/bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2/python/sglang/srt/managers/scheduler.py#L1645-L1654

This supports the PR's interpretation: the speedup is mainly from reducing request/message granularity before SGLang's internal continuous batching stage. The old Miles path paid HTTP/tokenizer/scheduler-dispatch overhead once per sampled response; the new path pays it once per prompt group, then lets SGLang expand the group inside the server.

Validation from proposed fix

• python3 -m py_compile miles/rollout/sglang_rollout.py
• Devpod smoke profile with no-op reward confirmed the batched path reduced rollout time from ~28s to ~14s for the same workload.