Dev/gtong/disagg pipeline

[gtong-nv]

Three-stage disaggregated inference

LingBot can load the encoder, DiT, and decoder on three independent GPUs:

GPU 0: UMT5 + image/VAE/camera encoder
         │ Mooncake GPU-memory transfer
GPU 1: scheduler + DiT + session-pinned KV cache
         │ Mooncake GPU-memory transfer
GPU 2: streaming VAE or LightTAE decoder

This is pipeline-stage disaggregation, analogous to SGLang's independently
scheduled prefill and decode pools but split at diffusion-native boundaries.
The evolving autoregressive KV cache stays on the DiT worker. Only one-shot
conditioning, per-step encoder features, and clean latents cross stage
boundaries. Moving the KV cache every chunk would add a much larger transfer
and break session affinity.

Install the optional Mooncake transport:

uv sync --package flashdreams-lingbot --extra dev --extra disagg

# The container/host runtime must also provide the RDMA userspace libraries.
apt-get install libibverbs1 ibverbs-providers librdmacm1 ibverbs-utils

Run the reproducible three-GPU benchmark:

uv run --package flashdreams-lingbot torchrun \
  --standalone --nproc_per_node=3 \
  -m lingbot.disagg.benchmark \
  --model lingbot-world-fast-taehv-window15-sink3 \
  --warmup-blocks 6 --measured-blocks 5 \
  --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 \
  --output-dir outputs/lingbot_disagg

Use all eight GPUs for concurrent sessions by keeping one encoder and one
decoder worker and assigning the other six GPUs to session-affine DiT workers:

uv run --package flashdreams-lingbot torchrun \
  --standalone --nproc_per_node=8 \
  -m lingbot.disagg.benchmark_replicated \
  --dit-replicas 6 \
  --model lingbot-world-fast-taehv-window15-sink3 \
  --warmup-blocks 6 --measured-blocks 5 \
  --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 \
  --output-dir outputs/lingbot_disagg_1e6d1d

The allocation is derived from the tracked 1:1:1 stage service times. It is a
throughput topology: each DiT replica owns a distinct session and KV cache.
It does not split one session's DiT computation over six GPUs.

To minimize one session's latency instead, make ranks 1–6 one context-parallel
DiT group. Rank 1 receives the Mooncake handoff, broadcasts the input within
the NCCL subgroup, and sends the gathered clean latent to rank 7:

TORCHINDUCTOR_COMPILE_THREADS=4 \
uv run --package flashdreams-lingbot torchrun \
  --standalone --nproc_per_node=8 \
  -m lingbot.disagg.benchmark_cp \
  --cp-ranks 6 --cp-method ring \
  --model lingbot-world-fast-taehv-window15-sink3 \
  --warmup-blocks 6 --measured-blocks 5 \
  --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 \
  --output-dir outputs/lingbot_disagg_cp6

CP6 must use ring attention for this 40-head model. Ulysses requires the head
count to be divisible by the context-parallel size, so CP6 Ulysses is rejected
(40 % 6 != 0). A CP4 Ulysses comparison uses six processes and leaves two
GPUs idle:

TORCHINDUCTOR_COMPILE_THREADS=4 \
uv run --package flashdreams-lingbot torchrun \
  --standalone --nproc_per_node=6 \
  -m lingbot.disagg.benchmark_cp \
  --cp-ranks 4 --cp-method ulysses \
  --model lingbot-world-fast-taehv-window15-sink3 \
  --warmup-blocks 6 --measured-blocks 5 \
  --output-dir outputs/lingbot_disagg_cp4

Cap Inductor compilation parallelism when several compiled DiT ranks share a
node. The default 32 workers per rank becomes 192 workers at CP6 and can
overcommit the host during cold compilation.

On the tested H100 node, CP6 ring was the minimum-latency allocation: 743.27 ms
median per 12-frame chunk and 15.90 generated FPS, a 3.01× latency speedup over
CP1. CP4 Ulysses reached 754.41 ms and 15.70 FPS. CP6 was 1.5% faster; CP4 had
higher scaling efficiency and left two GPUs available for other work.

For an aggregated baseline, put the complete pipeline on every GPU and use all
eight ranks as the DiT context-parallel WORLD group:

TORCHINDUCTOR_COMPILE_THREADS=4 \
uv run --package flashdreams-lingbot torchrun \
  --standalone --nproc_per_node=8 \
  -m lingbot.disagg.benchmark_aggregated \
  --cp-method ulysses \
  --model lingbot-world-fast-taehv-window15-sink3 \
  --pixel-width 832 --pixel-height 448 \
  --warmup-blocks 6 --measured-blocks 5 \
  --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 \
  --output-dir outputs/lingbot_aggregated_cp8

