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

### Checklist - [x] I searched related issues but found no solution. - [x] The bug persists in the latest version. - [x] Issues without environment info and a minimal reproducible demo are hard to resolve and may receive no feedback. - [x] 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. - [x] 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 | Branch | |---|---| | 1 | [`CjhHa1:fix/multi-output-rollout-trajectory-merge`](https://github.com/CjhHa1/sglang/tree/fix/multi-output-rollout-trajectory-merge) | | 2 | [`CjhHa1:fix/grouped-forward-residency-manager`](https://github.com/CjhHa1/sglang/tree/fix/grouped-forward-residency-manager) | | 3 | [`CjhHa1:fix/provided-latents-pack-and-latent-ids`](https://github.com/CjhHa1/sglang/tree/fix/provided-latents-pack-and-latent-ids) | --- ## 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`](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py#L882-L886) — `_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`](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py#L516) — `_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](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py#L698)) and `_collect_sequential_outputs` ([:430](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py#L430)) 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`](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py#L78), [`:40`](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py#L40)). 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__`](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py#L53) seeds `component_residency_manager = None`, and every `_execute_stages` run enters `_component_residency_request` → [`begin_component_residency_request`](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/pipelines_core/executors/pipeline_executor.py#L65), which dereferences it unguarded. Of the three forward entry points in `composed_pipeline_base.py`, two install the manager and one does not: | entry point | installs manager? | |---|---| | `forward` ([:1032-1035](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py#L1032-L1035)) | yes | | `forward_batch` ([:1065](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py#L1065)) | **no** | | `forward_batch_sequentially` ([:1083-1086](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py#L1083-L1086)) | yes | `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`](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py#L145-L170) runs different preparation depending on where the latents came from: ```python 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`](https://github.com/sgl-project/sglang/blob/3ed2a0adf3d87b0f527c173a500dfb40d64b572f/python/sglang/multimodal_gen/configs/pipeline_configs/flux.py#L606-L607) later does: ```python 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: ```python # 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`)