[feat] dreamverse: sequence parallelism for serving

[shaoxiongduan]

Add DREAMVERSE_SP_SIZE to run a single generation across multiple GPUs via sequence parallelism. Defaults to 1, which is byte-for-byte the existing single-GPU-per-session behavior.

• config: DREAMVERSE_SP_SIZE env knob
• gpu_pool: group visible GPUs into SP slots, expose the whole group to the worker (CUDA_VISIBLE_DEVICES="g0,g1,..."), keyed by the first GPU
• video_generation: drive engine num_gpus from DREAMVERSE_SP_SIZE
• ltx2: drop @torch.compiler.disable on LTXDistributedAttention so the all-to-all distributed attention compiles into the fullgraph at sp>1

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟠 PR merge requirements

Waiting for

• check-success=fastcheck-passed
• check-success=full-suite-passed

Waiting checks: fastcheck-passed, full-suite-passed.

• check-success=fastcheck-passed
• check-success=full-suite-passed
• #approved-reviews-by>=1
• check-success~=pre-commit
• title~=(?i)^\[(feat|feature|bugfix|fix|refactor|perf|ci|doc|docs|misc|chore|kernel|new.?model|skill|skills|infra)\]

[gemini-code-assist]

Code Review

This pull request introduces sequence-parallel (SP) support to Dreamverse by adding a new DREAMVERSE_SP_SIZE configuration. It updates the GPU pool manager to group available GPUs into sequence-parallel slots of this size, configures the video generation engine to use the specified number of GPUs, and removes the @torch.compiler.disable decorator from the forward pass in fastvideo/models/dits/ltx2.py. The review feedback suggests adding a warning log in gpu_pool.py to inform operators when some GPUs are left idle because they cannot be evenly grouped into sequence-parallel slots.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

When DREAMVERSE_SP_SIZE is greater than 1, any available GPUs that do not fit evenly into sequence-parallel groups of size sp_size are silently dropped. It would be highly beneficial to log a warning message to inform the operator that some GPUs will remain idle due to the grouping configuration.

        sp_size = DREAMVERSE_SP_SIZE
        groups = [gpu_ids[i:i + sp_size] for i in range(0, len(gpu_ids), sp_size)]
        unused_count = len(gpu_ids) % sp_size
        if unused_count > 0:
            print(f"[WARNING] {unused_count} GPU(s) will be unused because they cannot be grouped into sequence-parallel slots of size {sp_size}.")
        groups = [g for g in groups if len(g) == sp_size]

[mergify]

Pre-commit checks failed

Hi @shaoxiongduan, the pre-commit checks have failed. To fix them locally:

# Install pre-commit if you haven't already
uv pip install pre-commit
pre-commit install

# Run all checks and auto-fix what's possible
pre-commit run --all-files

Common fixes:

• yapf: yapf -i <file> (formatting)
• ruff: ruff check --fix <file> (linting)
• codespell: codespell --write-changes <file> (spelling)

After fixing, commit and push the changes. The checks will re-run automatically.

For future commits, pre-commit will run automatically on changed files before each commit.

[SolitaryThinker]

Hi @shaoxiongduan — automated review from Gob, one of @SolitaryThinker's AI reviewers. Findings below are advisory; pushback welcome, especially on the S2 and S3 calls which are judgment-dependent. The S1 mirrors Gemini's prior flag at the same HEAD.

---

Summary (verdict: ship-with-fixes)

Small, focused PR. Config guards are right, GPU grouping is contiguous and keyed by leader, the EngineConfig.num_gpus wiring works. Two issues worth addressing before merge:

1. S1: Leftover GPUs are silently dropped when len(gpu_ids) % SP_SIZE != 0. Operators won't know why a slot vanished.
2. S2: The @torch.compiler.disable removal on LTXAttention.forward has no regression test — the only validation signal is a live Dreamverse serving deploy.

One optional clarification (S3) on SP-vs-TP semantics.

---

[S1] gpu_pool.py:849-851 — silent GPU capacity loss

groups = [gpu_ids[i:i + sp_size] for i in range(0, len(gpu_ids), sp_size)]
groups = [g for g in groups if len(g) == sp_size]

With gpu_ids=[0,1,2,3,4] and SP_SIZE=2, the trailing [4] is silently dropped — 2 slots × 2 GPUs = 4 used, GPU 4 sits idle, no log line. The print at line 855 fires only when sp_size > 1 AND a slot was created, so it prints what was used, not what was dropped. Same flag Gemini left earlier on this PR.