The normal 832×464 grid has 4,524 tokens and cannot divide over CP8. The
nearest valid height is 448, which produces 4,368 tokens. On eight H100s, the
aggregated CP8 Ulysses run reached 393.33 ms median latency and 29.50 generated
FPS. It is the fastest tested single-session topology, but it replicates
encoder and decoder work and gives up independent stage placement. The
aggregated report and
comparison chart include the
resolution-normalized throughput and node-wide HBM tradeoff.

Validate the data plane without loading checkpoints:

uv run --package flashdreams-lingbot torchrun \
  --standalone --nproc_per_node=3 \
  -m lingbot.disagg.benchmark --transport-only \
  --bandwidth-probe-mib 256 --bandwidth-probe-iters 8

Both benchmarks write benchmark.json and a Markdown summary. They report:

• median and p90 encoder, DiT, finalize, decoder, and end-to-end chunk latency;
• generated FPS after excluding warmup;
• payload bytes, sender registration time, transfer time, full handoff time,
  and effective GB/s for encoder → DiT and DiT → decoder;
• a reusable 256 MiB transfer probe to distinguish link bandwidth from
  small-payload setup overhead;
• per-stage peak GPU memory and the exact software/hardware environment.

See the
three-stage H100 experiment report
for the tested stack, measurement method, detailed findings, Slurm
reproduction procedure, 1:1:1 versus 1:6:1 results, component chart, and
acceptance limitations, including the CP4/CP6 single-session comparison.

Mooncake is explicitly initialized with its rdma protocol. On a single node,
the engine may select a topology-local GPU path; the measured effective GB/s
is therefore authoritative for that allocation, while the protocol name alone
must not be presented as proof that traffic traversed an InfiniBand NIC.
Cross-node deployment additionally needs routable stage hostnames, RDMA-capable
NICs, GPUDirect RDMA, and a control plane that forwards the opaque
TensorTransferTicket between stage services.

The design follows the
LightX2V three-stage disaggregation study:
control-plane messages carry only tensor metadata and registered destination
addresses; Mooncake moves the tensor payload directly between device buffers.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[greptile-apps]

Greptile Summary

Adds three-stage LingBot inference disaggregation and benchmark tooling.

• Introduces independently deployable encoder, diffusion, and decoder stages with Mooncake-backed GPU tensor transfers.
• Adds replicated-DiT, context-parallel, aggregated, and transport-only benchmark topologies with reports and plotting utilities.
• Corrects ring attention to use subgroup-local rank and adds regression coverage.
• Adds the optional Mooncake CUDA 13 dependency and extensive benchmark documentation/results.

Confidence Score: 5/5

The PR appears safe to merge, with no concrete changed-code defect identified.

The stage lifecycle, tensor handoffs, subgroup rank handling, collective participation, and documented topology constraints are internally consistent, and the advisory leads do not establish a PR-introduced reachable vulnerability.

Important Files Changed

Filename	Overview
flashdreams/flashdreams/core/attention/cp.py	Uses subgroup-local rank for ring rotation, aligning gathered tensor indexing with non-WORLD process groups.
flashdreams/flashdreams/infra/pipeline/stages.py	Adds stage-specific wrappers while preserving the established generate/finalize cache lifecycle.
flashdreams/flashdreams/infra/transfer.py	Implements synchronous registered-memory tensor bundle transfers, ticket validation, metrics, and explicit cleanup.
integrations/lingbot/lingbot/disagg/stages.py	Splits LingBot conditioning, streaming encoder, DiT cache, and decoder responsibilities across deployable stages.
integrations/lingbot/lingbot/disagg/benchmark.py	Implements the fixed three-rank pipeline benchmark and coordinated Mooncake handoffs.
integrations/lingbot/lingbot/disagg/benchmark_cp.py	Implements a DiT-only context-parallel subgroup with leader handoffs and subgroup tensor fanout.
integrations/lingbot/lingbot/disagg/benchmark_replicated.py	Implements session-affine replicated DiT workers for concurrent throughput measurements.
integrations/lingbot/lingbot/disagg/benchmark_aggregated.py	Implements an aggregated full-pipeline CP8 baseline with token-layout validation and resource reporting.
integrations/lingbot/pyproject.toml	Adds Mooncake as an optional disaggregation dependency.

Sequence Diagram

sequenceDiagram
    participant E as Encoder GPU
    participant D as DiT GPU/CP Group
    participant V as Decoder GPU
    E->>E: Encode conditioning and controls
    E->>D: Mooncake tensor handoff
    D->>D: Denoise and update resident KV cache
    D->>V: Mooncake clean-latent handoff
    V->>V: Decode video frames

Loading

_{Reviews (1): Last reviewed commit: "Add aggregated LingBot CP8 baseline" | Re-trigger Greptile}