sgl-project / sgl-project/sglang
[Feature] Batched Candidate-Trajectory Ensembling for Stochastic Action Generation
- Dominant language
- Python
- Stars
- 36.1k
- Forks
- 9k
- Avg merge
- 1d 1h
- Merged PRs (30d)
- 219
Description
### Checklist
- [x] If this is not a feature request but a general question, please start a discussion at https://github.com/sgl-project/sglang/discussions. Otherwise, it will be closed.
- [x] Please use English. Otherwise, it will be closed.
### Motivation
#### Summary
Add a first-class, opt-in execution contract for generating multiple independent action trajectories from one conditioning context, executing the candidates as one physical batch, and reducing them to one served action result.
SGLang Diffusion already has most of the lower-level pieces:
- `Req.num_outputs_per_prompt` represents multiple outputs;
- latent preparation preserves per-request RNG streams and batches the deterministic packing/scaling work;
- pipeline stages can expand shared conditioning to the sample batch; and
- Cosmos3 is implemented as a composed pipeline with explicit latent, timestep, denoising, action, and decode stages.
What is missing is a model/runtime contract that says these outputs are *candidates for one logical prediction*, not independent public media outputs. Without that distinction, a world-action model has to implement sequential sampling or ad hoc aggregation inside model code, and the runtime cannot preserve candidate identity, enforce RNG independence, or measure the benefit of batching.
This proposal is intentionally narrow. It does not propose averaging images or videos, and it does not change default `num_outputs_per_prompt` behavior. It adds an opt-in path for models whose public action output has a validated candidate reducer.
#### Proposed execution contract
Illustrative types:
```python
@dataclass(frozen=True)
class CandidateTrajectorySpec:
count: int = 1
reducer: str = "none" # "none", "mean", or a registered model reducer
return_candidates: bool = False
seed_policy: str = "per_candidate"
@dataclass(frozen=True)
class CandidateGroup:
request_id: str
candidate_ids: tuple[int, ...]
public_output: str = "action"
```
The model/pipeline adapter would declare:
1. which tensor is the candidate action trajectory;
2. whether reduction occurs before or after denormalization;
3. the supported reducer and output dtype/shape;
4. which conditioning tensors are candidate-invariant; and
5. whether video or another auxiliary branch must execute for action correctness.
The runtime owns grouping, candidate identity, batching, failure handling, observability, and final response construction. The model owns the mathematical reduction rule.
#### Request and execution flow
```mermaid
flowchart LR
Req["One logical request
conditioning + candidate count N"]
Validate["Validate model capability
and candidate contract"]
Seeds["Create N independent
request-owned RNG streams"]
Shared["Compute candidate-invariant
conditioning once"]
Expand["Expand/fan out conditioning
and build candidate batch"]
Denoise["One batched denoising path
candidate axis = N"]
Action["Model-specific action
postprocessing"]
Reduce["Validated reducer
mean or registered reducer"]
Resp["One public action result
optional raw candidates"]
Req --> Validate --> Seeds --> Shared --> Expand --> Denoise --> Action --> Reduce --> Resp
```
#### RNG and equivalence rules
Candidate batching must not collapse stochastic independence or silently change seed semantics.
- Candidate `i` receives a stable stream derived from `(base_seed, i)` or an explicitly supplied generator.
- `count=1` must preserve the current output exactly.
- For a fixed candidate seed list, batched execution must produce the same candidate set as the sequential reference under the model's declared equivalence rule.
- Candidate ordering is stable across dynamic batching and output slicing.
- Retry/fallback may rerun the whole group, but must not reuse a partially advanced generator as though it were fresh.
The existing grouped latent-preparation logic is a useful precedent: it already draws request-owned random latents separately and only batches deterministic work afterward. The candidate path should retain the same property.
#### Scheduler and batching behavior
Phase 1 should admit a candidate group atomically when it fits the configured execution batch. If it does not fit, the runtime should fail clearly or use a documented sequential fallback; it should not silently truncate the group.
A later phase may split candidates into microbatches and aggregate when all slices finish, but the request state must then track:
- candidate IDs completed/pending;
- seed/generator state per candidate;
- partial reducer state; and
- cancellation semantics for the entire logical request.
Candidate-invariant conditioning should be encoded once where the stage contract permits it. An adapter may use an expanded view or an explicit attention fan-out; physical tensor replication is not required by this RFC.
#### Integration points
Suggested initial ownership:
- request fields and validation: diffusion sampling/request data;
- group identity and completion: request/scheduler state;
- independent raw noise: `LatentPreparationStage`;
- conditioning expansion: pipeline configuration's existing sample-batch expansion hook;
- candidate reducer: a small action-output capability implemented by opt-in world-model pipelines;
- metrics: candidate count, physical forwards, fallback reason, reducer time, and action-only branch time.
The first model integration could be Cosmos3 action generation because SGLang already has an explicit Cosmos3 composed pipeline. The framework abstraction should not contain Cosmos-specific tensor names.
#### Correctness and quality gates
Required tests:
1. `count=1` is identical to the existing path.
2. Fixed per-candidate seeds produce the same ordered candidate tensors in sequential and batched modes.
3. The reduced action matches a standalone reducer applied to captured candidate tensors.
4. Candidate IDs remain stable when unrelated requests are co-batched.
5. CFG serial and CFG-parallel paths preserve candidate membership.
6. TP/SP output reconstruction happens before reduction when the action tensor is sharded.
7. Cancellation and one-candidate failure do not return a partial aggregate as a valid result.
8. Auxiliary video decode can be skipped only when the model declares it irrelevant to the action output.
Quality validation must be model-specific. The framework should report candidate-level and reduced-output artifacts but should not hard-code a universal action metric or assume that a larger candidate count always improves quality.
#### Benchmark plan
Report, for `N in {1, 2, 4, 8, 10}`:
- sequential reference latency;
- batched latency and peak memory;
- number of physical transformer forwards;
- conditioning encode time;
- candidate-reduction time;
- candidate tensor parity under fixed seeds; and
- model-declared action-quality metrics.
Benchmark at more than one parallel topology. Wider TP may help the larger candidate batch while regressing `N=1`; the RFC should not encode one topology as universally optimal.
#### Alternatives considered
- **Keep aggregation in application code.** This forces repeated RPCs or model-specific server loops and prevents the runtime from preserving one logical request/candidate group.
- **Treat candidates as ordinary independent outputs.** This loses the reduction contract and makes failure/cancellation semantics ambiguous.
- **Always expose raw candidates.** This moves a model-system accuracy feature into every client and expands the public payload unnecessarily.
- **Average videos/latents generically.** Out of scope and not semantically valid for this proposal.
#### Rollout
1. Land internal types and validation with no model enabled.
2. Add a Cosmos3 action adapter and sequential-equivalence tests.
3. Add atomic candidate-group scheduling and metrics.
4. Benchmark before enabling any default.
5. Consider microbatching and additional reducers only after the initial contract is stable.
### Related resources
- Cosmos3 composed pipeline: https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py
- Latent preparation and per-request RNG handling: https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py
- Diffusion denoising stage and sample-batch expansion: https://github.com/sgl-project/sglang/blob/main/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
- vLLM-Omni request-selected output routing RFC (complementary output-branch contract): https://github.com/vllm-project/vllm-omni/issues/6302
Contributor guide
Research direction
Start by reading python/sglang/multimodal_gen/runtime/pipelines/cosmos3_pipeline.py, pipelines_core/stages/latent_preparation.py, and pipelines_core/stages/denoising.py to understand existing sample batching and RNG handling. Trace the diffusion sampling request and scheduler state entry points before deciding where the proposed contract belongs. Done means the opt-in candidate path, Cosmos3 integration, required equivalence and failure tests, metrics, and benchmark results satisfy the stated rollout and quality gates.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- ai, backend, distributed-systems
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100