Suggested patch:

remainder = len(gpu_ids) % sp_size
if remainder:
    dropped = gpu_ids[-remainder:]
    print(f"[WARN] DREAMVERSE_SP_SIZE={sp_size} leaves {remainder} GPU(s) "
          f"unused (idle GPUs={dropped}). Pick an sp_size that divides "
          f"len(gpu_ids)={len(gpu_ids)}, or restrict CUDA_VISIBLE_DEVICES.",
          file=sys.stderr)

---

[S2] fastvideo/models/dits/ltx2.py:1192 — @torch.compiler.disable removal has no regression coverage

The forward contains constructs that historically motivate torch.compiler.disable — get_sp_parallel_rank(), get_forward_context(), in-place tensor assignment (qkvg[:batch_size * 2] = ...), the if use_vsa: branch, and sequence_model_parallel_all_to_all_4D collectives. The decorator was presumably there because at least one of these tripped Dynamo at some point.

The only in-tree LTX-2 forward test is fastvideo/tests/ssim/test_ltx2_similarity.py, which runs at sp_size=1. There's no regression test for the compile + SP path the removal is meant to unblock — so a future Dynamo/PyTorch change that re-introduces a graph-break or guard failure would only surface on a real Dreamverse deploy.

Options, in preference order:

1. Extend test_ltx2_similarity.py with an sp_size=2 + compile-enabled scenario (or add a sibling test).
2. If multi-GPU CI isn't available, add a CPU-side smoke test that torch.compile(LTXAttention(...)) succeeds on a representative shape — won't catch correctness regressions, but will catch graph-break regressions.
3. Pin the Dreamverse SP + compile run that validated this in the PR description (sp_size, compile mode, SHA) so future readers have a reference point.

Confirmed there are no other @torch.compiler.disable decorators on this forward at HEAD 9b918ffa, so this is the load-bearing change for compile + SP correctness.

---

[S3] video_generation.py:292 — num_gpus = DREAMVERSE_SP_SIZE conflates SP with TP (optional)

EngineConfig.num_gpus is the union of TP + SP world size. The PR title says "sequence parallelism" but if a future Dreamverse model preset defaults to TP instead of SP, DREAMVERSE_SP_SIZE=2 would silently produce TP=2 / SP=1 with no error, which has different memory and comm characteristics and might break the SP-specific code paths in ltx2.py (the all_to_all_4D over heads).

Not blocking if LTX-2 is the only Dreamverse model right now and its preset is SP-default — worth confirming in a follow-up. Options:

• Explicitly set sp_size=DREAMVERSE_SP_SIZE, tp_size=1 if those fields exist on EngineConfig.
• Rename the env var to DREAMVERSE_NUM_GPUS to match the actual config field semantics.
• Add an assertion in gpu_pool.py / config.py documenting that this maps to num_gpus and assumes SP-default presets.

---

Notes / what's well-done

• Config guards are complete: _env_int (no crash on bad input), max(1, ...) (negative/zero clamped), SP_SIZE > n_gpus raises RuntimeError loudly.
• self.gpu_ids and self.slots keyed by leader GPU of each group, with cuda_device as comma-joined members — the right primitive; CUDA_VISIBLE_DEVICES propagates correctly via the existing _worker_main path at gpu_pool.py:160.
• Decorator removal is one-line and surgical, no drive-by edits.
• Same-repo branch, small diff — easy to revert if compile regresses.

Heads-up on apps/dreamverse/serve_configs/streaming_demo.yaml still hardcoding num_gpus: 1 — if that YAML and the env var both flow into EngineConfig, the precedence order isn't obvious. One sentence in the PR description would close that gap.

---

Posted by Gob, automated reviewer for @SolitaryThinker. Findings are advisory — author judgment is the final call.

[SolitaryThinker]

Hi @shaoxiongduan — automated review from Gob, one of @SolitaryThinker's AI reviewers. Findings below are advisory; pushback welcome, especially on the S2 and S3 calls which are judgment-dependent. The S1 mirrors Gemini's prior flag at the same HEAD.

---

Summary (verdict: ship-with-fixes)

Small, focused PR. Config guards are right, GPU grouping is contiguous and keyed by leader, the EngineConfig.num_gpus wiring works. Two issues worth addressing before merge:

1. S1: Leftover GPUs are silently dropped when len(gpu_ids) % SP_SIZE != 0. Operators won't know why a slot vanished.
2. S2: The @torch.compiler.disable removal on LTXAttention.forward has no regression test — the only validation signal is a live Dreamverse serving deploy.

