Skip to content

[WIP][Feature] Temporal Pipeline Parallelism & Stream Batch for Real-Time Video#3099

Draft
mnasser02 wants to merge 56 commits into
vllm-project:mainfrom
zzhang-fr:stream-diffusion
Draft

[WIP][Feature] Temporal Pipeline Parallelism & Stream Batch for Real-Time Video#3099
mnasser02 wants to merge 56 commits into
vllm-project:mainfrom
zzhang-fr:stream-diffusion

Conversation

@mnasser02
Copy link
Copy Markdown

@mnasser02 mnasser02 commented Apr 24, 2026 •
edited
Loading

Purpose

Implements #2280.

Targets 40+ FPS on 4×H100, sub-500 ms TTFF (Wan2.1-1.3B).

Stacks on #2322.

Item Status
SupportsStepExecution on Wan22Pipeline ✅
micro-step execution pipeline with async comms, double-buffering, compute/comm overlap ✅
StreamBatchScheduler driving micro-step execution pipeline ✅
SLO-adaptive batch size, motion-aware noise controller + config 🔲
StreamVAE 🔲
Rolling KV cache (after #2305), adaptive sink tokens, RoPE offset reset 🔲
Make sure temporal PP works with other types of parallelism 🔲
(optional) VAE-as-stage (after #2089) 🔲
(optional) Run-time adaptive scheduling of transformer blocks across PP stages 🔲
Benchmarks vs SP baseline, PP×SP hybrid, docs 🔲

Test Plan

(NOT FINALIZED)

pytest -v \
tests/diffusion/test_diffusion_micro_step_pipeline.py \
tests/diffusion/test_diffusion_scheduler.py::TestStreamBatchScheduler

Test Result

Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the test style doc
  • The test results. Please paste the results comparison before and after, or the e2e results.
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model. Please run mkdocs serve to sync the documentation editions to ./docs.
  • (Optional) Release notes update. If your change is user-facing, please update the release notes draft.

BEFORE SUBMITTING, PLEASE READ https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md (anything written below this line will be removed by GitHub Actions)

hadipash and others added 30 commits April 10, 2026 15:40
@hadipash
add pipeline parallel to Wan 2.2
a25e9a1 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
add async communication for last -> first rank
70fb785 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
bug fix
ea5aa03 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
bug fix
95c9db8 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
align isend with vLLM's implementation
8a62204 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
add CFG parallel support
79a83db 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
add unit tests and docs
0db0f8f 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
add online serving
7244874 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
remove unused pipeline
eb9d9ac 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
Merge branch 'main' into wan_pipe_parallel
eb1190b 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>

# Conflicts:
#	docs/user_guide/diffusion_features.md
@hadipash
add safeguard
154a4db 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
align UT with other tests
8c394da 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
bug fixes and align UT with other tests
2be4cc9 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
Merge branch 'main' into wan_pipe_parallel
169ad84 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>

# Conflicts:
#	vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py
#	vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
@hadipash
Merge branch 'main' into wan_pipe_parallel
46ac762 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>

# Conflicts:
#	docs/user_guide/diffusion_features.md
@hadipash
Merge branch 'main' into wan_pipe_parallel
7134edd 
# Conflicts:
#	docs/user_guide/diffusion_features.md
@hadipash
Merge branch 'main' into wan_pipe_parallel
34d4abb 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>

# Conflicts:
#	docs/user_guide/diffusion_features.md
#	vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
#	vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
#	vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py
#	vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
@hadipash
add NPU support
4f19702 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@hadipash
remove unused pipeline
e1abb75 
Signed-off-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>
@mnasser02
Add step execution and denoise_micro_step to Wan2.2 pipeline
f4d6732 
Implement SupportsStepExecution protocol on Wan22Pipeline, decomposing the monolithic forward() into prepare_encode, denoise_step,step_scheduler, and post_decode. Add denoise_micro_step for temporal PP.

Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
remove denoise_micro_step
59b15ab 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
add unit tests for wan22 step exec
15c31a8 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
add RankTask and track per-rank tasks in DiffusionSchedulerOutput
0c003a4 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
add num_chunks sampling param
004f5f2 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Add chunk related data structures to track its state.
cc4e7e0 
>> Different ranks work on different chunks. A context manager that views the req state of a rank as a chunk state allows benefitting from existing functionalities.

Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Add execute_micro_step execution path that works on a per-rank assign…
5716eb6 
…ed task

Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Add StreamBatchScheduler that controls the flow of chunks through the…
9c08c56 
… pipeline (B and T hardcoded for now)

Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Add StreamBatchScheduler tests
283ca17 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Add micro-step execution pipeline tests
0084f89 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Add stream_batch arg
def46ff 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
mnasser02 added 17 commits April 24, 2026 10:40
@mnasser02
Set reply rank to 0
255263b 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Fix blocking send/recv resulting in deadlock by registering metadata …
528f3be 
…instead of sync send/recv

Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
bugfix
5bc3c78 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
bugfix
f0af7a9 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
bugfix
017abee 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Implement fully async isend/recv_dicts and migrate PPMixin to use them
1a9b42c 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Rmv stale code
a9745ed 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
bugfix
a81b233 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Make sending dict schema unblocking
cdac8ec 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Add warmup for nccl comms on init
958d020 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Use batch_isend_irecv instead of plain isend/recv
aaa0609 
Plain P2P on size-2 PG triggers lazy sub-comm creation that requires the peer present.

Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Handle tensor receive as a task
462ffe9 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Chunk-aware buffers
cfd1cd3 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
fix
87059ee 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Improve and fix PP communication.
05f5020 
- Separate comms stream
- Double buffering
- Set rcv buffers for a new req
- Revert changes regarding Async structs

Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
bugfixes
07842fd 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
Remove stale tests
4c8e2b7 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@hsliuustc0106
Copy link
Copy Markdown
Collaborator

hsliuustc0106 commented Apr 24, 2026

Ready for full review when WIP status is removed. Preliminary scan available on request.

Note: Test plan and test results sections are currently empty. Please provide:

  1. Test scripts & commands for the new pipeline parallelism and stream batch scheduler
  2. E2e test results demonstrating the 40+ FPS target on 4×H100 with sub-500ms TTFF
  3. Benchmark comparisons against SP baseline and PP×SP hybrid

@mnasser02
Add unit tests for StreamBatchScheduler
98145a1 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
add SupportsMicroStepExecution
dac11a9 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
add unit tests for micro-step execution pipeline
7e4be1d 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
edit
b34cf19 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
mnasser02 added 3 commits May 6, 2026 10:45
@mnasser02
add slo-adaptive scheduling
c0ef080 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
bugfix
0d562d1 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
bugfix
897b065 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@wtomin wtomin mentioned this pull request May 7, 2026
Open
25 tasks
mnasser02 added 2 commits May 11, 2026 18:49
@mnasser02
fix and improve stream batching
7bec36c 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
@mnasser02
modify tests
1eeba3b 
Signed-off-by: Mahdi Nasser <94046147+mnasser02@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Isotr0py Isotr0py Awaiting requested review from Isotr0py Isotr0py will be requested when the pull request is marked ready for review Isotr0py is a code owner

@princepride princepride Awaiting requested review from princepride princepride will be requested when the pull request is marked ready for review princepride is a code owner

@SamitHuang SamitHuang Awaiting requested review from SamitHuang SamitHuang will be requested when the pull request is marked ready for review SamitHuang is a code owner

@wtomin wtomin Awaiting requested review from wtomin wtomin will be requested when the pull request is marked ready for review wtomin is a code owner

@ZJY0516 ZJY0516 Awaiting requested review from ZJY0516 ZJY0516 will be requested when the pull request is marked ready for review ZJY0516 is a code owner

@RuixiangMa RuixiangMa Awaiting requested review from RuixiangMa RuixiangMa will be requested when the pull request is marked ready for review RuixiangMa is a code owner

@david6666666 david6666666 Awaiting requested review from david6666666 david6666666 will be requested when the pull request is marked ready for review david6666666 is a code owner

@tzhouam tzhouam Awaiting requested review from tzhouam tzhouam will be requested when the pull request is marked ready for review tzhouam is a code owner

@yenuo26 yenuo26 Awaiting requested review from yenuo26 yenuo26 will be requested when the pull request is marked ready for review yenuo26 is a code owner

@hsliuustc0106 hsliuustc0106 Awaiting requested review from hsliuustc0106 hsliuustc0106 will be requested when the pull request is marked ready for review hsliuustc0106 is a code owner

@Gaohan123 Gaohan123 Awaiting requested review from Gaohan123 Gaohan123 will be requested when the pull request is marked ready for review Gaohan123 is a code owner

@ywang96 ywang96 Awaiting requested review from ywang96 ywang96 will be requested when the pull request is marked ready for review ywang96 is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@mnasser02 @hsliuustc0106 @hadipash