[Bug] Multi-output diffusion rollout: per-sample trajectories collapse to output 0, grouped forward AttributeError, provided latents skip packing

Checklist

• I searched related issues but found no solution.
• The bug persists in the latest version.
• Issues without environment info and a minimal reproducible demo are hard to resolve and may receive no feedback.
• If this is not a bug report but a general question, please start a discussion at https://github.com/sgl-project/sglang/discussions. Otherwise, it will be closed.
• Please use English. Otherwise, it will be closed.

Describe the bug

Three independent bugs on the diffusion multi-output / rollout path, all present on main at 3ed2a0adf3d87b0f527c173a500dfb40d64b572f. They surface together when a request expands to num_outputs_per_prompt > 1 with rollout=True (GRPO-style post-training), but each is independent and fixable on its own.

Bug 1 is the serious one: it silently produces wrong numbers rather than failing.

I have fixes for all three, each on its own branch with a regression test, and can open PRs — I'd like a maintainer's read on the intended semantics first, particularly for Bug 1. Branches:

Bug 1 — all K samples in a multi-output group report output 0's rollout trajectory

Severity: silent numerical corruption. No exception, no warning; training just stops learning.

A num_outputs_per_prompt=K request runs as K per-output forwards, each producing its own rollout_trajectory_data from its own x_T slice and per-step SDE noise. Two sites then collapse those K distinct trajectories into K copies of output 0's:

1. gpu_worker.py:882-886 — _merge_expanded_singletons keeps the trajectory of the first per-output batch only:

   python复制

   if (
       merged.rollout_trajectory_data is None
       and output_batch.rollout_trajectory_data is not None
   ):
       merged.rollout_trajectory_data = output_batch.rollout_trajectory_data

   Note the asymmetry: output, trajectory_latents, noise_pred and trajectory_decoded are all concatenated along the batch dim in _finalize_expanded_parts. Only rollout_trajectory_data is treated as a singleton.

2. diffusion_generator.py:516 — _result_common slices samples and metrics per output_index (metrics_list[output_index] at line 502) but hands the whole trajectory to every GenerationResult.

So every sample in the group carries output 0's rollout_log_probs. For GRPO the per-sample log-probs become identical, the mean-zero advantages cancel, and the gradient vanishes — reward stays flat with grad_norm ≈ 0 and nothing in the logs indicates a problem.

Both entry points into the merge are affected, including the recently added sequential path: _forward_group (:698) and _collect_sequential_outputs (:430) both call _merge_expanded_output_batches.

The HTTP rollout path is affected too, and reveals the intent. rollout_api._slice_rollout_trajectory_for_sample already slices per-sample — but every slice is guarded on shape[0] == batch_size (rollout_api.py:78, :40). Against the collapsed [1, ...] tensor that guard doesn't match, so _extract_single_sample_tensor returns the tensor whole and the slicing silently no-ops. That existing code expects a [K, ...] trajectory, which is what the merge should have produced.

Observed on my branch, K=3, T=4 (per-output values tagged 1/2/3):

                          merged log_probs   result[0]  result[1]  result[2]
before (3ed2a0ad)             (1, 4)             1          1          1     <- all output 0
after                         (3, 4)             1          2          3

Question for maintainers: was PR #22183 ("Sequential Per-Output Execution") meant to cover this? It concatenates the flat trajectory_latents / trajectory_decoded fields but doesn't touch rollout_trajectory_data, so as far as I can tell this bug survives it. My fix works at the GPUWorker merge layer and looks orthogonal to that PR's executor-layer change.

Bug 2 — AttributeError on the first grouped forward

PipelineExecutor.__init__ seeds component_residency_manager = None, and every _execute_stages run enters _component_residency_request → begin_component_residency_request, which dereferences it unguarded.

Of the three forward entry points in composed_pipeline_base.py, two install the manager and one does not:

forward_batch is the grouped path taken when a request expands to num_outputs_per_prompt > 1, so it reaches execute_group_with_profiling with the manager still None:

AttributeError: 'NoneType' object has no attribute 'begin_request'

That the grouped path is meant to be supported is visible in begin_component_residency_request itself, which explicitly handles a list payload (if isinstance(batch, list): batch = batch[0]). Only the install was missed. forward_batch_sequentially — added later — got it right, which suggests the two-way duplication is what let forward_batch drift.

Bug 3 — caller-provided initial latents skip latent-ids and packing

LatentPreparationStage.forward runs different preparation depending on where the latents came from:

if latents is None:
    latents = randn_tensor(...)
    latent_ids = server_args.pipeline_config.maybe_prepare_latent_ids(latents) ...
    if latent_ids is not None:
        batch.latent_ids = latent_ids.to(device=device)
    if spec.pack_latents:
        latents = server_args.pipeline_config.maybe_pack_latents(latents, batch_size, batch)