One optional clarification (S3) on SP-vs-TP semantics.

---

[S1] gpu_pool.py:849-851 — silent GPU capacity loss

groups = [gpu_ids[i:i + sp_size] for i in range(0, len(gpu_ids), sp_size)]
groups = [g for g in groups if len(g) == sp_size]

With gpu_ids=[0,1,2,3,4] and SP_SIZE=2, the trailing [4] is silently dropped — 2 slots × 2 GPUs = 4 used, GPU 4 sits idle, no log line. The print at line 855 fires only when sp_size > 1 AND a slot was created, so it prints what was used, not what was dropped. Same flag Gemini left earlier on this PR.

Suggested patch:

remainder = len(gpu_ids) % sp_size
if remainder:
    dropped = gpu_ids[-remainder:]
    print(f"[WARN] DREAMVERSE_SP_SIZE={sp_size} leaves {remainder} GPU(s) "
          f"unused (idle GPUs={dropped}). Pick an sp_size that divides "
          f"len(gpu_ids)={len(gpu_ids)}, or restrict CUDA_VISIBLE_DEVICES.",
          file=sys.stderr)

---

[S2] fastvideo/models/dits/ltx2.py:1192 — @torch.compiler.disable removal has no regression coverage

The forward contains constructs that historically motivate torch.compiler.disable — get_sp_parallel_rank(), get_forward_context(), in-place tensor assignment (qkvg[:batch_size * 2] = ...), the if use_vsa: branch, and sequence_model_parallel_all_to_all_4D collectives. The decorator was presumably there because at least one of these tripped Dynamo at some point.

The only in-tree LTX-2 forward test is fastvideo/tests/ssim/test_ltx2_similarity.py, which runs at sp_size=1. There's no regression test for the compile + SP path the removal is meant to unblock — so a future Dynamo/PyTorch change that re-introduces a graph-break or guard failure would only surface on a real Dreamverse deploy.

Options, in preference order:

1. Extend test_ltx2_similarity.py with an sp_size=2 + compile-enabled scenario (or add a sibling test).
2. If multi-GPU CI isn't available, add a CPU-side smoke test that torch.compile(LTXAttention(...)) succeeds on a representative shape — won't catch correctness regressions, but will catch graph-break regressions.
3. Pin the Dreamverse SP + compile run that validated this in the PR description (sp_size, compile mode, SHA) so future readers have a reference point.

Confirmed there are no other @torch.compiler.disable decorators on this forward at HEAD 9b918ffa, so this is the load-bearing change for compile + SP correctness.

---

[S3] video_generation.py:292 — num_gpus = DREAMVERSE_SP_SIZE conflates SP with TP (optional)

EngineConfig.num_gpus is the union of TP + SP world size. The PR title says "sequence parallelism" but if a future Dreamverse model preset defaults to TP instead of SP, DREAMVERSE_SP_SIZE=2 would silently produce TP=2 / SP=1 with no error, which has different memory and comm characteristics and might break the SP-specific code paths in ltx2.py (the all_to_all_4D over heads).

Not blocking if LTX-2 is the only Dreamverse model right now and its preset is SP-default — worth confirming in a follow-up. Options:

• Explicitly set sp_size=DREAMVERSE_SP_SIZE, tp_size=1 if those fields exist on EngineConfig.
• Rename the env var to DREAMVERSE_NUM_GPUS to match the actual config field semantics.
• Add an assertion in gpu_pool.py / config.py documenting that this maps to num_gpus and assumes SP-default presets.

---

Notes / what's well-done

• Config guards are complete: _env_int (no crash on bad input), max(1, ...) (negative/zero clamped), SP_SIZE > n_gpus raises RuntimeError loudly.
• self.gpu_ids and self.slots keyed by leader GPU of each group, with cuda_device as comma-joined members — the right primitive; CUDA_VISIBLE_DEVICES propagates correctly via the existing _worker_main path at gpu_pool.py:160.
• Decorator removal is one-line and surgical, no drive-by edits.
• Same-repo branch, small diff — easy to revert if compile regresses.

Heads-up on apps/dreamverse/serve_configs/streaming_demo.yaml still hardcoding num_gpus: 1 — if that YAML and the env var both flow into EngineConfig, the precedence order isn't obvious. One sentence in the PR description would close that gap.

---

Posted by Gob, automated reviewer for @SolitaryThinker. Findings are advisory — author judgment is the final call.