diff --git a/AGENTS.md b/AGENTS.md index 1c4624c1..cdf65762 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -27,6 +27,7 @@ Start here, then use the narrower docs for the task in front of you: - `skills/apply-inference-optimizations`: read before porting runtime speedups such as bounded K/V caches, overlap, compile, CUDA graphs, decoder layout changes, or presentation queue tuning into an integration. - `skills/validate-performance-quality`: read before adding benchmark sweeps, quality comparisons, profiler probes, performance summaries, or docs for a performance change. - `skills/flashdreams-postprocessing`: read before adding or modifying video post-processors, postprocess presets, `VideoPostprocessStream`, buffering/layout behavior, or runner postprocess wiring. +- `skills/use-slurm-gpu-job`: read before running builds, test suites, inference, benchmarks, or other resource-intensive work on the Slurm cluster; reuse one allocation and keep the login node lightweight. - `skills/python-docstring-style`: read before adding or polishing Python docstrings, field docstrings, module comments, or SPDX headers. - `skills/maintaining-oss-state`: read before dependency, license, NOTICE, REUSE, or OSS-release collateral changes. - When adding a new `skills//SKILL.md`, update this section so agents can discover when to use it. diff --git a/Dynamo.txt b/Dynamo.txt new file mode 100644 index 00000000..c77fe6c1 --- /dev/null +++ b/Dynamo.txt @@ -0,0 +1,1341 @@ +# Dynamo Review Notes + +This document summarizes the questions and answers from the current Dynamo architecture review. + +> This is a point-in-time architecture note. Versions, CLI flags, examples, and +> implementation paths should be checked against the Dynamo revision used for +> deployment. + +## 1. What is Dynamo? + +Dynamo is an open-source, datacenter-scale inference orchestration layer. It does not replace inference engines such as SGLang, vLLM, or TensorRT-LLM. Instead, it coordinates engine workers across GPUs and nodes. + +Its main responsibilities include: + +* An OpenAI-compatible frontend +* Worker discovery and request routing +* KV-cache-aware routing +* Disaggregated prefill and decode serving +* KV-cache transfer and multi-tier cache management +* SLA-aware capacity planning and autoscaling +* Kubernetes deployment and lifecycle management +* Health checking, graceful shutdown, and fault handling + +The implementation is split across several languages: + +* **Rust:** performance-sensitive runtime, protocols, routing, memory, and KV-cache components +* **Python:** frontend services, backend adapters, configuration, and extensibility +* **Go:** Kubernetes Operator +* **YAML/Helm:** recipes and production deployment definitions + +### High-level architecture + +Dynamo separates the system into three conceptual planes: + +1. **Request plane:** Frontend, Router, and inference workers execute and stream requests. +2. **Control plane:** Planner and Kubernetes Operator manage desired capacity and deployment state. +3. **Storage and events plane:** KV events, KVBM, and NIXL track and move cached inference state. + +A typical request path is: + + Client -> Dynamo Frontend -> Router -> engine worker(s) -> Frontend -> Client + +## 2. How does Dynamo work with inference engines such as SGLang? + +Dynamo uses an adapter-based integration. The inference engine remains a separately packaged dependency, while Dynamo runs it inside a Dynamo-managed worker process. + +For SGLang, the worker is started with: + + python -m dynamo.sglang + +The Dynamo SGLang adapter then: + +1. Parses native SGLang arguments together with Dynamo arguments. +2. Creates an upstream `sglang.Engine`. +3. Registers the worker and model with Dynamo discovery. +4. Advertises the worker's role, capacity, KV block size, and topology. +5. Converts Dynamo requests into SGLang engine calls. +6. Streams SGLang output through the Dynamo frontend. +7. Connects SGLang KV events and metrics to Dynamo routing and planning. +8. Adds health checks, graceful shutdown, cancellation, and distributed lifecycle controls. + +The corresponding adapter is under `components/src/dynamo/sglang/`. + +## 3. Does Dynamo include the SGLang source code? + +No. The Dynamo repository and Python wheel do not vendor SGLang's source code. + +There are three normal installation paths: + +### Prebuilt container + +The production CUDA image is based on an upstream SGLang runtime image. In the reviewed tree, the configured upstream base is: + + lmsysorg/sglang:v0.5.14-cu130-runtime + +Dynamo installs its own wheels into that image using `pip install --no-deps`, preserving the SGLang, CUDA, PyTorch, and related dependency stack already present in the upstream image. + +Kubernetes pulls the completed `sglang-runtime` image. It does not normally install SGLang during pod startup. + +### Python package installation + +Installing the SGLang extra installs a compatible, pinned SGLang package: + + uv pip install "ai-dynamo[sglang]" + +At the time of this review, the dependency is: + + sglang[diffusion] == 0.5.14 + +### Development installation + +Developers can clone SGLang separately and install it as an editable package. Dynamo and SGLang remain separate repositories. + +## 4. Must SGLang be modified to work with Dynamo? + +For the supported CUDA configuration, Dynamo does not require a maintained SGLang fork or a runtime patch set. It uses the upstream SGLang package and image. + +However, Dynamo relies on particular SGLang APIs and features, so arbitrary SGLang versions are not guaranteed to work. Dynamo pins a compatible release and maintains narrow compatibility shims for changes in SGLang's pre-1.0 internal APIs. + +Examples of required or consumed SGLang capabilities include: + +* `sglang.Engine` +* Prefill/decode modes +* Disaggregation bootstrap coordination +* KV-cache transfer +* KV-event publication +* Metrics and forward-pass information +* Cancellation and engine control APIs + +Compatibility handling is centralized in `components/src/dynamo/sglang/_compat.py` rather than applied as modifications to SGLang source. + +## 5. Does SGLang support disaggregated serving? + +Yes. SGLang supplies the engine-side prefill/decode separation, bootstrap handshake, and KV-cache transfer mechanisms. Dynamo supplies discovery, worker selection, routing, independent scaling, and lifecycle management around those mechanisms. + +### What are prefill and decode? + +LLM inference has two major phases: + +* **Prefill:** Processes all input tokens and constructs the attention KV cache. This phase is generally compute-intensive and highly parallel. +* **Decode:** Generates output one token at a time using the KV cache. This phase is generally constrained by memory bandwidth, concurrency, and KV-cache capacity. + +In aggregated serving, one worker performs both phases: + + GPU 0: prompt prefill -> token 1 -> token 2 -> token 3 -> ... + +In disaggregated serving, different workers perform the phases: + + GPU 0, prefill worker: prompt -> KV cache + | + transfer KV + v + GPU 1, decode worker: token 1 -> token 2 -> token 3 -> ... + +Both workers normally load the same model weights. The per-request KV cache, not the model, is transferred between them. + +## 6. Concrete SGLang disaggregation example + +The repository includes a two-GPU Qwen example in `examples/backends/sglang/launch/disagg.sh`: + +* GPU 0: SGLang prefill worker +* GPU 1: SGLang decode worker +* NIXL: KV-cache transfer +* Dynamo Frontend: OpenAI-compatible API on port 8000 + +Simplified commands are: + + # Frontend and routing + python -m dynamo.frontend + + # GPU 0: prefill only + CUDA_VISIBLE_DEVICES=0 python -m dynamo.sglang \ + --model-path Qwen/Qwen3-0.6B \ + --disaggregation-mode prefill \ + --disaggregation-bootstrap-port 12345 \ + --disaggregation-transfer-backend nixl + + # GPU 1: decode only + CUDA_VISIBLE_DEVICES=1 python -m dynamo.sglang \ + --model-path Qwen/Qwen3-0.6B \ + --disaggregation-mode decode \ + --disaggregation-bootstrap-port 12345 \ + --disaggregation-transfer-backend nixl + +For a request containing a 20,000-token document and asking for a 500-token summary: + +1. Dynamo selects a decode worker to own and stream the request. +2. Dynamo selects a prefill worker, potentially using load and cache overlap. +3. The decode worker reserves GPU pages for the incoming KV cache. +4. The prefill worker processes the 20,000-token prompt. +5. SGLang transfers the resulting KV cache to decode through NIXL, commonly over RDMA in a production multi-node cluster. +6. The decode worker generates the 500 output tokens. +7. Tokens stream back through the Dynamo Frontend. + +Dynamo contacts the decode side first for coordination, but model computation still occurs in the logical order of prefill followed by decode. + +## 7. Why disaggregate? + +Prefill and decode may need different amounts or configurations of hardware. Dynamo can scale their pools independently, for example: + + Long prompts, short answers: 4 prefill workers + 2 decode workers + Short prompts, long answers: 2 prefill workers + 6 decode workers + +Workers may also use different parallelism: + + Prefill pool: 2 replicas x tensor parallel size 4 + Decode pool: 4 replicas x tensor parallel size 2 + +Tensor parallelism and disaggregation are different: + +* **Tensor parallelism** splits one model forward pass across GPUs. +* **Disaggregation** assigns prefill and decode to separate worker pools. + +They can be used together. + +Disaggregation is most useful for large models, long prompts, high concurrency, or workloads whose input/output mix changes. It is not automatically faster for small models: transfer and coordination overhead can exceed the utilization benefit. + +## Relevant source locations + +* `README.md` +* `docs/design-docs/architecture.md` +* `docs/backends/sglang/README.md` +* `docs/backends/sglang/sglang-disaggregation.md` +* `components/src/dynamo/frontend/main.py` +* `components/src/dynamo/sglang/main.py` +* `components/src/dynamo/sglang/init_llm.py` +* `components/src/dynamo/sglang/register.py` +* `components/src/dynamo/sglang/_compat.py` +* `container/templates/sglang_runtime.Dockerfile` +* `examples/backends/sglang/launch/disagg.sh` + +## 8. What is NIXL? + +NIXL is the **NVIDIA Inference Transfer Library**. It is a data-movement layer for inference workloads, not a network protocol itself. It can register GPU or host memory and initiate asynchronous reads or writes through the best transport available in the environment. + +Depending on hardware and configuration, the underlying path can include: + +* GPU-to-GPU transfer within a machine +* GPUDirect RDMA over InfiniBand or RoCE +* UCX or libfabric transports +* Host-memory staging +* Slower TCP fallback +* Storage backends for other Dynamo data-movement use cases + +For disaggregated serving, NIXL allows a prefill worker to place a request's KV-cache tensors into memory reserved by a decode worker. With GPUDirect RDMA, the transfer can go between GPU memory on different hosts without copying the payload through application CPU memory. + +## 9. How large is a KV cache? + +For a conventional transformer, a useful approximation for the total KV cache is: + + bytes = tokens + x layers + x KV heads + x head dimension + x 2 # key and value + x bytes per element + +The model architecture and KV data type therefore matter greatly. Grouped-query attention reduces the number of KV heads; FP8 halves storage relative to BF16/FP16; MLA models use a different, compressed representation. + +Approximate BF16 examples: + +| Model shape | KV per token | KV for a 20K-token prompt | +| --- | --- | --- | +| 32 layers, 8 KV heads, head dimension 128 | 128 KiB | 2.44 GiB | +| 80 layers, 8 KV heads, head dimension 128 | 320 KiB | 6.10 GiB | + +These are aggregate full-model figures. Tensor parallelism shards the cache and transfer across ranks, changing the per-GPU amount but not eliminating the aggregate data volume. During disaggregation, the prompt KV is transferred; KV for subsequently generated output tokens grows locally on the decode side. + +## 10. Can KV transfer become the bottleneck? + +Yes. A rough lower bound is: + + transfer time = KV bytes / effective transfer bandwidth + setup overhead + +For example, moving a 6.10 GiB cache over a nominal 200 Gbit/s link has an ideal serialization time of roughly 0.26 seconds. Real effective time is higher because nominal link rate is not payload throughput and transfers have registration, coordination, contention, and topology costs. + +Transfer is more likely to dominate when: + +* Prompts are long. +* KV uses BF16/FP16 rather than a smaller representation. +* The model has many layers or KV heads. +* Workers are cross-node. +* RDMA is unavailable and the path falls back to TCP or host staging. +* Multiple prefill workers concurrently target the same network or decode worker. +* Prefill computation is short relative to the amount of KV data produced. + +This is why Dynamo's documentation treats validation of RDMA/UCX/NIXL as an early deployment step. Disaggregation is beneficial only when better phase-specific utilization outweighs the transfer cost. Aggregated serving can be faster for small models, short prompts, low concurrency, or weak interconnects. + +## 11. How does Dynamo know prompt and answer lengths? + +The two lengths have different certainty: + +* **Input sequence length (ISL):** Known for an individual request after the frontend tokenizes the prompt. The Router can immediately use it for request placement and load accounting. +* **Output sequence length (OSL):** Not known exactly until generation finishes. `max_tokens` is only an upper bound. A caller can optionally supply `nvext.agent_hints.osl` as an expected output length for routing, but this is a hint rather than a guarantee. + +The Planner does not normally resize the cluster independently for every incoming request. It sizes pools from aggregate workload observations and predictions. + +The frontend records Prometheus metrics including: + +* Started and completed request counts +* Average input sequence tokens +* Average output sequence tokens +* TTFT and ITL +* Request duration and concurrency +* KV hit rate and engine load signals when available + +The output-length metric comes from completed or progressing requests, so it describes observed workload history rather than foreknowledge of a new answer. + +## 12. How does Dynamo decide prefill and decode replica counts? + +For SLA-based throughput planning, the control loop is approximately: + + Observed request rate, ISL and OSL + | + v + Predict next request rate, ISL and OSL + | + v + Engine performance model under TTFT/ITL targets + | + v + Per-worker prefill RPS and decode RPS + | + v + ceil(predicted demand RPS / per-worker capacity RPS) + | + v + Apply minimum replicas and global GPU budget + +Dynamo can use constant, ARIMA, Kalman, or Prophet predictors for request count, ISL, and OSL. The engine-capacity model is bootstrapped from AIConfigurator estimates, profiler/self-benchmark data, or live ForwardPassMetrics. + +Prefill and decode are modeled separately: + +* **Prefill capacity** depends strongly on ISL, KV hit rate, engine configuration, and the TTFT target. +* **Decode capacity** depends on ISL + OSL context growth, concurrency/KV pressure, engine configuration, speculative acceptance when used, and the ITL target. + +An illustrative decision might be: + + Predicted traffic: 20 requests/second + Prefill capacity at predicted ISL: 4 requests/second/worker + Decode capacity at predicted OSL: 2 requests/second/worker + + Prefill replicas = ceil(20 / 4) = 5 + Decode replicas = ceil(20 / 2) = 10 + +The actual values come from model- and hardware-specific performance data rather than fixed ratios. + +Dynamo also has a faster load-based loop. It consumes live ForwardPassMetrics, queue pressure, KV utilization, and estimated latency to react to bursts. When both loops are active, predicted throughput establishes a capacity floor and live load scaling can temporarily scale above that floor. Default documented cadences are much slower for throughput planning (around 180 seconds) and faster for load reactions (around 5 seconds), because starting a GPU worker is expensive. + +Relevant implementation files include: + +* `docs/api/nixl-connect/README.md` +* `docs/features/disaggregated-serving/README.md` +* `components/src/dynamo/planner/monitoring/traffic_metrics.py` +* `components/src/dynamo/planner/plugins/builtins/local_planner.py` +* `components/src/dynamo/planner/core/throughput_scaling.py` + +## 13. Which SGLang APIs enable disaggregation and KV-aware routing? + +There is no single `enable_dynamo()` SGLang API. Dynamo composes several upstream SGLang configuration fields and engine APIs. + +### Engine construction and role configuration + +Dynamo parses upstream SGLang `ServerArgs` and constructs the normal engine: + + server_args = ServerArgs.from_cli_args(parsed_args) + engine = sglang.Engine(server_args=server_args) + +The important upstream arguments are: + + --disaggregation-mode prefill|decode + --disaggregation-transfer-backend nixl + --disaggregation-bootstrap-port 12345 + --kv-events-config '{"publisher":"zmq","topic":"kv-events","endpoint":"tcp://*:5557"}' + +The first three enable SGLang's prefill/decode role and its transfer path. `kv_events_config` asks the SGLang scheduler to publish cache lifecycle events. + +### Per-request disaggregation contract + +SGLang's main execution API is `Engine.async_generate()`. For disaggregated requests, both peers receive the same bootstrap triple: + + stream = await engine.async_generate( + input_ids=token_ids, + sampling_params=sampling_params, + stream=True, + bootstrap_host=bootstrap_host, + bootstrap_port=bootstrap_port, + bootstrap_room=bootstrap_room, + data_parallel_rank=selected_dp_rank, + rid=request_id, + ) + +The triple identifies the transfer rendezvous: + +* `bootstrap_host`: host on which SGLang's transfer/bootstrap service is reachable +* `bootstrap_port`: bootstrap service port +* `bootstrap_room`: unique per-request rendezvous identifier + +Dynamo obtains the prefill worker's address through SGLang's live engine state: + + engine.tokenizer_manager.server_args.disaggregation_bootstrap_port + +Dynamo generates or forwards a `bootstrap_room`, sends the same triple to the selected prefill and decode workers, and calls `async_generate()` on both. SGLang registers the room, allocates the decode-side cache pages, performs prefill, and executes the configured NIXL KV transfer. Dynamo must continue draining the prefill stream until the transfer finishes. + +`data_parallel_rank` is also important: Dynamo can choose a specific SGLang DP rank and pass that decision into `async_generate()`, keeping request placement aligned with per-rank cache events and load state. + +### KV-event API used for routing + +With `kv_events_config` enabled, SGLang schedulers publish KV block creation/deletion events over ZMQ. Dynamo imports SGLang's event helper: + + from sglang.srt.disaggregation.kv_events import ZmqEventPublisher + +For data-parallel workers, the helper derives the correct publisher endpoint for each rank: + + rank_endpoint = ZmqEventPublisher.offset_endpoint_port(base_endpoint, dp_rank) + +Dynamo subscribes to those endpoints, associates the events with the Dynamo worker and DP rank, and feeds them into its KV indexer. The indexer records which hashed prompt blocks are present on which workers. + +For a new request, Dynamo—not SGLang—then: + +1. Tokenizes the prompt. +2. Hashes it in KV-cache-sized blocks. +3. Looks up matching cached blocks in the event-maintained index. +4. Combines cached overlap with active prefill/decode load. +5. Selects the lowest-cost worker and DP rank. +6. Passes that rank to SGLang through `data_parallel_rank`. + +Thus SGLang provides cache events and rank-targeted execution, while Dynamo provides the cache-aware selection algorithm. + +Relevant source locations: + +* `components/src/dynamo/sglang/args.py` +* `components/src/dynamo/sglang/_disagg.py` +* `components/src/dynamo/sglang/llm_engine.py` +* `components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py` +* `components/src/dynamo/sglang/request_handlers/llm/decode_handler.py` +* `components/src/dynamo/sglang/publisher.py` +* `lib/llm/src/kv_router/` +* `examples/backends/sglang/launch/disagg_router.sh` + +## 14. What would be required to run FlashDreams behind Dynamo? + +### Current FlashDreams execution model + +FlashDreams is a video/world-model inference framework rather than an LLM token server. Its central API is a persistent autoregressive pipeline: + + cache = pipeline.initialize_cache(...) + for i in range(total_blocks): + video_chunk = pipeline.generate(i, cache, input=control) + pipeline.finalize(i, cache) + +`initialize_cache()` creates per-rollout state, `generate()` runs one autoregressive video chunk, and `finalize()` advances the transformer KV cache. + +FlashDreams already supports multi-GPU execution through `torchrun` and context parallelism. Every rank constructs a sharded pipeline and executes `generate()` / `finalize()` in lockstep. Rank zero alone performs final user-facing I/O. + +This is multi-GPU model execution, but it is not by itself a general multi-request batch server. FlashDreams documentation describes serving as integration-specific, and its LingBot reference server currently supports one active WebRTC session per process. + +### Recommended first architecture + + Client + | + v + Dynamo Frontend (/v1/videos) + | + | least-loaded or round-robin + v + FlashDreams worker-group leader (rank 0) + | + | broadcast request/control command + v + FlashDreams context-parallel ranks 0..N-1 + | + | all ranks run initialize_cache/generate/finalize + v + Rank 0 encodes/uploads result and returns response + +One Dynamo worker replica should represent one complete FlashDreams GPU group, not one GPU rank. Only rank zero should register the Dynamo endpoint. Otherwise Dynamo could route unrelated requests independently to ranks that must participate in the same collectives. + +### Required addition 1: Dynamo video backend adapter + +Add a small package, ideally outside the Dynamo repository initially, that wraps a selected FlashDreams pipeline. The proven path for video workloads is Dynamo's lower-level Python worker API, similar to the existing FastVideo custom worker. + +The adapter should: + +* Initialize the chosen FlashDreams pipeline once at process startup. +* Register `ModelInput.Text` and `ModelType.Videos`. +* Serve a `dynamo..generate` endpoint. +* Accept Dynamo's `/v1/videos` request fields: model, prompt, input reference, size, seconds, FPS, frame count, inference steps, seed, and response format. +* Translate those fields into a FlashDreams integration/pipeline configuration. +* Run `initialize_cache()`, repeated `generate()` / `finalize()`, and cleanup. +* Return an `NvVideosResponse` with a URL or base64 MP4. +* Implement health readiness only after model loading, compilation, and warmup complete. +* Implement cancellation and bounded queueing. + +For large video outputs, object storage plus a returned URL is preferable to transporting base64 MP4 through every service layer. + +### Required addition 2: distributed request command loop + +Under `torchrun`, all context-parallel ranks must enter the same collectives in the same order. A Dynamo request reaches rank zero only, so the adapter needs a group command protocol: + +1. Rank zero receives and validates the request. +2. Rank zero broadcasts a fixed request descriptor to all ranks. +3. Every rank constructs equivalent per-request cache state. +4. Every rank executes each autoregressive step in lockstep. +5. Rank zero gathers/owns the user-facing video output. +6. Errors, cancellation, shutdown, and the next command are broadcast to every rank. + +FlashDreams' multi-GPU WebRTC server provides a useful precedent for rank-zero networking with all-rank inference, but a Dynamo adapter needs this behavior behind a Dynamo endpoint rather than WebRTC. + +### Required addition 3: serving scheduler + +For the minimum viable integration, serialize requests with one in-flight request per GPU group. Scale throughput by deploying several identical groups and let Dynamo load-balance between their leaders. + +True engine-level batching requires additional FlashDreams work: + +* A bounded pending-request queue +* Admission control based on resolution, frames, rollout blocks, model variant, and memory +* Compatibility buckets for requests that can share one tensor batch and CUDA graph +* Batched prompt/image/control encoding +* A batched `StreamInferencePipelineCache`, or a scheduler that stacks compatible per-request caches +* Per-request active masks when rollouts have different lengths +* Result demultiplexing and per-request cancellation +* Fairness and timeout behavior + +Although FlashDreams tensor shapes include batch dimensions, the public runners shown in the repository drive one rollout at a time. Dynamo does not automatically combine concurrent requests into a tensor batch; that batching logic belongs inside the FlashDreams engine adapter/framework. + +### Required addition 4: runtime image and deployment + +Build a custom CUDA 13 image containing: + +* A pinned FlashDreams release and selected integration package +* The matching PyTorch/CUDA/Transformer Engine stack +* Dynamo runtime and Python package +* FFmpeg or another production video encoder +* Model/cache volume support +* The adapter entry point + +For a same-node four-GPU group, the Dynamo deployment should request four GPUs for one worker pod and launch approximately: + + torchrun --standalone --nproc_per_node=4 -m flashdreams_dynamo.worker + +Set `replicas` to the desired number of four-GPU groups. Multi-node context parallelism would additionally require gang scheduling, stable rendezvous, cross-node NCCL networking, and failure handling for the entire rank group; that should be a later phase. + +### Routing and scaling limitations + +Do not initially enable Dynamo's LLM KV-aware router for FlashDreams. FlashDreams' cache is per-video-rollout state, not the standardized token-prefix block cache expected by Dynamo's KV indexer. For independent batch jobs, use least-loaded or round-robin routing. For interactive world-model sessions, add session-affinity routing so every action returns to the worker group holding that session's cache. + +Likewise, the current Dynamo Planner is oriented around request rate, ISL/OSL, TTFT, ITL, token queues, and LLM KV utilization. A FlashDreams deployment should initially use fixed replica counts or queue/latency-based autoscaling. Native Planner integration would require video-specific capacity metrics or a custom Planner plugin using fields such as queued jobs, active sessions, pixels x frames x denoising steps, generation latency, and GPU memory. + +Prefill/decode disaggregation should not be an MVP goal. FlashDreams' pattern is closer to `initialize -> generate chunk -> finalize -> generate chunk`, with a session-private cache that evolves each step. Supporting cross-worker cache movement would require defining serialization and transfer contracts for every integration's encoder, transformer, and decoder caches. + +### Suggested delivery phases + +1. **MVP:** One model, one request per N-GPU group, `/v1/videos`, fixed replicas, round-robin/least-loaded routing. +2. **Production hardening:** readiness, cancellation, queue limits, URL-based outputs, metrics, Kubernetes manifests, and graceful group shutdown. +3. **Replica throughput:** multiple N-GPU groups with Dynamo load balancing. +4. **True batching:** compatible-shape microbatch scheduler and batch-aware cache handling. +5. **Interactive sessions:** session IDs, sticky routing, state expiry, reconnection, and admission limits. +6. **Optional advanced work:** custom Planner plugin, cache checkpoint/migration, or pipeline-stage disaggregation. + +External references reviewed: + +* +* +* +* +* + +## 15. How should FlashDreams refactor its serving API for autoregressive world models? + +### Do not force the LLM prefill/decode abstraction + +FlashDreams has a different state machine from an LLM: + + LLM request: prefill(prompt) -> decode token -> decode token -> done + World-model session: initialize(context) -> advance chunk -> advance chunk -> ... + +`initialize_cache()` is similar to prefill only in the broad sense that it creates persistent state. It is not necessarily a separable compute pool, and the cache evolves after every video chunk. FlashDreams should therefore model a stateful session explicitly rather than rename initialization to prefill or introduce a decode role prematurely. + +### Recommended engine contract + +Use a transport-neutral session API: + + class WorldModelEngine: + async def start(self) -> EngineCapabilities: ... + + async def create_session( + self, request: CreateSessionRequest + ) -> SessionHandle: ... + + async def advance( + self, session: SessionHandle, request: StepRequest + ) -> AsyncIterator[VideoChunk]: ... + + async def cancel_step(self, session: SessionHandle) -> None: ... + + async def close_session(self, session: SessionHandle) -> None: ... + + async def health(self) -> HealthStatus: ... + +`create_session()` owns FlashDreams' `initialize_cache()`. `advance()` should encapsulate both `pipeline.generate()` and `pipeline.finalize()` as one state transition: + + async def advance(session, request): + output = pipeline.generate(session.next_index, session.cache, request.control) + pipeline.finalize(session.next_index, session.cache) + session.next_index += 1 + return output + +The public serving layer should not ask clients or transports to invoke `finalize()` separately. The next step must not start until the previous cache update is committed, even if finalization is internally overlapped with encoding or network delivery. + +### Session state model + +Each session should have: + +* Opaque session ID +* Owning worker group and model/config identity +* Pipeline cache and next autoregressive index +* Resolution, FPS, chunk shape, seed, and batching class +* Lifecycle state: creating, ready, advancing, closing, closed, or failed +* Idle TTL and maximum lifetime +* Cancellation state and idempotency keys +* Optional checkpoint/migration metadata + +The cache should remain opaque to Dynamo and transport code. Routing should use session affinity: after creation, every step returns to the worker group that owns the cache. + +### Scheduler model + +The equivalent of continuous batching is step-level scheduling across sessions: + + Session A step 7 ┐ + Session B step 3 ├─ compatible microbatch -> one multi-GPU execution + Session C step 11 ┘ + +Requests can batch only when their execution shapes are compatible, including model variant, resolution, chunk frame count, denoising schedule, guidance mode, precision, and context-parallel group. The scheduler needs deadline-aware admission because an interactive session cannot wait indefinitely for a larger batch. + +A useful policy is: + +* Interactive queue: short batching window and latency deadline +* Batch-job queue: longer batching window for utilization +* Weighted fairness between the two classes +* Fixed shape buckets compatible with CUDA graphs +* Per-session active masks and output demultiplexing + +Initialization should be scheduled separately from `advance()` because it may include text/image encoding and large cache allocation. It can later be split into stateless preprocessing and stateful activation, but this should be driven by profiling rather than by analogy to LLM prefill/decode. + +### Two serving products over one engine + +#### Finite video-generation jobs + +Expose an asynchronous REST resource: + + POST /v1/videos create queued job + GET /v1/videos/{id} status/progress + GET /v1/videos/{id}/content MP4 or preview content + DELETE /v1/videos/{id} cancel/delete + GET /v1/videos/{id}/events optional SSE progress events + +The create call should return quickly with a job ID and state such as `queued`. Completed output should normally live in object storage and be returned as a signed URL. This matches the common shape used by OpenAI video jobs, Google Veo long-running operations, Amazon Bedrock async invocation, and Runway tasks. + +#### Persistent interactive sessions + +Expose a session control plane plus interchangeable streaming transports: + + POST /v1/world-model/sessions + GET /v1/world-model/sessions/{id} + DELETE /v1/world-model/sessions/{id} + +Then support: + +* WebRTC media tracks for low-latency browser playback +* WebRTC DataChannel for browser controls and step acknowledgements +* WebSocket for simpler bidirectional SDK clients +* Bidirectional gRPC for native simulation, robotics, and datacenter clients + +WebRTC signaling, ICE/TURN, codecs, congestion control, and browser integration belong in a transport adapter. They should not be embedded in the pipeline or scheduler API. + +### Should FlashDreams focus only on WebRTC? + +No. WebRTC is the right primary demo and browser transport for interactive world models because it standardizes real-time media and generic data channels. It is a poor sole API for offline batch generation, backend-to-backend calls, job orchestration, and many native robotics/simulation clients. + +Recommended product priority: + +1. Build and stabilize the transport-neutral session engine. +2. Provide asynchronous `/v1/videos` jobs for broad batch/API compatibility. +3. Keep WebRTC as the first-class interactive browser adapter. +4. Add WebSocket or bidirectional gRPC for native/headless clients. + +### Current industry pattern + +There is no single universal generative-video API standard. The strongest current conventions are: + +* **Finite generation:** asynchronous job creation, job ID, polling/webhook, progress state, cancellation, and downloadable/object-storage output. +* **Interactive generation:** persistent session with bidirectional controls and streamed media; WebRTC is standard for browser real-time media, while WebSocket/gRPC are common application transports outside the browser. +* **Request fields:** prompt, model, image/reference input, duration, resolution/aspect ratio, seed, sample count, and provider-specific generation settings. +* **Output:** MP4 or another encoded asset via URL/content endpoint, not a long-held synchronous JSON response. + +The OpenAI-style `/v1/videos` job resource is a practical compatibility target for finite generation. It should not be stretched to represent an endless interactive world session; that deserves a separate session API. + +External references reviewed: + +* +* +* +* +* + +## 16. Is “prefill, decode, and KV-cache-transfer worker pools” an accurate description? + +Broadly yes, with one correction: KV-cache transfer is normally a data movement path between the prefill and decode workers, not a third inference worker pool. + +A more precise statement is: + +> Disaggregated LLM serving places prefill and decode computation in independently scalable worker pools. Dynamo routes and coordinates requests across those pools, while the inference backend and a transport such as NIXL move the generated KV-cache state from the selected prefill worker to the selected decode worker. + +The responsibilities are: + +* **Prefill workers:** process prompt tokens and create the prompt's KV cache. +* **Decode workers:** consume that cache and generate output tokens autoregressively. +* **Dynamo:** provides the frontend, discovery, routing, coordination, metrics, and deployment/orchestration integration needed to select and connect the workers. +* **Inference backend:** SGLang, TensorRT-LLM, vLLM, or another backend performs the model computation and participates in the disaggregated protocol. +* **KV transport:** NIXL or a backend transport transfers/registers the cache memory; it is a mechanism rather than usually a pool that runs model inference. + +Some deployments may add separate cache-storage or offload services, but those are optional and should not be treated as part of the basic definition of prefill/decode disaggregation. + +## 17. Dynamo introduction: three slide bullets + +* **Distributed inference framework:** NVIDIA Dynamo turns inference engines such as SGLang, TensorRT-LLM, and vLLM into scalable, distributed serving systems. +* **Intelligent request orchestration:** It discovers workers and routes requests using load and KV-cache locality, including coordination of disaggregated prefill and decode. +* **Production infrastructure:** It provides APIs, observability, autoscaling, and Kubernetes deployment components while leaving model execution to the underlying inference engine. + +## 18. Does Dynamo support image and video generation? + +Yes. This is **media generation**, which is separate from VLM multimodal input/understanding. + +### SGLang + +Dynamo wraps SGLang's `DiffGenerator` with dedicated media workers: + +* `--image-diffusion-worker` exposes `POST /v1/images/generations`. +* `--video-generation-worker` exposes `POST /v1/videos`. +* Documented examples include FLUX text-to-image and Wan 2.1 text-to-video; the video handler also has image-to-video support. +* Output can be returned as base64 or written to local/S3-compatible storage and returned by URL. +* SGLang `DiffGenerator` can use tensor and data parallelism across GPUs. + +The important limitation is that Dynamo currently treats each SGLang image/video generator as an **aggregated media worker**. The whole generation pipeline belongs to that worker group. There is no LLM-style prefill/decode split, KV-cache transfer, or KV-aware routing for these jobs. Multiple worker replicas can still be registered for request-level scaling and load distribution. + +### Other generation backends + +* **vLLM-Omni:** text-to-image, text-to-video, image-to-video, and text-to-audio. It also supports some explicitly multi-stage pipelines, such as disaggregated GLM-Image. +* **TensorRT-LLM diffusion:** text-to-image and text-to-video; currently documented as experimental, with NVENC required for video output. +* **FastVideo:** a custom Dynamo worker for production-oriented text-to-video Kubernetes deployments. +* **Custom backends:** Dynamo's `DiffusionEngine` interface and image/video protocols allow another framework, such as FlashDreams, to register `ModelType.Images` or `ModelType.Videos` and use the same frontend endpoints. + +Current deployment caveat: the built-in SGLang, vLLM-Omni, and TensorRT-LLM diffusion integrations primarily ship CLI/local launch examples rather than ready-made Kubernetes recipes. FastVideo includes a Kubernetes path. + +## 19. How vLLM-Omni multi-stage disaggregation works in Dynamo + +This is **pipeline-stage disaggregation**, not the LLM prefill/decode form of disaggregation. A heterogeneous model is divided at its natural model boundaries, and every stage becomes an independently served Dynamo worker pool. + +The current concrete Dynamo example is GLM-Image: + + POST /v1/images/generations + | + Dynamo frontend + | + public Omni stage router + | + v + Stage 0: AR model, GPU 0 + generates prior image token IDs + | + | intermediate output through an OmniConnector + | router forwards only an opaque connector reference + v + Stage 1: ar2diffusion processor + DiT, GPU 1 + performs diffusion denoising and VAE decoding + | + v + router formats image response + +### Configuration and process layout + +The vLLM-Omni stage YAML describes: + +* the stage IDs and types, such as autoregressive LLM or diffusion; +* each stage's model/runtime and GPU configuration; +* `engine_input_source`, which declares upstream dependencies; +* the model-specific transition function, such as `ar2diffusion`; +* the connector used on each inter-stage edge. + +Dynamo launches one process per stage using `--stage-id N` and a separate coordinator using `--omni-router`. In the GLM-Image example, Stage 0 and Stage 1 are bound to GPUs 0 and 1 respectively. + +Each stage worker exposes a private endpoint: + + dyn:////generate + +Only the Omni router registers the public model endpoint with the Dynamo frontend. From the normal frontend's perspective, that router is one aggregated model worker, while the stage graph behind it is private. + +### Control plane versus data plane + +For each request, the router walks the stage graph: + +1. It sends the original request to Stage 0 using a Dynamo endpoint client. +2. Stage 0 runs its engine and writes the intermediate output to the configured vLLM-Omni connector. +3. Stage 0 returns only connector metadata—an address ticket—to the router. +4. The router forwards that opaque ticket to Stage 1 without reading or serializing the tensor payload. +5. Stage 1 calls `connector.get`, reconstructs the upstream output, applies `ar2diffusion`, and invokes its DiT engine. +6. The final image is returned to the router, which encodes/stores it and creates the API response. + +Connector references accumulate, allowing a later stage to consume outputs from more than one preceding stage. The stage YAML's `engine_input_source` selects which references it reads. + +This differs from LLM P/D disaggregation: + +| LLM P/D | vLLM-Omni stage disaggregation | +| --- | --- | +| Splits one transformer into prefill and decode | Splits a heterogeneous model at semantic stage boundaries | +| Moves attention KV cache | Moves stage outputs such as token IDs, embeddings, latents, or media tensors | +| Commonly uses NIXL KV connectors | Uses vLLM-Omni `OmniConnector` edges | +| Decode repeatedly generates tokens | Each stage executes its own engine and transition processor | + +### Independent scaling + +Every stage endpoint is independently discoverable. Additional workers with the same stage ID form replicas of that stage; the router uses round-robin selection within that stage pool. This makes it possible to assign more replicas or more intra-stage parallelism to the slow DiT stage without duplicating the AR stage at the same ratio. + +### Current Dynamo limitations + +* The tested Dynamo disaggregated example is currently GLM-Image, AR to DiT, and is documented as experimental. +* Inter-stage streaming (`async_chunk=true`) is not supported; each stage finishes before its connector reference is sent onward. +* If connector edges are omitted, Dynamo synthesizes shared-memory connectors. +* The final-stage-to-router result currently uses shared memory, which requires the final worker and router to be on the same host. Although connector types such as Mooncake can support remote inter-stage movement, the final hop is still a single-node limitation in the current implementation. +* This path does not publish KV events or use Dynamo's KV-aware LLM router. +* The local launcher creates processes directly; it is not yet a ready-made `DynamoGraphDeployment`. + +## 20. GLM-Image stages, video capability, and transfer size + +### What “GLM” means + +GLM originally stands for **General Language Model**, the model family introduced with autoregressive blank-infilling pretraining. Here, `GLM-Image` is Z.ai's image-generation member of that family. It initializes its autoregressive image generator from GLM-4-9B and adds visual tokens plus a diffusion decoder. + +### Dynamo does not invent the AR/DiT split + +GLM-Image itself is trained as a hybrid, two-stage model: + + prompt + | + v + 9B autoregressive generator + | discrete prior image tokens + v + 7B DiT diffusion decoder + glyph/text encoder + VAE + | + v + image + +Dynamo and vLLM-Omni expose that native model boundary as two independently served stages. This is useful because: + +* the AR stage and DiT stage use different execution engines and scheduling strategies; +* their weights and working memory do not have to fit on the same GPU; +* the stages have different bottlenecks and can be replicated independently; +* AR token generation can use vLLM's token scheduler while DiT uses diffusion-specific batching and parallelism. + +Published vLLM-Omni estimates for GLM-Image are roughly 18 GiB plus KV cache for Stage 0 and 20 GiB for Stage 1. + +### What each stage does + +**Stage 0: autoregressive prior generator** + +* Reads the prompt and, for editing, source image information. +* Uses a 9B GLM-derived causal transformer. +* Sequentially generates a compact visual plan followed by higher-resolution discrete visual tokens. +* For a square text-to-image request, the current processor expects a 16×16 preview token grid plus a target grid downsampled by 32, followed by EOS. + +**Transition processor: `ar2diffusion`** + +* Extracts the AR stage's cumulative token IDs. +* Selects the target-grid tokens and upsamples their token layout from 32× downsampling to the 16× layout expected by the DiT. +* Builds the DiT request containing the original prompt, resolution, prior token IDs, and optional image-editing tokens. + +**Stage 1: diffusion decoder** + +* Encodes prompt/glyph text used for accurate text rendering. +* Embeds the AR-generated prior tokens as semantic/layout conditioning. +* Starts from latent noise and applies the DiT denoising loop. +* Uses the VAE to decode the final latent into pixels. + +In short, the AR stage decides broadly **what and where**; the DiT stage renders **pixels, texture, detail, and text appearance**. + +### Can GLM-Image generate video? + +No. GLM-Image supports text-to-image and image-to-image/editing, not video. + +vLLM-Omni supports video-generation models such as Wan, and Dynamo exposes them through `/v1/videos`. However, Dynamo's currently tested multi-stage disaggregated vLLM-Omni example is GLM-Image AR→DiT. Existing Dynamo video examples are aggregated video workers; there is not yet a validated multi-stage disaggregated video recipe. + +### How much data moves between the GLM-Image AR and DiT workers? + +For text-to-image, the payload is small because Stage 0 sends **discrete token IDs**, not KV cache, hidden states, latents, or pixels. + +The current upstream processor computes the number of generated AR tokens as: + + large target tokens = floor(H / 32) × floor(W / 32) + small preview tokens ≈ 256 for a square image + total T2I tokens = large target tokens + preview tokens + 1 EOS + +Examples: + +| Resolution | AR tokens crossing the stage boundary | Packed as uint16 | As int32 | As int64 | +| --- | --- | --- | --- | --- | +| 1024×1024 | 1,281 | ~2.5 KiB | ~5 KiB | ~10 KiB | +| 1536×1536 | 2,561 | ~5 KiB | ~10 KiB | ~20 KiB | +| 2048×2048 | 4,353 | ~8.5 KiB | ~17 KiB | ~34 KiB | + +Dynamo's current connector serializes the containing vLLM request-output object rather than a bare packed array, so the real connector payload includes object and protocol metadata. A reasonable order-of-magnitude estimate for ordinary text-to-image is **tens of KiB per request**, and it should be measured for an exact build. + +This makes AR→DiT transfer dramatically smaller than multi-GiB LLM KV-cache transfers. It is unlikely to be a network bandwidth bottleneck; serialization, synchronization, stage imbalance, and inference time matter more. + +Two qualifications: + +* Image-to-image requests may also carry source-image-related VQ token IDs and source media, making the payload larger. +* The final DiT result sent to the router is an image rather than token IDs. A raw 1024×1024 RGB image is about 3 MiB (and may differ after serialization/compression). In the current Dynamo example, both inter-stage and final-result movement default to shared memory on one host rather than crossing the network. + +## 21. What KV-aware routing means + +KV-aware routing sends a request to the worker that can reuse the largest useful portion of the request's existing attention KV cache, while also accounting for current worker load. + +### Concrete example + +Suppose every request begins with a 4,000-token system prompt: + + Worker A: already caches the 4,000-token system prompt + Worker B: does not cache it + New request: same system prompt + 100 new user tokens + +Round-robin routing might choose Worker B, forcing it to prefill all 4,100 tokens. A KV-aware router normally chooses Worker A, which can reuse the cached prefix and prefill only the 100 new tokens. That reduces GPU computation and usually improves time to first token. + +### How Dynamo knows where the cache is + +1. The inference backend divides token sequences into cache blocks and derives stable hashes for their prefixes. +2. Workers publish `BlockStored` and `BlockRemoved` events when cache blocks are allocated or evicted. +3. Dynamo's event plane carries those events through ZMQ or NATS. +4. The router maintains an index mapping prefix-block hashes to workers. +5. For an incoming tokenized request, the router hashes its prefix, calculates the overlap with each worker, combines this with active prefill/decode load, and selects the lowest estimated-cost worker. + +The simplified cost is: + + score = prefill_load_scale × adjusted_uncached_prefill_blocks + + active_decode_blocks + +Cached prefix blocks reduce the prefill portion of the score. The router still accounts for active load, so an overloaded worker with a perfect cache hit can lose to a lightly loaded worker with a smaller hit. + +When precise cache events are unavailable, Dynamo can run an approximate mode that predicts cache placement from past routing decisions and expires those predictions by TTL. + +### KV-aware routing versus KV transfer + +* **KV-aware routing:** chooses a worker where useful cache blocks already exist. +* **KV transfer:** moves KV state between workers, for example from a prefill worker to a decode worker through NIXL. + +In an aggregated deployment, KV-aware routing chooses the worker that handles the complete request. In a P/D-disaggregated deployment, it is especially useful when selecting a prefill worker; after prefill, the required cache may still be transferred to the selected decode worker. + +This feature is mainly for transformer token inference with reusable prompt prefixes. It does not currently apply to the GLM-Image AR-to-DiT connector flow or ordinary image/video diffusion workers; vLLM-Omni stage workers do not publish the LLM KV events required by this router. + +## 22. Shared-memory transfer between GPU 0 and GPU 1 + +Yes, the AR output still has to move from the GPU 0 worker to the GPU 1 worker. “Shared memory on the same host” means that the current default connector uses **CPU host memory as the staging area** instead of sending the payload across a network or copying directly between GPU memories. + +The path is: + + GPU 0 + | AR inference produces token IDs + | device-to-host representation + v + Stage 0 Python process + | serialize vLLM output object + v + POSIX shared memory (/dev/shm in host RAM) + | Stage 1 opens the segment and reads the bytes + | deserialize and build the DiT input + v + Stage 1 Python process + | host-to-device copy + v + GPU 1 + +The stage router carries only a small handle such as the shared-memory segment name and byte count. It does not carry the token payload itself. + +“Shared” means both worker processes run under the same operating-system host and can map the same `/dev/shm` segment. It does **not** mean: + +* CUDA shared memory inside an SM; +* one GPU directly reading the other GPU's VRAM; +* CUDA peer-to-peer, NVLink, NCCL, or NIXL transfer; +* no data copy. + +The current vLLM-Omni connector documentation describes its stage transfers as D2H2D: device-to-host-to-device. This adds serialization, a host-memory write/read, and a host-to-GPU copy. For GLM-Image text-to-image, the AR-to-DiT payload is only on the order of tens of KiB, so this is usually negligible compared with running a 9B AR model and a 7B diffusion decoder. + +For stages on different machines, `/dev/shm` cannot work because each host has different physical RAM and a different OS namespace. A remote connector such as Mooncake over TCP or RDMA is required. Direct device-to-device OmniConnector transports are described as future work in the current vLLM-Omni design. + +## 23. Commands for launching vLLM and SGLang in P/D-disaggregated mode + +For local testing, first start Dynamo's discovery/supporting services: + + docker compose -f dev/docker-compose.yml up -d + +The repository presets launch a frontend, one decode worker, and one prefill worker on two GPUs: + + # vLLM + cd examples/backends/vllm + bash launch/disagg.sh + + # SGLang + cd examples/backends/sglang + bash launch/disagg.sh + +Both scripts also accept `--unified` to use the newer common backend entry point: + + bash launch/disagg.sh --unified + +### Expanded vLLM commands + + MODEL=Qwen/Qwen3-0.6B + + python -m dynamo.frontend & + + CUDA_VISIBLE_DEVICES=0 \ + DYN_SYSTEM_PORT=8081 \ + python -m dynamo.vllm \ + --model "$MODEL" \ + --enforce-eager \ + --disaggregation-mode decode \ + --kv-transfer-config \ + '{"kv_connector":"NixlConnector","kv_role":"kv_both"}' & + + CUDA_VISIBLE_DEVICES=1 \ + DYN_SYSTEM_PORT=8082 \ + VLLM_NIXL_SIDE_CHANNEL_PORT=20097 \ + python -m dynamo.vllm \ + --model "$MODEL" \ + --enforce-eager \ + --disaggregation-mode prefill \ + --kv-transfer-config \ + '{"kv_connector":"NixlConnector","kv_role":"kv_both"}' \ + --kv-events-config \ + '{"publisher":"zmq","topic":"kv-events","endpoint":"tcp://*:20081","enable_kv_cache_events":true}' + +`--enforce-eager` is included by the example for quick startup and should normally be removed when benchmarking or tuning production performance. + +### Expanded SGLang commands + + MODEL=Qwen/Qwen3-0.6B + BOOTSTRAP_PORT=12345 + + python -m dynamo.frontend & + + CUDA_VISIBLE_DEVICES=0 \ + DYN_SYSTEM_PORT=8081 \ + python -m dynamo.sglang \ + --enable-multimodal \ + --model-path "$MODEL" \ + --served-model-name "$MODEL" \ + --page-size 16 \ + --tp 1 \ + --trust-remote-code \ + --disaggregation-mode prefill \ + --disaggregation-bootstrap-port "$BOOTSTRAP_PORT" \ + --disaggregation-transfer-backend nixl \ + --host 0.0.0.0 \ + --port 40000 \ + --enable-metrics \ + --disable-piecewise-cuda-graph & + + CUDA_VISIBLE_DEVICES=1 \ + DYN_SYSTEM_PORT=8082 \ + python -m dynamo.sglang \ + --enable-multimodal \ + --model-path "$MODEL" \ + --served-model-name "$MODEL" \ + --page-size 16 \ + --tp 1 \ + --trust-remote-code \ + --disaggregation-mode decode \ + --disaggregation-bootstrap-port "$BOOTSTRAP_PORT" \ + --disaggregation-transfer-backend nixl \ + --host 0.0.0.0 \ + --enable-metrics \ + --disable-piecewise-cuda-graph + +The important backend-specific difference is that vLLM selects `NixlConnector` through `--kv-transfer-config`, while SGLang selects NIXL with `--disaggregation-transfer-backend nixl` and coordinates each request through the shared bootstrap port. + +For two prefill and two decode replicas with KV-aware routing, use: + + bash launch/disagg_router.sh + +from the respective backend example directory. These presets require four GPUs and launch the frontend with `--router-mode kv`. + +## 24. What underlying vLLM API does Dynamo use for disaggregated mode? + +Dynamo does **not** call a vLLM API such as `launch_disaggregated()`, and it does not launch `vllm serve` as a second HTTP service. It embeds ordinary vLLM Python engines and coordinates them as separate prefill and decode workers. + +There are three distinct layers: + +1. **Dynamo control plane:** `--disaggregation-mode prefill|decode` assigns the worker role, registers/discovers workers, routes requests, and relays the prefill result to the selected decoder. +2. **vLLM engine API:** Dynamo constructs each engine with `AsyncEngineArgs.create_engine_config(...)` and `AsyncLLM.from_vllm_config(...)`, then submits requests through `AsyncLLM.generate(...)`. +3. **vLLM KV-transfer API:** `--kv-transfer-config` becomes vLLM's `KVTransferConfig`, selecting `NixlConnector`. Per-request transfer metadata is passed as `SamplingParams.extra_args["kv_transfer_params"]`. + +Conceptually, engine creation looks like this: + + vllm_config = engine_args.create_engine_config( + usage_context=UsageContext.OPENAI_API_SERVER + ) + + engine = AsyncLLM.from_vllm_config( + vllm_config=vllm_config, + usage_context=UsageContext.OPENAI_API_SERVER, + ) + +The launch argument: + + --kv-transfer-config \ + '{"kv_connector":"NixlConnector","kv_role":"kv_both"}' + +populates vLLM's `KVTransferConfig`. vLLM then instantiates and invokes the NIXL connector internally; Dynamo does not directly call the connector for every GPU transfer. + +For the prefill request, Dynamo uses the ordinary vLLM generation API but marks the request for a remote decode: + + sampling_params.extra_args["kv_transfer_params"] = { + "do_remote_prefill": False, + "remote_engine_id": None, + "remote_block_ids": None, + "remote_host": None, + "remote_port": None, + "do_remote_decode": True, + } + sampling_params.max_tokens = 1 + sampling_params.min_tokens = 1 + + async for output in prefill_engine.generate( + prompt, sampling_params, request_id + ): + handoff = output.kv_transfer_params + +The one-token limit makes this engine perform prompt prefill and prepare the KV handoff rather than generate the full answer. vLLM returns `output.kv_transfer_params`, which describes where the produced KV blocks can be fetched. + +Dynamo relays that handoff to the decoder and invokes the same vLLM API: + + sampling_params.extra_args["kv_transfer_params"] = handoff + + async for output in decode_engine.generate( + prompt, sampling_params, request_id + ): + ... + +The decode-side vLLM scheduler and `NixlConnector` interpret the transfer metadata, retrieve the remote KV blocks, and then perform autoregressive decoding. In the current pull-based protocol, the decode worker initiates the NIXL read from the prefill worker. + +Therefore, the short answer is: + +> Dynamo uses vLLM's `AsyncLLM.from_vllm_config()` and `AsyncLLM.generate()` APIs, configured with `KVTransferConfig(kv_connector="NixlConnector", ...)`. The crucial per-request disaggregation hook is `SamplingParams.extra_args["kv_transfer_params"]`; Dynamo supplies the orchestration around it. + +This is a Python engine integration, not an OpenAI HTTP API. The KV-transfer interface is also an evolving/experimental vLLM integration surface, so it is more version-sensitive than the normal generation API. This checkout pins `vllm==0.24.0`. + +Relevant implementation: + +* `components/src/dynamo/vllm/main.py` — creates the embedded `AsyncLLM`. +* `components/src/dynamo/vllm/llm_engine.py` — constructs the engine and calls `generate()` for prefill/decode. +* `components/src/dynamo/vllm/handlers.py` — legacy prefill/decode request handling. +* `components/src/dynamo/vllm/kv_connector_protocols.py` — constructs and relays the NIXL transfer parameters. +* `components/src/dynamo/vllm/args.py` — validates Dynamo's disaggregation arguments. +* `pyproject.toml` — pins the vLLM version. + +## 25. What underlying SGLang API does Dynamo use for disaggregated mode? + +SGLang differs from vLLM because PD disaggregation is an explicit upstream +SGLang engine feature. The following are native SGLang `ServerArgs`, not +Dynamo inventions: + + --disaggregation-mode prefill|decode + --disaggregation-transfer-backend nixl + --disaggregation-bootstrap-port 12345 + +Dynamo parses the upstream arguments with `ServerArgs.add_cli_args(...)` and +`ServerArgs.from_cli_args(...)`, then embeds an SGLang engine: + + server_args = ServerArgs.from_cli_args(parsed_args) + engine = sgl.Engine(server_args=server_args) + +It does not launch a separate `sglang.launch_server` HTTP server. The installed +SGLang engine starts its own tokenizer-manager, scheduler, and detokenizer +processes internally. This checkout installs `sglang[diffusion]==0.5.14`. + +The per-request disaggregation interface is SGLang's ordinary +`Engine.async_generate()` method with three native bootstrap arguments: + + stream = await engine.async_generate( + input_ids=prompt_token_ids, + sampling_params=sampling_params, + stream=True, + rid=request_id, + bootstrap_host=prefill_host, + bootstrap_port=prefill_bootstrap_port, + bootstrap_room=request_room, + ) + +`bootstrap_host`, `bootstrap_port`, and `bootstrap_room` form a rendezvous key: + +* `bootstrap_host` identifies the prefill worker. +* `bootstrap_port` identifies that worker's SGLang disaggregation bootstrap service. +* `bootstrap_room` is a request-specific channel identifier so the matching prefill and decode operations find one another. + +For a prefill worker, Dynamo first calls `async_generate()` with this triple. +Because the engine was constructed with `disaggregation_mode="prefill"`, +SGLang registers the room and performs the prefill side of the protocol. +Dynamo then sends only the small bootstrap triple to the decode side: + + { + "bootstrap_host": "...", + "bootstrap_port": 12345, + "bootstrap_room": 428729103, + } + +The decoder calls its own SGLang engine with the same triple: + + decode_stream = await decode_engine.async_generate( + input_ids=prompt_token_ids, + sampling_params=sampling_params, + stream=True, + bootstrap_host=bootstrap_info["bootstrap_host"], + bootstrap_port=bootstrap_info["bootstrap_port"], + bootstrap_room=bootstrap_info["bootstrap_room"], + ) + +Because that engine was constructed with `disaggregation_mode="decode"`, +SGLang interprets the request as the decode peer. Its internal disaggregation +implementation uses the selected transfer backend, such as NIXL, to transfer +the actual KV-cache blocks. The large KV tensors do not travel through Dynamo +or inside the bootstrap dictionary. + +Compared with vLLM: + +| Concern | vLLM integration | SGLang integration | +| --- | --- | --- | +| Engine construction | `AsyncLLM.from_vllm_config()` | `sgl.Engine(server_args=...)` | +| Request API | `AsyncLLM.generate()` | `Engine.async_generate()` | +| Engine role | Dynamo role plus vLLM `KVTransferConfig` | Native SGLang `disaggregation_mode` | +| Transfer backend | `KVTransferConfig(kv_connector="NixlConnector")` | Native `disaggregation_transfer_backend="nixl"` | +| Per-request handoff | `SamplingParams.extra_args["kv_transfer_params"]` | `bootstrap_host`, `bootstrap_port`, `bootstrap_room` | +| Large KV transfer | vLLM connector internals | SGLang disaggregation internals | + +Unlike Dynamo's vLLM adapter, the SGLang adapter does not need to force +`max_tokens=1` to emulate a prefill request. The SGLang engine already knows it +is a prefill-only engine from `ServerArgs.disaggregation_mode`. + +No patched copy of SGLang is required for this mechanism. Dynamo installs and +embeds the upstream SGLang package, while its adapter supplies Dynamo routing, +worker registration, request translation, bootstrap coordination, monitoring, +and cancellation behavior. + +Relevant implementation: + +* `components/src/dynamo/sglang/args.py` — consumes native SGLang `ServerArgs`. +* `components/src/dynamo/sglang/llm_engine.py` — constructs `sgl.Engine`, calls `async_generate()`, and coordinates the bootstrap triple. +* `components/src/dynamo/sglang/_disagg.py` — resolves the advertised bootstrap address and warms up the prefill path. +* `components/src/dynamo/sglang/request_handlers/llm/prefill_handler.py` and `decode_handler.py` — equivalent legacy request paths. +* `pyproject.toml` — pins the SGLang package version. + +## 26. What do SGLang's bootstrap host, port, and room mean? + +These three values are a rendezvous address that lets the decode-side request +find the matching prefill-side request: + + bootstrap_host="10.0.0.12" + bootstrap_port=12345 + bootstrap_room=428729103 + +* **`bootstrap_host`** is the network-reachable address of the SGLang prefill worker that produced the KV cache. It must be reachable from the decode worker; `127.0.0.1` is only suitable when both engines are on the same host. +* **`bootstrap_port`** is the prefill worker's SGLang disaggregation bootstrap service port. It is used to establish and coordinate the transfer. It should not be confused with the public OpenAI/HTTP serving port. +* **`bootstrap_room`** is a request-specific channel or meeting ID. Many requests can use the same host and port concurrently, so the room tells SGLang which prefill request must be paired with which decode request. Dynamo normally generates a random 63-bit room when the router has not already supplied one. + +Both SGLang calls receive the same triple: + + # On the prefill engine + await prefill_engine.async_generate( + ..., + bootstrap_host="10.0.0.12", + bootstrap_port=12345, + bootstrap_room=428729103, + ) + + # On the decode engine + await decode_engine.async_generate( + ..., + bootstrap_host="10.0.0.12", + bootstrap_port=12345, + bootstrap_room=428729103, + ) + +Conceptually, the triple is comparable to: + + server address = 10.0.0.12:12345 + request/session key = 428729103 + +It is small control-plane metadata, not the KV cache. SGLang uses the rendezvous +to exchange the information needed by the selected transfer backend. NIXL then +moves the large KV tensors through the available data path, such as GPU memory +over NVLink/PCIe on one host or RDMA-capable networking across hosts. The +bootstrap port coordinates that operation; it should not be interpreted as the +Python application serializing the KV cache into the `async_generate()` call. + +If the host or port is unreachable, the peers cannot establish the transfer. If +the room differs, the decoder does not match the intended prefill operation and +the request will normally wait or fail rather than receiving that KV cache. + +## 27. Complete command to launch SGLang in prefill mode + +When SGLang is used under Dynamo, launch the prefill worker through Dynamo's +SGLang entry point: + + CUDA_VISIBLE_DEVICES=0 \ + OTEL_SERVICE_NAME=dynamo-worker-prefill \ + DYN_SYSTEM_PORT=8081 \ + python3 -m dynamo.sglang \ + --model-path Qwen/Qwen3-0.6B \ + --served-model-name Qwen/Qwen3-0.6B \ + --page-size 16 \ + --tp 1 \ + --trust-remote-code \ + --disaggregation-mode prefill \ + --disaggregation-transfer-backend nixl \ + --disaggregation-bootstrap-port 12345 \ + --host 0.0.0.0 \ + --port 40000 \ + --enable-metrics \ + --disable-piecewise-cuda-graph + +This launches an embedded `sgl.Engine` registered as a Dynamo prefill worker. +It is not the client-facing OpenAI endpoint. A usable Dynamo deployment also +needs the Dynamo runtime services, frontend/router, and at least one matching +decode worker. + +The repository's complete two-GPU example, including the frontend and decoder, +is: + + bash examples/backends/sglang/launch/disagg.sh + +To use the unified worker implementation: + + bash examples/backends/sglang/launch/disagg.sh --unified + +If the intention is to launch SGLang's own HTTP server without Dynamo, use the +upstream entry point instead: + + CUDA_VISIBLE_DEVICES=0 \ + python3 -m sglang.launch_server \ + --model-path Qwen/Qwen3-0.6B \ + --tp-size 1 \ + --disaggregation-mode prefill \ + --disaggregation-transfer-backend nixl \ + --disaggregation-bootstrap-port 12345 \ + --host 0.0.0.0 \ + --port 30000 + +The native SGLang command uses `--tp-size`; Dynamo's wrapper accepts the +upstream alias `--tp` used by the repository examples. A standalone prefill +server is only one half of a PD deployment: it still needs a decode server and +an orchestrator/load balancer such as SGLang's `mini_lb`. diff --git a/flashdreams/flashdreams/core/attention/cp.py b/flashdreams/flashdreams/core/attention/cp.py index 735a5daf..fb80728b 100644 --- a/flashdreams/flashdreams/core/attention/cp.py +++ b/flashdreams/flashdreams/core/attention/cp.py @@ -113,7 +113,10 @@ def _impl_ring(self, query: Tensor, key: Tensor, value: Tensor) -> Tensor: if self.device_mesh is None: return attn_op(query, key, value, return_lse=False)[0] - rank = self.device_mesh.get_rank() + # ``get_rank()`` is the global process rank. Ring rotation indexes the + # tuple returned by the subgroup all-gather, so it must use the rank + # local to that subgroup (for example 0..5 for global ranks 1..6). + rank = self.device_mesh.get_local_rank() world_size = self.device_mesh.size() group = self.device_mesh.get_group() if world_size == 1: diff --git a/flashdreams/flashdreams/infra/pipeline/__init__.py b/flashdreams/flashdreams/infra/pipeline/__init__.py index dde076d7..65770f27 100644 --- a/flashdreams/flashdreams/infra/pipeline/__init__.py +++ b/flashdreams/flashdreams/infra/pipeline/__init__.py @@ -20,9 +20,19 @@ StreamInferencePipelineCache, StreamInferencePipelineConfig, ) +from flashdreams.infra.pipeline.stages import ( + DecoderStage, + DiffusionStage, + DiffusionStageCache, + StreamingEncoderStage, +) __all__ = [ + "DecoderStage", + "DiffusionStage", + "DiffusionStageCache", "StreamInferencePipeline", "StreamInferencePipelineCache", "StreamInferencePipelineConfig", + "StreamingEncoderStage", ] diff --git a/flashdreams/flashdreams/infra/pipeline/stages.py b/flashdreams/flashdreams/infra/pipeline/stages.py new file mode 100644 index 00000000..d753e42d --- /dev/null +++ b/flashdreams/flashdreams/infra/pipeline/stages.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Independently deployable encoder, diffusion, and decoder pipeline stages.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Generic + +import torch +from torch import Tensor, nn + +from flashdreams.infra.decoder import ( + DecoderConfig, + StreamingDecoder, + StreamingDecoderCacheT, +) +from flashdreams.infra.diffusion.model import DiffusionModel, DiffusionModelConfig +from flashdreams.infra.diffusion.transformer import TransformerCacheT +from flashdreams.infra.encoder import ( + EncoderConfig, + StreamingEncoder, + StreamingEncoderCacheT, +) + + +class StreamingEncoderStage(nn.Module, Generic[StreamingEncoderCacheT]): + """Own only a pipeline's per-AR-step encoder.""" + + encoder: StreamingEncoder[StreamingEncoderCacheT] + + def __init__(self, config: EncoderConfig) -> None: + super().__init__() + self.encoder = config.setup() + + def initialize_cache(self, **context: Any) -> StreamingEncoderCacheT: + """Build the encoder's per-rollout cache.""" + return self.encoder.initialize_autoregressive_cache(**context) + + @torch.no_grad() + def encode( + self, + input: Any, + autoregressive_index: int, + cache: StreamingEncoderCacheT, + ) -> Any: + """Encode one raw control chunk.""" + return self.encoder( + input=input, + autoregressive_index=autoregressive_index, + cache=cache, + ) + + +@dataclass(kw_only=True) +class DiffusionStageCache(Generic[TransformerCacheT]): + """Per-rollout state owned by a diffusion stage worker.""" + + transformer_cache: TransformerCacheT + """Long-lived transformer cache pinned to this worker.""" + + final_state: DiffusionModel.FinalState[TransformerCacheT] | None = None + """Most recent denoising result consumed by :meth:`DiffusionStage.finalize`.""" + + autoregressive_index: int | None = None + """Most recent AR index, or ``None`` before generation starts.""" + + +class DiffusionStage(nn.Module, Generic[TransformerCacheT]): + """Own only the scheduler and denoising transformer (DiT).""" + + diffusion_model: DiffusionModel[TransformerCacheT] + + def __init__(self, config: DiffusionModelConfig) -> None: + super().__init__() + self.diffusion_model = config.setup() + + @property + def device(self) -> torch.device: + """Return the DiT device.""" + return self.diffusion_model.device + + def initialize_cache( + self, **context: Any + ) -> DiffusionStageCache[TransformerCacheT]: + """Build the transformer cache from encoder-stage context.""" + transformer_cache = ( + self.diffusion_model.transformer.initialize_autoregressive_cache(**context) + ) + return DiffusionStageCache(transformer_cache=transformer_cache) + + @torch.no_grad() + def generate( + self, + autoregressive_index: int, + cache: DiffusionStageCache[TransformerCacheT], + input: Any = None, + ) -> Tensor: + """Denoise one AR chunk and retain the state needed for finalization.""" + previous = cache.autoregressive_index + expected = previous + 1 if previous is not None else 0 + assert autoregressive_index == expected, ( + f"AR step out of order: previous step was {previous}, expected " + f"{expected}, got {autoregressive_index}." + ) + clean_latent, final_state = self.diffusion_model.generate( + autoregressive_index=autoregressive_index, + cache=cache.transformer_cache, + input=input, + ) + cache.autoregressive_index = autoregressive_index + cache.final_state = final_state + return clean_latent + + @torch.no_grad() + def finalize( + self, + autoregressive_index: int, + cache: DiffusionStageCache[TransformerCacheT], + ) -> None: + """Advance the resident DiT cache after one generated chunk.""" + assert cache.autoregressive_index == autoregressive_index, ( + f"autoregressive_index mismatch: generate() ran with " + f"{cache.autoregressive_index}, finalize() got {autoregressive_index}." + ) + assert cache.final_state is not None, ( + "finalize() called before generate() produced a final state." + ) + self.diffusion_model.finalize(cache.final_state) + cache.final_state = None + + +class DecoderStage(nn.Module, Generic[StreamingDecoderCacheT]): + """Own only a pipeline's streaming decoder.""" + + decoder: StreamingDecoder[StreamingDecoderCacheT] + + def __init__(self, config: DecoderConfig) -> None: + super().__init__() + self.decoder = config.setup() + + def initialize_cache(self, **context: Any) -> StreamingDecoderCacheT: + """Build the decoder's per-rollout cache.""" + return self.decoder.initialize_autoregressive_cache(**context) + + @torch.no_grad() + def decode( + self, + input: Tensor, + autoregressive_index: int, + cache: StreamingDecoderCacheT, + ) -> Tensor: + """Decode one clean latent chunk.""" + return self.decoder( + input=input, + autoregressive_index=autoregressive_index, + cache=cache, + ) diff --git a/flashdreams/flashdreams/infra/transfer.py b/flashdreams/flashdreams/infra/transfer.py new file mode 100644 index 00000000..1e647413 --- /dev/null +++ b/flashdreams/flashdreams/infra/transfer.py @@ -0,0 +1,730 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Registered GPU tensor transfer through Mooncake or NIXL.""" + +from __future__ import annotations + +import importlib +import socket +import threading +import time +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from typing import Any + +import torch +from torch import Tensor + +TensorBundle = dict[str, Tensor] +"""Flat, name-keyed tensor payload transferred between pipeline stages.""" + + +@dataclass(frozen=True, kw_only=True) +class TensorDescriptor: + """Shape and storage metadata for one transferred tensor.""" + + name: str + """Stable field name inside the tensor bundle.""" + + shape: tuple[int, ...] + """Tensor shape.""" + + dtype: torch.dtype + """Tensor element type.""" + + nbytes: int + """Contiguous payload size in bytes.""" + + +@dataclass(frozen=True, kw_only=True) +class TensorTransferTicket: + """Receiver-owned registered memory advertised to a sender.""" + + session_id: str + """Backend peer or agent identifier.""" + + descriptors: tuple[TensorDescriptor, ...] + """Expected bundle layout in transfer order.""" + + addresses: tuple[int, ...] = () + """Registered receiver addresses used by address-based transports.""" + + metadata: bytes | None = None + """Serialized receiver-agent metadata used by NIXL.""" + + remote_descriptors: bytes | None = None + """Serialized receiver transfer descriptors used by NIXL.""" + + +@dataclass(kw_only=True) +class TensorTransferHandle: + """One submitted tensor transfer that may still be in flight.""" + + backend: str + """Selected transfer backend.""" + + payload_bytes: int + """Total tensor bytes in the operation.""" + + registration_ms: float + """Sender memory-registration time for previously unseen allocations.""" + + submit_ms: float + """Host time spent submitting the transfer.""" + + submitted_at: float + """Monotonic timestamp immediately before backend submission.""" + + opaque_handle: Any = None + """Backend-specific asynchronous operation handle.""" + + source_bundle: Mapping[str, Tensor] | None = None + """Strong reference preserving source storage until completion.""" + + completed_stats: TransferStats | None = None + """Immediate result when the backend used a synchronous fallback.""" + + +@dataclass(frozen=True, kw_only=True) +class TransferStats: + """Observed metrics for one tensor-bundle transfer.""" + + backend: str + """Selected transfer backend and protocol.""" + + payload_bytes: int + """Total tensor bytes transferred.""" + + registration_ms: float + """Sender memory-registration time for previously unseen allocations.""" + + transfer_ms: float + """Submission-to-completion-observation wall time. + + For asynchronous operations this includes any useful work performed before + the caller invokes ``wait``; it is an in-flight residency window, not an + isolated copy-time measurement. + """ + + bandwidth_gbps: float + """Effective decimal GB/s computed from payload bytes and transfer time.""" + + submit_ms: float = 0.0 + """Host time spent submitting the operation.""" + + wait_ms: float = 0.0 + """Host time spent waiting after submission.""" + + asynchronous: bool = False + """Whether the backend used its asynchronous operation API.""" + + +@dataclass(frozen=True, kw_only=True) +class PooledTensorBuffer: + """One reusable registered receiver bundle and its transfer ticket.""" + + bundle: TensorBundle + """Registered receiver tensors.""" + + ticket: TensorTransferTicket + """Stable ticket advertising the registered tensors.""" + + bucket: tuple[tuple[TensorDescriptor, ...], str] + """Pool bucket key used when returning this buffer.""" + + +def describe_tensor_bundle( + bundle: Mapping[str, Tensor], +) -> tuple[TensorDescriptor, ...]: + """Describe a contiguous tensor bundle in stable insertion order.""" + descriptors: list[TensorDescriptor] = [] + for name, tensor in bundle.items(): + if not tensor.is_contiguous(): + raise ValueError(f"Tensor bundle field {name!r} must be contiguous.") + descriptors.append( + TensorDescriptor( + name=name, + shape=tuple(tensor.shape), + dtype=tensor.dtype, + nbytes=tensor.numel() * tensor.element_size(), + ) + ) + return tuple(descriptors) + + +class MooncakeTensorTransport: + """Move tensor bundles directly between registered GPU buffers.""" + + def __init__( + self, + *, + hostname: str | None = None, + device_name: str | None = None, + protocol: str = "rdma", + engine_factory: Callable[[], Any] | None = None, + ) -> None: + """Initialize a Mooncake endpoint. + + Args: + hostname: Routable address advertised to peers. Defaults to the + local hostname. + device_name: Optional RDMA device selection passed to Mooncake. + protocol: Mooncake transport protocol. Production disaggregation + should use ``"rdma"``. + engine_factory: Test hook returning a Transfer Engine-compatible + object. ``None`` imports ``mooncake.engine.TransferEngine``. + + Raises: + ImportError: Mooncake's Python package is unavailable. + RuntimeError: The transfer endpoint fails to initialize. + """ + if engine_factory is None: + try: + from mooncake.engine import TransferEngine + except ImportError as error: + raise ImportError( + "Mooncake transfer support requires the CUDA 13 package: " + "pip install mooncake-transfer-engine-cuda13." + ) from error + engine_factory = TransferEngine + + self.hostname = hostname or socket.gethostname() + self.protocol = protocol + self.engine = engine_factory() + result = self.engine.initialize( + self.hostname, + "P2PHANDSHAKE", + protocol, + device_name or "", + ) + if result != 0: + raise RuntimeError( + f"Mooncake Transfer Engine initialization failed with status {result}." + ) + self.session_id = f"{self.hostname}:{self.engine.get_rpc_port()}" + self._registered_addresses: set[int] = set() + self._registered_tensors: dict[int, Tensor] = {} + + @property + def backend(self) -> str: + """Return the backend label recorded in benchmark output.""" + return f"mooncake-{self.protocol}" + + def allocate( + self, + descriptors: tuple[TensorDescriptor, ...], + *, + device: torch.device, + ) -> TensorBundle: + """Allocate and register receiver VRAM for a described bundle.""" + bundle: TensorBundle = {} + for descriptor in descriptors: + tensor = torch.empty( + descriptor.shape, + dtype=descriptor.dtype, + device=device, + ).contiguous() + actual_nbytes = tensor.numel() * tensor.element_size() + if actual_nbytes != descriptor.nbytes: + raise ValueError( + f"Descriptor size mismatch for {descriptor.name!r}: expected " + f"{descriptor.nbytes}, allocated {actual_nbytes}." + ) + result = self.engine.register_memory(tensor.data_ptr(), actual_nbytes) + if result != 0: + raise RuntimeError( + f"Mooncake failed to register {descriptor.name!r} " + f"({actual_nbytes} bytes), status={result}." + ) + self._registered_addresses.add(tensor.data_ptr()) + self._registered_tensors[tensor.data_ptr()] = tensor + bundle[descriptor.name] = tensor + return bundle + + def make_ticket(self, bundle: Mapping[str, Tensor]) -> TensorTransferTicket: + """Advertise registered receiver buffers to a sender.""" + descriptors = describe_tensor_bundle(bundle) + missing = [ + descriptor.name + for descriptor in descriptors + if bundle[descriptor.name].data_ptr() not in self._registered_addresses + ] + if missing: + raise ValueError(f"Receiver buffers are not registered: {missing}.") + return TensorTransferTicket( + session_id=self.session_id, + descriptors=descriptors, + addresses=tuple(bundle[item.name].data_ptr() for item in descriptors), + ) + + def register(self, bundle: Mapping[str, Tensor]) -> float: + """Register unseen bundle allocations and return elapsed milliseconds.""" + started = time.perf_counter() + for descriptor in describe_tensor_bundle(bundle): + tensor = bundle[descriptor.name] + address = tensor.data_ptr() + if address in self._registered_addresses: + continue + result = self.engine.register_memory(address, descriptor.nbytes) + if result != 0: + raise RuntimeError( + f"Mooncake failed to register {descriptor.name!r} " + f"({descriptor.nbytes} bytes), status={result}." + ) + self._registered_addresses.add(address) + self._registered_tensors[address] = tensor + return (time.perf_counter() - started) * 1000.0 + + def unregister(self, bundle: Mapping[str, Tensor]) -> None: + """Unregister bundle allocations before their tensors are released.""" + for tensor in bundle.values(): + address = tensor.data_ptr() + if address not in self._registered_addresses: + continue + result = self.engine.unregister_memory(address) + if result != 0: + raise RuntimeError( + f"Mooncake failed to unregister address {address}, status={result}." + ) + self._registered_addresses.remove(address) + self._registered_tensors.pop(address, None) + + @staticmethod + def _wait_until_source_ready(bundle: Mapping[str, Tensor]) -> None: + """Wait only for the CUDA producer stream instead of the whole device.""" + cuda_tensors = [tensor for tensor in bundle.values() if tensor.is_cuda] + if not cuda_tensors: + return + devices = {tensor.device for tensor in cuda_tensors} + if len(devices) != 1: + raise ValueError("A tensor bundle must reside on one CUDA device.") + event = torch.cuda.Event() + event.record(torch.cuda.current_stream(cuda_tensors[0].device)) + event.synchronize() + + def send_async( + self, + bundle: Mapping[str, Tensor], + ticket: TensorTransferTicket, + ) -> TensorTransferHandle: + """Submit a tensor-bundle write and return without waiting for RDMA.""" + descriptors = describe_tensor_bundle(bundle) + if descriptors != ticket.descriptors: + raise ValueError( + "Sender bundle layout does not match the receiver transfer ticket." + ) + if not ticket.addresses: + raise ValueError("Mooncake transfer tickets require receiver addresses.") + + registration_ms = self.register(bundle) + self._wait_until_source_ready(bundle) + sources = [bundle[item.name].data_ptr() for item in descriptors] + lengths = [item.nbytes for item in descriptors] + payload_bytes = sum(lengths) + submitted_at = time.perf_counter() + + async_write = getattr(self.engine, "batch_transfer_async_write", None) + if callable(async_write): + batch_id = async_write( + ticket.session_id, + sources, + list(ticket.addresses), + lengths, + ) + submit_ms = (time.perf_counter() - submitted_at) * 1000.0 + if not isinstance(batch_id, int) or batch_id <= 0: + raise RuntimeError( + f"Mooncake async tensor transfer submission failed: {batch_id}." + ) + return TensorTransferHandle( + backend=self.backend, + payload_bytes=payload_bytes, + registration_ms=registration_ms, + submit_ms=submit_ms, + submitted_at=submitted_at, + opaque_handle=batch_id, + source_bundle=bundle, + ) + + result = self.engine.batch_transfer_sync_write( + ticket.session_id, + sources, + list(ticket.addresses), + lengths, + ) + elapsed_s = time.perf_counter() - submitted_at + if result != 0: + raise RuntimeError(f"Mooncake tensor transfer failed with status {result}.") + stats = TransferStats( + backend=self.backend, + payload_bytes=payload_bytes, + registration_ms=registration_ms, + transfer_ms=elapsed_s * 1000.0, + bandwidth_gbps=( + payload_bytes / elapsed_s / 1e9 if elapsed_s > 0.0 else float("inf") + ), + submit_ms=elapsed_s * 1000.0, + asynchronous=False, + ) + return TensorTransferHandle( + backend=self.backend, + payload_bytes=payload_bytes, + registration_ms=registration_ms, + submit_ms=stats.submit_ms, + submitted_at=submitted_at, + source_bundle=bundle, + completed_stats=stats, + ) + + def wait(self, handle: TensorTransferHandle) -> TransferStats: + """Wait for a submitted Mooncake transfer and return its metrics.""" + if handle.completed_stats is not None: + return handle.completed_stats + status = self.engine.get_batch_transfer_status([handle.opaque_handle]) + wait_finished = time.perf_counter() + if status != 0: + raise RuntimeError( + f"Mooncake async tensor transfer failed with status {status}." + ) + elapsed_s = wait_finished - handle.submitted_at + wait_ms = max(0.0, elapsed_s * 1000.0 - handle.submit_ms) + return TransferStats( + backend=self.backend, + payload_bytes=handle.payload_bytes, + registration_ms=handle.registration_ms, + transfer_ms=elapsed_s * 1000.0, + bandwidth_gbps=( + handle.payload_bytes / elapsed_s / 1e9 + if elapsed_s > 0.0 + else float("inf") + ), + submit_ms=handle.submit_ms, + wait_ms=wait_ms, + asynchronous=True, + ) + + def send( + self, + bundle: Mapping[str, Tensor], + ticket: TensorTransferTicket, + ) -> TransferStats: + """Write a tensor bundle and wait for completion.""" + return self.wait(self.send_async(bundle, ticket)) + + def close(self) -> None: + """Unregister every receiver allocation owned by this endpoint.""" + for address in tuple(self._registered_addresses): + result = self.engine.unregister_memory(address) + if result != 0: + raise RuntimeError( + f"Mooncake failed to unregister address {address}, status={result}." + ) + self._registered_addresses.remove(address) + self._registered_tensors.pop(address, None) + + +class NixlTensorTransport: + """Move registered tensor bundles through the NIXL Python API.""" + + def __init__( + self, + *, + agent_name: str | None = None, + agent_factory: Callable[[str, Any], Any] | None = None, + config_factory: Callable[..., Any] | None = None, + ) -> None: + """Initialize a NIXL transfer agent. + + Args: + agent_name: Unique agent name. Defaults to hostname plus process ID. + agent_factory: Test hook compatible with ``nixl_agent``. + config_factory: Test hook compatible with ``nixl_agent_config``. + + Raises: + ImportError: The NIXL Python package is unavailable. + """ + if agent_factory is None or config_factory is None: + try: + nixl_module = importlib.import_module("nixl") + except ImportError as error: + raise ImportError( + "NIXL transfer support requires the optional package: " + "pip install nixl." + ) from error + if not hasattr(nixl_module, "nixl_agent"): + nixl_module = importlib.import_module("nixl._api") + agent_factory = agent_factory or getattr(nixl_module, "nixl_agent") + config_factory = config_factory or getattr( + nixl_module, + "nixl_agent_config", + ) + + import os + + self.session_id = agent_name or f"{socket.gethostname()}-{os.getpid()}" + self.agent = agent_factory( + self.session_id, + config_factory(True, True, 0), + ) + self._registrations: dict[int, Any] = {} + self._registered_tensors: dict[int, Tensor] = {} + self._remote_agents: set[str] = set() + + @property + def backend(self) -> str: + """Return the backend label recorded in benchmark output.""" + return "nixl" + + def allocate( + self, + descriptors: tuple[TensorDescriptor, ...], + *, + device: torch.device, + ) -> TensorBundle: + """Allocate and register receiver memory for a described bundle.""" + bundle = { + descriptor.name: torch.empty( + descriptor.shape, + dtype=descriptor.dtype, + device=device, + ).contiguous() + for descriptor in descriptors + } + self.register(bundle) + return bundle + + def register(self, bundle: Mapping[str, Tensor]) -> float: + """Register unseen bundle allocations and return elapsed milliseconds.""" + started = time.perf_counter() + for tensor in bundle.values(): + address = tensor.data_ptr() + if address in self._registrations: + continue + registration = self.agent.register_memory(tensor) + if not registration: + raise RuntimeError( + f"NIXL failed to register tensor at address {address}." + ) + self._registrations[address] = registration + self._registered_tensors[address] = tensor + return (time.perf_counter() - started) * 1000.0 + + def make_ticket(self, bundle: Mapping[str, Tensor]) -> TensorTransferTicket: + """Serialize receiver metadata and tensor descriptors for a sender.""" + descriptors = describe_tensor_bundle(bundle) + missing = [ + item.name + for item in descriptors + if bundle[item.name].data_ptr() not in self._registrations + ] + if missing: + raise ValueError(f"Receiver buffers are not registered: {missing}.") + remote = self.agent.get_xfer_descs([bundle[item.name] for item in descriptors]) + if not remote: + raise RuntimeError("NIXL failed to create receiver transfer descriptors.") + return TensorTransferTicket( + session_id=self.session_id, + descriptors=descriptors, + metadata=self.agent.get_agent_metadata(), + remote_descriptors=self.agent.get_serialized_descs(remote), + ) + + def send_async( + self, + bundle: Mapping[str, Tensor], + ticket: TensorTransferTicket, + ) -> TensorTransferHandle: + """Submit a NIXL write and return its asynchronous handle.""" + descriptors = describe_tensor_bundle(bundle) + if descriptors != ticket.descriptors: + raise ValueError( + "Sender bundle layout does not match the receiver transfer ticket." + ) + if ticket.metadata is None or ticket.remote_descriptors is None: + raise ValueError("NIXL tickets require agent metadata and descriptors.") + registration_ms = self.register(bundle) + MooncakeTensorTransport._wait_until_source_ready(bundle) + if ticket.session_id not in self._remote_agents: + loaded_name = self.agent.add_remote_agent(ticket.metadata) + if loaded_name != ticket.session_id: + raise RuntimeError( + f"NIXL ticket names {ticket.session_id!r}, metadata loaded " + f"{loaded_name!r}." + ) + self._remote_agents.add(loaded_name) + + local = self.agent.get_xfer_descs([bundle[item.name] for item in descriptors]) + remote = self.agent.deserialize_descs(ticket.remote_descriptors) + xfer = self.agent.initialize_xfer( + "WRITE", + local, + remote, + ticket.session_id, + ) + submitted_at = time.perf_counter() + state = self.agent.transfer(xfer) + submit_ms = (time.perf_counter() - submitted_at) * 1000.0 + if state == "ERR": + self.agent.release_xfer_handle(xfer) + raise RuntimeError("NIXL tensor transfer submission failed.") + return TensorTransferHandle( + backend=self.backend, + payload_bytes=sum(item.nbytes for item in descriptors), + registration_ms=registration_ms, + submit_ms=submit_ms, + submitted_at=submitted_at, + opaque_handle=xfer, + source_bundle=bundle, + ) + + def wait( + self, + handle: TensorTransferHandle, + *, + timeout_s: float = 60.0, + ) -> TransferStats: + """Wait for a submitted NIXL transfer and return its metrics.""" + deadline = time.monotonic() + timeout_s + while True: + state = self.agent.check_xfer_state(handle.opaque_handle) + if state == "DONE": + break + if state == "ERR": + self.agent.release_xfer_handle(handle.opaque_handle) + raise RuntimeError("NIXL tensor transfer failed.") + if time.monotonic() >= deadline: + self.agent.release_xfer_handle(handle.opaque_handle) + raise TimeoutError( + f"NIXL tensor transfer exceeded {timeout_s:.1f} seconds." + ) + time.sleep(0) + finished = time.perf_counter() + self.agent.release_xfer_handle(handle.opaque_handle) + elapsed_s = finished - handle.submitted_at + return TransferStats( + backend=self.backend, + payload_bytes=handle.payload_bytes, + registration_ms=handle.registration_ms, + transfer_ms=elapsed_s * 1000.0, + bandwidth_gbps=( + handle.payload_bytes / elapsed_s / 1e9 + if elapsed_s > 0.0 + else float("inf") + ), + submit_ms=handle.submit_ms, + wait_ms=max(0.0, elapsed_s * 1000.0 - handle.submit_ms), + asynchronous=True, + ) + + def send( + self, + bundle: Mapping[str, Tensor], + ticket: TensorTransferTicket, + ) -> TransferStats: + """Write a tensor bundle and wait for completion.""" + return self.wait(self.send_async(bundle, ticket)) + + def unregister(self, bundle: Mapping[str, Tensor]) -> None: + """Deregister bundle allocations before their tensors are released.""" + for tensor in bundle.values(): + address = tensor.data_ptr() + registration = self._registrations.pop(address, None) + if registration is None: + continue + self.agent.deregister_memory(registration) + self._registered_tensors.pop(address, None) + + def close(self) -> None: + """Release remote metadata and every local memory registration.""" + for remote_agent in tuple(self._remote_agents): + self.agent.remove_remote_agent(remote_agent) + self._remote_agents.remove(remote_agent) + for address, registration in tuple(self._registrations.items()): + self.agent.deregister_memory(registration) + self._registrations.pop(address) + self._registered_tensors.pop(address, None) + + +class RegisteredTensorPool: + """Reuse fixed-shape registered receiver buffers across transfers.""" + + def __init__(self, transport: Any, *, max_buffers_per_bucket: int = 2) -> None: + """Initialize a registered buffer pool. + + Args: + transport: Tensor transport providing ``allocate`` and ``make_ticket``. + max_buffers_per_bucket: Maximum simultaneous leases for one shape bucket. + """ + if max_buffers_per_bucket < 1: + raise ValueError("max_buffers_per_bucket must be positive.") + self.transport = transport + self.max_buffers_per_bucket = max_buffers_per_bucket + self._available: dict[ + tuple[tuple[TensorDescriptor, ...], str], list[PooledTensorBuffer] + ] = {} + self._allocated: dict[ + tuple[tuple[TensorDescriptor, ...], str], list[PooledTensorBuffer] + ] = {} + self._leased: set[int] = set() + self._lock = threading.Lock() + + def acquire( + self, + descriptors: tuple[TensorDescriptor, ...], + *, + device: torch.device, + ) -> PooledTensorBuffer: + """Lease one registered buffer from a fixed-shape bucket.""" + bucket = (descriptors, str(device)) + with self._lock: + available = self._available.setdefault(bucket, []) + if available: + lease = available.pop() + else: + allocated = self._allocated.setdefault(bucket, []) + if len(allocated) >= self.max_buffers_per_bucket: + raise RuntimeError( + "Registered tensor bucket is exhausted; release a lease " + "or increase max_buffers_per_bucket." + ) + bundle = self.transport.allocate(descriptors, device=device) + lease = PooledTensorBuffer( + bundle=bundle, + ticket=self.transport.make_ticket(bundle), + bucket=bucket, + ) + allocated.append(lease) + self._leased.add(id(lease)) + return lease + + def release(self, lease: PooledTensorBuffer) -> None: + """Return one registered buffer lease to its shape bucket.""" + with self._lock: + if id(lease) not in self._leased: + raise ValueError("Registered tensor buffer is not currently leased.") + self._leased.remove(id(lease)) + self._available[lease.bucket].append(lease) + + def close(self) -> None: + """Unregister every pooled allocation after all leases return.""" + with self._lock: + if self._leased: + raise RuntimeError("Cannot close a pool with active buffer leases.") + for leases in self._allocated.values(): + for lease in leases: + self.transport.unregister(lease.bundle) + self._available.clear() + self._allocated.clear() diff --git a/flashdreams/tests/test_cp_attention_subgroup.py b/flashdreams/tests/test_cp_attention_subgroup.py new file mode 100644 index 00000000..56df7aca --- /dev/null +++ b/flashdreams/tests/test_cp_attention_subgroup.py @@ -0,0 +1,77 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU regression tests for context-parallel subgroup rank handling.""" + +from __future__ import annotations + +from typing import cast + +import pytest +import torch +from torch import Tensor +from torch.distributed.tensor.device_mesh import DeviceMesh + +import flashdreams.core.attention.cp as cp_module +from flashdreams.core.attention.cp import ContextParallelAttention + +pytestmark = pytest.mark.ci_cpu + + +class _FakeSubgroupMesh: + def get_rank(self) -> int: + raise AssertionError("ring rotation must not use the global rank") + + def get_local_rank(self) -> int: + return 1 + + def size(self) -> int: + return 3 + + def get_group(self) -> object: + return object() + + +def test_ring_rotation_indexes_all_gather_with_subgroup_local_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + visited_keys: list[float] = [] + + def fake_all_gather( + _local: Tensor, + *, + gather_dim: int, + group: object, + ) -> Tensor: + assert gather_dim == 0 + assert group is not None + return torch.tensor([0.0, 100.0, 1.0, 101.0, 2.0, 102.0]) + + def fake_attention( + query: Tensor, + key: Tensor, + value: Tensor, + *, + return_lse: bool, + ) -> tuple[Tensor, Tensor]: + assert return_lse + visited_keys.append(float(key.item())) + return torch.zeros_like(query), torch.zeros_like(query) + + monkeypatch.setattr(cp_module.funcol, "all_gather_tensor", fake_all_gather) + monkeypatch.setattr(cp_module, "torch_sdpa_cudnn", fake_attention) + + attention = ContextParallelAttention( + backend="cudnn", + method="ring", + convert_to_fp32=False, + ) + attention.device_mesh = cast(DeviceMesh, _FakeSubgroupMesh()) + query = torch.zeros(1, 1, 1, 1) + key = torch.ones(1, 1, 1, 1) + value = torch.full((1, 1, 1, 1), 101.0) + + output = attention._impl_ring(query, key, value) + + assert output.shape == query.shape + assert visited_keys == [1.0, 2.0, 0.0] diff --git a/flashdreams/tests/test_pipeline_stages.py b/flashdreams/tests/test_pipeline_stages.py new file mode 100644 index 00000000..03327c30 --- /dev/null +++ b/flashdreams/tests/test_pipeline_stages.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for independently deployable pipeline stages.""" + +from __future__ import annotations + +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch +from torch import Tensor, nn + +from flashdreams.infra.pipeline import ( + DecoderStage, + DiffusionStage, + StreamingEncoderStage, +) + +pytestmark = pytest.mark.ci_cpu + + +@dataclass +class _Config: + target: nn.Module + + def setup(self) -> nn.Module: + return self.target + + +class _Encoder(nn.Module): + def initialize_autoregressive_cache(self, *, value: int) -> dict[str, int]: + return {"value": value} + + def forward( + self, + *, + input: Tensor, + autoregressive_index: int, + cache: dict[str, int], + ) -> Tensor: + cache["value"] += autoregressive_index + return input + cache["value"] + + +class _Decoder(nn.Module): + def initialize_autoregressive_cache(self) -> list[int]: + return [] + + def forward( + self, + *, + input: Tensor, + autoregressive_index: int, + cache: list[int], + ) -> Tensor: + cache.append(autoregressive_index) + return input * 2 + + +class _Transformer: + device = torch.device("cpu") + + def initialize_autoregressive_cache(self, *, prompt: Tensor) -> dict[str, Tensor]: + return {"prompt": prompt} + + +class _DiffusionModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.transformer = _Transformer() + self.finalized: list[object] = [] + + @property + def device(self) -> torch.device: + return torch.device("cpu") + + def generate( + self, + *, + autoregressive_index: int, + cache: dict[str, Tensor], + input: Tensor, + ) -> tuple[Tensor, object]: + return input + cache["prompt"], SimpleNamespace(index=autoregressive_index) + + def finalize(self, final_state: object) -> None: + self.finalized.append(final_state) + + +def test_encoder_and_decoder_stages_own_only_their_component() -> None: + encoder_stage = StreamingEncoderStage(cast(Any, _Config(_Encoder()))) + encoder_cache = encoder_stage.initialize_cache(value=3) + encoded = encoder_stage.encode(torch.tensor(2), 1, encoder_cache) + assert encoded.item() == 6 + + decoder_stage = DecoderStage(cast(Any, _Config(_Decoder()))) + decoder_cache = decoder_stage.initialize_cache() + decoded = decoder_stage.decode(torch.tensor(4), 0, decoder_cache) + assert decoded.item() == 8 + assert decoder_cache == [0] + + +def test_diffusion_stage_keeps_finalization_state_on_dit_worker() -> None: + model = _DiffusionModel() + stage = DiffusionStage(cast(Any, _Config(model))) + cache = stage.initialize_cache(prompt=torch.tensor(5)) + + output = stage.generate(0, cache, input=torch.tensor(7)) + assert output.item() == 12 + assert cache.final_state is not None + + stage.finalize(0, cache) + assert cache.final_state is None + assert len(model.finalized) == 1 + + +def test_diffusion_stage_rejects_out_of_order_ar_steps() -> None: + stage = DiffusionStage(cast(Any, _Config(_DiffusionModel()))) + cache = stage.initialize_cache(prompt=torch.tensor(0)) + with pytest.raises(AssertionError, match="expected 0, got 1"): + stage.generate(1, cache, input=torch.tensor(0)) diff --git a/flashdreams/tests/test_transfer.py b/flashdreams/tests/test_transfer.py new file mode 100644 index 00000000..30269d4d --- /dev/null +++ b/flashdreams/tests/test_transfer.py @@ -0,0 +1,265 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for Mooncake-backed tensor-bundle transfer.""" + +from __future__ import annotations + +import ctypes + +import pytest +import torch + +from flashdreams.infra.transfer import ( + MooncakeTensorTransport, + NixlTensorTransport, + RegisteredTensorPool, + describe_tensor_bundle, +) + +pytestmark = pytest.mark.ci_cpu + + +class _FakeTransferEngine: + def __init__(self) -> None: + self.registered: set[int] = set() + self.registration_calls = 0 + self._next_batch_id = 1 + + def initialize( + self, + hostname: str, + metadata: str, + protocol: str, + device_name: str, + ) -> int: + assert hostname == "encoder" + assert metadata == "P2PHANDSHAKE" + assert protocol == "rdma" + assert device_name == "mlx5_0" + return 0 + + def get_rpc_port(self) -> int: + return 12345 + + def register_memory(self, address: int, length: int) -> int: + assert length > 0 + self.registered.add(address) + self.registration_calls += 1 + return 0 + + def unregister_memory(self, address: int) -> int: + self.registered.remove(address) + return 0 + + def batch_transfer_sync_write( + self, + session_id: str, + sources: list[int], + destinations: list[int], + lengths: list[int], + ) -> int: + assert session_id == "encoder:12345" + for source, destination, length in zip( + sources, destinations, lengths, strict=True + ): + ctypes.memmove(destination, source, length) + return 0 + + def batch_transfer_async_write( + self, + session_id: str, + sources: list[int], + destinations: list[int], + lengths: list[int], + ) -> int: + assert session_id == "encoder:12345" + for source, destination, length in zip( + sources, destinations, lengths, strict=True + ): + ctypes.memmove(destination, source, length) + batch_id = self._next_batch_id + self._next_batch_id += 1 + return batch_id + + def get_batch_transfer_status(self, batch_ids: list[int]) -> int: + assert batch_ids + return 0 + + +class _FakeNixlAgent: + agents: dict[str, "_FakeNixlAgent"] = {} + + def __init__(self, name: str, config: object) -> None: + del config + self.name = name + self.registrations: dict[int, torch.Tensor] = {} + self.remote_agents: set[str] = set() + self.handles: dict[int, tuple[list[torch.Tensor], list[torch.Tensor]]] = {} + self._next_handle = 1 + self.agents[name] = self + + def register_memory(self, tensor: torch.Tensor) -> tuple[int]: + self.registrations[tensor.data_ptr()] = tensor + return (tensor.data_ptr(),) + + def deregister_memory(self, registration: tuple[int]) -> None: + self.registrations.pop(registration[0]) + + def get_xfer_descs(self, tensors: list[torch.Tensor]) -> list[torch.Tensor]: + return tensors + + def get_serialized_descs(self, tensors: list[torch.Tensor]) -> list[int]: + return [tensor.data_ptr() for tensor in tensors] + + def deserialize_descs(self, addresses: list[int]) -> list[torch.Tensor]: + tensors: list[torch.Tensor] = [] + for agent in self.agents.values(): + for address in addresses: + if address in agent.registrations: + tensors.append(agent.registrations[address]) + return tensors + + def get_agent_metadata(self) -> bytes: + return self.name.encode() + + def add_remote_agent(self, metadata: bytes) -> str: + name = metadata.decode() + self.remote_agents.add(name) + return name + + def remove_remote_agent(self, name: str) -> None: + self.remote_agents.remove(name) + + def initialize_xfer( + self, + operation: str, + local: list[torch.Tensor], + remote: list[torch.Tensor], + remote_agent: str, + ) -> int: + assert operation == "WRITE" + assert remote_agent in self.remote_agents + handle = self._next_handle + self._next_handle += 1 + self.handles[handle] = (local, remote) + return handle + + def transfer(self, handle: int) -> str: + local, remote = self.handles[handle] + for source, destination in zip(local, remote, strict=True): + destination.copy_(source) + return "IN_PROGRESS" + + def check_xfer_state(self, handle: int) -> str: + assert handle in self.handles + return "DONE" + + def release_xfer_handle(self, handle: int) -> None: + self.handles.pop(handle) + + +def _transport() -> MooncakeTensorTransport: + return MooncakeTensorTransport( + hostname="encoder", + device_name="mlx5_0", + engine_factory=_FakeTransferEngine, + ) + + +def test_mooncake_transfer_copies_bundle_and_reports_bandwidth() -> None: + sender = _transport() + receiver = _transport() + source = { + "context": torch.arange(8, dtype=torch.float32), + "mask": torch.ones(2, dtype=torch.int64), + } + destination = receiver.allocate( + describe_tensor_bundle(source), + device=torch.device("cpu"), + ) + stats = sender.send(source, receiver.make_ticket(destination)) + + torch.testing.assert_close(destination["context"], source["context"]) + torch.testing.assert_close(destination["mask"], source["mask"]) + assert stats.backend == "mooncake-rdma" + assert stats.payload_bytes == 48 + assert stats.transfer_ms >= 0.0 + assert stats.registration_ms >= 0.0 + assert stats.bandwidth_gbps > 0.0 + assert stats.asynchronous + sender.unregister(source) + receiver.unregister(destination) + sender.close() + receiver.close() + + +def test_transfer_rejects_noncontiguous_and_mismatched_bundles() -> None: + transport = _transport() + with pytest.raises(ValueError, match="must be contiguous"): + describe_tensor_bundle({"x": torch.ones(2, 3).T}) + + receiver = transport.allocate( + describe_tensor_bundle({"x": torch.ones(2)}), + device=torch.device("cpu"), + ) + ticket = transport.make_ticket(receiver) + with pytest.raises(ValueError, match="does not match"): + transport.send({"x": torch.ones(3)}, ticket) + transport.close() + + +def test_registered_pool_reuses_receiver_registration() -> None: + transport = _transport() + pool = RegisteredTensorPool(transport, max_buffers_per_bucket=1) + descriptors = describe_tensor_bundle({"x": torch.ones(4)}) + + first = pool.acquire(descriptors, device=torch.device("cpu")) + registration_calls = transport.engine.registration_calls + pool.release(first) + second = pool.acquire(descriptors, device=torch.device("cpu")) + + assert second is first + assert transport.engine.registration_calls == registration_calls + pool.release(second) + pool.close() + transport.close() + + +def test_nixl_transfer_uses_serialized_receiver_descriptors() -> None: + _FakeNixlAgent.agents.clear() + sender = NixlTensorTransport( + agent_name="sender", + agent_factory=_FakeNixlAgent, + config_factory=lambda *_: object(), + ) + receiver = NixlTensorTransport( + agent_name="receiver", + agent_factory=_FakeNixlAgent, + config_factory=lambda *_: object(), + ) + source = {"x": torch.arange(8, dtype=torch.float32)} + destination = receiver.allocate( + describe_tensor_bundle(source), + device=torch.device("cpu"), + ) + + stats = sender.send(source, receiver.make_ticket(destination)) + + torch.testing.assert_close(destination["x"], source["x"]) + assert stats.backend == "nixl" + assert stats.asynchronous + sender.close() + receiver.close() diff --git a/integrations/flashvsr/scripts/benchmark_postprocess_session.py b/integrations/flashvsr/scripts/benchmark_postprocess_session.py new file mode 100644 index 00000000..e3a759f5 --- /dev/null +++ b/integrations/flashvsr/scripts/benchmark_postprocess_session.py @@ -0,0 +1,279 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Standalone multi-GPU FlashVSR postprocessor benchmark.""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import subprocess +import time +from pathlib import Path +from typing import Any + +_FIRST_CHUNK_FRAMES = 13 +"""Frame count for the initial call of the FlashVSR 16-frame chunk mode.""" + +_STEADY_CHUNK_FRAMES = 16 +"""Frame count for steady calls of the FlashVSR 16-frame chunk mode.""" + + +def _parse_args() -> argparse.Namespace: + """Parse benchmark arguments.""" + parser = argparse.ArgumentParser( + description=( + "Benchmark a warmed FlashVSR full-attention postprocessor session. " + "Launch this script with torchrun for multi-GPU context parallelism." + ) + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--height", type=int, default=640) + parser.add_argument("--width", type=int, default=1152) + parser.add_argument("--fps", type=float, default=30.0) + parser.add_argument("--warmup-steps", type=int, default=3) + parser.add_argument("--measured-steps", type=int, default=16) + parser.add_argument( + "--compile-cache-root", + type=Path, + default=None, + help=( + "Persistent compile-cache root. Defaults to " + "$FLASHVSR_BENCHMARK_CACHE_ROOT or " + "$FLASHDREAMS_CACHE_DIR/compile/flashvsr-postprocess." + ), + ) + args = parser.parse_args() + if args.height <= 0 or args.width <= 0: + parser.error("--height and --width must be positive") + if args.fps <= 0: + parser.error("--fps must be positive") + if args.warmup_steps < 0: + parser.error("--warmup-steps must be non-negative") + if args.measured_steps <= 0: + parser.error("--measured-steps must be positive") + return args + + +def _default_cache_root() -> Path: + """Return the default persistent FlashVSR compile-cache root.""" + explicit = os.environ.get("FLASHVSR_BENCHMARK_CACHE_ROOT") + if explicit: + return Path(explicit) + flashdreams_cache = os.environ.get("FLASHDREAMS_CACHE_DIR") + if flashdreams_cache: + return Path(flashdreams_cache) / "compile" / "flashvsr-postprocess" + xdg_cache = os.environ.get("XDG_CACHE_HOME") + cache_parent = Path(xdg_cache) if xdg_cache else Path.home() / ".cache" + return cache_parent / "flashdreams" / "compile" / "flashvsr-postprocess" + + +def _configure_compile_cache(cache_root: Path | None) -> tuple[dict[str, str], bool]: + """Set persistent, rank-scoped compiler cache environment variables. + + Args: + cache_root: Explicit cache root; ``None`` selects the persistent default. + + Returns: + Effective compiler cache environment variables and whether this rank's + Inductor cache contained artifacts before launch. + """ + root = (cache_root or _default_cache_root()).expanduser().resolve() + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + rank_suffix = Path(f"world-{world_size}") / f"rank-{local_rank}" + defaults = { + "TORCHINDUCTOR_CACHE_DIR": root / "torchinductor" / rank_suffix, + "TRITON_CACHE_DIR": root / "triton" / rank_suffix, + "TORCH_EXTENSIONS_DIR": root / "torch-extensions" / rank_suffix, + "CUDA_CACHE_PATH": root / "cuda" / rank_suffix, + } + inductor_path = Path( + os.environ.get("TORCHINDUCTOR_CACHE_DIR", defaults["TORCHINDUCTOR_CACHE_DIR"]) + ) + cache_preexisting = inductor_path.is_dir() and any(inductor_path.iterdir()) + effective: dict[str, str] = {} + for name, default in defaults.items(): + value = os.environ.setdefault(name, str(default)) + Path(value).mkdir(parents=True, exist_ok=True) + effective[name] = value + return effective, cache_preexisting + + +def _percentile(values: list[float], q: float) -> float: + """Return a linearly interpolated percentile from non-empty values.""" + ordered = sorted(values) + index = (len(ordered) - 1) * q + lower = int(index) + upper = min(lower + 1, len(ordered) - 1) + fraction = index - lower + return ordered[lower] * (1 - fraction) + ordered[upper] * fraction + + +def _git_commit() -> str: + """Return the current Git commit, or ``unknown`` outside a checkout.""" + result = subprocess.run( + ["git", "rev-parse", "HEAD"], + check=False, + capture_output=True, + text=True, + ) + return result.stdout.strip() if result.returncode == 0 else "unknown" + + +def main() -> None: + """Run the distributed FlashVSR postprocessor benchmark.""" + args = _parse_args() + process_started = time.perf_counter() + + # Configure caches before importing Torch or FlashVSR so Inductor and + # Triton observe the persistent paths during their module initialization. + cache_environment, cache_preexisting = _configure_compile_cache( + args.compile_cache_root + ) + + import torch + import torch.distributed as dist + from flashvsr.postprocess import POSTPROCESS_PRESET_FLASHVSR_V1_1_FULL_ATTN + + from flashdreams.infra.postprocess import VideoChunk, VideoSpec + + dist.init_process_group("nccl") + rank = dist.get_rank() + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = dist.get_world_size() + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + config = POSTPROCESS_PRESET_FLASHVSR_V1_1_FULL_ATTN + + try: + setup_started = time.perf_counter() + session = config.setup().start( + VideoSpec(height=args.height, width=args.width, fps=args.fps) + ) + session.prepare() + dist.barrier() + prepare_seconds = time.perf_counter() - setup_started + + first = torch.zeros( + (1, 3, _FIRST_CHUNK_FRAMES, args.height, args.width), + device=device, + dtype=torch.bfloat16, + ) + steady = torch.zeros( + (1, 3, _STEADY_CHUNK_FRAMES, args.height, args.width), + device=device, + dtype=torch.bfloat16, + ) + + assert session.reset() + session.process(VideoChunk(tensor=first, layout="bcthw")) + torch.cuda.synchronize() + for _ in range(args.warmup_steps): + session.process(VideoChunk(tensor=steady, layout="bcthw")) + torch.cuda.synchronize() + dist.barrier() + warmup_seconds = time.perf_counter() - setup_started - prepare_seconds + startup_seconds = time.perf_counter() - process_started + + torch.cuda.reset_peak_memory_stats(device) + records: list[dict[str, float | int]] = [] + output_shape: list[int] | None = None + for step in range(args.measured_steps): + dist.barrier() + torch.cuda.synchronize() + started = time.perf_counter() + outputs = session.process(VideoChunk(tensor=steady, layout="bcthw")) + torch.cuda.synchronize() + local_elapsed_ms = (time.perf_counter() - started) * 1000.0 + + # Report the slowest rank because every context-parallel call is + # gated by that rank even when host-side timers differ slightly. + elapsed = torch.tensor(local_elapsed_ms, device=device) + dist.all_reduce(elapsed, op=dist.ReduceOp.MAX) + elapsed_ms = float(elapsed.item()) + output_frames = sum(int(chunk.tensor.shape[2]) for chunk in outputs) + if outputs: + output_shape = list(outputs[-1].tensor.shape) + records.append( + { + "step": step, + "elapsed_ms": elapsed_ms, + "input_frames": _STEADY_CHUNK_FRAMES, + "output_frames": output_frames, + "fps": output_frames * 1000.0 / elapsed_ms, + } + ) + + peak_mib = torch.cuda.max_memory_allocated(device) / (1024 * 1024) + peaks: list[float | None] = [None] * world_size + gpu_names: list[str | None] = [None] * world_size + cache_environments: list[dict[str, str] | None] = [None] * world_size + cache_preexisting_by_rank: list[bool | None] = [None] * world_size + dist.all_gather_object(peaks, peak_mib) + dist.all_gather_object(gpu_names, torch.cuda.get_device_name(device)) + dist.all_gather_object(cache_environments, cache_environment) + dist.all_gather_object(cache_preexisting_by_rank, cache_preexisting) + + if rank == 0: + times = [float(record["elapsed_ms"]) for record in records] + total_output_frames = sum( + int(record["output_frames"]) for record in records + ) + result: dict[str, Any] = { + "kind": f"standalone_flashvsr_{world_size}gpu", + "commit": _git_commit(), + "world_size": world_size, + "gpu_names_by_rank": gpu_names, + "input": { + "height": args.height, + "width": args.width, + "frames": _STEADY_CHUNK_FRAMES, + }, + "output_tensor_shape": output_shape, + "dtype": "bfloat16", + "compile_network": config.compile_network, + "use_cuda_graph": config.use_cuda_graph, + "compile_cache_by_rank": cache_environments, + "compile_cache_preexisting_by_rank": cache_preexisting_by_rank, + "prepare_seconds_excluded": prepare_seconds, + "warmup_seconds_excluded": warmup_seconds, + "startup_seconds": startup_seconds, + "warmup_steps_excluded": args.warmup_steps, + "measured_steps": args.measured_steps, + "median_ms": statistics.median(times), + "p90_ms": _percentile(times, 0.90), + "median_fps": _STEADY_CHUNK_FRAMES * 1000.0 / statistics.median(times), + "aggregate_fps": total_output_frames * 1000.0 / sum(times), + "peak_memory_mib_by_rank": peaks, + "software": { + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + }, + "records": records, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(result, indent=2), encoding="utf-8") + print(json.dumps(result, indent=2), flush=True) + dist.barrier() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/README.md b/integrations/lingbot/README.md index b781e4a2..8dfb7c0e 100644 --- a/integrations/lingbot/README.md +++ b/integrations/lingbot/README.md @@ -143,6 +143,181 @@ for i in range(total_blocks): generated_chunks.append(video_chunk.cpu()) # each chunk is [T, C, H, W] ``` +## Three-stage disaggregated inference + +LingBot can load the encoder, DiT, and decoder on three independent GPUs: + +```text +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: + +```bash +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: + +```bash +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: + +```bash +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: + +```bash +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: + +```bash +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. + +The complete pipeline also fits on one H100 80 GB. At 832×464, the measured +single-GPU aggregated run reached **5.56 FPS** and **2157.51 ms median / 2166.25 +ms p90** latency per 12-frame chunk. Initialization peaked at **66.55 GiB** +allocated HBM; rollout peaked at **59.36 GiB**. See the +[single-H100 report](docs/benchmark_h100_aggregated_cp1/README.md). + +Running one complete CP1 pipeline and one session independently on each of +eight H100s reached **43.44 aggregate FPS**, **5.54 median FPS per session**, +and **2163.64 ms median** chunk latency. Rollout peak allocation was **59.35 +GiB per GPU / 474.84 GiB node-wide**. See the +[eight-replica report](docs/benchmark_h100_aggregated_8xcp1/README.md). + +For the eight-GPU aggregated baseline, put the complete pipeline on every GPU +and use all ranks as the DiT context-parallel WORLD group: + +```bash +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](docs/benchmark_h100_aggregated_cp8/README.md) and +[comparison chart](docs/aggregated_vs_disaggregated.svg) include the +resolution-normalized throughput and node-wide HBM tradeoff. + +Validate the data plane without loading checkpoints: + +```bash +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. + +Start with the concise +[disaggregation experiment summary](docs/disaggregation_experiment_summary.md) +for the deployment decision, headline data, limitations, and comparison chart. +The [full H100 experiment record](docs/disaggregated_inference_experiment.md) +contains the tested stack, measurement method, Slurm reproduction procedures, +stage breakdowns, and chronological optimization findings. + +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](https://light-ai.top/LightX2V-BLOG/posts/Disaggregation/): +control-plane messages carry only tensor metadata and registered destination +addresses; Mooncake moves the tensor payload directly between device buffers. + ## Run (WebRTC interactive demo) The `lingbot.webrtc` subpackage exposes a minimal WebRTC server that diff --git a/integrations/lingbot/docs/aggregated_vs_disaggregated.svg b/integrations/lingbot/docs/aggregated_vs_disaggregated.svg new file mode 100644 index 00000000..9d5dc4c7 --- /dev/null +++ b/integrations/lingbot/docs/aggregated_vs_disaggregated.svg @@ -0,0 +1,133 @@ + +LingBot disaggregated CP6 and aggregated CP8 performance comparison +Median per-chunk component wall time and per-rank peak allocated HBM on eight H100 GPUs. + + +LingBot: stage-local CP6 vs full-pipeline CP8 +8× H100 80 GB · BF16 · six warmup + five measured blocks · CP6 832×464 · CP8 832×448 +Median steady-state wall time + +0 + +100 + +200 + +300 + +400 + +500 + +600 + +700 + +800 +milliseconds +Disaggregated CP6 ring + + + +548 + +135 + + + +743 ms +Aggregated CP8 Ulysses + + +309 + +76 + + +393 ms + +Encoder + +Input handoff + +DiT denoise + +KV finalize + +Output handoff + +Decoder + +Coordination +Peak allocated HBM by rank +Stage-local CP6 totals 251.07 GiB; eight full-pipeline replicas total 327.03 GiB (+30.3%) + +0 + +10 + +20 + +30 + +40 + +50 +GiB + +13.7 +G0 + +39.2 +G1 + +39.2 +G2 + +39.2 +G3 + +39.2 +G4 + +39.2 +G5 + +39.2 +G6 + +2.3 +G7 +Disaggregated CP6 + +40.9 +G0 + +40.9 +G1 + +40.9 +G2 + +40.9 +G3 + +40.9 +G4 + +40.9 +G5 + +40.9 +G6 + +40.9 +G7 +Aggregated CP8 +Aggregated bars contain encoder + DiT + decoder on every GPU. + diff --git a/integrations/lingbot/docs/benchmark_h100_1e6d1d/README.md b/integrations/lingbot/docs/benchmark_h100_1e6d1d/README.md new file mode 100644 index 00000000..ca8bd673 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_1e6d1d/README.md @@ -0,0 +1,49 @@ +# LingBot replicated-DiT disaggregation benchmark + +## Result + +Topology: **1 encoder : 6 DiT : 1 decoder**. +Each DiT worker owns one concurrent session and its resident autoregressive KV cache. + +| Metric | Median | P90 | +| --- | ---: | ---: | +| Six-session wave latency | 2657.06 ms | 2671.08 ms | +| Encoder wave | 4.91 ms | 5.21 ms | +| DiT critical path | 2185.88 ms | 2209.48 ms | +| Decoder wave | 42.26 ms | 42.34 ms | +| Encoder → DiT handoff, each | 33.74 ms | 39.20 ms | +| DiT → decoder handoff, each | 31.51 ms | 41.58 ms | +| 256 MiB RDMA probes, all edges | 41.22 GB/s | 42.27 GB/s | + +- Aggregate throughput: **27.20 generated FPS** +- Per-session throughput: **4.53 generated FPS** +- Throughput versus tracked 1:1:1 baseline: **5.07×** +- Wave latency versus one-session baseline latency: **1.19×** +- GPU-normalized throughput versus the three-GPU baseline: **1.90×** + +The headline excludes 6 warmup waves and measures +5 waves. It represents six concurrent, session-affine +rollouts, not acceleration of one autoregressive session. + +## Peak allocated memory + +| Role | Peak | +| --- | ---: | +| Shared encoder | 18.77 GiB | +| DiT workers | 56.34–56.51 GiB each | +| Shared decoder | 2.65 GiB | + +## Reproduction + +```bash +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 integrations/lingbot/docs/benchmark_h100_1e6d1d +``` + +- Repository revision: `08d4c6c159321221c9a2d213c5ebb1359f443ef0` (modified worktree) +- Slurm: job `14628860` on `pool0-00205` +- GPU: `NVIDIA H100 80GB HBM3` × 8 +- Model: `lingbot-world-fast-taehv-window15-sink3` + +For the allocation method, component and memory chart, raw-result checks, and +the observed Mooncake deregistration warning, see the +[full experiment record](../disaggregated_inference_experiment.md). diff --git a/integrations/lingbot/docs/benchmark_h100_1e6d1d/benchmark.json b/integrations/lingbot/docs/benchmark_h100_1e6d1d/benchmark.json new file mode 100644 index 00000000..0992762e --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_1e6d1d/benchmark.json @@ -0,0 +1,2729 @@ +{ + "environment": { + "command": "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 integrations/lingbot/docs/benchmark_h100_1e6d1d", + "commit": "08d4c6c159321221c9a2d213c5ebb1359f443ef0", + "worktree_dirty": true, + "hostname": "pool0-00205", + "slurm_job_id": "14628860", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "encoder": 1, + "dit": 6, + "decoder": 1, + "rank_roles": { + "encoder": [ + 0 + ], + "dit": [ + 1, + 2, + 3, + 4, + 5, + 6 + ], + "decoder": [ + 7 + ] + } + }, + "sessions_per_wave": 6, + "peak_memory_gib_by_rank": [ + 18.772661209106445, + 56.50908803939819, + 56.51177167892456, + 56.33582067489624, + 56.33582067489624, + 56.50908803939819, + 56.51177167892456, + 2.6532506942749023 + ] + }, + "summary": { + "aggregate_fps": 27.198381414481183, + "per_session_fps": 4.5330635690801975, + "throughput_speedup": 5.073727277147145, + "gpu_normalized_speedup": 1.9026477289301793, + "wave_latency_ms": { + "median": 2657.061628997326, + "p90": 2671.0753187537193, + "min": 2595.9399454295635, + "max": 2678.1477704644203 + }, + "latency_vs_baseline": 1.1896047081286683, + "encoder_wave_ms": { + "median": 4.912608027458191, + "p90": 5.2090048551559445, + "min": 4.6869760155677795, + "max": 5.31222403049469 + }, + "dit_critical_path_ms": { + "median": 2185.8832397460938, + "p90": 2209.4816467285154, + "min": 2179.3236083984375, + "max": 2224.4627990722656 + }, + "dit_worker_total_ms": { + "median": 2168.326400756836, + "p90": 2211.0376068115233, + "min": 2137.8883361816406, + "max": 2224.4627990722656 + }, + "decoder_wave_ms": { + "median": 42.25612831115723, + "p90": 42.33598108291626, + "min": 42.21548795700073, + "max": 42.36956834793091 + }, + "encoder_to_dit": { + "payload_mib_each": 14.3583984375, + "copy_ms_each": { + "median": 1.1359388008713722, + "p90": 1.31047572940588, + "min": 1.0589156299829483, + "max": 1.3628657907247543 + }, + "handoff_ms_each": { + "median": 33.743989653885365, + "p90": 39.199621975421906, + "min": 31.6530279815197, + "max": 45.00754736363888 + }, + "aggregate_handoff_ms_per_wave": { + "median": 205.9357985854149, + "p90": 214.27173502743244, + "min": 197.3810624331236, + "max": 215.35456366837025 + } + }, + "dit_to_decoder": { + "payload_mib_each": 0.55224609375, + "copy_ms_each": { + "median": 19.954374991357327, + "p90": 23.790492117404938, + "min": 6.893867626786232, + "max": 35.75972095131874 + }, + "handoff_ms_each": { + "median": 31.50751441717148, + "p90": 41.58355575054885, + "min": 19.808784127235413, + "max": 47.63868823647499 + }, + "aggregate_handoff_ms_per_wave": { + "median": 184.62874926626682, + "p90": 210.51813438534737, + "min": 149.45011585950851, + "max": 213.1159007549286 + } + }, + "bandwidth_probe_gbps": { + "all_edges": { + "median": 41.22417133087569, + "p90": 42.27280309344343, + "min": 33.88831768133709, + "max": 42.486375597128884 + }, + "by_edge": { + "encoder_to_dit_1": { + "median": 41.20407098143717, + "p90": 41.24523356104776, + "min": 40.88172330666681, + "max": 41.250761906819704 + }, + "encoder_to_dit_2": { + "median": 41.1153298170234, + "p90": 41.208036839275344, + "min": 33.88831768133709, + "max": 41.22100513044626 + }, + "encoder_to_dit_3": { + "median": 41.15347778375615, + "p90": 41.26175209862822, + "min": 41.08912865895793, + "max": 41.284144799506784 + }, + "encoder_to_dit_4": { + "median": 41.71502625123036, + "p90": 41.87455345471129, + "min": 39.493816642451975, + "max": 42.164563169041266 + }, + "encoder_to_dit_5": { + "median": 42.420847111804804, + "p90": 42.477715592559804, + "min": 42.23524518063581, + "max": 42.486375597128884 + }, + "encoder_to_dit_6": { + "median": 41.55635369081118, + "p90": 41.59220801269506, + "min": 41.33382858217133, + "max": 41.62772619175502 + }, + "dit_1_to_decoder": { + "median": 42.25909268419748, + "p90": 42.34157330994636, + "min": 41.794954298903605, + "max": 42.35799167150823 + }, + "dit_2_to_decoder": { + "median": 41.42712402181323, + "p90": 41.48538846659678, + "min": 40.827763827879636, + "max": 41.526548710999016 + }, + "dit_3_to_decoder": { + "median": 41.44438939563786, + "p90": 41.498896048148026, + "min": 40.64272131333952, + "max": 41.524945355028485 + }, + "dit_4_to_decoder": { + "median": 40.947550965822124, + "p90": 41.07730217604183, + "min": 40.47969082389142, + "max": 41.109042173505195 + }, + "dit_5_to_decoder": { + "median": 41.02884266459423, + "p90": 41.118244776867236, + "min": 40.22953487886376, + "max": 41.19255104389842 + }, + "dit_6_to_decoder": { + "median": 41.08612508232443, + "p90": 41.180620312038585, + "min": 40.76362448463375, + "max": 41.18249841855751 + } + } + }, + "baseline": { + "fps": 5.360631332509123, + "latency_ms": 2233.5668401792645, + "topology": "1 encoder : 1 DiT : 1 decoder" + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "encoder_times_ms": [ + 631.0217895507812, + 128.63302612304688, + 128.86831665039062, + 129.03219604492188, + 128.77296447753906, + 128.4593963623047 + ], + "encoder_wave_ms": 1274.7876892089844, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.924667447805405, + "transfer_ms": 1.4190319925546646, + "bandwidth_gbps": 10.609959520993684, + "handoff_ms": 33.3186537027359 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.5836580991745, + "transfer_ms": 1.357870176434517, + "bandwidth_gbps": 11.087858221861511, + "handoff_ms": 29.98587116599083 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.75598630309105, + "transfer_ms": 1.306215301156044, + "bandwidth_gbps": 11.526332593619944, + "handoff_ms": 32.153552398085594 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 15.905359759926796, + "transfer_ms": 1.268099993467331, + "bandwidth_gbps": 11.872779810394245, + "handoff_ms": 34.24675390124321 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 16.053514555096626, + "transfer_ms": 1.4264080673456192, + "bandwidth_gbps": 10.555094537580146, + "handoff_ms": 34.73745100200176 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.553848326206207, + "transfer_ms": 1.2946855276823044, + "bandwidth_gbps": 11.628979916808397, + "handoff_ms": 31.542843207716942 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.728944972157478, + "transfer_ms": 1.2201089411973953, + "bandwidth_gbps": 0.4746067998089646, + "handoff_ms": 14.355633407831192 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.042856395244598, + "transfer_ms": 1.162813976407051, + "bandwidth_gbps": 0.49799195034553995, + "handoff_ms": 12.248149141669273 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.898024722933769, + "transfer_ms": 1.003015786409378, + "bandwidth_gbps": 0.5773308933381567, + "handoff_ms": 14.710228890180588 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.719827324151993, + "transfer_ms": 0.7893107831478119, + "bandwidth_gbps": 0.7336425807976817, + "handoff_ms": 12.467645108699799 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.9111507833004, + "transfer_ms": 0.9794607758522034, + "bandwidth_gbps": 0.591215099447106, + "handoff_ms": 12.947460636496544 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.8171598464250565, + "transfer_ms": 0.9855479001998901, + "bandwidth_gbps": 0.5875635267271655, + "handoff_ms": 12.422783300280571 + } + ], + "wave_latency_ms": 137820.007475093, + "decoder_times_ms": [ + 32914.04296875, + 21.73094367980957, + 21.409183502197266, + 21.358463287353516, + 21.36249542236328, + 21.35683250427246 + ], + "decoder_wave_ms": 33021.260887145996, + "output_frames": 54, + "dit_workers": [ + { + "dit_ms": 101758.0546875, + "finalize_ms": 301.6978454589844, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 100733.5546875, + "finalize_ms": 296.3050537109375, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 99900.609375, + "finalize_ms": 294.6186218261719, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 99903.03125, + "finalize_ms": 295.212158203125, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 100432.4765625, + "finalize_ms": 294.6019592285156, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 102934.0078125, + "finalize_ms": 297.0594482421875, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "encoder_times_ms": [ + 184.1183319091797, + 162.93174743652344, + 162.5849609375, + 162.46908569335938, + 162.36204528808594, + 162.39622497558594 + ], + "encoder_wave_ms": 996.8623962402344, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.331158250570297, + "transfer_ms": 1.644585281610489, + "bandwidth_gbps": 9.154813780928572, + "handoff_ms": 43.95357519388199 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.058930799365044, + "transfer_ms": 1.4903955161571503, + "bandwidth_gbps": 10.101930552515483, + "handoff_ms": 35.85589490830898 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.150215312838554, + "transfer_ms": 1.4233868569135666, + "bandwidth_gbps": 10.577498258377025, + "handoff_ms": 48.42544533312321 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 13.922607526183128, + "transfer_ms": 1.2638680636882782, + "bandwidth_gbps": 11.912534569520854, + "handoff_ms": 39.987143129110336 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 13.879535719752312, + "transfer_ms": 1.2350156903266907, + "bandwidth_gbps": 12.190834592568915, + "handoff_ms": 32.19422325491905 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 13.89598660171032, + "transfer_ms": 1.4283768832683563, + "bandwidth_gbps": 10.540545829578074, + "handoff_ms": 35.77662818133831 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.805658012628555, + "transfer_ms": 1.5971045941114426, + "bandwidth_gbps": 0.36257612816032864, + "handoff_ms": 15.008462592959404 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.048822447657585, + "transfer_ms": 1.6164742410182953, + "bandwidth_gbps": 0.35823150490490624, + "handoff_ms": 12.2495386749506 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.929427057504654, + "transfer_ms": 13.455847278237343, + "bandwidth_gbps": 0.04303497119327115, + "handoff_ms": 26.756135746836662 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.424411967396736, + "transfer_ms": 9.507454931735992, + "bandwidth_gbps": 0.06090715172017814, + "handoff_ms": 23.0734683573246 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.326820373535156, + "transfer_ms": 9.08975675702095, + "bandwidth_gbps": 0.06370599516348151, + "handoff_ms": 24.202557280659676 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 6.363740190863609, + "transfer_ms": 9.803596884012222, + "bandwidth_gbps": 0.05906730018085045, + "handoff_ms": 23.352965712547302 + } + ], + "wave_latency_ms": 85165.60726054013, + "decoder_times_ms": [ + 8.234496116638184, + 7.6821441650390625, + 22.41164779663086, + 7.062079906463623, + 7.018303871154785, + 7.0386881828308105 + ], + "decoder_wave_ms": 59.447360038757324, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 82974.953125, + "finalize_ms": 312.3064270019531, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 83423.875, + "finalize_ms": 310.3614196777344, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 83255.71875, + "finalize_ms": 311.4656066894531, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 81741.1328125, + "finalize_ms": 311.13555908203125, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 81733.1640625, + "finalize_ms": 309.6766662597656, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 81146.8984375, + "finalize_ms": 311.32684326171875, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "encoder_times_ms": [ + 163.10134887695312, + 162.64492797851562, + 162.41346740722656, + 162.7038116455078, + 162.38006591796875, + 162.30572509765625 + ], + "encoder_wave_ms": 975.5493469238281, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 15.930622816085815, + "transfer_ms": 1.440398395061493, + "bandwidth_gbps": 10.452574823479473, + "handoff_ms": 38.43693435192108 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.032278209924698, + "transfer_ms": 1.3814959675073624, + "bandwidth_gbps": 10.898238108624637, + "handoff_ms": 32.799992710351944 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.157861471176147, + "transfer_ms": 1.4112405478954315, + "bandwidth_gbps": 10.668536999204472, + "handoff_ms": 34.94852967560291 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.190837740898132, + "transfer_ms": 1.2358911335468292, + "bandwidth_gbps": 12.182199217492418, + "handoff_ms": 34.09750573337078 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.755478128790855, + "transfer_ms": 1.2002959847450256, + "bandwidth_gbps": 12.543466104486106, + "handoff_ms": 35.514894872903824 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.300039038062096, + "transfer_ms": 1.3662409037351608, + "bandwidth_gbps": 11.01992478693824, + "handoff_ms": 34.98373366892338 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.825718700885773, + "transfer_ms": 1.018887385725975, + "bandwidth_gbps": 0.5683375887389175, + "handoff_ms": 12.963667511940002 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.4388845562934875, + "transfer_ms": 12.119635939598083, + "bandwidth_gbps": 0.04777965302637658, + "handoff_ms": 23.137709125876427 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.843028262257576, + "transfer_ms": 10.622333735227585, + "bandwidth_gbps": 0.05451457414481182, + "handoff_ms": 25.98879486322403 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.64203767478466, + "transfer_ms": 8.036358281970024, + "bandwidth_gbps": 0.07205651859738227, + "handoff_ms": 19.288208335638046 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.314299672842026, + "transfer_ms": 8.220219984650612, + "bandwidth_gbps": 0.0704448300752638, + "handoff_ms": 19.162843003869057 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.163408815860748, + "transfer_ms": 20.271245390176773, + "bandwidth_gbps": 0.028566177797867912, + "handoff_ms": 33.20649266242981 + } + ], + "wave_latency_ms": 3133.5156485438347, + "decoder_times_ms": [ + 7.1367998123168945, + 7.032512187957764, + 7.027167797088623, + 7.018752098083496, + 7.018303871154785, + 7.004159927368164 + ], + "decoder_wave_ms": 42.23769569396973, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 1390.972412109375, + "finalize_ms": 335.7643127441406, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 1378.6815185546875, + "finalize_ms": 333.1183166503906, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 1362.6214599609375, + "finalize_ms": 328.4350280761719, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 1372.42822265625, + "finalize_ms": 330.5821533203125, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 1376.3056640625, + "finalize_ms": 331.93896484375, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 1384.0992431640625, + "finalize_ms": 333.9122009277344, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "encoder_times_ms": [ + 163.8098602294922, + 162.706298828125, + 162.3319091796875, + 162.5345001220703, + 162.57286071777344, + 162.98582458496094 + ], + "encoder_wave_ms": 976.9412536621094, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.20489326119423, + "transfer_ms": 1.3157520443201065, + "bandwidth_gbps": 11.442788225178003, + "handoff_ms": 86.56103163957596 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 13.996100053191185, + "transfer_ms": 1.4198459684848785, + "bandwidth_gbps": 10.603876993830648, + "handoff_ms": 32.400140538811684 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.179455116391182, + "transfer_ms": 1.2699998915195465, + "bandwidth_gbps": 11.855018335462807, + "handoff_ms": 34.04638729989529 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.085663482546806, + "transfer_ms": 1.2874789535999298, + "bandwidth_gbps": 11.694072324756968, + "handoff_ms": 33.163150772452354 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.08047042787075, + "transfer_ms": 1.3766884803771973, + "bandwidth_gbps": 10.93629547613768, + "handoff_ms": 33.00353325903416 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 15.803102403879166, + "transfer_ms": 1.3296771794557571, + "bandwidth_gbps": 11.322952843458166, + "handoff_ms": 42.50502027571201 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.825839772820473, + "transfer_ms": 11.893494054675102, + "bandwidth_gbps": 0.04868813128740566, + "handoff_ms": 23.84888008236885 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.126455634832382, + "transfer_ms": 9.779980406165123, + "bandwidth_gbps": 0.05920993457563202, + "handoff_ms": 20.444951951503754 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.887504503130913, + "transfer_ms": 10.193319991230965, + "bandwidth_gbps": 0.05680896905994905, + "handoff_ms": 21.815160289406776 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.6175941824913025, + "transfer_ms": 9.433532133698463, + "bandwidth_gbps": 0.06138443075117527, + "handoff_ms": 20.498700439929962 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.375075921416283, + "transfer_ms": 8.865894749760628, + "bandwidth_gbps": 0.0653145583547148, + "handoff_ms": 19.79709416627884 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.968138411641121, + "transfer_ms": 8.350299671292305, + "bandwidth_gbps": 0.0693474513245082, + "handoff_ms": 19.7947658598423 + } + ], + "wave_latency_ms": 3295.8768773823977, + "decoder_times_ms": [ + 7.147327899932861, + 7.035488128662109, + 7.015007972717285, + 6.9999680519104, + 7.036672115325928, + 6.990911960601807 + ], + "decoder_wave_ms": 42.22537612915039, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 1509.4464111328125, + "finalize_ms": 365.5462951660156, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 1499.294677734375, + "finalize_ms": 369.55841064453125, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 1483.052490234375, + "finalize_ms": 358.6845397949219, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 1498.6370849609375, + "finalize_ms": 362.59619140625, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 1492.728759765625, + "finalize_ms": 363.17803955078125, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 1505.5838623046875, + "finalize_ms": 368.54925537109375, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "encoder_times_ms": [ + 162.6070098876953, + 163.00469970703125, + 162.45033264160156, + 162.49151611328125, + 162.9448699951172, + 166.07859802246094 + ], + "encoder_wave_ms": 979.5770263671875, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.374008402228355, + "transfer_ms": 1.351144164800644, + "bandwidth_gbps": 11.143053711238457, + "handoff_ms": 39.2177514731884 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.262808486819267, + "transfer_ms": 1.4700740575790405, + "bandwidth_gbps": 10.241573832541768, + "handoff_ms": 36.08059696853161 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.34284821152687, + "transfer_ms": 1.2859497219324112, + "bandwidth_gbps": 11.70797873603905, + "handoff_ms": 34.426555037498474 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.2380241304636, + "transfer_ms": 1.2911055237054825, + "bandwidth_gbps": 11.661224991733857, + "handoff_ms": 43.33159141242504 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.174934476613998, + "transfer_ms": 1.3451240956783295, + "bandwidth_gbps": 11.192924168388723, + "handoff_ms": 32.659975811839104 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.05714824795723, + "transfer_ms": 1.3076215982437134, + "bandwidth_gbps": 11.513936463134115, + "handoff_ms": 35.60830466449261 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.79026697576046, + "transfer_ms": 10.369395837187767, + "bandwidth_gbps": 0.05584433356505438, + "handoff_ms": 22.792324423789978 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.098881036043167, + "transfer_ms": 9.963249787688255, + "bandwidth_gbps": 0.05812079515617167, + "handoff_ms": 21.107880398631096 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.875315353274345, + "transfer_ms": 25.770196691155434, + "bandwidth_gbps": 0.022470608468376292, + "handoff_ms": 36.9757916778326 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.595091566443443, + "transfer_ms": 16.17290824651718, + "bandwidth_gbps": 0.035805063082869, + "handoff_ms": 27.216577902436256 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.266144707798958, + "transfer_ms": 17.330363392829895, + "bandwidth_gbps": 0.03341372519860605, + "handoff_ms": 28.133375570178032 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.999438300728798, + "transfer_ms": 16.59962721168995, + "bandwidth_gbps": 0.0348846388304552, + "handoff_ms": 30.48177808523178 + } + ], + "wave_latency_ms": 3462.1315710246563, + "decoder_times_ms": [ + 7.139776229858398, + 7.04256010055542, + 7.016863822937012, + 7.002495765686035, + 7.017375946044922, + 7.060768127441406 + ], + "decoder_wave_ms": 42.27983999252319, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 1625.97314453125, + "finalize_ms": 411.241943359375, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 1615.9259033203125, + "finalize_ms": 409.80328369140625, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 1597.45263671875, + "finalize_ms": 399.94671630859375, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 1619.6253662109375, + "finalize_ms": 402.7091369628906, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 1606.3167724609375, + "finalize_ms": 403.0625, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 1622.154052734375, + "finalize_ms": 406.8531799316406, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "encoder_times_ms": [ + 1.308575987815857, + 1.0932159423828125, + 0.8324800133705139, + 0.8055679798126221, + 0.8092799782752991, + 0.8284800052642822 + ], + "encoder_wave_ms": 5.677599906921387, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.534078538417816, + "transfer_ms": 1.4062989503145218, + "bandwidth_gbps": 10.706025199430549, + "handoff_ms": 40.827520191669464 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.196760952472687, + "transfer_ms": 1.1651013046503067, + "bandwidth_gbps": 12.922371591128607, + "handoff_ms": 32.26822055876255 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.273058623075485, + "transfer_ms": 1.1318475008010864, + "bandwidth_gbps": 13.302032287338996, + "handoff_ms": 33.697087317705154 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 17.810095101594925, + "transfer_ms": 1.1323317885398865, + "bandwidth_gbps": 13.296343132267063, + "handoff_ms": 41.02505184710026 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.133481308817863, + "transfer_ms": 1.3029277324676514, + "bandwidth_gbps": 11.555416025634255, + "handoff_ms": 32.19245560467243 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.081116765737534, + "transfer_ms": 1.1410079896450043, + "bandwidth_gbps": 13.195238014664781, + "handoff_ms": 33.94414484500885 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.237618461251259, + "transfer_ms": 18.08556169271469, + "bandwidth_gbps": 0.03201846919873461, + "handoff_ms": 30.2664153277874 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.077039659023285, + "transfer_ms": 27.521099895238876, + "bandwidth_gbps": 0.021041019516090595, + "handoff_ms": 37.93354332447052 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.871148616075516, + "transfer_ms": 10.235585272312164, + "bandwidth_gbps": 0.056574390676654554, + "handoff_ms": 23.736413568258286 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.625637084245682, + "transfer_ms": 7.984304800629616, + "bandwidth_gbps": 0.07252628932131151, + "handoff_ms": 19.229505211114883 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.297537729144096, + "transfer_ms": 9.728806093335152, + "bandwidth_gbps": 0.05952138365638729, + "handoff_ms": 20.67144773900509 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.9652159214019775, + "transfer_ms": 8.403722196817398, + "bandwidth_gbps": 0.06890660905227237, + "handoff_ms": 20.492931827902794 + } + ], + "wave_latency_ms": 2626.0259561240673, + "decoder_times_ms": [ + 7.133024215698242, + 7.034336090087891, + 7.01632022857666, + 7.0062079429626465, + 7.030335903167725, + 6.999263763427734 + ], + "decoder_wave_ms": 42.2194881439209, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 1748.7366943359375, + "finalize_ms": 448.4148864746094, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 1739.1812744140625, + "finalize_ms": 439.6708068847656, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 1719.5787353515625, + "finalize_ms": 439.2612609863281, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 1731.141357421875, + "finalize_ms": 440.4485778808594, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 1741.504150390625, + "finalize_ms": 441.9870910644531, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 1744.4736328125, + "finalize_ms": 441.36785888671875, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "encoder_times_ms": [ + 1.0594559907913208, + 0.8283200263977051, + 0.8167999982833862, + 0.8398399949073792, + 0.948032021522522, + 0.8197759985923767 + ], + "encoder_wave_ms": 5.31222403049469, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.364119619131088, + "transfer_ms": 1.171741634607315, + "bandwidth_gbps": 12.849139738083698, + "handoff_ms": 33.94879028201103 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 18.70782859623432, + "transfer_ms": 1.127868890762329, + "bandwidth_gbps": 13.348955825728735, + "handoff_ms": 40.114378556609154 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 19.023790955543518, + "transfer_ms": 1.163853332400322, + "bandwidth_gbps": 12.936227942870506, + "handoff_ms": 39.49657268822193 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.136271551251411, + "transfer_ms": 1.3392921537160873, + "bandwidth_gbps": 11.241663708866655, + "handoff_ms": 35.39552353322506 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.097023755311966, + "transfer_ms": 1.3628657907247543, + "bandwidth_gbps": 11.047215435639838, + "handoff_ms": 32.44374319911003 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.09466378390789, + "transfer_ms": 1.3101845979690552, + "bandwidth_gbps": 11.491412754613682, + "handoff_ms": 33.95555540919304 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.8156119883060455, + "transfer_ms": 35.75972095131874, + "bandwidth_gbps": 0.016193414953889484, + "handoff_ms": 47.63868823647499 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.096759483218193, + "transfer_ms": 12.198098003864288, + "bandwidth_gbps": 0.04747231903011054, + "handoff_ms": 23.58475886285305 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.929840564727783, + "transfer_ms": 10.130221024155617, + "bandwidth_gbps": 0.05716281990483691, + "handoff_ms": 38.295913487672806 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.354747176170349, + "transfer_ms": 6.893867626786232, + "bandwidth_gbps": 0.08399813157856506, + "handoff_ms": 26.389088481664658 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.287254065275192, + "transfer_ms": 8.20172019302845, + "bandwidth_gbps": 0.0706037253614452, + "handoff_ms": 19.808784127235413 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.017247051000595, + "transfer_ms": 8.185800164937973, + "bandwidth_gbps": 0.07074103793546342, + "handoff_ms": 20.8461694419384 + } + ], + "wave_latency_ms": 2678.1477704644203, + "decoder_times_ms": [ + 7.150047779083252, + 7.061728000640869, + 7.033664226531982, + 7.024447917938232, + 7.008768081665039, + 7.006944179534912 + ], + "decoder_wave_ms": 42.28560018539429, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 1774.71240234375, + "finalize_ms": 448.3245849609375, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 1762.305908203125, + "finalize_ms": 443.3638610839844, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 1784.869384765625, + "finalize_ms": 439.5934143066406, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 1771.6845703125, + "finalize_ms": 441.2993469238281, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 1763.5460205078125, + "finalize_ms": 439.1201171875, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 1765.5919189453125, + "finalize_ms": 445.22943115234375, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "encoder_times_ms": [ + 1.0465600490570068, + 0.7767040133476257, + 0.9567360281944275, + 0.7522559762001038, + 0.7503679990768433, + 0.7715520262718201 + ], + "encoder_wave_ms": 5.054176092147827, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.292668551206589, + "transfer_ms": 1.2981556355953217, + "bandwidth_gbps": 11.597894418180084, + "handoff_ms": 34.245479851961136 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 16.13219454884529, + "transfer_ms": 1.1743362993001938, + "bandwidth_gbps": 12.820749906966208, + "handoff_ms": 33.694472163915634 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.18815366923809, + "transfer_ms": 1.136263832449913, + "bandwidth_gbps": 13.250331102713918, + "handoff_ms": 33.61496888101101 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.03358019888401, + "transfer_ms": 1.1538118124008179, + "bandwidth_gbps": 13.048810766352082, + "handoff_ms": 33.70863199234009 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 16.80907793343067, + "transfer_ms": 1.2906957417726517, + "bandwidth_gbps": 11.664927304495594, + "handoff_ms": 35.059716552495956 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.061834663152695, + "transfer_ms": 1.116424798965454, + "bandwidth_gbps": 13.485791442425562, + "handoff_ms": 33.504536375403404 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.7867633402347565, + "transfer_ms": 9.996671229600906, + "bandwidth_gbps": 0.05792648239599234, + "handoff_ms": 21.9383854418993 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 9.046686813235283, + "transfer_ms": 10.411426424980164, + "bandwidth_gbps": 0.05561889181780423, + "handoff_ms": 29.666466638445854 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.896808415651321, + "transfer_ms": 10.435864329338074, + "bandwidth_gbps": 0.05548864777516031, + "handoff_ms": 21.829405799508095 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 9.015770629048347, + "transfer_ms": 7.982032373547554, + "bandwidth_gbps": 0.07254693703311, + "handoff_ms": 24.857282638549805 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.388023167848587, + "transfer_ms": 15.200864523649216, + "bandwidth_gbps": 0.0380946754113288, + "handoff_ms": 29.80240248143673 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.101105198264122, + "transfer_ms": 9.39272902905941, + "bandwidth_gbps": 0.061651091840130345, + "handoff_ms": 21.35617285966873 + } + ], + "wave_latency_ms": 2595.9399454295635, + "decoder_times_ms": [ + 7.139743804931641, + 7.014080047607422, + 7.026879787445068, + 7.011712074279785, + 7.005152225494385, + 7.017920017242432 + ], + "decoder_wave_ms": 42.21548795700073, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 1733.865966796875, + "finalize_ms": 449.26220703125, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 1725.5482177734375, + "finalize_ms": 443.6260070800781, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 1711.588623046875, + "finalize_ms": 435.9664306640625, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 1714.28125, + "finalize_ms": 440.7195739746094, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 1718.1455078125, + "finalize_ms": 437.8485107421875, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 1732.28271484375, + "finalize_ms": 446.1881408691406, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "encoder_times_ms": [ + 0.9593920111656189, + 0.7674559950828552, + 0.7548800110816956, + 0.7544000148773193, + 0.9076160192489624, + 0.7688639760017395 + ], + "encoder_wave_ms": 4.912608027458191, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.243414625525475, + "transfer_ms": 1.1516641825437546, + "bandwidth_gbps": 13.073144262197275, + "handoff_ms": 45.00754736363888 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.113832265138626, + "transfer_ms": 1.1246632784605026, + "bandwidth_gbps": 13.387004171247822, + "handoff_ms": 31.6530279815197 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.23366367816925, + "transfer_ms": 1.0866131633520126, + "bandwidth_gbps": 13.855779138138962, + "handoff_ms": 33.39028172194958 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 15.89745283126831, + "transfer_ms": 1.1356137692928314, + "bandwidth_gbps": 13.25791603370183, + "handoff_ms": 34.44983996450901 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.162980020046234, + "transfer_ms": 1.0801292955875397, + "bandwidth_gbps": 13.938953476685688, + "handoff_ms": 32.13808685541153 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.133434742689133, + "transfer_ms": 1.0937117040157318, + "bandwidth_gbps": 13.765850675932274, + "handoff_ms": 36.00870817899704 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.837507382035255, + "transfer_ms": 9.807426482439041, + "bandwidth_gbps": 0.05904423561439623, + "handoff_ms": 21.44324779510498 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.157181829214096, + "transfer_ms": 10.267935693264008, + "bandwidth_gbps": 0.05639614595364909, + "handoff_ms": 24.083444848656654 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.913916811347008, + "transfer_ms": 24.672089144587517, + "bandwidth_gbps": 0.023470732316441673, + "handoff_ms": 40.415188297629356 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.312897264957428, + "transfer_ms": 20.074591040611267, + "bandwidth_gbps": 0.028846017277688332, + "handoff_ms": 34.6575602889061 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.441197961568832, + "transfer_ms": 20.095979794859886, + "bandwidth_gbps": 0.028815315596013587, + "handoff_ms": 31.322389841079712 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.974089562892914, + "transfer_ms": 20.98916843533516, + "bandwidth_gbps": 0.027589087284903353, + "handoff_ms": 32.70691819489002 + } + ], + "wave_latency_ms": 2644.465768709779, + "decoder_times_ms": [ + 7.118207931518555, + 7.026815891265869, + 7.021056175231934, + 7.032544136047363, + 7.024928092956543, + 7.1460161209106445 + ], + "decoder_wave_ms": 42.36956834793091, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 1741.7938232421875, + "finalize_ms": 445.2160949707031, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 1724.51416015625, + "finalize_ms": 442.96441650390625, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 1701.4180908203125, + "finalize_ms": 436.4702453613281, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 1713.4537353515625, + "finalize_ms": 440.8396911621094, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 1722.6947021484375, + "finalize_ms": 438.60516357421875, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 1732.0794677734375, + "finalize_ms": 442.53387451171875, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "encoder_times_ms": [ + 0.8568959832191467, + 0.7789760231971741, + 0.7565119862556458, + 0.7626559734344482, + 0.7631040215492249, + 0.7688320279121399 + ], + "encoder_wave_ms": 4.6869760155677795, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.250962063670158, + "transfer_ms": 1.140505075454712, + "bandwidth_gbps": 13.201056552946353, + "handoff_ms": 33.826472237706184 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.106879010796547, + "transfer_ms": 1.1196192353963852, + "bandwidth_gbps": 13.447314519091558, + "handoff_ms": 31.795300543308258 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.168970286846161, + "transfer_ms": 1.0834354907274246, + "bandwidth_gbps": 13.89641757987031, + "handoff_ms": 33.61984342336655 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.100492000579834, + "transfer_ms": 1.1082515120506287, + "bandwidth_gbps": 13.585248327017123, + "handoff_ms": 32.527538016438484 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.14940319955349, + "transfer_ms": 1.0589156299829483, + "bandwidth_gbps": 14.218197912748199, + "handoff_ms": 31.93596377968788 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.051154255867004, + "transfer_ms": 1.103455200791359, + "bandwidth_gbps": 13.644298372242446, + "handoff_ms": 33.675944432616234 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.8334039747715, + "transfer_ms": 9.909860789775848, + "bandwidth_gbps": 0.058433918728448456, + "handoff_ms": 21.95270173251629 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.208503291010857, + "transfer_ms": 24.36334826052189, + "bandwidth_gbps": 0.023768161658564893, + "handoff_ms": 41.02863185107708 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 20.864687860012054, + "transfer_ms": 20.068103447556496, + "bandwidth_gbps": 0.028855342584478662, + "handoff_ms": 47.45940305292606 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.34147210419178, + "transfer_ms": 21.016493439674377, + "bandwidth_gbps": 0.027553216794331793, + "handoff_ms": 34.459641203284264 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.332359880208969, + "transfer_ms": 20.118704065680504, + "bandwidth_gbps": 0.028782768418360014, + "handoff_ms": 35.80939956009388 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.003131926059723, + "transfer_ms": 20.104166120290756, + "bandwidth_gbps": 0.028803582129952335, + "handoff_ms": 32.40612335503101 + } + ], + "wave_latency_ms": 2657.061628997326, + "decoder_times_ms": [ + 7.125088214874268, + 7.033631801605225, + 7.018400192260742, + 7.028895854949951, + 7.0215678215026855, + 7.004159927368164 + ], + "decoder_wave_ms": 42.231743812561035, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 1733.860107421875, + "finalize_ms": 444.21160888671875, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 1720.5631103515625, + "finalize_ms": 439.8471984863281, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 1709.021240234375, + "finalize_ms": 438.38226318359375, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 1709.5009765625, + "finalize_ms": 440.72357177734375, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 1722.126953125, + "finalize_ms": 441.9761962890625, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 1734.1160888671875, + "finalize_ms": 445.20751953125, + "session_index": 5, + "rank": 6 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "encoder_times_ms": [ + 0.9499199986457825, + 0.7614719867706299, + 0.7561280131340027, + 0.7758079767227173, + 0.7519040107727051, + 0.8536959886550903 + ], + "encoder_wave_ms": 4.848927974700928, + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.258474111557007, + "transfer_ms": 1.1504162102937698, + "bandwidth_gbps": 13.087326017559626, + "handoff_ms": 34.33205001056194 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 22.159012034535408, + "transfer_ms": 1.1064503341913223, + "bandwidth_gbps": 13.607363597575278, + "handoff_ms": 39.16662745177746 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.360159635543823, + "transfer_ms": 1.0927990078926086, + "bandwidth_gbps": 13.777347793382667, + "handoff_ms": 33.77934731543064 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.131084084510803, + "transfer_ms": 1.1116117238998413, + "bandwidth_gbps": 13.544182448148204, + "handoff_ms": 32.398153096437454 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.077415689826012, + "transfer_ms": 1.3130959123373032, + "bandwidth_gbps": 11.465934710893004, + "handoff_ms": 32.34328143298626 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 14.123434200882912, + "transfer_ms": 1.1486038565635681, + "bandwidth_gbps": 13.107976186885415, + "handoff_ms": 33.91633927822113 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 15.624841675162315, + "transfer_ms": 23.726841434836388, + "bandwidth_gbps": 0.024405776958992563, + "handoff_ms": 46.57787084579468 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.092097282409668, + "transfer_ms": 21.46412804722786, + "bandwidth_gbps": 0.02697859417936097, + "handoff_ms": 32.09282457828522 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.884099587798119, + "transfer_ms": 20.07398195564747, + "bandwidth_gbps": 0.02884689252383671, + "handoff_ms": 31.35775215923786 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.322568118572235, + "transfer_ms": 19.840646535158157, + "bandwidth_gbps": 0.02918614567190988, + "handoff_ms": 31.75538033246994 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.375739023089409, + "transfer_ms": 20.71194536983967, + "bandwidth_gbps": 0.027958358795366147, + "handoff_ms": 31.657276675105095 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.969332367181778, + "transfer_ms": 20.296961069107056, + "bandwidth_gbps": 0.02852998525387011, + "handoff_ms": 33.18038024008274 + } + ], + "wave_latency_ms": 2660.466641187668, + "decoder_times_ms": [ + 7.133088111877441, + 7.052864074707031, + 7.023519992828369, + 7.02236795425415, + 7.013247966766357, + 7.011040210723877 + ], + "decoder_wave_ms": 42.25612831115723, + "output_frames": 72, + "dit_workers": [ + { + "dit_ms": 1738.6453857421875, + "finalize_ms": 447.23785400390625, + "session_index": 0, + "rank": 1 + }, + { + "dit_ms": 1719.7462158203125, + "finalize_ms": 438.96905517578125, + "session_index": 1, + "rank": 2 + }, + { + "dit_ms": 1709.420654296875, + "finalize_ms": 439.6869201660156, + "session_index": 2, + "rank": 3 + }, + { + "dit_ms": 1706.7230224609375, + "finalize_ms": 440.58642578125, + "session_index": 3, + "rank": 4 + }, + { + "dit_ms": 1723.7269287109375, + "finalize_ms": 443.4239807128906, + "session_index": 4, + "rank": 5 + }, + { + "dit_ms": 1733.6558837890625, + "finalize_ms": 444.4276428222656, + "session_index": 5, + "rank": 6 + } + ] + } + ], + "bandwidth_probe": { + "encoder_to_dit_1": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.004446133971214294, + "transfer_ms": 6.56614825129509, + "bandwidth_gbps": 40.88172330666681 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.00194646418094635, + "transfer_ms": 6.51608407497406, + "bandwidth_gbps": 41.195824503088325 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0017639249563217163, + "transfer_ms": 6.5086521208286285, + "bandwidth_gbps": 41.24286427000265 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0021867454051971436, + "transfer_ms": 6.516415625810623, + "bandwidth_gbps": 41.19372848729356 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014808028936386108, + "transfer_ms": 6.551502272486687, + "bandwidth_gbps": 40.973114994908286 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.00169314444065094, + "transfer_ms": 6.513476371765137, + "bandwidth_gbps": 41.212317459786014 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013355165719985962, + "transfer_ms": 6.507406011223793, + "bandwidth_gbps": 41.250761906819704 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001296401023864746, + "transfer_ms": 6.5111033618450165, + "bandwidth_gbps": 41.227337531305125 + } + ], + "encoder_to_dit_2": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0028014183044433594, + "transfer_ms": 6.530014798045158, + "bandwidth_gbps": 41.10793992080378 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002125278115272522, + "transfer_ms": 6.5121036022901535, + "bandwidth_gbps": 41.22100513044626 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015459954738616943, + "transfer_ms": 6.527667865157127, + "bandwidth_gbps": 41.12271971324302 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001514330506324768, + "transfer_ms": 7.921179756522179, + "bandwidth_gbps": 33.88831768133709 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0023562461137771606, + "transfer_ms": 6.554935127496719, + "bandwidth_gbps": 40.95165715278612 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0017285346984863281, + "transfer_ms": 6.525227800011635, + "bandwidth_gbps": 41.13809727830826 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0018086284399032593, + "transfer_ms": 6.515031680464745, + "bandwidth_gbps": 41.202479000202096 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0029262155294418335, + "transfer_ms": 6.537208333611488, + "bandwidth_gbps": 41.062704797064725 + } + ], + "encoder_to_dit_3": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002518296241760254, + "transfer_ms": 6.513189524412155, + "bandwidth_gbps": 41.214132491289284 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0018477439880371094, + "transfer_ms": 6.533004343509674, + "bandwidth_gbps": 41.08912865895793 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015739351511001587, + "transfer_ms": 6.5285321325063705, + "bandwidth_gbps": 41.1172757599563 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013336539268493652, + "transfer_ms": 6.518153473734856, + "bandwidth_gbps": 41.18274555542007 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016652047634124756, + "transfer_ms": 6.528481841087341, + "bandwidth_gbps": 41.11759250222424 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013262033462524414, + "transfer_ms": 6.527431309223175, + "bandwidth_gbps": 41.124210012092234 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013820827007293701, + "transfer_ms": 6.507186219096184, + "bandwidth_gbps": 41.25215522682312 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001300126314163208, + "transfer_ms": 6.502144038677216, + "bandwidth_gbps": 41.284144799506784 + } + ], + "encoder_to_dit_4": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.00478699803352356, + "transfer_ms": 6.366375833749771, + "bandwidth_gbps": 42.164563169041266 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0020917505025863647, + "transfer_ms": 6.546150892972946, + "bandwidth_gbps": 41.00660989775773 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.003164634108543396, + "transfer_ms": 6.436444818973541, + "bandwidth_gbps": 41.70554763534958 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016111880540847778, + "transfer_ms": 6.796898320317268, + "bandwidth_gbps": 39.493816642451975 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016968697309494019, + "transfer_ms": 6.458917632699013, + "bandwidth_gbps": 41.56043957597704 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015310943126678467, + "transfer_ms": 6.4295511692762375, + "bandwidth_gbps": 41.750263577141304 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001434236764907837, + "transfer_ms": 6.433520466089249, + "bandwidth_gbps": 41.72450486711114 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.00155717134475708, + "transfer_ms": 6.429582834243774, + "bandwidth_gbps": 41.75005796181993 + } + ], + "encoder_to_dit_5": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.00278279185295105, + "transfer_ms": 6.35572150349617, + "bandwidth_gbps": 42.23524518063581 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001821666955947876, + "transfer_ms": 6.326209753751755, + "bandwidth_gbps": 42.43227247417848 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014156103134155273, + "transfer_ms": 6.335984915494919, + "bandwidth_gbps": 42.36680793597373 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013243407011032104, + "transfer_ms": 6.318153813481331, + "bandwidth_gbps": 42.486375597128884 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0018477439880371094, + "transfer_ms": 6.323313340544701, + "bandwidth_gbps": 42.451708707649864 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016391277313232422, + "transfer_ms": 6.319994106888771, + "bandwidth_gbps": 42.4740041620302 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013150274753570557, + "transfer_ms": 6.33593462407589, + "bandwidth_gbps": 42.36714422209682 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015310943126678467, + "transfer_ms": 6.329618394374847, + "bandwidth_gbps": 42.40942174943113 + } + ], + "encoder_to_dit_6": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002726912498474121, + "transfer_ms": 6.462186574935913, + "bandwidth_gbps": 41.53941593719183 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001950189471244812, + "transfer_ms": 6.4943283796310425, + "bandwidth_gbps": 41.33382858217133 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014640390872955322, + "transfer_ms": 6.456347182393074, + "bandwidth_gbps": 41.57698593595507 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0017266720533370972, + "transfer_ms": 6.458705291152, + "bandwidth_gbps": 41.5618059501397 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001339241862297058, + "transfer_ms": 6.448477506637573, + "bandwidth_gbps": 41.62772619175502 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.005323439836502075, + "transfer_ms": 6.458412855863571, + "bandwidth_gbps": 41.563687858122655 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014025717973709106, + "transfer_ms": 6.477518007159233, + "bandwidth_gbps": 41.44109760919437 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015087425708770752, + "transfer_ms": 6.460400298237801, + "bandwidth_gbps": 41.55090143148266 + } + ], + "dit_1_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0031925737857818604, + "transfer_ms": 6.422676146030426, + "bandwidth_gbps": 41.794954298903605 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001994892954826355, + "transfer_ms": 6.3529908657073975, + "bandwidth_gbps": 42.253398702173335 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0027511268854141235, + "transfer_ms": 6.348870694637299, + "bandwidth_gbps": 42.28081952066521 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0018514692783355713, + "transfer_ms": 6.337303668260574, + "bandwidth_gbps": 42.35799167150823 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015813857316970825, + "transfer_ms": 6.340814754366875, + "bandwidth_gbps": 42.334536869276995 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015106052160263062, + "transfer_ms": 6.351279094815254, + "bandwidth_gbps": 42.26478666622164 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0012833625078201294, + "transfer_ms": 6.358066573739052, + "bandwidth_gbps": 42.2196673920856 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013820827007293701, + "transfer_ms": 6.354004144668579, + "bandwidth_gbps": 42.24666051331344 + } + ], + "dit_2_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0039637088775634766, + "transfer_ms": 6.574826315045357, + "bandwidth_gbps": 40.827763827879636 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0018365681171417236, + "transfer_ms": 6.483161821961403, + "bandwidth_gbps": 41.405021711888736 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0024996697902679443, + "transfer_ms": 6.464188918471336, + "bandwidth_gbps": 41.526548710999016 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015851110219955444, + "transfer_ms": 6.47493451833725, + "bandwidth_gbps": 41.45763254281275 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014659017324447632, + "transfer_ms": 6.494257599115372, + "bandwidth_gbps": 41.33427907703653 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014919787645339966, + "transfer_ms": 6.473354995250702, + "bandwidth_gbps": 41.46774836185297 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0012926757335662842, + "transfer_ms": 6.476247683167458, + "bandwidth_gbps": 41.44922633173773 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.006020069122314453, + "transfer_ms": 6.494706496596336, + "bandwidth_gbps": 41.33142215751832 + } + ], + "dit_3_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0040531158447265625, + "transfer_ms": 6.6047608852386475, + "bandwidth_gbps": 40.64272131333952 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0019278377294540405, + "transfer_ms": 6.478607654571533, + "bandwidth_gbps": 41.43412756452114 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0027213245630264282, + "transfer_ms": 6.472621113061905, + "bandwidth_gbps": 41.47245008027595 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016689300537109375, + "transfer_ms": 6.505731493234634, + "bandwidth_gbps": 41.26137948963131 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015255063772201538, + "transfer_ms": 6.470236927270889, + "bandwidth_gbps": 41.48773205948497 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013671815395355225, + "transfer_ms": 6.4754001796245575, + "bandwidth_gbps": 41.454651226754585 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013709068298339844, + "transfer_ms": 6.464438512921333, + "bandwidth_gbps": 41.524945355028485 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.00155717134475708, + "transfer_ms": 6.51620514690876, + "bandwidth_gbps": 41.19505907933912 + } + ], + "dit_4_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0031758099794387817, + "transfer_ms": 6.631361320614815, + "bandwidth_gbps": 40.47969082389142 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0017173588275909424, + "transfer_ms": 6.550529971718788, + "bandwidth_gbps": 40.979196669420844 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002469867467880249, + "transfer_ms": 6.567204371094704, + "bandwidth_gbps": 40.87514882002276 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015534460544586182, + "transfer_ms": 6.5928734838962555, + "bandwidth_gbps": 40.71600291673731 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001341104507446289, + "transfer_ms": 6.539035588502884, + "bandwidth_gbps": 41.051230317811815 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013224780559539795, + "transfer_ms": 6.5606627613306046, + "bandwidth_gbps": 40.915905262223404 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.00130385160446167, + "transfer_ms": 6.529839709401131, + "bandwidth_gbps": 41.109042173505195 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0012759119272232056, + "transfer_ms": 6.537050008773804, + "bandwidth_gbps": 41.06369931998611 + } + ], + "dit_5_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.006245449185371399, + "transfer_ms": 6.67259655892849, + "bandwidth_gbps": 40.22953487886376 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0026691704988479614, + "transfer_ms": 6.592424586415291, + "bandwidth_gbps": 40.71877538851984 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0028070062398910522, + "transfer_ms": 6.544835865497589, + "bandwidth_gbps": 41.01484919050624 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016409903764724731, + "transfer_ms": 6.559696048498154, + "bandwidth_gbps": 40.9219351042124 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013746321201324463, + "transfer_ms": 6.533438339829445, + "bandwidth_gbps": 41.08639923385387 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014137476682662964, + "transfer_ms": 6.540372967720032, + "bandwidth_gbps": 41.04283613868222 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014919787645339966, + "transfer_ms": 6.516601890325546, + "bandwidth_gbps": 41.19255104389842 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013243407011032104, + "transfer_ms": 6.5385326743125916, + "bandwidth_gbps": 41.05438779171064 + } + ], + "dit_6_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.005111098289489746, + "transfer_ms": 6.5851714462041855, + "bandwidth_gbps": 40.76362448463375 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002508983016014099, + "transfer_ms": 6.547164171934128, + "bandwidth_gbps": 41.000263465319556 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0029001384973526, + "transfer_ms": 6.54309056699276, + "bandwidth_gbps": 41.02578945707218 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016838312149047852, + "transfer_ms": 6.5181925892829895, + "bandwidth_gbps": 41.18249841855751 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016521662473678589, + "transfer_ms": 6.534557789564133, + "bandwidth_gbps": 41.079360630752824 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014994293451309204, + "transfer_ms": 6.532406434416771, + "bandwidth_gbps": 41.09288953389603 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015124678611755371, + "transfer_ms": 6.518617272377014, + "bandwidth_gbps": 41.17981540924476 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013504177331924438, + "transfer_ms": 6.53238408267498, + "bandwidth_gbps": 41.09303014070124 + } + ] + } +} diff --git a/integrations/lingbot/docs/benchmark_h100_1io7dit_optimized/README.md b/integrations/lingbot/docs/benchmark_h100_1io7dit_optimized/README.md new file mode 100644 index 00000000..844aca82 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_1io7dit_optimized/README.md @@ -0,0 +1,188 @@ + + +# LingBot 1-I/O : 7-DiT optimization benchmark + +## Result + +The optimized eight-H100 layout co-locates encoder and decoder on GPU 0 and +uses GPUs 1–7 as independent, session-affine DiT workers. Five measured +seven-session waves reached **35.15 aggregate generated FPS** and **5.02 FPS +per session**. + +| Metric | 1E:6DiT:1D tracked | 1IO:7DiT synchronous | 1IO:7DiT optimized | +| --- | ---: | ---: | ---: | +| Aggregate FPS | 27.20 | 31.57 | **35.15** | +| FPS per session | 4.53 | 4.51 | **5.02** | +| Median wave | 2657.06 ms | 2671.78 ms | **2358.51 ms** | +| P90 wave | 2671.08 ms | 2705.95 ms | **2454.10 ms** | +| Median 256 MiB RDMA probe | 41.22 GB/s | 42.28 GB/s | **41.97 GB/s** | +| I/O GPU peak HBM | encoder 18.77 + decoder 2.65 GiB on separate GPUs | 20.51 GiB | 20.59 GiB | +| DiT peak HBM, each | 56.34–56.51 GiB | 56.29 GiB | 56.29–56.51 GiB | + +The optimized path is **29.2% faster** than the tracked six-DiT topology and +**11.4% faster** than the same co-located seven-DiT topology using synchronous, +per-request allocation and registration. It is 10.9% above the earlier +unvalidated linear projection of 31.7 FPS. + +This is aggregate capacity for seven concurrent sessions. It does not make one +autoregressive session run at 35 FPS. Each individual session measured 5.02 +generated FPS; the aggregated CP8 baseline remains the one-session latency +choice at 29.50 FPS. + +## Wall time and memory + +![Optimized LingBot wall-time and HBM breakdown](../disaggregated_inference_optimized.svg) + +| Component | Median | P90 | Interpretation | +| --- | ---: | ---: | --- | +| Encoder compute, seven inputs | 5.91 ms | 6.08 ms | Sequential on co-located I/O GPU | +| DiT denoise, worker sample | 1718.78 ms | 1816.21 ms | Seven workers execute concurrently | +| DiT cache finalization, worker sample | 426.45 ms | 436.32 ms | Overlapped with clean-latent RDMA | +| DiT critical path | 2178.39 ms | 2256.77 ms | Slowest worker per wave | +| Decoder compute, seven outputs | 49.06 ms | 49.11 ms | Sequential on co-located I/O GPU | +| End-to-end seven-session wave | 2358.51 ms | 2454.10 ms | Barrier-to-barrier service time | + +The encoder and decoder together use only 20.59 GiB, so co-location releases +the eighth GPU for another DiT while leaving substantial HBM headroom. Each +DiT still uses about 56.3 GiB because its weights and autoregressive cache stay +resident and are never transferred. + +The asynchronous JSON fields named `transfer_ms` and `handoff_ms` cover the +whole interval from submission until the deliberately delayed wait. They are +**in-flight residency windows**, not isolated copy latency, and are +non-additive because useful compute happens inside those windows. In +particular, the 0.55 MiB DiT-to-decoder wait occurs after roughly 426 ms of +cache finalization. Use the reusable 256 MiB probe—not payload divided by that +residency window—to characterize the link. + +## What changed + +| Optimization | Implementation | Validation | +| --- | --- | --- | +| Co-locate encoder and decoder | `--co-locate-io`; both stage weights/caches live on rank 0 | Full seven-session rollout | +| Seven independent DiTs | ranks 1–7, one resident cache per session | Full rollout; 35.15 aggregate FPS | +| Async Mooncake | non-blocking batch writes plus explicit wait handles; CUDA event replaces device-wide synchronization | Full rollout; clean completion | +| Receiver/ticket pooling | fixed-shape `RegisteredTensorPool` buckets and stable remote tickets | CPU reuse test and full rollout | +| Finalization overlap | submit clean latent before `finalize`, wait before decoder | Full rollout | +| Session-aware routing | sticky placement by pool, queue prediction, shape/CP compatibility, HBM, rack/NIC locality, and verified RDMA | CPU policy tests | +| Direct CP input shards | patchify once on encoder and transfer each rank's token shard directly | Full CP6 model measurements completed, but handoff grew to 174.23 ms and FPS fell to 12.70; experimental only | +| NIXL transport | interchangeable `NixlTensorTransport` behind the same descriptor/ticket/handle contract | Fake-agent CPU round trip; NIXL was absent from the allocated image | +| DiT microbatch admission | groups independent sessions only when worker, shape, and CP size match | CPU scheduler test; fused model/cache execution remains deferred | +| Separate service pools | `aggregated-cp8` latency pool and `io-plus-7-dit` throughput pool | CPU scheduler test and measured topology comparison | + +The synchronous control completed at 31.57 FPS but emitted repeated Mooncake +`remote access error`, `local access violation`, rail-pause, and rail-recovery +messages while short-lived registrations were recycled. The pooled run emitted +none of those errors. Pooling is therefore a buffer-lifetime correctness fix as +well as a performance optimization; the synchronous control is retained only +as diagnostic evidence, not as a production-safe configuration. + +## Reproduction + +Allocate or reuse one eight-GPU node: + +```bash +squeue -u "$USER" +cd /home/gtong/work +export FLASHDREAMS_HOST_DIR=/home/gtong/lustre/flashdreams-dist +./srun.sh +# Reattach instead of allocating again: +./srun.sh 1 +``` + +Inside the node, verify the exact mounted checkout and fabric: + +```bash +cd /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist +git rev-parse --show-toplevel +git rev-parse HEAD +nvidia-smi -L +nvidia-smi topo -m +ibv_devices +``` + +The experiment used Slurm job `14761875`, node `pool0-01924`, eight NVIDIA +H100 80 GB HBM3 GPUs, driver 535.216.03, PyTorch 2.12.1+cu130, CUDA 13.0, +Mooncake 0.3.12.post1, and base revision +`b762d079245681e1db70f1ffc5728753ce2a90b8` plus this worktree change. +The container needed RDMA userspace libraries: + +```bash +apt-get update +apt-get install -y libibverbs1 ibverbs-providers rdma-core +``` + +Run focused CPU validation: + +```bash +uv run --no-sync --package flashdreams-lingbot pytest -q \ + flashdreams/tests/test_transfer.py \ + integrations/lingbot/tests/test_disagg_stages.py \ + integrations/lingbot/tests/test_disagg_scheduler.py +``` + +Probe every stage edge before model loading: + +```bash +env TORCHINDUCTOR_COMPILE_THREADS=1 \ +uv run --no-sync --package flashdreams-lingbot torchrun \ + --standalone --nproc_per_node=8 \ + -m lingbot.disagg.benchmark_replicated \ + --dit-replicas 7 \ + --co-locate-io \ + --pooled-async \ + --transport mooncake \ + --transport-only \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 +``` + +Check the logs for Mooncake's RDMA transport installation and RDMA-ready +handshakes. Reject TCP fallback. Then run the model: + +```bash +env GLOG_minloglevel=2 TORCHINDUCTOR_COMPILE_THREADS=1 \ +uv run --no-sync --package flashdreams-lingbot torchrun \ + --standalone --nproc_per_node=8 \ + -m lingbot.disagg.benchmark_replicated \ + --dit-replicas 7 \ + --co-locate-io \ + --pooled-async \ + --transport mooncake \ + --warmup-blocks 6 \ + --measured-blocks 5 \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 \ + --output-dir outputs/lingbot_disagg_1io7dit_optimized +``` + +Inspect the raw measurements: + +```bash +jq '.environment, .summary' \ + outputs/lingbot_disagg_1io7dit_optimized/benchmark.json +jq '.records[] | select(.warmup == false)' \ + outputs/lingbot_disagg_1io7dit_optimized/benchmark.json +``` + +Keep six warmup waves: block 5 changes the cache shape and can otherwise put +compilation in the measured set. The measured output must contain five waves, +seven DiT records and 84 decoded frames per wave. + +To exercise the NIXL adapter after installing a compatible NIXL release, use +`--transport nixl` first with `--transport-only`, confirm `Backend UCX was +instantiated`, and only then load the model. This allocation did not contain +NIXL, so no real NIXL bandwidth or model result is claimed. + +## Deployment decision + +Use the disaggregated 1IO:7DiT pool when traffic contains enough simultaneous +interactive sessions to keep the seven DiTs busy, independent stage scaling or +fault isolation matters, and the stage path has verified RDMA. Use aggregated +CP8 when one session needs minimum latency, concurrency is low, or a simple +single-pod deployment is more valuable than stage elasticity. Maintain both +fixed pools for mixed SLAs; model loading, compilation, and cache warmup make +hot-repartitioning an eight-GPU node too expensive. diff --git a/integrations/lingbot/docs/benchmark_h100_1io7dit_optimized/summary.json b/integrations/lingbot/docs/benchmark_h100_1io7dit_optimized/summary.json new file mode 100644 index 00000000..c78bfc24 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_1io7dit_optimized/summary.json @@ -0,0 +1,69 @@ +{ + "date": "2026-07-30", + "slurm_job_id": "14761875", + "hostname": "pool0-01924", + "base_commit": "b762d079245681e1db70f1ffc5728753ce2a90b8", + "topology": { + "io_ranks": [0], + "dit_ranks": [1, 2, 3, 4, 5, 6, 7], + "sessions_per_wave": 7, + "transport": "mooncake-rdma", + "pooled_async": true + }, + "workload": { + "model": "lingbot-world-fast-taehv-window15-sink3", + "resolution": [464, 832], + "warmup_waves": 6, + "measured_waves": 5, + "frames_per_wave": 84 + }, + "performance": { + "aggregate_fps": 35.1528845526397, + "per_session_fps": 5.0218406503771, + "wave_latency_ms": { + "median": 2358.5133550077444, + "p90": 2454.0991360030603 + }, + "encoder_wave_ms": { + "median": 5.907039999961853, + "p90": 6.0826752305030825 + }, + "dit_critical_path_ms": { + "median": 2178.39404296875, + "p90": 2256.7726928710936 + }, + "dit_denoise_worker_ms": { + "median": 1718.7764892578125, + "p90": 1816.21455078125 + }, + "dit_finalize_worker_ms": { + "median": 426.4505310058594, + "p90": 436.3228820800781 + }, + "decoder_wave_ms": { + "median": 49.05964708328247, + "p90": 49.10519666671753 + }, + "rdma_probe_gbps": { + "median": 41.968, + "p90": 42.662, + "min": 39.664, + "max": 42.788 + } + }, + "peak_allocated_gib_by_rank": [ + 20.589931964874268, + 56.33582067489624, + 56.33582067489624, + 56.5107626914978, + 56.28736877441406, + 56.335811138153076, + 56.335811138153076, + 56.28736877441406 + ], + "comparisons": { + "tracked_1e6d1d_fps": 27.2, + "same_topology_synchronous_fps": 31.569263045981906, + "aggregated_cp8_fps_at_832x448": 29.5 + } +} diff --git a/integrations/lingbot/docs/benchmark_h100_3stage/README.md b/integrations/lingbot/docs/benchmark_h100_3stage/README.md new file mode 100644 index 00000000..5c2cec9a --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_3stage/README.md @@ -0,0 +1,52 @@ +# LingBot three-stage disaggregation benchmark + +For the full tested configuration, methodology, findings, Slurm setup, and +limitations, see the +[experiment report](../disaggregated_inference_experiment.md). + +## Result + +| Metric | Median | P90 | +| --- | ---: | ---: | +| End-to-end chunk latency | 2233.57 ms | 2250.79 ms | +| Encoder compute | 1.08 ms | 1.14 ms | +| DiT denoise | 1734.84 ms | 1755.00 ms | +| DiT cache finalize | 444.79 ms | 446.72 ms | +| Decoder compute | 7.14 ms | 7.15 ms | +| Encoder → DiT handoff | 25.38 ms | 25.88 ms | +| DiT → decoder handoff | 12.05 ms | 14.73 ms | +| Encoder → DiT payload bandwidth | 11.12 GB/s | 11.45 GB/s | +| DiT → decoder payload bandwidth | 0.39 GB/s | 0.52 GB/s | +| 256 MiB encoder → DiT probe | 41.35 GB/s | 41.48 GB/s | +| 256 MiB DiT → decoder probe | 41.00 GB/s | 41.31 GB/s | + +Steady-state throughput: **5.36 generated FPS**. + +The headline excludes 6 warmup block(s). Mooncake was +configured with the RDMA protocol. Effective payload bandwidth includes the +synchronous transfer call but excludes receiver allocation and control-plane +ticket exchange; handoff timing in `benchmark.json` includes those costs. +The real payloads were 14.36 MiB +(encoder → DiT) and 0.55 MiB +(DiT → decoder). The two synchronous copy calls account for +0.13% of median +chunk latency; complete allocation, metadata, synchronization, and copy +handoffs account for 1.68%. + +## Reproduction + +```bash +uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=3 \ + -m lingbot.disagg.benchmark \ + --model lingbot-world-fast-taehv-window15-sink3 \ + --example-idx 0 --pixel-width 832 --pixel-height 464 --fps 16 \ + --warmup-blocks 6 --measured-blocks 5 \ + --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 \ + --output-dir integrations/lingbot/docs/benchmark_h100_3stage +``` + +- Repository base: `e580e27d408b3cf8bd8a549f990c361b94d3379f`; the implementation was the worktree change recorded with this report. +- Slurm: job `14621292` on `pool0-01299` +- GPU: `NVIDIA H100 80GB HBM3` × 3 +- Resolution: `832x464` +- Model: `lingbot-world-fast-taehv-window15-sink3` diff --git a/integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json b/integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json new file mode 100644 index 00000000..3a15949d --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json @@ -0,0 +1,548 @@ +{ + "environment": { + "command": "uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=3 -m lingbot.disagg.benchmark --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --pixel-width 832 --pixel-height 464 --fps 16 --warmup-blocks 6 --measured-blocks 5 --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 --output-dir integrations/lingbot/docs/benchmark_h100_3stage", + "commit": "e580e27d408b3cf8bd8a549f990c361b94d3379f", + "worktree_dirty": true, + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "unknown", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "peak_memory_gib_by_stage": { + "encoder": 13.719394207000732, + "dit": 56.28736877441406, + "decoder": 2.2908687591552734 + }, + "hostname": "pool0-01299", + "slurm_job_id": "14621292" + }, + "summary": { + "fps": 5.360631332509123, + "latency_ms": { + "median": 2233.5668401792645, + "p90": 2250.7897848263383, + "min": 2231.1054738238454, + "max": 2262.059194035828 + }, + "encoder_ms": { + "median": 1.0813119411468506, + "p90": 1.1356287956237794, + "min": 0.8738560080528259, + "max": 1.15065598487854 + }, + "dit_ms": { + "median": 1734.8382568359375, + "p90": 1754.9966796875, + "min": 1731.94189453125, + "max": 1765.748291015625 + }, + "finalize_ms": { + "median": 444.786865234375, + "p90": 446.72052001953125, + "min": 444.50189208984375, + "max": 447.21051025390625 + }, + "decoder_ms": { + "median": 7.14086389541626, + "p90": 7.14958086013794, + "min": 7.117055892944336, + "max": 7.154304027557373 + }, + "encoder_to_dit": { + "payload_mib": 14.3583984375, + "transfer_ms": { + "median": 1.354343257844448, + "p90": 1.601235382258892, + "min": 1.3001281768083572, + "max": 1.6657030209898949 + }, + "bandwidth_gbps": { + "median": 11.116732713656871, + "p90": 11.446511671164952, + "min": 9.038749291006622, + "max": 11.580298211027296 + }, + "handoff_ms": { + "median": 25.382738560438156, + "p90": 25.876259058713913, + "min": 25.090406648814678, + "max": 26.015397161245346 + } + }, + "dit_to_decoder": { + "payload_mib": 0.55224609375, + "transfer_ms": { + "median": 1.489844173192978, + "p90": 2.0092779770493507, + "min": 1.112041063606739, + "max": 2.2190073505043983 + }, + "bandwidth_gbps": { + "median": 0.38867957496451105, + "p90": 0.5182747045250334, + "min": 0.26095992871243634, + "max": 0.5207289721135535 + }, + "handoff_ms": { + "median": 12.049876153469086, + "p90": 14.728017896413803, + "min": 11.790870688855648, + "max": 16.113361343741417 + } + }, + "bandwidth_probe_gbps": { + "encoder_to_dit": { + "median": 41.34810593787971, + "p90": 41.48398153764155, + "min": 39.89549867186514, + "max": 41.533244658448375 + }, + "dit_to_decoder": { + "median": 40.99657050342333, + "p90": 41.30850823617201, + "min": 32.205209381266116, + "max": 41.40562841024771 + } + }, + "transfer_overhead_percent": { + "synchronous_copy": 0.12733836211542043, + "full_handoff": 1.675912000506908 + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "encoder_ms": 330.6899108886719, + "encoder_to_dit_handoff_ms": 27.6114484295249, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 9.935536421835423, + "transfer_ms": 1.3665007427334785, + "bandwidth_gbps": 11.017829357255232 + }, + "dit_to_decoder_handoff_ms": 18.7167814001441, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 5.7073924690485, + "transfer_ms": 1.2289220467209816, + "bandwidth_gbps": 0.47120319921437165 + }, + "end_to_end_ms": 14204.82233632356, + "dit_ms": 10014.931640625, + "finalize_ms": 278.7642822265625, + "decoder_ms": 3527.971435546875, + "output_frames": 9 + }, + { + "autoregressive_index": 1, + "warmup": true, + "encoder_ms": 179.03184509277344, + "encoder_to_dit_handoff_ms": 54.61540725082159, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.05543302744627, + "transfer_ms": 1.5367651358246803, + "bandwidth_gbps": 9.797119708809966 + }, + "dit_to_decoder_handoff_ms": 11.787611059844494, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.584078513085842, + "transfer_ms": 1.3692956417798996, + "bandwidth_gbps": 0.42289771641081436 + }, + "end_to_end_ms": 13454.238000325859, + "dit_ms": 12883.484375, + "finalize_ms": 312.16180419921875, + "decoder_ms": 7.982111930847168, + "output_frames": 12 + }, + { + "autoregressive_index": 2, + "warmup": true, + "encoder_ms": 162.63037109375, + "encoder_to_dit_handoff_ms": 25.766183622181416, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.137472301721573, + "transfer_ms": 1.392514444887638, + "bandwidth_gbps": 10.81200418083624 + }, + "dit_to_decoder_handoff_ms": 10.47912985086441, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.530191257596016, + "transfer_ms": 0.9602615609765053, + "bandwidth_gbps": 0.6030356972855734 + }, + "end_to_end_ms": 1935.4530731216073, + "dit_ms": 1388.372314453125, + "finalize_ms": 335.4577941894531, + "decoder_ms": 7.823647975921631, + "output_frames": 12 + }, + { + "autoregressive_index": 3, + "warmup": true, + "encoder_ms": 162.38783264160156, + "encoder_to_dit_handoff_ms": 26.856454089283943, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.054773651063442, + "transfer_ms": 1.2486670166254044, + "bandwidth_gbps": 12.057555616940515 + }, + "dit_to_decoder_handoff_ms": 10.370046831667423, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.5941191017627716, + "transfer_ms": 1.0182484984397888, + "bandwidth_gbps": 0.568694185051373 + }, + "end_to_end_ms": 2104.299475438893, + "dit_ms": 1507.4981689453125, + "finalize_ms": 370.435791015625, + "decoder_ms": 22.703296661376953, + "output_frames": 12 + }, + { + "autoregressive_index": 4, + "warmup": true, + "encoder_ms": 162.43638610839844, + "encoder_to_dit_handoff_ms": 25.723207741975784, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.20294614136219, + "transfer_ms": 1.5074005350470543, + "bandwidth_gbps": 9.987970449758413 + }, + "dit_to_decoder_handoff_ms": 10.693598538637161, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.5667959600687027, + "transfer_ms": 1.0530250146985054, + "bandwidth_gbps": 0.5499128623889299 + }, + "end_to_end_ms": 2248.7359000369906, + "dit_ms": 1628.2943115234375, + "finalize_ms": 410.56829833984375, + "decoder_ms": 7.130080223083496, + "output_frames": 12 + }, + { + "autoregressive_index": 5, + "warmup": true, + "encoder_ms": 0.9327359795570374, + "encoder_to_dit_handoff_ms": 28.35660893470049, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.2960414737463, + "transfer_ms": 1.2558028101921082, + "bandwidth_gbps": 11.989041494258807 + }, + "dit_to_decoder_handoff_ms": 12.713887728750706, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.6875084042549133, + "transfer_ms": 1.5450455248355865, + "bandwidth_gbps": 0.37479283988193224 + }, + "end_to_end_ms": 13046.747697517276, + "dit_ms": 12555.27734375, + "finalize_ms": 425.007080078125, + "decoder_ms": 7.144032001495361, + "output_frames": 12 + }, + { + "autoregressive_index": 6, + "warmup": false, + "encoder_ms": 1.15065598487854, + "encoder_to_dit_handoff_ms": 26.015397161245346, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.213454253971577, + "transfer_ms": 1.6657030209898949, + "bandwidth_gbps": 9.038749291006622 + }, + "dit_to_decoder_handoff_ms": 12.049876153469086, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.6088526248931885, + "transfer_ms": 1.489844173192978, + "bandwidth_gbps": 0.38867957496451105 + }, + "end_to_end_ms": 2262.059194035828, + "dit_ms": 1765.748291015625, + "finalize_ms": 444.786865234375, + "decoder_ms": 7.14086389541626, + "output_frames": 12 + }, + { + "autoregressive_index": 7, + "warmup": false, + "encoder_ms": 1.0813119411468506, + "encoder_to_dit_handoff_ms": 25.090406648814678, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.084053501486778, + "transfer_ms": 1.3387957587838173, + "bandwidth_gbps": 11.245831861371435 + }, + "dit_to_decoder_handoff_ms": 11.790870688855648, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.6583244800567627, + "transfer_ms": 1.1253003031015396, + "bandwidth_gbps": 0.5145933031422533 + }, + "end_to_end_ms": 2233.5668401792645, + "dit_ms": 1738.8692626953125, + "finalize_ms": 444.51544189453125, + "decoder_ms": 7.154304027557373, + "output_frames": 12 + }, + { + "autoregressive_index": 8, + "warmup": false, + "encoder_ms": 1.1130880117416382, + "encoder_to_dit_handoff_ms": 25.382738560438156, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.165871120989323, + "transfer_ms": 1.354343257844448, + "bandwidth_gbps": 11.116732713656871 + }, + "dit_to_decoder_handoff_ms": 16.113361343741417, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.647986799478531, + "transfer_ms": 1.112041063606739, + "bandwidth_gbps": 0.5207289721135535 + }, + "end_to_end_ms": 2233.8856710121036, + "dit_ms": 1731.94189453125, + "finalize_ms": 444.50189208984375, + "decoder_ms": 7.142496109008789, + "output_frames": 12 + }, + { + "autoregressive_index": 9, + "warmup": false, + "encoder_ms": 0.9545599818229675, + "encoder_to_dit_handoff_ms": 25.667551904916763, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.174481198191643, + "transfer_ms": 1.5045339241623878, + "bandwidth_gbps": 10.00700067855365 + }, + "dit_to_decoder_handoff_ms": 12.650002725422382, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.7020957097411156, + "transfer_ms": 2.2190073505043983, + "bandwidth_gbps": 0.26095992871243634 + }, + "end_to_end_ms": 2232.094327919185, + "dit_ms": 1734.0849609375, + "finalize_ms": 447.21051025390625, + "decoder_ms": 7.117055892944336, + "output_frames": 12 + }, + { + "autoregressive_index": 10, + "warmup": false, + "encoder_ms": 0.8738560080528259, + "encoder_to_dit_handoff_ms": 25.33851470798254, + "encoder_to_dit": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.107764042913914, + "transfer_ms": 1.3001281768083572, + "bandwidth_gbps": 11.580298211027296 + }, + "dit_to_decoder_handoff_ms": 12.027109041810036, + "dit_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 3.597964532673359, + "transfer_ms": 1.6946839168667793, + "bandwidth_gbps": 0.34169911818754894 + }, + "end_to_end_ms": 2231.1054738238454, + "dit_ms": 1734.8382568359375, + "finalize_ms": 445.98553466796875, + "decoder_ms": 7.117824077606201, + "output_frames": 12 + } + ], + "bandwidth_probe": { + "encoder_to_dit": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.004693865776062012, + "transfer_ms": 6.540342234075069, + "bandwidth_gbps": 41.043029002588874 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0023655593395233154, + "transfer_ms": 6.474116817116737, + "bandwidth_gbps": 41.462868771581476 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.003981404006481171, + "transfer_ms": 6.489843130111694, + "bandwidth_gbps": 41.36239514858351 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016130506992340088, + "transfer_ms": 6.498831324279308, + "bandwidth_gbps": 41.30518898023074 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015292316675186157, + "transfer_ms": 6.494330242276192, + "bandwidth_gbps": 41.3338167271759 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001475214958190918, + "transfer_ms": 6.47522509098053, + "bandwidth_gbps": 41.455772151289246 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013979151844978333, + "transfer_ms": 6.728464737534523, + "bandwidth_gbps": 39.89549867186514 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015543773770332336, + "transfer_ms": 6.463146768510342, + "bandwidth_gbps": 41.533244658448375 + } + ], + "dit_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0029364600777626038, + "transfer_ms": 6.546538323163986, + "bandwidth_gbps": 41.00418308866835 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0018831342458724976, + "transfer_ms": 6.544428877532482, + "bandwidth_gbps": 41.01739984088744 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0027604401111602783, + "transfer_ms": 6.504863500595093, + "bandwidth_gbps": 41.26688530442528 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016242265701293945, + "transfer_ms": 7.034887559711933, + "bandwidth_gbps": 38.15774647732848 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014677643775939941, + "transfer_ms": 8.335156366229057, + "bandwidth_gbps": 32.205209381266116 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014211982488632202, + "transfer_ms": 6.548970006406307, + "bandwidth_gbps": 40.9889579181783 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014873221516609192, + "transfer_ms": 6.483066827058792, + "bandwidth_gbps": 41.40562841024771 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013010576367378235, + "transfer_ms": 6.984638050198555, + "bandwidth_gbps": 38.43226435940644 + } + ] + } +} diff --git a/integrations/lingbot/docs/benchmark_h100_aggregated_8xcp1/README.md b/integrations/lingbot/docs/benchmark_h100_aggregated_8xcp1/README.md new file mode 100644 index 00000000..f49700b9 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_aggregated_8xcp1/README.md @@ -0,0 +1,77 @@ +# LingBot eight independent aggregated workers + +Eight H100s each ran one complete CP1 encoder + DiT + LightTAE decoder pipeline +and one independent session. All workers completed six warmup chunks, waited at +a shared barrier, and then ran five measured 12-frame chunks concurrently. + +## Result + +| Metric | Value | +| --- | ---: | +| Aggregate generated FPS | **43.44** | +| Per-session FPS, median / p90 | **5.54 / 5.59** | +| Chunk latency, median / p90 | **2163.64 / 2205.53 ms** | +| Shared measurement wall time | 11.050 s | +| Measurement start skew | 0.32 ms | +| Measured sessions / chunks / frames | 8 / 40 / 480 | +| Rollout peak allocated HBM per GPU | **59.35 GiB** | +| Initialization peak allocated HBM per GPU | **66.55 GiB** | +| Steady allocated HBM per GPU | **57.15 GiB** | +| Rollout peak allocated HBM, node total | **474.84 GiB** | +| Initialization peak allocated HBM, node total | **532.38 GiB** | +| Steady allocated HBM, node total | **457.16 GiB** | + +The measured aggregate is 97.7% of eight times the tracked 5.56 FPS +single-H100 result. The small gap includes the 0.32 ms start skew and 329.85 ms +finish skew across workers; summing the independently measured worker rates +gives 44.34 FPS, while the stricter shared-window calculation gives the 43.44 +FPS headline. + +## Serving comparison + +| Eight-GPU topology | Sessions | Aggregate FPS | FPS/session | Median latency | Rollout peak node HBM | +| --- | ---: | ---: | ---: | ---: | ---: | +| 1 I/O + 7 independent DiTs, pooled async | 7 | 35.15 | 5.02 | 2358.51 ms wave | 415.02 GiB | +| **8 independent full pipelines** | **8** | **43.44** | **5.54 median** | **2163.64 ms** | **474.84 GiB** | +| One aggregated CP8 pipeline | 1 | 29.50 | 29.50 | 393.33 ms | 327.03 GiB | + +Eight full replicas deliver 23.6% more aggregate FPS than 1 I/O + 7 DiTs and +10.4% more median FPS per session, while using 14.4% more rollout peak node +HBM. CP8 remains the single-session latency topology; the independent replicas +are the highest-throughput measured topology when eight full pipelines fit. + +## Reproduction + +From an eight-GPU Slurm node: + +```bash +GLOG_minloglevel=2 \ +uv run --package flashdreams-lingbot python \ + -m lingbot.disagg.benchmark_independent \ + --replicas 8 \ + --model lingbot-world-fast-taehv-window15-sink3 \ + --example-idx 0 \ + --pixel-width 832 --pixel-height 464 --fps 16 \ + --warmup-blocks 6 --measured-blocks 5 \ + --compile-threads-per-replica 4 \ + --timeout-s 3600 \ + --output-dir outputs/lingbot_aggregated_8xcp1 +``` + +The coordinator assigns one visible GPU to each subprocess. Every subprocess +runs ``benchmark_aggregated`` with one ``torchrun`` rank, creates a readiness +file after warmup, and waits for the common release file. Aggregate FPS is 480 +frames divided by the wall time from the earliest worker start to the latest +worker finish. + +Environment: + +- Repository base: `0c2d48a8249577fb617bb5280208dd77409d9b1a`, plus the benchmark worktree changes +- Slurm: job `14793417`, node `pool0-01151` +- GPU: 8 × NVIDIA H100 80 GB HBM3 +- Resolution: 832×464; BF16; seed 42; four diffusion steps; window 15; sink 3 +- PyTorch 2.12.1+cu130, CUDA 13.0, cuDNN 9.2, driver 535.216.03 + +The complete per-worker environment, timing records, stage breakdown, and +memory arrays are in [`benchmark.json`](benchmark.json). Worker logs remain in +the untracked output directory. diff --git a/integrations/lingbot/docs/benchmark_h100_aggregated_8xcp1/benchmark.json b/integrations/lingbot/docs/benchmark_h100_aggregated_8xcp1/benchmark.json new file mode 100644 index 00000000..daef042c --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_aggregated_8xcp1/benchmark.json @@ -0,0 +1,3846 @@ +{ + "environment": { + "command": "uv run --package flashdreams-lingbot python -m lingbot.disagg.benchmark_independent --replicas 8 --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --pixel-width 832 --pixel-height 464 --fps 16 --warmup-blocks 6 --measured-blocks 5 --compile-threads-per-replica 4 --timeout-s 3600 --output-dir outputs/lingbot_aggregated_8xcp1", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01151", + "slurm_job_id": "14793417", + "model": "lingbot-world-fast-taehv-window15-sink3", + "resolution": [ + 464, + 832 + ], + "warmup_blocks": 6, + "measured_blocks": 5, + "replicas": 8, + "gpus": [ + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3" + ], + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "precision": "bfloat16", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "worker_commands": [ + "/lustre/fsw/portfolios/healthcareeng/users/gtong/venvs/flashdreams/bin/python3 -m torch.distributed.run --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 0 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-0", + "/lustre/fsw/portfolios/healthcareeng/users/gtong/venvs/flashdreams/bin/python3 -m torch.distributed.run --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 1 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-1", + "/lustre/fsw/portfolios/healthcareeng/users/gtong/venvs/flashdreams/bin/python3 -m torch.distributed.run --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 2 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-2", + "/lustre/fsw/portfolios/healthcareeng/users/gtong/venvs/flashdreams/bin/python3 -m torch.distributed.run --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 3 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-3", + "/lustre/fsw/portfolios/healthcareeng/users/gtong/venvs/flashdreams/bin/python3 -m torch.distributed.run --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 4 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-4", + "/lustre/fsw/portfolios/healthcareeng/users/gtong/venvs/flashdreams/bin/python3 -m torch.distributed.run --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 5 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-5", + "/lustre/fsw/portfolios/healthcareeng/users/gtong/venvs/flashdreams/bin/python3 -m torch.distributed.run --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 6 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-6", + "/lustre/fsw/portfolios/healthcareeng/users/gtong/venvs/flashdreams/bin/python3 -m torch.distributed.run --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 7 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-7" + ] + }, + "summary": { + "aggregate_fps": 43.43749180420828, + "sum_of_worker_fps": 44.344999148883396, + "total_output_frames": 480, + "measurement_wall_s": 11.050361797213554, + "measurement_start_skew_ms": 0.3231372684240341, + "measurement_finish_skew_ms": 329.84739169478416, + "per_session_fps": { + "median": 5.544989798538596, + "p90": 5.59244924622212, + "min": 5.431233201432407, + "max": 5.5984819295074555 + }, + "per_session_median_latency_ms": { + "median": 2163.28211594373, + "p90": 2183.3966828882694, + "min": 2141.4667814970016, + "max": 2209.239514544606 + }, + "all_chunk_latency_ms": { + "median": 2163.643025793135, + "p90": 2205.5334428325295, + "min": 2137.736974284053, + "max": 2213.644528761506 + }, + "memory": { + "rollout_peak_gib_by_gpu": [ + 59.35440921783447, + 59.35440921783447, + 59.35440921783447, + 59.35440921783447, + 59.35440921783447, + 59.35440921783447, + 59.35440921783447, + 59.35440921783447 + ], + "rollout_peak_gib_per_gpu": { + "median": 59.35440921783447, + "p90": 59.35440921783447, + "min": 59.35440921783447, + "max": 59.35440921783447 + }, + "rollout_peak_gib_node_total": 474.8352737426758, + "initialization_peak_gib_by_gpu": [ + 66.54739189147949, + 66.54739189147949, + 66.54739189147949, + 66.54739189147949, + 66.54739189147949, + 66.54739189147949, + 66.54739189147949, + 66.54739189147949 + ], + "initialization_peak_gib_per_gpu": { + "median": 66.54739189147949, + "p90": 66.54739189147949, + "min": 66.54739189147949, + "max": 66.54739189147949 + }, + "initialization_peak_gib_node_total": 532.3791351318359, + "steady_allocated_gib_by_gpu": [ + 57.14538335800171, + 57.14538335800171, + 57.14538335800171, + 57.14538335800171, + 57.14538335800171, + 57.14538335800171, + 57.14538335800171, + 57.14538335800171 + ], + "steady_allocated_gib_node_total": 457.1630668640137 + } + }, + "workers": [ + { + "environment": { + "command": "env TORCHINDUCTOR_COMPILE_THREADS=4 uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 0 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-0", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01151", + "slurm_job_id": "14793417", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0 + ], + "dit_cp_group": [ + 0 + ], + "cp_size": 1, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42 + ], + "replica_id": 0, + "measurement_window": { + "started_at": 11145874.429103078, + "finished_at": 11145885.29418634, + "elapsed_s": 10.865083262324333 + } + }, + "summary": { + "fps": 5.523946310994633, + "latency_ms": { + "median": 2171.0722651332617, + "p90": 2176.1993937194347, + "min": 2168.278167024255, + "max": 2177.5269210338593 + }, + "encoder_ms": { + "median": 0.8008639812469482, + "p90": 1.000268816947937, + "min": 0.7934079766273499, + "max": 1.0645760297775269 + }, + "dit_ms": { + "median": 1736.72607421875, + "p90": 1738.9442626953125, + "min": 1726.501708984375, + "max": 1740.3426513671875 + }, + "decoder_ms": { + "median": 7.425727844238281, + "p90": 7.530271911621094, + "min": 7.380095958709717, + "max": 7.582431793212891 + }, + "finalize_ms": { + "median": 428.4263916015625, + "p90": 433.63258056640626, + "min": 425.04583740234375, + "max": 437.065673828125 + }, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 2082.5277592449766, + "cp_probe_gbps": { + "broadcast": { + "median": 20.301848368903027, + "p90": 20.301848368903027, + "min": 20.301848368903027, + "max": 20.301848368903027 + }, + "all_gather": { + "median": 34.04184274181, + "p90": 34.04184274181, + "min": 34.04184274181, + "max": 34.04184274181 + } + }, + "memory": { + "peak_gib_by_rank": [ + 59.35440921783447 + ], + "steady_allocated_gib_by_rank": [ + 57.14538335800171 + ], + "initialization_peak_gib_by_rank": [ + 66.54739189147949 + ], + "node_peak_gib": 59.35440921783447, + "node_steady_allocated_gib": 57.14538335800171, + "per_rank_peak_gib": { + "median": 59.35440921783447, + "p90": 59.35440921783447, + "min": 59.35440921783447, + "max": 59.35440921783447 + } + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 138200.74504986405, + "output_frames": 9, + "critical_rank": { + "encode_ms": 590.80517578125, + "diffuse_ms": 92299.890625, + "decode_ms": 45043.6328125, + "finalize_ms": 265.7388916015625 + }, + "per_rank": [ + { + "encode_ms": 590.80517578125, + "diffuse_ms": 92299.890625, + "decode_ms": 45043.6328125, + "finalize_ms": 265.7388916015625, + "total_ms": 138200.0675048828, + "total_ms_wo_finalize": 137934.32861328125, + "mem_alloc_gib": 57.087345123291016, + "mem_reserved_gib": 61.439453125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 138200.74504986405, + "output_frames": 9, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 80273.97294715047, + "output_frames": 12, + "critical_rank": { + "encode_ms": 181.9831085205078, + "diffuse_ms": 79769.765625, + "decode_ms": 8.339455604553223, + "finalize_ms": 312.95416259765625 + }, + "per_rank": [ + { + "encode_ms": 181.9831085205078, + "diffuse_ms": 79769.765625, + "decode_ms": 8.339455604553223, + "finalize_ms": 312.95416259765625, + "total_ms": 80273.04235172272, + "total_ms_wo_finalize": 79960.08818912506, + "mem_alloc_gib": 57.08803939819336, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 80273.97294715047, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 1890.9110315144062, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.99746704101562, + "diffuse_ms": 1380.8189697265625, + "decode_ms": 7.524799823760986, + "finalize_ms": 337.9274597167969 + }, + "per_rank": [ + { + "encode_ms": 163.99746704101562, + "diffuse_ms": 1380.8189697265625, + "decode_ms": 7.524799823760986, + "finalize_ms": 337.9274597167969, + "total_ms": 1890.268696308136, + "total_ms_wo_finalize": 1552.341236591339, + "mem_alloc_gib": 57.08770561218262, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 1890.9110315144062, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 2051.9016850739717, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.3020477294922, + "diffuse_ms": 1499.044921875, + "decode_ms": 19.27168083190918, + "finalize_ms": 368.6894226074219 + }, + "per_rank": [ + { + "encode_ms": 164.3020477294922, + "diffuse_ms": 1499.044921875, + "decode_ms": 19.27168083190918, + "finalize_ms": 368.6894226074219, + "total_ms": 2051.3080730438232, + "total_ms_wo_finalize": 1682.6186504364014, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.423828125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2051.9016850739717, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 2200.6694059818983, + "output_frames": 12, + "critical_rank": { + "encode_ms": 165.09075927734375, + "diffuse_ms": 1614.4637451171875, + "decode_ms": 7.398752212524414, + "finalize_ms": 413.0950927734375 + }, + "per_rank": [ + { + "encode_ms": 165.09075927734375, + "diffuse_ms": 1614.4637451171875, + "decode_ms": 7.398752212524414, + "finalize_ms": 413.0950927734375, + "total_ms": 2200.048349380493, + "total_ms_wo_finalize": 1786.9532566070557, + "mem_alloc_gib": 57.144843101501465, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2200.6694059818983, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 2192.485436797142, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.0230079889297485, + "diffuse_ms": 1742.4036865234375, + "decode_ms": 7.463903903961182, + "finalize_ms": 440.8986511230469 + }, + "per_rank": [ + { + "encode_ms": 1.0230079889297485, + "diffuse_ms": 1742.4036865234375, + "decode_ms": 7.463903903961182, + "finalize_ms": 440.8986511230469, + "total_ms": 2191.7892495393753, + "total_ms_wo_finalize": 1750.8905984163284, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2192.485436797142, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 2174.208102747798, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.0645760297775269, + "diffuse_ms": 1726.501708984375, + "decode_ms": 7.382016181945801, + "finalize_ms": 437.065673828125 + }, + "per_rank": [ + { + "encode_ms": 1.0645760297775269, + "diffuse_ms": 1726.501708984375, + "decode_ms": 7.382016181945801, + "finalize_ms": 437.065673828125, + "total_ms": 2172.0139750242233, + "total_ms_wo_finalize": 1734.9483011960983, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2174.208102747798, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 2168.278167024255, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9038079977035522, + "diffuse_ms": 1730.7530517578125, + "decode_ms": 7.582431793212891, + "finalize_ms": 428.4829406738281 + }, + "per_rank": [ + { + "encode_ms": 0.9038079977035522, + "diffuse_ms": 1730.7530517578125, + "decode_ms": 7.582431793212891, + "finalize_ms": 428.4829406738281, + "total_ms": 2167.722232222557, + "total_ms_wo_finalize": 1739.239291548729, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2168.278167024255, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 2170.7145366817713, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7939199805259705, + "diffuse_ms": 1736.8466796875, + "decode_ms": 7.452032089233398, + "finalize_ms": 425.04583740234375 + }, + "per_rank": [ + { + "encode_ms": 0.7939199805259705, + "diffuse_ms": 1736.8466796875, + "decode_ms": 7.452032089233398, + "finalize_ms": 425.04583740234375, + "total_ms": 2170.138469159603, + "total_ms_wo_finalize": 1745.0926317572594, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2170.7145366817713, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 2171.0722651332617, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8008639812469482, + "diffuse_ms": 1736.72607421875, + "decode_ms": 7.425727844238281, + "finalize_ms": 425.5943908691406 + }, + "per_rank": [ + { + "encode_ms": 0.8008639812469482, + "diffuse_ms": 1736.72607421875, + "decode_ms": 7.425727844238281, + "finalize_ms": 425.5943908691406, + "total_ms": 2170.547056913376, + "total_ms_wo_finalize": 1744.9526660442352, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2171.0722651332617, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 2177.5269210338593, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7934079766273499, + "diffuse_ms": 1740.3426513671875, + "decode_ms": 7.380095958709717, + "finalize_ms": 428.4263916015625 + }, + "per_rank": [ + { + "encode_ms": 0.7934079766273499, + "diffuse_ms": 1740.3426513671875, + "decode_ms": 7.380095958709717, + "finalize_ms": 428.4263916015625, + "total_ms": 2176.942546904087, + "total_ms_wo_finalize": 1748.5161553025246, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2177.5269210338593, + "output_frames": 12, + "rank": 0 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.05164928734302521, + "bandwidth_gbps": 20.301848368903027 + } + ], + "all_gather": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.030802562832832336, + "bandwidth_gbps": 34.04184274181 + } + ] + } + }, + { + "environment": { + "command": "env TORCHINDUCTOR_COMPILE_THREADS=4 uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 1 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-1", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01151", + "slurm_job_id": "14793417", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0 + ], + "dit_cp_group": [ + 0 + ], + "cp_size": 1, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42 + ], + "replica_id": 1, + "measurement_window": { + "started_at": 11145874.429260097, + "finished_at": 11145885.479436684, + "elapsed_s": 11.050176586955786 + } + }, + "summary": { + "fps": 5.431233201432407, + "latency_ms": { + "median": 2209.239514544606, + "p90": 2213.0047027021646, + "min": 2205.3810749202967, + "max": 2213.644528761506 + }, + "encoder_ms": { + "median": 0.8197439908981323, + "p90": 0.9503488302230835, + "min": 0.7971839904785156, + "max": 1.0180480480194092 + }, + "dit_ms": { + "median": 1763.7745361328125, + "p90": 1765.95751953125, + "min": 1756.140869140625, + "max": 1766.673095703125 + }, + "decoder_ms": { + "median": 7.529024124145508, + "p90": 7.707660865783692, + "min": 7.421440124511719, + "max": 7.752384185791016 + }, + "finalize_ms": { + "median": 436.60076904296875, + "p90": 443.597265625, + "min": 431.5830078125, + "max": 447.26544189453125 + }, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 2047.5749169400176, + "cp_probe_gbps": { + "broadcast": { + "median": 15.300463496352894, + "p90": 15.300463496352894, + "min": 15.300463496352894, + "max": 15.300463496352894 + }, + "all_gather": { + "median": 28.054916446791186, + "p90": 28.054916446791186, + "min": 28.054916446791186, + "max": 28.054916446791186 + } + }, + "memory": { + "peak_gib_by_rank": [ + 59.35440921783447 + ], + "steady_allocated_gib_by_rank": [ + 57.14538335800171 + ], + "initialization_peak_gib_by_rank": [ + 66.54739189147949 + ], + "node_peak_gib": 59.35440921783447, + "node_steady_allocated_gib": 57.14538335800171, + "per_rank_peak_gib": { + "median": 59.35440921783447, + "p90": 59.35440921783447, + "min": 59.35440921783447, + "max": 59.35440921783447 + } + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 175426.7371315509, + "output_frames": 9, + "critical_rank": { + "encode_ms": 613.5792236328125, + "diffuse_ms": 99693.546875, + "decode_ms": 74844.7734375, + "finalize_ms": 272.13897705078125 + }, + "per_rank": [ + { + "encode_ms": 613.5792236328125, + "diffuse_ms": 99693.546875, + "decode_ms": 74844.7734375, + "finalize_ms": 272.13897705078125, + "total_ms": 175424.0385131836, + "total_ms_wo_finalize": 175151.8995361328, + "mem_alloc_gib": 57.087345123291016, + "mem_reserved_gib": 61.439453125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 175426.7371315509, + "output_frames": 9, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 79537.70590759814, + "output_frames": 12, + "critical_rank": { + "encode_ms": 182.7596435546875, + "diffuse_ms": 79028.9453125, + "decode_ms": 8.16988754272461, + "finalize_ms": 316.3579406738281 + }, + "per_rank": [ + { + "encode_ms": 182.7596435546875, + "diffuse_ms": 79028.9453125, + "decode_ms": 8.16988754272461, + "finalize_ms": 316.3579406738281, + "total_ms": 79536.23278427124, + "total_ms_wo_finalize": 79219.87484359741, + "mem_alloc_gib": 57.08803939819336, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 79537.70590759814, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 1914.4557267427444, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.57635498046875, + "diffuse_ms": 1402.1595458984375, + "decode_ms": 7.56006383895874, + "finalize_ms": 339.5437316894531 + }, + "per_rank": [ + { + "encode_ms": 164.57635498046875, + "diffuse_ms": 1402.1595458984375, + "decode_ms": 7.56006383895874, + "finalize_ms": 339.5437316894531, + "total_ms": 1913.8396964073181, + "total_ms_wo_finalize": 1574.295964717865, + "mem_alloc_gib": 57.08770561218262, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 1914.4557267427444, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 2084.323525428772, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.3065643310547, + "diffuse_ms": 1524.3265380859375, + "decode_ms": 19.166112899780273, + "finalize_ms": 375.96490478515625 + }, + "per_rank": [ + { + "encode_ms": 164.3065643310547, + "diffuse_ms": 1524.3265380859375, + "decode_ms": 19.166112899780273, + "finalize_ms": 375.96490478515625, + "total_ms": 2083.7641201019287, + "total_ms_wo_finalize": 1707.7992153167725, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.423828125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2084.323525428772, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 2230.1153894513845, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.94998168945312, + "diffuse_ms": 1638.8804931640625, + "decode_ms": 7.443647861480713, + "finalize_ms": 418.2716979980469 + }, + "per_rank": [ + { + "encode_ms": 164.94998168945312, + "diffuse_ms": 1638.8804931640625, + "decode_ms": 7.443647861480713, + "finalize_ms": 418.2716979980469, + "total_ms": 2229.545820713043, + "total_ms_wo_finalize": 1811.2741227149963, + "mem_alloc_gib": 57.144843101501465, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2230.1153894513845, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 2257.7255573123693, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9730240106582642, + "diffuse_ms": 1801.062255859375, + "decode_ms": 7.568543910980225, + "finalize_ms": 447.5612487792969 + }, + "per_rank": [ + { + "encode_ms": 0.9730240106582642, + "diffuse_ms": 1801.062255859375, + "decode_ms": 7.568543910980225, + "finalize_ms": 447.5612487792969, + "total_ms": 2257.1650725603104, + "total_ms_wo_finalize": 1809.6038237810135, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2257.7255573123693, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 2213.644528761506, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.0180480480194092, + "diffuse_ms": 1756.140869140625, + "decode_ms": 7.752384185791016, + "finalize_ms": 447.26544189453125 + }, + "per_rank": [ + { + "encode_ms": 1.0180480480194092, + "diffuse_ms": 1756.140869140625, + "decode_ms": 7.752384185791016, + "finalize_ms": 447.26544189453125, + "total_ms": 2212.1767432689667, + "total_ms_wo_finalize": 1764.9113013744354, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2213.644528761506, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 2206.9047540426254, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.848800003528595, + "diffuse_ms": 1763.7745361328125, + "decode_ms": 7.640575885772705, + "finalize_ms": 434.0928955078125 + }, + "per_rank": [ + { + "encode_ms": 0.848800003528595, + "diffuse_ms": 1763.7745361328125, + "decode_ms": 7.640575885772705, + "finalize_ms": 434.0928955078125, + "total_ms": 2206.3568075299263, + "total_ms_wo_finalize": 1772.2639120221138, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2206.9047540426254, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 2205.3810749202967, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8197439908981323, + "diffuse_ms": 1764.8841552734375, + "decode_ms": 7.529024124145508, + "finalize_ms": 431.5830078125 + }, + "per_rank": [ + { + "encode_ms": 0.8197439908981323, + "diffuse_ms": 1764.8841552734375, + "decode_ms": 7.529024124145508, + "finalize_ms": 431.5830078125, + "total_ms": 2204.815931200981, + "total_ms_wo_finalize": 1773.2329233884811, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2205.3810749202967, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 2212.0449636131525, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8001919984817505, + "diffuse_ms": 1766.673095703125, + "decode_ms": 7.421440124511719, + "finalize_ms": 436.60076904296875 + }, + "per_rank": [ + { + "encode_ms": 0.8001919984817505, + "diffuse_ms": 1766.673095703125, + "decode_ms": 7.421440124511719, + "finalize_ms": 436.60076904296875, + "total_ms": 2211.495496869087, + "total_ms_wo_finalize": 1774.8947278261185, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2212.0449636131525, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 2209.239514544606, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7971839904785156, + "diffuse_ms": 1762.2459716796875, + "decode_ms": 7.487936019897461, + "finalize_ms": 438.0950012207031 + }, + "per_rank": [ + { + "encode_ms": 0.7971839904785156, + "diffuse_ms": 1762.2459716796875, + "decode_ms": 7.487936019897461, + "finalize_ms": 438.0950012207031, + "total_ms": 2208.6260929107666, + "total_ms_wo_finalize": 1770.5310916900635, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2209.239514544606, + "output_frames": 12, + "rank": 0 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.0685323029756546, + "bandwidth_gbps": 15.300463496352894 + } + ], + "all_gather": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.037375837564468384, + "bandwidth_gbps": 28.054916446791186 + } + ] + } + }, + { + "environment": { + "command": "env TORCHINDUCTOR_COMPILE_THREADS=4 uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 2 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-2", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01151", + "slurm_job_id": "14793417", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0 + ], + "dit_cp_group": [ + 0 + ], + "cp_size": 1, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42 + ], + "replica_id": 2, + "measurement_window": { + "started_at": 11145874.429111755, + "finished_at": 11145885.16577754, + "elapsed_s": 10.736665785312653 + } + }, + "summary": { + "fps": 5.589863810528405, + "latency_ms": { + "median": 2144.325716421008, + "p90": 2151.1040780693293, + "min": 2143.3675326406956, + "max": 2151.623649522662 + }, + "encoder_ms": { + "median": 0.8196160197257996, + "p90": 0.9688384056091309, + "min": 0.7701759934425354, + "max": 1.058784008026123 + }, + "dit_ms": { + "median": 1712.18017578125, + "p90": 1718.5796875, + "min": 1706.93115234375, + "max": 1720.412841796875 + }, + "decoder_ms": { + "median": 7.421631813049316, + "p90": 7.495385646820068, + "min": 7.307648181915283, + "max": 7.505280017852783 + }, + "finalize_ms": { + "median": 423.29205322265625, + "p90": 431.4473022460937, + "min": 418.8133850097656, + "max": 434.0836181640625 + }, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 2107.3786565692085, + "cp_probe_gbps": { + "broadcast": { + "median": 24.04741364465237, + "p90": 24.04741364465237, + "min": 24.04741364465237, + "max": 24.04741364465237 + }, + "all_gather": { + "median": 33.40156363007666, + "p90": 33.40156363007666, + "min": 33.40156363007666, + "max": 33.40156363007666 + } + }, + "memory": { + "peak_gib_by_rank": [ + 59.35440921783447 + ], + "steady_allocated_gib_by_rank": [ + 57.14538335800171 + ], + "initialization_peak_gib_by_rank": [ + 66.54739189147949 + ], + "node_peak_gib": 59.35440921783447, + "node_steady_allocated_gib": 57.14538335800171, + "per_rank_peak_gib": { + "median": 59.35440921783447, + "p90": 59.35440921783447, + "min": 59.35440921783447, + "max": 59.35440921783447 + } + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 124731.86967894435, + "output_frames": 9, + "critical_rank": { + "encode_ms": 583.7667236328125, + "diffuse_ms": 96032.53125, + "decode_ms": 27845.583984375, + "finalize_ms": 263.0707092285156 + }, + "per_rank": [ + { + "encode_ms": 583.7667236328125, + "diffuse_ms": 96032.53125, + "decode_ms": 27845.583984375, + "finalize_ms": 263.0707092285156, + "total_ms": 124724.95266723633, + "total_ms_wo_finalize": 124461.88195800781, + "mem_alloc_gib": 57.087345123291016, + "mem_reserved_gib": 61.439453125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 124731.86967894435, + "output_frames": 9, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 81199.26889054477, + "output_frames": 12, + "critical_rank": { + "encode_ms": 181.4916534423828, + "diffuse_ms": 80695.2421875, + "decode_ms": 8.290847778320312, + "finalize_ms": 312.3619689941406 + }, + "per_rank": [ + { + "encode_ms": 181.4916534423828, + "diffuse_ms": 80695.2421875, + "decode_ms": 8.290847778320312, + "finalize_ms": 312.3619689941406, + "total_ms": 81197.38665771484, + "total_ms_wo_finalize": 80885.0246887207, + "mem_alloc_gib": 57.08803939819336, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 81199.26889054477, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 1861.066561192274, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.9847412109375, + "diffuse_ms": 1364.094970703125, + "decode_ms": 7.538112163543701, + "finalize_ms": 324.8186950683594 + }, + "per_rank": [ + { + "encode_ms": 163.9847412109375, + "diffuse_ms": 1364.094970703125, + "decode_ms": 7.538112163543701, + "finalize_ms": 324.8186950683594, + "total_ms": 1860.4365191459656, + "total_ms_wo_finalize": 1535.6178240776062, + "mem_alloc_gib": 57.08770561218262, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 1861.066561192274, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 2014.5772360265255, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.66934204101562, + "diffuse_ms": 1473.288330078125, + "decode_ms": 19.154176712036133, + "finalize_ms": 357.7861633300781 + }, + "per_rank": [ + { + "encode_ms": 163.66934204101562, + "diffuse_ms": 1473.288330078125, + "decode_ms": 19.154176712036133, + "finalize_ms": 357.7861633300781, + "total_ms": 2013.8980121612549, + "total_ms_wo_finalize": 1656.1118488311768, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.423828125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2014.5772360265255, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 2168.692423030734, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.9001922607422, + "diffuse_ms": 1592.3995361328125, + "decode_ms": 7.360960006713867, + "finalize_ms": 403.35052490234375 + }, + "per_rank": [ + { + "encode_ms": 164.9001922607422, + "diffuse_ms": 1592.3995361328125, + "decode_ms": 7.360960006713867, + "finalize_ms": 403.35052490234375, + "total_ms": 2168.0112133026123, + "total_ms_wo_finalize": 1764.6606884002686, + "mem_alloc_gib": 57.144843101501465, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2168.692423030734, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 2178.6916963756084, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9667199850082397, + "diffuse_ms": 1732.883056640625, + "decode_ms": 7.317311763763428, + "finalize_ms": 436.9365234375 + }, + "per_rank": [ + { + "encode_ms": 0.9667199850082397, + "diffuse_ms": 1732.883056640625, + "decode_ms": 7.317311763763428, + "finalize_ms": 436.9365234375, + "total_ms": 2178.1036118268967, + "total_ms_wo_finalize": 1741.1670883893967, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2178.6916963756084, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 2151.623649522662, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.058784008026123, + "diffuse_ms": 1706.93115234375, + "decode_ms": 7.307648181915283, + "finalize_ms": 434.0836181640625 + }, + "per_rank": [ + { + "encode_ms": 1.058784008026123, + "diffuse_ms": 1706.93115234375, + "decode_ms": 7.307648181915283, + "finalize_ms": 434.0836181640625, + "total_ms": 2149.381202697754, + "total_ms_wo_finalize": 1715.2975845336914, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2151.623649522662, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 2144.0724804997444, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8339200019836426, + "diffuse_ms": 1707.73046875, + "decode_ms": 7.480544090270996, + "finalize_ms": 427.4928283691406 + }, + "per_rank": [ + { + "encode_ms": 0.8339200019836426, + "diffuse_ms": 1707.73046875, + "decode_ms": 7.480544090270996, + "finalize_ms": 427.4928283691406, + "total_ms": 2143.5377612113953, + "total_ms_wo_finalize": 1716.0449328422546, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2144.0724804997444, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 2144.325716421008, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7915840148925781, + "diffuse_ms": 1712.18017578125, + "decode_ms": 7.505280017852783, + "finalize_ms": 423.29205322265625 + }, + "per_rank": [ + { + "encode_ms": 0.7915840148925781, + "diffuse_ms": 1712.18017578125, + "decode_ms": 7.505280017852783, + "finalize_ms": 423.29205322265625, + "total_ms": 2143.7690930366516, + "total_ms_wo_finalize": 1720.4770398139954, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2144.325716421008, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 2150.32472088933, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8196160197257996, + "diffuse_ms": 1720.412841796875, + "decode_ms": 7.421631813049316, + "finalize_ms": 421.1144714355469 + }, + "per_rank": [ + { + "encode_ms": 0.8196160197257996, + "diffuse_ms": 1720.412841796875, + "decode_ms": 7.421631813049316, + "finalize_ms": 421.1144714355469, + "total_ms": 2149.768561065197, + "total_ms_wo_finalize": 1728.6540896296501, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2150.32472088933, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 2143.3675326406956, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7701759934425354, + "diffuse_ms": 1715.8299560546875, + "decode_ms": 7.379551887512207, + "finalize_ms": 418.8133850097656 + }, + "per_rank": [ + { + "encode_ms": 0.7701759934425354, + "diffuse_ms": 1715.8299560546875, + "decode_ms": 7.379551887512207, + "finalize_ms": 418.8133850097656, + "total_ms": 2142.793068945408, + "total_ms_wo_finalize": 1723.9796839356422, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2143.3675326406956, + "output_frames": 12, + "rank": 0 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.043604522943496704, + "bandwidth_gbps": 24.04741364465237 + } + ], + "all_gather": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.03139302134513855, + "bandwidth_gbps": 33.40156363007666 + } + ] + } + }, + { + "environment": { + "command": "env TORCHINDUCTOR_COMPILE_THREADS=4 uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 3 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-3", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01151", + "slurm_job_id": "14793417", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0 + ], + "dit_cp_group": [ + 0 + ], + "cp_size": 1, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42 + ], + "replica_id": 3, + "measurement_window": { + "started_at": 11145874.429096408, + "finished_at": 11145885.149589293, + "elapsed_s": 10.720492884516716 + } + }, + "summary": { + "fps": 5.5984819295074555, + "latency_ms": { + "median": 2141.4667814970016, + "p90": 2148.58845025301, + "min": 2137.736974284053, + "max": 2149.8332284390926 + }, + "encoder_ms": { + "median": 0.8732479810714722, + "p90": 1.1013311862945556, + "min": 0.8052160143852234, + "max": 1.227295994758606 + }, + "dit_ms": { + "median": 1709.76220703125, + "p90": 1713.335693359375, + "min": 1701.818603515625, + "max": 1715.22412109375 + }, + "decoder_ms": { + "median": 7.462048053741455, + "p90": 7.494681549072266, + "min": 7.314432144165039, + "max": 7.5155839920043945 + }, + "finalize_ms": { + "median": 427.0291442871094, + "p90": 429.3756530761719, + "min": 422.11553955078125, + "max": 429.9440612792969 + }, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 2110.627687424311, + "cp_probe_gbps": { + "broadcast": { + "median": 19.97764127262543, + "p90": 19.97764127262543, + "min": 19.97764127262543, + "max": 19.97764127262543 + }, + "all_gather": { + "median": 32.44294337375011, + "p90": 32.44294337375011, + "min": 32.44294337375011, + "max": 32.44294337375011 + } + }, + "memory": { + "peak_gib_by_rank": [ + 59.35440921783447 + ], + "steady_allocated_gib_by_rank": [ + 57.14538335800171 + ], + "initialization_peak_gib_by_rank": [ + 66.54739189147949 + ], + "node_peak_gib": 59.35440921783447, + "node_steady_allocated_gib": 57.14538335800171, + "per_rank_peak_gib": { + "median": 59.35440921783447, + "p90": 59.35440921783447, + "min": 59.35440921783447, + "max": 59.35440921783447 + } + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 191956.83398284018, + "output_frames": 9, + "critical_rank": { + "encode_ms": 630.0097045898438, + "diffuse_ms": 92541.125, + "decode_ms": 98508.3203125, + "finalize_ms": 271.2525329589844 + }, + "per_rank": [ + { + "encode_ms": 630.0097045898438, + "diffuse_ms": 92541.125, + "decode_ms": 98508.3203125, + "finalize_ms": 271.2525329589844, + "total_ms": 191950.70755004883, + "total_ms_wo_finalize": 191679.45501708984, + "mem_alloc_gib": 57.087345123291016, + "mem_reserved_gib": 61.439453125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 191956.83398284018, + "output_frames": 9, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 16647.290909662843, + "output_frames": 12, + "critical_rank": { + "encode_ms": 183.00636291503906, + "diffuse_ms": 16145.9697265625, + "decode_ms": 8.30463981628418, + "finalize_ms": 308.3145751953125 + }, + "per_rank": [ + { + "encode_ms": 183.00636291503906, + "diffuse_ms": 16145.9697265625, + "decode_ms": 8.30463981628418, + "finalize_ms": 308.3145751953125, + "total_ms": 16645.595304489136, + "total_ms_wo_finalize": 16337.280729293823, + "mem_alloc_gib": 57.08803939819336, + "mem_reserved_gib": 61.8515625, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 16647.290909662843, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 1868.4289250522852, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.72274780273438, + "diffuse_ms": 1366.1490478515625, + "decode_ms": 7.524960041046143, + "finalize_ms": 330.417236328125 + }, + "per_rank": [ + { + "encode_ms": 163.72274780273438, + "diffuse_ms": 1366.1490478515625, + "decode_ms": 7.524960041046143, + "finalize_ms": 330.417236328125, + "total_ms": 1867.813992023468, + "total_ms_wo_finalize": 1537.396755695343, + "mem_alloc_gib": 57.08770561218262, + "mem_reserved_gib": 61.8515625, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 1868.4289250522852, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 2023.038288578391, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.7052459716797, + "diffuse_ms": 1482.0423583984375, + "decode_ms": 19.439231872558594, + "finalize_ms": 357.0379943847656 + }, + "per_rank": [ + { + "encode_ms": 163.7052459716797, + "diffuse_ms": 1482.0423583984375, + "decode_ms": 19.439231872558594, + "finalize_ms": 357.0379943847656, + "total_ms": 2022.2248306274414, + "total_ms_wo_finalize": 1665.1868362426758, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.47265625, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2023.038288578391, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 2178.171617910266, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.34169006347656, + "diffuse_ms": 1600.6512451171875, + "decode_ms": 7.3787522315979, + "finalize_ms": 405.10858154296875 + }, + "per_rank": [ + { + "encode_ms": 164.34169006347656, + "diffuse_ms": 1600.6512451171875, + "decode_ms": 7.3787522315979, + "finalize_ms": 405.10858154296875, + "total_ms": 2177.4802689552307, + "total_ms_wo_finalize": 1772.371687412262, + "mem_alloc_gib": 57.144843101501465, + "mem_reserved_gib": 62.474609375, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2178.171617910266, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 23423.814419656992, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9950079917907715, + "diffuse_ms": 22997.19921875, + "decode_ms": 7.346911907196045, + "finalize_ms": 417.6161804199219 + }, + "per_rank": [ + { + "encode_ms": 0.9950079917907715, + "diffuse_ms": 22997.19921875, + "decode_ms": 7.346911907196045, + "finalize_ms": 417.6161804199219, + "total_ms": 23423.15731906891, + "total_ms_wo_finalize": 23005.541138648987, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.474609375, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 23423.814419656992, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 2149.8332284390926, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.227295994758606, + "diffuse_ms": 1709.76220703125, + "decode_ms": 7.314432144165039, + "finalize_ms": 429.9440612792969 + }, + "per_rank": [ + { + "encode_ms": 1.227295994758606, + "diffuse_ms": 1709.76220703125, + "decode_ms": 7.314432144165039, + "finalize_ms": 429.9440612792969, + "total_ms": 2148.2479964494705, + "total_ms_wo_finalize": 1718.3039351701736, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2149.8332284390926, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 2137.736974284053, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8732479810714722, + "diffuse_ms": 1701.818603515625, + "decode_ms": 7.462048053741455, + "finalize_ms": 427.0291442871094 + }, + "per_rank": [ + { + "encode_ms": 0.8732479810714722, + "diffuse_ms": 1701.818603515625, + "decode_ms": 7.462048053741455, + "finalize_ms": 427.0291442871094, + "total_ms": 2137.1830438375473, + "total_ms_wo_finalize": 1710.153899550438, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2137.736974284053, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 2141.4327062666416, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8181440234184265, + "diffuse_ms": 1704.036376953125, + "decode_ms": 7.422719955444336, + "finalize_ms": 428.5230407714844 + }, + "per_rank": [ + { + "encode_ms": 0.8181440234184265, + "diffuse_ms": 1704.036376953125, + "decode_ms": 7.422719955444336, + "finalize_ms": 428.5230407714844, + "total_ms": 2140.800281703472, + "total_ms_wo_finalize": 1712.2772409319878, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2141.4327062666416, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 2146.7212829738855, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9123839735984802, + "diffuse_ms": 1715.22412109375, + "decode_ms": 7.463327884674072, + "finalize_ms": 422.5871276855469 + }, + "per_rank": [ + { + "encode_ms": 0.9123839735984802, + "diffuse_ms": 1715.22412109375, + "decode_ms": 7.463327884674072, + "finalize_ms": 422.5871276855469, + "total_ms": 2146.1869606375694, + "total_ms_wo_finalize": 1723.5998329520226, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2146.7212829738855, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 2141.4667814970016, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8052160143852234, + "diffuse_ms": 1710.5030517578125, + "decode_ms": 7.5155839920043945, + "finalize_ms": 422.11553955078125 + }, + "per_rank": [ + { + "encode_ms": 0.8052160143852234, + "diffuse_ms": 1710.5030517578125, + "decode_ms": 7.5155839920043945, + "finalize_ms": 422.11553955078125, + "total_ms": 2140.9393913149834, + "total_ms_wo_finalize": 1718.8238517642021, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2141.4667814970016, + "output_frames": 12, + "rank": 0 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.05248747766017914, + "bandwidth_gbps": 19.97764127262543 + } + ], + "all_gather": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.032320618629455566, + "bandwidth_gbps": 32.44294337375011 + } + ] + } + }, + { + "environment": { + "command": "env TORCHINDUCTOR_COMPILE_THREADS=4 uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 4 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-4", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01151", + "slurm_job_id": "14793417", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0 + ], + "dit_cp_group": [ + 0 + ], + "cp_size": 1, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42 + ], + "replica_id": 4, + "measurement_window": { + "started_at": 11145874.429111345, + "finished_at": 11145885.253132468, + "elapsed_s": 10.824021123349667 + } + }, + "summary": { + "fps": 5.54487966051973, + "latency_ms": { + "median": 2166.7757872492075, + "p90": 2167.4971897155046, + "min": 2157.7277276664972, + "max": 2167.676465585828 + }, + "encoder_ms": { + "median": 0.856544017791748, + "p90": 1.0044992208480834, + "min": 0.7843199968338013, + "max": 1.0886080265045166 + }, + "dit_ms": { + "median": 1725.0692138671875, + "p90": 1734.29072265625, + "min": 1718.1109619140625, + "max": 1735.802978515625 + }, + "decoder_ms": { + "median": 7.404543876647949, + "p90": 7.512204837799072, + "min": 7.3921918869018555, + "max": 7.547520160675049 + }, + "finalize_ms": { + "median": 425.2576599121094, + "p90": 434.6324096679688, + "min": 422.6251220703125, + "max": 439.471435546875 + }, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 2090.4196320159385, + "cp_probe_gbps": { + "broadcast": { + "median": 19.047536911565285, + "p90": 19.047536911565285, + "min": 19.047536911565285, + "max": 19.047536911565285 + }, + "all_gather": { + "median": 28.24777728041106, + "p90": 28.24777728041106, + "min": 28.24777728041106, + "max": 28.24777728041106 + } + }, + "memory": { + "peak_gib_by_rank": [ + 59.35440921783447 + ], + "steady_allocated_gib_by_rank": [ + 57.14538335800171 + ], + "initialization_peak_gib_by_rank": [ + 66.54739189147949 + ], + "node_peak_gib": 59.35440921783447, + "node_steady_allocated_gib": 57.14538335800171, + "per_rank_peak_gib": { + "median": 59.35440921783447, + "p90": 59.35440921783447, + "min": 59.35440921783447, + "max": 59.35440921783447 + } + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 204641.40384458005, + "output_frames": 9, + "critical_rank": { + "encode_ms": 614.59423828125, + "diffuse_ms": 92576.8046875, + "decode_ms": 111181.4765625, + "finalize_ms": 266.4762268066406 + }, + "per_rank": [ + { + "encode_ms": 614.59423828125, + "diffuse_ms": 92576.8046875, + "decode_ms": 111181.4765625, + "finalize_ms": 266.4762268066406, + "total_ms": 204639.3517150879, + "total_ms_wo_finalize": 204372.87548828125, + "mem_alloc_gib": 57.087345123291016, + "mem_reserved_gib": 61.439453125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 204641.40384458005, + "output_frames": 9, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 13033.618155866861, + "output_frames": 12, + "critical_rank": { + "encode_ms": 182.26185607910156, + "diffuse_ms": 12527.947265625, + "decode_ms": 8.210687637329102, + "finalize_ms": 314.601318359375 + }, + "per_rank": [ + { + "encode_ms": 182.26185607910156, + "diffuse_ms": 12527.947265625, + "decode_ms": 8.210687637329102, + "finalize_ms": 314.601318359375, + "total_ms": 13033.021127700806, + "total_ms_wo_finalize": 12718.41980934143, + "mem_alloc_gib": 57.08803939819336, + "mem_reserved_gib": 61.984375, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 13033.618155866861, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 1881.8960543721914, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.98851013183594, + "diffuse_ms": 1375.63525390625, + "decode_ms": 7.641823768615723, + "finalize_ms": 334.04876708984375 + }, + "per_rank": [ + { + "encode_ms": 163.98851013183594, + "diffuse_ms": 1375.63525390625, + "decode_ms": 7.641823768615723, + "finalize_ms": 334.04876708984375, + "total_ms": 1881.3143548965454, + "total_ms_wo_finalize": 1547.2655878067017, + "mem_alloc_gib": 57.08770561218262, + "mem_reserved_gib": 61.984375, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 1881.8960543721914, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 2042.4561835825443, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.9012451171875, + "diffuse_ms": 1495.1202392578125, + "decode_ms": 19.154144287109375, + "finalize_ms": 363.7153625488281 + }, + "per_rank": [ + { + "encode_ms": 163.9012451171875, + "diffuse_ms": 1495.1202392578125, + "decode_ms": 19.154144287109375, + "finalize_ms": 363.7153625488281, + "total_ms": 2041.8909912109375, + "total_ms_wo_finalize": 1678.1756286621094, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.46875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2042.4561835825443, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 2196.0909999907017, + "output_frames": 12, + "critical_rank": { + "encode_ms": 165.92530822753906, + "diffuse_ms": 1614.0648193359375, + "decode_ms": 7.3546881675720215, + "finalize_ms": 408.1500549316406 + }, + "per_rank": [ + { + "encode_ms": 165.92530822753906, + "diffuse_ms": 1614.0648193359375, + "decode_ms": 7.3546881675720215, + "finalize_ms": 408.1500549316406, + "total_ms": 2195.494870662689, + "total_ms_wo_finalize": 1787.3448157310486, + "mem_alloc_gib": 57.144843101501465, + "mem_reserved_gib": 62.470703125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2196.0909999907017, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 23326.71902887523, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9790400266647339, + "diffuse_ms": 22897.310546875, + "decode_ms": 7.355807781219482, + "finalize_ms": 420.5391540527344 + }, + "per_rank": [ + { + "encode_ms": 0.9790400266647339, + "diffuse_ms": 22897.310546875, + "decode_ms": 7.355807781219482, + "finalize_ms": 420.5391540527344, + "total_ms": 23326.18454873562, + "total_ms_wo_finalize": 22905.645394682884, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.470703125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 23326.71902887523, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 2167.676465585828, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.0886080265045166, + "diffuse_ms": 1718.1109619140625, + "decode_ms": 7.394752025604248, + "finalize_ms": 439.471435546875 + }, + "per_rank": [ + { + "encode_ms": 1.0886080265045166, + "diffuse_ms": 1718.1109619140625, + "decode_ms": 7.394752025604248, + "finalize_ms": 439.471435546875, + "total_ms": 2166.0657575130463, + "total_ms_wo_finalize": 1726.5943219661713, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2167.676465585828, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 2161.385642364621, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.856544017791748, + "diffuse_ms": 1725.0692138671875, + "decode_ms": 7.547520160675049, + "finalize_ms": 427.3738708496094 + }, + "per_rank": [ + { + "encode_ms": 0.856544017791748, + "diffuse_ms": 1725.0692138671875, + "decode_ms": 7.547520160675049, + "finalize_ms": 427.3738708496094, + "total_ms": 2160.8471488952637, + "total_ms_wo_finalize": 1733.4732780456543, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2161.385642364621, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 2157.7277276664972, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8032320141792297, + "diffuse_ms": 1724.963623046875, + "decode_ms": 7.459231853485107, + "finalize_ms": 423.8592529296875 + }, + "per_rank": [ + { + "encode_ms": 0.8032320141792297, + "diffuse_ms": 1724.963623046875, + "decode_ms": 7.459231853485107, + "finalize_ms": 423.8592529296875, + "total_ms": 2157.085339844227, + "total_ms_wo_finalize": 1733.2260869145393, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2157.7277276664972, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 2167.22827591002, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8783360123634338, + "diffuse_ms": 1735.802978515625, + "decode_ms": 7.404543876647949, + "finalize_ms": 422.6251220703125 + }, + "per_rank": [ + { + "encode_ms": 0.8783360123634338, + "diffuse_ms": 1735.802978515625, + "decode_ms": 7.404543876647949, + "finalize_ms": 422.6251220703125, + "total_ms": 2166.710980474949, + "total_ms_wo_finalize": 1744.0858584046364, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2167.22827591002, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 2166.7757872492075, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7843199968338013, + "diffuse_ms": 1732.0223388671875, + "decode_ms": 7.3921918869018555, + "finalize_ms": 425.2576599121094 + }, + "per_rank": [ + { + "encode_ms": 0.7843199968338013, + "diffuse_ms": 1732.0223388671875, + "decode_ms": 7.3921918869018555, + "finalize_ms": 425.2576599121094, + "total_ms": 2165.4565106630325, + "total_ms_wo_finalize": 1740.1988507509232, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2166.7757872492075, + "output_frames": 12, + "rank": 0 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.055050477385520935, + "bandwidth_gbps": 19.047536911565285 + } + ], + "all_gather": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.03712065517902374, + "bandwidth_gbps": 28.24777728041106 + } + ] + } + }, + { + "environment": { + "command": "env TORCHINDUCTOR_COMPILE_THREADS=4 uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 5 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-5", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01151", + "slurm_job_id": "14793417", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0 + ], + "dit_cp_group": [ + 0 + ], + "cp_size": 1, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42 + ], + "replica_id": 5, + "measurement_window": { + "started_at": 11145874.429121751, + "finished_at": 11145885.288201544, + "elapsed_s": 10.859079793095589 + } + }, + "summary": { + "fps": 5.526838929943354, + "latency_ms": { + "median": 2172.3211836069822, + "p90": 2173.0400942265987, + "min": 2165.900409221649, + "max": 2173.2870899140835 + }, + "encoder_ms": { + "median": 0.8277760148048401, + "p90": 0.9301823973655701, + "min": 0.8144639730453491, + "max": 0.9792320132255554 + }, + "dit_ms": { + "median": 1733.2933349609375, + "p90": 1736.2165771484374, + "min": 1730.846435546875, + "max": 1736.8795166015625 + }, + "decoder_ms": { + "median": 7.442975997924805, + "p90": 7.5361793518066404, + "min": 7.399968147277832, + "max": 7.574848175048828 + }, + "finalize_ms": { + "median": 429.31097412109375, + "p90": 430.79743041992185, + "min": 426.1729736328125, + "max": 431.2060852050781 + }, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 2083.6182765886447, + "cp_probe_gbps": { + "broadcast": { + "median": 22.004845148001095, + "p90": 22.004845148001095, + "min": 22.004845148001095, + "max": 22.004845148001095 + }, + "all_gather": { + "median": 34.167877726469534, + "p90": 34.167877726469534, + "min": 34.167877726469534, + "max": 34.167877726469534 + } + }, + "memory": { + "peak_gib_by_rank": [ + 59.35440921783447 + ], + "steady_allocated_gib_by_rank": [ + 57.14538335800171 + ], + "initialization_peak_gib_by_rank": [ + 66.54739189147949 + ], + "node_peak_gib": 59.35440921783447, + "node_steady_allocated_gib": 57.14538335800171, + "per_rank_peak_gib": { + "median": 59.35440921783447, + "p90": 59.35440921783447, + "min": 59.35440921783447, + "max": 59.35440921783447 + } + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 179743.6931002885, + "output_frames": 9, + "critical_rank": { + "encode_ms": 602.4649047851562, + "diffuse_ms": 92286.421875, + "decode_ms": 86584.2265625, + "finalize_ms": 268.6270751953125 + }, + "per_rank": [ + { + "encode_ms": 602.4649047851562, + "diffuse_ms": 92286.421875, + "decode_ms": 86584.2265625, + "finalize_ms": 268.6270751953125, + "total_ms": 179741.74041748047, + "total_ms_wo_finalize": 179473.11334228516, + "mem_alloc_gib": 57.087345123291016, + "mem_reserved_gib": 61.439453125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 179743.6931002885, + "output_frames": 9, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 85268.3373503387, + "output_frames": 12, + "critical_rank": { + "encode_ms": 183.08604431152344, + "diffuse_ms": 84764.8359375, + "decode_ms": 8.193632125854492, + "finalize_ms": 310.8353271484375 + }, + "per_rank": [ + { + "encode_ms": 183.08604431152344, + "diffuse_ms": 84764.8359375, + "decode_ms": 8.193632125854492, + "finalize_ms": 310.8353271484375, + "total_ms": 85266.95094108582, + "total_ms_wo_finalize": 84956.11561393738, + "mem_alloc_gib": 57.08803939819336, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 85268.3373503387, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 1873.354883864522, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.9582061767578, + "diffuse_ms": 1372.1710205078125, + "decode_ms": 7.503200054168701, + "finalize_ms": 329.1710205078125 + }, + "per_rank": [ + { + "encode_ms": 163.9582061767578, + "diffuse_ms": 1372.1710205078125, + "decode_ms": 7.503200054168701, + "finalize_ms": 329.1710205078125, + "total_ms": 1872.8034472465515, + "total_ms_wo_finalize": 1543.632426738739, + "mem_alloc_gib": 57.08770561218262, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 1873.354883864522, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 2037.762951105833, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.39071655273438, + "diffuse_ms": 1489.8067626953125, + "decode_ms": 18.76563262939453, + "finalize_ms": 364.2565002441406 + }, + "per_rank": [ + { + "encode_ms": 164.39071655273438, + "diffuse_ms": 1489.8067626953125, + "decode_ms": 18.76563262939453, + "finalize_ms": 364.2565002441406, + "total_ms": 2037.219612121582, + "total_ms_wo_finalize": 1672.9631118774414, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.423828125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2037.762951105833, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 2178.1770400702953, + "output_frames": 12, + "critical_rank": { + "encode_ms": 166.12477111816406, + "diffuse_ms": 1598.2432861328125, + "decode_ms": 7.306079864501953, + "finalize_ms": 405.9862365722656 + }, + "per_rank": [ + { + "encode_ms": 166.12477111816406, + "diffuse_ms": 1598.2432861328125, + "decode_ms": 7.306079864501953, + "finalize_ms": 405.9862365722656, + "total_ms": 2177.660373687744, + "total_ms_wo_finalize": 1771.6741371154785, + "mem_alloc_gib": 57.144843101501465, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2178.1770400702953, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 2168.7360797077417, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8688639998435974, + "diffuse_ms": 1722.0150146484375, + "decode_ms": 7.41596794128418, + "finalize_ms": 437.7752380371094 + }, + "per_rank": [ + { + "encode_ms": 0.8688639998435974, + "diffuse_ms": 1722.0150146484375, + "decode_ms": 7.41596794128418, + "finalize_ms": 437.7752380371094, + "total_ms": 2168.0750846266747, + "total_ms_wo_finalize": 1730.2998465895653, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2168.7360797077417, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 2171.936895698309, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9792320132255554, + "diffuse_ms": 1730.888916015625, + "decode_ms": 7.574848175048828, + "finalize_ms": 431.2060852050781 + }, + "per_rank": [ + { + "encode_ms": 0.9792320132255554, + "diffuse_ms": 1730.888916015625, + "decode_ms": 7.574848175048828, + "finalize_ms": 431.2060852050781, + "total_ms": 2170.6490814089775, + "total_ms_wo_finalize": 1739.4429962038994, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2171.936895698309, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 2165.900409221649, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.856607973575592, + "diffuse_ms": 1730.846435546875, + "decode_ms": 7.478176116943359, + "finalize_ms": 426.1729736328125 + }, + "per_rank": [ + { + "encode_ms": 0.856607973575592, + "diffuse_ms": 1730.846435546875, + "decode_ms": 7.478176116943359, + "finalize_ms": 426.1729736328125, + "total_ms": 2165.3541932702065, + "total_ms_wo_finalize": 1739.181219637394, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2165.900409221649, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 2172.6696006953716, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8150079846382141, + "diffuse_ms": 1736.8795166015625, + "decode_ms": 7.401760101318359, + "finalize_ms": 427.0155944824219 + }, + "per_rank": [ + { + "encode_ms": 0.8150079846382141, + "diffuse_ms": 1736.8795166015625, + "decode_ms": 7.401760101318359, + "finalize_ms": 427.0155944824219, + "total_ms": 2172.111879169941, + "total_ms_wo_finalize": 1745.096284687519, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2172.6696006953716, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 2173.2870899140835, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8144639730453491, + "diffuse_ms": 1735.22216796875, + "decode_ms": 7.399968147277832, + "finalize_ms": 429.31097412109375 + }, + "per_rank": [ + { + "encode_ms": 0.8144639730453491, + "diffuse_ms": 1735.22216796875, + "decode_ms": 7.399968147277832, + "finalize_ms": 429.31097412109375, + "total_ms": 2172.747574210167, + "total_ms_wo_finalize": 1743.4366000890732, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2173.2870899140835, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 2172.3211836069822, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8277760148048401, + "diffuse_ms": 1733.2933349609375, + "decode_ms": 7.442975997924805, + "finalize_ms": 430.1844482421875 + }, + "per_rank": [ + { + "encode_ms": 0.8277760148048401, + "diffuse_ms": 1733.2933349609375, + "decode_ms": 7.442975997924805, + "finalize_ms": 430.1844482421875, + "total_ms": 2171.7485352158546, + "total_ms_wo_finalize": 1741.5640869736671, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2172.3211836069822, + "output_frames": 12, + "rank": 0 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.047652050852775574, + "bandwidth_gbps": 22.004845148001095 + } + ], + "all_gather": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.030688941478729248, + "bandwidth_gbps": 34.167877726469534 + } + ] + } + }, + { + "environment": { + "command": "env TORCHINDUCTOR_COMPILE_THREADS=4 uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 6 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-6", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01151", + "slurm_job_id": "14793417", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0 + ], + "dit_cp_group": [ + 0 + ], + "cp_size": 1, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42 + ], + "replica_id": 6, + "measurement_window": { + "started_at": 11145874.429398024, + "finished_at": 11145885.253023999, + "elapsed_s": 10.823625974357128 + } + }, + "summary": { + "fps": 5.545099936557462, + "latency_ms": { + "median": 2159.7884446382523, + "p90": 2171.7981200665236, + "min": 2158.867284655571, + "max": 2173.339394852519 + }, + "encoder_ms": { + "median": 0.8133440017700195, + "p90": 0.999020791053772, + "min": 0.7985600233078003, + "max": 1.0879679918289185 + }, + "dit_ms": { + "median": 1724.6712646484375, + "p90": 1736.9388427734375, + "min": 1710.7734375, + "max": 1737.7554931640625 + }, + "decoder_ms": { + "median": 7.494592189788818, + "p90": 7.506969451904297, + "min": 7.348320007324219, + "max": 7.512127876281738 + }, + "finalize_ms": { + "median": 426.65533447265625, + "p90": 435.1299621582031, + "min": 424.867431640625, + "max": 438.1396789550781 + }, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 2090.5026760821634, + "cp_probe_gbps": { + "broadcast": { + "median": 22.40329327528303, + "p90": 22.40329327528303, + "min": 22.40329327528303, + "max": 22.40329327528303 + }, + "all_gather": { + "median": 32.5592801284738, + "p90": 32.5592801284738, + "min": 32.5592801284738, + "max": 32.5592801284738 + } + }, + "memory": { + "peak_gib_by_rank": [ + 59.35440921783447 + ], + "steady_allocated_gib_by_rank": [ + 57.14538335800171 + ], + "initialization_peak_gib_by_rank": [ + 66.54739189147949 + ], + "node_peak_gib": 59.35440921783447, + "node_steady_allocated_gib": 57.14538335800171, + "per_rank_peak_gib": { + "median": 59.35440921783447, + "p90": 59.35440921783447, + "min": 59.35440921783447, + "max": 59.35440921783447 + } + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 194030.51247261465, + "output_frames": 9, + "critical_rank": { + "encode_ms": 583.8138427734375, + "diffuse_ms": 92297.609375, + "decode_ms": 100880.5, + "finalize_ms": 264.7221984863281 + }, + "per_rank": [ + { + "encode_ms": 583.8138427734375, + "diffuse_ms": 92297.609375, + "decode_ms": 100880.5, + "finalize_ms": 264.7221984863281, + "total_ms": 194026.64541625977, + "total_ms_wo_finalize": 193761.92321777344, + "mem_alloc_gib": 57.087345123291016, + "mem_reserved_gib": 61.439453125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 194030.51247261465, + "output_frames": 9, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 15083.65355990827, + "output_frames": 12, + "critical_rank": { + "encode_ms": 182.58604431152344, + "diffuse_ms": 14577.419921875, + "decode_ms": 8.33779239654541, + "finalize_ms": 313.6082458496094 + }, + "per_rank": [ + { + "encode_ms": 182.58604431152344, + "diffuse_ms": 14577.419921875, + "decode_ms": 8.33779239654541, + "finalize_ms": 313.6082458496094, + "total_ms": 15081.952004432678, + "total_ms_wo_finalize": 14768.343758583069, + "mem_alloc_gib": 57.08803939819336, + "mem_reserved_gib": 61.984375, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 15083.65355990827, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 1872.1758779138327, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.10211181640625, + "diffuse_ms": 1370.0150146484375, + "decode_ms": 7.51907205581665, + "finalize_ms": 329.8470764160156 + }, + "per_rank": [ + { + "encode_ms": 164.10211181640625, + "diffuse_ms": 1370.0150146484375, + "decode_ms": 7.51907205581665, + "finalize_ms": 329.8470764160156, + "total_ms": 1871.483274936676, + "total_ms_wo_finalize": 1541.6361985206604, + "mem_alloc_gib": 57.08770561218262, + "mem_reserved_gib": 61.984375, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 1872.1758779138327, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 2037.8571413457394, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.83395385742188, + "diffuse_ms": 1482.5743408203125, + "decode_ms": 27.436864852905273, + "finalize_ms": 363.39678955078125 + }, + "per_rank": [ + { + "encode_ms": 163.83395385742188, + "diffuse_ms": 1482.5743408203125, + "decode_ms": 27.436864852905273, + "finalize_ms": 363.39678955078125, + "total_ms": 2037.241949081421, + "total_ms_wo_finalize": 1673.8451595306396, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.46875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2037.8571413457394, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 2179.8755042254925, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.06253051757812, + "diffuse_ms": 1599.4937744140625, + "decode_ms": 7.367199897766113, + "finalize_ms": 407.2137451171875 + }, + "per_rank": [ + { + "encode_ms": 164.06253051757812, + "diffuse_ms": 1599.4937744140625, + "decode_ms": 7.367199897766113, + "finalize_ms": 407.2137451171875, + "total_ms": 2178.1372499465942, + "total_ms_wo_finalize": 1770.9235048294067, + "mem_alloc_gib": 57.144843101501465, + "mem_reserved_gib": 62.470703125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2179.8755042254925, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 23413.542149588466, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.0579520463943481, + "diffuse_ms": 22986.931640625, + "decode_ms": 7.357728004455566, + "finalize_ms": 417.3621826171875 + }, + "per_rank": [ + { + "encode_ms": 1.0579520463943481, + "diffuse_ms": 22986.931640625, + "decode_ms": 7.357728004455566, + "finalize_ms": 417.3621826171875, + "total_ms": 23412.709503293037, + "total_ms_wo_finalize": 22995.34732067585, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.470703125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 23413.542149588466, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 2158.8827166706324, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.0879679918289185, + "diffuse_ms": 1710.7734375, + "decode_ms": 7.348320007324219, + "finalize_ms": 438.1396789550781 + }, + "per_rank": [ + { + "encode_ms": 1.0879679918289185, + "diffuse_ms": 1710.7734375, + "decode_ms": 7.348320007324219, + "finalize_ms": 438.1396789550781, + "total_ms": 2157.3494044542313, + "total_ms_wo_finalize": 1719.2097254991531, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2158.8827166706324, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 2159.7884446382523, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8655999898910522, + "diffuse_ms": 1720.193115234375, + "decode_ms": 7.499231815338135, + "finalize_ms": 430.6153869628906 + }, + "per_rank": [ + { + "encode_ms": 0.8655999898910522, + "diffuse_ms": 1720.193115234375, + "decode_ms": 7.499231815338135, + "finalize_ms": 430.6153869628906, + "total_ms": 2159.173334002495, + "total_ms_wo_finalize": 1728.5579470396042, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2159.7884446382523, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 2158.867284655571, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8133440017700195, + "diffuse_ms": 1724.6712646484375, + "decode_ms": 7.512127876281738, + "finalize_ms": 425.1802978515625 + }, + "per_rank": [ + { + "encode_ms": 0.8133440017700195, + "diffuse_ms": 1724.6712646484375, + "decode_ms": 7.512127876281738, + "finalize_ms": 425.1802978515625, + "total_ms": 2158.1770343780518, + "total_ms_wo_finalize": 1732.9967365264893, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2158.867284655571, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 2169.4862078875303, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.810368001461029, + "diffuse_ms": 1735.7138671875, + "decode_ms": 7.494592189788818, + "finalize_ms": 424.867431640625 + }, + "per_rank": [ + { + "encode_ms": 0.810368001461029, + "diffuse_ms": 1735.7138671875, + "decode_ms": 7.494592189788818, + "finalize_ms": 424.867431640625, + "total_ms": 2168.886259019375, + "total_ms_wo_finalize": 1744.0188273787498, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2169.4862078875303, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 2173.339394852519, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7985600233078003, + "diffuse_ms": 1737.7554931640625, + "decode_ms": 7.436384201049805, + "finalize_ms": 426.65533447265625 + }, + "per_rank": [ + { + "encode_ms": 0.7985600233078003, + "diffuse_ms": 1737.7554931640625, + "decode_ms": 7.436384201049805, + "finalize_ms": 426.65533447265625, + "total_ms": 2172.6457718610764, + "total_ms_wo_finalize": 1745.99043738842, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2173.339394852519, + "output_frames": 12, + "rank": 0 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.04680454730987549, + "bandwidth_gbps": 22.40329327528303 + } + ], + "all_gather": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.03220513463020325, + "bandwidth_gbps": 32.5592801284738 + } + ] + } + }, + { + "environment": { + "command": "env TORCHINDUCTOR_COMPILE_THREADS=4 uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --warmup-blocks 6 --measured-blocks 5 --pixel-height 464 --pixel-width 832 --fps 16 --cp-method ulysses --bandwidth-probe-mib 1 --bandwidth-probe-iters 1 --comparison-json /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/no-comparison.json --replica-id 7 --measurement-barrier-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/barrier --output-dir /lustre/fsw/portfolios/healthcareeng/projects/healthcareeng_computervision/users/gtong/flashdreams-dist/outputs/lingbot_aggregated_8xcp1/worker-7", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01151", + "slurm_job_id": "14793417", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0 + ], + "dit_cp_group": [ + 0 + ], + "cp_size": 1, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42 + ], + "replica_id": 7, + "measurement_window": { + "started_at": 11145874.429074887, + "finished_at": 11145885.175903356, + "elapsed_s": 10.746828468516469 + } + }, + "summary": { + "fps": 5.584655369399949, + "latency_ms": { + "median": 2148.4167370945215, + "p90": 2151.7272770404816, + "min": 2144.309898838401, + "max": 2151.90726518631 + }, + "encoder_ms": { + "median": 0.8322880268096924, + "p90": 1.0315904140472412, + "min": 0.7863680124282837, + "max": 1.0977280139923096 + }, + "dit_ms": { + "median": 1717.7955322265625, + "p90": 1720.198095703125, + "min": 1704.934814453125, + "max": 1721.02099609375 + }, + "decoder_ms": { + "median": 7.437119960784912, + "p90": 7.498579216003418, + "min": 7.367008209228516, + "max": 7.50710391998291 + }, + "finalize_ms": { + "median": 424.7770690917969, + "p90": 431.455810546875, + "min": 420.665771484375, + "max": 432.703369140625 + }, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 2105.4150742637808, + "cp_probe_gbps": { + "broadcast": { + "median": 23.612682077988, + "p90": 23.612682077988, + "min": 23.612682077988, + "max": 23.612682077988 + }, + "all_gather": { + "median": 27.606411995945077, + "p90": 27.606411995945077, + "min": 27.606411995945077, + "max": 27.606411995945077 + } + }, + "memory": { + "peak_gib_by_rank": [ + 59.35440921783447 + ], + "steady_allocated_gib_by_rank": [ + 57.14538335800171 + ], + "initialization_peak_gib_by_rank": [ + 66.54739189147949 + ], + "node_peak_gib": 59.35440921783447, + "node_steady_allocated_gib": 57.14538335800171, + "per_rank_peak_gib": { + "median": 59.35440921783447, + "p90": 59.35440921783447, + "min": 59.35440921783447, + "max": 59.35440921783447 + } + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 162844.9953533709, + "output_frames": 9, + "critical_rank": { + "encode_ms": 592.1774291992188, + "diffuse_ms": 101462.015625, + "decode_ms": 60524.4921875, + "finalize_ms": 264.7309875488281 + }, + "per_rank": [ + { + "encode_ms": 592.1774291992188, + "diffuse_ms": 101462.015625, + "decode_ms": 60524.4921875, + "finalize_ms": 264.7309875488281, + "total_ms": 162843.41622924805, + "total_ms_wo_finalize": 162578.68524169922, + "mem_alloc_gib": 57.087345123291016, + "mem_reserved_gib": 61.439453125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 162844.9953533709, + "output_frames": 9, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 79477.71265730262, + "output_frames": 12, + "critical_rank": { + "encode_ms": 181.90762329101562, + "diffuse_ms": 78975.09375, + "decode_ms": 8.219840049743652, + "finalize_ms": 311.2512512207031 + }, + "per_rank": [ + { + "encode_ms": 181.90762329101562, + "diffuse_ms": 78975.09375, + "decode_ms": 8.219840049743652, + "finalize_ms": 311.2512512207031, + "total_ms": 79476.47246456146, + "total_ms_wo_finalize": 79165.22121334076, + "mem_alloc_gib": 57.08803939819336, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 79477.71265730262, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 1865.990050137043, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.68588256835938, + "diffuse_ms": 1364.6890869140625, + "decode_ms": 7.480703830718994, + "finalize_ms": 329.5069580078125 + }, + "per_rank": [ + { + "encode_ms": 163.68588256835938, + "diffuse_ms": 1364.6890869140625, + "decode_ms": 7.480703830718994, + "finalize_ms": 329.5069580078125, + "total_ms": 1865.3626313209534, + "total_ms_wo_finalize": 1535.8556733131409, + "mem_alloc_gib": 57.08770561218262, + "mem_reserved_gib": 61.857421875, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 1865.990050137043, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 2023.2099127024412, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.17910766601562, + "diffuse_ms": 1482.314208984375, + "decode_ms": 19.14019203186035, + "finalize_ms": 356.9726257324219 + }, + "per_rank": [ + { + "encode_ms": 164.17910766601562, + "diffuse_ms": 1482.314208984375, + "decode_ms": 19.14019203186035, + "finalize_ms": 356.9726257324219, + "total_ms": 2022.6061344146729, + "total_ms_wo_finalize": 1665.633508682251, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.423828125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2023.2099127024412, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 2183.0479446798563, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.3009033203125, + "diffuse_ms": 1606.981689453125, + "decode_ms": 7.380767822265625, + "finalize_ms": 403.021240234375 + }, + "per_rank": [ + { + "encode_ms": 164.3009033203125, + "diffuse_ms": 1606.981689453125, + "decode_ms": 7.380767822265625, + "finalize_ms": 403.021240234375, + "total_ms": 2181.684600830078, + "total_ms_wo_finalize": 1778.6633605957031, + "mem_alloc_gib": 57.144843101501465, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2183.0479446798563, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 2164.208112284541, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9376959800720215, + "diffuse_ms": 1718.5064697265625, + "decode_ms": 7.4245758056640625, + "finalize_ms": 436.6324157714844 + }, + "per_rank": [ + { + "encode_ms": 0.9376959800720215, + "diffuse_ms": 1718.5064697265625, + "decode_ms": 7.4245758056640625, + "finalize_ms": 436.6324157714844, + "total_ms": 2163.501157283783, + "total_ms_wo_finalize": 1726.8687415122986, + "mem_alloc_gib": 57.14517688751221, + "mem_reserved_gib": 62.42578125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2164.208112284541, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 2147.633533924818, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.0977280139923096, + "diffuse_ms": 1704.934814453125, + "decode_ms": 7.367008209228516, + "finalize_ms": 432.703369140625 + }, + "per_rank": [ + { + "encode_ms": 1.0977280139923096, + "diffuse_ms": 1704.934814453125, + "decode_ms": 7.367008209228516, + "finalize_ms": 432.703369140625, + "total_ms": 2146.102919816971, + "total_ms_wo_finalize": 1713.3995506763458, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2147.633533924818, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 2144.309898838401, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9323840141296387, + "diffuse_ms": 1705.7056884765625, + "decode_ms": 7.50710391998291, + "finalize_ms": 429.58447265625 + }, + "per_rank": [ + { + "encode_ms": 0.9323840141296387, + "diffuse_ms": 1705.7056884765625, + "decode_ms": 7.50710391998291, + "finalize_ms": 429.58447265625, + "total_ms": 2143.729649066925, + "total_ms_wo_finalize": 1714.145176410675, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2144.309898838401, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 2151.457294821739, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8322880268096924, + "diffuse_ms": 1717.7955322265625, + "decode_ms": 7.48579216003418, + "finalize_ms": 424.7770690917969 + }, + "per_rank": [ + { + "encode_ms": 0.8322880268096924, + "diffuse_ms": 1717.7955322265625, + "decode_ms": 7.48579216003418, + "finalize_ms": 424.7770690917969, + "total_ms": 2150.8906815052032, + "total_ms_wo_finalize": 1726.1136124134064, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2151.457294821739, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 2148.4167370945215, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8064320087432861, + "diffuse_ms": 1718.9637451171875, + "decode_ms": 7.437119960784912, + "finalize_ms": 420.665771484375 + }, + "per_rank": [ + { + "encode_ms": 0.8064320087432861, + "diffuse_ms": 1718.9637451171875, + "decode_ms": 7.437119960784912, + "finalize_ms": 420.665771484375, + "total_ms": 2147.8730685710907, + "total_ms_wo_finalize": 1727.2072970867157, + "mem_alloc_gib": 57.14571714401245, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2148.4167370945215, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 2151.90726518631, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7863680124282837, + "diffuse_ms": 1721.02099609375, + "decode_ms": 7.382271766662598, + "finalize_ms": 422.1673583984375 + }, + "per_rank": [ + { + "encode_ms": 0.7863680124282837, + "diffuse_ms": 1721.02099609375, + "decode_ms": 7.382271766662598, + "finalize_ms": 422.1673583984375, + "total_ms": 2151.3569942712784, + "total_ms_wo_finalize": 1729.1896358728409, + "mem_alloc_gib": 57.14538335800171, + "mem_reserved_gib": 62.580078125, + "mem_peak_gib": 59.35440921783447, + "wall_ms": 2151.90726518631, + "output_frames": 12, + "rank": 0 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.04440732300281525, + "bandwidth_gbps": 23.612682077988 + } + ], + "all_gather": [ + { + "payload_bytes": 1048576.0, + "transfer_ms": 0.037983059883117676, + "bandwidth_gbps": 27.606411995945077 + } + ] + } + } + ] +} diff --git a/integrations/lingbot/docs/benchmark_h100_aggregated_cp1/README.md b/integrations/lingbot/docs/benchmark_h100_aggregated_cp1/README.md new file mode 100644 index 00000000..8f7206cf --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_aggregated_cp1/README.md @@ -0,0 +1,66 @@ +# LingBot fully aggregated single-H100 benchmark + +The complete LingBot encoder, DiT, and LightTAE decoder ran in one process on +one H100 80 GB. There are no RDMA stage boundaries and CP1 has no inter-GPU +attention communication. + +## Result + +| Metric | Median | P90 | +| --- | ---: | ---: | +| End-to-end 12-frame chunk latency | **2157.51 ms** | **2166.25 ms** | +| Encoder | 0.88 ms | 0.94 ms | +| DiT denoise | 1726.27 ms | 1733.16 ms | +| DiT cache finalize | 424.76 ms | 425.28 ms | +| Decoder | 7.39 ms | 7.49 ms | + +- Generated throughput: **5.56 FPS** +- Initialization peak allocated HBM: **66.55 GiB** +- Measured-rollout peak allocated HBM: **59.36 GiB** +- Steady allocated HBM after rollout: **57.15 GiB** + +The five measured chunks were 2155.13–2171.30 ms and each emitted 12 frames. +Six preceding chunks were excluded for model compilation, autotuning, cache +fill, and the block-5 cache-shape transition. Inductor rejected several Triton +autotuning candidates that exceeded H100 shared-memory resources and selected +valid fallback kernels; the pipeline did not encounter an HBM out-of-memory +failure. + +Compared with the tracked three-GPU, stage-disaggregated CP1 result at the same +832×464 shape, aggregation improved FPS from 5.36 to 5.56 and reduced median +latency from 2233.57 to 2157.51 ms. Removing stage handoffs therefore saved +76.05 ms, or 3.4%, but did not materially change the DiT-dominated latency. + +## Reproduction + +```bash +./srun.sh + +cd /path/to/flashdreams +mkdir -p outputs/lingbot_aggregated_cp1 + +CUDA_VISIBLE_DEVICES=0 TORCHINDUCTOR_COMPILE_THREADS=4 GLOG_minloglevel=2 \ +uv run --package flashdreams-lingbot torchrun \ + --standalone --nproc_per_node=1 \ + -m lingbot.disagg.benchmark_aggregated \ + --model lingbot-world-fast-taehv-window15-sink3 \ + --example-idx 0 \ + --pixel-width 832 --pixel-height 464 --fps 16 \ + --cp-method ulysses \ + --warmup-blocks 6 --measured-blocks 5 \ + --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 \ + --comparison-json outputs/no-comparison.json \ + --output-dir outputs/lingbot_aggregated_cp1 +``` + +Environment: + +- Repository base: `0c2d48a8249577fb617bb5280208dd77409d9b1a`, plus the CP1 benchmark-harness change +- Slurm: job `14790020`, node `pool0-01083` +- GPU: one NVIDIA H100 80 GB HBM3 +- PyTorch 2.12.1+cu130, CUDA 13.0, cuDNN 9.2, driver 535.216.03 +- BF16, seed 42, four diffusion steps, window 15, sink 3 +- Checkpoint: `robbyant/lingbot-world-fast` + +The machine-readable result, including all warmup and measured records, is in +[`benchmark.json`](benchmark.json). diff --git a/integrations/lingbot/docs/benchmark_h100_aggregated_cp1/benchmark.json b/integrations/lingbot/docs/benchmark_h100_aggregated_cp1/benchmark.json new file mode 100644 index 00000000..9f86e6c2 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_aggregated_cp1/benchmark.json @@ -0,0 +1,530 @@ +{ + "environment": { + "command": "env TORCHINDUCTOR_COMPILE_THREADS=4 uv run --package flashdreams-lingbot torchrun --standalone --nproc_per_node=1 -m lingbot.disagg.benchmark_aggregated --model lingbot-world-fast-taehv-window15-sink3 --example-idx 0 --pixel-width 832 --pixel-height 464 --fps 16 --cp-method ulysses --warmup-blocks 6 --measured-blocks 5 --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 --comparison-json outputs/no-comparison.json --output-dir outputs/lingbot_aggregated_cp1", + "commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "worktree_dirty": true, + "hostname": "pool0-01083", + "slurm_job_id": "14790020", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0 + ], + "dit_cp_group": [ + 0 + ], + "cp_size": 1, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42 + ] + }, + "summary": { + "fps": 5.55553754894227, + "latency_ms": { + "median": 2157.513060141355, + "p90": 2166.248793900013, + "min": 2155.1325689069927, + "max": 2171.2969318032265 + }, + "encoder_ms": { + "median": 0.8756800293922424, + "p90": 0.9436095952987671, + "min": 0.770687997341156, + "max": 0.9571520090103149 + }, + "dit_ms": { + "median": 1726.2662353515625, + "p90": 1733.1642333984375, + "min": 1723.576171875, + "max": 1736.1300048828125 + }, + "decoder_ms": { + "median": 7.390175819396973, + "p90": 7.4898113250732425, + "min": 7.323999881744385, + "max": 7.543392181396484 + }, + "finalize_ms": { + "median": 424.7573547363281, + "p90": 425.27886962890625, + "min": 419.94976806640625, + "max": 425.31866455078125 + }, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 2094.4376559512357, + "cp_probe_gbps": { + "broadcast": { + "median": 10817.623328068143, + "p90": 11321.74250011288, + "min": 6100.113780988609, + "max": 11366.671641593679 + }, + "all_gather": { + "median": 1204.2908306532627, + "p90": 1208.7825052237476, + "min": 1179.869729222284, + "max": 1209.984367372116 + } + }, + "memory": { + "peak_gib_by_rank": [ + 59.3552770614624 + ], + "steady_allocated_gib_by_rank": [ + 57.14636754989624 + ], + "initialization_peak_gib_by_rank": [ + 66.54739189147949 + ], + "node_peak_gib": 59.3552770614624, + "node_steady_allocated_gib": 57.14636754989624, + "per_rank_peak_gib": { + "median": 59.3552770614624, + "p90": 59.3552770614624, + "min": 59.3552770614624, + "max": 59.3552770614624 + } + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 123287.52873605117, + "output_frames": 9, + "critical_rank": { + "encode_ms": 490.09832763671875, + "diffuse_ms": 87463.1171875, + "decode_ms": 35063.96875, + "finalize_ms": 268.6892395019531 + }, + "per_rank": [ + { + "encode_ms": 490.09832763671875, + "diffuse_ms": 87463.1171875, + "decode_ms": 35063.96875, + "finalize_ms": 268.6892395019531, + "total_ms": 123285.87350463867, + "total_ms_wo_finalize": 123017.18426513672, + "mem_alloc_gib": 57.08832931518555, + "mem_reserved_gib": 61.46875, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 123287.52873605117, + "output_frames": 9, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 78228.68528589606, + "output_frames": 12, + "critical_rank": { + "encode_ms": 182.13430786132812, + "diffuse_ms": 77730.203125, + "decode_ms": 8.143360137939453, + "finalize_ms": 308.2727966308594 + }, + "per_rank": [ + { + "encode_ms": 182.13430786132812, + "diffuse_ms": 77730.203125, + "decode_ms": 8.143360137939453, + "finalize_ms": 308.2727966308594, + "total_ms": 78228.75358963013, + "total_ms_wo_finalize": 77920.48079299927, + "mem_alloc_gib": 57.08878707885742, + "mem_reserved_gib": 61.861328125, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 78228.68528589606, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 1868.6057380400598, + "output_frames": 12, + "critical_rank": { + "encode_ms": 163.92034912109375, + "diffuse_ms": 1368.20361328125, + "decode_ms": 7.503551959991455, + "finalize_ms": 328.4642333984375 + }, + "per_rank": [ + { + "encode_ms": 163.92034912109375, + "diffuse_ms": 1368.20361328125, + "decode_ms": 7.503551959991455, + "finalize_ms": 328.4642333984375, + "total_ms": 1868.0917477607727, + "total_ms_wo_finalize": 1539.6275143623352, + "mem_alloc_gib": 57.08868980407715, + "mem_reserved_gib": 61.861328125, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 1868.6057380400598, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 2028.643268160522, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.41754150390625, + "diffuse_ms": 1478.8184814453125, + "decode_ms": 20.656095504760742, + "finalize_ms": 364.23333740234375 + }, + "per_rank": [ + { + "encode_ms": 164.41754150390625, + "diffuse_ms": 1478.8184814453125, + "decode_ms": 20.656095504760742, + "finalize_ms": 364.23333740234375, + "total_ms": 2028.1254558563232, + "total_ms_wo_finalize": 1663.8921184539795, + "mem_alloc_gib": 57.14592456817627, + "mem_reserved_gib": 62.328125, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 2028.643268160522, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 2172.1249939873815, + "output_frames": 12, + "critical_rank": { + "encode_ms": 164.94435119628906, + "diffuse_ms": 1594.80859375, + "decode_ms": 7.348447799682617, + "finalize_ms": 404.4841613769531 + }, + "per_rank": [ + { + "encode_ms": 164.94435119628906, + "diffuse_ms": 1594.80859375, + "decode_ms": 7.348447799682617, + "finalize_ms": 404.4841613769531, + "total_ms": 2171.585554122925, + "total_ms_wo_finalize": 1767.1013927459717, + "mem_alloc_gib": 57.145827293395996, + "mem_reserved_gib": 62.330078125, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 2172.1249939873815, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 2166.6380460374057, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8436160087585449, + "diffuse_ms": 1718.6199951171875, + "decode_ms": 7.415999889373779, + "finalize_ms": 439.20550537109375 + }, + "per_rank": [ + { + "encode_ms": 0.8436160087585449, + "diffuse_ms": 1718.6199951171875, + "decode_ms": 7.415999889373779, + "finalize_ms": 439.20550537109375, + "total_ms": 2166.0851163864136, + "total_ms_wo_finalize": 1726.8796110153198, + "mem_alloc_gib": 57.14592456817627, + "mem_reserved_gib": 62.330078125, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 2166.6380460374057, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 2171.2969318032265, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9232959747314453, + "diffuse_ms": 1736.1300048828125, + "decode_ms": 7.543392181396484, + "finalize_ms": 425.31866455078125 + }, + "per_rank": [ + { + "encode_ms": 0.9232959747314453, + "diffuse_ms": 1736.1300048828125, + "decode_ms": 7.543392181396484, + "finalize_ms": 425.31866455078125, + "total_ms": 2169.9153575897217, + "total_ms_wo_finalize": 1744.5966930389404, + "mem_alloc_gib": 57.14636754989624, + "mem_reserved_gib": 62.333984375, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 2171.2969318032265, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 2155.1325689069927, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9571520090103149, + "diffuse_ms": 1726.2662353515625, + "decode_ms": 7.409440040588379, + "finalize_ms": 419.94976806640625 + }, + "per_rank": [ + { + "encode_ms": 0.9571520090103149, + "diffuse_ms": 1726.2662353515625, + "decode_ms": 7.409440040588379, + "finalize_ms": 419.94976806640625, + "total_ms": 2154.5825954675674, + "total_ms_wo_finalize": 1734.6328274011612, + "mem_alloc_gib": 57.146464824676514, + "mem_reserved_gib": 62.5078125, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 2155.1325689069927, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 2158.6765870451927, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8756800293922424, + "diffuse_ms": 1728.715576171875, + "decode_ms": 7.328832149505615, + "finalize_ms": 421.2311096191406 + }, + "per_rank": [ + { + "encode_ms": 0.8756800293922424, + "diffuse_ms": 1728.715576171875, + "decode_ms": 7.328832149505615, + "finalize_ms": 421.2311096191406, + "total_ms": 2158.1511979699135, + "total_ms_wo_finalize": 1736.9200883507729, + "mem_alloc_gib": 57.14636754989624, + "mem_reserved_gib": 62.5078125, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 2158.6765870451927, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 2157.513060141355, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.770687997341156, + "diffuse_ms": 1724.068115234375, + "decode_ms": 7.390175819396973, + "finalize_ms": 424.7573547363281 + }, + "per_rank": [ + { + "encode_ms": 0.770687997341156, + "diffuse_ms": 1724.068115234375, + "decode_ms": 7.390175819396973, + "finalize_ms": 424.7573547363281, + "total_ms": 2156.9863337874413, + "total_ms_wo_finalize": 1732.2289790511131, + "mem_alloc_gib": 57.146464824676514, + "mem_reserved_gib": 62.5078125, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 2157.513060141355, + "output_frames": 12, + "rank": 0 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 2157.4158570729196, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7902719974517822, + "diffuse_ms": 1723.576171875, + "decode_ms": 7.323999881744385, + "finalize_ms": 425.21917724609375 + }, + "per_rank": [ + { + "encode_ms": 0.7902719974517822, + "diffuse_ms": 1723.576171875, + "decode_ms": 7.323999881744385, + "finalize_ms": 425.21917724609375, + "total_ms": 2156.90962100029, + "total_ms_wo_finalize": 1731.6904437541962, + "mem_alloc_gib": 57.14636754989624, + "mem_reserved_gib": 62.5078125, + "mem_peak_gib": 59.3552770614624, + "wall_ms": 2157.4158570729196, + "output_frames": 12, + "rank": 0 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.04400499165058136, + "bandwidth_gbps": 6100.113780988609 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.03661308437585831, + "bandwidth_gbps": 7331.681025404109 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.02770824357867241, + "bandwidth_gbps": 9687.927538164857 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.02585211768746376, + "bandwidth_gbps": 10383.499690246654 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.023857224732637405, + "bandwidth_gbps": 11251.746965889632 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.023750122636556625, + "bandwidth_gbps": 11302.487153763965 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.023616012185811996, + "bandwidth_gbps": 11366.671641593679 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.02376176416873932, + "bandwidth_gbps": 11296.949759022957 + } + ], + "all_gather": [ + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.22751279175281525, + "bandwidth_gbps": 1179.869729222284 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.22351323068141937, + "bandwidth_gbps": 1200.9823990357409 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.22252416610717773, + "bandwidth_gbps": 1206.3204671024778 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.22301124036312103, + "bandwidth_gbps": 1203.685767421034 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.2227872610092163, + "bandwidth_gbps": 1204.8958938854914 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.22216559946537018, + "bandwidth_gbps": 1208.2674214458755 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.22423221170902252, + "bandwidth_gbps": 1197.1315537320675 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.22185035049915314, + "bandwidth_gbps": 1209.984367372116 + } + ] + } +} diff --git a/integrations/lingbot/docs/benchmark_h100_aggregated_cp8/README.md b/integrations/lingbot/docs/benchmark_h100_aggregated_cp8/README.md new file mode 100644 index 00000000..02f74670 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_aggregated_cp8/README.md @@ -0,0 +1,48 @@ +# LingBot aggregated CP8 benchmark + +All ranks own the complete encoder, DiT, and decoder pipeline. The DiT token +axis is context-parallel across WORLD with ulysses attention. Encoder +and decoder work is replicated on every rank; there are no RDMA stage +boundaries in this topology. + +## Result + +| Metric | Median | P90 | +| --- | ---: | ---: | +| End-to-end chunk latency | 393.33 ms | 434.08 ms | +| Encoder critical-rank compute | 0.84 ms | 1.03 ms | +| DiT critical-rank denoise | 309.14 ms | 349.85 ms | +| Decoder critical-rank compute | 6.88 ms | 6.91 ms | +| DiT cache finalize | 76.06 ms | 76.76 ms | +| NCCL broadcast probe | 266.83 GB/s | 268.38 GB/s | +| NCCL all-gather probe | 360.38 GB/s | 361.75 GB/s | + +- Generated throughput: **29.50 FPS** +- DiT token throughput: **10739 token/s** +- Peak allocated HBM: **40.88–40.88 GiB per rank**, **327.03 GiB node total** +- Steady allocated HBM after rollout: **310.37 GiB node total** + +## Comparison with disaggregated CP + +| Metric | Disaggregated CP6 | Aggregated CP8 | Change | +| --- | ---: | ---: | ---: | +| Median chunk latency | 743.27 ms | 393.33 ms | 1.89× faster | +| Generated FPS | 15.90 | 29.50 | 1.86× | +| DiT token throughput | 5995 token/s | 10739 token/s | 1.79× | +| Node peak allocated HBM | 251.07 GiB | 327.03 GiB | 1.30× | + +The resolutions differ because the tracked 832×464 grid has 4,524 tokens, +which is not divisible by eight. CP8 uses 832×448 and 4,368 tokens (3.45% +fewer). Token throughput is therefore the fairest compute-rate comparison. + +## Reproduction + +```bash +env 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 --example-idx 0 --pixel-width 832 --pixel-height 448 --fps 16 --warmup-blocks 6 --measured-blocks 5 --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 --comparison-json integrations/lingbot/docs/benchmark_h100_cp6_single_session/benchmark.json --output-dir outputs/lingbot_aggregated_cp8 +``` + +- Repository revision: `bb67d2868babea53cda7a9027831f26ee948293f` (modified worktree) +- Slurm: job `14652956` on `pool0-01714` +- GPU: `NVIDIA H100 80GB HBM3` × 8 +- Resolution: `832x448` +- Warmup / measured blocks: 6 / 5 diff --git a/integrations/lingbot/docs/benchmark_h100_aggregated_cp8/benchmark.json b/integrations/lingbot/docs/benchmark_h100_aggregated_cp8/benchmark.json new file mode 100644 index 00000000..66fe6085 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_aggregated_cp8/benchmark.json @@ -0,0 +1,1673 @@ +{ + "environment": { + "command": "env 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 --example-idx 0 --pixel-width 832 --pixel-height 448 --fps 16 --warmup-blocks 6 --measured-blocks 5 --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 --comparison-json integrations/lingbot/docs/benchmark_h100_cp6_single_session/benchmark.json --output-dir outputs/lingbot_aggregated_cp8", + "commit": "bb67d2868babea53cda7a9027831f26ee948293f", + "worktree_dirty": true, + "hostname": "pool0-01714", + "slurm_job_id": "14652956", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 448, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "full_pipeline_replicas": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7 + ], + "dit_cp_group": [ + 0, + 1, + 2, + 3, + 4, + 5, + 6, + 7 + ], + "cp_size": 8, + "cp_method": "ulysses" + }, + "token_layout": { + "latent_height": 56, + "latent_width": 104, + "tokens_per_chunk": 4368, + "tokens_per_rank": 546 + }, + "transport": { + "stage_handoffs": "none", + "dit_collectives": "NCCL" + }, + "noise_seed_by_rank": [ + 42, + 43, + 44, + 45, + 46, + 47, + 48, + 49 + ] + }, + "summary": { + "fps": 29.503035463925226, + "latency_ms": { + "median": 393.33291398361325, + "p90": 434.0799169614911, + "min": 392.140940297395, + "max": 460.55883821099997 + }, + "encoder_ms": { + "median": 0.839136004447937, + "p90": 1.0282432079315185, + "min": 0.77183997631073, + "max": 1.0959680080413818 + }, + "dit_ms": { + "median": 309.1429138183594, + "p90": 349.8536071777344, + "min": 307.56024169921875, + "max": 376.3791809082031 + }, + "decoder_ms": { + "median": 6.880799770355225, + "p90": 6.912819290161133, + "min": 6.855487823486328, + "max": 6.922368049621582 + }, + "finalize_ms": { + "median": 76.05542755126953, + "p90": 76.7609146118164, + "min": 75.66960144042969, + "max": 77.17052459716797 + }, + "tokens_per_chunk": 4368, + "token_throughput_per_second": 10739.104908868781, + "cp_probe_gbps": { + "broadcast": { + "median": 266.82638995992966, + "p90": 268.38334320650205, + "min": 264.2735748919249, + "max": 270.44839017080693 + }, + "all_gather": { + "median": 360.3832523638345, + "p90": 361.7482704350743, + "min": 354.0439658375764, + "max": 362.0113760268079 + } + }, + "memory": { + "peak_gib_by_rank": [ + 40.879210472106934, + 40.879210472106934, + 40.879210472106934, + 40.879210472106934, + 40.879210472106934, + 40.879210472106934, + 40.879210472106934, + 40.879210472106934 + ], + "steady_allocated_gib_by_rank": [ + 38.795698165893555, + 38.795698165893555, + 38.795698165893555, + 38.795698165893555, + 38.795698165893555, + 38.795698165893555, + 38.795698165893555, + 38.795698165893555 + ], + "initialization_peak_gib_by_rank": [ + 48.26605987548828, + 48.26605987548828, + 48.26605987548828, + 48.26605987548828, + 48.26605987548828, + 48.26605987548828, + 48.26605987548828, + 48.26605987548828 + ], + "node_peak_gib": 327.03368377685547, + "node_steady_allocated_gib": 310.36558532714844, + "per_rank_peak_gib": { + "median": 40.879210472106934, + "p90": 40.879210472106934, + "min": 40.879210472106934, + "max": 40.879210472106934 + } + }, + "comparison": { + "topology": "1 encoder : CP6 DiT : 1 decoder", + "resolution": [ + 464, + 832 + ], + "fps": 15.901998021609447, + "latency_ms": 743.272824001906, + "tokens_per_chunk": 4524, + "token_throughput_per_second": 5995.053254146762, + "node_peak_gib": 251.07237005233765, + "latency_speedup": 1.8896786858596626, + "fps_ratio": 1.8553036809483399, + "token_throughput_ratio": 1.791327691950121, + "node_peak_memory_ratio": 1.3025474834554005 + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "end_to_end_ms": 25475.266749039292, + "output_frames": 9, + "critical_rank": { + "encode_ms": 328.1184997558594, + "diffuse_ms": 23222.494140625, + "decode_ms": 1902.6649169921875, + "finalize_ms": 80.23884582519531 + }, + "per_rank": [ + { + "encode_ms": 303.47943115234375, + "diffuse_ms": 23222.494140625, + "decode_ms": 1868.583740234375, + "finalize_ms": 80.23884582519531, + "total_ms": 25474.796157836914, + "total_ms_wo_finalize": 25394.55731201172, + "mem_alloc_gib": 38.739038944244385, + "mem_reserved_gib": 42.927734375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 25475.266749039292, + "output_frames": 9, + "rank": 0 + }, + { + "encode_ms": 328.1184997558594, + "diffuse_ms": 23197.83984375, + "decode_ms": 1902.6649169921875, + "finalize_ms": 46.160160064697266, + "total_ms": 25474.783420562744, + "total_ms_wo_finalize": 25428.623260498047, + "mem_alloc_gib": 38.739038944244385, + "mem_reserved_gib": 42.927734375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 25475.130709819496, + "output_frames": 9, + "rank": 1 + }, + { + "encode_ms": 307.8042297363281, + "diffuse_ms": 23218.111328125, + "decode_ms": 1902.533203125, + "finalize_ms": 46.28569412231445, + "total_ms": 25474.734455108643, + "total_ms_wo_finalize": 25428.448760986328, + "mem_alloc_gib": 38.739038944244385, + "mem_reserved_gib": 42.927734375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 25475.13401461765, + "output_frames": 9, + "rank": 2 + }, + { + "encode_ms": 304.86785888671875, + "diffuse_ms": 23221.158203125, + "decode_ms": 1872.94482421875, + "finalize_ms": 75.8653793334961, + "total_ms": 25474.836265563965, + "total_ms_wo_finalize": 25398.97088623047, + "mem_alloc_gib": 38.739038944244385, + "mem_reserved_gib": 42.927734375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 25475.137917790562, + "output_frames": 9, + "rank": 3 + }, + { + "encode_ms": 326.9462890625, + "diffuse_ms": 23199.01171875, + "decode_ms": 1891.94873046875, + "finalize_ms": 56.852928161621094, + "total_ms": 25474.75966644287, + "total_ms_wo_finalize": 25417.90673828125, + "mem_alloc_gib": 38.739038944244385, + "mem_reserved_gib": 42.927734375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 25475.163844879717, + "output_frames": 9, + "rank": 4 + }, + { + "encode_ms": 304.2637634277344, + "diffuse_ms": 23221.640625, + "decode_ms": 1896.371337890625, + "finalize_ms": 52.43507385253906, + "total_ms": 25474.7108001709, + "total_ms_wo_finalize": 25422.27572631836, + "mem_alloc_gib": 38.739038944244385, + "mem_reserved_gib": 42.927734375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 25475.138308946043, + "output_frames": 9, + "rank": 5 + }, + { + "encode_ms": 304.21649169921875, + "diffuse_ms": 23221.71875, + "decode_ms": 1901.421875, + "finalize_ms": 47.3870735168457, + "total_ms": 25474.744190216064, + "total_ms_wo_finalize": 25427.35711669922, + "mem_alloc_gib": 38.739038944244385, + "mem_reserved_gib": 42.927734375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 25475.128335878253, + "output_frames": 9, + "rank": 6 + }, + { + "encode_ms": 305.29571533203125, + "diffuse_ms": 23220.646484375, + "decode_ms": 1893.497314453125, + "finalize_ms": 55.31951904296875, + "total_ms": 25474.759033203125, + "total_ms_wo_finalize": 25419.439514160156, + "mem_alloc_gib": 38.739038944244385, + "mem_reserved_gib": 42.927734375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 25475.21718405187, + "output_frames": 9, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 1, + "warmup": true, + "end_to_end_ms": 16657.64493122697, + "output_frames": 12, + "critical_rank": { + "encode_ms": 198.2529296875, + "diffuse_ms": 16404.859375, + "decode_ms": 6.991583824157715, + "finalize_ms": 53.74016189575195 + }, + "per_rank": [ + { + "encode_ms": 196.2623748779297, + "diffuse_ms": 16400.19921875, + "decode_ms": 6.966911792755127, + "finalize_ms": 53.68182373046875, + "total_ms": 16657.110329151154, + "total_ms_wo_finalize": 16603.428505420685, + "mem_alloc_gib": 38.73953199386597, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 16657.49899391085, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 197.6308135986328, + "diffuse_ms": 16398.802734375, + "decode_ms": 6.986368179321289, + "finalize_ms": 53.62713623046875, + "total_ms": 16657.047052383423, + "total_ms_wo_finalize": 16603.419916152954, + "mem_alloc_gib": 38.73953199386597, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 16657.473276834935, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 192.8092498779297, + "diffuse_ms": 16403.603515625, + "decode_ms": 6.975232124328613, + "finalize_ms": 53.66969680786133, + "total_ms": 16657.05769443512, + "total_ms_wo_finalize": 16603.38799762726, + "mem_alloc_gib": 38.73953199386597, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 16657.41445729509, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 191.61875915527344, + "diffuse_ms": 16404.859375, + "decode_ms": 6.932799816131592, + "finalize_ms": 53.712223052978516, + "total_ms": 16657.123157024384, + "total_ms_wo_finalize": 16603.410933971405, + "mem_alloc_gib": 38.73953199386597, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 16657.64493122697, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 197.7554168701172, + "diffuse_ms": 16398.66796875, + "decode_ms": 6.924255847930908, + "finalize_ms": 53.71955108642578, + "total_ms": 16657.067192554474, + "total_ms_wo_finalize": 16603.347641468048, + "mem_alloc_gib": 38.73953199386597, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 16657.437787391245, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 198.2529296875, + "diffuse_ms": 16398.154296875, + "decode_ms": 6.948319911956787, + "finalize_ms": 53.7061767578125, + "total_ms": 16657.06172323227, + "total_ms_wo_finalize": 16603.355546474457, + "mem_alloc_gib": 38.73953199386597, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 16657.447980251163, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 192.0247344970703, + "diffuse_ms": 16404.423828125, + "decode_ms": 6.991583824157715, + "finalize_ms": 53.65951919555664, + "total_ms": 16657.099665641785, + "total_ms_wo_finalize": 16603.440146446228, + "mem_alloc_gib": 38.73953199386597, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 16657.427730038762, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 197.27430725097656, + "diffuse_ms": 16399.150390625, + "decode_ms": 6.902048110961914, + "finalize_ms": 53.74016189575195, + "total_ms": 16657.06690788269, + "total_ms_wo_finalize": 16603.32674598694, + "mem_alloc_gib": 38.73953199386597, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 16657.48513210565, + "output_frames": 12, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 2, + "warmup": true, + "end_to_end_ms": 494.97264716774225, + "output_frames": 12, + "critical_rank": { + "encode_ms": 158.442626953125, + "diffuse_ms": 269.3844299316406, + "decode_ms": 6.998688220977783, + "finalize_ms": 60.06905746459961 + }, + "per_rank": [ + { + "encode_ms": 158.11622619628906, + "diffuse_ms": 269.19219970703125, + "decode_ms": 6.930272102355957, + "finalize_ms": 60.05129623413086, + "total_ms": 494.28999423980713, + "total_ms_wo_finalize": 434.23869800567627, + "mem_alloc_gib": 38.73938703536987, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 494.97264716774225, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 158.442626953125, + "diffuse_ms": 268.8555603027344, + "decode_ms": 6.972832202911377, + "finalize_ms": 59.982303619384766, + "total_ms": 494.2533230781555, + "total_ms_wo_finalize": 434.27101945877075, + "mem_alloc_gib": 38.73938703536987, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 494.89049101248384, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 158.37353515625, + "diffuse_ms": 268.93408203125, + "decode_ms": 6.998688220977783, + "finalize_ms": 59.981536865234375, + "total_ms": 494.28784227371216, + "total_ms_wo_finalize": 434.3063054084778, + "mem_alloc_gib": 38.73938703536987, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 494.78159798309207, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 158.12265014648438, + "diffuse_ms": 269.173828125, + "decode_ms": 6.917503833770752, + "finalize_ms": 60.06905746459961, + "total_ms": 494.28303956985474, + "total_ms_wo_finalize": 434.2139821052551, + "mem_alloc_gib": 38.73938703536987, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 494.91724325343966, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 158.0963897705078, + "diffuse_ms": 269.2049255371094, + "decode_ms": 6.935967922210693, + "finalize_ms": 60.04390335083008, + "total_ms": 494.28118658065796, + "total_ms_wo_finalize": 434.2372832298279, + "mem_alloc_gib": 38.73938703536987, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 494.79188583791256, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 158.16925048828125, + "diffuse_ms": 269.12738037109375, + "decode_ms": 6.93228816986084, + "finalize_ms": 60.06175994873047, + "total_ms": 494.2906789779663, + "total_ms_wo_finalize": 434.22891902923584, + "mem_alloc_gib": 38.73938703536987, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 494.87607926130295, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 158.37799072265625, + "diffuse_ms": 268.926513671875, + "decode_ms": 6.969696044921875, + "finalize_ms": 60.01612854003906, + "total_ms": 494.2903289794922, + "total_ms_wo_finalize": 434.2742004394531, + "mem_alloc_gib": 38.73938703536987, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 494.8704931885004, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 157.90512084960938, + "diffuse_ms": 269.3844299316406, + "decode_ms": 6.944704055786133, + "finalize_ms": 60.03539276123047, + "total_ms": 494.2696475982666, + "total_ms_wo_finalize": 434.23425483703613, + "mem_alloc_gib": 38.73938703536987, + "mem_reserved_gib": 42.55859375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 494.75579615682364, + "output_frames": 12, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 3, + "warmup": true, + "end_to_end_ms": 583.7101927027106, + "output_frames": 12, + "critical_rank": { + "encode_ms": 158.21615600585938, + "diffuse_ms": 299.8464660644531, + "decode_ms": 57.398624420166016, + "finalize_ms": 76.19673919677734 + }, + "per_rank": [ + { + "encode_ms": 157.8217010498047, + "diffuse_ms": 299.8464660644531, + "decode_ms": 49.183807373046875, + "finalize_ms": 76.19673919677734, + "total_ms": 583.048713684082, + "total_ms_wo_finalize": 506.8519744873047, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.21875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 583.7101927027106, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 158.01280212402344, + "diffuse_ms": 299.6541442871094, + "decode_ms": 51.741153717041016, + "finalize_ms": 73.6022720336914, + "total_ms": 583.0103721618652, + "total_ms_wo_finalize": 509.4081001281738, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.21875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 583.5666051134467, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 158.21615600585938, + "diffuse_ms": 299.45184326171875, + "decode_ms": 57.398624420166016, + "finalize_ms": 67.98553466796875, + "total_ms": 583.0521583557129, + "total_ms_wo_finalize": 515.0666236877441, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.21875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 583.650553599, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 158.04742431640625, + "diffuse_ms": 299.6188659667969, + "decode_ms": 57.37459182739258, + "finalize_ms": 67.99798583984375, + "total_ms": 583.0388679504395, + "total_ms_wo_finalize": 515.0408821105957, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.21875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 583.5552681237459, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 157.95164489746094, + "diffuse_ms": 299.7125244140625, + "decode_ms": 57.3790397644043, + "finalize_ms": 67.99212646484375, + "total_ms": 583.0353355407715, + "total_ms_wo_finalize": 515.0432090759277, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.21875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 583.5753991268575, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 157.98252868652344, + "diffuse_ms": 299.6834411621094, + "decode_ms": 51.75603103637695, + "finalize_ms": 73.61929321289062, + "total_ms": 583.0412940979004, + "total_ms_wo_finalize": 509.42200088500977, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.21875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 583.6678519845009, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 158.0471649169922, + "diffuse_ms": 299.61871337890625, + "decode_ms": 50.561344146728516, + "finalize_ms": 74.81391906738281, + "total_ms": 583.0411415100098, + "total_ms_wo_finalize": 508.22722244262695, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.21875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 583.6512157693505, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 157.88243103027344, + "diffuse_ms": 299.7842712402344, + "decode_ms": 49.22844696044922, + "finalize_ms": 76.14179229736328, + "total_ms": 583.0369415283203, + "total_ms_wo_finalize": 506.89514923095703, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.21875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 583.5692649707198, + "output_frames": 12, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 4, + "warmup": true, + "end_to_end_ms": 558.2462050952017, + "output_frames": 12, + "critical_rank": { + "encode_ms": 158.65330505371094, + "diffuse_ms": 320.765380859375, + "decode_ms": 6.850368022918701, + "finalize_ms": 72.03926086425781 + }, + "per_rank": [ + { + "encode_ms": 158.10499572753906, + "diffuse_ms": 320.5745544433594, + "decode_ms": 6.82371187210083, + "finalize_ms": 72.01519775390625, + "total_ms": 557.5184597969055, + "total_ms_wo_finalize": 485.50326204299927, + "mem_alloc_gib": 38.795631885528564, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 558.2462050952017, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 158.19078063964844, + "diffuse_ms": 320.49066162109375, + "decode_ms": 6.844639778137207, + "finalize_ms": 71.96198272705078, + "total_ms": 557.4880647659302, + "total_ms_wo_finalize": 485.5260820388794, + "mem_alloc_gib": 38.795631885528564, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 557.9758360981941, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 158.65330505371094, + "diffuse_ms": 320.02685546875, + "decode_ms": 6.850368022918701, + "finalize_ms": 71.98486328125, + "total_ms": 557.5153918266296, + "total_ms_wo_finalize": 485.53052854537964, + "mem_alloc_gib": 38.795631885528564, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 558.2149280235171, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 157.9605712890625, + "diffuse_ms": 320.7204895019531, + "decode_ms": 6.834400177001953, + "finalize_ms": 72.00361633300781, + "total_ms": 557.5190773010254, + "total_ms_wo_finalize": 485.5154609680176, + "mem_alloc_gib": 38.795631885528564, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 558.013821952045, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 158.05711364746094, + "diffuse_ms": 320.6231994628906, + "decode_ms": 6.797567844390869, + "finalize_ms": 72.03926086425781, + "total_ms": 557.5171418190002, + "total_ms_wo_finalize": 485.47788095474243, + "mem_alloc_gib": 38.795631885528564, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 558.0291519872844, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 158.07737731933594, + "diffuse_ms": 320.5868225097656, + "decode_ms": 6.8061442375183105, + "finalize_ms": 72.03718566894531, + "total_ms": 557.5075297355652, + "total_ms_wo_finalize": 485.4703440666199, + "mem_alloc_gib": 38.795631885528564, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 558.1137449480593, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 158.02735900878906, + "diffuse_ms": 320.6547546386719, + "decode_ms": 6.833695888519287, + "finalize_ms": 72.00761413574219, + "total_ms": 557.5234236717224, + "total_ms_wo_finalize": 485.5158095359802, + "mem_alloc_gib": 38.795631885528564, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 558.0610232427716, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 157.9139862060547, + "diffuse_ms": 320.765380859375, + "decode_ms": 6.8133440017700195, + "finalize_ms": 72.02022552490234, + "total_ms": 557.512936592102, + "total_ms_wo_finalize": 485.4927110671997, + "mem_alloc_gib": 38.795631885528564, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 557.9717261716723, + "output_frames": 12, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 5, + "warmup": true, + "end_to_end_ms": 33764.66832496226, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.8915839791297913, + "diffuse_ms": 33677.94921875, + "decode_ms": 6.844160079956055, + "finalize_ms": 78.76217651367188 + }, + "per_rank": [ + { + "encode_ms": 0.879360020160675, + "diffuse_ms": 33677.91796875, + "decode_ms": 6.815936088562012, + "finalize_ms": 78.74748992919922, + "total_ms": 33764.36075478792, + "total_ms_wo_finalize": 33685.61326485872, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 33764.64233500883, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 0.8915839791297913, + "diffuse_ms": 33677.83203125, + "decode_ms": 6.835423946380615, + "finalize_ms": 78.69747161865234, + "total_ms": 33764.25651079416, + "total_ms_wo_finalize": 33685.55903917551, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 33764.61685588583, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 0.8871679902076721, + "diffuse_ms": 33677.796875, + "decode_ms": 6.844160079956055, + "finalize_ms": 78.7161636352539, + "total_ms": 33764.24436670542, + "total_ms_wo_finalize": 33685.528203070164, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 33764.54149186611, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 0.8829759955406189, + "diffuse_ms": 33677.94921875, + "decode_ms": 6.813375949859619, + "finalize_ms": 78.75424194335938, + "total_ms": 33764.39981263876, + "total_ms_wo_finalize": 33685.6455706954, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 33764.55949479714, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 0.8791040182113647, + "diffuse_ms": 33677.8359375, + "decode_ms": 6.798463821411133, + "finalize_ms": 78.76217651367188, + "total_ms": 33764.275681853294, + "total_ms_wo_finalize": 33685.51350533962, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 33764.563678298146, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 0.8814079761505127, + "diffuse_ms": 33677.78515625, + "decode_ms": 6.810495853424072, + "finalize_ms": 78.75846099853516, + "total_ms": 33764.23552107811, + "total_ms_wo_finalize": 33685.477060079575, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 33764.66832496226, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 0.8463039994239807, + "diffuse_ms": 33677.5859375, + "decode_ms": 6.83465576171875, + "finalize_ms": 78.72940826416016, + "total_ms": 33763.9963055253, + "total_ms_wo_finalize": 33685.26689726114, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 33764.288023114204, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 0.8852800130844116, + "diffuse_ms": 33677.84765625, + "decode_ms": 6.815104007720947, + "finalize_ms": 78.74201965332031, + "total_ms": 33764.290059924126, + "total_ms_wo_finalize": 33685.548040270805, + "mem_alloc_gib": 38.79577684402466, + "mem_reserved_gib": 44.220703125, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 33764.64367099106, + "output_frames": 12, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 6, + "warmup": false, + "end_to_end_ms": 460.55883821099997, + "output_frames": 12, + "critical_rank": { + "encode_ms": 1.0959680080413818, + "diffuse_ms": 376.3791809082031, + "decode_ms": 6.855487823486328, + "finalize_ms": 75.66960144042969 + }, + "per_rank": [ + { + "encode_ms": 1.0875840187072754, + "diffuse_ms": 376.3791809082031, + "decode_ms": 6.805024147033691, + "finalize_ms": 75.66960144042969, + "total_ms": 459.9413905143738, + "total_ms_wo_finalize": 384.2717890739441, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.68359375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 460.5527608655393, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 1.0926719903945923, + "diffuse_ms": 376.36676025390625, + "decode_ms": 6.855487823486328, + "finalize_ms": 75.58963012695312, + "total_ms": 459.9045501947403, + "total_ms_wo_finalize": 384.31492006778717, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.68359375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 460.4460019618273, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 1.0944639444351196, + "diffuse_ms": 376.3625183105469, + "decode_ms": 6.851424217224121, + "finalize_ms": 75.62223815917969, + "total_ms": 459.9306446313858, + "total_ms_wo_finalize": 384.3084064722061, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.68359375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 460.4845163412392, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 1.0959680080413818, + "diffuse_ms": 376.3589782714844, + "decode_ms": 6.831007957458496, + "finalize_ms": 75.64720153808594, + "total_ms": 459.9331557750702, + "total_ms_wo_finalize": 384.28595423698425, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.68359375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 460.4762657545507, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 1.0883519649505615, + "diffuse_ms": 376.3618469238281, + "decode_ms": 6.810783863067627, + "finalize_ms": 75.66006469726562, + "total_ms": 459.92104744911194, + "total_ms_wo_finalize": 384.2609827518463, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.68359375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 460.55883821099997, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 1.0828800201416016, + "diffuse_ms": 376.3715515136719, + "decode_ms": 6.809823989868164, + "finalize_ms": 75.66754913330078, + "total_ms": 459.9318046569824, + "total_ms_wo_finalize": 384.26425552368164, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.68359375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 460.4821247048676, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 1.0888320207595825, + "diffuse_ms": 376.3781433105469, + "decode_ms": 6.842656135559082, + "finalize_ms": 75.63442993164062, + "total_ms": 459.94406139850616, + "total_ms_wo_finalize": 384.30963146686554, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.68359375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 460.50783805549145, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 1.0880320072174072, + "diffuse_ms": 376.36810302734375, + "decode_ms": 6.834559917449951, + "finalize_ms": 75.63705444335938, + "total_ms": 459.9277493953705, + "total_ms_wo_finalize": 384.2906949520111, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.68359375, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 460.454027634114, + "output_frames": 12, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 7, + "warmup": false, + "end_to_end_ms": 394.3615350872278, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.9266560077667236, + "diffuse_ms": 310.06524658203125, + "decode_ms": 6.880799770355225, + "finalize_ms": 76.05542755126953 + }, + "per_rank": [ + { + "encode_ms": 0.9086719751358032, + "diffuse_ms": 310.0634460449219, + "decode_ms": 6.818175792694092, + "finalize_ms": 76.04524993896484, + "total_ms": 393.8355437517166, + "total_ms_wo_finalize": 317.79029381275177, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 394.3615350872278, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 0.9266560077667236, + "diffuse_ms": 310.0484619140625, + "decode_ms": 6.880799770355225, + "finalize_ms": 75.95174407958984, + "total_ms": 393.8076617717743, + "total_ms_wo_finalize": 317.85591769218445, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 394.29089799523354, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 0.9104639887809753, + "diffuse_ms": 310.0630798339844, + "decode_ms": 6.845151901245117, + "finalize_ms": 76.01606750488281, + "total_ms": 393.8347632288933, + "total_ms_wo_finalize": 317.81869572401047, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 394.30600916966796, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 0.912447988986969, + "diffuse_ms": 310.0645751953125, + "decode_ms": 6.830719947814941, + "finalize_ms": 76.0323486328125, + "total_ms": 393.8400917649269, + "total_ms_wo_finalize": 317.8077431321144, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 394.31095914915204, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 0.9106879830360413, + "diffuse_ms": 310.0616760253906, + "decode_ms": 6.804160118103027, + "finalize_ms": 76.05542755126953, + "total_ms": 393.8319516777992, + "total_ms_wo_finalize": 317.7765241265297, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 394.3494288250804, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 0.9143040180206299, + "diffuse_ms": 310.0634460449219, + "decode_ms": 6.818304061889648, + "finalize_ms": 76.04914855957031, + "total_ms": 393.84520268440247, + "total_ms_wo_finalize": 317.79605412483215, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 394.32757813483477, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 0.9136959910392761, + "diffuse_ms": 310.06524658203125, + "decode_ms": 6.833568096160889, + "finalize_ms": 76.03424072265625, + "total_ms": 393.84675139188766, + "total_ms_wo_finalize": 317.8125106692314, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 394.3099360913038, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 0.9118720293045044, + "diffuse_ms": 310.06439208984375, + "decode_ms": 6.810783863067627, + "finalize_ms": 76.05474853515625, + "total_ms": 393.84179651737213, + "total_ms_wo_finalize": 317.7870479822159, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 394.29685892537236, + "output_frames": 12, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 8, + "warmup": false, + "end_to_end_ms": 393.33291398361325, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.839136004447937, + "diffuse_ms": 308.13580322265625, + "decode_ms": 6.922368049621582, + "finalize_ms": 77.17052459716797 + }, + "per_rank": [ + { + "encode_ms": 0.7513279914855957, + "diffuse_ms": 308.117431640625, + "decode_ms": 6.802080154418945, + "finalize_ms": 77.17052459716797, + "total_ms": 392.8413643836975, + "total_ms_wo_finalize": 315.67083978652954, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.33291398361325, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 0.7589439749717712, + "diffuse_ms": 308.11248779296875, + "decode_ms": 6.9152960777282715, + "finalize_ms": 77.03132629394531, + "total_ms": 392.8180541396141, + "total_ms_wo_finalize": 315.7867278456688, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.2549571618438, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 0.7545920014381409, + "diffuse_ms": 308.1140441894531, + "decode_ms": 6.87667179107666, + "finalize_ms": 77.0962905883789, + "total_ms": 392.84159857034683, + "total_ms_wo_finalize": 315.7453079819679, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.30667117610574, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 0.839136004447937, + "diffuse_ms": 308.01910400390625, + "decode_ms": 6.822783946990967, + "finalize_ms": 77.15331268310547, + "total_ms": 392.8343366384506, + "total_ms_wo_finalize": 315.68102395534515, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.29038420692086, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 0.7474880218505859, + "diffuse_ms": 308.11724853515625, + "decode_ms": 6.810848236083984, + "finalize_ms": 77.16185760498047, + "total_ms": 392.8374423980713, + "total_ms_wo_finalize": 315.6755847930908, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.3010441251099, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 0.7458239793777466, + "diffuse_ms": 308.11773681640625, + "decode_ms": 6.82092809677124, + "finalize_ms": 77.15782165527344, + "total_ms": 392.8423105478287, + "total_ms_wo_finalize": 315.68448889255524, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.2732022367418, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 0.7377920150756836, + "diffuse_ms": 308.13580322265625, + "decode_ms": 6.922368049621582, + "finalize_ms": 77.05593872070312, + "total_ms": 392.85190200805664, + "total_ms_wo_finalize": 315.7959632873535, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.2799599133432, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 0.7356160283088684, + "diffuse_ms": 308.13262939453125, + "decode_ms": 6.8212480545043945, + "finalize_ms": 77.14665222167969, + "total_ms": 392.8361456990242, + "total_ms_wo_finalize": 315.6894934773445, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.23975006118417, + "output_frames": 12, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 9, + "warmup": false, + "end_to_end_ms": 393.29481683671474, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.7973759770393372, + "diffuse_ms": 309.1429138183594, + "decode_ms": 6.898496150970459, + "finalize_ms": 76.14649963378906 + }, + "per_rank": [ + { + "encode_ms": 0.7012479901313782, + "diffuse_ms": 309.1429138183594, + "decode_ms": 6.814208030700684, + "finalize_ms": 76.13990020751953, + "total_ms": 392.79827004671097, + "total_ms_wo_finalize": 316.65836983919144, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.29481683671474, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 0.7140480279922485, + "diffuse_ms": 309.1291809082031, + "decode_ms": 6.898496150970459, + "finalize_ms": 76.03142547607422, + "total_ms": 392.77315056324005, + "total_ms_wo_finalize": 316.74172508716583, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.19934509694576, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 0.7973759770393372, + "diffuse_ms": 309.0380859375, + "decode_ms": 6.861023902893066, + "finalize_ms": 76.09347534179688, + "total_ms": 392.7899611592293, + "total_ms_wo_finalize": 316.6964858174324, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.2195222005248, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 0.734112024307251, + "diffuse_ms": 309.10565185546875, + "decode_ms": 6.830080032348633, + "finalize_ms": 76.1266860961914, + "total_ms": 392.79653000831604, + "total_ms_wo_finalize": 316.66984391212463, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.2318752631545, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 0.7052479982376099, + "diffuse_ms": 309.1343688964844, + "decode_ms": 6.809088230133057, + "finalize_ms": 76.14649963378906, + "total_ms": 392.7952047586441, + "total_ms_wo_finalize": 316.64870512485504, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.26663920655847, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 0.7067520022392273, + "diffuse_ms": 309.1338806152344, + "decode_ms": 6.833759784698486, + "finalize_ms": 76.12553405761719, + "total_ms": 392.7999264597893, + "total_ms_wo_finalize": 316.6743924021721, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.26402405276895, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 0.701632022857666, + "diffuse_ms": 309.1421203613281, + "decode_ms": 6.851327896118164, + "finalize_ms": 76.10626983642578, + "total_ms": 392.80135011672974, + "total_ms_wo_finalize": 316.69508028030396, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.23107805103064, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 0.6990079879760742, + "diffuse_ms": 309.14007568359375, + "decode_ms": 6.814464092254639, + "finalize_ms": 76.14057922363281, + "total_ms": 392.7941269874573, + "total_ms_wo_finalize": 316.65354776382446, + "mem_alloc_gib": 38.79584312438965, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 393.23105197399855, + "output_frames": 12, + "rank": 7 + } + ] + }, + { + "autoregressive_index": 10, + "warmup": false, + "end_to_end_ms": 392.140940297395, + "output_frames": 12, + "critical_rank": { + "encode_ms": 0.77183997631073, + "diffuse_ms": 307.56024169921875, + "decode_ms": 6.8657917976379395, + "finalize_ms": 76.00300598144531 + }, + "per_rank": [ + { + "encode_ms": 0.7087680101394653, + "diffuse_ms": 307.551513671875, + "decode_ms": 6.816671848297119, + "finalize_ms": 76.00300598144531, + "total_ms": 391.0799595117569, + "total_ms_wo_finalize": 315.0769535303116, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 391.5406297892332, + "output_frames": 12, + "rank": 0 + }, + { + "encode_ms": 0.7086079716682434, + "diffuse_ms": 307.552001953125, + "decode_ms": 6.8657917976379395, + "finalize_ms": 75.92253112792969, + "total_ms": 391.04893285036087, + "total_ms_wo_finalize": 315.1264017224312, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 392.09823217242956, + "output_frames": 12, + "rank": 1 + }, + { + "encode_ms": 0.700543999671936, + "diffuse_ms": 307.56024169921875, + "decode_ms": 6.865407943725586, + "finalize_ms": 75.94831848144531, + "total_ms": 391.0745121240616, + "total_ms_wo_finalize": 315.1261936426163, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 392.0955271460116, + "output_frames": 12, + "rank": 2 + }, + { + "encode_ms": 0.7470719814300537, + "diffuse_ms": 307.5117492675781, + "decode_ms": 6.8230719566345215, + "finalize_ms": 75.99581146240234, + "total_ms": 391.07770466804504, + "total_ms_wo_finalize": 315.0818932056427, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 392.140940297395, + "output_frames": 12, + "rank": 3 + }, + { + "encode_ms": 0.7072640061378479, + "diffuse_ms": 307.5475158691406, + "decode_ms": 6.820256233215332, + "finalize_ms": 75.99501037597656, + "total_ms": 391.07004648447037, + "total_ms_wo_finalize": 315.0750361084938, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 391.5128209628165, + "output_frames": 12, + "rank": 4 + }, + { + "encode_ms": 0.77183997631073, + "diffuse_ms": 307.4792785644531, + "decode_ms": 6.83292818069458, + "finalize_ms": 75.99065399169922, + "total_ms": 391.07470071315765, + "total_ms_wo_finalize": 315.08404672145844, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 391.50249119848013, + "output_frames": 12, + "rank": 5 + }, + { + "encode_ms": 0.7010560035705566, + "diffuse_ms": 307.55877685546875, + "decode_ms": 6.837855815887451, + "finalize_ms": 75.98025512695312, + "total_ms": 391.0779438018799, + "total_ms_wo_finalize": 315.09768867492676, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 391.6066149249673, + "output_frames": 12, + "rank": 6 + }, + { + "encode_ms": 0.7002559900283813, + "diffuse_ms": 307.5577697753906, + "decode_ms": 6.83139181137085, + "finalize_ms": 75.98316955566406, + "total_ms": 391.0725871324539, + "total_ms_wo_finalize": 315.08941757678986, + "mem_alloc_gib": 38.795698165893555, + "mem_reserved_gib": 43.685546875, + "mem_peak_gib": 40.879210472106934, + "wall_ms": 391.47986518219113, + "output_frames": 12, + "rank": 7 + } + ] + } + ], + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 268435456.0, + "transfer_ms": 1.015748381614685, + "bandwidth_gbps": 264.2735748919249 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 1.0130791664123535, + "bandwidth_gbps": 264.9698709634097 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 1.004738211631775, + "bandwidth_gbps": 267.1695501299184 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 1.0077767372131348, + "bandwidth_gbps": 266.36401306733933 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 1.0035033226013184, + "bandwidth_gbps": 267.4983230789428 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 1.0073258876800537, + "bandwidth_gbps": 266.48322978994094 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.9925570487976074, + "bandwidth_gbps": 270.44839017080693 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 1.004172921180725, + "bandwidth_gbps": 267.3199509147973 + } + ], + "all_gather": [ + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.7515121102333069, + "bandwidth_gbps": 357.1937861608966 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.758198082447052, + "bandwidth_gbps": 354.0439658375764 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.7415111064910889, + "bandwidth_gbps": 362.0113760268079 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.7422817945480347, + "bandwidth_gbps": 361.63551089575986 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.7438980937004089, + "bandwidth_gbps": 360.8497699795254 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.7442696690559387, + "bandwidth_gbps": 360.66961635087756 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.7518082857131958, + "bandwidth_gbps": 357.05306938104735 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.745453417301178, + "bandwidth_gbps": 360.09688837679147 + } + ] + } +} diff --git a/integrations/lingbot/docs/benchmark_h100_cp4_single_session/README.md b/integrations/lingbot/docs/benchmark_h100_cp4_single_session/README.md new file mode 100644 index 00000000..ca2d801a --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_cp4_single_session/README.md @@ -0,0 +1,47 @@ +# LingBot CP4 single-session disaggregation benchmark + +## Result + +Topology: **1 encoder : 1 DiT group with CP4 (ulysses) : 1 decoder**. +The 4 DiT ranks cooperate on one autoregressive session. + +| Metric | Median | P90 | +| --- | ---: | ---: | +| End-to-end chunk latency | 754.41 ms | 786.46 ms | +| Encoder compute | 0.85 ms | 0.92 ms | +| Encoder → CP leader handoff | 31.70 ms | 32.61 ms | +| CP input fanout | 0.76 ms | 1.54 ms | +| CP DiT critical path | 696.76 ms | 725.92 ms | +| CP leader → decoder handoff | 12.38 ms | 13.01 ms | +| Decoder compute | 7.07 ms | 7.09 ms | +| 256 MiB Mooncake probes | 41.98 GB/s | — | +| 256 MiB-equivalent NCCL broadcast | 307.31 GB/s | 309.31 GB/s | +| 256 MiB-equivalent NCCL all-gather | 389.11 GB/s | 391.99 GB/s | + +- Single-session throughput: **15.70 generated FPS** +- Latency speedup versus tracked CP1 baseline: **2.96×** +- DiT critical-path speedup: **3.13×** +- CP scaling efficiency: **78.2%** + +The headline excludes 6 warmup blocks and measures +5 blocks. It accelerates one session; it does not represent +independent concurrent sessions. + +## Peak allocated memory + +| Role | Peak | +| --- | ---: | +| Encoder | 13.72 GiB | +| CP DiT ranks | 40.52–40.64 GiB each | +| Decoder | 2.37 GiB | + +## Reproduction + +```bash +env 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 --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 --output-dir integrations/lingbot/docs/benchmark_h100_cp4_single_session +``` + +- Repository revision: `66bcd32ece1d03b3362d71a7340691a0687a4069` (modified worktree) +- Slurm: job `14646820` on `pool0-01260` +- GPU: `NVIDIA H100 80GB HBM3` × 6 +- Model: `lingbot-world-fast-taehv-window15-sink3` diff --git a/integrations/lingbot/docs/benchmark_h100_cp4_single_session/benchmark.json b/integrations/lingbot/docs/benchmark_h100_cp4_single_session/benchmark.json new file mode 100644 index 00000000..267657ca --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_cp4_single_session/benchmark.json @@ -0,0 +1,958 @@ +{ + "environment": { + "command": "env 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 --bandwidth-probe-mib 256 --bandwidth-probe-iters 8 --output-dir integrations/lingbot/docs/benchmark_h100_cp4_single_session", + "commit": "66bcd32ece1d03b3362d71a7340691a0687a4069", + "worktree_dirty": true, + "hostname": "pool0-01260", + "slurm_job_id": "14646820", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "encoder": [ + 0 + ], + "dit_cp_group": [ + 1, + 2, + 3, + 4 + ], + "decoder": [ + 5 + ], + "cp_size": 4, + "cp_method": "ulysses" + }, + "noise_seed_by_cp_rank": [ + 42, + 43, + 44, + 45 + ], + "peak_memory_gib_by_rank": [ + 13.719008922576904, + 40.641592502593994, + 40.641592502593994, + 40.52054977416992, + 40.64252519607544, + 2.3715553283691406 + ] + }, + "summary": { + "fps": 15.695497779264775, + "latency_ms": { + "median": 754.4078070059186, + "p90": 786.4595750012086, + "min": 752.1404660001281, + "max": 805.1460170026985 + }, + "latency_speedup": 2.9606889263829443, + "fps_speedup": 2.9279196433600414, + "encoder_ms": { + "median": 0.8518400192260742, + "p90": 0.9241215944290161, + "min": 0.7815999984741211, + "max": 0.9528319835662842 + }, + "encoder_to_cp_leader": { + "payload_mib": 14.3583984375, + "copy_ms": { + "median": 1.3915720046497881, + "p90": 1.4601693954318762, + "min": 1.2229470012243837, + "max": 1.4713849959662184 + }, + "handoff_ms": { + "median": 31.70047700405121, + "p90": 32.611204001295846, + "min": 31.465993997699115, + "max": 33.13022400107002 + } + }, + "cp_input_fanout_ms": { + "median": 0.7586110004922375, + "p90": 1.5373780013760552, + "min": 0.7123480027075857, + "max": 2.052417999948375 + }, + "dit_ms": { + "median": 558.4301147460938, + "p90": 587.654345703125, + "min": 556.900634765625, + "max": 606.6038818359375 + }, + "finalize_ms": { + "median": 138.19427490234375, + "p90": 138.49698181152343, + "min": 137.53570556640625, + "max": 138.60121154785156 + }, + "dit_critical_path_ms": { + "median": 696.76318359375, + "p90": 725.924917602539, + "min": 694.4180908203125, + "max": 744.6621246337891 + }, + "dit_speedup": 3.128215114392654, + "cp_efficiency": 0.7820537785981635, + "cp_leader_to_decoder": { + "payload_mib": 0.55224609375, + "copy_ms": { + "median": 0.6824059964856133, + "p90": 0.8569654004531913, + "min": 0.6703260005451739, + "max": 0.9332609988632612 + }, + "handoff_ms": { + "median": 12.382764994981699, + "p90": 13.014173397095874, + "min": 12.201350000395905, + "max": 13.416604997473769 + } + }, + "decoder_ms": { + "median": 7.068831920623779, + "p90": 7.087052917480468, + "min": 7.055232048034668, + "max": 7.098688125610352 + }, + "mooncake_probe_gbps": { + "encoder_to_cp_leader": { + "median": 41.440791969662115, + "p90": 41.66010550567605, + "min": 33.04448130618013, + "max": 41.74742723770987 + }, + "cp_leader_to_decoder": { + "median": 42.51806596682346, + "p90": 42.5759180712987, + "min": 41.03074074711107, + "max": 42.61020933184513 + } + }, + "cp_probe_gbps": { + "broadcast": { + "median": 307.30739047654595, + "p90": 309.3124883855456, + "min": 301.9276905627407, + "max": 309.6337682734194 + }, + "all_gather": { + "median": 389.1070721288048, + "p90": 391.9884906736949, + "min": 73.74961070647561, + "max": 393.9417522633542 + } + }, + "baseline": { + "fps": 5.360631332509123, + "latency_ms": 2233.5668401792645, + "dit_critical_path_ms": 2179.6251220703125, + "topology": "1 encoder : 1 DiT : 1 decoder" + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "encoder_ms": 328.5774230957031, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 11.583731000428088, + "transfer_ms": 1.3928869957453571, + "bandwidth_gbps": 10.809112329994402 + }, + "encoder_to_cp_leader_handoff_ms": 32.75364900036948, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 6.440343000576831, + "transfer_ms": 1.0051130011561327, + "bandwidth_gbps": 0.5761262657372073 + }, + "cp_leader_to_decoder_handoff_ms": 16.200124999159016, + "end_to_end_ms": 123448.03048999893, + "cp_input_fanout_ms": 2.642430001287721, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.6733039990649559, + "dit_ms": 98526.9375, + "finalize_ms": 84.4049301147461, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 2.574874997662846, + "dit_ms": 98526.6015625, + "finalize_ms": 84.33760070800781, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 2.589513002021704, + "dit_ms": 98526.71875, + "finalize_ms": 84.4142074584961, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 2.642430001287721, + "dit_ms": 98527.140625, + "finalize_ms": 84.43023681640625, + "cp_rank": 3 + } + ], + "decoder_ms": 24447.322265625, + "output_frames": 9 + }, + { + "autoregressive_index": 1, + "warmup": true, + "encoder_ms": 183.88345336914062, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 34.161319003032986, + "transfer_ms": 1.5505840055993758, + "bandwidth_gbps": 9.709807366534893 + }, + "encoder_to_cp_leader_handoff_ms": 57.352356998308096, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.641785999410786, + "transfer_ms": 1.0275450040353462, + "bandwidth_gbps": 0.5635490394346568 + }, + "cp_leader_to_decoder_handoff_ms": 14.383066001755651, + "end_to_end_ms": 106482.99699999916, + "cp_input_fanout_ms": 2.2798829959356226, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.6927620052010752, + "dit_ms": 106112.4921875, + "finalize_ms": 96.51990509033203, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 2.2302859943010844, + "dit_ms": 106112.15625, + "finalize_ms": 96.49561309814453, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 2.215884000179358, + "dit_ms": 106112.3046875, + "finalize_ms": 96.51715087890625, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 2.2798829959356226, + "dit_ms": 106112.5234375, + "finalize_ms": 96.51939392089844, + "cp_rank": 3 + } + ], + "decoder_ms": 8.858112335205078, + "output_frames": 12 + }, + { + "autoregressive_index": 2, + "warmup": true, + "encoder_ms": 162.7545623779297, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.530125000921544, + "transfer_ms": 1.3899249970563687, + "bandwidth_gbps": 10.832147081235208 + }, + "encoder_to_cp_leader_handoff_ms": 40.56161800690461, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.615201993146911, + "transfer_ms": 0.889042996277567, + "bandwidth_gbps": 0.651343076121831 + }, + "cp_leader_to_decoder_handoff_ms": 19.472036001388915, + "end_to_end_ms": 804.2553280029097, + "cp_input_fanout_ms": 0.7735899998806417, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.6053379984223284, + "dit_ms": 458.44818115234375, + "finalize_ms": 107.85107421875, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.7146019997890107, + "dit_ms": 458.4527282714844, + "finalize_ms": 107.78377532958984, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.7029999978840351, + "dit_ms": 458.4573059082031, + "finalize_ms": 107.89606475830078, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.7735899998806417, + "dit_ms": 458.4759826660156, + "finalize_ms": 107.91206359863281, + "cp_rank": 3 + } + ], + "decoder_ms": 7.832032203674316, + "output_frames": 12 + }, + { + "autoregressive_index": 3, + "warmup": true, + "encoder_ms": 162.59071350097656, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.3245240029064, + "transfer_ms": 1.2906810006825253, + "bandwidth_gbps": 11.665060531640506 + }, + "encoder_to_cp_leader_handoff_ms": 31.744685999001376, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.625280998880044, + "transfer_ms": 0.8671780014992692, + "bandwidth_gbps": 0.6677660168948463 + }, + "cp_leader_to_decoder_handoff_ms": 12.751242997182999, + "end_to_end_ms": 856.4458559994819, + "cp_input_fanout_ms": 0.7571550013381056, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5959249974694103, + "dit_ms": 500.4807434082031, + "finalize_ms": 118.81871795654297, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6989710018387996, + "dit_ms": 500.48175048828125, + "finalize_ms": 118.79654693603516, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.6871869991300628, + "dit_ms": 500.489013671875, + "finalize_ms": 118.83296203613281, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.7571550013381056, + "dit_ms": 500.4973449707031, + "finalize_ms": 118.82921600341797, + "cp_rank": 3 + } + ], + "decoder_ms": 23.30944061279297, + "output_frames": 12 + }, + { + "autoregressive_index": 4, + "warmup": true, + "encoder_ms": 162.44419860839844, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.389432995405514, + "transfer_ms": 1.280101998418104, + "bandwidth_gbps": 11.761462772970756 + }, + "encoder_to_cp_leader_handoff_ms": 31.801788994926028, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.595009995682631, + "transfer_ms": 0.8645479974802583, + "bandwidth_gbps": 0.669797398973471 + }, + "cp_leader_to_decoder_handoff_ms": 12.68124500347767, + "end_to_end_ms": 894.6599800037802, + "cp_input_fanout_ms": 0.7668589969398454, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5944180011283606, + "dit_ms": 544.385986328125, + "finalize_ms": 129.6016387939453, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6959400052437559, + "dit_ms": 544.41552734375, + "finalize_ms": 129.60604858398438, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.7002049969742075, + "dit_ms": 544.3984375, + "finalize_ms": 129.63848876953125, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.7668589969398454, + "dit_ms": 544.4238891601562, + "finalize_ms": 129.6150360107422, + "cp_rank": 3 + } + ], + "decoder_ms": 7.077087879180908, + "output_frames": 12 + }, + { + "autoregressive_index": 5, + "warmup": true, + "encoder_ms": 0.8561279773712158, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.347429001238197, + "transfer_ms": 1.394298997183796, + "bandwidth_gbps": 10.798165981909072 + }, + "encoder_to_cp_leader_handoff_ms": 31.78755299450131, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.612675998941995, + "transfer_ms": 0.8892059995559976, + "bandwidth_gbps": 0.6512236762787755 + }, + "cp_leader_to_decoder_handoff_ms": 12.969750998308882, + "end_to_end_ms": 787.4171129951719, + "cp_input_fanout_ms": 0.7558500001323409, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5990170029690489, + "dit_ms": 586.8340454101562, + "finalize_ms": 141.0741729736328, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.7006710002315231, + "dit_ms": 586.8402709960938, + "finalize_ms": 140.9886016845703, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.6811769999330863, + "dit_ms": 586.8350830078125, + "finalize_ms": 141.03689575195312, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.7558500001323409, + "dit_ms": 586.8484497070312, + "finalize_ms": 141.1025848388672, + "cp_rank": 3 + } + ], + "decoder_ms": 7.082911968231201, + "output_frames": 12 + }, + { + "autoregressive_index": 6, + "warmup": false, + "encoder_ms": 0.881056010723114, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.341219997324515, + "transfer_ms": 1.4433459946303628, + "bandwidth_gbps": 10.43122858691673 + }, + "encoder_to_cp_leader_handoff_ms": 31.61951199581381, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.604931003996171, + "transfer_ms": 0.9332609988632612, + "bandwidth_gbps": 0.6204823738539662 + }, + "cp_leader_to_decoder_handoff_ms": 13.416604997473769, + "end_to_end_ms": 805.1460170026985, + "cp_input_fanout_ms": 0.750424005673267, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5924139986746013, + "dit_ms": 606.5610961914062, + "finalize_ms": 138.04669189453125, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6820019989390858, + "dit_ms": 606.5941162109375, + "finalize_ms": 138.04722595214844, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.697588999173604, + "dit_ms": 606.5676879882812, + "finalize_ms": 138.06198120117188, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.750424005673267, + "dit_ms": 606.6038818359375, + "finalize_ms": 138.05824279785156, + "cp_rank": 3 + } + ], + "decoder_ms": 7.098688125610352, + "output_frames": 12 + }, + { + "autoregressive_index": 7, + "warmup": false, + "encoder_ms": 0.9528319835662842, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.557260997709818, + "transfer_ms": 1.4713849959662184, + "bandwidth_gbps": 10.232449047173557 + }, + "encoder_to_cp_leader_handoff_ms": 33.13022400107002, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.532245002337731, + "transfer_ms": 0.7425220028380863, + "bandwidth_gbps": 0.7798718392002613 + }, + "cp_leader_to_decoder_handoff_ms": 12.382764994981699, + "end_to_end_ms": 758.4299119989737, + "cp_input_fanout_ms": 2.052417999948375, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.6273480030358769, + "dit_ms": 559.21533203125, + "finalize_ms": 138.56838989257812, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 1.983968002605252, + "dit_ms": 559.2216186523438, + "finalize_ms": 138.55628967285156, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 1.984201997402124, + "dit_ms": 559.2178955078125, + "finalize_ms": 138.60121154785156, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 2.052417999948375, + "dit_ms": 559.2300415039062, + "finalize_ms": 138.56492614746094, + "cp_rank": 3 + } + ], + "decoder_ms": 7.068831920623779, + "output_frames": 12 + }, + { + "autoregressive_index": 8, + "warmup": false, + "encoder_ms": 0.8232960104942322, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.405875000695232, + "transfer_ms": 1.3915720046497881, + "bandwidth_gbps": 10.819326595887544 + }, + "encoder_to_cp_leader_handoff_ms": 31.70047700405121, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.541837995930109, + "transfer_ms": 0.6824059964856133, + "bandwidth_gbps": 0.8485740204251095 + }, + "cp_leader_to_decoder_handoff_ms": 12.410525996529032, + "end_to_end_ms": 752.1404660001281, + "cp_input_fanout_ms": 0.7648180035175756, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5936609959462658, + "dit_ms": 556.8681030273438, + "finalize_ms": 137.49510192871094, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6949390008230694, + "dit_ms": 556.8894653320312, + "finalize_ms": 137.49679565429688, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.6988070017541759, + "dit_ms": 556.8823852539062, + "finalize_ms": 137.53570556640625, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.7648180035175756, + "dit_ms": 556.900634765625, + "finalize_ms": 137.51023864746094, + "cp_rank": 3 + } + ], + "decoder_ms": 7.0696001052856445, + "output_frames": 12 + }, + { + "autoregressive_index": 9, + "warmup": false, + "encoder_ms": 0.7815999984741211, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.366376002319157, + "transfer_ms": 1.3626879954244941, + "bandwidth_gbps": 11.048656809594855 + }, + "encoder_to_cp_leader_handoff_ms": 31.832674001634587, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.555463005090132, + "transfer_ms": 0.6703260005451739, + "bandwidth_gbps": 0.8638662375158396 + }, + "cp_leader_to_decoder_handoff_ms": 12.201350000395905, + "end_to_end_ms": 752.6280829988536, + "cp_input_fanout_ms": 0.7586110004922375, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.595359000726603, + "dit_ms": 557.02099609375, + "finalize_ms": 138.18722534179688, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6985649961279705, + "dit_ms": 557.0323486328125, + "finalize_ms": 138.18450927734375, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.6833459992776625, + "dit_ms": 557.0263061523438, + "finalize_ms": 138.1914520263672, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.7586110004922375, + "dit_ms": 557.0416870117188, + "finalize_ms": 138.19427490234375, + "cp_rank": 3 + } + ], + "decoder_ms": 7.067391872406006, + "output_frames": 12 + }, + { + "autoregressive_index": 10, + "warmup": false, + "encoder_ms": 0.8518400192260742, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.298101995838806, + "transfer_ms": 1.2229470012243837, + "bandwidth_gbps": 12.311140208796 + }, + "encoder_to_cp_leader_handoff_ms": 31.465993997699115, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.536013002507389, + "transfer_ms": 0.6793099964852445, + "bandwidth_gbps": 0.8524414523503604 + }, + "cp_leader_to_decoder_handoff_ms": 12.370764998195227, + "end_to_end_ms": 754.4078070059186, + "cp_input_fanout_ms": 0.7123480027075857, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5653779953718185, + "dit_ms": 558.4015502929688, + "finalize_ms": 138.29837036132812, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6622430009883828, + "dit_ms": 558.4152221679688, + "finalize_ms": 138.291748046875, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.6471220040111803, + "dit_ms": 558.4225463867188, + "finalize_ms": 138.34063720703125, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.7123480027075857, + "dit_ms": 558.4301147460938, + "finalize_ms": 138.3023681640625, + "cp_rank": 3 + } + ], + "decoder_ms": 7.055232048034668, + "output_frames": 12 + } + ], + "mooncake_probe": { + "encoder_to_cp_leader": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.005185000190977007, + "transfer_ms": 6.557585998962168, + "bandwidth_gbps": 40.93510264943286 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0022800013539381325, + "transfer_ms": 6.480381001892965, + "bandwidth_gbps": 41.422789172671806 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0019889994291588664, + "transfer_ms": 6.474753005022649, + "bandwidth_gbps": 41.45879476665242 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0023190004867501557, + "transfer_ms": 6.4527340000495315, + "bandwidth_gbps": 41.60026680131856 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0019609942683018744, + "transfer_ms": 6.4492590026929975, + "bandwidth_gbps": 41.622681906232984 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.005360001523513347, + "transfer_ms": 8.123457999317907, + "bandwidth_gbps": 33.04448130618013 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002687003870960325, + "transfer_ms": 7.159761000366416, + "bandwidth_gbps": 37.49223695962229 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001863001671154052, + "transfer_ms": 6.42998799594352, + "bandwidth_gbps": 41.74742723770987 + } + ], + "cp_leader_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.003088003722950816, + "transfer_ms": 6.384660002368037, + "bandwidth_gbps": 42.04381375052684 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.00183499651029706, + "transfer_ms": 6.308395997621119, + "bandwidth_gbps": 42.5520934483546 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016980047803372145, + "transfer_ms": 6.317300998489372, + "bandwidth_gbps": 42.492111119003155 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002481996489223093, + "transfer_ms": 6.542300994624384, + "bandwidth_gbps": 41.03074074711107 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.001515996700618416, + "transfer_ms": 6.299792003119364, + "bandwidth_gbps": 42.61020933184513 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0012460004654712975, + "transfer_ms": 6.3070429969229735, + "bandwidth_gbps": 42.561221816778804 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0013600001693703234, + "transfer_ms": 6.309593001788016, + "bandwidth_gbps": 42.54402081464376 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014600009308196604, + "transfer_ms": 6.329215000732802, + "bandwidth_gbps": 42.41212472145761 + } + ] + }, + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.8890720009803772, + "bandwidth_gbps": 301.9276905627407 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.8837230205535889, + "bandwidth_gbps": 303.7551922454668 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.8746740221977234, + "bandwidth_gbps": 306.89771181899715 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.8669450283050537, + "bandwidth_gbps": 309.6337682734194 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.8723450303077698, + "bandwidth_gbps": 307.7170691340948 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.8749210238456726, + "bandwidth_gbps": 306.81107058109666 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.8693100214004517, + "bandwidth_gbps": 308.79139707552497 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.868232011795044, + "bandwidth_gbps": 309.17479700502827 + } + ], + "all_gather": [ + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.6906840205192566, + "bandwidth_gbps": 388.65160916592527 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.6969090104103088, + "bandwidth_gbps": 385.18006223216605 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.6890689730644226, + "bandwidth_gbps": 389.5625350916843 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.6888560056686401, + "bandwidth_gbps": 389.68297262566847 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.6945549845695496, + "bandwidth_gbps": 386.4855367302026 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.6814090013504028, + "bandwidth_gbps": 393.9417522633542 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 3.639822006225586, + "bandwidth_gbps": 73.74961070647561 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.686269998550415, + "bandwidth_gbps": 391.1513785638409 + } + ] + } +} diff --git a/integrations/lingbot/docs/benchmark_h100_cp6_single_session/README.md b/integrations/lingbot/docs/benchmark_h100_cp6_single_session/README.md new file mode 100644 index 00000000..209be7d7 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_cp6_single_session/README.md @@ -0,0 +1,47 @@ +# LingBot CP6 single-session disaggregation benchmark + +## Result + +Topology: **1 encoder : 1 DiT group with CP6 (ring) : 1 decoder**. +The six DiT ranks cooperate on one autoregressive session. + +| Metric | Median | P90 | +| --- | ---: | ---: | +| End-to-end chunk latency | 743.27 ms | 780.60 ms | +| Encoder compute | 0.90 ms | 1.12 ms | +| Encoder → CP leader handoff | 30.33 ms | 36.77 ms | +| CP input fanout | 0.89 ms | 1.56 ms | +| CP DiT critical path | 683.16 ms | 713.95 ms | +| CP leader → decoder handoff | 12.59 ms | 15.92 ms | +| Decoder compute | 7.13 ms | 7.17 ms | +| 256 MiB Mooncake probes | 42.44 GB/s | — | +| 256 MiB-equivalent NCCL broadcast | 283.71 GB/s | 286.38 GB/s | +| 256 MiB-equivalent NCCL all-gather | 231.72 GB/s | 232.51 GB/s | + +- Single-session throughput: **15.90 generated FPS** +- Latency speedup versus tracked CP1 baseline: **3.01×** +- DiT critical-path speedup: **3.19×** +- CP scaling efficiency: **53.2%** + +The headline excludes 6 warmup blocks and measures +5 blocks. It accelerates one session; it does not represent +independent concurrent sessions. + +## Peak allocated memory + +| Role | Peak | +| --- | ---: | +| Encoder | 13.72 GiB | +| CP DiT ranks | 39.18–39.18 GiB each | +| Decoder | 2.29 GiB | + +## Reproduction + +```bash +env 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 integrations/lingbot/docs/benchmark_h100_cp6_single_session +``` + +- Repository revision: `66bcd32ece1d03b3362d71a7340691a0687a4069` (modified worktree) +- Slurm: job `14646820` on `pool0-01260` +- GPU: `NVIDIA H100 80GB HBM3` × 8 +- Model: `lingbot-world-fast-taehv-window15-sink3` diff --git a/integrations/lingbot/docs/benchmark_h100_cp6_single_session/benchmark.json b/integrations/lingbot/docs/benchmark_h100_cp6_single_session/benchmark.json new file mode 100644 index 00000000..6edc86bb --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_cp6_single_session/benchmark.json @@ -0,0 +1,1098 @@ +{ + "environment": { + "command": "env 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 integrations/lingbot/docs/benchmark_h100_cp6_single_session", + "commit": "66bcd32ece1d03b3362d71a7340691a0687a4069", + "worktree_dirty": true, + "hostname": "pool0-01260", + "slurm_job_id": "14646820", + "python": "3.12.13", + "model": "lingbot-world-fast-taehv-window15-sink3", + "checkpoint": "https://huggingface.co/robbyant/lingbot-world-fast/blob/main/diffusion_pytorch_model.safetensors.index.json", + "precision": "bfloat16", + "decoder_config": "TeahvVAEDecoderConfig", + "seed": 42, + "example_index": 0, + "example_url": "https://raw.githubusercontent.com/Robbyant/lingbot-world-v2/main/examples/00", + "prompt": "The video presents a soaring journey through a fantasy jungle. The wind whips past the rider's blue hands gripping the reins, causing the leather straps to vibrate. The ancient gothic castle approaches steadily, its stone details becoming clearer against the backdrop of floating islands and distant waterfalls.", + "resolution": [ + 464, + 832 + ], + "target_fps": 16, + "latent_frames_per_chunk": 3, + "window_size_t": 15, + "sink_size_t": 3, + "guidance_scale": 1.0, + "num_inference_steps": 4, + "compile_network": true, + "warmup_blocks": 6, + "measured_blocks": 5, + "torch": "2.12.1+cu130", + "cuda": "13.0", + "cudnn": 92000, + "driver": "535.216.03", + "mooncake": "0.3.12.post1", + "triton_cache_dir": "/lustre/fsw/portfolios/healthcareeng/users/gtong/triton-cache", + "torchinductor_compile_threads": "4", + "hf_home": "/lustre/fsw/portfolios/healthcareeng/users/gtong/hf", + "gpus": [ + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3", + "NVIDIA H100 80GB HBM3" + ], + "rdma_device": null, + "allocation": { + "encoder": [ + 0 + ], + "dit_cp_group": [ + 1, + 2, + 3, + 4, + 5, + 6 + ], + "decoder": [ + 7 + ], + "cp_size": 6, + "cp_method": "ring" + }, + "noise_seed_by_cp_rank": [ + 42, + 43, + 44, + 45, + 46, + 47 + ], + "peak_memory_gib_by_rank": [ + 13.719008922576904, + 39.17708206176758, + 39.17708206176758, + 39.17708206176758, + 39.17708206176758, + 39.17708206176758, + 39.17708206176758, + 2.2908687591552734 + ] + }, + "summary": { + "fps": 15.901998021609447, + "latency_ms": { + "median": 743.272824001906, + "p90": 780.5990148001001, + "min": 739.5753040036652, + "max": 802.4121859998559 + }, + "latency_speedup": 3.0050430582856036, + "fps_speedup": 2.9664412706714307, + "encoder_ms": { + "median": 0.8965439796447754, + "p90": 1.1200191974639893, + "min": 0.8386880159378052, + "max": 1.1643199920654297 + }, + "encoder_to_cp_leader": { + "payload_mib": 14.3583984375, + "copy_ms": { + "median": 1.0783469988382421, + "p90": 1.116982998792082, + "min": 1.0716870019678026, + "max": 1.130860997363925 + }, + "handoff_ms": { + "median": 30.32546600297792, + "p90": 36.77386300114449, + "min": 29.192604997660965, + "max": 40.76279900618829 + } + }, + "cp_input_fanout_ms": { + "median": 0.8855909982230514, + "p90": 1.5593643998727202, + "min": 0.8139759956975468, + "max": 1.9920600025216118 + }, + "dit_ms": { + "median": 547.8768310546875, + "p90": 578.8016479492187, + "min": 547.67626953125, + "max": 597.9907836914062 + }, + "finalize_ms": { + "median": 135.25392150878906, + "p90": 135.28230895996094, + "min": 135.10354614257812, + "max": 135.29872131347656 + }, + "dit_critical_path_ms": { + "median": 683.1607208251953, + "p90": 713.9518127441406, + "min": 682.9142608642578, + "max": 733.0852966308594 + }, + "dit_speedup": 3.190501232912697, + "cp_efficiency": 0.5317502054854495, + "cp_leader_to_decoder": { + "payload_mib": 0.55224609375, + "copy_ms": { + "median": 0.7526259942096658, + "p90": 0.8522303978679702, + "min": 0.7345539997913875, + "max": 0.9063379984581843 + }, + "handoff_ms": { + "median": 12.59425900207134, + "p90": 15.920846402877942, + "min": 12.365162001515273, + "max": 17.5781400030246 + } + }, + "decoder_ms": { + "median": 7.12608003616333, + "p90": 7.167359924316406, + "min": 7.104191780090332, + "max": 7.178239822387695 + }, + "mooncake_probe_gbps": { + "encoder_to_cp_leader": { + "median": 42.61892648972679, + "p90": 42.71346310150518, + "min": 42.11309136346771, + "max": 42.73283630120451 + }, + "cp_leader_to_decoder": { + "median": 42.25628792037267, + "p90": 42.28569924882899, + "min": 41.8635907422804, + "max": 42.308059370096046 + } + }, + "cp_probe_gbps": { + "broadcast": { + "median": 283.7137016661082, + "p90": 286.3811137793339, + "min": 280.2329841250139, + "max": 286.6822005305427 + }, + "all_gather": { + "median": 231.71582450799085, + "p90": 232.50550409607658, + "min": 100.15586831387085, + "max": 232.54694908667284 + } + }, + "baseline": { + "fps": 5.360631332509123, + "latency_ms": 2233.5668401792645, + "dit_critical_path_ms": 2179.6251220703125, + "topology": "1 encoder : 1 DiT : 1 decoder" + } + }, + "records": [ + { + "autoregressive_index": 0, + "warmup": true, + "encoder_ms": 322.0684814453125, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.245085002679843, + "transfer_ms": 1.2169600013294257, + "bandwidth_gbps": 12.371706533947489 + }, + "encoder_to_cp_leader_handoff_ms": 30.891644004441332, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.327522998210043, + "transfer_ms": 0.9407070028828457, + "bandwidth_gbps": 0.6155710526501914 + }, + "cp_leader_to_decoder_handoff_ms": 14.864731005218346, + "end_to_end_ms": 108229.41726999852, + "cp_input_fanout_ms": 3.8529209996340796, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.6306880022748373, + "dit_ms": 104103.7421875, + "finalize_ms": 68.08755493164062, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 3.6345599946798757, + "dit_ms": 104103.390625, + "finalize_ms": 68.05280303955078, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 3.687378004542552, + "dit_ms": 104103.390625, + "finalize_ms": 67.98947143554688, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 3.798120000283234, + "dit_ms": 104103.65625, + "finalize_ms": 68.02851104736328, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 3.7955410007270984, + "dit_ms": 104103.6796875, + "finalize_ms": 67.9978256225586, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 3.8529209996340796, + "dit_ms": 104103.9140625, + "finalize_ms": 68.03446197509766, + "cp_rank": 5 + } + ], + "decoder_ms": 3677.967041015625, + "output_frames": 9 + }, + { + "autoregressive_index": 1, + "warmup": true, + "encoder_ms": 179.2939453125, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 11.127067999041174, + "transfer_ms": 1.3310289941728115, + "bandwidth_gbps": 11.311453068200594 + }, + "encoder_to_cp_leader_handoff_ms": 34.32763800083194, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.409023000334855, + "transfer_ms": 1.0616360013955273, + "bandwidth_gbps": 0.5454524895904116 + }, + "cp_leader_to_decoder_handoff_ms": 14.75338899763301, + "end_to_end_ms": 139759.40580300085, + "cp_input_fanout_ms": 1.0192160043516196, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.7029449989204295, + "dit_ms": 139433.46875, + "finalize_ms": 81.27254486083984, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.8079170002019964, + "dit_ms": 139433.0625, + "finalize_ms": 81.30006408691406, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.8097389945760369, + "dit_ms": 139433.203125, + "finalize_ms": 81.35478210449219, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.9622339930501767, + "dit_ms": 139433.53125, + "finalize_ms": 81.35183715820312, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 0.9453659949940629, + "dit_ms": 139433.390625, + "finalize_ms": 81.32160186767578, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 1.0192160043516196, + "dit_ms": 139433.734375, + "finalize_ms": 81.28697967529297, + "cp_rank": 5 + } + ], + "decoder_ms": 8.011199951171875, + "output_frames": 12 + }, + { + "autoregressive_index": 2, + "warmup": true, + "encoder_ms": 162.93930053710938, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 11.216922001040075, + "transfer_ms": 1.1733100036508404, + "bandwidth_gbps": 12.831964232089172 + }, + "encoder_to_cp_leader_handoff_ms": 72.1355819987366, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.35779999679653, + "transfer_ms": 0.9203850058838725, + "bandwidth_gbps": 0.6291627919817102 + }, + "cp_leader_to_decoder_handoff_ms": 13.305078995472286, + "end_to_end_ms": 833.0178769974736, + "cp_input_fanout_ms": 0.8932590062613599, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.6093280026107095, + "dit_ms": 474.8171691894531, + "finalize_ms": 94.63977813720703, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6950169990886934, + "dit_ms": 474.8047180175781, + "finalize_ms": 94.63542175292969, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.7046709943097085, + "dit_ms": 474.7917785644531, + "finalize_ms": 94.7043228149414, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.8411980015807785, + "dit_ms": 474.79449462890625, + "finalize_ms": 94.77766418457031, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 0.8467049992759712, + "dit_ms": 474.82696533203125, + "finalize_ms": 94.73712158203125, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 0.8932590062613599, + "dit_ms": 474.8154296875, + "finalize_ms": 94.67286682128906, + "cp_rank": 5 + } + ], + "decoder_ms": 7.923935890197754, + "output_frames": 12 + }, + { + "autoregressive_index": 3, + "warmup": true, + "encoder_ms": 162.5497283935547, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 11.09844299935503, + "transfer_ms": 1.1059039970859885, + "bandwidth_gbps": 13.61408588780907 + }, + "encoder_to_cp_leader_handoff_ms": 29.335072002140805, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.367657995317131, + "transfer_ms": 0.8612400051788427, + "bandwidth_gbps": 0.6723700670172091 + }, + "cp_leader_to_decoder_handoff_ms": 13.213771999289747, + "end_to_end_ms": 870.8467379983631, + "cp_input_fanout_ms": 0.8225310011766851, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5755200036219321, + "dit_ms": 522.888427734375, + "finalize_ms": 108.86710357666016, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6737930016242899, + "dit_ms": 522.8917236328125, + "finalize_ms": 108.87321472167969, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.6564969953615218, + "dit_ms": 522.8981323242188, + "finalize_ms": 108.88716888427734, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.7676740060560405, + "dit_ms": 522.9008178710938, + "finalize_ms": 108.91696166992188, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 0.7652640051674098, + "dit_ms": 522.8975830078125, + "finalize_ms": 108.88742065429688, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 0.8225310011766851, + "dit_ms": 522.8972778320312, + "finalize_ms": 108.86297607421875, + "cp_rank": 5 + } + ], + "decoder_ms": 21.86777687072754, + "output_frames": 12 + }, + { + "autoregressive_index": 4, + "warmup": true, + "encoder_ms": 162.5436553955078, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 11.081044998718426, + "transfer_ms": 1.1266889996477403, + "bandwidth_gbps": 13.362935117594317 + }, + "encoder_to_cp_leader_handoff_ms": 29.348579999350477, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.401886995765381, + "transfer_ms": 0.864851004735101, + "bandwidth_gbps": 0.66956273026169 + }, + "cp_leader_to_decoder_handoff_ms": 13.14943300531013, + "end_to_end_ms": 919.7116880022804, + "cp_input_fanout_ms": 0.8869459998095408, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.6023640016792342, + "dit_ms": 577.4309692382812, + "finalize_ms": 122.73359680175781, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.7082009979058057, + "dit_ms": 577.4498291015625, + "finalize_ms": 122.748291015625, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.7084039971232414, + "dit_ms": 577.44580078125, + "finalize_ms": 122.76534271240234, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.8443309998256154, + "dit_ms": 577.4500122070312, + "finalize_ms": 122.7576675415039, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 0.8253650012193248, + "dit_ms": 577.4526977539062, + "finalize_ms": 122.75145721435547, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 0.8869459998095408, + "dit_ms": 577.4497680664062, + "finalize_ms": 122.7537612915039, + "cp_rank": 5 + } + ], + "decoder_ms": 7.172512054443359, + "output_frames": 12 + }, + { + "autoregressive_index": 5, + "warmup": true, + "encoder_ms": 1.050976037979126, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 11.193180005648173, + "transfer_ms": 1.103852002415806, + "bandwidth_gbps": 13.639393656984696 + }, + "encoder_to_cp_leader_handoff_ms": 29.30819299945142, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.387084998597857, + "transfer_ms": 0.9164600051008165, + "bandwidth_gbps": 0.6318573606889679 + }, + "cp_leader_to_decoder_handoff_ms": 13.342461003048811, + "end_to_end_ms": 812.876291005523, + "cp_input_fanout_ms": 0.8991159993456677, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5880129974684678, + "dit_ms": 617.8521728515625, + "finalize_ms": 137.089599609375, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6983849962125532, + "dit_ms": 617.8724975585938, + "finalize_ms": 137.1097869873047, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.6973109993850812, + "dit_ms": 617.8749389648438, + "finalize_ms": 137.1744384765625, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.8421500024269335, + "dit_ms": 617.8759765625, + "finalize_ms": 137.154052734375, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 0.83458999870345, + "dit_ms": 617.8741455078125, + "finalize_ms": 137.10704040527344, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 0.8991159993456677, + "dit_ms": 617.8782958984375, + "finalize_ms": 137.0811767578125, + "cp_rank": 5 + } + ], + "decoder_ms": 7.168479919433594, + "output_frames": 12 + }, + { + "autoregressive_index": 6, + "warmup": false, + "encoder_ms": 1.0535680055618286, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 11.598237993894145, + "transfer_ms": 1.0961660009343177, + "bandwidth_gbps": 13.73502917182901 + }, + "encoder_to_cp_leader_handoff_ms": 40.76279900618829, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.386552995129023, + "transfer_ms": 0.9063379984581843, + "bandwidth_gbps": 0.6389139603382927 + }, + "cp_leader_to_decoder_handoff_ms": 13.434906002657954, + "end_to_end_ms": 802.4121859998559, + "cp_input_fanout_ms": 0.87535099737579, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5783199958386831, + "dit_ms": 597.9799194335938, + "finalize_ms": 135.01052856445312, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6802579955547117, + "dit_ms": 597.9730834960938, + "finalize_ms": 135.0397186279297, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.6833650040789507, + "dit_ms": 597.9851684570312, + "finalize_ms": 135.10012817382812, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.8178429998224601, + "dit_ms": 597.9724731445312, + "finalize_ms": 135.10354614257812, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 0.8125800013658591, + "dit_ms": 597.9907836914062, + "finalize_ms": 135.0647430419922, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 0.87535099737579, + "dit_ms": 597.976318359375, + "finalize_ms": 135.02490234375, + "cp_rank": 5 + } + ], + "decoder_ms": 7.178239822387695, + "output_frames": 12 + }, + { + "autoregressive_index": 7, + "warmup": false, + "encoder_ms": 1.1643199920654297, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 11.142417002702132, + "transfer_ms": 1.130860997363925, + "bandwidth_gbps": 13.313636278106454 + }, + "encoder_to_cp_leader_handoff_ms": 30.19312100514071, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.31898999522673, + "transfer_ms": 0.7526259942096658, + "bandwidth_gbps": 0.7694020728158409 + }, + "cp_leader_to_decoder_handoff_ms": 12.365162001515273, + "end_to_end_ms": 743.272824001906, + "cp_input_fanout_ms": 1.9920600025216118, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.6100310056353919, + "dit_ms": 550.0137939453125, + "finalize_ms": 135.1255645751953, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 1.7673210022621788, + "dit_ms": 549.99169921875, + "finalize_ms": 135.12037658691406, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 1.7932339978870004, + "dit_ms": 550.00830078125, + "finalize_ms": 135.20230102539062, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 1.930327998707071, + "dit_ms": 549.993896484375, + "finalize_ms": 135.2576904296875, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 1.937477005412802, + "dit_ms": 550.0179443359375, + "finalize_ms": 135.20953369140625, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 1.9920600025216118, + "dit_ms": 549.9854736328125, + "finalize_ms": 135.17864990234375, + "cp_rank": 5 + } + ], + "decoder_ms": 7.151040077209473, + "output_frames": 12 + }, + { + "autoregressive_index": 8, + "warmup": false, + "encoder_ms": 0.8965439796447754, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 10.990686001605354, + "transfer_ms": 1.0783469988382421, + "bandwidth_gbps": 13.961991841420668 + }, + "encoder_to_cp_leader_handoff_ms": 30.790458993578795, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.3321010016370565, + "transfer_ms": 0.771068996982649, + "bandwidth_gbps": 0.7509989407770608 + }, + "cp_leader_to_decoder_handoff_ms": 12.431761999323498, + "end_to_end_ms": 739.9711980033317, + "cp_input_fanout_ms": 0.8855909982230514, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5862039979547262, + "dit_ms": 547.658447265625, + "finalize_ms": 135.15501403808594, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6804029981140047, + "dit_ms": 547.65869140625, + "finalize_ms": 135.1744384765625, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.7181150067481212, + "dit_ms": 547.67626953125, + "finalize_ms": 135.23228454589844, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.8443400001851842, + "dit_ms": 547.6729125976562, + "finalize_ms": 135.24134826660156, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 0.837255000078585, + "dit_ms": 547.670654296875, + "finalize_ms": 135.21142578125, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 0.8855909982230514, + "dit_ms": 547.6614379882812, + "finalize_ms": 135.1871337890625, + "cp_rank": 5 + } + ], + "decoder_ms": 7.115744113922119, + "output_frames": 12 + }, + { + "autoregressive_index": 9, + "warmup": false, + "encoder_ms": 0.8770880103111267, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 11.073223999119364, + "transfer_ms": 1.0750720030046068, + "bandwidth_gbps": 14.004524308996897 + }, + "encoder_to_cp_leader_handoff_ms": 29.192604997660965, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.371476003143471, + "transfer_ms": 0.7390309983748011, + "bandwidth_gbps": 0.7835557659603372 + }, + "cp_leader_to_decoder_handoff_ms": 17.5781400030246, + "end_to_end_ms": 747.8792580004665, + "cp_input_fanout_ms": 0.8139759956975468, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.57275099970866, + "dit_ms": 547.8679809570312, + "finalize_ms": 135.21629333496094, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6605269954889081, + "dit_ms": 547.8587646484375, + "finalize_ms": 135.17510986328125, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.6536399960168637, + "dit_ms": 547.8768310546875, + "finalize_ms": 135.24073791503906, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.752116997318808, + "dit_ms": 547.8619995117188, + "finalize_ms": 135.29872131347656, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 0.7600410026498139, + "dit_ms": 547.8764038085938, + "finalize_ms": 135.27622985839844, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 0.8139759956975468, + "dit_ms": 547.8638305664062, + "finalize_ms": 135.25442504882812, + "cp_rank": 5 + } + ], + "decoder_ms": 7.12608003616333, + "output_frames": 12 + }, + { + "autoregressive_index": 10, + "warmup": false, + "encoder_ms": 0.8386880159378052, + "encoder_to_cp_leader": { + "backend": "mooncake-rdma", + "payload_bytes": 15055872, + "registration_ms": 12.28495500254212, + "transfer_ms": 1.0716870019678026, + "bandwidth_gbps": 14.048758613620222 + }, + "encoder_to_cp_leader_handoff_ms": 30.32546600297792, + "cp_leader_to_decoder": { + "backend": "mooncake-rdma", + "payload_bytes": 579072, + "registration_ms": 4.406019994348753, + "transfer_ms": 0.7345539997913875, + "bandwidth_gbps": 0.7883314231008963 + }, + "cp_leader_to_decoder_handoff_ms": 12.59425900207134, + "end_to_end_ms": 739.5753040036652, + "cp_input_fanout_ms": 0.910320995899383, + "cp_workers": [ + { + "cp_input_fanout_ms": 0.5813440002384596, + "dit_ms": 547.7222900390625, + "finalize_ms": 135.18223571777344, + "cp_rank": 0 + }, + { + "cp_input_fanout_ms": 0.6931000025360845, + "dit_ms": 547.7235717773438, + "finalize_ms": 135.15213012695312, + "cp_rank": 1 + }, + { + "cp_input_fanout_ms": 0.7049069972708821, + "dit_ms": 547.728271484375, + "finalize_ms": 135.20755004882812, + "cp_rank": 2 + }, + { + "cp_input_fanout_ms": 0.8326189999934286, + "dit_ms": 547.728759765625, + "finalize_ms": 135.25392150878906, + "cp_rank": 3 + }, + { + "cp_input_fanout_ms": 0.849386997288093, + "dit_ms": 547.7361450195312, + "finalize_ms": 135.2175750732422, + "cp_rank": 4 + }, + { + "cp_input_fanout_ms": 0.910320995899383, + "dit_ms": 547.7199096679688, + "finalize_ms": 135.18972778320312, + "cp_rank": 5 + } + ], + "decoder_ms": 7.104191780090332, + "output_frames": 12 + } + ], + "mooncake_probe": { + "encoder_to_cp_leader": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.005959998816251755, + "transfer_ms": 6.374156997480895, + "bandwidth_gbps": 42.11309136346771 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002316002792213112, + "transfer_ms": 6.306149996817112, + "bandwidth_gbps": 42.567248818294324 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002224995114374906, + "transfer_ms": 6.285784998908639, + "bandwidth_gbps": 42.70516030163404 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0027510031941346824, + "transfer_ms": 6.302879999566358, + "bandwidth_gbps": 42.58933313318174 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0029850052669644356, + "transfer_ms": 6.3149230045382865, + "bandwidth_gbps": 42.50811226155663 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015350015019066632, + "transfer_ms": 6.292646001384128, + "bandwidth_gbps": 42.65859797944378 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016379999578930438, + "transfer_ms": 6.281713998760097, + "bandwidth_gbps": 42.73283630120451 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014979959814809263, + "transfer_ms": 6.294132996117696, + "bandwidth_gbps": 42.648519846271846 + } + ], + "cp_leader_to_decoder": [ + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0033650067052803934, + "transfer_ms": 6.412146001821384, + "bandwidth_gbps": 41.8635907422804 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.005699999746866524, + "transfer_ms": 6.353461001708638, + "bandwidth_gbps": 42.25027208442919 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0018470018403604627, + "transfer_ms": 6.350560004648287, + "bandwidth_gbps": 42.26957241621509 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.002125001628883183, + "transfer_ms": 6.352262003929354, + "bandwidth_gbps": 42.258246878663435 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0015190016711130738, + "transfer_ms": 6.344783003441989, + "bandwidth_gbps": 42.308059370096046 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0014920005924068391, + "transfer_ms": 6.360911997035146, + "bandwidth_gbps": 42.2007812912864 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0012029995559714735, + "transfer_ms": 6.35285099997418, + "bandwidth_gbps": 42.25432896208191 + }, + { + "backend": "mooncake-rdma", + "payload_bytes": 268435456, + "registration_ms": 0.0016019985196180642, + "transfer_ms": 6.349577000946738, + "bandwidth_gbps": 42.27611633971454 + } + ] + }, + "cp_probe": { + "broadcast": [ + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.9579010009765625, + "bandwidth_gbps": 280.2329841250139 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.9531909823417664, + "bandwidth_gbps": 281.6177040833067 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.936352014541626, + "bandwidth_gbps": 286.6822005305427 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.9378129839897156, + "bandwidth_gbps": 286.2355934314338 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.9416670203208923, + "bandwidth_gbps": 285.0640940027029 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.9377589821815491, + "bandwidth_gbps": 286.2520766002444 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.9529169797897339, + "bandwidth_gbps": 281.6986806754474 + }, + { + "payload_bytes": 268435456.0, + "transfer_ms": 0.9506739974021912, + "bandwidth_gbps": 282.36330932951347 + } + ], + "all_gather": [ + { + "payload_bytes": 268435452.0, + "transfer_ms": 2.6801769733428955, + "bandwidth_gbps": 100.15586831387085 + }, + { + "payload_bytes": 268435452.0, + "transfer_ms": 1.1686650514602661, + "bandwidth_gbps": 229.69408699659968 + }, + { + "payload_bytes": 268435452.0, + "transfer_ms": 1.154621958732605, + "bandwidth_gbps": 232.48774195724963 + }, + { + "payload_bytes": 268435452.0, + "transfer_ms": 1.1543279886245728, + "bandwidth_gbps": 232.54694908667284 + }, + { + "payload_bytes": 268435452.0, + "transfer_ms": 1.1633909940719604, + "bandwidth_gbps": 230.73537045396466 + }, + { + "payload_bytes": 268435452.0, + "transfer_ms": 1.1605249643325806, + "bandwidth_gbps": 231.30519398553187 + }, + { + "payload_bytes": 268435452.0, + "transfer_ms": 1.156419038772583, + "bandwidth_gbps": 232.12645503044982 + }, + { + "payload_bytes": 268435452.0, + "transfer_ms": 1.155879020690918, + "bandwidth_gbps": 232.23490278380928 + } + ] + } +} diff --git a/integrations/lingbot/docs/benchmark_h100_pipeline_3x2/README.md b/integrations/lingbot/docs/benchmark_h100_pipeline_3x2/README.md new file mode 100644 index 00000000..8fb836b8 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_pipeline_3x2/README.md @@ -0,0 +1,90 @@ + + +# LingBot three-group, two-stage DiT benchmark + +## Result + +This experiment used GPU 0 for the shared encoder and decoder, three two-rank +pipeline-parallel DiT groups on GPUs 1–6, and left GPU 7 spare. Each DiT rank +owns 20 of the model's 40 transformer blocks and only the KV cache for those +blocks. + +| Metric | Batch 1 | Fixed batch 2, rerun | Double-buffered batch 2 | Full-DiT replica reference | +| --- | ---: | ---: | ---: | ---: | +| Sessions | 3 | 6 | 6 | 7 | +| Aggregate generated FPS | 16.80 | 16.77 | **21.16** | **35.15** | +| Generated FPS per session | **5.60** | 2.79 | **3.53** | 5.02 | +| Median wave latency | **2140.73 ms** | 4286.95 ms | **3401.46 ms** | 2358.51 ms | +| P90 wave latency | **2147.37 ms** | 4313.83 ms | **3406.93 ms** | 2454.10 ms | +| Median DiT denoise | 1647.70 ms | 3307.89 ms | **2615.76 ms** | 1718.78 ms | +| Median finalization | 408.37 ms | 819.46 ms | **646.98 ms** | 426.45 ms | +| Maximum DiT-rank required HBM | **28.70 GiB** | 39.63 GiB | **39.47 GiB** | 56.51 GiB | +| Node-wide required HBM | **188.44 GiB** | 257.29 GiB | **256.21 GiB** | 415.02 GiB | +| Median 256 MiB intra-pair P2P | 345.43 GB/s | 345.33 GB/s | 344.40 GB/s | — | + +Pipeline sharding reduces the largest DiT-rank footprint by 30.1% with double +buffering and 49.2% at batch one relative to the full-DiT replica. The +double-buffered layout uses 42.70 GiB per session; six independent full +pipelines require 399.30 GiB, so the equal-session memory saving is 35.8%. + +Double buffering is a material improvement over the matched fixed-batch run: +aggregate throughput rises **26.2%**, median wave latency falls **20.7%**, DiT +denoise time falls **20.9%**, and finalization falls **21.0%**. Maximum DiT-rank +HBM changes by only -0.16 GiB. The implementation keeps two separate batch-one +session caches, sends session 0 to stage 1, computes session 1 on stage 0 while +stage 1 processes session 0, then drains both outputs. The unchanged 344–345 +GB/s P2P result confirms that the gain comes from compute overlap, not a +transfer change. + +This still does not beat full-DiT replicas for maximum H100-node throughput. +The double-buffered layout is 39.8% below the seven-replica reference, while +reducing the largest DiT-rank footprint by 30.1%. It is therefore useful when +the full DiT and its session cache do not fit on a smaller GPU, or when lower +per-rank HBM matters more than maximum throughput. Each four-step denoise and +the final cache-update forward still pays one fill and one drain; a deeper +pipeline or more slots would increase bubbles and is not automatically better. + +## Reproduction + +The double-buffer and matched fixed runs used Slurm job `15002340` on +`pool0-01858`, eight H100 80 GB GPUs with NV18 connectivity between every pair, +revision `542190c0ae4134cf5e0da24342687a5328e93ac9` plus this worktree change, +PyTorch 2.12.1+cu130, CUDA 13.0, and Mooncake RDMA. Install the missing RDMA +userspace libraries in the allocation's writable container layer if needed: + +```bash +apt-get update +apt-get install -y libibverbs1 ibverbs-providers rdma-core +``` + +From the mounted repository checkout, run the double-buffered experiment: + +```bash +env GLOG_minloglevel=2 \ + TORCHINDUCTOR_COMPILE_THREADS=1 \ + TORCH_NCCL_SHOW_EAGER_INIT_P2P_SERIALIZATION_WARNING=0 \ +uv run --no-sync --package flashdreams-lingbot torchrun \ + --standalone --nproc_per_node=8 \ + -m lingbot.disagg.benchmark_pipeline \ + --sessions-per-group 2 \ + --double-buffered \ + --compile-network \ + --warmup-blocks 6 \ + --measured-blocks 5 \ + --bandwidth-probe-iters 10 \ + --output-dir outputs/lingbot_disagg_pipeline_3x2_double_buffered +``` + +Remove `--double-buffered`, retain `--sessions-per-group 2`, and use +`outputs/lingbot_disagg_pipeline_3x2_fixed_rerun` for the matched control. Six +warmup blocks are required because block 5 changes the cache shape and otherwise +puts compilation in the measurement window. Verify that `nvidia-smi topo -m` +reports NV18, Mooncake logs `installTransport, type=rdma`, and each result +contains three P2P pairs, five measured records, and 72 decoded frames in every +measured wave. + +The exact extracted values are in [summary.json](summary.json). Full per-wave +records remain in the generated `outputs/` benchmark directories. diff --git a/integrations/lingbot/docs/benchmark_h100_pipeline_3x2/summary.json b/integrations/lingbot/docs/benchmark_h100_pipeline_3x2/summary.json new file mode 100644 index 00000000..99a007c7 --- /dev/null +++ b/integrations/lingbot/docs/benchmark_h100_pipeline_3x2/summary.json @@ -0,0 +1,152 @@ +{ + "date": "2026-07-31", + "slurm_job_id": "14799647", + "hostname": "pool0-01062", + "base_commit": "0c2d48a8249577fb617bb5280208dd77409d9b1a", + "double_buffer_experiment": { + "date": "2026-08-03", + "slurm_job_id": "15002340", + "hostname": "pool0-01858", + "base_commit": "542190c0ae4134cf5e0da24342687a5328e93ac9" + }, + "topology": { + "io_rank": 0, + "dit_groups": [[1, 2], [3, 4], [5, 6]], + "spare_ranks": [7], + "pipeline_layers": [[0, 20], [20, 40]], + "stage_transport": "mooncake-rdma", + "dit_internal_transport": "nccl-p2p-nvlink" + }, + "workload": { + "model": "lingbot-world-fast-taehv-window15-sink3", + "resolution": [464, 832], + "warmup_waves": 6, + "measured_waves": 5, + "frames_per_session_per_wave": 12, + "compiled": true, + "cuda_graph": false + }, + "batch_1": { + "sessions": 3, + "aggregate_fps": 16.800616024879105, + "per_session_fps": 5.600205341626368, + "wave_latency_ms": { + "median": 2140.726614743471, + "p90": 2147.373372875154 + }, + "dit_ms": { + "median": 1647.7008056640625, + "p90": 1653.2711669921875 + }, + "finalize_ms": { + "median": 408.3681945800781, + "p90": 411.4204345703125 + }, + "max_dit_required_gib": 28.7034330368042, + "io_required_gib": 16.304168224334717, + "node_required_gib": 188.43816900253296, + "required_gib_per_session": 62.81272300084432, + "p2p_probe_gbps": { + "median": 345.42995899403087, + "p90": 348.4906482238796 + } + }, + "batch_2": { + "sessions": 6, + "aggregate_fps": 17.23536955355762, + "per_session_fps": 2.8725615922596037, + "wave_latency_ms": { + "median": 4177.556775510311, + "p90": 4183.501142449677 + }, + "dit_ms": { + "median": 3249.234619140625, + "p90": 3265.1302734375 + }, + "finalize_ms": { + "median": 803.6607055664062, + "p90": 810.2365600585938 + }, + "max_dit_required_gib": 39.54760932922363, + "io_required_gib": 19.629554271697998, + "node_required_gib": 256.9132375717163, + "required_gib_per_session": 42.818872928619385, + "p2p_probe_gbps": { + "median": 344.6416852808726, + "p90": 347.3234014628046 + } + }, + "fixed_batch_2_rerun": { + "sessions": 6, + "aggregate_fps": 16.766955159781013, + "per_session_fps": 2.794492526630169, + "wave_latency_ms": { + "median": 4286.9481993839145, + "p90": 4313.834723830223 + }, + "dit_ms": { + "median": 3307.89306640625, + "p90": 3346.546435546875 + }, + "finalize_ms": { + "median": 819.4578857421875, + "p90": 828.49990234375 + }, + "max_dit_required_gib": 39.63343858718872, + "io_required_gib": 19.629554271697998, + "node_required_gib": 257.28546237945557, + "required_gib_per_session": 42.88091039657593, + "p2p_probe_gbps": { + "median": 345.3327388451691, + "p90": 347.8951340927382 + } + }, + "double_buffered_batch_2": { + "sessions": 6, + "aggregate_fps": 21.16041714965187, + "per_session_fps": 3.526736191608645, + "wave_latency_ms": { + "median": 3401.4590242877603, + "p90": 3406.9314090535045 + }, + "dit_ms": { + "median": 2615.755615234375, + "p90": 2647.992138671875 + }, + "finalize_ms": { + "median": 646.9761962890625, + "p90": 657.2934326171875 + }, + "max_dit_required_gib": 39.4732871055603, + "io_required_gib": 19.629554271697998, + "node_required_gib": 256.2064833641052, + "required_gib_per_session": 42.701080560684204, + "p2p_probe_gbps": { + "median": 344.40037308204535, + "p90": 347.2302920849744 + } + }, + "replicated_dit_reference": { + "sessions": 7, + "aggregate_fps": 35.1528845526397, + "per_session_fps": 5.0218406503771, + "wave_latency_ms": { + "median": 2358.5133550077444, + "p90": 2454.0991360030603 + }, + "max_dit_peak_gib": 56.5107626914978, + "node_peak_gib": 415.0186958312988 + }, + "comparisons": { + "batch_2_max_dit_memory_reduction": 0.30017562238327966, + "batch_2_throughput_reduction": 0.5097025529228331, + "batch_2_equal_session_full_pipeline_memory_reduction": 0.3565909402160873, + "batch_2_vs_batch_1_throughput_gain": 0.025877235897503903, + "double_buffer_vs_fixed_rerun_throughput_gain": 0.2620309977573911, + "double_buffer_vs_fixed_rerun_latency_reduction": 0.20655467104160707, + "double_buffer_vs_fixed_rerun_dit_reduction": 0.20923815772673226, + "double_buffer_vs_fixed_rerun_finalize_reduction": 0.21048267696747757, + "double_buffer_max_dit_memory_reduction": 0.30149080944010753, + "double_buffer_throughput_reduction_vs_replicated_dit": 0.3980460659504287 + } +} diff --git a/integrations/lingbot/docs/disaggregated_inference.md b/integrations/lingbot/docs/disaggregated_inference.md new file mode 100644 index 00000000..233cad20 --- /dev/null +++ b/integrations/lingbot/docs/disaggregated_inference.md @@ -0,0 +1,212 @@ + + +# LingBot disaggregated inference design + +## Stage contract + +The monolithic `initialize_cache → generate → finalize` lifecycle becomes: + +1. The encoder worker creates session conditioning and its streaming VAE/camera + cache. It transfers text/image embeddings to the DiT worker once. +2. For every autoregressive block, the encoder worker transfers the I2V latent, + injection mask, and Plücker features to the DiT worker. +3. The DiT worker denoises and finalizes its resident autoregressive KV cache. + It transfers only the clean, unpatchified latent to the decoder. +4. The decoder worker advances its own streaming cache and produces pixels. + +The DiT worker must have session affinity. Its KV cache is not a standardized +LLM prefix cache, is mutated by `finalize`, and is deliberately excluded from +the transfer protocol. + +## Data plane + +`MooncakeTensorTransport` and `NixlTensorTransport` share a receiver-owned +buffer contract: + +1. The sender publishes tensor names, shapes, dtypes, and byte counts. +2. The receiver leases a fixed-shape contiguous device buffer from + `RegisteredTensorPool`; allocation and VRAM registration happen only when a + bucket is first created. +3. The receiver returns a reusable ticket. Mooncake tickets contain a session + and registered addresses; NIXL tickets contain agent metadata and serialized + transfer descriptors. +4. The producer records a CUDA event, submits a batched asynchronous write, + and retains the source storage until the transfer completes. +5. The consumer waits only at its dependency boundary. In the replicated + benchmark, encoder request N+1 overlaps encoder-to-DiT transfer N, and DiT + cache finalization overlaps clean-latent transfer to the decoder. + +No tensor is serialized into the Python control message and there is no +device-to-host-to-device staging in FlashDreams. + +## Scheduling + +The three-stage baseline is a fixed 1 encoder : 1 DiT : 1 decoder topology. +The first eight-GPU benchmark implemented a 1 encoder : 6 DiT : 1 decoder wave +scheduler. It starts with one replica per stage, then assigns each remaining +GPU to the stage with the largest measured service-time-per-replica: + +```text +stage capacity = replicas / median service time +system capacity = minimum stage capacity +``` + +The tracked baseline assigns all five additional GPUs to DiT because denoising +and cache finalization account for 97.58% of one-session latency. Each DiT +replica owns a distinct session and resident KV cache; the shared encoder feeds +six inputs, the DiTs execute concurrently, and the shared decoder drains six +latents. This scales concurrent-session throughput rather than one session's +latency. A production service still needs asynchronous bounded queues, request +IDs, cancellation, and failure recovery in place of benchmark-wide barriers; +the optimized benchmark now supplies the pooled registered-buffer and +transfer-ticket data path. + +The optimized eight-GPU topology co-locates encoder and decoder on GPU 0 and +uses GPUs 1–7 for seven session-affine DiT replicas. The measured I/O-stage +footprint fits safely on one H100, and each DiT retains its own cache. A +`SessionAwareScheduler` supplies the production placement policy: it filters by +service pool, shape, CP compatibility, free HBM, and verified RDMA support, +then scores compatible workers by predicted queue delay and topology locality. +Once placed, a session remains sticky; a missing resident worker is an explicit +recovery event, not permission to route a later chunk to a different cache. + +The scheduler separates two fixed pools: + +- `aggregated-cp8` for the minimum-latency service class; +- `io-plus-7-dit` for concurrent-session throughput. + +It also forms shape- and CP-compatible microbatch plans. Actual fused LingBot +DiT execution is intentionally not claimed yet: the model runtime still needs +batched cache gather/scatter and HBM/latency admission control before those +plans can share one kernel launch. + +The single-session topology uses the same physical allocation but groups ranks +1–6 into one context-parallel DiT: + +```text +GPU 0 encoder ── Mooncake ──▶ GPU 1 CP leader ══ NCCL CP6 ══ GPUs 2–6 +GPU 7 decoder ◀─ Mooncake ─── GPU 1 gathered clean latent +``` + +All six DiT ranks cooperate on every denoising step and retain their local +shard of the session cache. This can reduce one session's latency, unlike six +independent DiT replicas. The leader distributes full per-block input tensors +within the subgroup; the DiT's existing context-parallel path shards token +sequences and reconstructs its output before the decoder handoff. + +LingBot's 14B DiT has 40 attention heads. CP6 therefore uses ring attention: +Ulysses requires `num_heads % cp_size == 0`, and `40 % 6 != 0`. CP4 Ulysses is +a compatible alternative, but it uses only six of the allocated GPUs after +reserving separate encoder and decoder ranks. Ring rotation must use subgroup +local ranks because the DiT subgroup is global ranks 1–6. + +The aggregated baseline instead constructs the entire pipeline on all eight +ranks and makes WORLD the CP8 DiT group: + +```text +GPU 0 full pipeline ═══════════════════════════════════╗ +GPU 1–7 full pipeline ══ NCCL CP8 Ulysses ════════════╣ one session +each rank redundantly encodes and decodes; no RDMA stage handoffs ╝ +``` + +This topology can use Ulysses because 40 heads divide over eight ranks. It also +assigns all eight GPUs to DiT rather than reserving encoder and decoder ranks. +The cost is eight copies of the complete pipeline and no independent stage +scheduling. The 832×464 token grid is not CP8-compatible, so the tracked +aggregated experiment uses 832×448 and reports token throughput alongside FPS. + +## Measurement rules + +- Discard compilation, cache-fill, and connection warmup blocks. +- Report median and p90 compute and end-to-end chunk latency. +- Report generated frames divided by measured wall time, not target playback + FPS. +- Record sender memory registration separately from the Mooncake transfer. +- Report both real handoff payload bandwidth and a large-buffer link probe. +- Record the selected GPU/NIC topology. `protocol=rdma` is configuration, not + by itself evidence that an InfiniBand port carried the bytes. +- Compare decoded output against the aggregated pipeline with matched prompt, + first frame, camera path, checkpoint, seed, and decoder before recommending + the path as a default. + +## H100 validation + +The implementation was exercised in Slurm job `14621292` on `pool0-01299` +with three NVIDIA H100 80 GB HBM3 GPUs. Mooncake discovered the node's nine +`mlx5` HCAs, installed its RDMA transport, and completed RDMA ready handshakes. +After six warmup blocks, five measured LingBot blocks produced: + +- 5.36 generated FPS and 2233.57 ms median / 2250.79 ms p90 chunk latency; +- 1734.84 ms DiT denoise and 444.79 ms DiT cache finalization; +- 25.38 ms encoder → DiT handoff for 14.36 MiB and 12.05 ms DiT → decoder + handoff for 0.55 MiB; +- 0.13% of median latency in the two synchronous RDMA copy calls, or 1.68% + including allocation, metadata exchange, and synchronization; +- 41.35 GB/s encoder → DiT and 41.00 GB/s DiT → decoder on reusable 256 MiB + RDMA probes; +- 13.72 / 56.29 / 2.29 GiB peak allocated memory for encoder / DiT / decoder. + +The complete per-block measurements and reproduction command are in the +[H100 benchmark report](benchmark_h100_3stage/README.md). This path remains +opt-in: the stage boundaries and tensor round trips are covered by CPU tests +and the real LingBot rollout completed, but a matched-seed decoded-output +comparison with the original aggregated runner is still required before +making it the default serving path. + +The eight-GPU follow-up ran in Slurm job `14628860` on `pool0-00205`. Five +measured six-session waves produced **27.20 aggregate generated FPS**, a +**5.07× throughput gain** over the tracked 1:1:1 result and **1.90× higher +throughput per allocated GPU**. Median wave latency was 2657.06 ms, or 1.19× +the single-session baseline latency; median per-session throughput was +4.53 FPS. The twelve reusable 256 MiB Mooncake probes measured 41.22 GB/s +median across all stage edges. Peak allocations were 18.77 GiB on the shared +encoder, 56.34–56.51 GiB on each DiT, and 2.65 GiB on the shared decoder. +See the [eight-GPU report](benchmark_h100_1e6d1d/README.md) and the +[wall-time and memory chart](disaggregated_inference_breakdown.svg). + +The minimum-single-session experiment ran in Slurm job `14646820` on +`pool0-01260`. One CP6 ring DiT group on GPUs 1–6 reduced median chunk latency +to **743.27 ms** and raised one-session throughput to **15.90 generated FPS**: +a **3.01× end-to-end speedup** and **3.19× DiT speedup** over CP1. The CP4 +Ulysses comparison reached 754.41 ms and 15.70 FPS. CP6 was 1.5% faster and +used all eight GPUs; CP4 achieved higher per-DiT-GPU scaling efficiency (78.2% +versus 53.2%) and left two GPUs idle. See the +[CP6 report](benchmark_h100_cp6_single_session/README.md), +[CP4 report](benchmark_h100_cp4_single_session/README.md), and +[single-session wall-time and memory chart](disaggregated_inference_single_session.svg). + +The aggregated follow-up ran in Slurm job `14652956` on `pool0-01714`. Eight +full-pipeline replicas formed one CP8 Ulysses group. At 832×448, five measured +blocks reached **393.33 ms median / 434.08 ms p90** and **29.50 generated FPS**. +That is 1.89× lower latency and 1.79× higher DiT token throughput than the +tracked stage-local CP6 run. Rollout peak allocation was 40.88 GiB per GPU and +327.03 GiB node-wide; the one-shot initialization peak was 48.27 GiB per GPU. +The NCCL probes measured 266.83 GB/s broadcast and 360.38 GB/s all-gather +effective bandwidth. See the +[aggregated report](benchmark_h100_aggregated_cp8/README.md) and +[aggregated-versus-disaggregated chart](aggregated_vs_disaggregated.svg). + +The +[full experiment record](disaggregated_inference_experiment.md) +documents the tested stack, warmup behavior, stage and transfer breakdowns, +Slurm reproduction procedure, interpretation, and deferred validation. + +The optimized follow-up ran in Slurm job `14761875` on `pool0-01924`. With the +encoder and decoder co-located, seven independent DiTs, pooled registrations, +reusable tickets, asynchronous Mooncake writes, and clean-latent transfer +overlapped with cache finalization, it reached **35.15 aggregate generated +FPS** and **5.02 FPS per session**. That is 29.2% above the tracked 1:6:1 +result and 11.4% above the same 1:7 topology using per-request synchronous +handoffs. The fourteen 256 MiB stage-edge probes measured **41.97 GB/s** +median. GPU 0 peaked at 20.59 GiB and the DiT GPUs at 56.29–56.51 GiB. See the +[optimized report](benchmark_h100_1io7dit_optimized/README.md) and +[optimization chart](disaggregated_inference_optimized.svg). + +## References + +- [LightX2V: Breaking the Memory and Throughput Bottlenecks of Diffusion Model Inference](https://light-ai.top/LightX2V-BLOG/posts/Disaggregation/) +- [Mooncake Transfer Engine](https://github.com/kvcache-ai/Mooncake) +- [NIXL architecture](https://github.com/ai-dynamo/nixl/blob/main/docs/nixl.md) diff --git a/integrations/lingbot/docs/disaggregated_inference_breakdown.svg b/integrations/lingbot/docs/disaggregated_inference_breakdown.svg new file mode 100644 index 00000000..85268a3f --- /dev/null +++ b/integrations/lingbot/docs/disaggregated_inference_breakdown.svg @@ -0,0 +1,109 @@ + +LingBot disaggregated inference wall time and GPU memory +Stacked wall-time comparison for one versus six DiT workers and per-rank peak allocated GPU memory. + + +LingBot stage allocation: wall time and memory +H100 80 GB · BF16 · 832×464 · six warmup and five measured waves +Median steady-state wall time + +0 + +500 + +1000 + +1500 + +2000 + +2500 + +3000 +milliseconds +1E : 1D : 1V + + + +2180 + + + +2234 ms +1E : 6D : 1V + + +206 + +2186 + +185 + + +2657 ms + +Encoder + +Encoder → DiT + +DiT + +DiT → decoder + +Decoder + +Coordination +Peak allocated GPU memory by rank +E = encoder, D = DiT, V = decoder; the gap to 80 GiB is headroom, not free schedulable memory + +0 + +20 + +40 + +60 + +80 +GiB +1E : 1D : 1V + +13.7 +E0 + +56.3 +D1 + +2.3 +V2 +1E : 6D : 1V + +18.8 +E0 + +56.5 +D1 + +56.5 +D2 + +56.3 +D3 + +56.3 +D4 + +56.5 +D5 + +56.5 +D6 + +2.7 +V7 + diff --git a/integrations/lingbot/docs/disaggregated_inference_experiment.md b/integrations/lingbot/docs/disaggregated_inference_experiment.md new file mode 100644 index 00000000..499d46e2 --- /dev/null +++ b/integrations/lingbot/docs/disaggregated_inference_experiment.md @@ -0,0 +1,1200 @@ + + +# LingBot three-stage disaggregation experiment + +## Status + +**Useful opt-in.** The experiment validates that encoder, DiT, and decoder +components can run on separate H100 GPUs and exchange their real LingBot tensor +payloads through Mooncake's RDMA transport. The eight-GPU follow-up validates +six concurrent, session-affine DiT workers behind a shared encoder and decoder. +It does not establish output-quality equivalence with the aggregated runner, +cross-node performance, or production scheduler behavior, so it is not a new +default serving path. + +The machine-readable measurements are in +[`benchmark_h100_3stage/benchmark.json`](benchmark_h100_3stage/benchmark.json); +the generated compact table is in +[`benchmark_h100_3stage/README.md`](benchmark_h100_3stage/README.md). The +eight-GPU measurements are in +[`benchmark_h100_1e6d1d/benchmark.json`](benchmark_h100_1e6d1d/benchmark.json) +with its +[`generated summary`](benchmark_h100_1e6d1d/README.md). The combined +[wall-time and memory chart](disaggregated_inference_breakdown.svg) is generated +from those two JSON files. The single-session CP4 and CP6 results are in +[`benchmark_h100_cp4_single_session`](benchmark_h100_cp4_single_session/README.md) +and +[`benchmark_h100_cp6_single_session`](benchmark_h100_cp6_single_session/README.md); +their [comparison chart](disaggregated_inference_single_session.svg) is generated +directly from the raw JSON documents. The full-pipeline CP8 baseline is in +[`benchmark_h100_aggregated_cp8`](benchmark_h100_aggregated_cp8/README.md), with +its [stage-local CP6 comparison chart](aggregated_vs_disaggregated.svg). The +optimized co-located 1-I/O : 7-DiT result is in +[`benchmark_h100_1io7dit_optimized`](benchmark_h100_1io7dit_optimized/README.md), +with its [wall-time and memory chart](disaggregated_inference_optimized.svg). + +## Question and design + +The experiment asked whether the monolithic LingBot inference pipeline could be +split into independently placed services without making inter-GPU transfer the +new bottleneck. The tested topology was: + +```text +rank 0 / GPU 0 rank 1 / GPU 1 rank 2 / GPU 2 +text + image + VAE + camera ───▶ scheduler + DiT + KV cache ───▶ LightTAE + 14.36 MiB / block 0.55 MiB / block +``` + +The evolving autoregressive KV cache stays on the DiT worker. Session +conditioning crosses the encoder-to-DiT boundary once; each block then sends +the I2V latent, mask, and Plücker features to DiT and sends only the clean latent +to the decoder. Mooncake uses receiver-allocated, registered VRAM and a batched +synchronous write, so tensor contents are not serialized through the Python +control plane or staged through host memory. + +This is diffusion pipeline-stage disaggregation, analogous in scheduling intent +to separating prefill and decode pools in an LLM server, but the boundaries and +resident state are diffusion-native. + +## Tested configuration + +| Item | Value | +| --- | --- | +| Date | 2026-07-29 | +| Slurm allocation | Job `14621292`, node `pool0-01299` | +| Repository base | `e580e27d408b3cf8bd8a549f990c361b94d3379f`; the implementation under test was the worktree change recorded with this report | +| Container | `flashdreams-base-v0.3-20260429-af40a4f.sqsh` | +| GPUs used | 3 × NVIDIA H100 80 GB HBM3, one process per GPU | +| GPU topology | GPU 0/1/2 connected by NVLink; node exposed nine `mlx5` HCAs | +| Driver | Not captured by the original harness; the reproduction checklist below captures it | +| Python / PyTorch | Python 3.12, PyTorch `2.12.1+cu130` | +| CUDA / cuDNN | CUDA 13.0, cuDNN 92000 | +| Mooncake | `mooncake-transfer-engine-cuda13==0.3.12.post1` | +| Model | `lingbot-world-fast-taehv-window15-sink3` | +| DiT checkpoint | `robbyant/lingbot-world-fast`, `diffusion_pytorch_model.safetensors.index.json` | +| Precision | BF16 model and image tensors; FP32 camera tensors | +| Scheduler | Four-step distilled flow matching, CFG 1.0, seed 42 | +| Streaming layout | 3 latent frames per chunk, temporal window 15, sink 3 | +| Decoder | LightTAE / TAEHV | +| Input | Upstream LingBot example `00`: `image.jpg`, `intrinsics.npy`, `poses.npy`, and `prompt.txt` | +| Resolution / target playback | 832 × 464, 16 FPS | +| Process state | Fresh `torchrun` process with a persistent, previously populated Triton cache | +| Measurement policy | 6 warmup blocks, then 5 measured blocks; 8 iterations per 256 MiB link probe | + +The exact prompt was: + +> The video presents a soaring journey through a fantasy jungle. The wind +> whips past the rider's blue hands gripping the reins, causing the leather +> straps to vibrate. The ancient gothic castle approaches steadily, its stone +> details becoming clearer against the backdrop of floating islands and distant +> waterfalls. + +## Method + +1. Each rank constructed only its stage-local weights. Model construction + happened before `torch.distributed` initialization because LingBot otherwise + interprets a three-rank process group as three-way context parallelism. +2. The ranks initialized a Gloo control group and independent Mooncake P2P + endpoints with `protocol=rdma`. +3. Mooncake discovered nine `mlx5` HCAs, installed its RDMA transport, and + completed RDMA-ready handshakes. +4. Before model measurements, each edge ran one connection warmup followed by + eight synchronous 256 MiB transfers through reusable registered buffers. +5. The model ran eleven autoregressive blocks. Blocks 0–5 were excluded because + they cover initial model/compiler warmup and the second DiT cache-shape + transition at block 5. Blocks 6–10 were measured. +6. CUDA events measured stage compute. Host wall clocks measured synchronous + transfer calls, complete handoffs, and barrier-to-barrier chunk latency. +7. Throughput was computed from the 60 generated measured frames divided by the + sum of the five measured chunk latencies. It is generated throughput, not the + configured 16 FPS playback target. + +Warmup must remain at six blocks or more for this preset. Shorter trials observed +one-time DiT compilation at block 5 and produced misleading 13–24 second +outliers in the measured set. + +## Findings + +### Steady-state compute and latency + +| Component | Median | P90 | Share of median chunk latency | +| --- | ---: | ---: | ---: | +| Encoder compute | 1.08 ms | 1.14 ms | 0.05% | +| Encoder → DiT full handoff | 25.38 ms | 25.88 ms | 1.14% | +| DiT denoise | 1734.84 ms | 1755.00 ms | 77.67% | +| DiT cache finalization | 444.79 ms | 446.72 ms | 19.91% | +| DiT → decoder full handoff | 12.05 ms | 14.73 ms | 0.54% | +| Decoder compute | 7.14 ms | 7.15 ms | 0.32% | +| End-to-end chunk | 2233.57 ms | 2250.79 ms | 100% | + +The five measured blocks generated 12 frames each. Total generated throughput +was **5.36 FPS**. DiT denoising plus cache finalization consumed 97.58% of median +chunk latency, so additional DiT replicas—not encoder or decoder replicas—are +the first resource-scaling lever for concurrent sessions. + +### Transfer behavior + +| Edge | Real payload | Copy median | Payload bandwidth median | Full handoff median | 256 MiB probe median | +| --- | ---: | ---: | ---: | ---: | ---: | +| Encoder → DiT | 14.36 MiB | 1.35 ms | 11.12 GB/s | 25.38 ms | 41.35 GB/s | +| DiT → decoder | 0.55 MiB | 1.49 ms | 0.39 GB/s | 12.05 ms | 41.00 GB/s | + +The two transfer-engine copy calls used 0.13% of median chunk latency. Complete +handoffs—including receiver allocation and registration, sender registration, +metadata broadcasts, barriers, and the copy—used 1.68%. The small clean-latent +payload cannot saturate the link, which explains its low payload GB/s despite +the 41 GB/s large-buffer probe. + +The full handoff gap is mostly setup and synchronization rather than byte +movement. Production serving should pool and reuse registered destination +buffers and replace global barriers/object broadcasts with request-scoped +control messages. + +### Peak allocated GPU memory + +| Stage GPU | Peak allocated memory | Share of three-stage total | +| --- | ---: | ---: | +| Encoder | 13.72 GiB | 18.98% | +| DiT | 56.29 GiB | 77.86% | +| Decoder | 2.29 GiB | 3.17% | + +Disaggregation makes the imbalance explicit: encoder and decoder GPUs have +substantial unused capacity while the DiT GPU owns most weights and resident +state. A production scheduler should allow encoder and decoder workers to serve +multiple session-affine DiT workers. + +## Eight-GPU stage allocation + +### Allocation rule + +Starting with one GPU per stage, the scheduler repeatedly assigns an available +GPU to the stage with the largest baseline service time divided by its current +replica count. The baseline service-time inputs include stage compute and the +handoff that feeds that stage: + +| Stage | Baseline service time | Replicas after each of five assignments | Final allocation | +| --- | ---: | --- | ---: | +| Encoder | 26.46 ms | 1, 1, 1, 1, 1 | 1 | +| DiT + finalize | 2179.63 ms | 2, 3, 4, 5, 6 | 6 | +| Decoder | 19.19 ms | 1, 1, 1, 1, 1 | 1 | + +This produces **1 encoder : 6 DiT : 1 decoder**. Each DiT GPU owns one +independent session and persistent KV cache. The benchmark processes a wave by +encoding six inputs on GPU 0, running the six DiTs concurrently on GPUs 1–6, +and decoding six outputs on GPU 7. It measures throughput scaling for concurrent +sessions; it does not tensor-parallelize a single DiT or reduce a single +session's autoregressive dependency. + +### Eight-GPU tested configuration + +The follow-up used the same model, prompt, resolution, precision, scheduler, +warmup policy, and transfer probe as the three-GPU baseline, with these +differences: + +| Item | Value | +| --- | --- | +| Date | 2026-07-29 | +| Slurm allocation | Job `14628860`, node `pool0-00205` | +| Repository revision | `08d4c6c159321221c9a2d213c5ebb1359f443ef0` plus the replicated-benchmark worktree change | +| GPUs used | 8 × NVIDIA H100 80 GB HBM3, one process per GPU | +| Driver | 535.216.03 | +| Topology | GPU 0 encoder; GPUs 1–6 DiT workers; GPU 7 decoder | +| Sessions | Six independent sessions per wave | +| Measurements | 6 warmup waves followed by 5 measured waves | + +### Throughput and latency + +| Metric | 1:1:1 baseline | 1:6:1 wave | Change | +| --- | ---: | ---: | ---: | +| Aggregate generated FPS | 5.36 | 27.20 | **5.07×** | +| Generated FPS per allocated GPU | 1.79 | 3.40 | **1.90×** | +| Session/wave median latency | 2233.57 ms | 2657.06 ms | 1.19× | +| Session/wave p90 latency | 2250.79 ms | 2671.08 ms | 1.19× | +| Per-session generated FPS | 5.36 | 4.53 | 0.85× | + +Six-way DiT replication converts the dominant serial service into concurrent +capacity. Aggregate throughput scales to 84.6% of the ideal six-replica gain. +The remaining gap is visible in the sequential wave scheduler: six +encoder-to-DiT handoffs consume 205.94 ms per median wave and six +DiT-to-decoder handoffs consume 184.63 ms. A production implementation should +overlap those request-scoped transfers with DiT execution and reuse registered +buffers. + +### Eight-GPU wall-time and memory + +![LingBot component wall-time and GPU-memory breakdown](disaggregated_inference_breakdown.svg) + +| Component | 1:1:1 median | 1:6:1 median wave | +| --- | ---: | ---: | +| Encoder compute | 1.08 ms | 4.91 ms | +| Encoder → DiT handoff | 25.38 ms | 205.94 ms total | +| DiT critical path | 2179.63 ms | 2185.88 ms | +| DiT → decoder handoff | 12.05 ms | 184.63 ms total | +| Decoder compute | 7.14 ms | 42.26 ms | +| Coordination / residual | 8.29 ms | 33.40 ms | +| End-to-end | 2233.57 ms | 2657.06 ms | + +| Rank and role | Peak allocated memory | +| --- | ---: | +| GPU 0, shared encoder | 18.77 GiB | +| GPU 1–6, DiT workers | 56.34–56.51 GiB each | +| GPU 7, shared decoder | 2.65 GiB | + +The shared encoder grows by about 5 GiB relative to the baseline because it +owns six streaming encoder caches and participates in all six input edges. The +DiT ranks remain near the baseline's 56.29 GiB, confirming that each session's +resident state stays local rather than being copied between workers. + +### Eight-GPU transfer behavior + +All twelve 256 MiB probes—GPU 0 to each DiT and each DiT to GPU 7—completed +through Mooncake's configured RDMA transport. The combined median was +**41.22 GB/s**, p90 was **42.27 GB/s**, and the observed range was +33.89–42.49 GB/s. The one low sample occurred on encoder → DiT rank 2; that +edge's median remained 41.12 GB/s. + +During teardown after the successful report was written, Mooncake emitted +non-fatal `remote access error` and rail-pause messages while registrations +were being removed. `torchrun` exited with status 0 and every measured transfer +call had returned success, but this is still a buffer-lifetime warning: +production code must pool registrations and introduce an explicit drain before +unregistering memory or closing endpoints. The current benchmark result should +not be interpreted as failure-recovery validation. + +## Minimum single-session latency + +### Topologies and compatibility + +The replicated 1:6:1 topology above cannot accelerate one session because each +DiT rank owns a different request. The latency experiment instead made the DiT +ranks one context-parallel group. The encoder sends one Mooncake handoff to the +CP leader, the leader broadcasts the full step input within the NCCL subgroup, +and the existing Wan context-parallel path splits the 4524-token sequence. +Only the leader sends the reconstructed clean latent to the decoder. + +Two viable candidates were measured: + +| Candidate | Processes used | Compatibility | +| --- | ---: | --- | +| 1 encoder : CP4 DiT : 1 decoder | 6 of 8 | Ulysses; 40 heads and 4524 tokens are divisible by 4 | +| 1 encoder : CP6 DiT : 1 decoder | 8 of 8 | Ring; 4524 tokens are divisible by 6 | + +CP6 Ulysses was rejected before measurement because LingBot has 40 attention +heads and Ulysses requires `num_heads % cp_size == 0`; `40 % 6 != 0`. CP5 is +also unsuitable at this resolution because `4524 % 5 != 0`. CP4 is therefore +the largest configuration that satisfies both Ulysses head partitioning and +the token-sharding constraint while reserving distinct encoder and decoder +GPUs. + +The DiT subgroup occupies global ranks 1–6 for CP6, not ranks 0–5. The +experiment exposed and fixed a ring-rotation bug that used global +`DeviceMesh.get_rank()` to index subgroup all-gather output. Ring attention now +uses `get_local_rank()`, so every rank visits each subgroup K/V shard exactly +once. + +### Tested configuration + +Both trials ran in Slurm job `14646820` on `pool0-01260` with the same model, +prompt, resolution, precision, scheduler, six warmup blocks, five measured +blocks, and eight 256 MiB transfer probes as the baseline. CP6 used GPUs 0–7; +CP4 used GPUs 0–5 and left GPUs 6–7 idle. The software stack was Python 3.12.13, +PyTorch 2.12.1+cu130, CUDA 13.0, driver 535.216.03, and Mooncake +0.3.12.post1. The repository base was +`66bcd32ece1d03b3362d71a7340691a0687a4069` plus the worktree changes recorded +by this report. + +Each DiT rank initially launched 32 Inductor compiler workers. At CP6 that +created 192 workers and overcommitted the host during cold compilation. The +reported runs set `TORCHINDUCTOR_COMPILE_THREADS=4`, limiting the six DiT ranks +to 24 compiler workers. Compile/autotune time is excluded from the measured +steady-state blocks. + +### Latency and throughput + +![LingBot single-session context-parallel wall-time and memory breakdown](disaggregated_inference_single_session.svg) + +| Metric | CP1 baseline | CP4 Ulysses | CP6 ring | +| --- | ---: | ---: | ---: | +| Median chunk latency | 2233.57 ms | 754.41 ms | **743.27 ms** | +| P90 chunk latency | 2250.79 ms | 786.46 ms | **780.60 ms** | +| Generated FPS | 5.36 | 15.70 | **15.90** | +| End-to-end speedup | 1.00× | 2.96× | **3.01×** | +| DiT critical path | 2179.63 ms | 696.76 ms | **683.16 ms** | +| DiT speedup | 1.00× | 3.13× | **3.19×** | +| CP efficiency | — | **78.2%** | 53.2% | + +CP6 ring is the measured minimum: it is 11.13 ms, or 1.5%, faster than CP4 +Ulysses and uses all eight GPUs. CP4 is substantially more efficient per DiT +GPU and leaves two GPUs for other work, so it is the better capacity choice +when an 11 ms latency reduction is not worth two additional GPUs. Neither +configuration approaches ideal linear scaling because attention collectives, +duplicated non-attention work, and synchronization remain on every diffusion +step. + +The result answers the distinction behind the replicated experiment: +**1 encoder : 6 independent DiTs : 1 decoder** reached 27.20 aggregate FPS but +only 4.53 FPS per session, while **1 encoder : one CP6 DiT group : 1 decoder** +reached 15.90 FPS for one session. Replication increases request capacity; +context parallelism shortens the critical path of a single request. + +### Component, transfer, and memory breakdown + +| Component | CP1 | CP4 Ulysses | CP6 ring | +| --- | ---: | ---: | ---: | +| Encoder compute | 1.08 ms | 0.85 ms | 0.90 ms | +| Encoder → leader handoff | 25.38 ms | 31.70 ms | 30.33 ms | +| CP input fanout | — | 0.76 ms | 0.89 ms | +| DiT + finalize | 2179.63 ms | 696.76 ms | 683.16 ms | +| Leader → decoder handoff | 12.05 ms | 12.38 ms | 12.59 ms | +| Decoder compute | 7.14 ms | 7.07 ms | 7.13 ms | + +| Probe / allocation | CP4 Ulysses | CP6 ring | +| --- | ---: | ---: | +| Mooncake encoder → leader, 256 MiB | 41.44 GB/s | 42.62 GB/s | +| Mooncake leader → decoder, 256 MiB | 42.52 GB/s | 42.26 GB/s | +| NCCL broadcast, 256 MiB effective per rank | 307.31 GB/s | 283.71 GB/s | +| NCCL all-gather, 256 MiB effective per rank | 389.11 GB/s | 231.72 GB/s | +| Encoder peak allocated memory | 13.72 GiB | 13.72 GiB | +| DiT peak allocated memory per rank | 40.52–40.64 GiB | 39.18 GiB | +| Decoder peak allocated memory | 2.37 GiB | 2.29 GiB | + +The Mooncake stage boundaries remain small relative to the DiT critical path. +The CP probe values are effective materialized bytes per rank divided by +collective wall time; they are not raw physical-link throughput. CP6's lower +all-gather rate and lower scaling efficiency quantify the communication cost of +using all six DiT ranks. + +The five measured blocks in each trial emitted 12 frames and excluded both +initial compilation and the block-5 cache-shape transition. No decoded +aggregated-versus-CP quality comparison was performed, so these are performance +results, not output-parity evidence. The measured runs wrote their reports but +then stalled during process-group teardown. The final implementation drains the +NCCL subgroup with a CUDA synchronization and subgroup barrier before destroying +it, then destroys the Gloo control group. A subsequent one-block full-model +smoke and a transport-only run both exited normally. This cleanup-only change +does not affect the five-block timing samples above. + +## Fully aggregated single-H100 baseline + +The complete encoder, DiT, and LightTAE decoder were also measured in one +process on one H100 80 GB at the original 832×464 resolution. The run used six +warmup chunks followed by five measured 12-frame chunks. + +| Metric | Fully aggregated CP1 | +| --- | ---: | +| Generated FPS | **5.56** | +| Median / p90 chunk latency | **2157.51 / 2166.25 ms** | +| Initialization peak allocated HBM | **66.55 GiB** | +| Rollout peak allocated HBM | **59.36 GiB** | +| Steady allocated HBM | **57.15 GiB** | + +This proves the tested LingBot Fast configuration fits on one H100. Against the +same-resolution three-GPU stage-disaggregated CP1 result, aggregation improved +FPS by 3.6% and reduced median latency by 76.05 ms, or 3.4%. The small change +confirms that DiT compute and cache finalization, rather than stage handoffs, +dominate CP1 latency. See the +[single-H100 report](benchmark_h100_aggregated_cp1/README.md) and +[raw result](benchmark_h100_aggregated_cp1/benchmark.json). + +## Eight independent aggregated workers + +The next capacity experiment ran eight complete CP1 pipelines concurrently, +one per H100 and one per session. A coordinator held every worker after six +warmup chunks and released all eight into the same five-chunk measurement +window. The start skew was 0.32 ms. + +| Metric | Eight independent full pipelines | +| --- | ---: | +| Aggregate generated FPS | **43.44** | +| Median per-session FPS | **5.54** | +| Median / p90 chunk latency | **2163.64 / 2205.53 ms** | +| Rollout / initialization peak HBM per GPU | **59.35 / 66.55 GiB** | +| Rollout / initialization peak HBM node total | **474.84 / 532.38 GiB** | + +The shared-window result is 97.7% of linear scaling from the single-H100 run. +It is 23.6% faster in aggregate than 1 I/O + 7 DiTs, but uses 14.4% more +rollout peak node HBM because every GPU owns the encoder, full DiT, decoder, and +session cache. This makes full replication the best measured capacity topology +when every complete pipeline fits in one GPU. Disaggregation remains relevant +when it does not fit, or when stages need independent placement and scaling. +See the [eight-replica report](benchmark_h100_aggregated_8xcp1/README.md) and +[raw result](benchmark_h100_aggregated_8xcp1/benchmark.json). + +## Aggregated eight-GPU baseline + +### Topology and resolution + +The next experiment removed the stage boundaries and constructed the complete +encoder + DiT + decoder pipeline on every rank. WORLD is one CP8 Ulysses group: + +```text +GPU 0–7: full encoder + full DiT weights + full decoder + ║ + NCCL CP8 Ulysses + ║ + one session +``` + +Every rank executes the encoder and decoder redundantly. There is no Mooncake +or NIXL data path because no tensor crosses an independently deployed stage +boundary; all inter-GPU movement is inside the DiT's NCCL context-parallel +collectives. This is the conventional aggregated baseline against which the +stage-local deployment should be compared. + +The original 832×464 resolution produces 4,524 post-patch tokens. CP8 requires +an even split, and `4524 % 8 = 4`, so the benchmark rejects that shape. The +closest valid height is 448: 832×448 produces 4,368 tokens, or 546 tokens per +rank. This is 3.45% fewer tokens than the earlier experiments. FPS and latency +are reported directly, while DiT token throughput is the resolution-normalized +comparison. + +### Tested configuration + +| Item | Value | +| --- | --- | +| Date | 2026-07-29 | +| Slurm allocation | Job `14652956`, node `pool0-01714` | +| Repository base | `bb67d2868babea53cda7a9027831f26ee948293f` plus the aggregated-benchmark worktree change | +| GPUs | 8 × NVIDIA H100 80 GB HBM3, all-to-all NV18 connectivity | +| Driver / runtime | Driver 535.216.03; PyTorch 2.12.1+cu130; CUDA 13.0 | +| Topology | Eight full-pipeline replicas; WORLD CP8 Ulysses | +| Resolution / tokens | 832×448; 4,368 tokens per chunk; 546 tokens per rank | +| Model / precision / scheduler | Same LingBot Fast + LightTAE, BF16, four-step configuration | +| Compilation | Persistent Triton cache; `TORCHINDUCTOR_COMPILE_THREADS=4` | +| Measurement | Six warmup blocks, then five measured 12-frame blocks | + +The first attempt completed all model measurements but hit a report-only +missing comparison label and then blocked in distributed teardown. After adding +the label fallback, the benchmark was rerun from a fresh `torchrun` process. +Only the successful second run is recorded. Its warmup blocks include cached +code loading and the block-5 cache-shape compilation; none of those times enter +the headline. + +### Performance and resource comparison + +![Aggregated CP8 versus disaggregated CP6 wall time and HBM](aggregated_vs_disaggregated.svg) + +| Metric | Stage-local CP6 ring | Aggregated CP8 Ulysses | Change | +| --- | ---: | ---: | ---: | +| Resolution | 832×464 | 832×448 | 3.45% fewer tokens | +| Median chunk latency | 743.27 ms | **393.33 ms** | **1.89× faster** | +| P90 chunk latency | 780.60 ms | **434.08 ms** | **1.80× faster** | +| Generated FPS | 15.90 | **29.50** | **1.86×** | +| DiT token throughput | 5,995 token/s | **10,739 token/s** | **1.79×** | +| FPS per allocated GPU | 1.99 | **3.69** | **1.86×** | +| GPU-seconds per chunk | 5.95 | **3.15** | **47.1% lower** | +| Rollout peak HBM, node total | 251.07 GiB | 327.03 GiB | **30.3% higher** | + +The aggregated run is the new measured minimum single-session latency. Two +factors dominate the gain. First, it assigns all eight GPUs to DiT and can use +Ulysses because `40 % 8 == 0`; stage-local CP6 reserves two ranks for encoder +and decoder and must use ring attention. Second, it removes 43.81 ms of +encoder-to-leader fanout and leader-to-decoder handoff from the CP6 critical +path. The handoff removal alone is only 5.9% of CP6 latency, so most of the +1.89× gain comes from the larger, Ulysses-compatible DiT group rather than +from eliminating RDMA. + +| Component / probe | Stage-local CP6 | Aggregated CP8 | +| --- | ---: | ---: | +| Encoder critical-rank compute | 0.90 ms | 0.84 ms | +| Encoder handoff + CP fanout | 31.21 ms | — | +| DiT denoise | 547.88 ms | 309.14 ms | +| KV-cache finalize | 135.25 ms | 76.06 ms | +| Decoder compute | 7.13 ms | 6.88 ms | +| Output handoff | 12.59 ms | — | +| 256 MiB NCCL broadcast probe | 283.71 GB/s | 266.83 GB/s | +| 256 MiB NCCL all-gather probe | 231.72 GB/s | 360.38 GB/s | + +Probe bandwidth is effective materialized bytes per rank divided by collective +wall time, not raw NVLink wire bandwidth. It characterizes the allocation but +does not directly measure every all-to-all used by Ulysses. The node was fully +NVLink-connected; these numbers are not cross-node RDMA results. + +Each aggregated rank peaked at 40.88 GiB during the measured rollout and held +38.80 GiB after it. Replicating that footprint eight times consumes 327.03 GiB +of rollout peak allocation and 310.37 GiB steady allocation node-wide. The +one-shot text/model initialization peak was higher at 48.27 GiB per rank, or +386.13 GiB across the node. The stage-local CP6 layout uses only 251.07 GiB +node-wide because encoder and decoder weights exist on one rank each. + +This is a latency topology, not a blanket serving winner. It dedicates the +whole node to one session and cannot independently autoscale or place encoder, +DiT, and decoder services. The 1:6:1 topology remains the concurrent-session +choice: it serves six sessions at 27.20 aggregate FPS, while aggregated CP8 +serves one session at 29.50 FPS. Production scheduling can choose between +replicated stage-local groups for capacity and aggregated CP groups for a +premium low-latency class. + +## Disaggregation optimization follow-up + +### Scope and implementation + +The next experiment implemented the highest-value changes behind the earlier +recommendations: + +1. Encoder and decoder are co-located on GPU 0; GPUs 1–7 become seven + independent, session-affine DiT replicas. +2. Mooncake exposes `send_async` and an explicit transfer handle. A CUDA event + waits for only the producer stream, replacing the previous device-wide + synchronization. +3. `RegisteredTensorPool` allocates and registers fixed-shape receiver buckets + once. Their transfer tickets are exchanged once and reused for all blocks. +4. The encoder submits a transfer after each session's encode and continues + encoding later sessions before waiting. A DiT submits its clean latent + before cache finalization, then the decoder waits at its actual dependency. +5. The CP benchmark can patchify once on the encoder and RDMA the corresponding + token shard directly to each DiT rank, removing leader-mediated input + broadcast and its full-input materialization on every rank. +6. An interchangeable `NixlTensorTransport` implements the same + descriptor/ticket/handle contract. Its real backend remains optional. +7. `SessionAwareScheduler` adds fixed latency/throughput pools, sticky cache + placement, predicted queue-time scoring, shape/CP/HBM compatibility, + rack/NIC locality, and mandatory RDMA capability. +8. Compatible-session microbatch formation is implemented at the scheduler + layer. Fused model execution is not: LingBot still needs cache + gather/scatter, mask batching, and latency/HBM admission inside DiT. + +This follows the design principle reported by the +[LightX2V disaggregation study](https://light-ai.top/LightX2V-BLOG/posts/Disaggregation/): +retain large mutable state at its compute stage, transfer boundary tensors +instead of model state, and recover utilization by independently scaling the +dominant stage. The transport design also follows +[NIXL's recommendation](https://github.com/ai-dynamo/nixl/blob/main/docs/nixl.md) +to register transfer memory during initialization and reuse its metadata. + +### Reused allocation and method + +All follow-up trials reused Slurm job `14761875` on `pool0-01924`; no second +node was allocated. The node exposed eight H100 80 GB HBM3 GPUs and nine +`mlx5` HCAs. The container was missing `libibverbs.so.1`, so `libibverbs1`, +`ibverbs-providers`, and `rdma-core` were installed in its writable layer. +Mooncake then discovered all nine HCAs, installed `type=rdma`, completed +RDMA-ready handshakes, and did not fall back to TCP. + +The model, 832×464 input, seed, four-step scheduler, six warmup waves, five +measured waves, and 256 MiB eight-iteration probes match the tracked +disaggregated workload. The base revision was +`b762d079245681e1db70f1ffc5728753ce2a90b8` plus the changes recorded here. +`TORCHINDUCTOR_COMPILE_THREADS=1` limited host oversubscription. + +### Throughput result + +![Optimized LingBot component wall-time and HBM breakdown](disaggregated_inference_optimized.svg) + +| Metric | Tracked 1E:6DiT:1D | 1IO:7DiT synchronous | 1IO:7DiT pooled async | +| --- | ---: | ---: | ---: | +| Aggregate generated FPS | 27.20 | 31.57 | **35.15** | +| Per-session generated FPS | 4.53 | 4.51 | **5.02** | +| Median seven-session wave | — | 2671.78 ms | **2358.51 ms** | +| P90 seven-session wave | — | 2705.95 ms | **2454.10 ms** | +| Median DiT critical path | 2185.88 ms | 2177.66 ms | 2178.39 ms | +| Median encoder wave | 4.91 ms for six | 5.92 ms | 5.91 ms | +| Median decoder wave | 42.26 ms for six | 49.20 ms | 49.06 ms | +| Median 256 MiB stage-edge probe | 41.22 GB/s | 42.28 GB/s | 41.97 GB/s | + +Co-location plus the seventh DiT raises aggregate throughput 29.2% over the +tracked 1:6:1 result. Pooling and overlap then raise throughput 11.4% over the +same 1:7 topology's synchronous path. The optimized result is 10.9% above the +earlier 31.7 FPS linear projection. The stable five-wave range was +2354.46–2514.78 ms; p90 includes one slower 2.51 s measured wave. + +The DiT critical path did not become faster; it remains about 2.18 seconds. +The gain comes from removing serialized handoff work from the wave's critical +path and adding one independent session. This result therefore increases +concurrent interactive-session capacity, not one session's FPS. The aggregated +CP8 topology remains the measured one-session latency winner. + +| Optimized component | Median | P90 | +| --- | ---: | ---: | +| Encoder compute, seven inputs | 5.91 ms | 6.08 ms | +| DiT denoise, per-worker samples | 1718.78 ms | 1816.21 ms | +| DiT cache finalization, per-worker samples | 426.45 ms | 436.32 ms | +| DiT critical path, slowest worker | 2178.39 ms | 2256.77 ms | +| Decoder compute, seven outputs | 49.06 ms | 49.11 ms | +| End-to-end wave | 2358.51 ms | 2454.10 ms | + +GPU 0 peaked at 20.59 GiB with both I/O stages and seven pairs of streaming +caches. DiT ranks peaked at 56.29–56.51 GiB each. Co-location therefore fits +with about 59 GiB of HBM headroom on the I/O GPU while keeping every DiT below +57 GiB. + +### Transfer interpretation and correctness + +The optimized fourteen-edge large-buffer probe measured 41.97 GB/s median, +42.66 GB/s p90, and 39.66–42.79 GB/s range. This is the reliable link +measurement. For asynchronous model transfers, `transfer_ms` spans submission +until the later wait. Useful encoder work or cache finalization happens inside +that interval, so it is an in-flight residency window, not isolated byte-copy +time or an additive latency component. The 0.55 MiB output's roughly 442 ms +window, for example, intentionally contains roughly 426 ms of finalization. + +The synchronous control emitted repeated Mooncake `remote access error`, +`local access violation`, rail-pause, and rail-recovery messages when +per-request registered buffers were quickly unregistered and their addresses +reused. It still wrote a report and exited with status zero, but it is not a +production-safe control. The pooled run completed without those errors. +Persistent registration is therefore both the performance optimization and +the required receiver-buffer lifetime fix. + +### Direct CP shard result + +The direct CP6 input contract also completed all model measurements. The +encoder patchified once, split 4,524 tokens into six 754-token bundles, and +each CP rank consumed its own bundle without a leader broadcast. DiT compute +was unchanged, confirming the sharded tensor contract: + +| CP6 ring metric | Leader handoff + NCCL fanout | Six direct Mooncake shards | +| --- | ---: | ---: | +| Median chunk latency | **743.27 ms** | 939.36 ms | +| Generated FPS | **15.90** | 12.70 | +| Encoder input handoff | **30.33 ms** | 174.23 ms | +| NCCL input fanout | 0.89 ms | **0.00 ms** | +| DiT critical path | 683.16 ms | 683.19 ms | +| DiT peak HBM per rank | 39.18 GiB | **39.16 GiB** | + +The direct layout saves a full input replica and 0.89 ms of NCCL fanout, but +six separately allocated, registered, ticketed, and synchronized 2.39 MiB +Mooncake handoffs cost 174.23 ms. It regresses latency by 26.4% and throughput +by 20.2%, so `--direct-cp-input` is an experimental correctness path, not a +recommended optimization on Mooncake 0.3.12. + +The full direct run wrote its five-wave report, then stalled in distributed +teardown and was manually terminated. Attempts to reuse registered destination +shards—both concurrent and sequential, including an explicit +post-registration barrier—failed with Mooncake status `-1`, `remote access +error`, and `local access violation`. That attempted pooling code was not kept. +A production direct-shard path needs a transport with reliable long-lived +multi-destination registrations, or a Mooncake fix, before it can outperform +the leader + NCCL path. + +The allocated image did not contain NIXL, so the adapter was validated with a +fake-agent CPU round trip but not benchmarked on H100. No NIXL bandwidth result +is claimed. NVIDIA's +[Dynamo communication guide](https://docs.nvidia.com/dynamo/kubernetes-deployment/operate/disagg-communication) +recommends independently checking GPU-to-GPU bandwidth and verifying that the +NIXL/UCX backend is instantiated rather than silently accepting TCP; the same +acceptance rule is encoded by the scheduler. + +## Two-stage pipeline-parallel DiT follow-up + +### Implementation and topology + +The next prototype partitions each 40-block DiT across a fixed two-GPU group. +The first rank retains blocks 0–19 and the input patch embedding; the second +retains blocks 20–39 and the output head. Each rank constructs only its local +autoregressive cache. Unowned blocks and endpoint modules are removed before +the module moves to CUDA, so this is model-state sharding rather than a full +DiT replica on both ranks. + +```text +GPU 0 shared encoder + decoder +GPU 1–2 DiT group A: blocks 0–19 -> blocks 20–39 +GPU 3–4 DiT group B: blocks 0–19 -> blocks 20–39 +GPU 5–6 DiT group C: blocks 0–19 -> blocks 20–39 +GPU 7 spare +``` + +GPU 0 sends each group's batched encoder payload to the group leader through +the pooled asynchronous Mooncake path. The two ranks fan out the small common +input with NCCL, then exchange the DiT boundary activation directly through +NCCL point-to-point operations over NVLink/NVSwitch. The fixed batch-two path +sends an 88.36 MiB forward activation; each double-buffer slot sends 44.18 MiB. +The returned output is 0.55 MiB per slot. The clean latent returns from each +leader to GPU 0 through Mooncake. Large mutable KV state never crosses a rank +boundary. + +`LingbotDiTStage.configure_pipeline_parallel()` exposes the stage contract, +and the transformer rejects configuration after cache initialization or while +CUDA graph capture is enabled. TorchInductor compilation remains supported; +NCCL calls form graph breaks. The benchmark fixes session placement for the +whole rollout and reuses registered receive buffers and transfer tickets. + +### Double-buffered six-session result + +The new schedule keeps two independent, session-affine batch-one caches in +each DiT group. For every DiT forward, stage 0 computes session 0 and sends its +hidden activation. It then computes session 1 while stage 1 processes session +0. A bidirectional NCCL exchange hands session 1 to stage 1 while returning the +session-0 output, after which the pipeline drains session 1. The same schedule +also overlaps the final cache-update forward. + +| Metric | Fixed batch 2 | Double-buffered batch 2 | 1 I/O + 7 full DiTs | +| --- | ---: | ---: | ---: | +| Concurrent sessions | 6 | 6 | 7 | +| Aggregate generated FPS | 16.77 | **21.16** | **35.15** | +| Generated FPS per session | 2.79 | **3.53** | 5.02 | +| Median / p90 wave latency | 4286.95 / 4313.83 ms | **3401.46 / 3406.93 ms** | 2358.51 / 2454.10 ms | +| Median DiT denoise / finalization | 3307.89 / 819.46 ms | **2615.76 / 646.98 ms** | 1718.78 / 426.45 ms | +| Maximum DiT-rank required HBM | 39.63 GiB | **39.47 GiB** | 56.51 GiB | +| I/O-rank required HBM | 19.63 GiB | 19.63 GiB | 20.59 GiB | +| Node-wide required HBM | 257.29 GiB | **256.21 GiB** | 415.02 GiB | +| Median 256 MiB intra-pair P2P bandwidth | 345.33 GB/s | 344.40 GB/s | — | + +Against the matched fixed run on the same node, double buffering improves +aggregate and per-session FPS by **26.2%** and reduces median wave latency by +**20.7%**. Median denoise and finalization times fall **20.9%** and **21.0%**, +respectively. Peak memory is effectively unchanged: the maximum DiT rank drops +from 39.63 to 39.47 GiB and node-wide required HBM drops 0.4%. + +The double-buffered topology requires **256.21 GiB** node-wide, or 42.70 GiB +per session. Six independent full pipelines require 399.30 GiB at the measured +66.55 GiB initialization peak, so the equal-session memory saving is **35.8%**. +The largest DiT rank uses **30.1%** less HBM than a full-DiT replica. The +rollout peak, not the 38.81 GiB cache-initialization peak, determines minimum +GPU capacity. + +This remains a memory/placement tradeoff, not the fastest H100-node layout. +Double buffering is 39.8% below the seven-full-DiT aggregate reference and +29.8% below it per session. The reference also serves one more session. The +schedule hides the middle pipeline bubble but every four-step denoise and the +final cache update still require a fill and drain. An isolated 256 MiB P2P +probe remained at 344.40 GB/s and input fanout remained 0.33 ms, so NVLink is +not the limiting resource. + +GPU 7 was intentionally left spare. Moving encoder or decoder work to it would +not materially improve this result because their combined measured work is +only about 46 ms of a 3.40-second wave. Use this topology to fit memory-limited +GPUs; use full replicas when the model fits and maximum node throughput is the +goal. + +### Reproduction + +The matched experiment reused Slurm job `15002340` on `pool0-01858`, with eight +H100 80 GB GPUs and all GPU pairs reported as NV18. The same RDMA userspace +packages listed above were installed, and Mooncake discovered nine HCAs and +installed its RDMA transport. From revision +`542190c0ae4134cf5e0da24342687a5328e93ac9` plus this worktree change, run: + +```bash +env GLOG_minloglevel=2 \ + TORCHINDUCTOR_COMPILE_THREADS=1 \ + TORCH_NCCL_SHOW_EAGER_INIT_P2P_SERIALIZATION_WARNING=0 \ +uv run --no-sync --package flashdreams-lingbot torchrun \ + --standalone --nproc_per_node=8 \ + -m lingbot.disagg.benchmark_pipeline \ + --sessions-per-group 2 \ + --double-buffered \ + --compile-network \ + --warmup-blocks 6 \ + --measured-blocks 5 \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 10 \ + --output-dir outputs/lingbot_disagg_pipeline_3x2_double_buffered +``` + +Remove `--double-buffered` and write to +`outputs/lingbot_disagg_pipeline_3x2_fixed_rerun` for the matched control. Six +warmup blocks cover the cache-shape transition at block 5. Each output must +contain five non-warmup records, 72 decoded frames per measured wave, six +nonzero DiT-rank memory records, and three P2P probe pairs (`1->2`, `3->4`, and +`5->6`). See the tracked +[pipeline-parallel benchmark report](benchmark_h100_pipeline_3x2/README.md) +and [summary](benchmark_h100_pipeline_3x2/summary.json). + +### Deployment decision and remaining work + +Use disaggregation when enough simultaneous interactive sessions exist to keep +multiple DiTs occupied, I/O and DiT need independent scaling, cache-affine +routing or stage fault isolation matters, and a verified fast data plane is +available. Use aggregated CP8 when a single session needs minimum latency, +traffic is sparse, or operational simplicity outweighs independent stage +placement. Mixed deployments should keep two warm pools: `aggregated-cp8` for +premium latency and `io-plus-7-dit` for throughput. Do not hot-repartition a +node between requests because weight loading, compilation, cache construction, +and warmup are expensive. + +The next optimization items, in order, are: + +1. Replace the benchmark's wave barriers with request-scoped bounded queues so + encoder N+1, DiT N, and decoder N−1 overlap across whole stage invocations. +2. Add completion timestamps from Mooncake/NIXL rather than treating delayed + waits as copy latency. +3. Generalize the two-slot schedule to bounded, shape-compatible session + queues without combining their session-affine caches. +4. Remove remaining cache-layout materialization inside LingBot finalization; + its 426 ms median is still 18% of the optimized wave. +5. Complete direct rank-to-rank CP output shards where the decoder layout + permits it, instead of gathering the clean latent on a leader. +6. Install and benchmark NIXL/UCX with registration caching, multi-rail + selection, GPUDirect RDMA verification, and telemetry. +7. Feed measured GPU/NIC affinity and topology domains into the production + worker snapshots, rather than relying on benchmark-wide HCA autodiscovery. +8. Add cancellation drain, worker-loss cache recovery, bounded backpressure, + and matched-seed decoded-output quality regression. + +## How to reproduce + +The commands below assume the repository is at +`/home/gtong/lustre/flashdreams-dist` and the cluster helper is +`/home/gtong/work/srun.sh`. + +### 1. Allocate or reuse a compute node + +Run on the login node: + +```bash +squeue -u gtong +cd /home/gtong/work +export FLASHDREAMS_HOST_DIR=/home/gtong/lustre/flashdreams-dist +./srun.sh +``` + +If `squeue` already shows a running job, attach to that exact allocation: + +```bash +./srun.sh 1 +``` + +Do not run model loading or package synchronization on the login node. + +### 2. Verify the mounted checkout and hardware + +Inside the container shell: + +```bash +hostname +pwd +git rev-parse --show-toplevel +git rev-parse HEAD +git status --short +nvidia-smi -L +nvidia-smi --query-gpu=name,driver_version --format=csv,noheader +nvidia-smi topo -m +ibv_devices +``` + +The Git toplevel and revision must match the intended checkout, and at least +three GPUs must be visible. In job `14646820`, `/workspace/flashdreams` pointed +at a different checkout; the experiment explicitly changed to the canonical +Lustre checkout before running. Do not assume a convenient container mount is +the requested branch. Preserve the driver and topology output with the +benchmark artifacts. + +### 3. Install RDMA userspace support and the optional transport + +The writable container used for the experiment needed: + +```bash +apt-get update +apt-get install -y libibverbs1 ibverbs-providers librdmacm1 ibverbs-utils +uv sync --package flashdreams-lingbot --extra dev --extra disagg +``` + +The optional `disagg` extra installs the CUDA 13 Mooncake wheel. A missing +`libibverbs.so.1` means the OS packages above are absent. + +### 4. Validate the transport before loading checkpoints + +```bash +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 +``` + +Confirm that Mooncake logs `installTransport, type=rdma`, discovers the expected +HCAs, completes RDMA-ready handshakes, and reports finite bandwidth in both +directions. Protocol configuration alone is not evidence that bytes traversed +the intended RDMA path. + +For the eight-GPU topology, probe all twelve stage edges: + +```bash +uv run --package flashdreams-lingbot torchrun \ + --standalone --nproc_per_node=8 \ + -m lingbot.disagg.benchmark_replicated \ + --transport-only \ + --dit-replicas 6 \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 +``` + +For the single-session CP6 topology, probe both Mooncake edges and the NCCL +subgroup collectives: + +```bash +uv run --package flashdreams-lingbot torchrun \ + --standalone --nproc_per_node=8 \ + -m lingbot.disagg.benchmark_cp \ + --transport-only \ + --cp-ranks 6 --cp-method ring \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 +``` + +### 5. Run the model benchmark + +Keep the default persistent `TRITON_CACHE_DIR` supplied by `srun.sh`, then run: + +```bash +uv run --package flashdreams-lingbot torchrun \ + --standalone --nproc_per_node=3 \ + -m lingbot.disagg.benchmark \ + --model lingbot-world-fast-taehv-window15-sink3 \ + --example-idx 0 \ + --pixel-width 832 \ + --pixel-height 464 \ + --fps 16 \ + --warmup-blocks 6 \ + --measured-blocks 5 \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 \ + --output-dir outputs/lingbot_disagg_h100 +``` + +The first run downloads checkpoints and example assets and may compile kernels. +For a warm-cache steady-state comparison, rerun in a fresh `torchrun` process +without clearing `TRITON_CACHE_DIR`. For a cold-start study, use a new explicit +cache directory and report startup/compile latency separately; do not mix it +into the steady-state rows. + +Run the eight-GPU concurrent-session benchmark in the same allocation: + +```bash +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 \ + --example-idx 0 \ + --pixel-width 832 \ + --pixel-height 464 \ + --fps 16 \ + --warmup-blocks 6 \ + --measured-blocks 5 \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 \ + --baseline-json integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json \ + --output-dir outputs/lingbot_disagg_1e6d1d +``` + +The command requires exactly `dit-replicas + 2` ranks. It rejects an allocation +that disagrees with the greedy recommendation derived from `--baseline-json`. +Model loading fans the DiT checkpoint out to six ranks and may take several +minutes from shared storage; that cold-start time is not included in the +steady-state throughput result. + +Run the optimized co-located topology in the same allocation: + +```bash +env GLOG_minloglevel=2 TORCHINDUCTOR_COMPILE_THREADS=1 \ +uv run --no-sync --package flashdreams-lingbot torchrun \ + --standalone --nproc_per_node=8 \ + -m lingbot.disagg.benchmark_replicated \ + --dit-replicas 7 \ + --co-locate-io \ + --pooled-async \ + --transport mooncake \ + --warmup-blocks 6 \ + --measured-blocks 5 \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 \ + --output-dir outputs/lingbot_disagg_1io7dit_optimized +``` + +Use `--transport-only` first on a new node and reject the run unless the logs +show an RDMA transport and the fourteen edge probes complete. The +`--transport nixl` option selects the NIXL adapter, but requires a separately +installed compatible NIXL release. + +Run the all-eight-GPU, minimum-latency candidate: + +```bash +env 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 \ + --example-idx 0 \ + --pixel-width 832 \ + --pixel-height 464 \ + --fps 16 \ + --warmup-blocks 6 \ + --measured-blocks 5 \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 \ + --baseline-json integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json \ + --output-dir outputs/lingbot_disagg_cp6 +``` + +Run the largest Ulysses-compatible comparison: + +```bash +env 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 \ + --example-idx 0 \ + --pixel-width 832 \ + --pixel-height 464 \ + --fps 16 \ + --warmup-blocks 6 \ + --measured-blocks 5 \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 \ + --baseline-json integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json \ + --output-dir outputs/lingbot_disagg_cp4 +``` + +The CP benchmark requires exactly `cp-ranks + 2` processes. Do not remove the +six warmup blocks: block 5 changes the cache shape and can otherwise leave +compile/autotune latency in the measured set. + +To reproduce the direct-shard diagnostic, add `--direct-cp-input` and use a +different output directory. Expect lower performance on Mooncake 0.3.12 as +reported above; do not use that flag as the production default. + +Run the eight-GPU aggregated baseline. The height must remain CP8-compatible +unless another resolution is explicitly validated: + +```bash +env 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 \ + --example-idx 0 \ + --pixel-width 832 \ + --pixel-height 448 \ + --fps 16 \ + --warmup-blocks 6 \ + --measured-blocks 5 \ + --bandwidth-probe-mib 256 \ + --bandwidth-probe-iters 8 \ + --comparison-json integrations/lingbot/docs/benchmark_h100_cp6_single_session/benchmark.json \ + --output-dir outputs/lingbot_aggregated_cp8 +``` + +All eight ranks instantiate the full pipeline, so expect eight simultaneous +checkpoint reads and up to 48.27 GiB of initialization allocation per GPU. +The harness requires exactly eight ranks and rejects token grids that do not +divide over CP8. No Mooncake installation is needed for this command. + +### 6. Inspect and preserve results + +```bash +sed -n '1,240p' outputs/lingbot_disagg_h100/README.md +jq '.environment, .summary' outputs/lingbot_disagg_h100/benchmark.json +jq '.records[] | select(.warmup == false)' \ + outputs/lingbot_disagg_h100/benchmark.json +``` + +The output directory contains the generated Markdown summary and raw per-block +JSON. Check that exactly five records have `warmup=false`, that every measured +block emits 12 frames, and that no measured latency contains a compile outlier. + +For the eight-GPU output, verify five measured records, 72 output frames per +wave, six `dit_workers` per record, and eight memory entries: + +```bash +jq '.environment.allocation, .environment.peak_memory_gib_by_rank, .summary' \ + outputs/lingbot_disagg_1e6d1d/benchmark.json +jq '.records[] | select(.warmup == false) | + {autoregressive_index, wave_latency_ms, output_frames, + dit_workers: (.dit_workers | length)}' \ + outputs/lingbot_disagg_1e6d1d/benchmark.json +``` + +For a CP result, verify the subgroup allocation, five measured records, one +worker record per CP rank, and one 12-frame output per block: + +```bash +jq '.environment.allocation, .environment.peak_memory_gib_by_rank, .summary' \ + outputs/lingbot_disagg_cp6/benchmark.json +jq '.records[] | select(.warmup == false) | + {autoregressive_index, end_to_end_ms, output_frames, + cp_workers: (.cp_workers | length)}' \ + outputs/lingbot_disagg_cp6/benchmark.json +``` + +For the aggregated result, verify five measured records, eight per-rank timing +entries, the CP8 token layout, and both rollout and initialization memory: + +```bash +jq '.environment.allocation, .environment.token_layout, .summary.memory, + .summary.cp_probe_gbps' \ + outputs/lingbot_aggregated_cp8/benchmark.json +jq '.records[] | select(.warmup == false) | + {autoregressive_index, end_to_end_ms, output_frames, + rank_records: (.per_rank | length)}' \ + outputs/lingbot_aggregated_cp8/benchmark.json +``` + +Regenerate the checked-in chart from the two raw result documents: + +```bash +python integrations/lingbot/scripts/plot_disagg_breakdown.py \ + integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json \ + integrations/lingbot/docs/benchmark_h100_1e6d1d/benchmark.json \ + integrations/lingbot/docs/disaggregated_inference_breakdown.svg + +python integrations/lingbot/scripts/plot_disagg_single_session.py \ + integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json \ + integrations/lingbot/docs/benchmark_h100_cp4_single_session/benchmark.json \ + integrations/lingbot/docs/benchmark_h100_cp6_single_session/benchmark.json \ + integrations/lingbot/docs/disaggregated_inference_single_session.svg + +python integrations/lingbot/scripts/plot_aggregated_comparison.py \ + integrations/lingbot/docs/benchmark_h100_cp6_single_session/benchmark.json \ + integrations/lingbot/docs/benchmark_h100_aggregated_cp8/benchmark.json \ + integrations/lingbot/docs/aggregated_vs_disaggregated.svg +``` + +### 7. Run focused CPU validation and release the node + +```bash +uv run --package flashdreams-lingbot pytest \ + flashdreams/tests/test_cp_attention_subgroup.py \ + flashdreams/tests/test_pipeline_stages.py \ + flashdreams/tests/test_transfer.py \ + integrations/lingbot/tests/test_disagg_stages.py \ + integrations/lingbot/tests/test_disagg_scheduler.py \ + integrations/lingbot/tests/test_disagg_replicated.py \ + integrations/lingbot/tests/test_disagg_cp.py \ + integrations/lingbot/tests/test_disagg_aggregated.py \ + integrations/lingbot/tests/test_disagg_independent.py + +uv lock --check +exit +``` + +The original focused validation completed with 17 passing tests; the later +aggregated CP1 and independent-replica harness validation completed with six +passing tests. A repository-wide +``pytest -m ci_cpu`` collection was also attempted, but the local environment +lacked the unrelated Omnidreams optional dependencies ``pyvirtualdisplay`` and +``flip_evaluator``. The changed LingBot and transport modules had already +completed collection and passed in the focused run above. + +## Acceptance and limitations + +| Decision | Status | Evidence / missing evidence | +| --- | --- | --- | +| Three independent GPU stages | Useful opt-in | Real LingBot rollout completed with stage-local weights and state | +| Mooncake RDMA data plane | Useful opt-in | RDMA transport/handshakes observed; 41 GB/s single-node probes | +| Eight-GPU throughput scaling | Useful opt-in | Six concurrent DiTs reached 27.20 FPS, 5.07× the 1:1:1 result | +| Default serving path | Deferred | Needs scheduler, bounded queues, cancellation, buffer pooling, and failure recovery | +| Output-quality equivalence | Deferred | No matched-seed aggregated-vs-disaggregated decoded comparison yet | +| Cross-node efficiency | Deferred | Only one eight-H100 node was measured | +| Single-session acceleration | Useful opt-in | CP6 ring reached 15.90 FPS and 743.27 ms, a 3.01× latency speedup | +| Aggregated CP8 baseline | Useful opt-in | 29.50 FPS and 393.33 ms at 832×448; 30.3% more node HBM than stage-local CP6 | +| Eight independent full pipelines | Useful opt-in | Eight sessions reached 43.44 aggregate FPS and 5.54 median FPS/session at 59.35 GiB rollout peak per GPU | +| Clean CP teardown | Useful opt-in | Post-fix full-model and transport-only lifecycle smokes exited normally | +| Co-located 1IO:7DiT topology | Useful opt-in | 35.15 aggregate FPS, 5.02 FPS/session, and 20.59 GiB on the I/O GPU | +| Pooled async Mooncake | Useful opt-in | 11.4% faster than same-topology synchronous control; no RDMA access errors | +| NIXL backend | Experimental | Contract and fake-agent CPU round trip pass; no NIXL package or H100 result on the allocated image | +| Session-aware scheduler | Experimental | Placement, affinity, compatibility, topology, and TCP-rejection tests pass; not wired to a production request frontend | +| Fused DiT microbatch | Deferred | Scheduler forms compatible groups; model/cache batch execution is not implemented | + +Results apply to this H100/CUDA 13/Mooncake stack and should not be generalized +to other GPU, NIC, driver, topology, or model configurations without repeating +the experiment. diff --git a/integrations/lingbot/docs/disaggregated_inference_optimized.svg b/integrations/lingbot/docs/disaggregated_inference_optimized.svg new file mode 100644 index 00000000..f94536ae --- /dev/null +++ b/integrations/lingbot/docs/disaggregated_inference_optimized.svg @@ -0,0 +1,78 @@ + + LingBot optimized disaggregated wall-time and HBM breakdown + Two-panel chart showing a 2358.51 millisecond seven-session wave dominated by DiT denoising and finalization, and peak allocated memory of 20.59 GiB on the co-located encoder decoder GPU and approximately 56.3 GiB on each of seven DiT GPUs. + + LingBot 1 I/O + 7 DiT optimization + 8× H100 · seven concurrent sessions · pooled async Mooncake RDMA + + + Median seven-session wave wall time + 2358.51 ms + + + + + + + + + + Encoder 5.91 ms + + DiT denoise 1718.78 ms + + Finalize 426.45 ms + + Decoder 49.06 ms + + Coordination 158.31 ms + DiT component values are medians across workers; finalization overlaps clean-latent RDMA, so transfer residency windows are not added. + + + Peak allocated HBM by GPU + GiB · 80 GiB device capacity + + + + + + + 0 + 20 + 40 + 60 + + + + + + + + + + + + + 20.59 + 56.34 + 56.34 + 56.51 + 56.29 + 56.34 + 56.34 + 56.29 + + + GPU 0 + GPU 1 + GPU 2 + GPU 3 + GPU 4 + GPU 5 + GPU 6 + GPU 7 + + encoder + decoder + seven session-affine DiT workers + Measured 2026-07-30 · Slurm 14761875 + diff --git a/integrations/lingbot/docs/disaggregated_inference_single_session.svg b/integrations/lingbot/docs/disaggregated_inference_single_session.svg new file mode 100644 index 00000000..17fcfb22 --- /dev/null +++ b/integrations/lingbot/docs/disaggregated_inference_single_session.svg @@ -0,0 +1,133 @@ + +LingBot single-session context-parallel wall time and GPU memory +Median wall-time components and per-rank peak allocated memory for CP1, CP4 Ulysses, and CP6 ring. + + +LingBot minimum single-session latency +H100 80 GB · BF16 · 832×464 · six warmup and five measured blocks +Median steady-state wall time + +0 + +500 + +1000 + +1500 + +2000 + +2500 +milliseconds +CP1 + + + +2180 + + + +2234 ms +CP4 Ulysses + + + +697 + + + +754 ms +CP6 ring + + + +683 + + + +743 ms + +Encoder + +Input handoff + +DiT + +Output handoff + +Decoder + +Coordination +Peak allocated GPU memory by rank +E = encoder, D = DiT, V = decoder; CP4 leaves two GPUs available for other work + +0 + +20 + +40 + +60 + +80 +GiB +CP1 + +13.7 +E0 + +56.3 +D1 + +2.3 +V2 +CP4 Ulysses + +13.7 +E0 + +40.6 +D1 + +40.6 +D2 + +40.5 +D3 + +40.6 +D4 + +2.4 +V5 +CP6 ring + +13.7 +E0 + +39.2 +D1 + +39.2 +D2 + +39.2 +D3 + +39.2 +D4 + +39.2 +D5 + +39.2 +D6 + +2.3 +V7 + diff --git a/integrations/lingbot/docs/disaggregation_experiment_summary.json b/integrations/lingbot/docs/disaggregation_experiment_summary.json new file mode 100644 index 00000000..5f5e7f0f --- /dev/null +++ b/integrations/lingbot/docs/disaggregation_experiment_summary.json @@ -0,0 +1,135 @@ +{ + "model": "lingbot-world-fast-taehv-window15-sink3", + "hardware": "NVIDIA H100 80GB HBM3", + "precision": "bfloat16", + "default_resolution": [464, 832], + "warmup_chunks": 6, + "measured_chunks": 5, + "single_session": [ + { + "topology": "three-stage-disaggregated-cp1", + "gpu_count": 3, + "sessions": 1, + "resolution": [464, 832], + "aggregate_fps": 5.360631332509123, + "median_latency_ms": 2233.5668401792645, + "p90_latency_ms": 2250.7897848263383, + "max_required_gib_per_gpu": 56.28736877441406 + }, + { + "topology": "aggregated-cp1", + "gpu_count": 1, + "sessions": 1, + "resolution": [464, 832], + "aggregate_fps": 5.55553754894227, + "median_latency_ms": 2157.513060141355, + "p90_latency_ms": 2166.248793900013, + "max_required_gib_per_gpu": 66.54739189147949, + "memory_basis": "initialization peak" + }, + { + "topology": "stage-disaggregated-cp4-ulysses", + "gpu_count": 6, + "sessions": 1, + "resolution": [464, 832], + "aggregate_fps": 15.695497779264775, + "median_latency_ms": 754.4078070059186, + "p90_latency_ms": 786.4595750012086, + "max_required_gib_per_gpu": 40.64252519607544, + "cp_efficiency": 0.7820537785981635 + }, + { + "topology": "stage-disaggregated-cp6-ring", + "gpu_count": 8, + "sessions": 1, + "resolution": [464, 832], + "aggregate_fps": 15.901998021609447, + "median_latency_ms": 743.272824001906, + "p90_latency_ms": 780.5990148001001, + "max_required_gib_per_gpu": 39.17708206176758, + "node_required_gib": 251.07237005233765, + "cp_efficiency": 0.5317502054854495 + }, + { + "topology": "aggregated-cp8-ulysses", + "gpu_count": 8, + "sessions": 1, + "resolution": [448, 832], + "aggregate_fps": 29.503035463925226, + "median_latency_ms": 393.33291398361325, + "p90_latency_ms": 434.0799169614911, + "max_required_gib_per_gpu": 48.26605987548828, + "memory_basis": "initialization peak", + "rollout_peak_gib_per_gpu": 40.879210472106934, + "node_rollout_peak_gib": 327.03368377685547, + "tokens_per_chunk": 4368 + } + ], + "multi_session": [ + { + "topology": "one-encoder-six-dits-one-decoder", + "gpu_count": 8, + "sessions": 6, + "resolution": [464, 832], + "aggregate_fps": 27.198381414481183, + "per_session_fps": 4.5330635690801975, + "median_wave_latency_ms": 2657.061628997326, + "p90_wave_latency_ms": 2671.0753187537193, + "max_required_gib_per_gpu": 56.5107626914978 + }, + { + "topology": "one-io-seven-dits-pooled-async", + "gpu_count": 8, + "sessions": 7, + "resolution": [464, 832], + "aggregate_fps": 35.1528845526397, + "per_session_fps": 5.0218406503771, + "median_wave_latency_ms": 2358.5133550077444, + "p90_wave_latency_ms": 2454.0991360030603, + "max_required_gib_per_gpu": 56.5107626914978, + "node_required_gib": 415.0186958312988, + "rdma_probe_gbps": 41.968 + }, + { + "topology": "eight-independent-aggregated-cp1", + "gpu_count": 8, + "sessions": 8, + "resolution": [464, 832], + "aggregate_fps": 43.43749180420828, + "per_session_fps": 5.544989798538596, + "median_wave_latency_ms": 2163.643025793135, + "p90_wave_latency_ms": 2205.5334428325295, + "max_required_gib_per_gpu": 66.54739189147949, + "memory_basis": "initialization peak", + "node_required_gib": 532.3791351318359, + "rollout_peak_gib_per_gpu": 59.35440921783447, + "node_rollout_peak_gib": 474.8352737426758 + }, + { + "topology": "three-two-rank-dit-groups-double-buffered", + "gpu_count": 8, + "sessions": 6, + "resolution": [464, 832], + "aggregate_fps": 21.16041714965187, + "per_session_fps": 3.526736191608645, + "median_wave_latency_ms": 3401.4590242877603, + "p90_wave_latency_ms": 3406.9314090535045, + "max_required_gib_per_gpu": 39.4732871055603, + "node_required_gib": 256.2064833641052, + "nvlink_probe_gbps": 344.40037308204535, + "fixed_batch_control_fps": 16.766955159781013, + "fixed_batch_control_latency_ms": 4286.9481993839145 + } + ], + "transport": { + "mooncake_large_buffer_probe_gbps": 41.968, + "pipeline_nvlink_probe_gbps": 344.40037308204535, + "nixl_gpu_validation": "not run; NIXL was absent from the allocated image" + }, + "comparability_notes": [ + "All headline values exclude six warmup chunks and use five measured chunks.", + "Aggregated CP8 uses 832x448 and 4368 tokens; the other headline runs use 832x464 and 4524 tokens.", + "Multi-session aggregate FPS rows contain different numbers of concurrent sessions.", + "Strict matched-seed decoded-quality equivalence across every topology remains deferred." + ] +} diff --git a/integrations/lingbot/docs/disaggregation_experiment_summary.md b/integrations/lingbot/docs/disaggregation_experiment_summary.md new file mode 100644 index 00000000..103dba90 --- /dev/null +++ b/integrations/lingbot/docs/disaggregation_experiment_summary.md @@ -0,0 +1,190 @@ + + +# LingBot disaggregated inference: experiment summary + +## Executive conclusion + +Disaggregation works, but it is not a general latency optimization for the +current LingBot interactive-video pipeline. + +- **Minimum single-session latency:** use the aggregated CP8 pipeline. It + reached **29.50 FPS** and **393.33 ms** median chunk latency. The comparable + stage-disaggregated CP6 run reached 15.90 FPS and 743.27 ms. +- **Maximum eight-H100 throughput:** use eight independent aggregated workers + when every GPU can hold the complete pipeline. They reached **43.44 aggregate + FPS** across eight sessions. +- **Shared-stage, multi-session serving:** the optimized 1 I/O + 7 DiT layout + reached **35.15 aggregate FPS** and 5.02 FPS per session. It is useful for + independent stage scaling and session-affine placement, but it is slower than + eight full replicas when those replicas fit. +- **Memory-constrained serving:** the double-buffered pipeline-parallel DiT + reduced the largest DiT rank to **39.47 GiB** and node-wide HBM to 256.21 GiB + for six sessions. The tradeoff is lower throughput: **21.16 aggregate FPS**. + +![LingBot latency, throughput, and memory tradeoffs](disaggregation_experiment_summary.svg) + +The chart deliberately separates single-session latency from concurrent-session +capacity. Aggregate FPS from multiple sessions must not be interpreted as one +session's frame rate. + +## What was implemented and tested + +The original `generate()` operation was separated into stateful encoder, DiT, +and decoder stages. Each worker retains only its session-affine state: + +- the encoder owns streaming VAE/camera state and produces I2V conditioning; +- the DiT owns the autoregressive KV cache and performs denoising/finalization; +- the decoder owns streaming VAE state and produces pixels. + +Small control messages stay outside the tensor data path. Large tensors use +Mooncake RDMA between stages, while context-parallel and pipeline-parallel DiT +ranks use NCCL over NVLink/NVSwitch. + +| Experiment | Purpose | Headline result | What it established | +| --- | --- | ---: | --- | +| Three-stage CP1 | Prove encoder → DiT → decoder separation | 5.36 FPS, 2233.57 ms | Stage separation is functional; DiT dominates latency | +| Aggregated CP1 | Same-shape control without stage handoffs | 5.56 FPS, 2157.51 ms | Disaggregation added 76.05 ms, or 3.4%, to one session | +| Stage-local CP4 Ulysses | Accelerate one session inside the DiT | 15.70 FPS, 754.41 ms | CP can accelerate a session, with 78.2% scaling efficiency | +| Stage-local CP6 ring | Use all six available DiT ranks | 15.90 FPS, 743.27 ms | Two extra ranks barely beat CP4 because ring efficiency fell to 53.2% | +| Aggregated CP8 Ulysses | Minimum-latency eight-GPU control | 29.50 FPS, 393.33 ms | Whole-pipeline CP is the measured latency winner | +| 1 encoder + 6 DiTs + 1 decoder | Scale independent sessions | 27.20 aggregate FPS | Replicated session-affine DiTs scale concurrent capacity | +| 1 co-located I/O + 7 DiTs | Use all eight GPUs and overlap transfers | 35.15 aggregate FPS | Pooling, asynchronous handoff, and the seventh DiT improved capacity | +| Eight full aggregated workers | Maximum-throughput control | 43.44 aggregate FPS | Replication wins when 66.55 GiB per GPU is available | +| Three 2-rank DiT pipelines | Reduce DiT memory per rank | 21.16 aggregate FPS, 39.47 GiB/rank | Pipeline sharding trades throughput for a lower memory floor | + +The final pipeline experiment also replaced a fixed batch with a two-slot +double-buffered schedule. Stage 0 processes session N+1 while stage 1 processes +session N. Against a matched fixed-batch run, this raised throughput from 16.77 +to **21.16 FPS** (+26.2%) and reduced median wave latency from 4286.95 to +**3401.46 ms** (-20.7%) without increasing peak HBM. + +## Measured serving tradeoffs + +### One interactive session + +| Topology | GPUs used | Generated FPS | Median latency | Key caveat | +| --- | ---: | ---: | ---: | --- | +| Three-stage disaggregated CP1 | 3 | 5.36 | 2233.57 ms | Adds stage handoffs but does not parallelize DiT | +| Fully aggregated CP1 | 1 | 5.56 | 2157.51 ms | Complete pipeline requires 66.55 GiB at initialization | +| Stage-disaggregated CP4 Ulysses | 6 | 15.70 | 754.41 ms | Four GPUs cooperate on one DiT | +| Stage-disaggregated CP6 ring | 8 | 15.90 | 743.27 ms | More ranks, but lower CP efficiency | +| Fully aggregated CP8 Ulysses | 8 | **29.50** | **393.33 ms** | Measured at 832x448, 3.45% fewer tokens than the 832x464 runs | + +Disaggregation itself did not reduce one-session latency. Context parallelism +did, but applying it to the entire aggregated pipeline was substantially faster +than reserving separate encoder and decoder GPUs. + +### Concurrent sessions on one eight-H100 node + +| Topology | Sessions | Aggregate FPS | FPS/session | Median wave | Maximum required HBM/GPU | Node HBM | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | +| Three 2-rank DiT groups, double buffered | 6 | 21.16 | 3.53 | 3401.46 ms | **39.47 GiB** | **256.21 GiB** | +| 1 co-located I/O + 7 full DiTs | 7 | 35.15 | 5.02 | 2358.51 ms | 56.51 GiB | 415.02 GiB | +| Eight independent full pipelines | 8 | **43.44** | **5.54** | **2163.64 ms** | 66.55 GiB initialization | 532.38 GiB initialization | + +The rows serve different session counts, so aggregate FPS is a node-capacity +metric rather than a strict per-request comparison. The direction is clear: +replication provides the best throughput, while deeper sharding lowers the +memory required by any one GPU. + +## Transfer and scheduling findings + +- Mooncake's reusable 256 MiB probe sustained about **42 GB/s**. The actual + small-payload handoff time was dominated by registration, metadata, and + control overhead rather than link bandwidth. +- Co-locating encoder and decoder, pooling registered buffers/tickets, and + using asynchronous handoffs improved the same seven-DiT topology from 31.57 + to **35.15 FPS** (+11.4%). Pooling also eliminated observed registration + lifetime errors. +- Pipeline-parallel DiT ranks sustained **344.40 GB/s** on the 256 MiB NVLink + probe. The fixed schedule was compute-bubble limited, not bandwidth limited; + double buffering recovered 26.2% throughput. +- Sending six direct CP input shards regressed CP6 from 15.90 to 12.70 FPS and + raised input handoff to 174.23 ms. That path remains experimental. +- A NIXL transport adapter exists behind the same tensor descriptor/ticket + contract, but the allocated image did not contain NIXL. No real NIXL GPU or + RDMA performance is claimed. + +## When disaggregation is useful + +Use it when: + +- the complete model and session cache cannot fit on the available GPU; +- many concurrent interactive sessions can keep multiple DiTs occupied; +- encoder, DiT, and decoder need independent scaling or fault isolation; +- session-affine placement can keep large KV/streaming caches resident; +- the deployment has verified RDMA, topology-aware placement, and stable + registered-buffer lifetimes. + +Avoid it when: + +- one session needs the lowest possible interactive latency; +- the complete pipeline fits on every GPU and maximum throughput is the goal; +- traffic is sparse, so stage workers would remain idle; +- operational simplicity matters more than independent stage scaling. + +## Limitations + +1. Results are steady-state measurements on one eight-H100 NVSwitch node. They + do not establish inter-node performance or behavior on smaller GPUs. +2. Six warmup chunks were excluded to cover compilation, autotuning, cache + fill, and the block-5 cache-shape transition. Startup latency is not part of + the headline numbers. +3. Most experiments use 832x464. Aggregated CP8 uses 832x448 because the larger + token count is not divisible by eight; token throughput is the fairest CP6 + versus CP8 compute comparison. +4. Full autoregressive decode completed successfully, but there is no strict + matched-seed visual-quality equivalence study across all topologies. +5. Benchmarks use fixed shapes, sticky session placement, and synchronized + waves. A production queue with variable arrival times and backpressure was + not measured. +6. The two-stage DiT schedule still pays fill and drain costs during every + denoise step and finalization. GPU 7 is unused in that topology. +7. Dynamo does not natively understand FlashDreams' encoder/DiT/decoder state + or three-stage tensor contract. Integration requires custom stateful workers, + a sticky session coordinator, and a direct NIXL/Mooncake tensor data plane. + +## Deployment recommendation + +Maintain separate, prewarmed pools instead of repartitioning a node at request +time: + +| Status | Request class | Recommended topology | Reason | +| --- | --- | --- | --- | +| **Recommended** | Premium single-session latency | Aggregated CP8 | Fastest measured interactive chunk latency | +| **Recommended** | Maximum H100-node throughput | Eight aggregated CP1 workers | Highest measured aggregate and per-session FPS | +| **Useful opt-in** | Shared-stage concurrent serving | 1 I/O + 7 DiTs | Independent scaling and cache-affine routing with moderate memory savings | +| **Useful opt-in** | Memory-constrained GPUs | Pipeline-parallel DiT with double buffering | Lowest measured per-DiT-rank HBM; accept lower throughput | +| **Rejected** | Direct CP input sharding | Six direct Mooncake shards | 20.1% lower FPS and 174.23 ms input handoff | +| **Deferred** | NIXL production data plane | NIXL/UCX adapter | Requires real GPU, GPUDirect RDMA, and inter-node validation | +| **Deferred** | Strict quality acceptance | Matched-seed decoded comparison | Full rollouts passed, but topology-wide visual equivalence was not measured | + +Kubernetes can deploy these fixed worker groups. Dynamo becomes valuable at +larger scale for worker discovery, cache-aware routing, autoscaling signals, +cancellation, and graceful shutdown. It does not provide the FlashDreams stage +split or make the video model faster by itself. + +## Evidence and reproduction + +The consolidated machine-readable values are in +[disaggregation_experiment_summary.json](disaggregation_experiment_summary.json). +Exact commands, commits, Slurm jobs, environments, warmup policies, and raw +records are retained in the individual benchmark reports: + +- [three-stage CP1](benchmark_h100_3stage/README.md) +- [aggregated CP1](benchmark_h100_aggregated_cp1/README.md) +- [stage-local CP4](benchmark_h100_cp4_single_session/README.md) and + [stage-local CP6](benchmark_h100_cp6_single_session/README.md) +- [aggregated CP8](benchmark_h100_aggregated_cp8/README.md) +- [six replicated DiTs](benchmark_h100_1e6d1d/README.md) +- [optimized seven-DiT serving](benchmark_h100_1io7dit_optimized/README.md) +- [eight independent aggregated workers](benchmark_h100_aggregated_8xcp1/README.md) +- [double-buffered pipeline-parallel DiT](benchmark_h100_pipeline_3x2/README.md) + +The full chronological engineering record remains in +[disaggregated_inference_experiment.md](disaggregated_inference_experiment.md). +The three-stage API and direct tensor-data-plane design follow the +[LightX2V disaggregation study](https://light-ai.top/LightX2V-BLOG/posts/Disaggregation/). diff --git a/integrations/lingbot/docs/disaggregation_experiment_summary.svg b/integrations/lingbot/docs/disaggregation_experiment_summary.svg new file mode 100644 index 00000000..f23a14c0 --- /dev/null +++ b/integrations/lingbot/docs/disaggregation_experiment_summary.svg @@ -0,0 +1,101 @@ + + LingBot disaggregation latency, throughput, and memory tradeoffs + The upper chart compares single-session median latency for disaggregated and aggregated topologies. The lower chart plots multi-session aggregate throughput against the maximum required memory on any GPU. + + + + LingBot serving tradeoff: latency, throughput, and memory + H100 80 GB · BF16 · six warmup + five measured chunks · generated FPS + + Disaggregated + + Aggregated + + Single-session median chunk latency + lower is better + + + + + + + + + + + 0 + 400 + 800 + 1200 + 1600 + 2000 + 2400 ms + + Stage-disagg CP1 + + 2234 ms + + Aggregated CP1 + + 2158 ms + + Stage-disagg CP6 ring + + 743 ms + + Aggregated CP8 Ulysses + + 393 ms + CP8 used 832×448 (3.45% fewer tokens); all other rows used 832×464. + + Concurrent-session capacity versus maximum required HBM per GPU + higher and farther left is better + + + + + + + + + + + 50 + 40 + 30 + 20 + 10 + 0 + aggregate generated FPS + + + + + + + 30 + 40 + 50 + 60 + 70 + maximum required HBM on any GPU (GiB) + + + 2-stage DiT pipeline + 6 sessions · 21.16 FPS · 39.47 GiB + + + 1 I/O + 7 DiTs + 7 sessions · 35.15 FPS · 56.51 GiB + + + 8 full pipelines + 8 sessions · 43.44 FPS · 66.55 GiB + diff --git a/integrations/lingbot/lingbot/disagg/__init__.py b/integrations/lingbot/lingbot/disagg/__init__.py new file mode 100644 index 00000000..86d48fb0 --- /dev/null +++ b/integrations/lingbot/lingbot/disagg/__init__.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Three-stage LingBot inference for disaggregated GPU serving.""" + +from lingbot.disagg.stages import ( + LingbotConditioning, + LingbotDecoderStage, + LingbotDiTStage, + LingbotEncoderStage, + conditioning_from_bundle, + conditioning_to_bundle, + encoder_output_from_bundle, + encoder_output_to_bundle, + encoder_output_to_cp_bundles, +) + +__all__ = [ + "LingbotConditioning", + "LingbotDecoderStage", + "LingbotDiTStage", + "LingbotEncoderStage", + "conditioning_from_bundle", + "conditioning_to_bundle", + "encoder_output_from_bundle", + "encoder_output_to_bundle", + "encoder_output_to_cp_bundles", +] diff --git a/integrations/lingbot/lingbot/disagg/benchmark.py b/integrations/lingbot/lingbot/disagg/benchmark.py new file mode 100644 index 00000000..384a34e4 --- /dev/null +++ b/integrations/lingbot/lingbot/disagg/benchmark.py @@ -0,0 +1,743 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Three-GPU LingBot disaggregation benchmark with Mooncake transfers.""" + +from __future__ import annotations + +import argparse +import importlib.metadata +import json +import os +import platform +import shlex +import socket +import statistics +import subprocess +import sys +import time +from collections.abc import Callable +from functools import partial +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch import Tensor + +from flashdreams.infra.transfer import ( + MooncakeTensorTransport, + TensorBundle, + TensorTransferTicket, + TransferStats, + describe_tensor_bundle, +) +from lingbot.config import PIPELINE_CONFIGS +from lingbot.disagg.stages import ( + LingbotDecoderStage, + LingbotDiTStage, + LingbotEncoderStage, + conditioning_from_bundle, + conditioning_to_bundle, + encoder_output_from_bundle, + encoder_output_to_bundle, +) +from lingbot.encoder.camctrl import CamCtrlInput +from lingbot.encoder.utils import get_Ks_transformed, preprocess_example_poses +from lingbot.runner import ( + _INTRINSICS_REFERENCE_HEIGHT, + _INTRINSICS_REFERENCE_WIDTH, + EXAMPLE_DATA_BASE_URL, + ensure_example_data_downloaded, +) +from lingbot.transformer import LingbotWorldTransformerConfig + +_ENCODER_RANK = 0 +"""Rank that owns text/image/VAE/camera encoders.""" + +_DIT_RANK = 1 +"""Rank that owns the scheduler, DiT, and session KV cache.""" + +_DECODER_RANK = 2 +"""Rank that owns the streaming VAE decoder.""" + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model", + choices=sorted(PIPELINE_CONFIGS), + default="lingbot-world-fast-taehv-window15-sink3", + ) + parser.add_argument("--example-idx", type=int, default=0) + parser.add_argument("--warmup-blocks", type=int, default=6) + parser.add_argument("--measured-blocks", type=int, default=5) + parser.add_argument("--pixel-height", type=int, default=464) + parser.add_argument("--pixel-width", type=int, default=832) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--rdma-device", default=None) + parser.add_argument("--bandwidth-probe-mib", type=int, default=256) + parser.add_argument("--bandwidth-probe-iters", type=int, default=8) + parser.add_argument( + "--transport-only", + action="store_true", + help="Run Mooncake bandwidth probes without loading model weights.", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("outputs/lingbot_disagg"), + ) + return parser.parse_args() + + +def _broadcast_object(value: Any, *, source: int) -> Any: + payload = [value if torch.distributed.get_rank() == source else None] + torch.distributed.broadcast_object_list(payload, src=source) + return payload[0] + + +def _timed_cuda(call: Callable[[], Any]) -> tuple[Any, float]: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + output = call() + end.record() + end.synchronize() + return output, start.elapsed_time(end) + + +def _transfer_bundle( + transport: MooncakeTensorTransport, + *, + source: int, + destination: int, + source_bundle: TensorBundle | None, + device: torch.device, +) -> tuple[TensorBundle | None, TransferStats, float]: + rank = torch.distributed.get_rank() + started = time.perf_counter() + source_descriptors = None + if rank == source: + assert source_bundle is not None + source_descriptors = describe_tensor_bundle(source_bundle) + descriptors = _broadcast_object( + source_descriptors, + source=source, + ) + destination_bundle = ( + transport.allocate(descriptors, device=device) if rank == destination else None + ) + ticket = _broadcast_object( + transport.make_ticket(destination_bundle) + if rank == destination and destination_bundle is not None + else None, + source=destination, + ) + stats = ( + transport.send(source_bundle, ticket) + if rank == source and source_bundle is not None + else None + ) + torch.distributed.barrier() + if rank == source and source_bundle is not None: + transport.unregister(source_bundle) + stats = _broadcast_object(stats, source=source) + source_handoff_ms = _broadcast_object( + (time.perf_counter() - started) * 1000.0 if rank == source else None, + source=source, + ) + return destination_bundle, stats, source_handoff_ms + + +def _bandwidth_probe( + transport: MooncakeTensorTransport, + *, + source: int, + destination: int, + size_mib: int, + iterations: int, + device: torch.device, +) -> list[TransferStats]: + rank = torch.distributed.get_rank() + numel = size_mib * 1024 * 1024 // torch.empty((), dtype=torch.uint8).element_size() + source_bundle = ( + {"probe": torch.empty(numel, dtype=torch.uint8, device=device)} + if rank == source + else None + ) + descriptors = _broadcast_object( + describe_tensor_bundle(source_bundle) if source_bundle is not None else None, + source=source, + ) + destination_bundle = ( + transport.allocate(descriptors, device=device) if rank == destination else None + ) + ticket: TensorTransferTicket = _broadcast_object( + transport.make_ticket(destination_bundle) + if destination_bundle is not None + else None, + source=destination, + ) + sender_stats: list[TransferStats] | None = None + if source_bundle is not None: + transport.register(source_bundle) + transport.send(source_bundle, ticket) # connection warmup + sender_stats = [ + transport.send(source_bundle, ticket) for _ in range(iterations) + ] + torch.distributed.barrier() + if source_bundle is not None: + transport.unregister(source_bundle) + if destination_bundle is not None: + transport.unregister(destination_bundle) + return _broadcast_object(sender_stats, source=source) + + +def _load_encoder_inputs( + args: argparse.Namespace, + *, + device: torch.device, +) -> tuple[str, Tensor, Tensor, Tensor, float]: + from flashdreams.infra.runner_io import load_first_frame_tensor + + example_dir = ensure_example_data_downloaded( + is_rank_zero=True, + example_idx=args.example_idx, + ) + prompt_path = example_dir / "prompt.txt" + prompt = ( + prompt_path.read_text().splitlines()[0].strip() if prompt_path.exists() else "" + ) + image = load_first_frame_tensor( + example_dir / "image.jpg", + pixel_height=args.pixel_height, + pixel_width=args.pixel_width, + device=device, + dtype=torch.bfloat16, + interpolation="cubic", + install_hint="Install the lingbot plugin.", + ) + intrinsics = torch.from_numpy(np.load(example_dir / "intrinsics.npy")).to( + device=device, + dtype=torch.float32, + ) + intrinsics = get_Ks_transformed( + intrinsics, + height_org=_INTRINSICS_REFERENCE_HEIGHT, + width_org=_INTRINSICS_REFERENCE_WIDTH, + height_resize=args.pixel_height, + width_resize=args.pixel_width, + height_final=args.pixel_height, + width_final=args.pixel_width, + ) + poses, world_scale = preprocess_example_poses(np.load(example_dir / "poses.npy")) + poses_tensor = torch.from_numpy(poses).to(device=device, dtype=torch.float32) + return prompt, image, intrinsics, poses_tensor, float(world_scale) + + +def _percentile(values: list[float], percentile: float) -> float: + return float(np.percentile(np.asarray(values, dtype=np.float64), percentile)) + + +def _metric_summary(values: list[float]) -> dict[str, float]: + return { + "median": statistics.median(values), + "p90": _percentile(values, 90), + "min": min(values), + "max": max(values), + } + + +def _environment( + args: argparse.Namespace, + *, + prompt: str | None, + world_size: int = 3, + module_name: str | None = None, +) -> dict[str, Any]: + config = PIPELINE_CONFIGS[args.model] + transformer = config.diffusion_model.transformer + assert isinstance(transformer, LingbotWorldTransformerConfig) + scheduler = config.diffusion_model.scheduler + try: + commit = subprocess.check_output( + ["git", "rev-parse", "HEAD"], + text=True, + ).strip() + worktree_dirty = bool( + subprocess.check_output( + ["git", "status", "--porcelain"], + text=True, + ).strip() + ) + except (OSError, subprocess.CalledProcessError): + commit = "unknown" + worktree_dirty = None + try: + driver = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader", + ], + text=True, + ).splitlines()[0] + except (OSError, subprocess.CalledProcessError, IndexError): + driver = "unknown" + try: + mooncake_version = importlib.metadata.version("mooncake-transfer-engine-cuda13") + except importlib.metadata.PackageNotFoundError: + mooncake_version = "unknown" + if module_name is None: + module_name = ( + __spec__.name if __spec__ is not None else "lingbot.disagg.benchmark" + ) + command_parts = [ + "uv", + "run", + "--package", + "flashdreams-lingbot", + "torchrun", + "--standalone", + f"--nproc_per_node={world_size}", + "-m", + module_name, + *sys.argv[1:], + ] + compile_threads = os.environ.get("TORCHINDUCTOR_COMPILE_THREADS") + if compile_threads is not None: + command_parts[:0] = [ + "env", + f"TORCHINDUCTOR_COMPILE_THREADS={compile_threads}", + ] + command = shlex.join(command_parts) + return { + "command": command, + "commit": commit, + "worktree_dirty": worktree_dirty, + "hostname": socket.gethostname(), + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "python": platform.python_version(), + "model": args.model, + "checkpoint": transformer.checkpoint_path, + "precision": str(transformer.dtype).removeprefix("torch."), + "decoder_config": type(config.decoder).__name__, + "seed": config.diffusion_model.seed, + "example_index": args.example_idx, + "example_url": f"{EXAMPLE_DATA_BASE_URL}/{args.example_idx:02d}", + "prompt": prompt, + "resolution": [args.pixel_height, args.pixel_width], + "target_fps": args.fps, + "latent_frames_per_chunk": transformer.len_t, + "window_size_t": transformer.window_size_t, + "sink_size_t": transformer.sink_size_t, + "guidance_scale": transformer.guidance_scale, + "num_inference_steps": getattr(scheduler, "num_inference_steps", None), + "compile_network": transformer.compile_network, + "warmup_blocks": args.warmup_blocks, + "measured_blocks": args.measured_blocks, + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "driver": driver, + "mooncake": mooncake_version, + "triton_cache_dir": os.environ.get("TRITON_CACHE_DIR"), + "torchinductor_compile_threads": compile_threads, + "hf_home": os.environ.get("HF_HOME"), + "gpus": [torch.cuda.get_device_name(index) for index in range(world_size)], + "rdma_device": args.rdma_device, + } + + +def _write_report( + args: argparse.Namespace, + *, + records: list[dict[str, Any]], + probe: dict[str, list[TransferStats]], + environment: dict[str, Any], +) -> None: + measured = [item for item in records if not item["warmup"]] + frame_count = sum(item["output_frames"] for item in measured) + total_seconds = sum(item["end_to_end_ms"] for item in measured) / 1000.0 + summary = { + "fps": frame_count / total_seconds, + "latency_ms": _metric_summary([item["end_to_end_ms"] for item in measured]), + "encoder_ms": _metric_summary([item["encoder_ms"] for item in measured]), + "dit_ms": _metric_summary([item["dit_ms"] for item in measured]), + "finalize_ms": _metric_summary([item["finalize_ms"] for item in measured]), + "decoder_ms": _metric_summary([item["decoder_ms"] for item in measured]), + "encoder_to_dit": { + "payload_mib": measured[0]["encoder_to_dit"]["payload_bytes"] / 2**20, + "transfer_ms": _metric_summary( + [item["encoder_to_dit"]["transfer_ms"] for item in measured] + ), + "bandwidth_gbps": _metric_summary( + [item["encoder_to_dit"]["bandwidth_gbps"] for item in measured] + ), + "handoff_ms": _metric_summary( + [item["encoder_to_dit_handoff_ms"] for item in measured] + ), + }, + "dit_to_decoder": { + "payload_mib": measured[0]["dit_to_decoder"]["payload_bytes"] / 2**20, + "transfer_ms": _metric_summary( + [item["dit_to_decoder"]["transfer_ms"] for item in measured] + ), + "bandwidth_gbps": _metric_summary( + [item["dit_to_decoder"]["bandwidth_gbps"] for item in measured] + ), + "handoff_ms": _metric_summary( + [item["dit_to_decoder_handoff_ms"] for item in measured] + ), + }, + "bandwidth_probe_gbps": { + edge: _metric_summary([item.bandwidth_gbps for item in samples]) + for edge, samples in probe.items() + }, + } + median_latency_ms = summary["latency_ms"]["median"] + summary["transfer_overhead_percent"] = { + "synchronous_copy": 100.0 + * ( + summary["encoder_to_dit"]["transfer_ms"]["median"] + + summary["dit_to_decoder"]["transfer_ms"]["median"] + ) + / median_latency_ms, + "full_handoff": 100.0 + * ( + summary["encoder_to_dit"]["handoff_ms"]["median"] + + summary["dit_to_decoder"]["handoff_ms"]["median"] + ) + / median_latency_ms, + } + args.output_dir.mkdir(parents=True, exist_ok=True) + raw_path = args.output_dir / "benchmark.json" + raw_path.write_text( + json.dumps( + { + "environment": environment, + "summary": summary, + "records": records, + "bandwidth_probe": { + edge: [vars(item) for item in samples] + for edge, samples in probe.items() + }, + }, + indent=2, + ) + + "\n" + ) + revision_label = ( + "Repository base" if environment.get("worktree_dirty") else "Commit" + ) + dirty_note = ( + "; the benchmark ran from a modified worktree" + if environment.get("worktree_dirty") + else "" + ) + markdown = f"""# LingBot three-stage disaggregation benchmark + +For the full tested configuration, methodology, findings, Slurm setup, and +limitations, see the +[experiment report](../disaggregated_inference_experiment.md). + +## Result + +| Metric | Median | P90 | +| --- | ---: | ---: | +| End-to-end chunk latency | {summary["latency_ms"]["median"]:.2f} ms | {summary["latency_ms"]["p90"]:.2f} ms | +| Encoder compute | {summary["encoder_ms"]["median"]:.2f} ms | {summary["encoder_ms"]["p90"]:.2f} ms | +| DiT denoise | {summary["dit_ms"]["median"]:.2f} ms | {summary["dit_ms"]["p90"]:.2f} ms | +| DiT cache finalize | {summary["finalize_ms"]["median"]:.2f} ms | {summary["finalize_ms"]["p90"]:.2f} ms | +| Decoder compute | {summary["decoder_ms"]["median"]:.2f} ms | {summary["decoder_ms"]["p90"]:.2f} ms | +| Encoder → DiT handoff | {summary["encoder_to_dit"]["handoff_ms"]["median"]:.2f} ms | {summary["encoder_to_dit"]["handoff_ms"]["p90"]:.2f} ms | +| DiT → decoder handoff | {summary["dit_to_decoder"]["handoff_ms"]["median"]:.2f} ms | {summary["dit_to_decoder"]["handoff_ms"]["p90"]:.2f} ms | +| Encoder → DiT payload bandwidth | {summary["encoder_to_dit"]["bandwidth_gbps"]["median"]:.2f} GB/s | {summary["encoder_to_dit"]["bandwidth_gbps"]["p90"]:.2f} GB/s | +| DiT → decoder payload bandwidth | {summary["dit_to_decoder"]["bandwidth_gbps"]["median"]:.2f} GB/s | {summary["dit_to_decoder"]["bandwidth_gbps"]["p90"]:.2f} GB/s | +| 256 MiB encoder → DiT probe | {summary["bandwidth_probe_gbps"]["encoder_to_dit"]["median"]:.2f} GB/s | {summary["bandwidth_probe_gbps"]["encoder_to_dit"]["p90"]:.2f} GB/s | +| 256 MiB DiT → decoder probe | {summary["bandwidth_probe_gbps"]["dit_to_decoder"]["median"]:.2f} GB/s | {summary["bandwidth_probe_gbps"]["dit_to_decoder"]["p90"]:.2f} GB/s | + +Steady-state throughput: **{summary["fps"]:.2f} generated FPS**. + +The headline excludes {args.warmup_blocks} warmup block(s). Mooncake was +configured with the RDMA protocol. Effective payload bandwidth includes the +synchronous transfer call but excludes receiver allocation and control-plane +ticket exchange; handoff timing in `benchmark.json` includes those costs. +The real payloads were {summary["encoder_to_dit"]["payload_mib"]:.2f} MiB +(encoder → DiT) and {summary["dit_to_decoder"]["payload_mib"]:.2f} MiB +(DiT → decoder). The two synchronous copy calls account for +{summary["transfer_overhead_percent"]["synchronous_copy"]:.2f}% of median +chunk latency; complete allocation, metadata, synchronization, and copy +handoffs account for {summary["transfer_overhead_percent"]["full_handoff"]:.2f}%. + +## Reproduction + +```bash +{environment["command"]} +``` + +- {revision_label}: `{environment["commit"]}`{dirty_note} +- Slurm: job `{environment["slurm_job_id"]}` on `{environment["hostname"]}` +- GPU: `{environment["gpus"][0]}` × 3 +- Resolution: `{args.pixel_width}x{args.pixel_height}` +- Model: `{args.model}` +""" + (args.output_dir / "README.md").write_text(markdown) + + +def main() -> None: + """Run the fixed three-rank encoder → DiT → decoder benchmark.""" + args = _parse_args() + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if world_size != 3: + raise ValueError( + "Launch with exactly three processes: torchrun --nproc_per_node=3 " + "-m lingbot.disagg.benchmark ..." + ) + if args.warmup_blocks < 0 or args.measured_blocks <= 0: + raise ValueError("warmup-blocks must be >= 0 and measured-blocks must be > 0.") + + torch.cuda.set_device(rank) + device = torch.device(f"cuda:{rank}") + config = PIPELINE_CONFIGS[args.model] + + if args.transport_only: + torch.distributed.init_process_group("gloo") + transport = MooncakeTensorTransport(device_name=args.rdma_device) + probe = { + "encoder_to_dit": _bandwidth_probe( + transport, + source=_ENCODER_RANK, + destination=_DIT_RANK, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ), + "dit_to_decoder": _bandwidth_probe( + transport, + source=_DIT_RANK, + destination=_DECODER_RANK, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ), + } + if rank == _ENCODER_RANK: + print( + json.dumps( + { + edge: _metric_summary([item.bandwidth_gbps for item in samples]) + for edge, samples in probe.items() + }, + indent=2, + ) + ) + transport.close() + torch.distributed.destroy_process_group() + return + + # Build stage-local weights before initializing torch.distributed. LingBot's + # transformer interprets an initialized process group as context parallelism, + # while these three ranks are independent pipeline stages. + encoder_stage = ( + LingbotEncoderStage(config).to(device).eval() if rank == _ENCODER_RANK else None + ) + dit_stage = LingbotDiTStage(config).to(device).eval() if rank == _DIT_RANK else None + decoder_stage = ( + LingbotDecoderStage(config).to(device).eval() if rank == _DECODER_RANK else None + ) + + encoder_inputs = ( + _load_encoder_inputs(args, device=device) if rank == _ENCODER_RANK else None + ) + torch.distributed.init_process_group("gloo") + transport = MooncakeTensorTransport(device_name=args.rdma_device) + + conditioning_bundle = None + encoder_cache = None + height_width = None + prompt = None + if encoder_stage is not None and encoder_inputs is not None: + prompt, image, intrinsics, poses, world_scale = encoder_inputs + (encoder_cache, conditioning), _ = _timed_cuda( + lambda: encoder_stage.initialize_cache(text=[prompt], image=image) + ) + conditioning_bundle = conditioning_to_bundle(conditioning) + height_width = (conditioning.height, conditioning.width) + height_width = _broadcast_object(height_width, source=_ENCODER_RANK) + received_context, _, _ = _transfer_bundle( + transport, + source=_ENCODER_RANK, + destination=_DIT_RANK, + source_bundle=conditioning_bundle, + device=device, + ) + + dit_cache = None + if dit_stage is not None and received_context is not None: + conditioning = conditioning_from_bundle( + received_context, + height=height_width[0], + width=height_width[1], + ) + dit_cache = dit_stage.initialize_cache(conditioning) + decoder_cache = ( + decoder_stage.initialize_cache() if decoder_stage is not None else None + ) + + probe = { + "encoder_to_dit": _bandwidth_probe( + transport, + source=_ENCODER_RANK, + destination=_DIT_RANK, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ), + "dit_to_decoder": _bandwidth_probe( + transport, + source=_DIT_RANK, + destination=_DECODER_RANK, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ), + } + + total_blocks = args.warmup_blocks + args.measured_blocks + records: list[dict[str, Any]] = [] + frame_start = 0 + for autoregressive_index in range(total_blocks): + torch.distributed.barrier() + step_started = time.perf_counter() + local: dict[str, Any] = {} + + encoded_bundle = None + if encoder_stage is not None and encoder_cache is not None: + num_input_frames = encoder_stage.get_num_input_frames(autoregressive_index) + frame_end = frame_start + num_input_frames + if frame_end > poses.shape[0]: + raise RuntimeError( + f"Example camera trajectory ended at frame {poses.shape[0]} " + f"before AR block {autoregressive_index}." + ) + control = CamCtrlInput( + intrinsics=intrinsics[frame_start:frame_end], + poses=poses[frame_start:frame_end], + world_scale=world_scale, + ) + encoded, local["encoder_ms"] = _timed_cuda( + partial( + encoder_stage.encode, + autoregressive_index=autoregressive_index, + cache=encoder_cache, + input=control, + ) + ) + encoded_bundle = encoder_output_to_bundle(encoded) + frame_start = frame_end + + received_encoded, encoder_transfer, encoder_handoff_ms = _transfer_bundle( + transport, + source=_ENCODER_RANK, + destination=_DIT_RANK, + source_bundle=encoded_bundle, + device=device, + ) + local["encoder_to_dit_handoff_ms"] = encoder_handoff_ms + local["encoder_to_dit"] = vars(encoder_transfer) + + clean_bundle = None + if dit_stage is not None and dit_cache is not None and received_encoded: + encoded = encoder_output_from_bundle(received_encoded) + clean_latent, local["dit_ms"] = _timed_cuda( + partial( + dit_stage.generate, + autoregressive_index=autoregressive_index, + cache=dit_cache, + input=encoded, + ) + ) + _, local["finalize_ms"] = _timed_cuda( + partial( + dit_stage.finalize, + autoregressive_index=autoregressive_index, + cache=dit_cache, + ) + ) + clean_bundle = {"clean_latent": clean_latent.contiguous()} + transport.unregister(received_encoded) + + received_clean, decoder_transfer, decoder_handoff_ms = _transfer_bundle( + transport, + source=_DIT_RANK, + destination=_DECODER_RANK, + source_bundle=clean_bundle, + device=device, + ) + local["dit_to_decoder_handoff_ms"] = decoder_handoff_ms + local["dit_to_decoder"] = vars(decoder_transfer) + + if decoder_stage is not None and decoder_cache is not None and received_clean: + decoded, local["decoder_ms"] = _timed_cuda( + partial( + decoder_stage.decode, + input=received_clean["clean_latent"], + autoregressive_index=autoregressive_index, + cache=decoder_cache, + ) + ) + local["output_frames"] = decoded.shape[-4] + transport.unregister(received_clean) + + torch.distributed.barrier() + local["end_to_end_ms"] = (time.perf_counter() - step_started) * 1000.0 + gathered: list[dict[str, Any] | None] = [None] * world_size + torch.distributed.all_gather_object(gathered, local) + if rank == _ENCODER_RANK: + record: dict[str, Any] = { + "autoregressive_index": autoregressive_index, + "warmup": autoregressive_index < args.warmup_blocks, + } + for rank_record in gathered: + assert rank_record is not None + record.update(rank_record) + records.append(record) + + peak_memory = torch.cuda.max_memory_allocated(device) / 2**30 + peak_memory_by_rank: list[float | None] = [None] * world_size + torch.distributed.all_gather_object(peak_memory_by_rank, peak_memory) + if rank == _ENCODER_RANK: + environment = _environment(args, prompt=prompt) + environment["peak_memory_gib_by_stage"] = { + "encoder": peak_memory_by_rank[_ENCODER_RANK], + "dit": peak_memory_by_rank[_DIT_RANK], + "decoder": peak_memory_by_rank[_DECODER_RANK], + } + _write_report( + args, + records=records, + probe=probe, + environment=environment, + ) + + transport.close() + torch.distributed.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/lingbot/disagg/benchmark_aggregated.py b/integrations/lingbot/lingbot/disagg/benchmark_aggregated.py new file mode 100644 index 00000000..8f12aa6e --- /dev/null +++ b/integrations/lingbot/lingbot/disagg/benchmark_aggregated.py @@ -0,0 +1,529 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Benchmark the replicated full LingBot pipeline with WORLD context parallelism.""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from pathlib import Path +from typing import Any + +import torch +from torch.distributed import ProcessGroup + +from flashdreams.core.distributed import init as init_distributed +from flashdreams.core.distributed import shutdown as shutdown_distributed +from flashdreams.infra.config import derive_config +from lingbot.config import PIPELINE_CONFIGS +from lingbot.disagg.benchmark import ( + _environment, + _load_encoder_inputs, + _metric_summary, +) +from lingbot.disagg.benchmark_cp import _cp_collective_probe +from lingbot.encoder.camctrl import CamCtrlInput +from lingbot.pipeline import LingbotWorldInferencePipeline +from lingbot.transformer import LingbotWorldTransformerConfig + +_DEFAULT_COMPARISON = Path( + "integrations/lingbot/docs/benchmark_h100_cp6_single_session/benchmark.json" +) +_WAN_SPATIAL_COMPRESSION = 8 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model", + choices=sorted(PIPELINE_CONFIGS), + default="lingbot-world-fast-taehv-window15-sink3", + ) + parser.add_argument("--example-idx", type=int, default=0) + parser.add_argument("--warmup-blocks", type=int, default=6) + parser.add_argument("--measured-blocks", type=int, default=5) + parser.add_argument( + "--pixel-height", + type=int, + default=448, + help="448 is the closest height to 464 whose LingBot token grid divides by CP8.", + ) + parser.add_argument("--pixel-width", type=int, default=832) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument( + "--cp-method", + choices=("ring", "ulysses"), + default="ulysses", + ) + parser.add_argument("--bandwidth-probe-mib", type=int, default=256) + parser.add_argument("--bandwidth-probe-iters", type=int, default=8) + parser.add_argument("--rdma-device", default=None, help=argparse.SUPPRESS) + parser.add_argument("--replica-id", type=int, default=0) + parser.add_argument( + "--measurement-barrier-dir", + type=Path, + default=None, + help="Optional shared directory that synchronizes independent replicas after warmup.", + ) + parser.add_argument("--comparison-json", type=Path, default=_DEFAULT_COMPARISON) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("outputs/lingbot_aggregated_cp8"), + ) + return parser.parse_args() + + +def _wait_for_measurement_release( + barrier_dir: Path, + *, + replica_id: int, + timeout_s: float = 3600.0, +) -> None: + """Advertise readiness and wait for the independent-replica release file. + + Args: + barrier_dir: Shared directory containing readiness and release files. + replica_id: Unique independent-replica index. + timeout_s: Maximum time to wait for the coordinator. + + Raises: + TimeoutError: The coordinator does not release measurement in time. + """ + barrier_dir.mkdir(parents=True, exist_ok=True) + (barrier_dir / f"ready-{replica_id}").write_text("ready\n") + release = barrier_dir / "release" + deadline = time.monotonic() + timeout_s + while not release.is_file(): + if time.monotonic() >= deadline: + raise TimeoutError( + f"Replica {replica_id} did not receive measurement release within " + f"{timeout_s:.0f} seconds." + ) + time.sleep(0.001) + + +def _token_layout( + *, + pixel_height: int, + pixel_width: int, + len_t: int, + patch_size: tuple[int, int, int], + cp_size: int, +) -> dict[str, int]: + """Return the LingBot token layout after validating CP divisibility.""" + kt, kh, kw = patch_size + pixel_patch_height = _WAN_SPATIAL_COMPRESSION * kh + pixel_patch_width = _WAN_SPATIAL_COMPRESSION * kw + if pixel_height % pixel_patch_height or pixel_width % pixel_patch_width: + raise ValueError( + f"Pixel resolution {pixel_width}x{pixel_height} must be divisible by " + f"{pixel_patch_width}x{pixel_patch_height} for the Wan VAE and DiT patch." + ) + if len_t % kt: + raise ValueError(f"len_t={len_t} must be divisible by temporal patch {kt}.") + latent_height = pixel_height // _WAN_SPATIAL_COMPRESSION + latent_width = pixel_width // _WAN_SPATIAL_COMPRESSION + total_tokens = (len_t // kt) * (latent_height // kh) * (latent_width // kw) + if total_tokens % cp_size: + raise ValueError( + f"Resolution {pixel_width}x{pixel_height} produces {total_tokens} tokens, " + f"which cannot be evenly sharded over CP{cp_size}." + ) + return { + "latent_height": latent_height, + "latent_width": latent_width, + "tokens_per_chunk": total_tokens, + "tokens_per_rank": total_tokens // cp_size, + } + + +def _read_comparison(path: Path) -> dict[str, Any] | None: + if not path.is_file(): + return None + return json.loads(path.read_text()) + + +def _summarize( + *, + records: list[dict[str, Any]], + tokens_per_chunk: int, + cp_probe: dict[str, list[dict[str, float]]], + peak_memory_gib_by_rank: list[float], + steady_memory_gib_by_rank: list[float], + initialization_peak_gib_by_rank: list[float], + comparison: dict[str, Any] | None, +) -> dict[str, Any]: + measured = [record for record in records if not record["warmup"]] + latency_values = [record["end_to_end_ms"] for record in measured] + elapsed_s = sum(latency_values) / 1000.0 + output_frames = sum(record["output_frames"] for record in measured) + summary: dict[str, Any] = { + "fps": output_frames / elapsed_s, + "latency_ms": _metric_summary(latency_values), + "encoder_ms": _metric_summary( + [record["critical_rank"]["encode_ms"] for record in measured] + ), + "dit_ms": _metric_summary( + [record["critical_rank"]["diffuse_ms"] for record in measured] + ), + "decoder_ms": _metric_summary( + [record["critical_rank"]["decode_ms"] for record in measured] + ), + "finalize_ms": _metric_summary( + [record["critical_rank"]["finalize_ms"] for record in measured] + ), + "tokens_per_chunk": tokens_per_chunk, + "token_throughput_per_second": tokens_per_chunk * len(measured) / elapsed_s, + "cp_probe_gbps": { + collective: _metric_summary( + [sample["bandwidth_gbps"] for sample in samples] + ) + for collective, samples in cp_probe.items() + }, + "memory": { + "peak_gib_by_rank": peak_memory_gib_by_rank, + "steady_allocated_gib_by_rank": steady_memory_gib_by_rank, + "initialization_peak_gib_by_rank": initialization_peak_gib_by_rank, + "node_peak_gib": sum(peak_memory_gib_by_rank), + "node_steady_allocated_gib": sum(steady_memory_gib_by_rank), + "per_rank_peak_gib": _metric_summary(peak_memory_gib_by_rank), + }, + } + if comparison is not None: + previous = comparison["summary"] + previous_environment = comparison["environment"] + previous_tokens = ( + previous_environment["latent_frames_per_chunk"] + * (previous_environment["resolution"][0] // 16) + * (previous_environment["resolution"][1] // 16) + ) + previous_memory = previous_environment["peak_memory_gib_by_rank"] + allocation = previous_environment.get("allocation", {}) + summary["comparison"] = { + "topology": ( + f"1 encoder : CP{allocation.get('cp_size', '?')} DiT : 1 decoder" + ), + "resolution": previous_environment["resolution"], + "fps": previous["fps"], + "latency_ms": previous["latency_ms"]["median"], + "tokens_per_chunk": previous_tokens, + "token_throughput_per_second": previous_tokens * previous["fps"] / 12.0, + "node_peak_gib": sum(previous_memory), + "latency_speedup": previous["latency_ms"]["median"] + / summary["latency_ms"]["median"], + "fps_ratio": summary["fps"] / previous["fps"], + "token_throughput_ratio": summary["token_throughput_per_second"] + / (previous_tokens * previous["fps"] / 12.0), + "node_peak_memory_ratio": summary["memory"]["node_peak_gib"] + / sum(previous_memory), + } + return summary + + +def _write_report( + args: argparse.Namespace, + *, + records: list[dict[str, Any]], + cp_probe: dict[str, list[dict[str, float]]], + environment: dict[str, Any], + summary: dict[str, Any], +) -> None: + args.output_dir.mkdir(parents=True, exist_ok=True) + (args.output_dir / "benchmark.json").write_text( + json.dumps( + { + "environment": environment, + "summary": summary, + "records": records, + "cp_probe": cp_probe, + }, + indent=2, + ) + + "\n" + ) + memory = summary["memory"] + comparison = summary.get("comparison") + comparison_section = "" + if comparison is not None: + comparison_rows = "\n".join( + ( + f"| Median chunk latency | {comparison['latency_ms']:.2f} ms | " + f"{summary['latency_ms']['median']:.2f} ms | " + f"{comparison['latency_speedup']:.2f}× faster |", + f"| Generated FPS | {comparison['fps']:.2f} | " + f"{summary['fps']:.2f} | {comparison['fps_ratio']:.2f}× |", + f"| DiT token throughput | " + f"{comparison['token_throughput_per_second']:.0f} token/s | " + f"{summary['token_throughput_per_second']:.0f} token/s | " + f"{comparison['token_throughput_ratio']:.2f}× |", + f"| Node peak allocated HBM | " + f"{comparison['node_peak_gib']:.2f} GiB | " + f"{memory['node_peak_gib']:.2f} GiB | " + f"{comparison['node_peak_memory_ratio']:.2f}× |", + ) + ) + comparison_section = f"""\ +## Comparison with disaggregated CP + +| Metric | {comparison["topology"]} | Aggregated CP{environment["allocation"]["cp_size"]} | Change | +| --- | ---: | ---: | ---: | +{comparison_rows} + +The resolutions may differ when the token grid cannot divide evenly over both +CP groups. Token throughput is the fairest compute-rate comparison in that +case. +""" + markdown = f"""# LingBot aggregated CP{environment["allocation"]["cp_size"]} benchmark + +All ranks own the complete encoder, DiT, and decoder pipeline. The DiT token +axis is context-parallel across WORLD with {args.cp_method} attention. Encoder +and decoder work is replicated on every rank; there are no RDMA stage +boundaries in this topology. + +## Result + +| Metric | Median | P90 | +| --- | ---: | ---: | +| End-to-end chunk latency | {summary["latency_ms"]["median"]:.2f} ms | {summary["latency_ms"]["p90"]:.2f} ms | +| Encoder critical-rank compute | {summary["encoder_ms"]["median"]:.2f} ms | {summary["encoder_ms"]["p90"]:.2f} ms | +| DiT critical-rank denoise | {summary["dit_ms"]["median"]:.2f} ms | {summary["dit_ms"]["p90"]:.2f} ms | +| Decoder critical-rank compute | {summary["decoder_ms"]["median"]:.2f} ms | {summary["decoder_ms"]["p90"]:.2f} ms | +| DiT cache finalize | {summary["finalize_ms"]["median"]:.2f} ms | {summary["finalize_ms"]["p90"]:.2f} ms | +| NCCL broadcast probe | {summary["cp_probe_gbps"]["broadcast"]["median"]:.2f} GB/s | {summary["cp_probe_gbps"]["broadcast"]["p90"]:.2f} GB/s | +| NCCL all-gather probe | {summary["cp_probe_gbps"]["all_gather"]["median"]:.2f} GB/s | {summary["cp_probe_gbps"]["all_gather"]["p90"]:.2f} GB/s | + +- Generated throughput: **{summary["fps"]:.2f} FPS** +- DiT token throughput: **{summary["token_throughput_per_second"]:.0f} token/s** +- Peak allocated HBM: **{memory["per_rank_peak_gib"]["min"]:.2f}–{memory["per_rank_peak_gib"]["max"]:.2f} GiB per rank**, **{memory["node_peak_gib"]:.2f} GiB node total** +- Steady allocated HBM after rollout: **{memory["node_steady_allocated_gib"]:.2f} GiB node total** + +{comparison_section} + +## Reproduction + +```bash +{environment["command"]} +``` + +- Repository revision: `{environment["commit"]}`{" (modified worktree)" if environment["worktree_dirty"] else ""} +- Slurm: job `{environment["slurm_job_id"]}` on `{environment["hostname"]}` +- GPU: `{environment["gpus"][0]}` × {len(environment["gpus"])} +- Resolution: `{args.pixel_width}x{args.pixel_height}` +- Warmup / measured blocks: {args.warmup_blocks} / {args.measured_blocks} +""" + (args.output_dir / "README.md").write_text(markdown) + + +def main() -> None: + """Run the complete pipeline on every rank with WORLD context parallelism.""" + args = _parse_args() + if args.warmup_blocks < 0 or args.measured_blocks <= 0: + raise ValueError("warmup-blocks must be >= 0 and measured-blocks must be > 0.") + + init_distributed() + rank = torch.distributed.get_rank() + world_size = torch.distributed.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + device = torch.device(f"cuda:{local_rank}") + + base_config = PIPELINE_CONFIGS[args.model] + transformer_config = base_config.diffusion_model.transformer + assert isinstance(transformer_config, LingbotWorldTransformerConfig) + layout = _token_layout( + pixel_height=args.pixel_height, + pixel_width=args.pixel_width, + len_t=transformer_config.len_t, + patch_size=transformer_config.network.patch_size, + cp_size=world_size, + ) + if ( + args.cp_method == "ulysses" + and transformer_config.network.num_heads % world_size + ): + raise ValueError( + f"Ulysses requires {transformer_config.network.num_heads} attention heads " + f"to divide CP{world_size}." + ) + + base_seed = base_config.diffusion_model.seed + config = derive_config( + base_config, + diffusion_model={ + "seed": None if base_seed is None else base_seed + rank, + "transformer": {"network": {"cp_method": args.cp_method}}, + }, + ) + pipeline = config.setup().to(device).eval() + assert isinstance(pipeline, LingbotWorldInferencePipeline) + prompt, image, intrinsics, poses, world_scale = _load_encoder_inputs( + args, + device=device, + ) + cache = pipeline.initialize_cache(text=[prompt], image=image) + + world_group = torch.distributed.group.WORLD + assert isinstance(world_group, ProcessGroup) + cp_probe = _cp_collective_probe( + cp_ranks=tuple(range(world_size)), + cp_group=world_group, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ) + initialization_peak = torch.cuda.max_memory_allocated(device) / 2**30 + torch.cuda.reset_peak_memory_stats(device) + + records: list[dict[str, Any]] = [] + frame_start = 0 + total_blocks = args.warmup_blocks + args.measured_blocks + measurement_started_at: float | None = None + for autoregressive_index in range(total_blocks): + if autoregressive_index == args.warmup_blocks: + torch.cuda.synchronize(device) + if args.measurement_barrier_dir is not None: + if world_size != 1: + raise ValueError( + "The independent-replica measurement barrier requires CP1." + ) + _wait_for_measurement_release( + args.measurement_barrier_dir, + replica_id=args.replica_id, + ) + measurement_started_at = time.perf_counter() + + num_input_frames = pipeline.get_num_input_frames(autoregressive_index) + frame_end = frame_start + num_input_frames + if frame_end > poses.shape[0]: + raise RuntimeError( + f"Camera trajectory ended at frame {poses.shape[0]} before " + f"AR block {autoregressive_index}." + ) + control = CamCtrlInput( + intrinsics=intrinsics[frame_start:frame_end], + poses=poses[frame_start:frame_end], + world_scale=world_scale, + ) + frame_start = frame_end + + torch.distributed.barrier() + started = time.perf_counter() + video = pipeline.generate( + autoregressive_index=autoregressive_index, + cache=cache, + input=control, + ) + stats = pipeline.finalize( + autoregressive_index=autoregressive_index, + cache=cache, + ) + assert stats is not None + local_record = { + **stats, + "wall_ms": (time.perf_counter() - started) * 1000.0, + "output_frames": video.shape[-4], + "rank": rank, + } + per_rank: list[dict[str, Any] | None] = [None] * world_size + torch.distributed.all_gather_object(per_rank, local_record) + if rank == 0: + rank_records = [item for item in per_rank if item is not None] + records.append( + { + "autoregressive_index": autoregressive_index, + "warmup": autoregressive_index < args.warmup_blocks, + "end_to_end_ms": max(item["wall_ms"] for item in rank_records), + "output_frames": rank_records[0]["output_frames"], + "critical_rank": { + metric: max(item[metric] for item in rank_records) + for metric in ( + "encode_ms", + "diffuse_ms", + "decode_ms", + "finalize_ms", + ) + }, + "per_rank": rank_records, + } + ) + + torch.cuda.synchronize(device) + measurement_finished_at = time.perf_counter() + assert measurement_started_at is not None + peak_memory = torch.cuda.max_memory_allocated(device) / 2**30 + steady_memory = torch.cuda.memory_allocated(device) / 2**30 + resource_record = { + "peak": peak_memory, + "steady": steady_memory, + "initialization_peak": initialization_peak, + } + resources: list[dict[str, float] | None] = [None] * world_size + torch.distributed.all_gather_object(resources, resource_record) + if rank == 0: + rank_resources = [item for item in resources if item is not None] + environment = _environment( + args, + prompt=prompt, + world_size=world_size, + module_name="lingbot.disagg.benchmark_aggregated", + ) + environment["allocation"] = { + "full_pipeline_replicas": list(range(world_size)), + "dit_cp_group": list(range(world_size)), + "cp_size": world_size, + "cp_method": args.cp_method, + } + environment["token_layout"] = layout + environment["transport"] = { + "stage_handoffs": "none", + "dit_collectives": "NCCL", + } + environment["noise_seed_by_rank"] = [ + None if base_seed is None else base_seed + local_rank + for local_rank in range(world_size) + ] + environment["replica_id"] = args.replica_id + environment["measurement_window"] = { + "started_at": measurement_started_at, + "finished_at": measurement_finished_at, + "elapsed_s": measurement_finished_at - measurement_started_at, + } + summary = _summarize( + records=records, + tokens_per_chunk=layout["tokens_per_chunk"], + cp_probe=cp_probe, + peak_memory_gib_by_rank=[item["peak"] for item in rank_resources], + steady_memory_gib_by_rank=[item["steady"] for item in rank_resources], + initialization_peak_gib_by_rank=[ + item["initialization_peak"] for item in rank_resources + ], + comparison=_read_comparison(args.comparison_json), + ) + _write_report( + args, + records=records, + cp_probe=cp_probe, + environment=environment, + summary=summary, + ) + + shutdown_distributed(synchronize=True, terminate_process=True) + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/lingbot/disagg/benchmark_cp.py b/integrations/lingbot/lingbot/disagg/benchmark_cp.py new file mode 100644 index 00000000..b2e3bccd --- /dev/null +++ b/integrations/lingbot/lingbot/disagg/benchmark_cp.py @@ -0,0 +1,854 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LingBot single-session benchmark with one context-parallel DiT group.""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import time +from functools import partial +from pathlib import Path +from typing import Any + +import torch +from torch.distributed import ProcessGroup + +from flashdreams.infra.config import derive_config +from flashdreams.infra.transfer import ( + TensorBundle, + TransferStats, + describe_tensor_bundle, +) +from lingbot.config import PIPELINE_CONFIGS +from lingbot.disagg.benchmark import ( + _bandwidth_probe, + _broadcast_object, + _environment, + _load_encoder_inputs, + _metric_summary, + _timed_cuda, + _transfer_bundle, +) +from lingbot.disagg.stages import ( + LingbotDecoderStage, + LingbotDiTStage, + LingbotEncoderStage, + conditioning_from_bundle, + conditioning_to_bundle, + encoder_output_from_bundle, + encoder_output_to_bundle, + encoder_output_to_cp_bundles, +) +from lingbot.encoder.camctrl import CamCtrlInput +from lingbot.transformer import LingbotWorldTransformerConfig + +_DEFAULT_BASELINE = Path( + "integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json" +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model", + choices=sorted(PIPELINE_CONFIGS), + default="lingbot-world-fast-taehv-window15-sink3", + ) + parser.add_argument("--example-idx", type=int, default=0) + parser.add_argument("--warmup-blocks", type=int, default=6) + parser.add_argument("--measured-blocks", type=int, default=5) + parser.add_argument("--pixel-height", type=int, default=464) + parser.add_argument("--pixel-width", type=int, default=832) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--cp-ranks", type=int, default=6) + parser.add_argument( + "--cp-method", + choices=("ring", "ulysses"), + default="ring", + help="DiT attention collective; CP6 requires ring for the 40-head model.", + ) + parser.add_argument("--rdma-device", default=None) + parser.add_argument( + "--direct-cp-input", + action="store_true", + help="Patchify on the encoder and RDMA each token shard to its CP rank.", + ) + parser.add_argument("--bandwidth-probe-mib", type=int, default=256) + parser.add_argument("--bandwidth-probe-iters", type=int, default=8) + parser.add_argument( + "--transport-only", + action="store_true", + help="Probe Mooncake stage edges and DiT NCCL collectives without weights.", + ) + parser.add_argument("--baseline-json", type=Path, default=_DEFAULT_BASELINE) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("outputs/lingbot_disagg_cp6"), + ) + return parser.parse_args() + + +def _read_baseline(path: Path) -> dict[str, Any]: + if not path.is_file(): + raise FileNotFoundError(f"Baseline benchmark not found: {path}") + return json.loads(path.read_text()) + + +def _create_cp_group(cp_ranks: tuple[int, ...]) -> ProcessGroup | None: + group = torch.distributed.new_group( + ranks=list(cp_ranks), + backend="nccl", + device_id=torch.device("cuda", torch.cuda.current_device()), + ) + if torch.distributed.get_rank() not in cp_ranks: + return None + assert isinstance(group, ProcessGroup) + return group + + +def _empty_bundle_like_descriptors( + descriptors: Any, + *, + device: torch.device, +) -> TensorBundle: + return { + descriptor.name: torch.empty( + descriptor.shape, + dtype=descriptor.dtype, + device=device, + ).contiguous() + for descriptor in descriptors + } + + +def _fanout_bundle_to_cp( + *, + source_bundle: TensorBundle | None, + source: int, + cp_ranks: tuple[int, ...], + cp_group: ProcessGroup | None, + device: torch.device, +) -> tuple[TensorBundle | None, float | None]: + """Broadcast one leader-owned full bundle across the DiT CP subgroup.""" + rank = torch.distributed.get_rank() + started = time.perf_counter() + source_descriptors = None + if rank == source: + assert source_bundle is not None + source_descriptors = describe_tensor_bundle(source_bundle) + descriptors = _broadcast_object( + source_descriptors, + source=source, + ) + if rank not in cp_ranks: + return None, None + assert cp_group is not None + if rank == source: + assert source_bundle is not None + local_bundle = source_bundle + else: + local_bundle = _empty_bundle_like_descriptors(descriptors, device=device) + for descriptor in descriptors: + torch.distributed.broadcast( + local_bundle[descriptor.name], + src=source, + group=cp_group, + ) + torch.cuda.synchronize(device) + return local_bundle, (time.perf_counter() - started) * 1000.0 + + +def _max_elapsed_ms(elapsed_ms: float, cp_group: ProcessGroup) -> float: + elapsed = torch.tensor(elapsed_ms, device=torch.cuda.current_device()) + torch.distributed.all_reduce( + elapsed, + op=torch.distributed.ReduceOp.MAX, + group=cp_group, + ) + return float(elapsed.item()) + + +def _time_cp_collective( + call: Any, + *, + cp_group: ProcessGroup, +) -> float: + torch.distributed.barrier(group=cp_group) + torch.cuda.synchronize() + started = time.perf_counter() + call() + torch.cuda.synchronize() + return _max_elapsed_ms((time.perf_counter() - started) * 1000.0, cp_group) + + +def _cp_collective_probe( + *, + cp_ranks: tuple[int, ...], + cp_group: ProcessGroup | None, + size_mib: int, + iterations: int, + device: torch.device, +) -> dict[str, list[dict[str, float]]]: + """Measure CP broadcast and all-gather effective per-rank bandwidth.""" + rank = torch.distributed.get_rank() + leader = cp_ranks[0] + leader_samples: dict[str, list[dict[str, float]]] | None = None + if rank in cp_ranks: + assert cp_group is not None + payload_bytes = size_mib * 2**20 + broadcast_buffer = torch.empty(payload_bytes, dtype=torch.uint8, device=device) + shard_bytes = payload_bytes // len(cp_ranks) + shard = torch.empty(shard_bytes, dtype=torch.uint8, device=device) + gathered = torch.empty( + shard_bytes * len(cp_ranks), + dtype=torch.uint8, + device=device, + ) + + _time_cp_collective( + lambda: torch.distributed.broadcast( + broadcast_buffer, + src=leader, + group=cp_group, + ), + cp_group=cp_group, + ) + _time_cp_collective( + lambda: torch.distributed.all_gather_into_tensor( + gathered, + shard, + group=cp_group, + ), + cp_group=cp_group, + ) + + broadcast_samples = [] + all_gather_samples = [] + for _ in range(iterations): + broadcast_ms = _time_cp_collective( + lambda: torch.distributed.broadcast( + broadcast_buffer, + src=leader, + group=cp_group, + ), + cp_group=cp_group, + ) + all_gather_ms = _time_cp_collective( + lambda: torch.distributed.all_gather_into_tensor( + gathered, + shard, + group=cp_group, + ), + cp_group=cp_group, + ) + broadcast_samples.append( + { + "payload_bytes": float(payload_bytes), + "transfer_ms": broadcast_ms, + "bandwidth_gbps": payload_bytes / (broadcast_ms / 1000.0) / 1e9, + } + ) + all_gather_samples.append( + { + "payload_bytes": float(gathered.numel()), + "transfer_ms": all_gather_ms, + "bandwidth_gbps": gathered.numel() / (all_gather_ms / 1000.0) / 1e9, + } + ) + if rank == leader: + leader_samples = { + "broadcast": broadcast_samples, + "all_gather": all_gather_samples, + } + return _broadcast_object(leader_samples, source=leader) + + +def _shutdown( + transport: Any, + *, + cp_group: ProcessGroup | None, + device: torch.device, +) -> None: + """Drain the NCCL subgroup before destroying it and the control group.""" + if cp_group is not None: + torch.cuda.synchronize(device) + torch.distributed.barrier(group=cp_group) + torch.distributed.barrier() + transport.close() + if cp_group is not None: + torch.distributed.destroy_process_group(cp_group) + torch.distributed.barrier() + torch.distributed.destroy_process_group() + + +def _summarize( + *, + records: list[dict[str, Any]], + mooncake_probe: dict[str, list[TransferStats]], + cp_probe: dict[str, list[dict[str, float]]], + baseline: dict[str, Any], + cp_size: int, +) -> dict[str, Any]: + measured = [record for record in records if not record["warmup"]] + frame_count = sum(record["output_frames"] for record in measured) + elapsed_s = sum(record["end_to_end_ms"] for record in measured) / 1000.0 + baseline_summary = baseline["summary"] + baseline_latency = baseline_summary["latency_ms"]["median"] + baseline_dit = ( + baseline_summary["dit_ms"]["median"] + baseline_summary["finalize_ms"]["median"] + ) + aggregate_fps = frame_count / elapsed_s + cp_critical = [ + max(worker["dit_ms"] + worker["finalize_ms"] for worker in record["cp_workers"]) + for record in measured + ] + return { + "fps": aggregate_fps, + "latency_ms": _metric_summary([record["end_to_end_ms"] for record in measured]), + "latency_speedup": baseline_latency + / statistics.median([record["end_to_end_ms"] for record in measured]), + "fps_speedup": aggregate_fps / baseline_summary["fps"], + "encoder_ms": _metric_summary([record["encoder_ms"] for record in measured]), + "encoder_to_cp_leader": { + "payload_mib": measured[0]["encoder_to_cp_leader"]["payload_bytes"] / 2**20, + "copy_ms": _metric_summary( + [record["encoder_to_cp_leader"]["transfer_ms"] for record in measured] + ), + "handoff_ms": _metric_summary( + [record["encoder_to_cp_leader_handoff_ms"] for record in measured] + ), + }, + "cp_input_fanout_ms": _metric_summary( + [record["cp_input_fanout_ms"] for record in measured] + ), + "dit_ms": _metric_summary( + [ + max(worker["dit_ms"] for worker in record["cp_workers"]) + for record in measured + ] + ), + "finalize_ms": _metric_summary( + [ + max(worker["finalize_ms"] for worker in record["cp_workers"]) + for record in measured + ] + ), + "dit_critical_path_ms": _metric_summary(cp_critical), + "dit_speedup": baseline_dit / statistics.median(cp_critical), + "cp_efficiency": (baseline_dit / statistics.median(cp_critical)) / cp_size, + "cp_leader_to_decoder": { + "payload_mib": measured[0]["cp_leader_to_decoder"]["payload_bytes"] / 2**20, + "copy_ms": _metric_summary( + [record["cp_leader_to_decoder"]["transfer_ms"] for record in measured] + ), + "handoff_ms": _metric_summary( + [record["cp_leader_to_decoder_handoff_ms"] for record in measured] + ), + }, + "decoder_ms": _metric_summary([record["decoder_ms"] for record in measured]), + "mooncake_probe_gbps": { + edge: _metric_summary([sample.bandwidth_gbps for sample in samples]) + for edge, samples in mooncake_probe.items() + }, + "cp_probe_gbps": { + collective: _metric_summary( + [sample["bandwidth_gbps"] for sample in samples] + ) + for collective, samples in cp_probe.items() + }, + "baseline": { + "fps": baseline_summary["fps"], + "latency_ms": baseline_latency, + "dit_critical_path_ms": baseline_dit, + "topology": "1 encoder : 1 DiT : 1 decoder", + }, + } + + +def _write_report( + args: argparse.Namespace, + *, + records: list[dict[str, Any]], + mooncake_probe: dict[str, list[TransferStats]], + cp_probe: dict[str, list[dict[str, float]]], + baseline: dict[str, Any], + environment: dict[str, Any], +) -> None: + summary = _summarize( + records=records, + mooncake_probe=mooncake_probe, + cp_probe=cp_probe, + baseline=baseline, + cp_size=args.cp_ranks, + ) + args.output_dir.mkdir(parents=True, exist_ok=True) + (args.output_dir / "benchmark.json").write_text( + json.dumps( + { + "environment": environment, + "summary": summary, + "records": records, + "mooncake_probe": { + edge: [vars(sample) for sample in samples] + for edge, samples in mooncake_probe.items() + }, + "cp_probe": cp_probe, + }, + indent=2, + ) + + "\n" + ) + peak_memory = environment["peak_memory_gib_by_rank"] + cp_memory = peak_memory[1:-1] + markdown = f"""# LingBot CP{args.cp_ranks} single-session disaggregation benchmark + +## Result + +Topology: **1 encoder : 1 DiT group with CP{args.cp_ranks} ({args.cp_method}) : 1 decoder**. +The {args.cp_ranks} DiT ranks cooperate on one autoregressive session. + +| Metric | Median | P90 | +| --- | ---: | ---: | +| End-to-end chunk latency | {summary["latency_ms"]["median"]:.2f} ms | {summary["latency_ms"]["p90"]:.2f} ms | +| Encoder compute | {summary["encoder_ms"]["median"]:.2f} ms | {summary["encoder_ms"]["p90"]:.2f} ms | +| Encoder → CP leader handoff | {summary["encoder_to_cp_leader"]["handoff_ms"]["median"]:.2f} ms | {summary["encoder_to_cp_leader"]["handoff_ms"]["p90"]:.2f} ms | +| CP input fanout | {summary["cp_input_fanout_ms"]["median"]:.2f} ms | {summary["cp_input_fanout_ms"]["p90"]:.2f} ms | +| CP DiT critical path | {summary["dit_critical_path_ms"]["median"]:.2f} ms | {summary["dit_critical_path_ms"]["p90"]:.2f} ms | +| CP leader → decoder handoff | {summary["cp_leader_to_decoder"]["handoff_ms"]["median"]:.2f} ms | {summary["cp_leader_to_decoder"]["handoff_ms"]["p90"]:.2f} ms | +| Decoder compute | {summary["decoder_ms"]["median"]:.2f} ms | {summary["decoder_ms"]["p90"]:.2f} ms | +| 256 MiB Mooncake probes | {statistics.median([value["median"] for value in summary["mooncake_probe_gbps"].values()]):.2f} GB/s | — | +| 256 MiB-equivalent NCCL broadcast | {summary["cp_probe_gbps"]["broadcast"]["median"]:.2f} GB/s | {summary["cp_probe_gbps"]["broadcast"]["p90"]:.2f} GB/s | +| 256 MiB-equivalent NCCL all-gather | {summary["cp_probe_gbps"]["all_gather"]["median"]:.2f} GB/s | {summary["cp_probe_gbps"]["all_gather"]["p90"]:.2f} GB/s | + +- Single-session throughput: **{summary["fps"]:.2f} generated FPS** +- Latency speedup versus tracked CP1 baseline: **{summary["latency_speedup"]:.2f}×** +- DiT critical-path speedup: **{summary["dit_speedup"]:.2f}×** +- CP scaling efficiency: **{summary["cp_efficiency"] * 100.0:.1f}%** + +The headline excludes {args.warmup_blocks} warmup blocks and measures +{args.measured_blocks} blocks. It accelerates one session; it does not represent +independent concurrent sessions. + +## Peak allocated memory + +| Role | Peak | +| --- | ---: | +| Encoder | {peak_memory[0]:.2f} GiB | +| CP DiT ranks | {min(cp_memory):.2f}–{max(cp_memory):.2f} GiB each | +| Decoder | {peak_memory[-1]:.2f} GiB | + +## Reproduction + +```bash +{environment["command"]} +``` + +- Repository revision: `{environment["commit"]}`{" (modified worktree)" if environment["worktree_dirty"] else ""} +- Slurm: job `{environment["slurm_job_id"]}` on `{environment["hostname"]}` +- GPU: `{environment["gpus"][0]}` × {len(environment["gpus"])} +- Model: `{args.model}` +""" + (args.output_dir / "README.md").write_text(markdown) + + +def main() -> None: + """Run one encoder, one CP-sharded DiT session, and one decoder.""" + args = _parse_args() + rank = int(os.environ.get("RANK", "0")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + expected_world_size = args.cp_ranks + 2 + if world_size != expected_world_size: + raise ValueError( + f"Launch with {expected_world_size} processes for one encoder, " + f"CP{args.cp_ranks} DiT, and one decoder; got {world_size}." + ) + if args.cp_ranks < 2: + raise ValueError("cp-ranks must be at least two.") + if args.warmup_blocks < 0 or args.measured_blocks <= 0: + raise ValueError("warmup-blocks must be >= 0 and measured-blocks must be > 0.") + + encoder_rank = 0 + cp_ranks = tuple(range(1, 1 + args.cp_ranks)) + cp_leader = cp_ranks[0] + decoder_rank = world_size - 1 + + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + base_config = PIPELINE_CONFIGS[args.model] + + encoder_stage = ( + LingbotEncoderStage(base_config).to(device).eval() + if rank == encoder_rank and not args.transport_only + else None + ) + dit_stage = None + if rank in cp_ranks and not args.transport_only: + base_seed = base_config.diffusion_model.seed + local_cp_rank = rank - cp_leader + config = derive_config( + base_config, + diffusion_model={ + "seed": None if base_seed is None else base_seed + local_cp_rank, + "transformer": {"network": {"cp_method": args.cp_method}}, + }, + ) + dit_stage = LingbotDiTStage(config).to(device).eval() + decoder_stage = ( + LingbotDecoderStage(base_config).to(device).eval() + if rank == decoder_rank and not args.transport_only + else None + ) + encoder_inputs = ( + _load_encoder_inputs(args, device=device) if encoder_stage is not None else None + ) + + torch.distributed.init_process_group("gloo") + cp_group = _create_cp_group(cp_ranks) + if dit_stage is not None: + assert cp_group is not None + dit_stage.set_context_parallel_group(cp_group) + + from flashdreams.infra.transfer import MooncakeTensorTransport + + transport = MooncakeTensorTransport(device_name=args.rdma_device) + mooncake_probe = { + "encoder_to_cp_leader": _bandwidth_probe( + transport, + source=encoder_rank, + destination=cp_leader, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ), + "cp_leader_to_decoder": _bandwidth_probe( + transport, + source=cp_leader, + destination=decoder_rank, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ), + } + cp_probe = _cp_collective_probe( + cp_ranks=cp_ranks, + cp_group=cp_group, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ) + if args.transport_only: + if rank == encoder_rank: + print( + json.dumps( + { + "mooncake_gbps": { + edge: _metric_summary( + [sample.bandwidth_gbps for sample in samples] + ) + for edge, samples in mooncake_probe.items() + }, + "cp_collective_gbps": { + collective: _metric_summary( + [sample["bandwidth_gbps"] for sample in samples] + ) + for collective, samples in cp_probe.items() + }, + }, + indent=2, + ) + ) + _shutdown(transport, cp_group=cp_group, device=device) + return + + conditioning_bundle = None + encoder_cache = None + height_width = None + prompt = None + image = None + intrinsics = None + poses = None + world_scale = None + if encoder_stage is not None and encoder_inputs is not None: + prompt, image, intrinsics, poses, world_scale = encoder_inputs + encoder_cache, conditioning = encoder_stage.initialize_cache( + text=[prompt], + image=image, + ) + conditioning_bundle = conditioning_to_bundle(conditioning) + height_width = (conditioning.height, conditioning.width) + height_width = _broadcast_object(height_width, source=encoder_rank) + received_context, _, _ = _transfer_bundle( + transport, + source=encoder_rank, + destination=cp_leader, + source_bundle=conditioning_bundle, + device=device, + ) + cp_context, _ = _fanout_bundle_to_cp( + source_bundle=received_context, + source=cp_leader, + cp_ranks=cp_ranks, + cp_group=cp_group, + device=device, + ) + torch.distributed.barrier() + + dit_cache = None + if dit_stage is not None and cp_context is not None: + conditioning = conditioning_from_bundle( + cp_context, + height=height_width[0], + width=height_width[1], + ) + dit_cache = dit_stage.initialize_cache(conditioning) + decoder_cache = ( + decoder_stage.initialize_cache() if decoder_stage is not None else None + ) + + total_blocks = args.warmup_blocks + args.measured_blocks + records: list[dict[str, Any]] = [] + frame_start = 0 + for autoregressive_index in range(total_blocks): + torch.distributed.barrier() + step_started = time.perf_counter() + local: dict[str, Any] = {} + + encoded_bundle: TensorBundle | None = None + encoded_cp_bundles: tuple[TensorBundle, ...] | None = None + if encoder_stage is not None and encoder_cache is not None: + assert intrinsics is not None + assert poses is not None + assert world_scale is not None + num_input_frames = encoder_stage.get_num_input_frames(autoregressive_index) + frame_end = frame_start + num_input_frames + if frame_end > poses.shape[0]: + raise RuntimeError( + f"Example camera trajectory ended at frame {poses.shape[0]} " + f"before AR block {autoregressive_index}." + ) + control = CamCtrlInput( + intrinsics=intrinsics[frame_start:frame_end], + poses=poses[frame_start:frame_end], + world_scale=world_scale, + ) + encoded, local["encoder_ms"] = _timed_cuda( + partial( + encoder_stage.encode, + autoregressive_index=autoregressive_index, + cache=encoder_cache, + input=control, + ) + ) + if args.direct_cp_input: + transformer_config = base_config.diffusion_model.transformer + assert isinstance( + transformer_config, + LingbotWorldTransformerConfig, + ) + encoded_cp_bundles = encoder_output_to_cp_bundles( + encoded, + cp_size=args.cp_ranks, + patch_size=transformer_config.network.patch_size, + ) + else: + encoded_bundle = encoder_output_to_bundle(encoded) + frame_start = frame_end + + cp_encoded = None + if args.direct_cp_input: + direct_transfers: list[TransferStats] = [] + direct_handoff_ms = 0.0 + for cp_index, cp_rank in enumerate(cp_ranks): + received_shard, transfer, handoff_ms = _transfer_bundle( + transport, + source=encoder_rank, + destination=cp_rank, + source_bundle=( + encoded_cp_bundles[cp_index] + if rank == encoder_rank and encoded_cp_bundles is not None + else None + ), + device=device, + ) + if rank == cp_rank: + cp_encoded = received_shard + if rank == encoder_rank: + direct_transfers.append(transfer) + direct_handoff_ms += handoff_ms + if rank == encoder_rank: + total_payload = sum(item.payload_bytes for item in direct_transfers) + total_transfer_ms = sum(item.transfer_ms for item in direct_transfers) + local["encoder_to_cp_leader"] = { + **vars(direct_transfers[0]), + "payload_bytes": total_payload, + "transfer_ms": total_transfer_ms, + "bandwidth_gbps": ( + total_payload / (total_transfer_ms / 1000.0) / 1e9 + ), + "direct_shards": args.cp_ranks, + } + local["encoder_to_cp_leader_handoff_ms"] = direct_handoff_ms + if rank in cp_ranks: + local["cp_input_fanout_ms"] = 0.0 + else: + received_encoded, encoder_transfer, encoder_handoff_ms = _transfer_bundle( + transport, + source=encoder_rank, + destination=cp_leader, + source_bundle=encoded_bundle, + device=device, + ) + if rank == encoder_rank: + local["encoder_to_cp_leader"] = dict(vars(encoder_transfer)) + local["encoder_to_cp_leader_handoff_ms"] = encoder_handoff_ms + cp_encoded, fanout_ms = _fanout_bundle_to_cp( + source_bundle=received_encoded, + source=cp_leader, + cp_ranks=cp_ranks, + cp_group=cp_group, + device=device, + ) + if rank in cp_ranks: + assert fanout_ms is not None + local["cp_input_fanout_ms"] = fanout_ms + torch.distributed.barrier() + + clean_bundle = None + if dit_stage is not None and dit_cache is not None and cp_encoded is not None: + encoded = encoder_output_from_bundle( + cp_encoded, + patchified=args.direct_cp_input, + ) + clean_latent, local["dit_ms"] = _timed_cuda( + partial( + dit_stage.generate, + autoregressive_index=autoregressive_index, + cache=dit_cache, + input=encoded, + ) + ) + _, local["finalize_ms"] = _timed_cuda( + partial( + dit_stage.finalize, + autoregressive_index=autoregressive_index, + cache=dit_cache, + ) + ) + local["cp_rank"] = rank - cp_leader + if args.direct_cp_input: + transport.unregister(cp_encoded) + if rank == cp_leader: + clean_bundle = {"clean_latent": clean_latent.contiguous()} + if not args.direct_cp_input: + transport.unregister(cp_encoded) + + received_clean, decoder_transfer, decoder_handoff_ms = _transfer_bundle( + transport, + source=cp_leader, + destination=decoder_rank, + source_bundle=clean_bundle, + device=device, + ) + if rank == encoder_rank: + local["cp_leader_to_decoder"] = dict(vars(decoder_transfer)) + local["cp_leader_to_decoder_handoff_ms"] = decoder_handoff_ms + + if decoder_stage is not None and decoder_cache is not None: + assert received_clean is not None + decoded, local["decoder_ms"] = _timed_cuda( + partial( + decoder_stage.decode, + input=received_clean["clean_latent"], + autoregressive_index=autoregressive_index, + cache=decoder_cache, + ) + ) + local["output_frames"] = decoded.shape[-4] + transport.unregister(received_clean) + + torch.distributed.barrier() + if rank == encoder_rank: + local["end_to_end_ms"] = (time.perf_counter() - step_started) * 1000.0 + gathered: list[dict[str, Any] | None] = [None] * world_size + torch.distributed.all_gather_object(gathered, local) + if rank == encoder_rank: + encoder_record = gathered[encoder_rank] + decoder_record = gathered[decoder_rank] + assert encoder_record is not None + assert decoder_record is not None + cp_workers = [] + for cp_rank in cp_ranks: + worker = gathered[cp_rank] + assert worker is not None + cp_workers.append(worker) + records.append( + { + "autoregressive_index": autoregressive_index, + "warmup": autoregressive_index < args.warmup_blocks, + **encoder_record, + "cp_input_fanout_ms": max( + worker["cp_input_fanout_ms"] for worker in cp_workers + ), + "cp_workers": cp_workers, + "decoder_ms": decoder_record["decoder_ms"], + "output_frames": decoder_record["output_frames"], + } + ) + + peak_memory = torch.cuda.max_memory_allocated(device) / 2**30 + peak_memory_by_rank: list[float | None] = [None] * world_size + torch.distributed.all_gather_object(peak_memory_by_rank, peak_memory) + if rank == encoder_rank: + environment = _environment( + args, + prompt=prompt, + world_size=world_size, + module_name="lingbot.disagg.benchmark_cp", + ) + environment["allocation"] = { + "encoder": [encoder_rank], + "dit_cp_group": list(cp_ranks), + "decoder": [decoder_rank], + "cp_size": args.cp_ranks, + "cp_method": args.cp_method, + "direct_cp_input": args.direct_cp_input, + } + environment["noise_seed_by_cp_rank"] = [ + None + if base_config.diffusion_model.seed is None + else base_config.diffusion_model.seed + cp_rank + for cp_rank in range(args.cp_ranks) + ] + environment["peak_memory_gib_by_rank"] = peak_memory_by_rank + _write_report( + args, + records=records, + mooncake_probe=mooncake_probe, + cp_probe=cp_probe, + baseline=_read_baseline(args.baseline_json), + environment=environment, + ) + + _shutdown(transport, cp_group=cp_group, device=device) + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/lingbot/disagg/benchmark_independent.py b/integrations/lingbot/lingbot/disagg/benchmark_independent.py new file mode 100644 index 00000000..a04c8230 --- /dev/null +++ b/integrations/lingbot/lingbot/disagg/benchmark_independent.py @@ -0,0 +1,363 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Concurrent independent-GPU LingBot benchmark coordinator.""" + +from __future__ import annotations + +import argparse +import json +import os +import shlex +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import Any + +from lingbot.config import PIPELINE_CONFIGS +from lingbot.disagg.benchmark import _metric_summary + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--replicas", type=int, default=8) + parser.add_argument( + "--model", + choices=sorted(PIPELINE_CONFIGS), + default="lingbot-world-fast-taehv-window15-sink3", + ) + parser.add_argument("--example-idx", type=int, default=0) + parser.add_argument("--warmup-blocks", type=int, default=6) + parser.add_argument("--measured-blocks", type=int, default=5) + parser.add_argument("--pixel-height", type=int, default=464) + parser.add_argument("--pixel-width", type=int, default=832) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--compile-threads-per-replica", type=int, default=4) + parser.add_argument("--timeout-s", type=float, default=3600.0) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("outputs/lingbot_aggregated_8xcp1"), + ) + return parser.parse_args() + + +def _child_command( + args: argparse.Namespace, + *, + replica_id: int, + barrier_dir: Path, + output_dir: Path, +) -> list[str]: + """Build one isolated CP1 worker command. + + Args: + args: Coordinator arguments. + replica_id: GPU and session index. + barrier_dir: Shared post-warmup synchronization directory. + output_dir: Worker-specific result directory. + + Returns: + Argument vector for one ``torchrun`` child process. + """ + return [ + sys.executable, + "-m", + "torch.distributed.run", + "--standalone", + "--nproc_per_node=1", + "-m", + "lingbot.disagg.benchmark_aggregated", + "--model", + args.model, + "--example-idx", + str(args.example_idx), + "--warmup-blocks", + str(args.warmup_blocks), + "--measured-blocks", + str(args.measured_blocks), + "--pixel-height", + str(args.pixel_height), + "--pixel-width", + str(args.pixel_width), + "--fps", + str(args.fps), + "--cp-method", + "ulysses", + "--bandwidth-probe-mib", + "1", + "--bandwidth-probe-iters", + "1", + "--comparison-json", + str(args.output_dir / "no-comparison.json"), + "--replica-id", + str(replica_id), + "--measurement-barrier-dir", + str(barrier_dir), + "--output-dir", + str(output_dir), + ] + + +def _summarize(worker_documents: list[dict[str, Any]]) -> dict[str, Any]: + """Summarize independent workers over their shared measurement window. + + Args: + worker_documents: Raw ``benchmark.json`` documents from every worker. + + Returns: + Aggregate throughput, latency, memory, and synchronization metrics. + """ + starts = [ + document["environment"]["measurement_window"]["started_at"] + for document in worker_documents + ] + finishes = [ + document["environment"]["measurement_window"]["finished_at"] + for document in worker_documents + ] + measured_records = [ + record + for document in worker_documents + for record in document["records"] + if not record["warmup"] + ] + total_frames = sum(record["output_frames"] for record in measured_records) + measurement_wall_s = max(finishes) - min(starts) + worker_fps = [document["summary"]["fps"] for document in worker_documents] + worker_latency_medians = [ + document["summary"]["latency_ms"]["median"] for document in worker_documents + ] + rollout_peak = [ + document["summary"]["memory"]["peak_gib_by_rank"][0] + for document in worker_documents + ] + initialization_peak = [ + document["summary"]["memory"]["initialization_peak_gib_by_rank"][0] + for document in worker_documents + ] + steady = [ + document["summary"]["memory"]["steady_allocated_gib_by_rank"][0] + for document in worker_documents + ] + return { + "aggregate_fps": total_frames / measurement_wall_s, + "sum_of_worker_fps": sum(worker_fps), + "total_output_frames": total_frames, + "measurement_wall_s": measurement_wall_s, + "measurement_start_skew_ms": (max(starts) - min(starts)) * 1000.0, + "measurement_finish_skew_ms": (max(finishes) - min(finishes)) * 1000.0, + "per_session_fps": _metric_summary(worker_fps), + "per_session_median_latency_ms": _metric_summary(worker_latency_medians), + "all_chunk_latency_ms": _metric_summary( + [record["end_to_end_ms"] for record in measured_records] + ), + "memory": { + "rollout_peak_gib_by_gpu": rollout_peak, + "rollout_peak_gib_per_gpu": _metric_summary(rollout_peak), + "rollout_peak_gib_node_total": sum(rollout_peak), + "initialization_peak_gib_by_gpu": initialization_peak, + "initialization_peak_gib_per_gpu": _metric_summary(initialization_peak), + "initialization_peak_gib_node_total": sum(initialization_peak), + "steady_allocated_gib_by_gpu": steady, + "steady_allocated_gib_node_total": sum(steady), + }, + } + + +def _write_report( + args: argparse.Namespace, + *, + worker_documents: list[dict[str, Any]], + worker_commands: list[list[str]], + summary: dict[str, Any], +) -> None: + """Write machine-readable and Markdown coordinator reports. + + Args: + args: Coordinator arguments. + worker_documents: Raw result from each independent worker. + worker_commands: Exact subprocess argument vectors. + summary: Aggregate benchmark metrics. + """ + first_environment = worker_documents[0]["environment"] + command = shlex.join( + [ + "uv", + "run", + "--package", + "flashdreams-lingbot", + "python", + "-m", + "lingbot.disagg.benchmark_independent", + *sys.argv[1:], + ] + ) + environment = { + "command": command, + "commit": first_environment["commit"], + "worktree_dirty": first_environment["worktree_dirty"], + "hostname": socket.gethostname(), + "slurm_job_id": os.environ.get("SLURM_JOB_ID"), + "model": args.model, + "resolution": [args.pixel_height, args.pixel_width], + "warmup_blocks": args.warmup_blocks, + "measured_blocks": args.measured_blocks, + "replicas": args.replicas, + "gpus": [document["environment"]["gpus"][0] for document in worker_documents], + "torch": first_environment["torch"], + "cuda": first_environment["cuda"], + "cudnn": first_environment["cudnn"], + "driver": first_environment["driver"], + "precision": first_environment["precision"], + "checkpoint": first_environment["checkpoint"], + "worker_commands": [shlex.join(item) for item in worker_commands], + } + document = { + "environment": environment, + "summary": summary, + "workers": worker_documents, + } + (args.output_dir / "benchmark.json").write_text( + json.dumps(document, indent=2) + "\n" + ) + + memory = summary["memory"] + markdown = f"""# LingBot eight independent aggregated workers + +Eight H100s each own one complete CP1 encoder + DiT + decoder pipeline and one +session. The workers synchronize after warmup and run their five measured +chunks concurrently. + +## Result + +| Metric | Value | +| --- | ---: | +| Aggregate generated FPS | **{summary["aggregate_fps"]:.2f}** | +| Per-session FPS, median / p90 | {summary["per_session_fps"]["median"]:.2f} / {summary["per_session_fps"]["p90"]:.2f} | +| Chunk latency, median / p90 | {summary["all_chunk_latency_ms"]["median"]:.2f} / {summary["all_chunk_latency_ms"]["p90"]:.2f} ms | +| Shared measurement wall time | {summary["measurement_wall_s"]:.3f} s | +| Measurement start skew | {summary["measurement_start_skew_ms"]:.2f} ms | +| Rollout peak HBM per GPU, min–max | {memory["rollout_peak_gib_per_gpu"]["min"]:.2f}–{memory["rollout_peak_gib_per_gpu"]["max"]:.2f} GiB | +| Initialization peak HBM per GPU, min–max | {memory["initialization_peak_gib_per_gpu"]["min"]:.2f}–{memory["initialization_peak_gib_per_gpu"]["max"]:.2f} GiB | +| Rollout peak HBM, node total | {memory["rollout_peak_gib_node_total"]:.2f} GiB | + +## Reproduction + +```bash +{command} +``` + +- Repository revision: `{environment["commit"]}`{" (modified worktree)" if environment["worktree_dirty"] else ""} +- Slurm: job `{environment["slurm_job_id"]}` on `{environment["hostname"]}` +- GPU: `{environment["gpus"][0]}` × {len(environment["gpus"])} +- Resolution: `{args.pixel_width}x{args.pixel_height}` +- Warmup / measured blocks per worker: {args.warmup_blocks} / {args.measured_blocks} +""" + (args.output_dir / "README.md").write_text(markdown) + + +def main() -> None: + """Launch independent CP1 workers and summarize concurrent serving.""" + args = _parse_args() + if args.replicas <= 0: + raise ValueError("replicas must be positive.") + if args.replicas > 8: + raise ValueError("replicas cannot exceed the eight GPUs in one node.") + + args.output_dir = args.output_dir.resolve() + args.output_dir.mkdir(parents=True, exist_ok=True) + barrier_dir = args.output_dir / "barrier" + barrier_dir.mkdir(exist_ok=True) + for stale_file in barrier_dir.iterdir(): + stale_file.unlink() + + processes: list[subprocess.Popen[bytes]] = [] + log_files: list[Any] = [] + worker_commands: list[list[str]] = [] + try: + for replica_id in range(args.replicas): + worker_dir = args.output_dir / f"worker-{replica_id}" + worker_dir.mkdir(exist_ok=True) + command = _child_command( + args, + replica_id=replica_id, + barrier_dir=barrier_dir, + output_dir=worker_dir, + ) + environment = os.environ.copy() + environment["CUDA_VISIBLE_DEVICES"] = str(replica_id) + environment["TORCHINDUCTOR_COMPILE_THREADS"] = str( + args.compile_threads_per_replica + ) + log_file = (worker_dir / "run.log").open("wb") + log_files.append(log_file) + worker_commands.append(command) + processes.append( + subprocess.Popen( + command, + env=environment, + stdout=log_file, + stderr=subprocess.STDOUT, + ) + ) + + deadline = time.monotonic() + args.timeout_s + ready_files = [barrier_dir / f"ready-{index}" for index in range(args.replicas)] + while not all(path.is_file() for path in ready_files): + failed = [ + (index, process.returncode) + for index, process in enumerate(processes) + if process.poll() is not None and process.returncode != 0 + ] + if failed: + raise RuntimeError(f"Workers failed before measurement: {failed}.") + if time.monotonic() >= deadline: + raise TimeoutError("Independent workers did not finish warmup in time.") + time.sleep(0.1) + + (barrier_dir / "release").write_text("release\n") + return_codes = [process.wait(timeout=args.timeout_s) for process in processes] + failed = [ + (index, return_code) + for index, return_code in enumerate(return_codes) + if return_code != 0 + ] + if failed: + raise RuntimeError(f"Workers failed during measurement: {failed}.") + finally: + for process in processes: + if process.poll() is None: + process.terminate() + for log_file in log_files: + log_file.close() + + worker_documents = [ + json.loads((args.output_dir / f"worker-{index}" / "benchmark.json").read_text()) + for index in range(args.replicas) + ] + summary = _summarize(worker_documents) + _write_report( + args, + worker_documents=worker_documents, + worker_commands=worker_commands, + summary=summary, + ) + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/lingbot/disagg/benchmark_pipeline.py b/integrations/lingbot/lingbot/disagg/benchmark_pipeline.py new file mode 100644 index 00000000..fab2e118 --- /dev/null +++ b/integrations/lingbot/lingbot/disagg/benchmark_pipeline.py @@ -0,0 +1,1080 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Eight-GPU LingBot benchmark with two-rank pipeline-parallel DiT groups.""" + +from __future__ import annotations + +import argparse +import json +import os +import time +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from typing import Any, cast + +import torch +from torch import Tensor +from torch.distributed import ProcessGroup + +from flashdreams.infra.config import derive_config +from flashdreams.infra.transfer import ( + RegisteredTensorPool, + TensorBundle, + describe_tensor_bundle, +) +from flashdreams.recipes.wan.autoencoder.i2v import I2VCtrl +from lingbot.config import PIPELINE_CONFIGS +from lingbot.disagg.benchmark import ( + _broadcast_object, + _environment, + _load_encoder_inputs, + _metric_summary, + _timed_cuda, + _transfer_bundle, +) +from lingbot.disagg.benchmark_replicated import ( + _TransferChannel, + _create_transport, + _finish_channel, + _setup_channel, + _submit_channel, +) +from lingbot.disagg.stages import ( + LingbotConditioning, + LingbotDecoderStage, + LingbotDiTStage, + LingbotEncoderStage, + conditioning_from_bundle, + conditioning_to_bundle, + encoder_output_from_bundle, + encoder_output_to_bundle, +) +from lingbot.encoder.camctrl import CamCtrlInput, I2VCamCtrlEmbeddings + +_DEFAULT_REPLICATED_BASELINE = Path( + "integrations/lingbot/docs/benchmark_h100_1io7dit_optimized/summary.json" +) +"""Tracked seven-session, one-GPU-per-DiT comparison.""" + + +@dataclass(frozen=True, kw_only=True) +class PipelineTopology: + """Fixed I/O, DiT-group, and spare-rank allocation.""" + + io_rank: int + """Rank hosting the shared encoder and decoder.""" + + dit_groups: tuple[tuple[int, int], ...] + """Ordered two-rank DiT pipeline groups.""" + + spare_ranks: tuple[int, ...] + """Ranks intentionally left without model weights.""" + + sessions_per_group: int + """Fixed session microbatch held by each DiT group.""" + + @property + def session_count(self) -> int: + """Return the number of concurrently generated sessions.""" + return len(self.dit_groups) * self.sessions_per_group + + @property + def dit_ranks(self) -> tuple[int, ...]: + """Return all ranks that own a DiT layer partition.""" + return tuple(rank for group in self.dit_groups for rank in group) + + +def build_pipeline_topology( + *, + world_size: int, + sessions_per_group: int, +) -> PipelineTopology: + """Build the supported one-I/O plus three-pair topology. + + Args: + world_size: Number of local ranks in the launch. + sessions_per_group: Fixed session batch assigned to each DiT group. + + Returns: + Eight-rank topology with GPU 7 left spare. + + Raises: + ValueError: The launch is not the required eight-rank layout. + """ + if world_size != 8: + raise ValueError(f"Pipeline benchmark requires eight ranks, got {world_size}.") + if sessions_per_group < 1: + raise ValueError("sessions_per_group must be positive.") + return PipelineTopology( + io_rank=0, + dit_groups=((1, 2), (3, 4), (5, 6)), + spare_ranks=(7,), + sessions_per_group=sessions_per_group, + ) + + +def stack_conditioning(items: list[LingbotConditioning]) -> LingbotConditioning: + """Stack per-session conditioning into one fixed DiT microbatch. + + Args: + items: Session conditioning records with identical spatial layout. + + Returns: + Batched conditioning for one DiT group. + + Raises: + ValueError: No items were supplied or optional fields are inconsistent. + """ + if not items: + raise ValueError("At least one conditioning record is required.") + height = items[0].height + width = items[0].width + if any(item.height != height or item.width != width for item in items): + raise ValueError("All microbatched sessions must share one spatial layout.") + + def stack_optional(name: str) -> torch.Tensor | None: + values = [getattr(item, name) for item in items] + if all(value is None for value in values): + return None + if any(value is None for value in values): + raise ValueError(f"Optional conditioning field {name} is inconsistent.") + return torch.cat(cast(list[torch.Tensor], values), dim=0) + + return LingbotConditioning( + height=height, + width=width, + text_embeddings=torch.cat([item.text_embeddings for item in items], dim=0), + negative_text_embeddings=stack_optional("negative_text_embeddings"), + image_embeddings=stack_optional("image_embeddings"), + ) + + +def stack_encoder_outputs( + items: list[I2VCamCtrlEmbeddings], +) -> I2VCamCtrlEmbeddings: + """Stack per-session encoder outputs into one fixed DiT microbatch. + + Args: + items: Unpatchified encoder outputs for one DiT group. + + Returns: + Batched I2V and camera-control payload. + + Raises: + ValueError: The input list is empty or already patchified. + """ + if not items: + raise ValueError("At least one encoder output is required.") + if any(item._is_patchified or item.i2v._is_patchified for item in items): + raise ValueError("Pipeline input batching requires unpatchified tensors.") + return I2VCamCtrlEmbeddings( + i2v=I2VCtrl( + latent=torch.stack([item.i2v.latent for item in items], dim=0), + mask=torch.stack([item.i2v.mask for item in items], dim=0), + _is_patchified=False, + ), + plucker=torch.stack([item.plucker for item in items], dim=0), + _is_patchified=False, + ) + + +def split_conditioning( + conditioning: LingbotConditioning, +) -> tuple[LingbotConditioning, LingbotConditioning]: + """Split one batch-two conditioning record into session records. + + Args: + conditioning: Conditioning whose leading batch dimension is two. + + Returns: + Two batch-one conditioning records. + + Raises: + ValueError: A conditioning tensor does not have batch size two. + """ + + def split_optional(name: str) -> tuple[Tensor | None, Tensor | None]: + value = getattr(conditioning, name) + if value is None: + return None, None + if value.shape[0] != 2: + raise ValueError(f"Conditioning field {name} must have batch size two.") + return value.narrow(0, 0, 1), value.narrow(0, 1, 1) + + if conditioning.text_embeddings.shape[0] != 2: + raise ValueError("Text conditioning must have batch size two.") + negative = split_optional("negative_text_embeddings") + image = split_optional("image_embeddings") + return ( + LingbotConditioning( + height=conditioning.height, + width=conditioning.width, + text_embeddings=conditioning.text_embeddings.narrow(0, 0, 1), + negative_text_embeddings=negative[0], + image_embeddings=image[0], + ), + LingbotConditioning( + height=conditioning.height, + width=conditioning.width, + text_embeddings=conditioning.text_embeddings.narrow(0, 1, 1), + negative_text_embeddings=negative[1], + image_embeddings=image[1], + ), + ) + + +def split_encoder_outputs( + output: I2VCamCtrlEmbeddings, +) -> tuple[I2VCamCtrlEmbeddings, I2VCamCtrlEmbeddings]: + """Split one batch-two encoder payload into session payloads. + + Args: + output: Unpatchified encoder payload with leading batch size two. + + Returns: + Two batch-one payloads. + + Raises: + ValueError: The payload is patchified or does not have batch size two. + """ + if output._is_patchified or output.i2v._is_patchified: + raise ValueError("Double buffering requires unpatchified encoder output.") + tensors = (output.i2v.latent, output.i2v.mask, output.plucker) + if any(tensor.shape[0] != 2 for tensor in tensors): + raise ValueError("Encoder payload fields must have batch size two.") + + def session(index: int) -> I2VCamCtrlEmbeddings: + return I2VCamCtrlEmbeddings( + i2v=I2VCtrl( + latent=output.i2v.latent.narrow(0, index, 1), + mask=output.i2v.mask.narrow(0, index, 1), + _is_patchified=False, + ), + plucker=output.plucker.narrow(0, index, 1), + _is_patchified=False, + ) + + return session(0), session(1) + + +def validate_double_buffered_schedule(*, sessions_per_group: int) -> None: + """Validate the two-slot pipeline scheduling contract. + + Args: + sessions_per_group: Number of session slots assigned to each DiT pair. + + Raises: + ValueError: The session count is not exactly two. + """ + if sessions_per_group != 2: + raise ValueError("Double buffering requires exactly two sessions per group.") + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model", + choices=sorted(PIPELINE_CONFIGS), + default="lingbot-world-fast-taehv-window15-sink3", + ) + parser.add_argument("--example-idx", type=int, default=0) + parser.add_argument("--warmup-blocks", type=int, default=6) + parser.add_argument("--measured-blocks", type=int, default=5) + parser.add_argument("--pixel-height", type=int, default=464) + parser.add_argument("--pixel-width", type=int, default=832) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--sessions-per-group", type=int, default=2) + parser.add_argument( + "--double-buffered", + action="store_true", + help=( + "Split batch-two input into session-affine microbatches and overlap " + "the two DiT stages." + ), + ) + parser.add_argument( + "--compile-network", + action=argparse.BooleanOptionalAction, + default=True, + ) + parser.add_argument( + "--transport", + choices=("mooncake", "nixl"), + default="mooncake", + ) + parser.add_argument("--rdma-device", default=None) + parser.add_argument("--bandwidth-probe-mib", type=int, default=256) + parser.add_argument("--bandwidth-probe-iters", type=int, default=8) + parser.add_argument( + "--baseline-json", + type=Path, + default=_DEFAULT_REPLICATED_BASELINE, + ) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("outputs/lingbot_disagg_pipeline_3x2"), + ) + return parser.parse_args() + + +def _create_pipeline_groups( + topology: PipelineTopology, + *, + device: torch.device, +) -> dict[tuple[int, int], ProcessGroup]: + """Create every NCCL pair in global rank order.""" + rank = torch.distributed.get_rank() + local_groups: dict[tuple[int, int], ProcessGroup] = {} + for ranks in topology.dit_groups: + group = torch.distributed.new_group( + ranks=list(ranks), + backend="nccl", + device_id=device, + ) + if rank in ranks: + assert isinstance(group, ProcessGroup) + local_groups[ranks] = group + return local_groups + + +def _empty_bundle( + descriptors: Any, + *, + device: torch.device, +) -> TensorBundle: + return { + descriptor.name: torch.empty( + descriptor.shape, + dtype=descriptor.dtype, + device=device, + ).contiguous() + for descriptor in descriptors + } + + +def _fanout_bundle( + *, + source_bundle: TensorBundle | None, + ranks: tuple[int, int], + group: ProcessGroup | None, + device: torch.device, +) -> tuple[TensorBundle | None, float | None]: + """Broadcast a leader-owned bundle to the second DiT pipeline rank.""" + rank = torch.distributed.get_rank() + leader = ranks[0] + descriptors = _broadcast_object( + describe_tensor_bundle(source_bundle) + if rank == leader and source_bundle is not None + else None, + source=leader, + ) + if rank not in ranks: + return None, None + assert group is not None + local_bundle = ( + source_bundle if rank == leader else _empty_bundle(descriptors, device=device) + ) + assert local_bundle is not None + + torch.distributed.barrier(group=group) + torch.cuda.synchronize(device) + started = time.perf_counter() + for descriptor in descriptors: + torch.distributed.broadcast( + local_bundle[descriptor.name], + src=leader, + group=group, + ) + torch.cuda.synchronize(device) + elapsed = torch.tensor( + (time.perf_counter() - started) * 1000.0, + device=device, + ) + torch.distributed.all_reduce( + elapsed, + op=torch.distributed.ReduceOp.MAX, + group=group, + ) + return local_bundle, float(elapsed.item()) + + +def _probe_pipeline_links( + topology: PipelineTopology, + *, + local_groups: dict[tuple[int, int], ProcessGroup], + size_mib: int, + iterations: int, + device: torch.device, +) -> dict[str, list[dict[str, float]]]: + """Measure one-way NCCL point-to-point bandwidth for every DiT pair.""" + rank = torch.distributed.get_rank() + payload_bytes = size_mib * 2**20 + local_results: dict[str, list[dict[str, float]]] = {} + for ranks in topology.dit_groups: + torch.distributed.barrier() + if rank in ranks: + group = local_groups[ranks] + buffer = torch.empty(payload_bytes, dtype=torch.uint8, device=device) + for _ in range(2): + if rank == ranks[0]: + torch.distributed.send(buffer, dst=ranks[1], group=group) + else: + torch.distributed.recv(buffer, src=ranks[0], group=group) + samples: list[dict[str, float]] = [] + for _ in range(iterations): + torch.distributed.barrier(group=group) + torch.cuda.synchronize(device) + started = time.perf_counter() + if rank == ranks[0]: + torch.distributed.send(buffer, dst=ranks[1], group=group) + else: + torch.distributed.recv(buffer, src=ranks[0], group=group) + torch.cuda.synchronize(device) + elapsed = torch.tensor( + (time.perf_counter() - started) * 1000.0, + device=device, + ) + torch.distributed.all_reduce( + elapsed, + op=torch.distributed.ReduceOp.MAX, + group=group, + ) + if rank == ranks[0]: + elapsed_ms = float(elapsed.item()) + samples.append( + { + "payload_bytes": float(payload_bytes), + "transfer_ms": elapsed_ms, + "bandwidth_gbps": payload_bytes + / (elapsed_ms / 1000.0) + / 1e9, + } + ) + if rank == ranks[0]: + local_results[f"{ranks[0]}->{ranks[1]}"] = samples + torch.distributed.barrier() + + gathered: list[dict[str, list[dict[str, float]]] | None] = [ + None + ] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(gathered, local_results) + if rank != topology.io_rank: + return {} + merged: dict[str, list[dict[str, float]]] = {} + for result in gathered: + if result: + merged.update(result) + return merged + + +def _summarize( + *, + records: list[dict[str, Any]], + topology: PipelineTopology, + p2p_probe: dict[str, list[dict[str, float]]], + memory: dict[str, list[float]], + baseline: dict[str, Any], +) -> dict[str, Any]: + measured = [record for record in records if not record["warmup"]] + total_frames = sum(record["output_frames"] for record in measured) + total_wall_s = sum(record["wave_latency_ms"] for record in measured) / 1000.0 + aggregate_fps = total_frames / total_wall_s + probe_bandwidth = [ + sample["bandwidth_gbps"] for samples in p2p_probe.values() for sample in samples + ] + return { + "aggregate_fps": aggregate_fps, + "per_session_fps": aggregate_fps / topology.session_count, + "wave_latency_ms": _metric_summary( + [record["wave_latency_ms"] for record in measured] + ), + "encoder_wave_ms": _metric_summary( + [record["encoder_wave_ms"] for record in measured] + ), + "dit_group_ms": _metric_summary( + [ + worker["dit_ms"] + for record in measured + for worker in record["dit_group_leaders"] + ] + ), + "finalize_group_ms": _metric_summary( + [ + worker["finalize_ms"] + for record in measured + for worker in record["dit_group_leaders"] + ] + ), + "decoder_wave_ms": _metric_summary( + [record["decoder_wave_ms"] for record in measured] + ), + "pair_fanout_ms": _metric_summary( + [value for record in measured for value in record["pair_fanout_ms"]] + ), + "p2p_probe_gbps": { + "all_pairs": _metric_summary(probe_bandwidth), + "by_pair": { + pair: _metric_summary([sample["bandwidth_gbps"] for sample in samples]) + for pair, samples in p2p_probe.items() + }, + }, + "memory": memory, + "baseline": { + "topology": "1 I/O + 7 full DiT replicas", + "sessions": baseline["topology"]["sessions_per_wave"], + "aggregate_fps": baseline["performance"]["aggregate_fps"], + "per_session_fps": baseline["performance"]["per_session_fps"], + "max_dit_peak_gib": max(baseline["peak_allocated_gib_by_rank"][1:]), + }, + } + + +def _write_report( + args: argparse.Namespace, + *, + environment: dict[str, Any], + records: list[dict[str, Any]], + summary: dict[str, Any], + p2p_probe: dict[str, list[dict[str, float]]], +) -> None: + args.output_dir.mkdir(parents=True, exist_ok=True) + document = { + "environment": environment, + "summary": summary, + "records": records, + "p2p_probe": p2p_probe, + } + (args.output_dir / "benchmark.json").write_text( + json.dumps(document, indent=2) + "\n" + ) + memory = summary["memory"] + baseline = summary["baseline"] + markdown = f"""# LingBot two-stage DiT pipeline benchmark + +## Result + +| Metric | Pipeline-parallel result | Replicated DiT reference | +| --- | ---: | ---: | +| Concurrent sessions | {environment["sessions"]} | {baseline["sessions"]} | +| Aggregate generated FPS | **{summary["aggregate_fps"]:.2f}** | {baseline["aggregate_fps"]:.2f} | +| Generated FPS per session | **{summary["per_session_fps"]:.2f}** | {baseline["per_session_fps"]:.2f} | +| Median wave latency | **{summary["wave_latency_ms"]["median"]:.2f} ms** | 2358.51 ms | +| Maximum DiT-rank capacity | **{max(memory["required_capacity_gib_by_rank"][1:7]):.2f} GiB** | {baseline["max_dit_peak_gib"]:.2f} GiB | +| Median 256 MiB NCCL P2P bandwidth | **{summary["p2p_probe_gbps"]["all_pairs"]["median"]:.2f} GB/s** | — | + +## Topology + +```text +GPU 0 shared encoder + decoder +GPU 1–2 DiT group A +GPU 3–4 DiT group B +GPU 5–6 DiT group C +GPU 7 spare +``` + +Each DiT group holds {args.sessions_per_group} sessions using the +{("double-buffered fill/drain schedule" if args.double_buffered else "fixed batched schedule")}. +The first rank owns layers 0–19; the second owns layers 20–39 and the output +head. CUDA graph capture is disabled because the graph boundary contains NCCL +point-to-point operations. + +## Reproduction + +```bash +{environment["command"]} +``` + +- Repository revision: `{environment["commit"]}`{" (modified worktree)" if environment["worktree_dirty"] else ""} +- Slurm job: `{environment["slurm_job_id"]}` on `{environment["hostname"]}` +- GPU: `{environment["gpus"][0]}` × {len(environment["gpus"])} +- Warmup / measured blocks: {args.warmup_blocks} / {args.measured_blocks} +""" + (args.output_dir / "README.md").write_text(markdown) + + +def main() -> None: + """Run three two-rank DiT pipelines with shared I/O on one node.""" + args = _parse_args() + rank = int(os.environ.get("RANK", "0")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + topology = build_pipeline_topology( + world_size=world_size, + sessions_per_group=args.sessions_per_group, + ) + if args.warmup_blocks < 0 or args.measured_blocks <= 0: + raise ValueError("warmup-blocks must be >= 0 and measured-blocks must be > 0.") + if args.double_buffered: + validate_double_buffered_schedule(sessions_per_group=args.sessions_per_group) + + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + base_config = PIPELINE_CONFIGS[args.model] + dit_batch_size = 1 if args.double_buffered else args.sessions_per_group + dit_config = derive_config( + base_config, + diffusion_model=dict( + transformer=dict( + batch_shape=(dit_batch_size,), + compile_network=args.compile_network, + use_cuda_graph=False, + ) + ), + ) + + encoder_stage = ( + LingbotEncoderStage(base_config).to(device).eval() + if rank == topology.io_rank + else None + ) + decoder_stage = ( + LingbotDecoderStage(base_config).to(device).eval() + if rank == topology.io_rank + else None + ) + dit_stage = LingbotDiTStage(dit_config) if rank in topology.dit_ranks else None + encoder_inputs = ( + _load_encoder_inputs(args, device=device) if rank == topology.io_rank else None + ) + + torch.distributed.init_process_group("gloo") + local_groups = _create_pipeline_groups(topology, device=device) + local_group_ranks = next( + (ranks for ranks in topology.dit_groups if rank in ranks), + None, + ) + if dit_stage is not None and local_group_ranks is not None: + stage_index = local_group_ranks.index(rank) + dit_stage.configure_pipeline_parallel( + stage_index=stage_index, + stage_count=2, + group=local_groups[local_group_ranks], + ranks=local_group_ranks, + ) + dit_stage = dit_stage.to(device).eval() + + weight_memory_gib = torch.cuda.memory_allocated(device) / 2**30 + torch.cuda.reset_peak_memory_stats(device) + transport = _create_transport(args, rank=rank) + pool = RegisteredTensorPool(transport, max_buffers_per_bucket=16) + + encoder_caches: list[Any] = [] + conditioning_bundles: list[TensorBundle] = [] + prompt = None + image = None + intrinsics = None + poses = None + world_scale = None + if encoder_stage is not None and encoder_inputs is not None: + prompt, image, intrinsics, poses, world_scale = encoder_inputs + conditionings: list[LingbotConditioning] = [] + for _ in range(topology.session_count): + cache, conditioning = encoder_stage.initialize_cache( + text=[prompt], + image=image, + ) + encoder_caches.append(cache) + conditionings.append(conditioning) + for group_index in range(len(topology.dit_groups)): + start = group_index * topology.sessions_per_group + end = start + topology.sessions_per_group + conditioning_bundles.append( + conditioning_to_bundle(stack_conditioning(conditionings[start:end])) + ) + + height_width = ( + (conditionings[0].height, conditionings[0].width) + if rank == topology.io_rank + else None + ) + height_width = _broadcast_object(height_width, source=topology.io_rank) + dit_caches: list[Any] = [] + for group_index, ranks in enumerate(topology.dit_groups): + leader = ranks[0] + received_context, _, _ = _transfer_bundle( + transport, + source=topology.io_rank, + destination=leader, + source_bundle=( + conditioning_bundles[group_index] if rank == topology.io_rank else None + ), + device=device, + ) + fanned_context, _ = _fanout_bundle( + source_bundle=received_context if rank == leader else None, + ranks=ranks, + group=local_groups.get(ranks), + device=device, + ) + if rank in ranks: + assert dit_stage is not None + assert fanned_context is not None + conditioning = conditioning_from_bundle( + fanned_context, + height=height_width[0], + width=height_width[1], + ) + if args.double_buffered: + dit_caches.extend( + dit_stage.initialize_cache(session_conditioning) + for session_conditioning in split_conditioning(conditioning) + ) + else: + dit_caches.append(dit_stage.initialize_cache(conditioning)) + if rank == leader and received_context is not None: + transport.unregister(received_context) + + decoder_caches = ( + [decoder_stage.initialize_cache() for _ in range(topology.session_count)] + if decoder_stage is not None + else [] + ) + cache_initialization_peak_gib = torch.cuda.max_memory_allocated(device) / 2**30 + torch.cuda.reset_peak_memory_stats(device) + + p2p_probe = _probe_pipeline_links( + topology, + local_groups=local_groups, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ) + torch.cuda.reset_peak_memory_stats(device) + + total_blocks = args.warmup_blocks + args.measured_blocks + frame_starts = [0] * topology.session_count + encoder_channels: list[_TransferChannel | None] = [None] * len(topology.dit_groups) + decoder_channels: list[_TransferChannel | None] = [None] * len(topology.dit_groups) + records: list[dict[str, Any]] = [] + for autoregressive_index in range(total_blocks): + torch.distributed.barrier() + wave_started = time.perf_counter() + local: dict[str, Any] = {} + encoder_times: list[float] = [] + group_encoder_bundles: list[TensorBundle] = [] + if encoder_stage is not None: + assert intrinsics is not None + assert poses is not None + assert world_scale is not None + encoded_sessions: list[I2VCamCtrlEmbeddings] = [] + for session_index in range(topology.session_count): + num_input_frames = encoder_stage.get_num_input_frames( + autoregressive_index + ) + frame_start = frame_starts[session_index] + frame_end = frame_start + num_input_frames + if frame_end > poses.shape[0]: + raise RuntimeError( + f"Camera trajectory ended before AR block {autoregressive_index}." + ) + control = CamCtrlInput( + intrinsics=intrinsics[frame_start:frame_end], + poses=poses[frame_start:frame_end], + world_scale=world_scale, + ) + encoded, encoder_ms = _timed_cuda( + partial( + encoder_stage.encode, + autoregressive_index=autoregressive_index, + cache=encoder_caches[session_index], + input=control, + ) + ) + encoded_sessions.append(encoded) + encoder_times.append(encoder_ms) + frame_starts[session_index] = frame_end + for group_index in range(len(topology.dit_groups)): + start = group_index * topology.sessions_per_group + end = start + topology.sessions_per_group + group_encoder_bundles.append( + encoder_output_to_bundle( + stack_encoder_outputs(encoded_sessions[start:end]) + ) + ) + + pending_encoder = [] + received_leader_bundle = None + for group_index, ranks in enumerate(topology.dit_groups): + leader = ranks[0] + source_bundle = ( + group_encoder_bundles[group_index] if rank == topology.io_rank else None + ) + channel = encoder_channels[group_index] + if channel is None: + channel = _setup_channel( + transport, + pool, + source=topology.io_rank, + destination=leader, + source_bundle=source_bundle, + device=device, + ) + encoder_channels[group_index] = channel + handle, started = _submit_channel(transport, channel, source_bundle) + pending_encoder.append((channel, handle, started, source_bundle)) + if rank == leader: + assert channel.receiver is not None + received_leader_bundle = channel.receiver.bundle + + encoder_transfers = [] + for channel, handle, started, source_bundle in pending_encoder: + stats, handoff_ms = _finish_channel( + transport, + channel, + handle, + started, + ) + if rank == topology.io_rank: + assert source_bundle is not None + transport.unregister(source_bundle) + encoder_transfers.append({**vars(stats), "handoff_ms": handoff_ms}) + + received_encoded = None + for ranks in topology.dit_groups: + fanned, fanout_ms = _fanout_bundle( + source_bundle=(received_leader_bundle if rank == ranks[0] else None), + ranks=ranks, + group=local_groups.get(ranks), + device=device, + ) + if rank in ranks: + received_encoded = fanned + assert fanout_ms is not None + local["pair_fanout_ms"] = fanout_ms + + clean_bundle = None + if dit_stage is not None: + assert received_encoded is not None + encoded_batch = encoder_output_from_bundle(received_encoded) + if args.double_buffered: + assert len(dit_caches) == 2 + session_inputs = split_encoder_outputs(encoded_batch) + clean_latents, local["dit_ms"] = _timed_cuda( + partial( + dit_stage.generate_double_buffered, + autoregressive_index=autoregressive_index, + caches=(dit_caches[0], dit_caches[1]), + inputs=session_inputs, + ) + ) + clean_latent = torch.cat(clean_latents, dim=0) + else: + assert len(dit_caches) == 1 + clean_latent, local["dit_ms"] = _timed_cuda( + partial( + dit_stage.generate, + autoregressive_index=autoregressive_index, + cache=dit_caches[0], + input=encoded_batch, + ) + ) + assert local_group_ranks is not None + local["group"] = list(local_group_ranks) + local["stage_index"] = local_group_ranks.index(rank) + local["schedule"] = ( + "double-buffered" if args.double_buffered else "fixed-batch" + ) + if rank == local_group_ranks[0]: + clean_bundle = {"clean_latent": clean_latent.contiguous()} + + torch.distributed.barrier() + pending_decoder = [] + decoder_inputs: list[TensorBundle] = [] + for group_index, ranks in enumerate(topology.dit_groups): + leader = ranks[0] + channel = decoder_channels[group_index] + if channel is None: + channel = _setup_channel( + transport, + pool, + source=leader, + destination=topology.io_rank, + source_bundle=clean_bundle if rank == leader else None, + device=device, + ) + decoder_channels[group_index] = channel + handle, started = _submit_channel( + transport, + channel, + clean_bundle if rank == leader else None, + ) + pending_decoder.append( + (channel, handle, started, clean_bundle if rank == leader else None) + ) + if rank == topology.io_rank: + assert channel.receiver is not None + decoder_inputs.append(channel.receiver.bundle) + + if dit_stage is not None: + if args.double_buffered: + assert len(dit_caches) == 2 + _, local["finalize_ms"] = _timed_cuda( + partial( + dit_stage.finalize_double_buffered, + autoregressive_index=autoregressive_index, + caches=(dit_caches[0], dit_caches[1]), + ) + ) + else: + assert len(dit_caches) == 1 + _, local["finalize_ms"] = _timed_cuda( + partial( + dit_stage.finalize, + autoregressive_index=autoregressive_index, + cache=dit_caches[0], + ) + ) + + decoder_transfers = [] + for channel, handle, started, source_bundle in pending_decoder: + stats, handoff_ms = _finish_channel( + transport, + channel, + handle, + started, + ) + if rank == channel.source: + assert source_bundle is not None + transport.unregister(source_bundle) + if rank == topology.io_rank: + decoder_transfers.append({**vars(stats), "handoff_ms": handoff_ms}) + + if decoder_stage is not None: + decoder_times: list[float] = [] + output_frames = 0 + session_index = 0 + for group_bundle in decoder_inputs: + clean_batch = group_bundle["clean_latent"] + for clean_session in clean_batch.split(1, dim=0): + decoded, decoder_ms = _timed_cuda( + partial( + decoder_stage.decode, + input=clean_session, + autoregressive_index=autoregressive_index, + cache=decoder_caches[session_index], + ) + ) + decoder_times.append(decoder_ms) + output_frames += decoded.shape[-4] + session_index += 1 + local["decoder_wave_ms"] = sum(decoder_times) + local["output_frames"] = output_frames + + torch.distributed.barrier() + if rank == topology.io_rank: + local["encoder_wave_ms"] = sum(encoder_times) + local["encoder_to_dit"] = encoder_transfers + local["dit_to_decoder"] = decoder_transfers + local["wave_latency_ms"] = (time.perf_counter() - wave_started) * 1000.0 + + gathered: list[dict[str, Any] | None] = [None] * world_size + torch.distributed.all_gather_object(gathered, local) + if rank == topology.io_rank: + io_record = gathered[topology.io_rank] + assert io_record is not None + group_leaders = [] + pair_fanout_ms = [] + for ranks in topology.dit_groups: + leader_record = gathered[ranks[0]] + follower_record = gathered[ranks[1]] + assert leader_record is not None and follower_record is not None + group_leaders.append(leader_record) + pair_fanout_ms.append( + max( + leader_record["pair_fanout_ms"], + follower_record["pair_fanout_ms"], + ) + ) + records.append( + { + "autoregressive_index": autoregressive_index, + "warmup": autoregressive_index < args.warmup_blocks, + **io_record, + "dit_group_leaders": group_leaders, + "pair_fanout_ms": pair_fanout_ms, + } + ) + + rollout_peak_gib = torch.cuda.max_memory_allocated(device) / 2**30 + memory_records = { + "weight_gib": weight_memory_gib, + "cache_initialization_peak_gib": cache_initialization_peak_gib, + "rollout_peak_gib": rollout_peak_gib, + "required_capacity_gib": max(cache_initialization_peak_gib, rollout_peak_gib), + } + gathered_memory: list[dict[str, float] | None] = [None] * world_size + torch.distributed.all_gather_object(gathered_memory, memory_records) + + if rank == topology.io_rank: + baseline = json.loads(args.baseline_json.read_text()) + memory = { + key + "_by_rank": [ + cast(dict[str, float], item)[key] for item in gathered_memory + ] + for key in memory_records + } + summary = _summarize( + records=records, + topology=topology, + p2p_probe=p2p_probe, + memory=memory, + baseline=baseline, + ) + environment = _environment( + args, + prompt=prompt, + world_size=world_size, + module_name="lingbot.disagg.benchmark_pipeline", + ) + environment.update( + { + "allocation": { + "io_rank": topology.io_rank, + "dit_groups": [list(group) for group in topology.dit_groups], + "spare_ranks": list(topology.spare_ranks), + }, + "sessions": topology.session_count, + "sessions_per_group": topology.sessions_per_group, + "schedule": ( + "double-buffered" if args.double_buffered else "fixed-batch" + ), + "compile_network": args.compile_network, + "cuda_graph": False, + "pipeline_layers": [[0, 20], [20, 40]], + "dit_internal_transport": "NCCL P2P over NVLink/NVSwitch", + "stage_transport": args.transport, + } + ) + _write_report( + args, + environment=environment, + records=records, + summary=summary, + p2p_probe=p2p_probe, + ) + + for channel in (*encoder_channels, *decoder_channels): + if channel is not None and channel.receiver is not None: + pool.release(channel.receiver) + pool.close() + transport.close() + torch.distributed.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/lingbot/disagg/benchmark_replicated.py b/integrations/lingbot/lingbot/disagg/benchmark_replicated.py new file mode 100644 index 00000000..07dcb699 --- /dev/null +++ b/integrations/lingbot/lingbot/disagg/benchmark_replicated.py @@ -0,0 +1,961 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Eight-GPU LingBot benchmark with replicated session-affine DiT workers.""" + +from __future__ import annotations + +import argparse +import json +import os +import statistics +import time +from dataclasses import dataclass +from functools import partial +from pathlib import Path +from typing import Any + +import torch + +from flashdreams.infra.transfer import ( + PooledTensorBuffer, + RegisteredTensorPool, + TensorBundle, + TensorTransferHandle, + TensorTransferTicket, + TransferStats, + describe_tensor_bundle, +) +from lingbot.config import PIPELINE_CONFIGS +from lingbot.disagg.benchmark import ( + _bandwidth_probe, + _broadcast_object, + _environment, + _load_encoder_inputs, + _metric_summary, + _timed_cuda, + _transfer_bundle, +) +from lingbot.disagg.stages import ( + LingbotDecoderStage, + LingbotDiTStage, + LingbotEncoderStage, + conditioning_from_bundle, + conditioning_to_bundle, + encoder_output_from_bundle, + encoder_output_to_bundle, +) +from lingbot.encoder.camctrl import CamCtrlInput + +_DEFAULT_BASELINE = Path( + "integrations/lingbot/docs/benchmark_h100_3stage/benchmark.json" +) +"""Tracked 1 encoder : 1 DiT : 1 decoder baseline.""" + + +@dataclass(frozen=True, kw_only=True) +class StageAllocation: + """Replica counts derived from measured per-session service time.""" + + encoder_replicas: int + """Number of encoder workers.""" + + dit_replicas: int + """Number of session-affine DiT workers.""" + + decoder_replicas: int + """Number of decoder workers.""" + + @property + def total_gpus(self) -> int: + """Return the total number of stage GPUs.""" + return self.encoder_replicas + self.dit_replicas + self.decoder_replicas + + +@dataclass(kw_only=True) +class _TransferChannel: + """One fixed-shape receiver allocation and reusable remote ticket.""" + + source: int + destination: int + ticket: TensorTransferTicket + receiver: PooledTensorBuffer | None + + +def allocate_stage_replicas( + *, + total_gpus: int, + encoder_service_ms: float, + dit_service_ms: float, + decoder_service_ms: float, +) -> StageAllocation: + """Greedily allocate extra GPUs to the highest service-time-per-replica stage.""" + if total_gpus < 3: + raise ValueError("At least three GPUs are required for stage disaggregation.") + service_ms = { + "encoder": encoder_service_ms, + "dit": dit_service_ms, + "decoder": decoder_service_ms, + } + if any(value <= 0.0 for value in service_ms.values()): + raise ValueError("Every stage service time must be positive.") + + replicas = {"encoder": 1, "dit": 1, "decoder": 1} + for _ in range(total_gpus - 3): + bottleneck = max( + service_ms, + key=lambda stage: service_ms[stage] / replicas[stage], + ) + replicas[bottleneck] += 1 + return StageAllocation( + encoder_replicas=replicas["encoder"], + dit_replicas=replicas["dit"], + decoder_replicas=replicas["decoder"], + ) + + +def allocation_from_baseline( + baseline: dict[str, Any], + *, + total_gpus: int, +) -> StageAllocation: + """Derive a stage allocation from a three-stage benchmark document.""" + summary = baseline["summary"] + encoder_service_ms = ( + summary["encoder_ms"]["median"] + + summary["encoder_to_dit"]["handoff_ms"]["median"] + ) + dit_service_ms = summary["dit_ms"]["median"] + summary["finalize_ms"]["median"] + decoder_service_ms = ( + summary["decoder_ms"]["median"] + + summary["dit_to_decoder"]["handoff_ms"]["median"] + ) + return allocate_stage_replicas( + total_gpus=total_gpus, + encoder_service_ms=encoder_service_ms, + dit_service_ms=dit_service_ms, + decoder_service_ms=decoder_service_ms, + ) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model", + choices=sorted(PIPELINE_CONFIGS), + default="lingbot-world-fast-taehv-window15-sink3", + ) + parser.add_argument("--example-idx", type=int, default=0) + parser.add_argument("--warmup-blocks", type=int, default=6) + parser.add_argument("--measured-blocks", type=int, default=5) + parser.add_argument("--pixel-height", type=int, default=464) + parser.add_argument("--pixel-width", type=int, default=832) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--dit-replicas", type=int, default=6) + parser.add_argument( + "--co-locate-io", + action="store_true", + help="Host encoder and decoder on rank 0, leaving seven GPUs for DiTs.", + ) + parser.add_argument( + "--pooled-async", + action="store_true", + help="Reuse receiver registrations/tickets and submit non-blocking writes.", + ) + parser.add_argument( + "--transport", + choices=("mooncake", "nixl"), + default="mooncake", + ) + parser.add_argument("--rdma-device", default=None) + parser.add_argument("--bandwidth-probe-mib", type=int, default=256) + parser.add_argument("--bandwidth-probe-iters", type=int, default=8) + parser.add_argument( + "--transport-only", + action="store_true", + help="Probe every stage edge without loading model weights.", + ) + parser.add_argument("--baseline-json", type=Path, default=_DEFAULT_BASELINE) + parser.add_argument( + "--output-dir", + type=Path, + default=Path("outputs/lingbot_disagg_1e6d1d"), + ) + return parser.parse_args() + + +def _read_baseline(path: Path) -> dict[str, Any]: + if not path.is_file(): + raise FileNotFoundError(f"Baseline benchmark not found: {path}") + return json.loads(path.read_text()) + + +def _create_transport(args: argparse.Namespace, *, rank: int) -> Any: + """Construct the selected optional tensor transport.""" + if args.transport == "nixl": + from flashdreams.infra.transfer import NixlTensorTransport + + return NixlTensorTransport(agent_name=f"lingbot-rank-{rank}") + from flashdreams.infra.transfer import MooncakeTensorTransport + + return MooncakeTensorTransport(device_name=args.rdma_device) + + +def _setup_channel( + transport: Any, + pool: RegisteredTensorPool, + *, + source: int, + destination: int, + source_bundle: TensorBundle | None, + device: torch.device, +) -> _TransferChannel: + """Allocate/register one receiver bucket and exchange its ticket once.""" + rank = torch.distributed.get_rank() + descriptors = _broadcast_object( + describe_tensor_bundle(source_bundle) + if rank == source and source_bundle is not None + else None, + source=source, + ) + receiver = pool.acquire(descriptors, device=device) if rank == destination else None + ticket = _broadcast_object( + receiver.ticket if receiver is not None else None, + source=destination, + ) + return _TransferChannel( + source=source, + destination=destination, + ticket=ticket, + receiver=receiver, + ) + + +def _submit_channel( + transport: Any, + channel: _TransferChannel, + source_bundle: TensorBundle | None, +) -> tuple[TensorTransferHandle | None, float]: + """Submit one channel write without introducing a process-group barrier.""" + started = time.perf_counter() + if torch.distributed.get_rank() != channel.source: + return None, started + assert source_bundle is not None + return transport.send_async(source_bundle, channel.ticket), started + + +def _finish_channel( + transport: Any, + channel: _TransferChannel, + handle: TensorTransferHandle | None, + started: float, +) -> tuple[TransferStats, float]: + """Wait on the source and share transfer telemetry with every rank.""" + rank = torch.distributed.get_rank() + stats = transport.wait(handle) if rank == channel.source and handle else None + handoff_ms = ( + (time.perf_counter() - started) * 1000.0 if rank == channel.source else None + ) + stats = _broadcast_object(stats, source=channel.source) + handoff_ms = _broadcast_object(handoff_ms, source=channel.source) + return stats, handoff_ms + + +def _probe_edges( + transport: Any, + *, + encoder_rank: int, + dit_ranks: tuple[int, ...], + decoder_rank: int, + size_mib: int, + iterations: int, + device: torch.device, +) -> dict[str, list[TransferStats]]: + probes: dict[str, list[TransferStats]] = {} + for dit_rank in dit_ranks: + probes[f"encoder_to_dit_{dit_rank}"] = _bandwidth_probe( + transport, + source=encoder_rank, + destination=dit_rank, + size_mib=size_mib, + iterations=iterations, + device=device, + ) + for dit_rank in dit_ranks: + probes[f"dit_{dit_rank}_to_decoder"] = _bandwidth_probe( + transport, + source=dit_rank, + destination=decoder_rank, + size_mib=size_mib, + iterations=iterations, + device=device, + ) + return probes + + +def _summarize( + *, + records: list[dict[str, Any]], + probes: dict[str, list[TransferStats]], + baseline: dict[str, Any], + dit_replicas: int, + total_gpus: int, +) -> dict[str, Any]: + measured = [record for record in records if not record["warmup"]] + frame_count = sum(record["output_frames"] for record in measured) + elapsed_s = sum(record["wave_latency_ms"] for record in measured) / 1000.0 + dit_worker_total_ms = [ + worker["dit_ms"] + worker["finalize_ms"] + for record in measured + for worker in record["dit_workers"] + ] + encoder_transfers = [ + transfer for record in measured for transfer in record["encoder_to_dit"] + ] + decoder_transfers = [ + transfer for record in measured for transfer in record["dit_to_decoder"] + ] + probe_bandwidth = [ + sample.bandwidth_gbps for samples in probes.values() for sample in samples + ] + aggregate_fps = frame_count / elapsed_s + baseline_fps = baseline["summary"]["fps"] + baseline_latency_ms = baseline["summary"]["latency_ms"]["median"] + return { + "aggregate_fps": aggregate_fps, + "per_session_fps": aggregate_fps / dit_replicas, + "throughput_speedup": aggregate_fps / baseline_fps, + "gpu_normalized_speedup": (aggregate_fps / total_gpus) / (baseline_fps / 3), + "wave_latency_ms": _metric_summary( + [record["wave_latency_ms"] for record in measured] + ), + "latency_vs_baseline": statistics.median( + [record["wave_latency_ms"] for record in measured] + ) + / baseline_latency_ms, + "encoder_wave_ms": _metric_summary( + [record["encoder_wave_ms"] for record in measured] + ), + "dit_critical_path_ms": _metric_summary( + [ + max( + worker["dit_ms"] + worker["finalize_ms"] + for worker in record["dit_workers"] + ) + for record in measured + ] + ), + "dit_worker_total_ms": _metric_summary(dit_worker_total_ms), + "decoder_wave_ms": _metric_summary( + [record["decoder_wave_ms"] for record in measured] + ), + "encoder_to_dit": { + "payload_mib_each": encoder_transfers[0]["payload_bytes"] / 2**20, + "copy_ms_each": _metric_summary( + [transfer["transfer_ms"] for transfer in encoder_transfers] + ), + "handoff_ms_each": _metric_summary( + [transfer["handoff_ms"] for transfer in encoder_transfers] + ), + "aggregate_handoff_ms_per_wave": _metric_summary( + [ + sum(item["handoff_ms"] for item in record["encoder_to_dit"]) + for record in measured + ] + ), + }, + "dit_to_decoder": { + "payload_mib_each": decoder_transfers[0]["payload_bytes"] / 2**20, + "copy_ms_each": _metric_summary( + [transfer["transfer_ms"] for transfer in decoder_transfers] + ), + "handoff_ms_each": _metric_summary( + [transfer["handoff_ms"] for transfer in decoder_transfers] + ), + "aggregate_handoff_ms_per_wave": _metric_summary( + [ + sum(item["handoff_ms"] for item in record["dit_to_decoder"]) + for record in measured + ] + ), + }, + "bandwidth_probe_gbps": { + "all_edges": _metric_summary(probe_bandwidth), + "by_edge": { + edge: _metric_summary([sample.bandwidth_gbps for sample in samples]) + for edge, samples in probes.items() + }, + }, + "baseline": { + "fps": baseline_fps, + "latency_ms": baseline_latency_ms, + "topology": "1 encoder : 1 DiT : 1 decoder", + }, + } + + +def _write_report( + args: argparse.Namespace, + *, + records: list[dict[str, Any]], + probes: dict[str, list[TransferStats]], + baseline: dict[str, Any], + environment: dict[str, Any], +) -> None: + summary = _summarize( + records=records, + probes=probes, + baseline=baseline, + dit_replicas=args.dit_replicas, + total_gpus=len(environment["peak_memory_gib_by_rank"]), + ) + args.output_dir.mkdir(parents=True, exist_ok=True) + (args.output_dir / "benchmark.json").write_text( + json.dumps( + { + "environment": environment, + "summary": summary, + "records": records, + "bandwidth_probe": { + edge: [vars(item) for item in samples] + for edge, samples in probes.items() + }, + }, + indent=2, + ) + + "\n" + ) + allocation = environment["allocation"] + peak_memory = environment["peak_memory_gib_by_rank"] + dit_ranks = allocation["rank_roles"]["dit"] + dit_memory = [peak_memory[index] for index in dit_ranks] + io_rank = allocation["rank_roles"]["encoder"][0] + co_located = ( + allocation["rank_roles"]["encoder"] == allocation["rank_roles"]["decoder"] + ) + io_role = "Co-located encoder + decoder" if co_located else "Shared encoder" + handoff_label = ( + "submission → wait window" if environment["pooled_async"] else "handoff" + ) + async_note = "" + if environment["pooled_async"]: + async_note = """ +The asynchronous submission-to-wait windows include useful work performed +before the delayed wait. They are not isolated copy times and must not be added +to the wave latency. Use the 256 MiB probes for link bandwidth. +""" + decoder_memory_row = "" + if not co_located: + decoder_rank = allocation["rank_roles"]["decoder"][0] + decoder_memory_row = f"| Shared decoder | {peak_memory[decoder_rank]:.2f} GiB |" + markdown = f"""# LingBot replicated-DiT disaggregation benchmark + +## Result + +Topology: **{allocation["encoder"]} encoder : {allocation["dit"]} DiT : {allocation["decoder"]} decoder**. +Each DiT worker owns one concurrent session and its resident autoregressive KV cache. + +| Metric | Median | P90 | +| --- | ---: | ---: | +| {args.dit_replicas}-session wave latency | {summary["wave_latency_ms"]["median"]:.2f} ms | {summary["wave_latency_ms"]["p90"]:.2f} ms | +| Encoder wave | {summary["encoder_wave_ms"]["median"]:.2f} ms | {summary["encoder_wave_ms"]["p90"]:.2f} ms | +| DiT critical path | {summary["dit_critical_path_ms"]["median"]:.2f} ms | {summary["dit_critical_path_ms"]["p90"]:.2f} ms | +| Decoder wave | {summary["decoder_wave_ms"]["median"]:.2f} ms | {summary["decoder_wave_ms"]["p90"]:.2f} ms | +| Encoder → DiT {handoff_label}, each | {summary["encoder_to_dit"]["handoff_ms_each"]["median"]:.2f} ms | {summary["encoder_to_dit"]["handoff_ms_each"]["p90"]:.2f} ms | +| DiT → decoder {handoff_label}, each | {summary["dit_to_decoder"]["handoff_ms_each"]["median"]:.2f} ms | {summary["dit_to_decoder"]["handoff_ms_each"]["p90"]:.2f} ms | +| 256 MiB RDMA probes, all edges | {summary["bandwidth_probe_gbps"]["all_edges"]["median"]:.2f} GB/s | {summary["bandwidth_probe_gbps"]["all_edges"]["p90"]:.2f} GB/s | + +- Aggregate throughput: **{summary["aggregate_fps"]:.2f} generated FPS** +- Per-session throughput: **{summary["per_session_fps"]:.2f} generated FPS** +- Throughput versus tracked 1:1:1 baseline: **{summary["throughput_speedup"]:.2f}×** +- Wave latency versus one-session baseline latency: **{summary["latency_vs_baseline"]:.2f}×** +- GPU-normalized throughput versus the three-GPU baseline: **{summary["gpu_normalized_speedup"]:.2f}×** + +The headline excludes {args.warmup_blocks} warmup waves and measures +{args.measured_blocks} waves. It represents {args.dit_replicas} concurrent, session-affine +rollouts, not acceleration of one autoregressive session. +{async_note} + +## Peak allocated memory + +| Role | Peak | +| --- | ---: | +| {io_role} | {peak_memory[io_rank]:.2f} GiB | +| DiT workers | {min(dit_memory):.2f}–{max(dit_memory):.2f} GiB each | +{decoder_memory_row} + +## Reproduction + +```bash +{environment["command"]} +``` + +- Repository revision: `{environment["commit"]}`{" (modified worktree)" if environment["worktree_dirty"] else ""} +- Slurm: job `{environment["slurm_job_id"]}` on `{environment["hostname"]}` +- GPU: `{environment["gpus"][0]}` × {len(environment["gpus"])} +- Model: `{args.model}` +""" + (args.output_dir / "README.md").write_text(markdown) + + +def main() -> None: + """Run one encoder, replicated DiTs, and one decoder on a single node.""" + args = _parse_args() + rank = int(os.environ.get("RANK", "0")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + expected_world_size = args.dit_replicas + (1 if args.co_locate_io else 2) + if world_size != expected_world_size: + raise ValueError( + f"Launch with {expected_world_size} processes for the selected topology; " + f"got {world_size}." + ) + if args.dit_replicas < 1: + raise ValueError("dit-replicas must be positive.") + if args.warmup_blocks < 0 or args.measured_blocks <= 0: + raise ValueError("warmup-blocks must be >= 0 and measured-blocks must be > 0.") + + baseline = _read_baseline(args.baseline_json) + if not args.co_locate_io: + allocation = allocation_from_baseline(baseline, total_gpus=world_size) + expected = StageAllocation( + encoder_replicas=1, + dit_replicas=args.dit_replicas, + decoder_replicas=1, + ) + if allocation != expected: + raise ValueError( + f"Measured service times recommend {allocation}, but the launch " + f"requests {expected}." + ) + + encoder_rank = 0 + dit_ranks = tuple(range(1, 1 + args.dit_replicas)) + decoder_rank = encoder_rank if args.co_locate_io else world_size - 1 + session_count = len(dit_ranks) + + torch.cuda.set_device(local_rank) + device = torch.device(f"cuda:{local_rank}") + config = PIPELINE_CONFIGS[args.model] + + if args.transport_only: + torch.distributed.init_process_group("gloo") + transport = _create_transport(args, rank=rank) + probes = _probe_edges( + transport, + encoder_rank=encoder_rank, + dit_ranks=dit_ranks, + decoder_rank=decoder_rank, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ) + if rank == encoder_rank: + print( + json.dumps( + { + edge: _metric_summary([item.bandwidth_gbps for item in samples]) + for edge, samples in probes.items() + }, + indent=2, + ) + ) + transport.close() + torch.distributed.destroy_process_group() + return + + encoder_stage = ( + LingbotEncoderStage(config).to(device).eval() if rank == encoder_rank else None + ) + dit_stage = LingbotDiTStage(config).to(device).eval() if rank in dit_ranks else None + decoder_stage = ( + LingbotDecoderStage(config).to(device).eval() if rank == decoder_rank else None + ) + encoder_inputs = ( + _load_encoder_inputs(args, device=device) if rank == encoder_rank else None + ) + + torch.distributed.init_process_group("gloo") + transport = _create_transport(args, rank=rank) + pool = RegisteredTensorPool(transport, max_buffers_per_bucket=16) + + encoder_caches: list[Any] = [] + conditioning_bundle = None + height_width = None + prompt = None + image = None + intrinsics = None + poses = None + world_scale = None + if encoder_stage is not None and encoder_inputs is not None: + prompt, image, intrinsics, poses, world_scale = encoder_inputs + for _ in range(session_count): + cache, conditioning = encoder_stage.initialize_cache( + text=[prompt], + image=image, + ) + encoder_caches.append(cache) + if conditioning_bundle is None: + conditioning_bundle = conditioning_to_bundle(conditioning) + height_width = (conditioning.height, conditioning.width) + + height_width = _broadcast_object(height_width, source=encoder_rank) + dit_cache = None + for dit_rank in dit_ranks: + received_context, _, _ = _transfer_bundle( + transport, + source=encoder_rank, + destination=dit_rank, + source_bundle=conditioning_bundle, + device=device, + ) + if rank == dit_rank and dit_stage is not None and received_context is not None: + conditioning = conditioning_from_bundle( + received_context, + height=height_width[0], + width=height_width[1], + ) + dit_cache = dit_stage.initialize_cache(conditioning) + transport.unregister(received_context) + + decoder_caches = ( + [decoder_stage.initialize_cache() for _ in range(session_count)] + if decoder_stage is not None + else [] + ) + probes = _probe_edges( + transport, + encoder_rank=encoder_rank, + dit_ranks=dit_ranks, + decoder_rank=decoder_rank, + size_mib=args.bandwidth_probe_mib, + iterations=args.bandwidth_probe_iters, + device=device, + ) + + total_blocks = args.warmup_blocks + args.measured_blocks + frame_starts = [0] * session_count + records: list[dict[str, Any]] = [] + encoder_channels: list[_TransferChannel | None] = [None] * session_count + decoder_channels: list[_TransferChannel | None] = [None] * session_count + for autoregressive_index in range(total_blocks): + torch.distributed.barrier() + wave_started = time.perf_counter() + local: dict[str, Any] = {} + encoder_times: list[float] = [] + encoder_transfers: list[dict[str, Any]] = [] + received_encoded = None + pending_encoder: list[ + tuple[ + _TransferChannel, + TensorTransferHandle | None, + float, + TensorBundle | None, + ] + ] = [] + + for session_index, dit_rank in enumerate(dit_ranks): + encoded_bundle = None + if encoder_stage is not None: + assert intrinsics is not None + assert poses is not None + assert world_scale is not None + num_input_frames = encoder_stage.get_num_input_frames( + autoregressive_index + ) + frame_start = frame_starts[session_index] + frame_end = frame_start + num_input_frames + if frame_end > poses.shape[0]: + raise RuntimeError( + f"Example camera trajectory ended at frame {poses.shape[0]} " + f"before AR block {autoregressive_index}." + ) + control = CamCtrlInput( + intrinsics=intrinsics[frame_start:frame_end], + poses=poses[frame_start:frame_end], + world_scale=world_scale, + ) + encoded, encoder_ms = _timed_cuda( + partial( + encoder_stage.encode, + autoregressive_index=autoregressive_index, + cache=encoder_caches[session_index], + input=control, + ) + ) + encoded_bundle = encoder_output_to_bundle(encoded) + encoder_times.append(encoder_ms) + frame_starts[session_index] = frame_end + + if args.pooled_async: + channel = encoder_channels[session_index] + if channel is None: + channel = _setup_channel( + transport, + pool, + source=encoder_rank, + destination=dit_rank, + source_bundle=encoded_bundle, + device=device, + ) + encoder_channels[session_index] = channel + handle, transfer_started = _submit_channel( + transport, + channel, + encoded_bundle, + ) + pending_encoder.append( + (channel, handle, transfer_started, encoded_bundle) + ) + if rank == dit_rank: + assert channel.receiver is not None + received_encoded = channel.receiver.bundle + else: + received, stats, handoff_ms = _transfer_bundle( + transport, + source=encoder_rank, + destination=dit_rank, + source_bundle=encoded_bundle, + device=device, + ) + if rank == dit_rank: + received_encoded = received + if rank == encoder_rank: + transfer_record = dict(vars(stats)) + transfer_record["handoff_ms"] = handoff_ms + encoder_transfers.append(transfer_record) + + if args.pooled_async: + for channel, handle, transfer_started, source_bundle in pending_encoder: + stats, handoff_ms = _finish_channel( + transport, + channel, + handle, + transfer_started, + ) + if rank == encoder_rank: + assert source_bundle is not None + transport.unregister(source_bundle) + transfer_record = dict(vars(stats)) + transfer_record["handoff_ms"] = handoff_ms + encoder_transfers.append(transfer_record) + torch.distributed.barrier() + + clean_bundle = None + if rank in dit_ranks: + assert dit_stage is not None + assert dit_cache is not None + assert received_encoded is not None + encoded = encoder_output_from_bundle(received_encoded) + clean_latent, local["dit_ms"] = _timed_cuda( + partial( + dit_stage.generate, + autoregressive_index=autoregressive_index, + cache=dit_cache, + input=encoded, + ) + ) + if not args.pooled_async: + _, local["finalize_ms"] = _timed_cuda( + partial( + dit_stage.finalize, + autoregressive_index=autoregressive_index, + cache=dit_cache, + ) + ) + local["session_index"] = dit_ranks.index(rank) + local["rank"] = rank + clean_bundle = {"clean_latent": clean_latent.contiguous()} + if not args.pooled_async: + transport.unregister(received_encoded) + + torch.distributed.barrier() + decoder_inputs: list[Any] = [] + decoder_transfers: list[dict[str, Any]] = [] + pending_decoder: list[ + tuple[ + _TransferChannel, + TensorTransferHandle | None, + float, + TensorBundle | None, + ] + ] = [] + for session_index, dit_rank in enumerate(dit_ranks): + if args.pooled_async: + channel = decoder_channels[session_index] + if channel is None: + channel = _setup_channel( + transport, + pool, + source=dit_rank, + destination=decoder_rank, + source_bundle=clean_bundle if rank == dit_rank else None, + device=device, + ) + decoder_channels[session_index] = channel + handle, transfer_started = _submit_channel( + transport, + channel, + clean_bundle if rank == dit_rank else None, + ) + pending_decoder.append( + ( + channel, + handle, + transfer_started, + clean_bundle if rank == dit_rank else None, + ) + ) + if rank == dit_rank: + assert dit_stage is not None + assert dit_cache is not None + _, local["finalize_ms"] = _timed_cuda( + partial( + dit_stage.finalize, + autoregressive_index=autoregressive_index, + cache=dit_cache, + ) + ) + if rank == decoder_rank: + assert channel.receiver is not None + decoder_inputs.append(channel.receiver.bundle) + else: + received_clean, stats, handoff_ms = _transfer_bundle( + transport, + source=dit_rank, + destination=decoder_rank, + source_bundle=clean_bundle if rank == dit_rank else None, + device=device, + ) + if rank == decoder_rank: + assert received_clean is not None + decoder_inputs.append(received_clean) + if rank == encoder_rank: + transfer_record = dict(vars(stats)) + transfer_record["handoff_ms"] = handoff_ms + decoder_transfers.append(transfer_record) + + if args.pooled_async: + for channel, handle, transfer_started, source_bundle in pending_decoder: + stats, handoff_ms = _finish_channel( + transport, + channel, + handle, + transfer_started, + ) + if rank == channel.source: + assert source_bundle is not None + transport.unregister(source_bundle) + if rank == encoder_rank: + transfer_record = dict(vars(stats)) + transfer_record["handoff_ms"] = handoff_ms + decoder_transfers.append(transfer_record) + torch.distributed.barrier() + + if decoder_stage is not None: + decoder_times: list[float] = [] + output_frames = 0 + for session_index, received_clean in enumerate(decoder_inputs): + decoded, decoder_ms = _timed_cuda( + partial( + decoder_stage.decode, + input=received_clean["clean_latent"], + autoregressive_index=autoregressive_index, + cache=decoder_caches[session_index], + ) + ) + decoder_times.append(decoder_ms) + output_frames += decoded.shape[-4] + if not args.pooled_async: + transport.unregister(received_clean) + local["decoder_times_ms"] = decoder_times + local["decoder_wave_ms"] = sum(decoder_times) + local["output_frames"] = output_frames + + torch.distributed.barrier() + if rank == encoder_rank: + local["encoder_times_ms"] = encoder_times + local["encoder_wave_ms"] = sum(encoder_times) + local["encoder_to_dit"] = encoder_transfers + local["dit_to_decoder"] = decoder_transfers + local["wave_latency_ms"] = (time.perf_counter() - wave_started) * 1000.0 + + gathered: list[dict[str, Any] | None] = [None] * world_size + torch.distributed.all_gather_object(gathered, local) + if rank == encoder_rank: + encoder_record = gathered[encoder_rank] + decoder_record = gathered[decoder_rank] + assert encoder_record is not None + assert decoder_record is not None + dit_worker_records: list[dict[str, Any]] = [] + for dit_rank in dit_ranks: + dit_worker_record = gathered[dit_rank] + assert dit_worker_record is not None + dit_worker_records.append(dit_worker_record) + record = { + "autoregressive_index": autoregressive_index, + "warmup": autoregressive_index < args.warmup_blocks, + **encoder_record, + "decoder_times_ms": decoder_record["decoder_times_ms"], + "decoder_wave_ms": decoder_record["decoder_wave_ms"], + "output_frames": decoder_record["output_frames"], + "dit_workers": dit_worker_records, + } + records.append(record) + + peak_memory = torch.cuda.max_memory_allocated(device) / 2**30 + peak_memory_by_rank: list[float | None] = [None] * world_size + torch.distributed.all_gather_object(peak_memory_by_rank, peak_memory) + if rank == encoder_rank: + environment = _environment( + args, + prompt=prompt, + world_size=world_size, + module_name="lingbot.disagg.benchmark_replicated", + ) + environment["allocation"] = { + "encoder": 1, + "dit": args.dit_replicas, + "decoder": 1, + "physical_gpus": world_size, + "co_located_io": args.co_locate_io, + "rank_roles": { + "encoder": [encoder_rank], + "dit": list(dit_ranks), + "decoder": [decoder_rank], + }, + } + environment["transport"] = args.transport + environment["pooled_async"] = args.pooled_async + environment["sessions_per_wave"] = session_count + environment["peak_memory_gib_by_rank"] = peak_memory_by_rank + _write_report( + args, + records=records, + probes=probes, + baseline=baseline, + environment=environment, + ) + + for channel in (*encoder_channels, *decoder_channels): + if channel is not None and channel.receiver is not None: + pool.release(channel.receiver) + pool.close() + transport.close() + torch.distributed.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/lingbot/disagg/scheduler.py b/integrations/lingbot/lingbot/disagg/scheduler.py new file mode 100644 index 00000000..e526b8be --- /dev/null +++ b/integrations/lingbot/lingbot/disagg/scheduler.py @@ -0,0 +1,207 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Session-affine scheduling policies for LingBot disaggregated workers.""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from enum import Enum + + +class ServiceClass(str, Enum): + """Deployment pool selected for an interactive session.""" + + LATENCY = "latency" + THROUGHPUT = "throughput" + + +@dataclass(frozen=True, kw_only=True) +class SessionRequest: + """Placement constraints supplied when a session opens.""" + + session_id: str + """Stable interactive-session identifier.""" + + shape: tuple[int, ...] + """Latent or pixel shape used for fixed-shape compatibility.""" + + cp_size: int + """Required context-parallel group size.""" + + service_class: ServiceClass = ServiceClass.THROUGHPUT + """Latency or aggregate-throughput deployment pool.""" + + preferred_rack: str | None = None + """Optional rack locality hint.""" + + preferred_nic: str | None = None + """Optional NIC or fabric locality hint.""" + + +@dataclass(frozen=True, kw_only=True) +class WorkerSnapshot: + """Current DiT worker-group state used for one placement decision.""" + + worker_id: str + """Stable worker or CP-group identifier.""" + + pool: str + """Deployment pool, such as ``aggregated-cp8`` or ``io-plus-7-dit``.""" + + queue_depth: int + """Queued autoregressive chunks.""" + + predicted_chunk_ms: float + """Expected service time for one compatible chunk.""" + + free_hbm_gib: float + """Currently available device memory.""" + + supported_shapes: frozenset[tuple[int, ...]] + """Shapes for which this worker has compiled kernels and buffers.""" + + supported_cp_sizes: frozenset[int] + """Context-parallel group sizes hosted by this worker.""" + + resident_sessions: frozenset[str] = frozenset() + """Sessions whose autoregressive KV state is already resident.""" + + rack: str | None = None + """Rack locality, when known.""" + + nic: str | None = None + """GPU-local NIC or fabric rail, when known.""" + + rdma_capable: bool = True + """Whether the selected path is verified as RDMA rather than TCP.""" + + +@dataclass(frozen=True, kw_only=True) +class ScheduledMicrobatch: + """Compatible independent sessions that may share one DiT launch.""" + + worker_id: str + shape: tuple[int, ...] + cp_size: int + session_ids: tuple[str, ...] + + +class SessionAwareScheduler: + """Place once, preserve cache affinity, and reject non-RDMA fallbacks.""" + + def __init__( + self, + *, + min_free_hbm_gib: float = 0.0, + require_rdma: bool = True, + ) -> None: + self.min_free_hbm_gib = min_free_hbm_gib + self.require_rdma = require_rdma + self._placements: dict[str, str] = {} + + @staticmethod + def pool_for(service_class: ServiceClass) -> str: + """Return the recommended fixed deployment pool.""" + if service_class is ServiceClass.LATENCY: + return "aggregated-cp8" + return "io-plus-7-dit" + + def assign( + self, + request: SessionRequest, + workers: Sequence[WorkerSnapshot], + ) -> str: + """Return a sticky compatible worker placement for ``request``.""" + existing = self._placements.get(request.session_id) + if existing is not None: + if any(worker.worker_id == existing for worker in workers): + return existing + raise RuntimeError( + f"Session {request.session_id!r} lost resident worker {existing!r}; " + "restore its cache before rerouting." + ) + + target_pool = self.pool_for(request.service_class) + compatible = [ + worker + for worker in workers + if worker.pool == target_pool + and request.shape in worker.supported_shapes + and request.cp_size in worker.supported_cp_sizes + and worker.free_hbm_gib >= self.min_free_hbm_gib + and (worker.rdma_capable or not self.require_rdma) + ] + if not compatible: + raise RuntimeError( + f"No compatible {target_pool!r} worker for shape={request.shape}, " + f"CP{request.cp_size}; TCP fallback is disabled={self.require_rdma}." + ) + + def score(worker: WorkerSnapshot) -> tuple[float, int, int, str]: + locality_penalty = 0 + if request.preferred_rack is not None: + locality_penalty += worker.rack != request.preferred_rack + if request.preferred_nic is not None: + locality_penalty += worker.nic != request.preferred_nic + predicted_wait_ms = (worker.queue_depth + 1) * worker.predicted_chunk_ms + residency_penalty = request.session_id not in worker.resident_sessions + return ( + predicted_wait_ms, + locality_penalty, + residency_penalty, + worker.worker_id, + ) + + selected = min(compatible, key=score) + self._placements[request.session_id] = selected.worker_id + return selected.worker_id + + def release(self, session_id: str) -> None: + """Forget placement after all stage caches have been destroyed.""" + self._placements.pop(session_id, None) + + +def build_microbatches( + assignments: Iterable[tuple[SessionRequest, str]], + *, + max_batch_size: int, +) -> tuple[ScheduledMicrobatch, ...]: + """Group independent compatible sessions for a future fused DiT launch. + + The scheduler preserves session identity; the model runtime must still + implement batched cache gather/scatter before these groups can share one + kernel launch. + """ + if max_batch_size < 1: + raise ValueError(f"max_batch_size must be positive, got {max_batch_size}.") + groups: dict[tuple[str, tuple[int, ...], int], list[str]] = defaultdict(list) + for request, worker_id in assignments: + groups[(worker_id, request.shape, request.cp_size)].append(request.session_id) + + batches: list[ScheduledMicrobatch] = [] + for (worker_id, shape, cp_size), session_ids in sorted(groups.items()): + for start in range(0, len(session_ids), max_batch_size): + batches.append( + ScheduledMicrobatch( + worker_id=worker_id, + shape=shape, + cp_size=cp_size, + session_ids=tuple(session_ids[start : start + max_batch_size]), + ) + ) + return tuple(batches) diff --git a/integrations/lingbot/lingbot/disagg/stages.py b/integrations/lingbot/lingbot/disagg/stages.py new file mode 100644 index 00000000..7aa2503a --- /dev/null +++ b/integrations/lingbot/lingbot/disagg/stages.py @@ -0,0 +1,531 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""LingBot-specific encoder, DiT, and decoder service stages.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from einops import rearrange +from torch import Tensor, nn +from torch.distributed import ProcessGroup + +from flashdreams.infra.decoder import StreamingVideoDecoder +from flashdreams.infra.diffusion.scheduler import FlowMatchScheduler +from flashdreams.infra.pipeline import DecoderStage, DiffusionStage, DiffusionStageCache +from flashdreams.infra.transfer import TensorBundle +from flashdreams.recipes.wan.autoencoder.i2v import I2VCtrl +from flashdreams.recipes.wan.autoencoder.vae import WanVAECache +from flashdreams.recipes.wan.transformer.constants import NEGATIVE_PROMPT +from lingbot.encoder.camctrl import ( + CamCtrlInput, + I2VCamCtrlEmbeddings, + I2VCamCtrlEncoder, + I2VCamCtrlEncoderCache, + I2VCamCtrlInput, +) +from lingbot.pipeline import LingbotWorldInferencePipelineConfig +from lingbot.transformer import ( + LingbotWorldTransformer, + LingbotWorldTransformerCache, + LingbotWorldTransformerConfig, +) + + +@dataclass(kw_only=True) +class LingbotConditioning: + """One-shot encoder output transferred into the DiT worker.""" + + height: int + """Pre-patchify latent height.""" + + width: int + """Pre-patchify latent width.""" + + text_embeddings: Tensor + """Positive UMT5 prompt embeddings.""" + + negative_text_embeddings: Tensor | None = None + """Negative UMT5 embeddings when classifier-free guidance is enabled.""" + + image_embeddings: Tensor | None = None + """Optional CLIP first-frame embeddings.""" + + +@dataclass(kw_only=True) +class LingbotEncoderStageCache: + """Per-session state retained by the LingBot encoder worker.""" + + encoder_cache: I2VCamCtrlEncoderCache + """Streaming VAE and camera-control cache.""" + + image: Tensor + """First-frame pixels used to construct every I2V input chunk.""" + + +def conditioning_to_bundle(conditioning: LingbotConditioning) -> TensorBundle: + """Flatten one-shot conditioning into a transferable tensor bundle.""" + bundle = {"text_embeddings": conditioning.text_embeddings.contiguous()} + if conditioning.negative_text_embeddings is not None: + bundle["negative_text_embeddings"] = ( + conditioning.negative_text_embeddings.contiguous() + ) + if conditioning.image_embeddings is not None: + bundle["image_embeddings"] = conditioning.image_embeddings.contiguous() + return bundle + + +def conditioning_from_bundle( + bundle: TensorBundle, + *, + height: int, + width: int, +) -> LingbotConditioning: + """Reconstruct one-shot conditioning after a transfer.""" + return LingbotConditioning( + height=height, + width=width, + text_embeddings=bundle["text_embeddings"], + negative_text_embeddings=bundle.get("negative_text_embeddings"), + image_embeddings=bundle.get("image_embeddings"), + ) + + +def encoder_output_to_bundle(output: I2VCamCtrlEmbeddings) -> TensorBundle: + """Flatten per-step LingBot encoder output for transport.""" + assert not output._is_patchified, ( + "Encoder output must cross the service boundary before DiT patchification." + ) + return { + "i2v.latent": output.i2v.latent.contiguous(), + "i2v.mask": output.i2v.mask.contiguous(), + "plucker": output.plucker.contiguous(), + } + + +def encoder_output_from_bundle( + bundle: TensorBundle, + *, + patchified: bool = False, +) -> I2VCamCtrlEmbeddings: + """Reconstruct per-step LingBot encoder output after transport.""" + return I2VCamCtrlEmbeddings( + i2v=I2VCtrl( + latent=bundle["i2v.latent"], + mask=bundle["i2v.mask"], + _is_patchified=patchified, + ), + plucker=bundle["plucker"], + _is_patchified=patchified, + ) + + +def encoder_output_to_cp_bundles( + output: I2VCamCtrlEmbeddings, + *, + cp_size: int, + patch_size: tuple[int, int, int], +) -> tuple[TensorBundle, ...]: + """Patchify once on the encoder and return one direct-transfer shard per rank. + + This removes the DiT leader's input broadcast/fan-out. Each destination + rank receives only its token shard and reconstructs the payload with + ``patchified=True``. + """ + if cp_size < 1: + raise ValueError(f"cp_size must be positive, got {cp_size}.") + if output._is_patchified or output.i2v._is_patchified: + raise ValueError("Expected raw encoder output before patchification.") + + def patchify(tensor: Tensor) -> Tensor: + kt, kh, kw = patch_size + return rearrange( + tensor, + "... (t kt) c (h kh) (w kw) -> ... (t h w) (c kt kh kw)", + kt=kt, + kh=kh, + kw=kw, + ).contiguous() + + patched = { + "i2v.latent": patchify(output.i2v.latent), + "i2v.mask": patchify(output.i2v.mask), + "plucker": patchify(output.plucker), + } + token_counts = {name: tensor.shape[-2] for name, tensor in patched.items()} + if len(set(token_counts.values())) != 1: + raise ValueError(f"Patchified token counts do not match: {token_counts}.") + token_count = next(iter(token_counts.values())) + if token_count % cp_size: + raise ValueError( + f"Patchified token count {token_count} is not divisible by CP{cp_size}." + ) + + per_field = { + name: tensor.chunk(cp_size, dim=-2) for name, tensor in patched.items() + } + return tuple( + {name: shards[rank].contiguous() for name, shards in per_field.items()} + for rank in range(cp_size) + ) + + +class LingbotEncoderStage(nn.Module): + """Own LingBot's text, image, streaming VAE, and camera encoders.""" + + def __init__(self, config: LingbotWorldInferencePipelineConfig) -> None: + super().__init__() + assert config.encoder is not None, "LingBot requires an I2V control encoder." + encoder = config.encoder.setup() + assert isinstance(encoder, I2VCamCtrlEncoder) + self.encoder = encoder + self.text_encoder = ( + config.text_encoder.setup() if config.text_encoder is not None else None + ) + self.image_encoder = ( + config.image_encoder.setup() if config.image_encoder is not None else None + ) + transformer_config = config.diffusion_model.transformer + assert isinstance(transformer_config, LingbotWorldTransformerConfig) + self.transformer_config = transformer_config + + @torch.no_grad() + def initialize_cache( + self, + *, + text: list[str], + image: Tensor, + ) -> tuple[LingbotEncoderStageCache, LingbotConditioning]: + """Encode session-level context and initialize streaming encoder state.""" + assert self.text_encoder is not None, "LingBot text encoder is not configured." + assert image.shape[-4] == 1, ( + f"image must contain exactly one frame, got shape {tuple(image.shape)}." + ) + spatial_ratio = self.encoder.spatial_compression_ratio + pixel_height, pixel_width = image.shape[-2:] + assert pixel_height % spatial_ratio == 0 + assert pixel_width % spatial_ratio == 0 + + text_embeddings = self.text_encoder(text) + negative_text_embeddings = None + if self.transformer_config.guidance_scale > 1.0: + negative_text_embeddings = self.text_encoder([NEGATIVE_PROMPT] * len(text)) + image_embeddings = None + if self.image_encoder is not None: + image_embeddings = self.image_encoder(image.squeeze(-4)) + + cache = LingbotEncoderStageCache( + encoder_cache=self.encoder.initialize_autoregressive_cache(), + image=image, + ) + conditioning = LingbotConditioning( + height=pixel_height // spatial_ratio, + width=pixel_width // spatial_ratio, + text_embeddings=text_embeddings, + negative_text_embeddings=negative_text_embeddings, + image_embeddings=image_embeddings, + ) + return cache, conditioning + + def get_num_input_frames(self, autoregressive_index: int) -> int: + """Return the pixel frames consumed by one encoder step.""" + return self.encoder.get_input_temporal_size( + autoregressive_index, + self.transformer_config.len_t, + ) + + def _preprocess_i2v_input( + self, + autoregressive_index: int, + image: Tensor, + ) -> Tensor: + """Build the first-frame-plus-padding chunk expected by the Wan VAE.""" + expected_frames = self.get_num_input_frames(autoregressive_index) + if autoregressive_index == 0: + return F.pad(image, (0, 0, 0, 0, 0, 0, 0, expected_frames - 1)) + return torch.zeros( + *image.shape[:-4], + expected_frames, + 3, + image.shape[-2], + image.shape[-1], + device=image.device, + dtype=image.dtype, + ) + + @torch.no_grad() + def encode( + self, + *, + autoregressive_index: int, + cache: LingbotEncoderStageCache, + input: CamCtrlInput, + ) -> I2VCamCtrlEmbeddings: + """Encode one camera-control and I2V chunk.""" + i2v_input = self._preprocess_i2v_input( + autoregressive_index, + cache.image, + ) + return self.encoder( + input=I2VCamCtrlInput(i2v=i2v_input, camctrl=input), + autoregressive_index=autoregressive_index, + cache=cache.encoder_cache, + ) + + +class LingbotDiTStage(DiffusionStage[LingbotWorldTransformerCache]): + """Own LingBot's scheduler, DiT weights, and evolving KV cache.""" + + def __init__(self, config: LingbotWorldInferencePipelineConfig) -> None: + super().__init__(config.diffusion_model) + + def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: + """Bind a DiT-only context-parallel group before cache construction.""" + transformer = self.diffusion_model.transformer + assert isinstance(transformer, LingbotWorldTransformer) + transformer.set_context_parallel_group(cp_group) + + def configure_pipeline_parallel( + self, + *, + stage_index: int, + stage_count: int, + group: ProcessGroup, + ranks: tuple[int, ...], + ) -> None: + """Partition the DiT and bind its ordered NCCL rank group. + + Args: + stage_index: Zero-based position inside the pipeline group. + stage_count: Number of pipeline stages. + group: NCCL process group containing the pipeline ranks. + ranks: Global ranks ordered from input to output stage. + """ + transformer = self.diffusion_model.transformer + assert isinstance(transformer, LingbotWorldTransformer) + transformer.configure_pipeline_parallel( + stage_index=stage_index, + stage_count=stage_count, + group=group, + ranks=ranks, + ) + + def initialize_cache( + self, + conditioning: LingbotConditioning, + ) -> DiffusionStageCache[LingbotWorldTransformerCache]: + """Build the resident DiT cache from encoder-stage conditioning.""" + return super().initialize_cache( + height=conditioning.height, + width=conditioning.width, + text_embeddings=conditioning.text_embeddings, + negative_text_embeddings=conditioning.negative_text_embeddings, + image_embeddings=conditioning.image_embeddings, + ) + + @torch.no_grad() + def generate_double_buffered( + self, + *, + autoregressive_index: int, + caches: tuple[ + DiffusionStageCache[LingbotWorldTransformerCache], + DiffusionStageCache[LingbotWorldTransformerCache], + ], + inputs: tuple[I2VCamCtrlEmbeddings, I2VCamCtrlEmbeddings], + ) -> tuple[Tensor, Tensor]: + """Denoise two sessions with an overlapped two-stage DiT schedule. + + Args: + autoregressive_index: Shared autoregressive chunk index. + caches: Two session-affine DiT caches. + inputs: Two unpatchified encoder outputs. + + Returns: + Two clean unpatchified latents in session order. + + Raises: + RuntimeError: The configured scheduler or diffusion mode is unsupported. + AssertionError: Either session is generated out of order. + """ + model = self.diffusion_model + transformer = model.transformer + if not isinstance(transformer, LingbotWorldTransformer): + raise RuntimeError("Double buffering requires LingbotWorldTransformer.") + scheduler = model.scheduler + if not isinstance(scheduler, FlowMatchScheduler): + raise RuntimeError("Double buffering requires FlowMatchScheduler.") + if model.config.noise_in_unpatchified_shape: + raise RuntimeError("Double buffering requires patchified scheduler noise.") + + for cache in caches: + previous = cache.autoregressive_index + expected = previous + 1 if previous is not None else 0 + assert autoregressive_index == expected, ( + f"AR step out of order: previous step was {previous}, expected " + f"{expected}, got {autoregressive_index}." + ) + + patchified_inputs_list: list[I2VCamCtrlEmbeddings] = [] + for input in inputs: + patchified = transformer.patchify_and_maybe_split_cp(input) + assert isinstance(patchified, I2VCamCtrlEmbeddings) + patchified_inputs_list.append(patchified) + patchified_inputs = (patchified_inputs_list[0], patchified_inputs_list[1]) + + transformer_caches = ( + caches[0].transformer_cache, + caches[1].transformer_cache, + ) + for cache in transformer_caches: + cache.start(autoregressive_index) + + input_dtype = model.dtype + noisy_latents = ( + torch.randn( + model.latent_shape, + device=model.device, + dtype=input_dtype, + generator=model.rng, + ), + torch.randn( + model.latent_shape, + device=model.device, + dtype=input_dtype, + generator=model.rng, + ), + ) + clean_latents: tuple[Tensor, Tensor] | None = None + for step_index in range(scheduler.denoising_step_list.shape[0]): + sigma = scheduler.denoising_sigmas[step_index] + timestep = scheduler.denoising_step_list[step_index].to(dtype=input_dtype) + if step_index > 0: + assert clean_latents is not None + noises = ( + torch.empty_like(noisy_latents[0]).normal_(generator=model.rng), + torch.empty_like(noisy_latents[1]).normal_(generator=model.rng), + ) + noisy_latents = ( + ((1.0 - sigma) * clean_latents[0] + sigma * noises[0]).to( + input_dtype + ), + ((1.0 - sigma) * clean_latents[1] + sigma * noises[1]).to( + input_dtype + ), + ) + flows = transformer.predict_flow_double_buffered( + noisy_latents=noisy_latents, + timesteps=(timestep, timestep), + caches=transformer_caches, + inputs=patchified_inputs, + ) + clean_latents = ( + noisy_latents[0] - sigma * flows[0], + noisy_latents[1] - sigma * flows[1], + ) + assert clean_latents is not None + + patchified_clean = ( + transformer.postprocess_clean_latent( + clean_latent=clean_latents[0].to(input_dtype), + cache=transformer_caches[0], + input=patchified_inputs[0].i2v, + ), + transformer.postprocess_clean_latent( + clean_latent=clean_latents[1].to(input_dtype), + cache=transformer_caches[1], + input=patchified_inputs[1].i2v, + ), + ) + for index, cache in enumerate(caches): + cache.autoregressive_index = autoregressive_index + cache.final_state = model.FinalState( + clean_latent=patchified_clean[index], + autoregressive_index=autoregressive_index, + cache=transformer_caches[index], + input=patchified_inputs[index], + ) + + return ( + transformer.unpatchify_and_maybe_gather_cp(patchified_clean[0]), + transformer.unpatchify_and_maybe_gather_cp(patchified_clean[1]), + ) + + @torch.no_grad() + def finalize_double_buffered( + self, + *, + autoregressive_index: int, + caches: tuple[ + DiffusionStageCache[LingbotWorldTransformerCache], + DiffusionStageCache[LingbotWorldTransformerCache], + ], + ) -> None: + """Advance two session caches with an overlapped pipeline forward. + + Args: + autoregressive_index: Shared autoregressive chunk index. + caches: Two session-affine DiT caches after generation. + + Raises: + RuntimeError: Context-noise finalization is enabled. + AssertionError: A cache does not contain the matching final state. + """ + model = self.diffusion_model + transformer = model.transformer + assert isinstance(transformer, LingbotWorldTransformer) + if model.config.context_noise != 0: + raise RuntimeError("Double buffering requires context_noise == 0.") + + final_states = (caches[0].final_state, caches[1].final_state) + for cache, final_state in zip(caches, final_states): + assert cache.autoregressive_index == autoregressive_index + assert final_state is not None + + state_0 = final_states[0] + state_1 = final_states[1] + assert state_0 is not None and state_1 is not None + timestep = torch.tensor(0, device=model.device, dtype=model.dtype) + inputs = (state_0.input, state_1.input) + assert isinstance(inputs[0], I2VCamCtrlEmbeddings) + assert isinstance(inputs[1], I2VCamCtrlEmbeddings) + transformer.predict_flow_double_buffered( + noisy_latents=(state_0.clean_latent, state_1.clean_latent), + timesteps=(timestep, timestep), + caches=(state_0.cache, state_1.cache), + inputs=(inputs[0], inputs[1]), + ) + state_0.cache.finalize(autoregressive_index) + state_1.cache.finalize(autoregressive_index) + caches[0].final_state = None + caches[1].final_state = None + + +class LingbotDecoderStage(DecoderStage[WanVAECache]): + """Own LingBot's streaming pixel decoder.""" + + decoder: StreamingVideoDecoder[WanVAECache] + + def __init__(self, config: LingbotWorldInferencePipelineConfig) -> None: + assert config.decoder is not None, "LingBot requires a video decoder." + super().__init__(config.decoder) + assert isinstance(self.decoder, StreamingVideoDecoder) + + def get_num_output_frames(self, autoregressive_index: int, len_t: int) -> int: + """Return decoded pixel frames produced by one latent chunk.""" + return self.decoder.get_output_temporal_size(autoregressive_index, len_t) diff --git a/integrations/lingbot/lingbot/transformer/__init__.py b/integrations/lingbot/lingbot/transformer/__init__.py index f85fe4a6..82a653df 100644 --- a/integrations/lingbot/lingbot/transformer/__init__.py +++ b/integrations/lingbot/lingbot/transformer/__init__.py @@ -18,10 +18,12 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import overload +from collections.abc import Callable +from typing import cast, overload import torch from torch import Tensor +from torch.distributed import ProcessGroup from flashdreams.recipes.wan.transformer.wan21 import ( Wan21Transformer, @@ -70,7 +72,7 @@ class LingbotWorldTransformerConfig(Wan21TransformerConfig): :class:`Wan21TransformerConfig`). """ - _target: type["LingbotWorldTransformer"] = field( + _target: type[LingbotWorldTransformer] = field( default_factory=lambda: LingbotWorldTransformer ) @@ -85,6 +87,82 @@ class LingbotWorldTransformer(Wan21Transformer): def __init__(self, config: LingbotWorldTransformerConfig) -> None: super().__init__(config) + self._pipeline_double_buffer_call: ( + Callable[..., tuple[Tensor, Tensor]] | None + ) = None + + def set_context_parallel_group(self, cp_group: ProcessGroup | None) -> None: + """Bind a DiT-only context-parallel group before cache construction. + + The normal runner uses the global process group that exists when the + transformer is constructed. A disaggregated deployment constructs the + stage-local weights first, then creates a subgroup containing only DiT + ranks. Rebinding is safe until a rollout cache or CUDA graph has bound + shapes and storage. + + Args: + cp_group: DiT-only process group, or ``None`` to disable CP. + + Raises: + RuntimeError: A rollout cache has already been initialized. + """ + if self._output_height is not None or self._output_width is not None: + raise RuntimeError( + "Context parallelism must be configured before initializing " + "a LingBot rollout cache." + ) + self._cp_group = cp_group + self._cp_size = cp_group.size() if cp_group is not None else 1 + network = getattr(self.network, "_orig_mod", self.network) + network.set_context_parallel_group(cp_group) + if self._use_cuda_graph: + self._cuda_graph_dispatch.reset() + + def configure_pipeline_parallel( + self, + *, + stage_index: int, + stage_count: int, + group: ProcessGroup, + ranks: tuple[int, ...], + ) -> None: + """Partition DiT layers and bind a fixed NCCL pipeline group. + + Args: + stage_index: Zero-based position inside the pipeline group. + stage_count: Number of pipeline stages. + group: NCCL process group containing the pipeline ranks. + ranks: Global ranks ordered from input to output stage. + + Raises: + RuntimeError: A rollout cache exists or CUDA graphs are enabled. + """ + if self._output_height is not None or self._output_width is not None: + raise RuntimeError( + "Pipeline parallelism must be configured before cache initialization." + ) + if self._use_cuda_graph: + raise RuntimeError( + "Pipeline-parallel DiT stages do not support CUDA graph capture." + ) + network = getattr(self.network, "_orig_mod", self.network) + assert isinstance(network, LingbotWorldDiTNetwork) + network.configure_pipeline_parallel( + stage_index=stage_index, + stage_count=stage_count, + group=group, + ranks=ranks, + ) + double_buffer_call = network.forward_pipeline_double_buffered + if self.config.compile_network: + double_buffer_call = torch.compile( + double_buffer_call, + mode="max-autotune-no-cudagraphs", + ) + self._pipeline_double_buffer_call = cast( + Callable[..., tuple[Tensor, Tensor]], + double_buffer_call, + ) @torch.no_grad() def replace_text_embeddings( @@ -114,6 +192,70 @@ def predict_flow( network_extra_kwargs={"plucker": input.plucker}, ) + def predict_flow_double_buffered( + self, + *, + noisy_latents: tuple[Tensor, Tensor], + timesteps: tuple[Tensor, Tensor], + caches: tuple[LingbotWorldTransformerCache, LingbotWorldTransformerCache], + inputs: tuple[I2VCamCtrlEmbeddings, I2VCamCtrlEmbeddings], + ) -> tuple[Tensor, Tensor]: + """Predict two session flows with a fill-and-drain pipeline schedule. + + Args: + noisy_latents: Per-session noisy latent tokens. + timesteps: Per-session diffusion timesteps. + caches: Per-session transformer caches. + inputs: Per-session patchified I2V and camera-control payloads. + + Returns: + Per-session flow predictions in input order. + + Raises: + RuntimeError: Pipeline parallelism is not configured or CFG is enabled. + """ + if self._pipeline_double_buffer_call is None: + raise RuntimeError("Double buffering requires pipeline parallelism.") + if any(cache.network_cache_uncond is not None for cache in caches): + raise RuntimeError( + "Double-buffered LingBot inference does not support CFG." + ) + + autoregressive_indices = tuple(cache.autoregressive_index for cache in caches) + if autoregressive_indices[0] != autoregressive_indices[1]: + raise RuntimeError( + "Double-buffered sessions must have the same autoregressive index." + ) + autoregressive_index = autoregressive_indices[0] + if autoregressive_index < 0: + raise RuntimeError("Call cache.start() before predicting flow.") + + network_inputs = tuple( + self._build_network_input(noisy_latent, input.i2v) + for noisy_latent, input in zip(noisy_latents, inputs) + ) + network_timesteps = tuple( + self._maybe_build_per_token_timestep( + timestep=timestep, + input=input.i2v, + autoregressive_index=autoregressive_index, + ) + for timestep, input in zip(timesteps, inputs) + ) + rope_freqs = tuple(cache.rope_freqs for cache in caches) + if any(value is None for value in rope_freqs): + raise RuntimeError("cache.start() must populate RoPE frequencies.") + + return self._pipeline_double_buffer_call( + pluckers=(inputs[0].plucker, inputs[1].plucker), + xs=(network_inputs[0], network_inputs[1]), + timesteps=(network_timesteps[0], network_timesteps[1]), + caches=(caches[0].network_cache, caches[1].network_cache), + rope_freqs=cast(tuple[Tensor, Tensor], rope_freqs), + current_chunk_idx=autoregressive_index, + eager_mode=False, + ) + @overload def patchify_and_maybe_split_cp(self, x: Tensor) -> Tensor: ... @overload diff --git a/integrations/lingbot/lingbot/transformer/impl/network.py b/integrations/lingbot/lingbot/transformer/impl/network.py index ea67774f..35e19db8 100644 --- a/integrations/lingbot/lingbot/transformer/impl/network.py +++ b/integrations/lingbot/lingbot/transformer/impl/network.py @@ -18,17 +18,24 @@ from __future__ import annotations from dataclasses import dataclass, field +from math import prod from typing import Literal +import torch import torch.nn as nn import torch.nn.functional as F from torch import Tensor +from torch.distributed import ProcessGroup from flashdreams.recipes.wan.transformer.impl.network import ( WanDiTNetwork, WanDiTNetworkCache, WanDiTNetworkConfig, ) +from flashdreams.recipes.wan.transformer.impl.modules import ( + Block, + sinusoidal_embedding_1d, +) from .modules import CamCtrlBlock @@ -68,12 +75,54 @@ class LingbotWorldDiTNetwork14BConfig(LingbotWorldDiTNetworkConfig): num_layers: int = 40 +def pipeline_partition_bounds( + num_layers: int, + *, + stage_index: int, + stage_count: int, +) -> tuple[int, int]: + """Return the balanced half-open layer range for one pipeline stage. + + Args: + num_layers: Total transformer-block count. + stage_index: Zero-based pipeline-stage index. + stage_count: Number of pipeline stages. + + Returns: + Global ``(start, end)`` layer indices assigned to the stage. + + Raises: + ValueError: The stage layout is empty or outside the valid range. + """ + if num_layers < 1: + raise ValueError(f"num_layers must be positive, got {num_layers}.") + if stage_count < 1 or stage_count > num_layers: + raise ValueError( + f"stage_count must be in [1, {num_layers}], got {stage_count}." + ) + if stage_index < 0 or stage_index >= stage_count: + raise ValueError( + f"stage_index must be in [0, {stage_count}), got {stage_index}." + ) + + base, remainder = divmod(num_layers, stage_count) + start = stage_index * base + min(stage_index, remainder) + end = start + base + (1 if stage_index < remainder else 0) + return start, end + + class LingbotWorldDiTNetwork(WanDiTNetwork): """Lingbot World DiT diffusion backbone for text-to-video and image-to-video.""" def __init__(self, config: LingbotWorldDiTNetworkConfig) -> None: super().__init__(config) + self._pipeline_stage_index: int | None = None + self._pipeline_stage_count = 1 + self._pipeline_group: ProcessGroup | None = None + self._pipeline_ranks: tuple[int, ...] = () + self._global_layer_range = (0, config.num_layers) + if config.control_type == "cam": control_dim = 6 elif config.control_type == "act": @@ -91,6 +140,61 @@ def __init__(self, config: LingbotWorldDiTNetworkConfig) -> None: self.c2ws_hidden_states_layer1 = nn.Linear(self.dim, self.dim) self.c2ws_hidden_states_layer2 = nn.Linear(self.dim, self.dim) + def configure_pipeline_parallel( + self, + *, + stage_index: int, + stage_count: int, + group: ProcessGroup, + ranks: tuple[int, ...], + ) -> None: + """Partition resident layers and bind the NCCL pipeline group. + + The checkpoint is loaded on CPU before this method runs. Unowned blocks + and endpoint-only modules are removed before the stage is moved to its + GPU, so peak device memory reflects the local partition. + + Args: + stage_index: Zero-based position in the pipeline group. + stage_count: Number of ranks in the pipeline group. + group: NCCL process group spanning the pipeline ranks. + ranks: Global ranks in pipeline order. + + Raises: + ValueError: The rank layout does not describe a two-stage group. + RuntimeError: Pipeline parallelism was already configured. + """ + if self._pipeline_stage_index is not None: + raise RuntimeError("Pipeline parallelism is already configured.") + if stage_count != 2 or len(ranks) != stage_count: + raise ValueError( + "LingBot pipeline parallelism currently requires two ranks." + ) + if torch.distributed.get_rank() != ranks[stage_index]: + raise ValueError( + f"Global rank {torch.distributed.get_rank()} does not match " + f"stage_index={stage_index} in ranks={ranks}." + ) + + global_num_layers = len(self.blocks) + start, end = pipeline_partition_bounds( + global_num_layers, + stage_index=stage_index, + stage_count=stage_count, + ) + self.blocks = nn.ModuleList(list(self.blocks[start:end])) + self.num_layers = len(self.blocks) + self._global_layer_range = (start, end) + self._pipeline_stage_index = stage_index + self._pipeline_stage_count = stage_count + self._pipeline_group = group + self._pipeline_ranks = ranks + + if stage_index == 0: + del self.head + else: + del self.patch_embedding + def _build_block(self, layer_idx: int) -> CamCtrlBlock: return CamCtrlBlock( dim=self.dim, @@ -117,6 +221,303 @@ def replace_text_embeddings( assert isinstance(block, CamCtrlBlock) block_cache.cross_attn.text = block.cross_attn.compute_kv(context_text) + def _pipeline_embeddings( + self, + *, + x: Tensor, + plucker: Tensor, + timesteps: Tensor, + ) -> tuple[Tensor, Tensor, Tensor]: + """Build local camera and timestep embeddings for a pipeline stage.""" + batch_shape = x.shape[:-2] + token_count = x.shape[-2] + per_token_timestep = ( + timesteps.ndim > len(batch_shape) and timesteps.shape[-1] == token_count + ) + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timesteps).type_as(x) + ) + e0 = self.time_projection(e).unflatten(-1, (6, self.dim)) + if per_token_timestep: + block_e_shape = batch_shape + (token_count, 6, self.dim) + head_e = torch.broadcast_to( + e, + batch_shape + (token_count, self.dim), + ).unsqueeze(-2) + else: + block_e_shape = batch_shape + (6, self.dim) + head_e = ( + torch.broadcast_to( + e, + batch_shape + (self.dim,), + ) + .unsqueeze(-2) + .unsqueeze(-2) + ) + block_e = torch.broadcast_to(e0, block_e_shape) + + plucker_embedding = self.patch_embedding_wancamctrl(plucker) + plucker_hidden_states = self.c2ws_hidden_states_layer2( + F.silu(self.c2ws_hidden_states_layer1(plucker_embedding)) + ) + return plucker_embedding + plucker_hidden_states, block_e, head_e + + def _pipeline_output_buffer(self, x: Tensor) -> Tensor: + return torch.empty( + *x.shape[:-1], + self.out_dim * prod(self.patch_size), + dtype=x.dtype, + device=x.device, + ) + + def _pipeline_local_forward( + self, + *, + plucker: Tensor, + x: Tensor, + timesteps: Tensor, + cache: LingbotWorldDiTNetworkCache, + rope_freqs: Tensor, + current_chunk_idx: int, + eager_mode: bool, + hidden: Tensor | None, + ) -> tuple[Tensor, Tensor]: + plucker_embedding, block_e, head_e = self._pipeline_embeddings( + x=x, + plucker=plucker, + timesteps=timesteps, + ) + if self._pipeline_stage_index == 0: + assert hidden is None + if self.patch_embedding_type == "linear": + hidden = self.patch_embedding(x) + elif self.patch_embedding_type == "conv3d": + weight = self.patch_embedding.weight.reshape(self.dim, -1) + hidden = F.linear(x, weight, self.patch_embedding.bias) + else: + raise ValueError( + f"Invalid patch embedding type: {self.patch_embedding_type}" + ) + else: + assert hidden is not None + + if eager_mode: + cache.before_update(current_chunk_idx) + for block_idx, block in enumerate(self.blocks): + assert isinstance(block, Block) + hidden = block( + x=hidden, + e=block_e, + rope_freqs=rope_freqs, + cache=cache[block_idx], + plucker_embedding=plucker_embedding, + ) + if eager_mode: + cache.after_update(current_chunk_idx) + return hidden, head_e + + def _forward_pipeline( + self, + *, + plucker: Tensor, + x: Tensor, + timesteps: Tensor, + cache: LingbotWorldDiTNetworkCache, + rope_freqs: Tensor, + current_chunk_idx: int, + eager_mode: bool, + ) -> Tensor: + """Run one local layer partition and exchange boundary activations.""" + assert self._pipeline_stage_index is not None + assert self._pipeline_group is not None + peer_rank = self._pipeline_ranks[1 - self._pipeline_stage_index] + + if self._pipeline_stage_index == 0: + hidden = None + else: + hidden = torch.empty( + *x.shape[:-1], + self.dim, + dtype=x.dtype, + device=x.device, + ) + torch.distributed.recv( + hidden, + src=peer_rank, + group=self._pipeline_group, + ) + + hidden, head_e = self._pipeline_local_forward( + plucker=plucker, + x=x, + timesteps=timesteps, + cache=cache, + rope_freqs=rope_freqs, + current_chunk_idx=current_chunk_idx, + eager_mode=eager_mode, + hidden=hidden, + ) + + if self._pipeline_stage_index == 0: + torch.distributed.send( + hidden.contiguous(), + dst=peer_rank, + group=self._pipeline_group, + ) + output = self._pipeline_output_buffer(x) + torch.distributed.recv( + output, + src=peer_rank, + group=self._pipeline_group, + ) + return output + + output = self.head(hidden, head_e) + torch.distributed.send( + output.contiguous(), + dst=peer_rank, + group=self._pipeline_group, + ) + return output + + def forward_pipeline_double_buffered( + self, + *, + pluckers: tuple[Tensor, Tensor], + xs: tuple[Tensor, Tensor], + timesteps: tuple[Tensor, Tensor], + caches: tuple[ + LingbotWorldDiTNetworkCache, + LingbotWorldDiTNetworkCache, + ], + rope_freqs: tuple[Tensor, Tensor], + current_chunk_idx: int, + eager_mode: bool = False, + ) -> tuple[Tensor, Tensor]: + """Pipeline two session microbatches through both layer partitions. + + Stage 0 computes the second microbatch while stage 1 processes the + first. A bidirectional NCCL exchange then hands off the second hidden + state and returns the first output before the pipeline drains. + + Args: + pluckers: Per-session camera-control token tensors. + xs: Per-session DiT input token tensors. + timesteps: Per-session diffusion timesteps. + caches: Per-session local network caches. + rope_freqs: Per-session RoPE frequency tensors. + current_chunk_idx: Current autoregressive chunk index. + eager_mode: Whether to run cache update hooks inside the network. + + Returns: + Two output tensors in the same order as ``xs``. + + Raises: + RuntimeError: Pipeline parallelism is not configured. + """ + if self._pipeline_stage_index is None or self._pipeline_group is None: + raise RuntimeError("Double buffering requires pipeline parallelism.") + peer_rank = self._pipeline_ranks[1 - self._pipeline_stage_index] + + def local_forward(index: int, hidden: Tensor | None) -> tuple[Tensor, Tensor]: + return self._pipeline_local_forward( + plucker=pluckers[index], + x=xs[index], + timesteps=timesteps[index], + cache=caches[index], + rope_freqs=rope_freqs[index], + current_chunk_idx=current_chunk_idx, + eager_mode=eager_mode, + hidden=hidden, + ) + + if self._pipeline_stage_index == 0: + hidden_0, _ = local_forward(0, None) + torch.distributed.send( + hidden_0.contiguous(), + dst=peer_rank, + group=self._pipeline_group, + ) + del hidden_0 + + hidden_1, _ = local_forward(1, None) + hidden_1 = hidden_1.contiguous() + output_0 = self._pipeline_output_buffer(xs[0]) + requests = torch.distributed.batch_isend_irecv( + [ + torch.distributed.P2POp( + torch.distributed.isend, + hidden_1, + peer_rank, + self._pipeline_group, + ), + torch.distributed.P2POp( + torch.distributed.irecv, + output_0, + peer_rank, + self._pipeline_group, + ), + ] + ) + for request in requests: + request.wait() + + output_1 = self._pipeline_output_buffer(xs[1]) + torch.distributed.recv( + output_1, + src=peer_rank, + group=self._pipeline_group, + ) + return output_0, output_1 + + hidden_0 = torch.empty( + *xs[0].shape[:-1], + self.dim, + dtype=xs[0].dtype, + device=xs[0].device, + ) + torch.distributed.recv( + hidden_0, + src=peer_rank, + group=self._pipeline_group, + ) + hidden_0, head_e_0 = local_forward(0, hidden_0) + output_0 = self.head(hidden_0, head_e_0).contiguous() + + hidden_1 = torch.empty( + *xs[1].shape[:-1], + self.dim, + dtype=xs[1].dtype, + device=xs[1].device, + ) + requests = torch.distributed.batch_isend_irecv( + [ + torch.distributed.P2POp( + torch.distributed.irecv, + hidden_1, + peer_rank, + self._pipeline_group, + ), + torch.distributed.P2POp( + torch.distributed.isend, + output_0, + peer_rank, + self._pipeline_group, + ), + ] + ) + for request in requests: + request.wait() + + hidden_1, head_e_1 = local_forward(1, hidden_1) + output_1 = self.head(hidden_1, head_e_1) + torch.distributed.send( + output_1.contiguous(), + dst=peer_rank, + group=self._pipeline_group, + ) + return output_0, output_1 + def forward( self, plucker: Tensor, @@ -148,6 +549,17 @@ def forward( "We expect to have called update_parameters_after_loading_checkpoint() after loading the checkpoint" ) + if self._pipeline_stage_index is not None: + return self._forward_pipeline( + plucker=plucker, + x=x, + timesteps=timesteps, + cache=cache, + rope_freqs=rope_freqs, + current_chunk_idx=current_chunk_idx, + eager_mode=eager_mode, + ) + plucker_embedding = self.patch_embedding_wancamctrl(plucker) plucker_hidden_states = self.c2ws_hidden_states_layer2( F.silu(self.c2ws_hidden_states_layer1(plucker_embedding)) diff --git a/integrations/lingbot/pyproject.toml b/integrations/lingbot/pyproject.toml index f9e39cd9..e5c726ca 100644 --- a/integrations/lingbot/pyproject.toml +++ b/integrations/lingbot/pyproject.toml @@ -40,6 +40,9 @@ dev = [ "pytest>=8.0", "pytest-asyncio>=0.23", ] +disagg = [ + "mooncake-transfer-engine-cuda13>=0.3.12", +] # Each entry registers one ``runner_name`` slug with ``flashdreams-run``. # The discovery layer (``flashdreams.plugins.registry.discover_runners``) diff --git a/integrations/lingbot/scripts/plot_aggregated_comparison.py b/integrations/lingbot/scripts/plot_aggregated_comparison.py new file mode 100644 index 00000000..3e2ddb73 --- /dev/null +++ b/integrations/lingbot/scripts/plot_aggregated_comparison.py @@ -0,0 +1,250 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Render disaggregated CP6 versus aggregated CP8 wall time and HBM as SVG.""" + +from __future__ import annotations + +import argparse +import html +import json +from pathlib import Path +from typing import Any, TypedDict + +_COLORS = { + "Encoder": "#59a14f", + "Input handoff": "#8cd17d", + "DiT denoise": "#4e79a7", + "KV finalize": "#b07aa1", + "Output handoff": "#76b7b2", + "Decoder": "#f28e2b", + "Coordination": "#bab0ac", +} + + +class _WallRow(TypedDict): + """One stacked wall-time row.""" + + label: str + total: float + parts: dict[str, float] + + +def _args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("disaggregated", type=Path) + parser.add_argument("aggregated", type=Path) + parser.add_argument("output", type=Path) + return parser.parse_args() + + +def _read(path: Path) -> dict[str, Any]: + return json.loads(path.read_text()) + + +def _wall_rows( + disaggregated: dict[str, Any], + aggregated: dict[str, Any], +) -> list[_WallRow]: + disagg = disaggregated["summary"] + agg = aggregated["summary"] + rows: list[_WallRow] = [ + { + "label": "Disaggregated CP6 ring", + "total": disagg["latency_ms"]["median"], + "parts": { + "Encoder": disagg["encoder_ms"]["median"], + "Input handoff": ( + disagg["encoder_to_cp_leader"]["handoff_ms"]["median"] + + disagg["cp_input_fanout_ms"]["median"] + ), + "DiT denoise": disagg["dit_ms"]["median"], + "KV finalize": disagg["finalize_ms"]["median"], + "Output handoff": disagg["cp_leader_to_decoder"]["handoff_ms"][ + "median" + ], + "Decoder": disagg["decoder_ms"]["median"], + }, + }, + { + "label": "Aggregated CP8 Ulysses", + "total": agg["latency_ms"]["median"], + "parts": { + "Encoder": agg["encoder_ms"]["median"], + "DiT denoise": agg["dit_ms"]["median"], + "KV finalize": agg["finalize_ms"]["median"], + "Decoder": agg["decoder_ms"]["median"], + }, + }, + ] + for row in rows: + row["parts"]["Coordination"] = max( + 0.0, + row["total"] - sum(row["parts"].values()), + ) + return rows + + +def _svg( + disaggregated: dict[str, Any], + aggregated: dict[str, Any], +) -> str: + width, height = 1200, 720 + rows = _wall_rows(disaggregated, aggregated) + max_wall = 800.0 + chart_x, chart_w = 245.0, 860.0 + wall_y = [145.0, 225.0] + bar_h = 48.0 + pieces = [ + ( + f'' + ), + ( + 'LingBot disaggregated CP6 and aggregated CP8 ' + "performance comparison" + ), + ( + 'Median per-chunk component wall time and ' + "per-rank peak allocated HBM on eight H100 GPUs." + ), + "", + '', + 'LingBot: stage-local CP6 vs full-pipeline CP8', + ( + '8× H100 80 GB · BF16 · ' + "six warmup + five measured blocks · CP6 832×464 · CP8 832×448" + ), + 'Median steady-state wall time', + ] + for tick in range(0, 801, 100): + x = chart_x + chart_w * tick / max_wall + pieces.extend( + [ + f'', + ( + f'{tick}' + ), + ] + ) + pieces.append( + 'milliseconds' + ) + for row, y in zip(rows, wall_y): + pieces.append( + f'' + f"{html.escape(row['label'])}" + ) + x = chart_x + for label, value in row["parts"].items(): + segment_w = chart_w * value / max_wall + pieces.append( + f'' + ) + if segment_w >= 45: + pieces.append( + f'{value:.0f}' + ) + x += segment_w + pieces.append( + f'' + f"{row['total']:.0f} ms" + ) + + for index, label in enumerate(_COLORS): + x = 75 + index * 154 + pieces.extend( + [ + f'', + f'{html.escape(label)}', + ] + ) + + pieces.extend( + [ + 'Peak allocated HBM by rank', + ( + 'Stage-local CP6 totals ' + "251.07 GiB; eight full-pipeline replicas total 327.03 GiB (+30.3%)" + ), + ] + ) + mem_base_y, mem_top_y = 626.0, 448.0 + mem_h = mem_base_y - mem_top_y + for tick in (0, 10, 20, 30, 40, 50): + y = mem_base_y - mem_h * tick / 50.0 + pieces.extend( + [ + f'', + f'{tick}', + ] + ) + pieces.append( + 'GiB' + ) + + disagg_memory = disaggregated["environment"]["peak_memory_gib_by_rank"] + agg_memory = aggregated["summary"]["memory"]["peak_gib_by_rank"] + memory_groups = ( + ( + "Disaggregated CP6", + disagg_memory, + ["Encoder", *(["DiT denoise"] * 6), "Decoder"], + 105.0, + ), + ( + "Aggregated CP8", + agg_memory, + ["DiT denoise"] * 8, + 650.0, + ), + ) + for group, values, stages, start_x in memory_groups: + for rank, (value, stage) in enumerate(zip(values, stages)): + x = start_x + rank * 53.0 + bar_height = mem_h * value / 50.0 + y = mem_base_y - bar_height + pieces.extend( + [ + f'', + ( + f'{value:.1f}' + ), + ( + f'G{rank}' + ), + ] + ) + center = start_x + (len(values) - 1) * 53.0 / 2 + 17.0 + pieces.append( + f'' + f"{html.escape(group)}" + ) + pieces.append( + 'Aggregated bars contain encoder + DiT + decoder on every GPU.' + ) + pieces.append("") + return "\n".join(pieces) + "\n" + + +def main() -> None: + """Write the SVG comparison.""" + args = _args() + args.output.write_text(_svg(_read(args.disaggregated), _read(args.aggregated))) + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/scripts/plot_disagg_breakdown.py b/integrations/lingbot/scripts/plot_disagg_breakdown.py new file mode 100644 index 00000000..2519e88f --- /dev/null +++ b/integrations/lingbot/scripts/plot_disagg_breakdown.py @@ -0,0 +1,248 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Render the tracked three-stage wall-time and GPU-memory comparison as SVG.""" + +from __future__ import annotations + +import argparse +import html +import json +from pathlib import Path +from typing import Any + +_COLORS = { + "Encoder": "#59a14f", + "Encoder → DiT": "#8cd17d", + "DiT": "#4e79a7", + "DiT → decoder": "#76b7b2", + "Decoder": "#f28e2b", + "Coordination": "#bab0ac", +} + + +def _args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("baseline", type=Path) + parser.add_argument("scaled", type=Path) + parser.add_argument("output", type=Path) + return parser.parse_args() + + +def _read(path: Path) -> dict[str, Any]: + return json.loads(path.read_text()) + + +def _walltime(baseline: dict[str, Any], scaled: dict[str, Any]) -> list[dict[str, Any]]: + base = baseline["summary"] + base_parts = { + "Encoder": base["encoder_ms"]["median"], + "Encoder → DiT": base["encoder_to_dit"]["handoff_ms"]["median"], + "DiT": base["dit_ms"]["median"] + base["finalize_ms"]["median"], + "DiT → decoder": base["dit_to_decoder"]["handoff_ms"]["median"], + "Decoder": base["decoder_ms"]["median"], + } + scaled_summary = scaled["summary"] + scaled_parts = { + "Encoder": scaled_summary["encoder_wave_ms"]["median"], + "Encoder → DiT": scaled_summary["encoder_to_dit"][ + "aggregate_handoff_ms_per_wave" + ]["median"], + "DiT": scaled_summary["dit_critical_path_ms"]["median"], + "DiT → decoder": scaled_summary["dit_to_decoder"][ + "aggregate_handoff_ms_per_wave" + ]["median"], + "Decoder": scaled_summary["decoder_wave_ms"]["median"], + } + rows: list[dict[str, Any]] = [ + { + "label": "1E : 1D : 1V", + "total": base["latency_ms"]["median"], + "parts": base_parts, + }, + { + "label": "1E : 6D : 1V", + "total": scaled_summary["wave_latency_ms"]["median"], + "parts": scaled_parts, + }, + ] + for row in rows: + row["parts"]["Coordination"] = max( + 0.0, row["total"] - sum(row["parts"].values()) + ) + return rows + + +def _memory( + baseline: dict[str, Any], scaled: dict[str, Any] +) -> list[tuple[str, float, str]]: + base = baseline["environment"]["peak_memory_gib_by_stage"] + scaled_memory = scaled["environment"]["peak_memory_gib_by_rank"] + return [ + ("E0", base["encoder"], "Encoder"), + ("D1", base["dit"], "DiT"), + ("V2", base["decoder"], "Decoder"), + ("E0", scaled_memory[0], "Encoder"), + *[ + (f"D{rank}", scaled_memory[rank], "DiT") + for rank in range(1, len(scaled_memory) - 1) + ], + ("V7", scaled_memory[-1], "Decoder"), + ] + + +def _svg(baseline: dict[str, Any], scaled: dict[str, Any]) -> str: + width, height = 1200, 720 + wall_rows = _walltime(baseline, scaled) + max_wall = 3000.0 + chart_x, chart_w = 165.0, 965.0 + wall_y = [142.0, 232.0] + bar_h = 52.0 + pieces = [ + ( + f'' + ), + 'LingBot disaggregated inference wall time and GPU memory', + ( + 'Stacked wall-time comparison for one versus six ' + "DiT workers and per-rank peak allocated GPU memory." + ), + "", + '', + 'LingBot stage allocation: wall time and memory', + ( + 'H100 80 GB · BF16 · 832×464 · ' + "six warmup and five measured waves" + ), + 'Median steady-state wall time', + ] + for tick in range(0, 3001, 500): + x = chart_x + chart_w * tick / max_wall + pieces.extend( + [ + f'', + f'{tick}', + ] + ) + pieces.append( + 'milliseconds' + ) + for row, y in zip(wall_rows, wall_y): + pieces.append( + f'' + f"{html.escape(row['label'])}" + ) + x = chart_x + for label, value in row["parts"].items(): + segment_w = chart_w * value / max_wall + pieces.append( + f'' + ) + if segment_w >= 48: + pieces.append( + f'{value:.0f}' + ) + x += segment_w + pieces.append( + f'' + f"{row['total']:.0f} ms" + ) + + legend_x = 165 + for index, label in enumerate(_COLORS): + x = legend_x + index * 157 + pieces.extend( + [ + f'', + f'{html.escape(label)}', + ] + ) + + pieces.extend( + [ + 'Peak allocated GPU memory by rank', + ( + 'E = encoder, D = DiT, V = decoder; ' + "the gap to 80 GiB is headroom, not free schedulable memory" + ), + ] + ) + memory = _memory(baseline, scaled) + groups = [(memory[:3], 155.0, "1E : 1D : 1V"), (memory[3:], 600.0, "1E : 6D : 1V")] + mem_base_y = 650.0 + mem_top_y = 455.0 + mem_h = mem_base_y - mem_top_y + for tick in (0, 20, 40, 60, 80): + y = mem_base_y - mem_h * tick / 80.0 + pieces.extend( + [ + f'', + f'{tick}', + ] + ) + pieces.append( + 'GiB' + ) + for rows, start_x, group_label in groups: + gap = 68.0 if len(rows) > 3 else 96.0 + bar_width = 43.0 + center = start_x + ((len(rows) - 1) * gap + bar_width) / 2 + pieces.append( + f'{group_label}' + ) + for index, (rank, value, stage) in enumerate(rows): + x = start_x + index * gap + bar_height = mem_h * value / 80.0 + y = mem_base_y - bar_height + pieces.extend( + [ + ( + f'' + ), + ( + f'{value:.1f}' + ), + ( + f'{rank}' + ), + ] + ) + pieces.append("") + return "\n".join(pieces) + "\n" + + +def main() -> None: + """Read benchmark JSON documents and write an SVG comparison.""" + args = _args() + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(_svg(_read(args.baseline), _read(args.scaled))) + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/scripts/plot_disagg_single_session.py b/integrations/lingbot/scripts/plot_disagg_single_session.py new file mode 100644 index 00000000..23209996 --- /dev/null +++ b/integrations/lingbot/scripts/plot_disagg_single_session.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Render LingBot CP1, CP4, and CP6 single-session results as SVG.""" + +from __future__ import annotations + +import argparse +import html +import json +from pathlib import Path +from typing import Any + +_COLORS = { + "Encoder": "#59a14f", + "Input handoff": "#8cd17d", + "DiT": "#4e79a7", + "Output handoff": "#76b7b2", + "Decoder": "#f28e2b", + "Coordination": "#bab0ac", +} + + +def _args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("baseline", type=Path) + parser.add_argument("cp4", type=Path) + parser.add_argument("cp6", type=Path) + parser.add_argument("output", type=Path) + return parser.parse_args() + + +def _read(path: Path) -> dict[str, Any]: + return json.loads(path.read_text()) + + +def _walltime( + baseline: dict[str, Any], + cp4: dict[str, Any], + cp6: dict[str, Any], +) -> list[dict[str, Any]]: + base = baseline["summary"] + rows: list[dict[str, Any]] = [ + { + "label": "CP1", + "total": base["latency_ms"]["median"], + "parts": { + "Encoder": base["encoder_ms"]["median"], + "Input handoff": base["encoder_to_dit"]["handoff_ms"]["median"], + "DiT": (base["dit_ms"]["median"] + base["finalize_ms"]["median"]), + "Output handoff": base["dit_to_decoder"]["handoff_ms"]["median"], + "Decoder": base["decoder_ms"]["median"], + }, + } + ] + for label, document in (("CP4 Ulysses", cp4), ("CP6 ring", cp6)): + summary = document["summary"] + rows.append( + { + "label": label, + "total": summary["latency_ms"]["median"], + "parts": { + "Encoder": summary["encoder_ms"]["median"], + "Input handoff": ( + summary["encoder_to_cp_leader"]["handoff_ms"]["median"] + + summary["cp_input_fanout_ms"]["median"] + ), + "DiT": summary["dit_critical_path_ms"]["median"], + "Output handoff": summary["cp_leader_to_decoder"]["handoff_ms"][ + "median" + ], + "Decoder": summary["decoder_ms"]["median"], + }, + } + ) + for row in rows: + row["parts"]["Coordination"] = max( + 0.0, + row["total"] - sum(row["parts"].values()), + ) + return rows + + +def _memory_rows( + baseline: dict[str, Any], + cp4: dict[str, Any], + cp6: dict[str, Any], +) -> list[tuple[str, list[tuple[str, float, str]]]]: + base = baseline["environment"]["peak_memory_gib_by_stage"] + + def cp_rows(document: dict[str, Any]) -> list[tuple[str, float, str]]: + memory = document["environment"]["peak_memory_gib_by_rank"] + return [ + ("E0", memory[0], "Encoder"), + *[(f"D{rank}", memory[rank], "DiT") for rank in range(1, len(memory) - 1)], + (f"V{len(memory) - 1}", memory[-1], "Decoder"), + ] + + return [ + ( + "CP1", + [ + ("E0", base["encoder"], "Encoder"), + ("D1", base["dit"], "DiT"), + ("V2", base["decoder"], "Decoder"), + ], + ), + ("CP4 Ulysses", cp_rows(cp4)), + ("CP6 ring", cp_rows(cp6)), + ] + + +def _svg( + baseline: dict[str, Any], + cp4: dict[str, Any], + cp6: dict[str, Any], +) -> str: + width, height = 1200, 780 + rows = _walltime(baseline, cp4, cp6) + max_wall = 2500.0 + chart_x, chart_w = 170.0, 940.0 + wall_y = [135.0, 207.0, 279.0] + bar_h = 43.0 + pieces = [ + ( + f'' + ), + ( + 'LingBot single-session context-parallel wall time ' + "and GPU memory" + ), + ( + 'Median wall-time components and per-rank peak ' + "allocated memory for CP1, CP4 Ulysses, and CP6 ring." + ), + "", + '', + ( + 'LingBot minimum single-session ' + "latency" + ), + ( + 'H100 80 GB · BF16 · 832×464 · ' + "six warmup and five measured blocks" + ), + 'Median steady-state wall time', + ] + for tick in range(0, 2501, 500): + x = chart_x + chart_w * tick / max_wall + pieces.extend( + [ + f'', + ( + f'{tick}' + ), + ] + ) + pieces.append( + 'milliseconds' + ) + for row, y in zip(rows, wall_y): + pieces.append( + f'' + f"{html.escape(row['label'])}" + ) + x = chart_x + for label, value in row["parts"].items(): + segment_w = chart_w * value / max_wall + pieces.append( + f'' + ) + if segment_w >= 48: + pieces.append( + f'{value:.0f}' + ) + x += segment_w + pieces.append( + f'' + f"{row['total']:.0f} ms" + ) + + for index, label in enumerate(_COLORS): + x = 170 + index * 156 + pieces.extend( + [ + f'', + f'{html.escape(label)}', + ] + ) + + pieces.extend( + [ + 'Peak allocated GPU memory by rank', + ( + 'E = encoder, D = DiT, ' + "V = decoder; CP4 leaves two GPUs available for other work" + ), + ] + ) + mem_base_y, mem_top_y = 701.0, 478.0 + mem_h = mem_base_y - mem_top_y + for tick in (0, 20, 40, 60, 80): + y = mem_base_y - mem_h * tick / 80.0 + pieces.extend( + [ + f'', + f'{tick}', + ] + ) + pieces.append( + 'GiB' + ) + layouts = [(125.0, 58.0), (395.0, 47.0), (765.0, 43.0)] + for (group_label, memory), (start_x, gap) in zip( + _memory_rows(baseline, cp4, cp6), + layouts, + ): + bar_width = 31.0 + center = start_x + ((len(memory) - 1) * gap + bar_width) / 2 + pieces.append( + f'{html.escape(group_label)}' + ) + for index, (rank, value, stage) in enumerate(memory): + x = start_x + index * gap + bar_height = mem_h * value / 80.0 + y = mem_base_y - bar_height + pieces.extend( + [ + ( + f'' + ), + ( + f'{value:.1f}' + ), + ( + f'{rank}' + ), + ] + ) + pieces.append("") + return "\n".join(pieces) + "\n" + + +def main() -> None: + """Read benchmark documents and write the single-session SVG.""" + args = _args() + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + _svg( + _read(args.baseline), + _read(args.cp4), + _read(args.cp6), + ) + ) + + +if __name__ == "__main__": + main() diff --git a/integrations/lingbot/tests/test_disagg_aggregated.py b/integrations/lingbot/tests/test_disagg_aggregated.py new file mode 100644 index 00000000..30089c0c --- /dev/null +++ b/integrations/lingbot/tests/test_disagg_aggregated.py @@ -0,0 +1,106 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the aggregated LingBot benchmark report.""" + +from __future__ import annotations + +import pytest +from lingbot.disagg.benchmark_aggregated import _summarize, _token_layout + +pytestmark = pytest.mark.ci_cpu + + +def test_cp1_layout_accepts_tracked_464_height() -> None: + layout = _token_layout( + pixel_height=464, + pixel_width=832, + len_t=3, + patch_size=(1, 2, 2), + cp_size=1, + ) + + assert layout == { + "latent_height": 58, + "latent_width": 104, + "tokens_per_chunk": 4524, + "tokens_per_rank": 4524, + } + + +def test_cp8_layout_uses_nearest_valid_height() -> None: + layout = _token_layout( + pixel_height=448, + pixel_width=832, + len_t=3, + patch_size=(1, 2, 2), + cp_size=8, + ) + + assert layout == { + "latent_height": 56, + "latent_width": 104, + "tokens_per_chunk": 4368, + "tokens_per_rank": 546, + } + + +def test_cp8_layout_rejects_tracked_464_height() -> None: + with pytest.raises(ValueError, match="4524 tokens"): + _token_layout( + pixel_height=464, + pixel_width=832, + len_t=3, + patch_size=(1, 2, 2), + cp_size=8, + ) + + +def test_aggregated_summary_uses_critical_rank_and_node_memory() -> None: + records = [ + { + "warmup": False, + "output_frames": 12, + "end_to_end_ms": 600.0, + "critical_rank": { + "encode_ms": 2.0, + "diffuse_ms": 500.0, + "decode_ms": 8.0, + "finalize_ms": 60.0, + }, + } + ] + sample = { + "payload_bytes": float(256 * 2**20), + "transfer_ms": 1.0, + "bandwidth_gbps": 100.0, + } + + summary = _summarize( + records=records, + tokens_per_chunk=4368, + cp_probe={"broadcast": [sample], "all_gather": [sample]}, + peak_memory_gib_by_rank=[50.0] * 8, + steady_memory_gib_by_rank=[49.0] * 8, + initialization_peak_gib_by_rank=[55.0] * 8, + comparison={ + "summary": { + "fps": 10.0, + "latency_ms": {"median": 1200.0}, + }, + "environment": { + "resolution": [464, 832], + "latent_frames_per_chunk": 3, + "peak_memory_gib_by_rank": [25.0] * 8, + "allocation": {"cp_size": 6}, + }, + }, + ) + + assert summary["fps"] == pytest.approx(20.0) + assert summary["token_throughput_per_second"] == pytest.approx(7280.0) + assert summary["dit_ms"]["median"] == 500.0 + assert summary["memory"]["node_peak_gib"] == 400.0 + assert summary["memory"]["node_steady_allocated_gib"] == 392.0 + assert summary["comparison"]["topology"] == "1 encoder : CP6 DiT : 1 decoder" + assert summary["comparison"]["latency_speedup"] == pytest.approx(2.0) diff --git a/integrations/lingbot/tests/test_disagg_cp.py b/integrations/lingbot/tests/test_disagg_cp.py new file mode 100644 index 00000000..e86ac900 --- /dev/null +++ b/integrations/lingbot/tests/test_disagg_cp.py @@ -0,0 +1,151 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for DiT-only context-parallel disaggregation.""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any, cast + +import pytest +from lingbot.disagg.benchmark_cp import _summarize +from lingbot.transformer import LingbotWorldTransformer +from torch.distributed import ProcessGroup + +from flashdreams.infra.transfer import TransferStats + +pytestmark = pytest.mark.ci_cpu + + +class _FakeGroup: + def __init__(self, size: int) -> None: + self._size = size + + def size(self) -> int: + return self._size + + +class _FakeNetwork: + def __init__(self) -> None: + self.group: ProcessGroup | None = None + + def set_context_parallel_group(self, group: ProcessGroup | None) -> None: + self.group = group + + +class _FakeDispatch: + def __init__(self) -> None: + self.reset_count = 0 + + def reset(self) -> None: + self.reset_count += 1 + + +def test_lingbot_transformer_can_bind_dit_only_cp_group_before_cache() -> None: + group = cast(ProcessGroup, _FakeGroup(6)) + network = _FakeNetwork() + dispatch = _FakeDispatch() + transformer = SimpleNamespace( + _output_height=None, + _output_width=None, + _cp_group=None, + _cp_size=1, + network=network, + _use_cuda_graph=True, + _cuda_graph_dispatch=dispatch, + ) + + LingbotWorldTransformer.set_context_parallel_group( + cast(LingbotWorldTransformer, transformer), + group, + ) + + assert transformer._cp_group is group + assert transformer._cp_size == 6 + assert network.group is group + assert dispatch.reset_count == 1 + + +def test_lingbot_transformer_rejects_cp_rebind_after_cache_init() -> None: + transformer = SimpleNamespace(_output_height=58, _output_width=104) + with pytest.raises(RuntimeError, match="before initializing"): + LingbotWorldTransformer.set_context_parallel_group( + cast(LingbotWorldTransformer, transformer), + cast(ProcessGroup, _FakeGroup(6)), + ) + + +def _baseline() -> dict[str, Any]: + return { + "summary": { + "fps": 6.0, + "latency_ms": {"median": 2000.0}, + "dit_ms": {"median": 1600.0}, + "finalize_ms": {"median": 400.0}, + } + } + + +def test_cp_summary_reports_single_session_latency_scaling() -> None: + transfer = { + "payload_bytes": 1024, + "transfer_ms": 1.0, + "handoff_ms": 2.0, + } + records = [ + { + "warmup": False, + "output_frames": 12, + "end_to_end_ms": 500.0, + "encoder_ms": 1.0, + "encoder_to_cp_leader": transfer, + "encoder_to_cp_leader_handoff_ms": 2.0, + "cp_input_fanout_ms": 3.0, + "cp_workers": [{"dit_ms": 300.0, "finalize_ms": 50.0} for _ in range(6)], + "cp_leader_to_decoder": transfer, + "cp_leader_to_decoder_handoff_ms": 2.0, + "decoder_ms": 7.0, + } + ] + mooncake_sample = TransferStats( + backend="mooncake-rdma", + payload_bytes=256 * 2**20, + registration_ms=0.0, + transfer_ms=6.4, + bandwidth_gbps=41.0, + ) + cp_sample = { + "payload_bytes": float(256 * 2**20), + "transfer_ms": 2.0, + "bandwidth_gbps": 100.0, + } + + summary = _summarize( + records=records, + mooncake_probe={ + "encoder_to_cp_leader": [mooncake_sample], + "cp_leader_to_decoder": [mooncake_sample], + }, + cp_probe={"broadcast": [cp_sample], "all_gather": [cp_sample]}, + baseline=_baseline(), + cp_size=6, + ) + + assert summary["fps"] == pytest.approx(24.0) + assert summary["latency_speedup"] == pytest.approx(4.0) + assert summary["fps_speedup"] == pytest.approx(4.0) + assert summary["dit_speedup"] == pytest.approx(2000.0 / 350.0) + assert summary["cp_efficiency"] == pytest.approx((2000.0 / 350.0) / 6) diff --git a/integrations/lingbot/tests/test_disagg_independent.py b/integrations/lingbot/tests/test_disagg_independent.py new file mode 100644 index 00000000..c6e083f9 --- /dev/null +++ b/integrations/lingbot/tests/test_disagg_independent.py @@ -0,0 +1,100 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for independent aggregated LingBot benchmark reporting.""" + +from __future__ import annotations + +from argparse import Namespace +from pathlib import Path + +import pytest +from lingbot.disagg.benchmark_independent import _child_command, _summarize + +pytestmark = pytest.mark.ci_cpu + + +def _worker_document( + *, + started_at: float, + finished_at: float, + latency_ms: float, + fps: float, + memory_gib: float, +) -> dict: + return { + "environment": { + "measurement_window": { + "started_at": started_at, + "finished_at": finished_at, + } + }, + "summary": { + "fps": fps, + "latency_ms": {"median": latency_ms}, + "memory": { + "peak_gib_by_rank": [memory_gib], + "initialization_peak_gib_by_rank": [memory_gib + 5.0], + "steady_allocated_gib_by_rank": [memory_gib - 2.0], + }, + }, + "records": [ + { + "warmup": False, + "output_frames": 12, + "end_to_end_ms": latency_ms, + } + ], + } + + +def test_summarize_uses_shared_measurement_window() -> None: + summary = _summarize( + [ + _worker_document( + started_at=10.0, + finished_at=12.0, + latency_ms=2000.0, + fps=6.0, + memory_gib=60.0, + ), + _worker_document( + started_at=10.01, + finished_at=12.01, + latency_ms=2010.0, + fps=5.97, + memory_gib=61.0, + ), + ] + ) + + assert summary["aggregate_fps"] == pytest.approx(24 / 2.01) + assert summary["sum_of_worker_fps"] == pytest.approx(11.97) + assert summary["measurement_start_skew_ms"] == pytest.approx(10.0) + assert summary["all_chunk_latency_ms"]["median"] == pytest.approx(2005.0) + assert summary["memory"]["rollout_peak_gib_node_total"] == 121.0 + assert summary["memory"]["initialization_peak_gib_node_total"] == 131.0 + + +def test_child_command_launches_one_isolated_rank(tmp_path: Path) -> None: + args = Namespace( + model="lingbot-world-fast-taehv-window15-sink3", + example_idx=0, + warmup_blocks=6, + measured_blocks=5, + pixel_height=464, + pixel_width=832, + fps=16, + output_dir=tmp_path, + ) + + command = _child_command( + args, + replica_id=3, + barrier_dir=tmp_path / "barrier", + output_dir=tmp_path / "worker-3", + ) + + assert "--nproc_per_node=1" in command + assert command[command.index("--replica-id") + 1] == "3" + assert command[command.index("--pixel-height") + 1] == "464" diff --git a/integrations/lingbot/tests/test_disagg_pipeline.py b/integrations/lingbot/tests/test_disagg_pipeline.py new file mode 100644 index 00000000..187e4890 --- /dev/null +++ b/integrations/lingbot/tests/test_disagg_pipeline.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU tests for the two-stage LingBot DiT pipeline benchmark.""" + +from __future__ import annotations + +import pytest +import torch +from flashdreams.recipes.wan.autoencoder.i2v import I2VCtrl +from lingbot.disagg.benchmark_pipeline import ( + _summarize, + build_pipeline_topology, + split_conditioning, + split_encoder_outputs, + stack_conditioning, + stack_encoder_outputs, + validate_double_buffered_schedule, +) +from lingbot.disagg.stages import LingbotConditioning +from lingbot.encoder.camctrl import I2VCamCtrlEmbeddings +from lingbot.transformer.impl.network import pipeline_partition_bounds + +pytestmark = pytest.mark.ci_cpu + + +def test_pipeline_topology_assigns_three_pairs_and_one_spare() -> None: + topology = build_pipeline_topology(world_size=8, sessions_per_group=2) + + assert topology.io_rank == 0 + assert topology.dit_groups == ((1, 2), (3, 4), (5, 6)) + assert topology.dit_ranks == (1, 2, 3, 4, 5, 6) + assert topology.spare_ranks == (7,) + assert topology.session_count == 6 + + +@pytest.mark.parametrize( + ("stage_index", "expected"), + [(0, (0, 20)), (1, (20, 40))], +) +def test_pipeline_partition_splits_lingbot_layers_evenly( + stage_index: int, + expected: tuple[int, int], +) -> None: + assert ( + pipeline_partition_bounds( + 40, + stage_index=stage_index, + stage_count=2, + ) + == expected + ) + + +def test_stack_conditioning_preserves_session_batch() -> None: + items = [ + LingbotConditioning( + height=58, + width=104, + text_embeddings=torch.full((1, 2, 3), float(index)), + ) + for index in range(2) + ] + + result = stack_conditioning(items) + + assert result.text_embeddings.shape == (2, 2, 3) + assert result.text_embeddings[:, 0, 0].tolist() == [0.0, 1.0] + + +def test_stack_encoder_outputs_preserves_session_batch() -> None: + items = [ + I2VCamCtrlEmbeddings( + i2v=I2VCtrl( + latent=torch.full((1, 3, 2, 2, 2), float(index)), + mask=torch.ones(1, 3, 1, 2, 2), + _is_patchified=False, + ), + plucker=torch.zeros(1, 3, 6, 2, 2), + _is_patchified=False, + ) + for index in range(2) + ] + + result = stack_encoder_outputs(items) + + assert result.i2v.latent.shape == (2, 1, 3, 2, 2, 2) + assert result.i2v.latent[:, 0, 0, 0, 0, 0].tolist() == [0.0, 1.0] + + +def test_split_conditioning_restores_two_batch_one_sessions() -> None: + batched = stack_conditioning( + [ + LingbotConditioning( + height=58, + width=104, + text_embeddings=torch.full((1, 2, 3), float(index)), + ) + for index in range(2) + ] + ) + + sessions = split_conditioning(batched) + + assert [item.text_embeddings.shape for item in sessions] == [(1, 2, 3)] * 2 + assert [item.text_embeddings[0, 0, 0].item() for item in sessions] == [0.0, 1.0] + + +def test_split_encoder_outputs_restores_two_batch_one_sessions() -> None: + batched = stack_encoder_outputs( + [ + I2VCamCtrlEmbeddings( + i2v=I2VCtrl( + latent=torch.full((1, 3, 2, 2, 2), float(index)), + mask=torch.ones(1, 3, 1, 2, 2), + _is_patchified=False, + ), + plucker=torch.zeros(1, 3, 6, 2, 2), + _is_patchified=False, + ) + for index in range(2) + ] + ) + + sessions = split_encoder_outputs(batched) + + assert [item.i2v.latent.shape for item in sessions] == [(1, 1, 3, 2, 2, 2)] * 2 + assert [item.i2v.latent[0, 0, 0, 0, 0, 0].item() for item in sessions] == [ + 0.0, + 1.0, + ] + + +def test_double_buffered_schedule_requires_two_sessions() -> None: + validate_double_buffered_schedule(sessions_per_group=2) + + with pytest.raises(ValueError, match="exactly two"): + validate_double_buffered_schedule(sessions_per_group=1) + + +def test_summary_reports_required_capacity_and_pair_bandwidth() -> None: + topology = build_pipeline_topology(world_size=8, sessions_per_group=2) + records = [ + { + "warmup": False, + "output_frames": 72, + "wave_latency_ms": 2000.0, + "encoder_wave_ms": 20.0, + "decoder_wave_ms": 30.0, + "pair_fanout_ms": [1.0, 1.1, 1.2], + "dit_group_leaders": [{"dit_ms": 1800.0, "finalize_ms": 100.0}] * 3, + } + ] + memory = { + "required_capacity_gib_by_rank": [20.0, 32.0, 34.0, 32.0, 34.0, 32.0, 34.0, 0.0] + } + baseline = { + "topology": {"sessions_per_wave": 7}, + "performance": {"aggregate_fps": 35.15, "per_session_fps": 5.02}, + "peak_allocated_gib_by_rank": [20.0] + [56.3] * 7, + } + + summary = _summarize( + records=records, + topology=topology, + p2p_probe={"1->2": [{"bandwidth_gbps": 300.0}]}, + memory=memory, + baseline=baseline, + ) + + assert summary["aggregate_fps"] == pytest.approx(36.0) + assert summary["per_session_fps"] == pytest.approx(6.0) + assert summary["p2p_probe_gbps"]["all_pairs"]["median"] == 300.0 + assert summary["baseline"]["max_dit_peak_gib"] == 56.3 diff --git a/integrations/lingbot/tests/test_disagg_replicated.py b/integrations/lingbot/tests/test_disagg_replicated.py new file mode 100644 index 00000000..ac2cae03 --- /dev/null +++ b/integrations/lingbot/tests/test_disagg_replicated.py @@ -0,0 +1,115 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for replicated-stage allocation and benchmark summaries.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from lingbot.disagg.benchmark_replicated import ( + StageAllocation, + _summarize, + allocate_stage_replicas, + allocation_from_baseline, +) + +from flashdreams.infra.transfer import TransferStats + +pytestmark = pytest.mark.ci_cpu + + +def _baseline() -> dict[str, Any]: + return { + "summary": { + "fps": 6.0, + "latency_ms": {"median": 2000.0}, + "encoder_ms": {"median": 1.0}, + "encoder_to_dit": {"handoff_ms": {"median": 25.0}}, + "dit_ms": {"median": 1700.0}, + "finalize_ms": {"median": 450.0}, + "decoder_ms": {"median": 7.0}, + "dit_to_decoder": {"handoff_ms": {"median": 12.0}}, + } + } + + +def test_dit_dominated_eight_gpu_allocation_is_one_six_one() -> None: + allocation = allocation_from_baseline(_baseline(), total_gpus=8) + assert allocation == StageAllocation( + encoder_replicas=1, + dit_replicas=6, + decoder_replicas=1, + ) + assert allocation.total_gpus == 8 + + +def test_allocation_requires_one_positive_service_time_per_stage() -> None: + with pytest.raises(ValueError, match="At least three GPUs"): + allocate_stage_replicas( + total_gpus=2, + encoder_service_ms=1.0, + dit_service_ms=1.0, + decoder_service_ms=1.0, + ) + with pytest.raises(ValueError, match="must be positive"): + allocate_stage_replicas( + total_gpus=3, + encoder_service_ms=1.0, + dit_service_ms=0.0, + decoder_service_ms=1.0, + ) + + +def test_replicated_summary_reports_aggregate_and_gpu_normalized_scaling() -> None: + transfer = { + "payload_bytes": 1024, + "transfer_ms": 1.0, + "handoff_ms": 2.0, + } + records = [ + { + "warmup": False, + "output_frames": 72, + "wave_latency_ms": 2400.0, + "encoder_wave_ms": 10.0, + "decoder_wave_ms": 20.0, + "dit_workers": [{"dit_ms": 1900.0, "finalize_ms": 200.0} for _ in range(6)], + "encoder_to_dit": [transfer.copy() for _ in range(6)], + "dit_to_decoder": [transfer.copy() for _ in range(6)], + } + ] + probe = TransferStats( + backend="mooncake-rdma", + payload_bytes=256 * 2**20, + registration_ms=0.0, + transfer_ms=6.4, + bandwidth_gbps=41.0, + ) + + summary = _summarize( + records=records, + probes={"encoder_to_dit_1": [probe]}, + baseline=_baseline(), + dit_replicas=6, + total_gpus=8, + ) + + assert summary["aggregate_fps"] == pytest.approx(30.0) + assert summary["per_session_fps"] == pytest.approx(5.0) + assert summary["throughput_speedup"] == pytest.approx(5.0) + assert summary["gpu_normalized_speedup"] == pytest.approx(1.875) + assert summary["dit_critical_path_ms"]["median"] == pytest.approx(2100.0) diff --git a/integrations/lingbot/tests/test_disagg_scheduler.py b/integrations/lingbot/tests/test_disagg_scheduler.py new file mode 100644 index 00000000..fb4a14d0 --- /dev/null +++ b/integrations/lingbot/tests/test_disagg_scheduler.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for LingBot session-affine scheduling.""" + +from __future__ import annotations + +import pytest +from lingbot.disagg.scheduler import ( + ServiceClass, + SessionAwareScheduler, + SessionRequest, + WorkerSnapshot, + build_microbatches, +) + +pytestmark = pytest.mark.ci_cpu + + +def _worker( + worker_id: str, + *, + pool: str = "io-plus-7-dit", + queue_depth: int = 0, + predicted_chunk_ms: float = 400.0, + rack: str = "rack-a", + nic: str = "mlx5_0", + rdma_capable: bool = True, +) -> WorkerSnapshot: + return WorkerSnapshot( + worker_id=worker_id, + pool=pool, + queue_depth=queue_depth, + predicted_chunk_ms=predicted_chunk_ms, + free_hbm_gib=20.0, + supported_shapes=frozenset({(12, 52, 104)}), + supported_cp_sizes=frozenset({1}), + rack=rack, + nic=nic, + rdma_capable=rdma_capable, + ) + + +def test_scheduler_uses_predicted_wait_and_preserves_session_affinity() -> None: + scheduler = SessionAwareScheduler(min_free_hbm_gib=4.0) + request = SessionRequest( + session_id="session-a", + shape=(12, 52, 104), + cp_size=1, + ) + workers = ( + _worker("slow", queue_depth=0, predicted_chunk_ms=800.0), + _worker("busy-fast", queue_depth=1, predicted_chunk_ms=300.0), + ) + + assert scheduler.assign(request, workers) == "busy-fast" + changed = ( + _worker("slow", queue_depth=0, predicted_chunk_ms=100.0), + _worker("busy-fast", queue_depth=9, predicted_chunk_ms=300.0), + ) + assert scheduler.assign(request, changed) == "busy-fast" + + +def test_scheduler_selects_latency_pool_and_rejects_tcp_fallback() -> None: + scheduler = SessionAwareScheduler() + latency = SessionRequest( + session_id="premium", + shape=(12, 52, 104), + cp_size=1, + service_class=ServiceClass.LATENCY, + ) + + with pytest.raises(RuntimeError, match="No compatible"): + scheduler.assign( + latency, + ( + _worker("throughput"), + _worker( + "latency-tcp", + pool="aggregated-cp8", + rdma_capable=False, + ), + ), + ) + + +def test_microbatches_only_group_compatible_sessions() -> None: + assignments = [ + ( + SessionRequest(session_id=f"s{index}", shape=(12, 52, 104), cp_size=1), + "dit-0", + ) + for index in range(3) + ] + assignments.append( + ( + SessionRequest(session_id="other", shape=(12, 60, 104), cp_size=1), + "dit-0", + ) + ) + + batches = build_microbatches(assignments, max_batch_size=2) + + assert [batch.session_ids for batch in batches] == [ + ("s0", "s1"), + ("s2",), + ("other",), + ] diff --git a/integrations/lingbot/tests/test_disagg_stages.py b/integrations/lingbot/tests/test_disagg_stages.py new file mode 100644 index 00000000..ba68a19a --- /dev/null +++ b/integrations/lingbot/tests/test_disagg_stages.py @@ -0,0 +1,124 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CPU tests for LingBot disaggregation payload boundaries.""" + +from __future__ import annotations + +import pytest +import torch +from einops import rearrange +from lingbot.disagg.stages import ( + LingbotConditioning, + conditioning_from_bundle, + conditioning_to_bundle, + encoder_output_from_bundle, + encoder_output_to_bundle, + encoder_output_to_cp_bundles, +) +from lingbot.encoder.camctrl import I2VCamCtrlEmbeddings + +from flashdreams.recipes.wan.autoencoder.i2v import I2VCtrl + +pytestmark = pytest.mark.ci_cpu + + +def test_conditioning_bundle_preserves_optional_fields_and_spatial_metadata() -> None: + conditioning = LingbotConditioning( + height=58, + width=104, + text_embeddings=torch.randn(1, 4, 8), + negative_text_embeddings=None, + image_embeddings=torch.randn(1, 2, 3), + ) + bundle = conditioning_to_bundle(conditioning) + restored = conditioning_from_bundle(bundle, height=58, width=104) + + assert tuple(bundle) == ("text_embeddings", "image_embeddings") + assert restored.height == 58 + assert restored.width == 104 + assert restored.negative_text_embeddings is None + torch.testing.assert_close( + restored.text_embeddings, + conditioning.text_embeddings, + ) + + +def test_encoder_output_round_trips_without_crossing_patchify_boundary() -> None: + output = I2VCamCtrlEmbeddings( + i2v=I2VCtrl( + latent=torch.randn(3, 16, 2, 2), + mask=torch.ones(3, 16, 2, 2), + ), + plucker=torch.randn(3, 384, 2, 2), + ) + restored = encoder_output_from_bundle(encoder_output_to_bundle(output)) + + assert not restored._is_patchified + assert not restored.i2v._is_patchified + torch.testing.assert_close(restored.i2v.latent, output.i2v.latent) + torch.testing.assert_close(restored.i2v.mask, output.i2v.mask) + torch.testing.assert_close(restored.plucker, output.plucker) + + +def test_encoder_output_direct_cp_shards_match_global_patchify() -> None: + output = I2VCamCtrlEmbeddings( + i2v=I2VCtrl( + latent=torch.arange(4 * 2 * 4 * 6).reshape(4, 2, 4, 6), + mask=torch.ones(4, 2, 4, 6), + ), + plucker=torch.arange(4 * 3 * 4 * 6).reshape(4, 3, 4, 6), + ) + + shards = encoder_output_to_cp_bundles( + output, + cp_size=3, + patch_size=(2, 2, 2), + ) + restored = [ + encoder_output_from_bundle(bundle, patchified=True) for bundle in shards + ] + + assert len(restored) == 3 + assert all(item._is_patchified and item.i2v._is_patchified for item in restored) + assert all(item.i2v.latent.shape == (4, 16) for item in restored) + expected = rearrange( + output.i2v.latent, + "... (t kt) c (h kh) (w kw) -> ... (t h w) (c kt kh kw)", + kt=2, + kh=2, + kw=2, + ) + torch.testing.assert_close( + torch.cat([item.i2v.latent for item in restored], dim=-2), + expected, + ) + + +def test_encoder_output_direct_cp_shards_reject_uneven_tokens() -> None: + output = I2VCamCtrlEmbeddings( + i2v=I2VCtrl( + latent=torch.randn(2, 2, 2, 2), + mask=torch.ones(2, 2, 2, 2), + ), + plucker=torch.randn(2, 3, 2, 2), + ) + + with pytest.raises(ValueError, match="not divisible by CP3"): + encoder_output_to_cp_bundles( + output, + cp_size=3, + patch_size=(1, 1, 1), + ) diff --git a/pyproject.toml b/pyproject.toml index 250048b8..b714c373 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,7 +104,12 @@ invalid-method-override = "ignore" # These packages have no useful type stubs (Triton kernels rewrite their # call signatures via ``@triton.jit``) or are unavailable in CI; treat # imports as Any. -replace-imports-with-any = ["flash_attn.**", "transformer_engine.**", "triton.**"] +replace-imports-with-any = [ + "flash_attn.**", + "mooncake.**", + "transformer_engine.**", + "triton.**", +] [tool.pytest.ini_options] addopts = "--import-mode=importlib -p flashdreams._pytest_plugins.marker_enforcement" diff --git a/skills/use-slurm-gpu-job/SKILL.md b/skills/use-slurm-gpu-job/SKILL.md new file mode 100644 index 00000000..8ce524a7 --- /dev/null +++ b/skills/use-slurm-gpu-job/SKILL.md @@ -0,0 +1,81 @@ +--- +name: use-slurm-gpu-job +description: Run FlashDreams builds, tests, inference, benchmarks, and other nontrivial commands on the Slurm cluster's 8-GPU compute node through srun.sh instead of burdening the login node. Use whenever work in this FlashDreams checkout may consume meaningful CPU, memory, GPU, compilation time, test time, or download bandwidth, or when attaching to an existing Slurm allocation. +--- + +# Use a Slurm GPU job + +Treat the login node as a control plane. Limit work there to file inspection, search, +small edits, Git metadata, shell syntax checks, and allocation management. Run package +installation, builds, test suites, model loading, generation, benchmarks, and other +heavy commands inside the allocation. + +## Start or attach to a job + +Always inspect the user's jobs before starting or attaching: + +```bash +squeue -u "$USER" +``` + +If a job is running, reuse its job ID for every command and test in the task. Do not +request a second allocation. From the directory containing the site-provided +`srun.sh` launcher (commonly `$HOME/work`), attach through the script in a PTY: + +```bash +cd "${FLASHDREAMS_SLURM_LAUNCH_DIR:-$HOME/work}" +./srun.sh 1 +``` + +If no job is running, start the script in a PTY without a job ID: + +```bash +cd "${FLASHDREAMS_SLURM_LAUNCH_DIR:-$HOME/work}" +./srun.sh +``` + +The script requests one interactive node for four hours with eight GPUs. It mounts the +canonical Lustre checkout at `/workspace/flashdreams` inside the container. It also +keeps the uv environment, caches, Hugging Face data, and Triton cache on Lustre. Do not +replace the script with a raw `srun` command or create a working copy under `/home`. + +The site launcher should inspect the current user's queue as a guard. Unless given an +explicit job ID, it should reuse a running job owned by that user and request a new +allocation only when none is running. Inspect the launcher before first use because +accounts, partitions, images, and default checkout paths are site-specific. + +The first positional argument is the node count and the second is the job ID. Prefer +one node unless the user explicitly asks for multi-node work. `SLURM_JOB_ID` and +`SLURM_JOB_NUM_NODES` may also supply those values. + +When operating through Codex, launch the script with a TTY and a short yield time. Keep +the returned terminal session ID. If the allocation is queued, poll that same session; +do not start duplicate allocations. Once the shell is ready, send every heavy command +and test to the same session for the rest of the task. + +## Verify the shell + +Before substantive work, verify that execution is on the allocated node and in the +mounted project: + +```bash +hostname +pwd +nvidia-smi -L +git rev-parse HEAD +``` + +Expect `pwd` to be `/workspace/flashdreams` and eight GPUs to be visible. Compare the +commit to the login-node checkout when there is any doubt about the mount. Stop and fix +the mount if it differs; never test a stale checkout. + +## Run work + +Run commands from `/workspace/flashdreams`. Reuse the allocated terminal for the whole +task, including follow-up tests. Use FlashDreams' narrowest relevant test command first, +then broaden only when useful. Keep large generated data and outputs on Lustre rather +than in `/home`. + +Remember that the allocation has a four-hour wall-clock limit. Preserve useful logs or +results before it expires. When finished, send `exit` to release the allocation. Report +whether validation ran on the compute node and name the commands used. diff --git a/skills/use-slurm-gpu-job/agents/openai.yaml b/skills/use-slurm-gpu-job/agents/openai.yaml new file mode 100644 index 00000000..397b8444 --- /dev/null +++ b/skills/use-slurm-gpu-job/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Use Slurm GPU Job" + short_description: "Run FlashDreams work on an allocated GPU node" + default_prompt: "Use $use-slurm-gpu-job to run this FlashDreams command on an 8-GPU Slurm node." diff --git a/uv.lock b/uv.lock index 023ec81e..3e3c1a69 100644 --- a/uv.lock +++ b/uv.lock @@ -1480,6 +1480,9 @@ dev = [ { name = "pytest" }, { name = "pytest-asyncio" }, ] +disagg = [ + { name = "mooncake-transfer-engine-cuda13" }, +] [package.metadata] requires-dist = [ @@ -1487,12 +1490,13 @@ requires-dist = [ { name = "aiortc", specifier = ">=1.9" }, { name = "flashdreams", editable = "flashdreams" }, { name = "mediapy", specifier = ">=1.1" }, + { name = "mooncake-transfer-engine-cuda13", marker = "extra == 'disagg'", specifier = ">=0.3.12" }, { name = "opencv-python-headless", specifier = ">=4.5" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" }, { name = "pytest-asyncio", marker = "extra == 'dev'", specifier = ">=0.23" }, { name = "scipy", specifier = ">=1.11" }, ] -provides-extras = ["dev"] +provides-extras = ["dev", "disagg"] [[package]] name = "flashdreams-omnidreams" @@ -2860,6 +2864,26 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/ad/3f/3d42e9a78fe5edf792a83c074b13b9b770092a4fbf3462872f4303135f09/ml_dtypes-0.5.4-cp314-cp314t-win_arm64.whl", hash = "sha256:11942cbf2cf92157db91e5022633c0d9474d4dfd813a909383bd23ce828a4b7d", size = 168825, upload-time = "2025-11-17T22:32:23.766Z" }, ] +[[package]] +name = "mooncake-transfer-engine-cuda13" +version = "0.3.12.post1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "aiohttp" }, + { name = "msgpack" }, + { name = "requests" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/94/9c/3ab3e96585242862fb8dd44fa11f1a8c9220cbb9e010eacc3d9b5089fdce/mooncake_transfer_engine_cuda13-0.3.12.post1-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:16c7871e0589f511dfe5645621c557d787e38dc341c493ba614c18dc1f07b190", size = 23264365, upload-time = "2026-07-25T04:16:51.746Z" }, + { url = "https://files.pythonhosted.org/packages/66/5e/4f4861b15f2f34980105d7337499d96601aec6ec387f5bd7a86a8f50f657/mooncake_transfer_engine_cuda13-0.3.12.post1-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:01f1212623631d09137d1118761ed97cf2cef32a1c286427849cb932dd0bdae4", size = 53502565, upload-time = "2026-07-25T04:16:54.577Z" }, + { url = "https://files.pythonhosted.org/packages/5f/9b/684708dac88ec6e977ac76bb5c31337df110c1e4a11031220d6cc3ef12f5/mooncake_transfer_engine_cuda13-0.3.12.post1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:c0b43587a874115a1e4b44644c09ef4823ca2fc38e802778bd11b7cbc0dc6e1e", size = 23265541, upload-time = "2026-07-25T04:17:01.102Z" }, + { url = "https://files.pythonhosted.org/packages/0a/40/9cf112327fcd7fbdc7e30bf62c48f2bfe2f5adc96b315d9287b12a915e72/mooncake_transfer_engine_cuda13-0.3.12.post1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:ed931f03d894067569ed019394f812028f6ce2836431b267526799af5840fd4b", size = 53520866, upload-time = "2026-07-25T04:17:03.419Z" }, + { url = "https://files.pythonhosted.org/packages/62/ce/5a7c55a1fdf1282e2dda8adce737b1094405255a21a7dc868ef97ecf82ab/mooncake_transfer_engine_cuda13-0.3.12.post1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:31664c1aedfcb5938c984ce132f95c42ef523709b7e104b23010c8cd2fec9cbb", size = 23263485, upload-time = "2026-07-25T04:17:05.88Z" }, + { url = "https://files.pythonhosted.org/packages/78/fa/7927b3b03d4840919dc055bf146152ffed8f0e9b06393cc4215293830ce0/mooncake_transfer_engine_cuda13-0.3.12.post1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:4173798950732749113625c2a1f6dd0518bd15e82598a7ecb2b7d5bc8e590b76", size = 53535430, upload-time = "2026-07-25T04:17:08.761Z" }, + { url = "https://files.pythonhosted.org/packages/b0/d1/8dee00882064e091efc355852cee44717aa14610af644817f6d89fed2760/mooncake_transfer_engine_cuda13-0.3.12.post1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5bafc5c30359e073a9f05581ba502ea0204e8c9ee26e10680d1073bb517e9dfd", size = 23266057, upload-time = "2026-07-25T04:17:11.181Z" }, + { url = "https://files.pythonhosted.org/packages/bd/93/6fd44ee1a88f2574f2a921a740313a20ab37604cdb219d4405e7957776c7/mooncake_transfer_engine_cuda13-0.3.12.post1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:280344764b65751f17975cefb824afe73a5064e7cf890d21ced691089f9c3be1", size = 53538915, upload-time = "2026-07-25T04:17:13.896Z" }, +] + [[package]] name = "moviepy" version = "1.0.3" @@ -2885,6 +2909,79 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, ] +[[package]] +name = "msgpack" +version = "1.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/31/f9/c0a1c127f9049db9155afc316952ea571720dd01833ff5e4d7e8e6352dbb/msgpack-1.2.1.tar.gz", hash = "sha256:04c721c2c7448767e9e3f2520a475663d8ee0f09c31890f6d2bd70fd636a9647", size = 183960, upload-time = "2026-06-18T16:13:52.594Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/16/f70100614b69feb3ade7285f08c9c52d6cda0a5c03f3f5e2facd63acb211/msgpack-1.2.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8c7b398c56ff125feae96c2737abfec5595f1fa0aa186df60c56040b8accb95c", size = 82926, upload-time = "2026-06-18T16:12:31.531Z" }, + { url = "https://files.pythonhosted.org/packages/e4/3c/08ecd5cdfe4e2de43aec79062028ad0f7b2d9b1fea5430068c198ba570da/msgpack-1.2.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1548006a91aa93c5da81f3bdcebc1a0d10cea2d25969754fbe848da622b2b895", size = 82730, upload-time = "2026-06-18T16:12:32.894Z" }, + { url = "https://files.pythonhosted.org/packages/19/9f/a70c9cb1a04ecc134005149367dcfe35d167284e8f65035a1e4156ad17b5/msgpack-1.2.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1dabedcd0f23559f3596428c6589c1cd8c6eaed3a0d720795b07b0225d769203", size = 400729, upload-time = "2026-06-18T16:12:34.052Z" }, + { url = "https://files.pythonhosted.org/packages/fa/7f/5ce020168cf0439041526e95aa068c722c016aee21624e331aeabeee2e8e/msgpack-1.2.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:83efa1c898e0fc5380fc0cabbf75164c52e3b5cbb45973710d75821928380c73", size = 407625, upload-time = "2026-06-18T16:12:35.239Z" }, + { url = "https://files.pythonhosted.org/packages/79/70/fb7668ce0386819303047057aef6fc1da73b584291d9cff82b821744e2ef/msgpack-1.2.1-cp310-cp310-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:01e2dd6c9b19d333a00282330cc8a73d38d8dabc306dc5b42cd668c3ac82e833", size = 377891, upload-time = "2026-06-18T16:12:36.684Z" }, + { url = "https://files.pythonhosted.org/packages/3d/dc/9ebe654a73c3aed2e40aa6b52e3c2a02b5f53ef0085fa235a45d5b367f87/msgpack-1.2.1-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:350cb813d0af6e65d2f7ef0d729f7ff5be5a8bce03665892f43e5883d4ecc1b8", size = 391987, upload-time = "2026-06-18T16:12:37.839Z" }, + { url = "https://files.pythonhosted.org/packages/42/eb/b67cf64218a2fa25e1c671fe1d3dbb06cbeb973e71bc4b822da079862d0b/msgpack-1.2.1-cp310-cp310-musllinux_1_2_riscv64.whl", hash = "sha256:ee1d9ed27d0497b848923746cf762ed2e7db24f4be7eec8e5cbe8c766aa707b7", size = 374603, upload-time = "2026-06-18T16:12:39.221Z" }, + { url = "https://files.pythonhosted.org/packages/a2/2e/9ee200cde32fd1a0101b4006202fde554c1860adfb9bf7bff31ea4c08df8/msgpack-1.2.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:633727297ed063441fd1cda2288865487f33ad14eeb8831afb5f0c396a62cfce", size = 405121, upload-time = "2026-06-18T16:12:40.524Z" }, + { url = "https://files.pythonhosted.org/packages/43/b6/f10117be7ca7a51e8feed699a907b8e663a8cd66e115ae6b4fb30cc7945c/msgpack-1.2.1-cp310-cp310-win32.whl", hash = "sha256:298872ecf9e61950f1c6af4ca969b859ee91783bb920ef6e6172697d0c8aad74", size = 64088, upload-time = "2026-06-18T16:12:41.762Z" }, + { url = "https://files.pythonhosted.org/packages/ba/93/89976c696fb0224662239d952c47b4d1661b34d79a332ef5584facaa8579/msgpack-1.2.1-cp310-cp310-win_amd64.whl", hash = "sha256:2ff164c1b0bcb740b073b99e945234d0212852fa378e44a208c425379140dbeb", size = 70113, upload-time = "2026-06-18T16:12:42.78Z" }, + { url = "https://files.pythonhosted.org/packages/f4/6b/e9b1cdc042c4458801d2545ed782a95f3d6ba8e270cce8745b8603c7f748/msgpack-1.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:29a3f6e9667868429d8240dfd063ea5ffdc1321c13d783aa23827a38de0dcb22", size = 82812, upload-time = "2026-06-18T16:12:45.022Z" }, + { url = "https://files.pythonhosted.org/packages/0c/3a/dd518a1bf78ed1e9ad8afe57307c079a00eafe4b3068932a27ca1ea56b4f/msgpack-1.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:aded5bdf32609dc7987a49bbbd15a8ef096193f96dd8bbeb791de729e650acf5", size = 82739, upload-time = "2026-06-18T16:12:46.025Z" }, + { url = "https://files.pythonhosted.org/packages/70/e0/7ba9e1542bf0771a27b8b37c1316e3f95ae9d748fd765284655c476ad4ef/msgpack-1.2.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:146ee4e9ce80b365c6d4c47073da9da7bcec473e58194ceee5dd7620ace77e06", size = 414233, upload-time = "2026-06-18T16:12:47.029Z" }, + { url = "https://files.pythonhosted.org/packages/03/8d/671d81534ea0e2b0e8a121be100020da09eb78861fe3aa8f3ef7dcd3bed1/msgpack-1.2.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a28d076ca7c82b9c8728ad90b7147489449557038bed50e4241eb832395169b4", size = 423843, upload-time = "2026-06-18T16:12:48.19Z" }, + { url = "https://files.pythonhosted.org/packages/d2/b6/e5c737515ed1f166664b87601b532f58cbb73d8aa6a90b99f7c2c5037e8e/msgpack-1.2.1-cp311-cp311-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7d31c0ac0c640f877804c67cb2bc9f4e23dc2db97e96c2e67fa27d38283b41f8", size = 390772, upload-time = "2026-06-18T16:12:49.624Z" }, + { url = "https://files.pythonhosted.org/packages/a8/46/62ed8c2e87d7021eab19921594d961ef3aa3794eec76c716dc30f3bfd433/msgpack-1.2.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8ff92d7feeaf5bc26c51495b69e2f99ed97ab79346fb6555f44be7dd2ac6503b", size = 409559, upload-time = "2026-06-18T16:12:50.936Z" }, + { url = "https://files.pythonhosted.org/packages/70/ff/59aa3887b860bbf43532835e192b1c388a17590d6068ae4f8b2bc74c906e/msgpack-1.2.1-cp311-cp311-musllinux_1_2_riscv64.whl", hash = "sha256:779197a6513bab3c3632265e3d0f7cb3227e62510841a6f34f1eaa37efbb345e", size = 387838, upload-time = "2026-06-18T16:12:52.161Z" }, + { url = "https://files.pythonhosted.org/packages/09/11/f8563e471093420cf6478cb3271a0175d8402b82d879783d4035d2d03360/msgpack-1.2.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:67f6dd22fa72a93752643f07889796d62739a13415ee630169a8ce764f86cf9f", size = 421732, upload-time = "2026-06-18T16:12:53.556Z" }, + { url = "https://files.pythonhosted.org/packages/57/cf/e673683c4c6c90c1022b24c65af4b03eda72b182a1176ef6449069d66acc/msgpack-1.2.1-cp311-cp311-win32.whl", hash = "sha256:91054a783328e0ea7954b8771095705c8d2243b814743fbaadf14552c9c52c5d", size = 64091, upload-time = "2026-06-18T16:12:54.821Z" }, + { url = "https://files.pythonhosted.org/packages/3f/07/ca212739d179f9083bff2c7c08c24101c3555a334fadc2b876b18768a3ae/msgpack-1.2.1-cp311-cp311-win_amd64.whl", hash = "sha256:2eda0b7ebb1283a98d3e4492ac933c8af6aff59fd3df1c3ed024f536af4b1dc8", size = 70462, upload-time = "2026-06-18T16:12:55.898Z" }, + { url = "https://files.pythonhosted.org/packages/6d/be/6798347b425e26f35db82e69dd83c09716c856a3714e7bffc4c0860fd830/msgpack-1.2.1-cp311-cp311-win_arm64.whl", hash = "sha256:6ee967f7c7e1df2890c671ff2ee51a28ded0efc95da3e507176dee881ce36c66", size = 65059, upload-time = "2026-06-18T16:12:57.053Z" }, + { url = "https://files.pythonhosted.org/packages/bc/dd/9e8cbd8f5582ca4b590336f2b91ee5662f6a6ca562b565abaf696a0f81ff/msgpack-1.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:2ef59c659f289eddf8aa6623823f19fa2f40a4029266889eac7a2505dd210c35", size = 83531, upload-time = "2026-06-18T16:12:58.249Z" }, + { url = "https://files.pythonhosted.org/packages/50/2e/ebdb85a8da151397a2790363676b7ed7c125924fe618e4c6d8befb0cc62c/msgpack-1.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d3567748a5107cb40cdf66a275430c2f87c07777698f4bfd25c35f44d533258c", size = 82657, upload-time = "2026-06-18T16:12:59.396Z" }, + { url = "https://files.pythonhosted.org/packages/26/aa/753ad8b007b464e1d8aa0c8e650b9c5f4f725e658fc5ac8a7635c55b7f6e/msgpack-1.2.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:60926b75d00c8e816ef98f3034f484a8bc64242d66839cef4cf7e503142316a0", size = 410634, upload-time = "2026-06-18T16:13:00.383Z" }, + { url = "https://files.pythonhosted.org/packages/6a/fd/6adabd4f6d5e686f97dd02ce7fce3fe4cf672cbac36b8f67ff4040e8ad8b/msgpack-1.2.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:020e881a764b20d8d7ca1a54fc01b8175519d108e3c3f194fddc200bda95951a", size = 419989, upload-time = "2026-06-18T16:13:01.776Z" }, + { url = "https://files.pythonhosted.org/packages/5a/cc/85039b7b0eb168aaad7383a23c97e291a11f08351cb45a606ce865e4e3f1/msgpack-1.2.1-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4202c74688ca06591f78cb18988228bd4cca2cc75d57b60008372892d2f1e6e6", size = 377544, upload-time = "2026-06-18T16:13:03.637Z" }, + { url = "https://files.pythonhosted.org/packages/ed/bf/35963899493b32030c85fc513b723ae66144ac70c11ebc52e889e16e3d99/msgpack-1.2.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:8b267ce94efb76fbd1b3373511420074ee3187f0f7811bf394531de13294735a", size = 400842, upload-time = "2026-06-18T16:13:05.012Z" }, + { url = "https://files.pythonhosted.org/packages/a6/df/8e2ac970c8f99264cd9997d1c73df5466bc19da3301d7dc5500862a9b089/msgpack-1.2.1-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:e4f1d0f8f98ade9634e01fb704a408f9336c0a8f1117b369f5db83dc7551d8b1", size = 374108, upload-time = "2026-06-18T16:13:06.232Z" }, + { url = "https://files.pythonhosted.org/packages/17/dd/fa8bd265110dfa51c20cb529f9e6d240a16fafe7e645004c6af2d01353ba/msgpack-1.2.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f02cf17a6ca1abe29b5f980644f7551f94d71f2011509b26d8625ce038f0df64", size = 414939, upload-time = "2026-06-18T16:13:07.478Z" }, + { url = "https://files.pythonhosted.org/packages/2e/b9/8377a5ad8953fc0437c70cc98d9ae29f27fe5ac5109fbec0812085865735/msgpack-1.2.1-cp312-cp312-win32.whl", hash = "sha256:0c0d9802354507bcba62af19c17918e3eb437cc25e6f50657d511b5856a77aac", size = 64504, upload-time = "2026-06-18T16:13:08.822Z" }, + { url = "https://files.pythonhosted.org/packages/57/7f/ce1e377df7e62461fefd9eb23bfb93a4a523f40a517b377b8f844d836828/msgpack-1.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:5c24aa15d5963051e1a5c62b12c50cd705992502b5ec1f3bece6046f33c9fc24", size = 71421, upload-time = "2026-06-18T16:13:09.828Z" }, + { url = "https://files.pythonhosted.org/packages/8f/32/ebfe84c9929f08f188d56c7a2fd913406a9ddad76a634697c1c43b8112e6/msgpack-1.2.1-cp312-cp312-win_arm64.whl", hash = "sha256:4227224aaec8f7fbcbfbd4272319347b2bb4030366502600f8c45588c5187b07", size = 64775, upload-time = "2026-06-18T16:13:11.056Z" }, + { url = "https://files.pythonhosted.org/packages/b0/ac/dcddcab6f6c20ecb387ca5e980371cdb3f87ff69aeca388be97eebc4c074/msgpack-1.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:0a70e3cf2804a300d921bb0940426e35f4e489a23adfb77a808892241db0a064", size = 83151, upload-time = "2026-06-18T16:13:12.173Z" }, + { url = "https://files.pythonhosted.org/packages/64/71/fbcfa83a1d6a9c6091942d1cfd070962244664b87427a9a49a6897b1b219/msgpack-1.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:491cc39455ca765fad51fb451bf2915eb2cf41192ab5801ce8d67c1d614fe056", size = 82351, upload-time = "2026-06-18T16:13:13.194Z" }, + { url = "https://files.pythonhosted.org/packages/e3/10/ddf7b06db879e8792d13934ddda09ff20bd2a583fd84c9b59aae9b0e650b/msgpack-1.2.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f310233ef7fb9c14e201c93639fe5f5260b005f56f0b29048e999c30935596cc", size = 407518, upload-time = "2026-06-18T16:13:14.233Z" }, + { url = "https://files.pythonhosted.org/packages/79/d3/36a46a8ed992b781acbc05928bd5bee3c810cb0c3563bf81a7b0c04a1a76/msgpack-1.2.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:787c9bebb5833e8f6fc8abca3c0597683d8d87f56a8842b6b89c75a5f3176e2d", size = 416405, upload-time = "2026-06-18T16:13:15.435Z" }, + { url = "https://files.pythonhosted.org/packages/f9/84/e8e9598b557c0ba6ddae901a73780a4c75ac667dddf59414b1e56a42fb34/msgpack-1.2.1-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:dc871b997a9370d855b7394465f2f350e847a5b806dd38dcc9c989e7d87da155", size = 376257, upload-time = "2026-06-18T16:13:17.022Z" }, + { url = "https://files.pythonhosted.org/packages/40/16/738fe6d875ad7e2a9429c165322a4ec088f4f273cdfae63d96a89c467961/msgpack-1.2.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:85f57e960d877f2977f6430896191b04a21f8901b3b4baf2e4604329f4db5402", size = 397469, upload-time = "2026-06-18T16:13:18.287Z" }, + { url = "https://files.pythonhosted.org/packages/ca/be/6d5952df75a7f24f35833af764c3a6860780364cb3a0030beb8099e1b2b4/msgpack-1.2.1-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:1233ee2dd0cefba127583de50ea654677277047d238303521db35def3d7b2e7c", size = 372802, upload-time = "2026-06-18T16:13:19.685Z" }, + { url = "https://files.pythonhosted.org/packages/e1/39/e2ef7dbf0473bcb8dc7c50bf782a892d67414877b63e47fc88eb189ef5e6/msgpack-1.2.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:e3dc2feb0876209d9c38aa56cb1de169bd6c4348f1aa48271f241226590993e6", size = 411273, upload-time = "2026-06-18T16:13:21.028Z" }, + { url = "https://files.pythonhosted.org/packages/ef/c5/133f4512a56e983a93445c836c9d94d88f3bc2e0980ff4b9e577bd8416ce/msgpack-1.2.1-cp313-cp313-win32.whl", hash = "sha256:6d09badf350af2be9d189184e04e64cf54ad93569ab3d96fca58bd3e84aad707", size = 64471, upload-time = "2026-06-18T16:13:22.293Z" }, + { url = "https://files.pythonhosted.org/packages/e2/98/577e10b055096a7dd40732358cabaf7180a20c79ed1dcdbb618e4b9deac7/msgpack-1.2.1-cp313-cp313-win_amd64.whl", hash = "sha256:33f14fba63278b714efe6ad07e50ea5f03d91537aa6a1c5f1ceca4cf44013ca9", size = 71274, upload-time = "2026-06-18T16:13:23.455Z" }, + { url = "https://files.pythonhosted.org/packages/ba/ee/0c0048e7cfbef23c6a94791b8959ab28155232e7956de8a305b5ff588f05/msgpack-1.2.1-cp313-cp313-win_arm64.whl", hash = "sha256:afc5febcd4c99effbc02b528e49d6fd0760b2b7d48c05239e345a5fa6e743d9a", size = 64795, upload-time = "2026-06-18T16:13:24.687Z" }, + { url = "https://files.pythonhosted.org/packages/77/58/cce442852c6b9e1639c7c8ac8fd9143121cb32dab0f308df4d1426a8eb9c/msgpack-1.2.1-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:05f340e47e7e47d2da8db9b53e1bb1d294369e9ef45a747441309f6650b8351d", size = 83610, upload-time = "2026-06-18T16:13:25.724Z" }, + { url = "https://files.pythonhosted.org/packages/60/5c/15b4c7a0182f75ffa90751958ba36a9c01cafee367d49a3edc10ed140b01/msgpack-1.2.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:810b916696c86ef0deb3b74588480224df4c1b071136c34183e4a2a4284d7ac7", size = 83138, upload-time = "2026-06-18T16:13:26.781Z" }, + { url = "https://files.pythonhosted.org/packages/b8/a6/99e58722feaffc5f2fbcc0c8c0d1451ab9f84097f7af87291b46af2390f4/msgpack-1.2.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ca0dacff965c47afdc3749a8469d7302a8f801d6a28758d55120d75e66ce6889", size = 406090, upload-time = "2026-06-18T16:13:28.072Z" }, + { url = "https://files.pythonhosted.org/packages/19/03/8c63e8cf52958534ef688625965ab04c269a6cadd8caef16758b380a821a/msgpack-1.2.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e2bf9280bceb5efca998435904b5d3e9fdbcc11d90dc9df30aec7973252b720", size = 412106, upload-time = "2026-06-18T16:13:29.427Z" }, + { url = "https://files.pythonhosted.org/packages/63/d2/155d9e71b40e41fd934bc0c48b9b2770f22263e1ac20aad8e29fdca7be3f/msgpack-1.2.1-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:aa6c4be5d1c02a42b066ca6ddb71adf36432868fdcdb6ee87e634e86e0674190", size = 374851, upload-time = "2026-06-18T16:13:30.631Z" }, + { url = "https://files.pythonhosted.org/packages/98/48/deaf2326262a8d5ea3295ce9649912ecd3f551ba7ec8e33c665d2ba583f3/msgpack-1.2.1-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ec0e675d59150a6269ddc9139087c722292664a37d071a849c05c473350f1f2d", size = 396168, upload-time = "2026-06-18T16:13:31.977Z" }, + { url = "https://files.pythonhosted.org/packages/10/2a/b4410f906c2ec0008f1608d3ab5143afc3ad3f4e6da0fed3ea2231d0bef4/msgpack-1.2.1-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:dd3bfe82d53edfe4b7fc9a7ec9761e23a7a5b1dac22264505af428253c29ed24", size = 371959, upload-time = "2026-06-18T16:13:33.282Z" }, + { url = "https://files.pythonhosted.org/packages/59/86/1edc67270099a528fa2093ea60fe191233cd238e4bd30cfacf7db79fc959/msgpack-1.2.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:5ad5467fc3f68b5468e06c5f788d712e9f8ffc8b0cd1bcb160c105c1ee92dae7", size = 408457, upload-time = "2026-06-18T16:13:34.567Z" }, + { url = "https://files.pythonhosted.org/packages/82/90/8b630fef07d8c5ab457b71ff2c217910c83d333c7a68472c186e87cc504a/msgpack-1.2.1-cp314-cp314-win32.whl", hash = "sha256:98b58bdb89c46190e4609bb36abe17c6d4105ad13f9c5f8f6f64d320f8ced3fb", size = 65942, upload-time = "2026-06-18T16:13:36.056Z" }, + { url = "https://files.pythonhosted.org/packages/16/f1/467b81e98b24dd3885d7b1857728797b4ffc76a7a7483af4fb321a07de3c/msgpack-1.2.1-cp314-cp314-win_amd64.whl", hash = "sha256:74847557e28ce71bd3c438a447ca90e4b507e997ddbdef8a12a7b283b86c156b", size = 72627, upload-time = "2026-06-18T16:13:37.079Z" }, + { url = "https://files.pythonhosted.org/packages/a7/1d/5d8c4c89985feb6acefb82a09e501c60392261856d2408d20bfe4f0360b1/msgpack-1.2.1-cp314-cp314-win_arm64.whl", hash = "sha256:b50b727bd652bdc37d950336c848ef20ec54a4cafc38dce19b1cd86ad625d0f7", size = 66908, upload-time = "2026-06-18T16:13:38.23Z" }, + { url = "https://files.pythonhosted.org/packages/1b/02/ad2afb678b4de94496cd432b581759b756a92c1192d8c767edd6b132efdc/msgpack-1.2.1-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:8d00f177ca88a77c1cf848d204a38f249751650b601cb6532acc68805d8a8273", size = 86000, upload-time = "2026-06-18T16:13:39.44Z" }, + { url = "https://files.pythonhosted.org/packages/54/74/0b797484013128837f3b1cbb6cea019277c4de4e377dc512b4d9a0f92940/msgpack-1.2.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5bb9c386f0a329c035ddbab4b72d1028bf9627add8dda41070288563d57ed1b1", size = 86544, upload-time = "2026-06-18T16:13:40.447Z" }, + { url = "https://files.pythonhosted.org/packages/a9/b4/b774d7eb95561739907fec675582f83203cf41c597a418c2589b4bfb8e9d/msgpack-1.2.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:20466cca18c49c7292a8984bc15d65857b171e7264bdcb5f96baf8be238791fc", size = 427661, upload-time = "2026-06-18T16:13:41.574Z" }, + { url = "https://files.pythonhosted.org/packages/b2/f9/3243191dc9937e00756c8bc1b0272fed8f23758e43df2a3b46f533e5090f/msgpack-1.2.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:196300e7e5d6e74d50f1607ab9c06c4a1484c383cd22defd727902591f7e8dde", size = 426375, upload-time = "2026-06-18T16:13:42.936Z" }, + { url = "https://files.pythonhosted.org/packages/23/c7/1693111db9944ba4ad4b67a1e788400d78a0b6af7a6523dc7e4e58f8274b/msgpack-1.2.1-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:575957e79cd51903a4e8495a242442949641e08f1efd5197b43bebd3ea7682b4", size = 380495, upload-time = "2026-06-18T16:13:44.306Z" }, + { url = "https://files.pythonhosted.org/packages/3e/2b/92f86956a0c13e8662f7e2ad630c4eb4db07497b967589bd5245e018b2c1/msgpack-1.2.1-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:8c2ed1e48cc0f460bf3c7780e7137ff21a4e18433451916f2442c1b21036cd7d", size = 410897, upload-time = "2026-06-18T16:13:45.629Z" }, + { url = "https://files.pythonhosted.org/packages/da/ea/1479f72d200313a76fc2f823a79d1e07ed052ab7b8a0280640aa7b95de42/msgpack-1.2.1-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:5f6277e5f783c36786a145e0247fc189a03f35f84b251646e53592d2bc12b355", size = 378519, upload-time = "2026-06-18T16:13:46.998Z" }, + { url = "https://files.pythonhosted.org/packages/f5/4d/fa006060ffa1011d32bfae826fe766fe73e02982183601633b7121058ab3/msgpack-1.2.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:f9389552ecf4784886345ead0647e4edc96bee37cbab05b75540f542f766c48c", size = 419815, upload-time = "2026-06-18T16:13:48.205Z" }, + { url = "https://files.pythonhosted.org/packages/2f/e1/aab6c946570496b78e67804721f3d5e2d62a93081b9b37df77764ef56347/msgpack-1.2.1-cp314-cp314t-win32.whl", hash = "sha256:c1c79a604a2969a868a78b6ebd27a887e00c624f14f66b3038e0590cb23332d1", size = 70914, upload-time = "2026-06-18T16:13:49.385Z" }, + { url = "https://files.pythonhosted.org/packages/13/0a/e608956488a2af014cfe6e3d665e090b8ee42aa14b07f8f95b8880d66b09/msgpack-1.2.1-cp314-cp314t-win_amd64.whl", hash = "sha256:f12038a35fabd52e56a3547bab42401af49a45caa6dd00b34c44de235bc93ee2", size = 77999, upload-time = "2026-06-18T16:13:50.467Z" }, + { url = "https://files.pythonhosted.org/packages/d2/8a/27e2e57055176e366a46b85d02d68e7a5bcfbdd8474c9706375d965f24d3/msgpack-1.2.1-cp314-cp314t-win_arm64.whl", hash = "sha256:0adcf06ffde0777c0e1a9b771a2b1c4226ba1bbf748c8efcc02fcdeca3299107", size = 71160, upload-time = "2026-06-18T16:13:51.498Z" }, +] + [[package]] name = "multidict" version = "6.7.1"