else:
    latents = latents.to(device)     # <- no latent_ids, no packing

For packed models this is fatal. Flux2PipelineConfig supplies both hooks, and FluxPipelineConfig.get_freqs_cis later does:

img_ids = batch.latent_ids
if img_ids.ndim == 3:

so an injected x_T reaches the denoising loop with latent_ids still None → AttributeError: 'NoneType' object has no attribute 'ndim'. The latents also stay unpacked (B, C, H, W) where the transformer expects packed (B, H*W, C).

Observed with a packed-model config stub:

                        latent_ids     packed?   latents shape
before  randn            (1, 64, 4)     True     (1, 64, 4)
before  provided x_T           None    False     (1, 4, 8, 8)   <- unusable downstream
after   provided x_T     (1, 64, 4)     True     (1, 64, 4)

Note the grouped path already funnels provided latents into forward: run_grouped_requests falls back to per-request self(batch, server_args) when any(batch.latents is not None ...), so fixing forward covers both paths.

This one is arguably "provided latents were only ever expected in already-prepared form" rather than a bug — I'd like a maintainer to confirm the intended contract. My reading is that the randn branch defines it (unpacked in, packed out), since maybe_pack_latents / maybe_prepare_latent_ids both expect unpacked (B, C, H, W).

Reproduction

All three require the diffusion multi-output rollout path. The common trigger:

# num_outputs_per_prompt > 1 with rollout enabled, on a FlowMatch scheduler model (e.g. SD3)
sampling_params_kwargs = {
    "prompt": "...",
    "num_outputs_per_prompt": 4,      # -> grouped forward
    "rollout": True,                  # -> rollout_trajectory_data populated
    "rollout_return_dit_trajectory": True,
}

• Bug 2 fires first, at the initial grouped forward (AttributeError: 'NoneType' object has no attribute 'begin_request').
• Bug 1 then shows up as identical rollout_log_probs across the K results of one group. Direct check: with K per-output batches carrying distinct trajectories, GPUWorker._merge_expanded_output_batches returns batch dim 1 instead of K, and DiffGenerator._result_common(req, merged, t, idx) returns the same trajectory for every idx.
• Bug 3 needs a packed model (FLUX.2-family) plus a caller-supplied Req.latents; it fails in get_freqs_cis with AttributeError: 'NoneType' object has no attribute 'ndim'.

How I verified. I could not run sglang's pytest suite against main on the machine I had. The sglang install available to me is 0.5.12.post1, whose tree predates the layout on main (it still has runtime/loader/weights_updater.py and lacks runtime/post_training/weights_updater.py and runtime/managers/memory_managers/), so it cannot exercise main-based changes; and the only box I had is a CPU login node where importing sglang.multimodal_gen off network storage takes >10 minutes, which made iterating on a real test run impractical.

Instead I extracted the real functions under test from both main (3ed2a0adf) and each fix branch via git show, executed them side by side against stubs, and confirmed for each bug: it reproduces before, is fixed after, and the unrelated paths (K=1 groups, non-rollout requests, the randn latent branch, the output/trajectory_latents merges) are unchanged. The numeric tables above come from that harness.

Each branch also carries a unit test under python/sglang/multimodal_gen/test/unit/, but those tests have not been executed under a real sglang install — they'd need CI or a maintainer to confirm. Please treat the committed tests as unvalidated and the tables above as the evidence.

One deliberate behavior change in the Bug 1 fix worth flagging: when only some outputs in a group carry a trajectory, the merge now returns None instead of the single present row. Labeling one sample's trajectory as the whole group's is the same broadcast bug in a different disguise, so dropping it seemed safer than misaligning it — but I'm happy to change that if you'd prefer it preserved.

Environment

I hit these while running diffusion RL post-training against sglang[diffusion]==0.5.12.post1, and re-verified every finding by reading and executing code from main at 3ed2a0adf3d87b0f527c173a500dfb40d64b572f (2026-08-07). All three bugs are present on that commit — the analysis above cites it exclusively, not the older pin.

python3 -m sglang.check_env output isn't informative here: the box I verified on is a CPU login node whose sglang install is 0.5.12.post1, not the main tree the analysis targets, so its environment dump would describe neither the runtime where the symptoms appeared nor the code being fixed. The bugs are all in Python control flow on the multi-output path and are independent of GPU/driver/kernel versions — the permalinks above should be sufficient to confirm each by inspection.

Original runtime where the symptoms appeared:

sglang[diffusion] 0.5.12.post1
torch 2.11.0+cu130
GPU: NVIDIA H20
model: Stable Diffusion 3 (FlowMatchEulerDiscreteScheduler), GRPO-style post-training
symptom: reward flat, grad_norm ~ 0, no error raised

cc @mickqian @ping1jing2 (per CODEOWNERS for /python/sglang/multimodal_gen)