diff --git a/integrations_v2/swiftvr/PERFORMANCE_RECOMMENDATIONS.md b/integrations_v2/swiftvr/PERFORMANCE_RECOMMENDATIONS.md new file mode 100644 index 000000000..3af332366 --- /dev/null +++ b/integrations_v2/swiftvr/PERFORMANCE_RECOMMENDATIONS.md @@ -0,0 +1,320 @@ + + +# SwiftVR inference performance recommendations + +This note records the current SwiftVR performance profile and a prioritized set +of experiments. The measurements are diagnostic results from one GB300 system, +not performance or quality guarantees for other GPUs. + +## Summary + +The validated opt-in path is transformer plus ReAE decoder compilation. It +reduced steady eight-frame processing latency from 96.47 ms to 79.60 ms at +2560x1408 output, a 17.5% latency reduction and a throughput increase from 82.9 +to 100.5 FPS. The tradeoff is approximately 142 seconds of cold preparation +with fresh compiler caches on the measured stack. + +The ReAE decoder provided most of the gain: decoder-only compilation reduced +latency by 14.0%. Transformer-only and encoder-only compilation improved it by +1.8% and 4.3%, respectively. Encoder plus decoder matched the selected path +within measurement noise but prepared more slowly, so the preset deliberately +combines the compiled transformer with only the compiled decoder. FP8 +transformer linear layers remain a possible follow-up. The ReAE layout +experiment is complete: explicit layout changes hurt eager execution, while a +compiler-only decoder path cut isolated decoder latency from 18.93 ms to 9.83 +ms without a visible output change. + +Increasing chunk size, changing the SDPA backend, enabling cuDNN benchmarking, +or optimizing adapter format conversions did not show enough potential to be +prioritized. + +## Measurement context + +- FlashDreams base: `origin/main` at `37e208298d105e54fa5e24edb4fe2b05ed516918` +- SwiftVR checkpoint revision: `743ed2530c550764905400f38eb6cc41af5abc80` +- GPU: one NVIDIA GB300, 256703 MiB +- Driver: 595.71.05 +- PyTorch: 2.12.1+cu130 +- PyTorch CUDA runtime: 13.0 +- cuDNN: 9.20.0 +- Precision: BF16 +- Attention: PyTorch dense SDPA, 16x16 shifted windows +- DiT overlap: 0 +- Input: eight 1280x704 RGB frames +- Output: eight 2560x1408 RGB frames in steady state +- Timing: CUDA events with synchronization; medians reported after warmup +- Compiler measurements used fresh processes and fresh TorchInductor caches + +The stage measurements exclude OmniDreams, presentation, encoding, and other +application work. They exercise the same `SwiftVRPipeline` and streaming state +used by the postprocessor. ReAE compilation was validated with matched real +video frames as described below. + +## Earlier diagnostic profile + +The following stage profile predates the mainline SwiftVR refactor and explains +which compile boundaries were investigated. Use the matched current-main +benchmark below for performance decisions. + +| Stage | Eager median | Compiled median | Eager share | +| --- | ---: | ---: | ---: | +| Resize and preprocessing | 0.40 ms | 0.39 ms | 0.3% | +| ReAE encoder | 9.20 ms | 9.25 ms | 7.5% | +| WAN transformer | 81.41 ms | 55.20 ms | 66.3% | +| ReAE decoder | 31.79 ms | 32.06 ms | 25.9% | +| Total per eight frames | 122.77 ms | 96.9 ms | 100% | +| Effective throughput | 65.1 FPS | 82.6 FPS | | + +The operator profile of one warmed eager chunk showed: + +- 333 `copy_` calls for `[1, 7040, 3072]` tensors consumed 17.94 ms. +- NCHW-to-NHWC and NHWC-to-NCHW kernels consumed 9.60 ms and 5.85 ms. +- Dense 3072-to-14336 and 14336-to-3072 FFN projections were prominent. +- SDPA consumed 2.73 ms and window-index gathers consumed 2.72 ms. +- Input conversion took 0.12 ms; output value/layout conversion and the + single-chunk concatenation together remained below 1 ms. + +These results make transformer fusion, dense projections, and ReAE convolution +layout more important than the attention kernel or postprocessor adapter. + +## Investigation order + +### 1. Compile transformer blocks + +**Status:** Useful only with the compiled decoder on the current mainline stack. + +The implementation supports `compile_blocks`, but it defaults to `False`, +including in the `swiftvr-2x` preset. In the current matched benchmark, +transformer-only compilation reduced total latency by 1.8%. Combining it with +the compiled decoder was the fastest tested configuration. + +Cold preparation results were: + +| Operation | Wall time | +| --- | ---: | +| Compile first steady chunk shape | 25.8 s | +| Compile buffered tail shape | 10.0 s | +| First chunk in a replacement stream | 97 ms | + +Recommended experiment: + +1. Enable `compile_blocks=True` only for the Interactive Drive SwiftVR preset. +2. Trigger preparation before the first rollout rather than from the first + postprocessed output. +3. Prewarm both the steady and tail shapes. +4. Keep a stable TorchInductor cache across application launches where the + deployment environment permits it. +5. Keep the eager preset available because cold compilation remains expensive. + +Also test `dynamic=True` as a separate candidate. It may avoid compiling the +tail token length independently, but it must retain the steady-state speedup. + +### 2. Compile ReAE encoder and decoder compute + +**Status:** Decoder recommended as part of the compiled opt-in preset; encoder +available for investigation but not recommended in combination. + +The implementation binds one resident callable for the encoder and one for the +decoder. Causal dictionaries and frame buffers remain explicit inputs and +outputs, so streams share compiled code without sharing temporal state. + +Fresh-process results used 24 real 1280x704 frames, five warmup chunks, twenty +measured chunks, fresh TorchInductor caches, and 2560x1408 output: + +| Candidate | Median | p90 | FPS | Cold prepare | Peak allocated | +| --- | ---: | ---: | ---: | ---: | ---: | +| Fully eager | 96.47 ms | 96.86 ms | 82.9 | 3.2 s | 18.65 GiB | +| ReAE encoder compiled | 92.31 ms | 92.66 ms | 86.7 | 24.7 s | 18.65 GiB | +| ReAE decoder compiled | 82.96 ms | 83.26 ms | 96.4 | 121.2 s | 20.45 GiB | +| ReAE encoder + decoder compiled | 79.91 ms | 80.57 ms | 100.1 | 145.5 s | 20.45 GiB | +| Transformer compiled | 94.69 ms | 95.34 ms | 84.5 | 30.0 s | 18.65 GiB | +| Transformer + decoder compiled | 79.60 ms | 79.97 ms | 100.5 | 142.0 s | 20.45 GiB | + +The registered `swiftvr-2x-compiled` postprocessor also passed an integration +smoke test: 24 inputs produced exactly 24 finite, non-black outputs at +2560x1408, including buffered startup and flush. Preparation took 16.3 seconds +with a warm compiler cache. + +Quality comparison between transformer-only and transformer-plus-decoder over +the 24-frame moving clip produced 54.22 dB PSNR, 0.00101 mean absolute error, +0.00195 RMSE, and 0.00108 temporal-delta MAE in `[0, 1]`. Per-frame mean error +stayed between 0.00092 and 0.00115, with no increasing drift. Normal-scale +contact-sheet inspection showed no visible difference; a 10x +absolute-difference view showed low-amplitude changes around edges and texture. + +Rejected compile boundaries: + +- Compiling 24 individual `_MemoryBlock`, `_TemporalPool`, and `_TemporalGrow` + modules slowed the eager path to 133.46 ms because graph-launch overhead + outweighed fusion. +- `mode="reduce-overhead"` failed because CUDA-graph output storage was reused + and overwritten between block calls. +- Enabling encoder and decoder stage compilation together did not beat the + selected transformer-plus-decoder path and prepared slightly more slowly. + +The reproducible harness is `scripts/benchmark_reae_compile.py`. Use +`--compile-encoder`, `--compile-decoder`, and `--compile-transformer` to isolate +each candidate in a fresh process. + +### ReAE memory layout follow-up + +**Status:** Implemented for the compiled decoder; rejected for eager execution +and the encoder. + +The experiment preserved `channels_last` through ordinary Conv2d regions, +tested `channels_last_3d` independently at SwiftVR's temporal Conv3d +boundaries, and returned to contiguous layout before pixel shuffle. Each case +ran in a fresh process and compiler cache on the same 24-frame input as the +compile sweep. + +| Stage and mode | Contiguous | Conv2d `channels_last` | Conv3d `channels_last_3d` | Combined | +| --- | ---: | ---: | ---: | ---: | +| Encoder eager | **9.05 ms** | 17.10 ms | n/a | n/a | +| Encoder compiled | **4.87 ms** | 4.94 ms | n/a | n/a | +| Decoder eager | **32.03 ms** | 44.17 ms | 32.87 ms | 41.15 ms | +| Decoder compiled | 18.93 ms | 20.54 ms | 10.91 ms | **9.83 ms** | + +The compiled combined result repeated at 9.83 ms against an 18.93 ms repeated +contiguous baseline, a 48.1% isolated decoder reduction. Inductor can absorb +the layout boundaries and select substantially faster temporal Conv3d kernels; +eager execution pays those conversions as separate kernels, which erases or +reverses the gain. Conv2d-only layout changes did not help either stage. + +The production optimization is therefore deliberately coupled to +`compile_reae_decoder=True`; the eager decoder remains unchanged. Direct +compiled decoder outputs were bit-identical in the isolated sweep. A full +24-frame pipeline comparison against the previous compiled decoder measured +0.000048 MAE, 0.000350 RMSE, 69.13 dB PSNR, and no visible difference at normal +scale. The reproducible isolated harness is +`scripts/benchmark_reae_layout.py`. + +### 3. Quantize transformer projections and FFNs to FP8 + +**Status:** Deferred experiment; potentially high impact and higher quality risk. + +The remaining transformer workload is dominated by dense projections across 30 +WAN blocks. FlashDreams already provides `QuantizedNonPersistentLinear`, so an +FP8 prototype does not require another dependency. + +Recommended experiment: + +1. Keep ReAE, normalization, residual accumulation, and initially SDPA in BF16. +2. Start with FP8 E4M3 weights and slice-scaled activations for FFN projections. +3. Add QKV and attention output projections as a separate candidate. +4. Do not prioritize FP8 SDPA: BF16 SDPA is only about 2.2% of current time. +5. Record quantization overhead as part of the measured path. + +Validate a short matched clip first, followed by a motion-heavy long rollout. +Generated detail, flicker, and temporal drift matter more than a small PSNR +difference for this experiment. + +### 4. Overlap OmniDreams and SwiftVR + +**Status:** Deferred architecture change. + +Interactive Drive currently generates one model chunk and then synchronously +postprocesses it. A queued postprocessing worker could process chunk N while the +world model generates chunk N+1. + +This is most promising when SwiftVR runs on a second GPU. On one GPU, concurrent +dense workloads may contend for the same compute and memory bandwidth and must +be benchmarked rather than assumed faster. The design also adds one chunk of +latency and needs bounded queues, ordered flush behavior, and clean error +propagation. + +Do this only if the local SwiftVR optimizations leave end-to-end throughput below +the target. + +### 5. Consider newer upstream runtime knobs selectively + +**Status:** Low priority. + +Newer upstream SwiftVR exposes `torch_compile` and selectable attention +backends. The compile option is valuable. Attention backend selection has a low +ceiling in the current profile because SDPA consumes only 2.73 ms per chunk. + +Synchronize individual runtime improvements rather than replacing the adapted +streaming and postprocessor contracts wholesale. + +## Experiments not worth prioritizing + +### Larger chunks + +| Chunk size | Median latency | Effective throughput | +| --- | ---: | ---: | +| 8 | 122.77 ms | 65.1 FPS | +| 16 | 242.65 ms | 65.9 FPS | +| 24 | 363.13 ms | 66.1 FPS | + +The throughput difference is approximately 1.5%, while larger chunks increase +latency and buffering. Keep eight frames for Interactive Drive. + +### cuDNN benchmark mode + +Enabling `torch.backends.cudnn.benchmark` improved eager throughput by less than +1% on the measured system and added several seconds of cold autotuning. It may +remain useful on a different GPU, but it should not be the next optimization. + +### Attention window or backend changes + +Completely eliminating the measured SDPA time would save only about 2.2% of the +chunk. Smaller attention windows may also change restoration quality. Preserve +the trained 16x16 behavior until larger bottlenecks are addressed. + +### Adapter layout conversion + +The float-to-uint8 input conversion, restored-output conversion, and generic +single-chunk concatenation together consume less than 1 ms. Removing these +copies may simplify the data path later, but it will not materially change +SwiftVR throughput on the measured system. + +## Validation checklist for every candidate + +Run baseline and candidate in separate fresh processes and record: + +- exact commit and flags; +- compiler-cache state and prewarm policy; +- first-visible frame time; +- median and p90 stage and total chunk latency after at least five warmup chunks; +- aggregate FPS over a long rollout; +- peak CUDA allocated and reserved memory; +- replacement-session startup and first chunk; +- flush latency and exact input/output frame count; +- failed compilations, graph breaks, or fallback kernels. + +For quality, use identical inputs and weights and save baseline and candidate +videos. Compare at least: + +- maximum and mean absolute error for compile/layout-only changes; +- PSNR or RMSE for deterministic matched-output checks; +- temporal MAE or an amplified frame-difference video; +- worst-frame crops and side-by-side playback; +- a long, motion-heavy Interactive Drive rollout for flicker and drift. + +Suggested acceptance criteria: + +- at least a 10% steady-state latency reduction for a nontrivial optimization; +- no missing, duplicated, black, or reordered frames; +- no regression in restart/replacement-session behavior; +- no material visual regression on matched clips or long rollouts; +- startup cost documented and absorbed before user-visible generation. + +## End-to-end command template + +Use the same output resolution and rollout length for every candidate: + +```bash +uv run --no-sync flashdreams-run-v2 interactive-drive-omnidreams \ + --mode mp4 \ + --output-path /tmp/interactive-drive-swiftvr-candidate.mp4 \ + --stats-path /tmp/interactive-drive-swiftvr-candidate-stats.json \ + -- --total-blocks 100 --no-ui --width 1280 --height 704 \ + --postprocess-preset swiftvr-2x-compiled +``` + +Rename the output and stats paths for each candidate and retain the eager run as +the baseline. The existing results and quality artifacts are described in +`BENCHMARK.md`. diff --git a/integrations_v2/swiftvr/README.md b/integrations_v2/swiftvr/README.md index 16e0f6289..40c64b659 100644 --- a/integrations_v2/swiftvr/README.md +++ b/integrations_v2/swiftvr/README.md @@ -52,11 +52,22 @@ postprocessor = SwiftVRPostProcessorConfig( Set `checkpoint` to a local directory for offline use. `chunk_size` must be a multiple of four. `dit_overlap=0` is the upstream throughput path; `dit_overlap=1` trades speed for latent overlap blending. `compile_blocks` is -off by default because it adds a long one-time compilation phase. Loading the -upstream FP32 transformer remaps roughly 19 GiB of weights in host memory before -moving them to the selected CUDA dtype; both single-file and standard sharded +off by default because it adds a long one-time compilation phase. +`compile_reae_encoder` and `compile_reae_decoder` independently compile the +ReAE compute paths while keeping causal state explicit. Loading the upstream +FP32 transformer remaps roughly 19 GiB of weights in host memory before moving +them to the selected CUDA dtype; both single-file and standard sharded safetensors checkpoints are accepted. +For long-running 2x streams, the opt-in `swiftvr-2x-compiled` preset compiles +the transformer and ReAE decoder. The compiled decoder also selects optimized +Conv2d and temporal Conv3d memory layouts; eager execution keeps its original +layout because explicit conversions are slower there. On one GB300 at +2560x1408, the initial compile experiment reduced the steady eight-frame +median from 96.47 ms to 79.60 ms, while increasing cold preparation from 3.2 +seconds to about 142 seconds with fresh compiler caches. The regular +`swiftvr-2x` preset remains the startup-friendly fallback. + ## End-to-end V2V demo The integration binds the reusable V2V application: @@ -82,6 +93,18 @@ the source or a same-resolution reference, and report speed only from complete steady-state 8-frame chunks; checkpoint loading and prewarm are intentionally outside those samples. +## Interactive Drive demo + +The shared OmniDreams application accepts SwiftVR through its postprocessor +option. Run OmniDreams at 1280x704 and present 2560x1408 output with: + +```bash +uv run --no-sync flashdreams-run-v2 interactive-drive-omnidreams \ + --mode webrtc --host 0.0.0.0 --port 8089 -- \ + --width 1280 --height 704 \ + --postprocess-preset swiftvr-2x-compiled +``` + ## Validation ```bash @@ -91,3 +114,7 @@ uv run --package flashdreams-swiftvr --extra dev \ The adapted code is pinned in attribution to upstream SwiftVR commit `5ca168cef6ca7200f135fdfea85e5e13d12c5b53` (Apache-2.0). + +See [REAE_COMPILE_BENCHMARK.md](REAE_COMPILE_BENCHMARK.md) for the compile +configuration matrix, GB300 results, quality checks, and RTX PRO 6000 repeat +procedure. diff --git a/integrations_v2/swiftvr/REAE_COMPILE_BENCHMARK.md b/integrations_v2/swiftvr/REAE_COMPILE_BENCHMARK.md new file mode 100644 index 000000000..a58b5090a --- /dev/null +++ b/integrations_v2/swiftvr/REAE_COMPILE_BENCHMARK.md @@ -0,0 +1,216 @@ + + +# SwiftVR ReAE compile benchmark + +This report records the encoder/decoder compile experiment and provides a +matched procedure for repeating it on another GPU. The measurements below are +from one GB300 and should not be treated as performance guarantees for other +hardware. + +## Result + +SwiftVR now exposes independent `compile_reae_encoder` and +`compile_reae_decoder` controls. The compiled functions are bound once to the +resident pipeline, while each stream keeps its own causal encoder and decoder +state. + +The selected `swiftvr-2x-compiled` preset compiles the transformer and ReAE +decoder. It does not compile the encoder: compiling both ReAE directions +together did not improve throughput over the selected path on the measured +stack. + +| Path | Decision | Reason | +| --- | --- | --- | +| Eager | Default fallback | Fast startup and lowest memory use | +| ReAE encoder compiled | Not selected | Only 4.3% lower latency with 21.5 s extra preparation | +| ReAE decoder compiled | Useful component | 14.0% lower latency in isolation | +| ReAE encoder + decoder compiled | Not selected | No gain over transformer + decoder | +| Transformer compiled | Useful component | 1.8% lower latency in isolation | +| Transformer + ReAE decoder compiled | Selected opt-in | 17.5% lower latency; visual validation passed | + +The configuration chart below predates the compiler-only decoder layout +optimization. A later isolated sweep reduced compiled decoder compute from +18.93 ms to 9.83 ms; see +[Compiled decoder layout follow-up](#compiled-decoder-layout-follow-up). + +## GB300 configuration + +| Setting | Value | +| --- | --- | +| Base commit | `origin/main` at `37e208298d105e54fa5e24edb4fe2b05ed516918` plus the compile change containing this report | +| SwiftVR checkpoint | `H-oliday/SwiftVR` revision `743ed2530c550764905400f38eb6cc41af5abc80` | +| GPU | NVIDIA GB300, 256703 MiB | +| Driver | 595.71.05 | +| PyTorch / CUDA / cuDNN | 2.12.1+cu130 / 13.0 / 9.20.0 | +| Precision | BF16 | +| Input | 24 real RGB frames, uint8 THWC, 1280x704 | +| Output | 2560x1408, 2x | +| Chunk / attention window / overlap | 8 frames / 16x16 / 0 | +| Warmup / samples | 5 chunks / 20 chunks | +| Timing | CUDA events with synchronization, fresh process and compiler cache per candidate | + +The 24-frame input tensor had SHA-256 +`f36d724c383358ee9d962caf8521f63a80ca33675d2b7a7a06b4d102b9f1fceb`. +It was center-cropped from a 1280x720 moving driving clip. The benchmark is the +SwiftVR streaming stage only; it excludes OmniDreams, presentation, and video +encoding. + +## Configuration and performance chart + +| Candidate | Transformer | Encoder | Decoder | Median | p90 | Effective FPS | Cold prepare | Peak allocation | +| --- | :---: | :---: | :---: | ---: | ---: | ---: | ---: | ---: | +| Eager | off | off | off | 96.47 ms | 96.86 ms | 82.9 | 3.2 s | 18.65 GiB | +| Encoder | off | on | off | 92.31 ms | 92.66 ms | 86.7 | 24.7 s | 18.65 GiB | +| Decoder | off | off | on | 82.96 ms | 83.26 ms | 96.4 | 121.2 s | 20.45 GiB | +| Encoder + decoder | off | on | on | 79.91 ms | 80.57 ms | 100.1 | 145.5 s | 20.45 GiB | +| Transformer | on | off | off | 94.69 ms | 95.34 ms | 84.5 | 30.0 s | 18.65 GiB | +| Transformer + decoder | on | off | on | **79.60 ms** | **79.97 ms** | **100.5** | 142.0 s | 20.45 GiB | + +## Compiled decoder layout follow-up + +The compiled decoder now preserves `channels_last` through its ordinary +Conv2d regions and uses `channels_last_3d` for SwiftVR's temporal Conv3d +boundaries. The eager path is unchanged because explicit layout conversions +made it slower. + +| Decoder candidate | Eager median | Compiled median | Compiled decision | +| --- | ---: | ---: | --- | +| Contiguous | **32.03 ms** | 18.93 ms | Baseline | +| Conv2d `channels_last` | 44.17 ms | 20.54 ms | Reject | +| Conv3d `channels_last_3d` | 32.87 ms | 10.91 ms | Useful | +| Combined | 41.15 ms | **9.83 ms** | Selected | + +The selected result reproduced in a second fresh-cache run. Direct compiled +decoder outputs were bit-identical across layouts. Through the complete +24-frame pipeline, the selected layout differed from the previous compiled +decoder by 0.000048 MAE and 69.13 dB PSNR, with no visible normal-scale +difference in frames 1, 12, or 24. + +Run the isolated layout sweep with `scripts/benchmark_reae_layout.py`. Keep +baseline and candidate in separate fresh processes and set a fresh +`TORCHINDUCTOR_CACHE_DIR` for every compiled case. + +## Output validation + +The quality comparison isolated decoder compilation by comparing the compiled +transformer against the compiled transformer plus decoder on the same 24 input +frames and weights. + +| Metric | Result | +| --- | ---: | +| Maximum absolute difference | 0.06836 | +| Mean absolute difference | 0.001012 | +| RMSE | 0.001946 | +| PSNR | 54.22 dB | +| Temporal-delta MAE | 0.001083 | +| Per-frame mean-error range | 0.000917-0.001147 | + +Normal-scale side-by-side inspection showed no visible difference through +frame 24 and no increasing temporal drift. A 10x absolute-difference view showed +only low-amplitude changes around texture and edges. The registered +`swiftvr-2x-compiled` postprocessor also passed a 24-frame integration smoke: +it returned exactly 24 finite, non-black 2560x1408 frames through buffered +startup and flush. + +## Repeat on RTX PRO 6000 + +Use the same code revision, checkpoint, input tensor, and environment settings +where possible. Record `git rev-parse HEAD`, the input SHA-256, and the output +JSON files with the result. + +Install the integration and confirm the GPU stack: + +```bash +uv sync --package flashdreams-swiftvr --extra dev \ + --extra interactive-drive --inexact +nvidia-smi --query-gpu=name,driver_version,memory.total \ + --format=csv,noheader +uv run --no-sync python -c \ + 'import torch; print(torch.__version__, torch.version.cuda, torch.backends.cudnn.version())' +``` + +The input must be a `torch.uint8` RGB tensor shaped `[T, 704, 1280, 3]`, with +`T` a positive multiple of eight. For the strictest comparison, copy the +24-frame tensor identified by the checksum above. Otherwise use a real moving +clip, report its checksum, and use the same tensor for every candidate. + +Run each candidate in a fresh process and fresh TorchInductor cache: + +```bash +export SWIFTVR_BENCH_INPUT=/path/to/input-24x704x1280.pt +export SWIFTVR_BENCH_RESULTS=/tmp/swiftvr-compile-rtx6000 +mkdir -p "$SWIFTVR_BENCH_RESULTS" + +run_swiftvr_case() { + local swiftvr_label="$1" + shift + local swiftvr_cache + swiftvr_cache="$(mktemp -d "/tmp/swiftvr-${swiftvr_label}.XXXXXX")" + TORCHINDUCTOR_CACHE_DIR="$swiftvr_cache" uv run --no-sync python \ + integrations_v2/swiftvr/scripts/benchmark_reae_compile.py \ + --label "$swiftvr_label" \ + --input-tensor "$SWIFTVR_BENCH_INPUT" \ + --output-dir "$SWIFTVR_BENCH_RESULTS" "$@" +} + +run_swiftvr_case eager +run_swiftvr_case encoder --compile-encoder +run_swiftvr_case decoder --compile-decoder +run_swiftvr_case encoder-decoder --compile-encoder --compile-decoder +run_swiftvr_case transformer --compile-transformer +run_swiftvr_case transformer-decoder \ + --compile-transformer --compile-decoder +``` + +Each run writes raw timing samples, stack metadata, peak CUDA allocation, and +the restored tensor to `$SWIFTVR_BENCH_RESULTS`. Allow roughly 4 GiB for the +six 24-frame outputs. + +Compare the isolated decoder-compile outputs numerically: + +```bash +uv run --no-sync python - <<'PY' +import math +import os +from pathlib import Path + +import torch + +root = Path(os.environ["SWIFTVR_BENCH_RESULTS"]) +reference = torch.load(root / "transformer.pt", weights_only=True).float() +candidate = torch.load(root / "transformer-decoder.pt", weights_only=True).float() +delta = candidate - reference +rmse = delta.square().mean().sqrt().item() +temporal = ((candidate[:, 1:] - candidate[:, :-1]) - + (reference[:, 1:] - reference[:, :-1])).abs().mean().item() +print({ + "max_abs": delta.abs().max().item(), + "mean_abs": delta.abs().mean().item(), + "rmse": rmse, + "psnr_db": -20 * math.log10(rmse), + "temporal_delta_mae": temporal, +}) +PY +``` + +Finally, exercise the production integration with a fixed world-model seed: + +```bash +uv run --no-sync flashdreams-run-v2 \ + interactive-drive-omnidreams \ + --mode mp4 \ + --output-path /tmp/interactive-drive-swiftvr-compiled.mp4 \ + --stats-path /tmp/interactive-drive-swiftvr-compiled.json -- \ + --total-blocks 100 --no-ui --world-model-seed 42 \ + --width 1280 --height 704 \ + --postprocess-preset swiftvr-2x-compiled +``` + +For the RTX PRO 6000 result, report the same configuration/performance table, +startup time, first steady chunk, median, p90, effective FPS, peak allocation, +frame count, and whether side-by-side playback shows flicker or drift. Keep the +regular `swiftvr-2x` preset as the fallback if compilation is slower or its +startup cost is unsuitable on that stack. diff --git a/integrations_v2/swiftvr/config.py b/integrations_v2/swiftvr/config.py index a7ba67688..a788d399a 100644 --- a/integrations_v2/swiftvr/config.py +++ b/integrations_v2/swiftvr/config.py @@ -36,6 +36,8 @@ def build_swiftvr_pipeline( dtype: torch.dtype = torch.bfloat16, attention_window: tuple[int, int] = (16, 16), compile_blocks: bool = False, + compile_reae_encoder: bool = False, + compile_reae_decoder: bool = False, chunk_size: int = 8, name: str = "swiftvr", ) -> SwiftVRPipelineConfig: @@ -52,6 +54,7 @@ def build_swiftvr_pipeline( encoder=SwiftVREncoderConfig( checkpoint_path=reae_checkpoint, dtype=dtype, + use_compile=compile_reae_encoder, ), diffusion_model=DiffusionModelConfig( transformer=SwiftVRTransformerConfig( @@ -75,6 +78,7 @@ def build_swiftvr_pipeline( decoder=SwiftVRDecoderConfig( checkpoint_path=reae_checkpoint, dtype=dtype, + use_compile=compile_reae_decoder, ), ) diff --git a/integrations_v2/swiftvr/impl/decoder/__init__.py b/integrations_v2/swiftvr/impl/decoder/__init__.py index 6e6f2d0e2..a498dd25e 100644 --- a/integrations_v2/swiftvr/impl/decoder/__init__.py +++ b/integrations_v2/swiftvr/impl/decoder/__init__.py @@ -34,6 +34,9 @@ class SwiftVRDecoderConfig(DecoderConfig): dtype: torch.dtype = torch.bfloat16 """Decoder compute dtype.""" + use_compile: bool = False + """Compile the ReAE decoder compute path.""" + @dataclass(kw_only=True) class SwiftVRDecoderCache(StreamingDecoderCache): @@ -50,7 +53,10 @@ class SwiftVRDecoder(StreamingDecoder[SwiftVRDecoderCache]): def __init__(self, config: SwiftVRDecoderConfig) -> None: super().__init__(config) self.config: SwiftVRDecoderConfig = config - self.network = SwiftVRTAEHV(config.checkpoint_path) + self.network = SwiftVRTAEHV( + config.checkpoint_path, + use_compile=config.use_compile, + ) if config.checkpoint_path is None: self.network.to_empty(device="cpu") for module in self.network.modules(): diff --git a/integrations_v2/swiftvr/impl/decoder/network.py b/integrations_v2/swiftvr/impl/decoder/network.py index 71784ccf1..f79c93bdc 100644 --- a/integrations_v2/swiftvr/impl/decoder/network.py +++ b/integrations_v2/swiftvr/impl/decoder/network.py @@ -7,12 +7,16 @@ from __future__ import annotations +from typing import cast + import torch import torch.nn.functional as F from torch import Tensor, nn +from flashdreams.infra.compile import compile_module +from flashdreams.infra.cuda_graph import set_or_copy from flashdreams.recipes.taehv.checkpoint import legacy_to_blocks_keys -from flashdreams.recipes.taehv.impl import TAEHV, TGrow +from flashdreams.recipes.taehv.impl import TAEHV, Decoder, MemBlock, TGrow class SwiftVRTemporalGrow(nn.Module): @@ -54,10 +58,123 @@ def forward(self, tensor: Tensor) -> Tensor: ) +def _memblock_step_channels_last( + block: MemBlock, + tensor: Tensor, + state: dict[int, Tensor], + batch: int, +) -> Tensor: + key = id(block) + bt, channels, height, width = tensor.shape + time_steps = bt // batch + video = tensor.reshape(batch, time_steps, channels, height, width) + past = torch.cat([state[key], video[:, :-1]], dim=1).reshape_as(tensor) + set_or_copy(state, key, video[:, -1:]) + return block( + tensor.contiguous(memory_format=torch.channels_last), + past.contiguous(memory_format=torch.channels_last), + ) + + +def _temporal_grow_channels_last_3d( + block: SwiftVRTemporalGrow, tensor: Tensor +) -> Tensor: + if block.stride == 1: + assert block.proj is not None + return block.proj(tensor) + assert block.conv3d is not None + frames, channels, height, width = tensor.shape + tensor = F.interpolate( + tensor.unsqueeze(2), + size=(block.stride, height, width), + mode="nearest", + ).contiguous(memory_format=torch.channels_last_3d) + tensor = block.conv3d(tensor) + return tensor.permute(0, 2, 1, 3, 4).reshape( + frames * block.stride, channels, height, width + ) + + +class _SwiftVRCompiledDecoder(nn.Module): + """SwiftVR decoder compute with compiler-friendly convolution layouts.""" + + def __init__( + self, + decoder: Decoder, + *, + channels_last: bool = True, + channels_last_3d: bool = True, + ) -> None: + super().__init__() + self.decoder = decoder + self.channels_last = channels_last + self.channels_last_3d = channels_last_3d + + temporal_convs = { + id(child) + for block in decoder.blocks + if isinstance(block, SwiftVRTemporalGrow) + for child in block.modules() + if isinstance(child, nn.Conv2d) + } + if channels_last: + for child in decoder.modules(): + if isinstance(child, nn.Conv2d) and id(child) not in temporal_convs: + child.weight.data = child.weight.data.contiguous( + memory_format=torch.channels_last + ) + if channels_last_3d: + for child in decoder.modules(): + if isinstance(child, nn.Conv3d): + child.weight.data = child.weight.data.contiguous( + memory_format=torch.channels_last_3d + ) + + @torch.no_grad() + def initialize_state( + self, + z_shape: tuple[int, int, int, int, int], + dtype: torch.dtype, + device: torch.device, + state: dict[int, Tensor], + ) -> None: + """Initialize causal state through the wrapped decoder.""" + self.decoder.initialize_state(z_shape, dtype, device, state) + + def forward(self, tensor: Tensor, state: dict[int, Tensor], batch: int) -> Tensor: + """Decode a latent chunk while preserving convolution memory formats.""" + _, time_steps, channels, height, width = tensor.shape + tensor = tensor.reshape(batch * time_steps, channels, height, width) + if self.channels_last: + tensor = tensor.contiguous(memory_format=torch.channels_last) + for block in self.decoder.blocks: + if isinstance(block, MemBlock) and self.channels_last: + tensor = _memblock_step_channels_last(block, tensor, state, batch) + elif isinstance(block, MemBlock): + tensor = block.cache_step(tensor, state, batch) + elif isinstance(block, SwiftVRTemporalGrow): + if self.channels_last_3d: + tensor = _temporal_grow_channels_last_3d(block, tensor.contiguous()) + else: + tensor = block(tensor.contiguous()) + if self.channels_last: + tensor = tensor.contiguous(memory_format=torch.channels_last) + else: + tensor = block(tensor) + tensor = tensor.contiguous() + _, channels, height, width = tensor.shape + return tensor.reshape(batch, -1, channels, height, width) + + class SwiftVRTAEHV(TAEHV): """Shared TAEHV configured for SwiftVR's ReAE checkpoint.""" - def __init__(self, checkpoint_path: str | None) -> None: + def __init__( + self, + checkpoint_path: str | None, + *, + use_compile: bool = False, + ) -> None: super().__init__( checkpoint_path=None, model_type="wan22", @@ -77,6 +194,11 @@ def __init__(self, checkpoint_path: str | None) -> None: checkpoint_path, state_dict_transform=legacy_to_blocks_keys, ) + if use_compile: + self.decoder = cast( + Decoder, + compile_module(_SwiftVRCompiledDecoder(self.decoder)), + ) __all__ = ["SwiftVRTAEHV", "SwiftVRTemporalGrow"] diff --git a/integrations_v2/swiftvr/impl/encoder/__init__.py b/integrations_v2/swiftvr/impl/encoder/__init__.py index e709f9010..0fb2748d2 100644 --- a/integrations_v2/swiftvr/impl/encoder/__init__.py +++ b/integrations_v2/swiftvr/impl/encoder/__init__.py @@ -7,7 +7,9 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass, field +from functools import partial from typing import Any import torch @@ -24,6 +26,8 @@ from flashdreams.recipes.taehv.impl import Encoder as TAEHVEncoder from flashdreams.recipes.taehv.impl import MemBlock +_EncodeCompleteGroups = Callable[[Tensor, dict[int, Tensor]], Tensor] + def _encode_complete_groups( network: TAEHVEncoder, @@ -63,6 +67,9 @@ class SwiftVREncoderConfig(EncoderConfig): dtype: torch.dtype = torch.bfloat16 """Encoder compute dtype.""" + use_compile: bool = False + """Compile the complete-group ReAE encoder compute path.""" + @dataclass(kw_only=True) class SwiftVREncoderCache(StreamingEncoderCache): @@ -102,6 +109,15 @@ def __init__(self, config: SwiftVREncoderConfig) -> None: } self.network.load_state_dict(encoder_state, strict=True) self.network.to(dtype=config.dtype).eval().requires_grad_(False) + encode_complete_groups: _EncodeCompleteGroups = partial( + _encode_complete_groups, + self.network, + ) + self._encode_complete_groups = ( + torch.compile(encode_complete_groups, mode="default", fullgraph=False) + if config.use_compile + else encode_complete_groups + ) @property def device(self) -> torch.device: @@ -181,7 +197,7 @@ def forward( # type: ignore[override] cache.tail = None if tensor.shape[1] == 0: return None - return _encode_complete_groups(self.network, tensor, cache.state) + return self._encode_complete_groups(tensor, cache.state) def flush(self, cache: SwiftVREncoderCache) -> Tensor | None: """Replicate-pad and encode the final partial temporal group.""" @@ -194,7 +210,7 @@ def flush(self, cache: SwiftVREncoderCache) -> Tensor | None: tensor = torch.cat( [tensor, tensor[:, -1:].expand(-1, padding, -1, -1, -1)], dim=1 ) - return _encode_complete_groups(self.network, tensor, cache.state) + return self._encode_complete_groups(tensor, cache.state) __all__ = [ diff --git a/integrations_v2/swiftvr/impl/pipeline.py b/integrations_v2/swiftvr/impl/pipeline.py index 112947f45..0c976fcc7 100644 --- a/integrations_v2/swiftvr/impl/pipeline.py +++ b/integrations_v2/swiftvr/impl/pipeline.py @@ -243,6 +243,8 @@ def from_pretrained( dtype: torch.dtype, attention_window: tuple[int, int], compile_blocks: bool, + compile_reae_encoder: bool = False, + compile_reae_decoder: bool = False, chunk_size: int = 8, ) -> "SwiftVRPipeline": """Resolve a checkpoint and construct the configured pipeline.""" @@ -257,6 +259,8 @@ def from_pretrained( dtype=dtype, attention_window=attention_window, compile_blocks=compile_blocks, + compile_reae_encoder=compile_reae_encoder, + compile_reae_decoder=compile_reae_decoder, chunk_size=chunk_size, ).setup() assert isinstance(pipeline, cls) diff --git a/integrations_v2/swiftvr/impl/postprocess.py b/integrations_v2/swiftvr/impl/postprocess.py index ddff0d466..647decc2a 100644 --- a/integrations_v2/swiftvr/impl/postprocess.py +++ b/integrations_v2/swiftvr/impl/postprocess.py @@ -116,6 +116,12 @@ class SwiftVRPostProcessorConfig(VideoPostProcessorConfig): compile_blocks: bool = False """Compile transformer blocks. Disabled by default to avoid long startup.""" + compile_reae_encoder: bool = False + """Compile ReAE encoder compute while keeping causal stream state explicit.""" + + compile_reae_decoder: bool = False + """Compile ReAE decoder compute while keeping causal stream state explicit.""" + prewarm: bool = True """Warm model kernels before the first measured rollout chunk.""" @@ -365,6 +371,8 @@ def _load_swiftvr_pipeline(config: SwiftVRPostProcessorConfig) -> SwiftVRPipelin dtype=_resolve_dtype(config.dtype), attention_window=config.attention_window, compile_blocks=config.compile_blocks, + compile_reae_encoder=config.compile_reae_encoder, + compile_reae_decoder=config.compile_reae_decoder, chunk_size=config.chunk_size, ) @@ -383,9 +391,18 @@ def _resolve_dtype(name: _DTypeName) -> torch.dtype: POSTPROCESS_PRESET_SWIFTVR_2X = SwiftVRPostProcessorConfig(scale=2, chunk_size=8) """SwiftVR 2x preset with an 8-frame streaming chunk.""" +POSTPROCESS_PRESET_SWIFTVR_2X_COMPILED = SwiftVRPostProcessorConfig( + scale=2, + chunk_size=8, + compile_blocks=True, + compile_reae_decoder=True, +) +"""Opt-in SwiftVR 2x preset with compiled transformer and ReAE decoder.""" + __all__ = [ "POSTPROCESS_PRESET_SWIFTVR_2X", + "POSTPROCESS_PRESET_SWIFTVR_2X_COMPILED", "POSTPROCESS_PRESET_SWIFTVR_4X", "SwiftVRPostProcessor", "SwiftVRPostProcessorConfig", diff --git a/integrations_v2/swiftvr/pyproject.toml b/integrations_v2/swiftvr/pyproject.toml index 65c2e0c3e..13736e321 100644 --- a/integrations_v2/swiftvr/pyproject.toml +++ b/integrations_v2/swiftvr/pyproject.toml @@ -32,6 +32,7 @@ dev = ["pytest>=8.0"] [project.entry-points."flashdreams.postprocess_presets"] "swiftvr-2x" = "swiftvr.impl.postprocess:POSTPROCESS_PRESET_SWIFTVR_2X" +"swiftvr-2x-compiled" = "swiftvr.impl.postprocess:POSTPROCESS_PRESET_SWIFTVR_2X_COMPILED" "swiftvr-4x" = "swiftvr.impl.postprocess:POSTPROCESS_PRESET_SWIFTVR_4X" [tool.setuptools] diff --git a/integrations_v2/swiftvr/scripts/benchmark_reae_compile.py b/integrations_v2/swiftvr/scripts/benchmark_reae_compile.py new file mode 100644 index 000000000..67fb6fdc1 --- /dev/null +++ b/integrations_v2/swiftvr/scripts/benchmark_reae_compile.py @@ -0,0 +1,204 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark one SwiftVR compile configuration in a fresh process.""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +from pathlib import Path +from typing import Any + +import torch +from swiftvr.impl.pipeline import SwiftVRPipeline +from swiftvr.impl.postprocess import SwiftVRStream +from torch import Tensor + +_CHECKPOINT_REVISION = "743ed2530c550764905400f38eb6cc41af5abc80" + + +def _arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--label", required=True) + parser.add_argument("--input-tensor", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--checkpoint", default="H-oliday/SwiftVR") + parser.add_argument("--output-height", type=int, default=1408) + parser.add_argument("--output-width", type=int, default=2560) + parser.add_argument("--chunk-size", type=int, default=8) + parser.add_argument("--warmup-chunks", type=int, default=5) + parser.add_argument("--measured-chunks", type=int, default=20) + parser.add_argument("--compile-transformer", action="store_true") + parser.add_argument("--compile-encoder", action="store_true") + parser.add_argument("--compile-decoder", action="store_true") + return parser.parse_args() + + +def _load_frames(path: Path, chunk_size: int) -> Tensor: + frames = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(frames, Tensor): + raise TypeError(f"expected a tensor in {path}, got {type(frames).__name__}") + if frames.ndim != 4 or frames.shape[-1] != 3 or frames.dtype != torch.uint8: + raise ValueError( + "input tensor must be uint8 [T, H, W, 3], got " + f"shape={tuple(frames.shape)} dtype={frames.dtype}" + ) + if frames.shape[0] < chunk_size or frames.shape[0] % chunk_size: + raise ValueError( + f"input frame count must be a positive multiple of {chunk_size}, " + f"got {frames.shape[0]}" + ) + return frames.pin_memory() + + +def _chunk(frames: Tensor, index: int, chunk_size: int, device: torch.device) -> Tensor: + chunks = frames.shape[0] // chunk_size + start = index % chunks * chunk_size + return frames[start : start + chunk_size].to(device, non_blocking=True) + + +def _percentile90(samples: list[float]) -> float: + return sorted(samples)[max(0, int(0.9 * len(samples)) - 1)] + + +def _restore_visual( + pipeline: SwiftVRPipeline, + frames: Tensor, + *, + output_height: int, + output_width: int, + chunk_size: int, +) -> Tensor: + stream = SwiftVRStream( + pipeline, + output_height=output_height, + output_width=output_width, + overlap=0, + ) + outputs = [] + for start in range(0, frames.shape[0], chunk_size): + output = stream.step( + frames[start : start + chunk_size].to(pipeline.device, non_blocking=True) + ) + if output is not None: + outputs.append(output) + # A complete four-frame group leaves SwiftVR's three-frame causal startup + # trim outstanding. Feed one replicated frame so flush emits those frames. + output = stream.step(frames[-1:].to(pipeline.device, non_blocking=True)) + if output is not None: + outputs.append(output) + output = stream.flush() + if output is not None: + outputs.append(output) + restored = torch.cat(outputs, dim=1)[:, : frames.shape[0]] + if restored.shape[1] != frames.shape[0]: + raise RuntimeError( + f"restored {restored.shape[1]} frames for {frames.shape[0]} inputs" + ) + return restored + + +def main() -> None: + args = _arguments() + if min(args.chunk_size, args.warmup_chunks, args.measured_chunks) <= 0: + raise ValueError("chunk and iteration counts must be positive") + if args.chunk_size % 4: + raise ValueError("chunk size must be a multiple of four") + + args.output_dir.mkdir(parents=True, exist_ok=True) + frames = _load_frames(args.input_tensor, args.chunk_size) + device = torch.device("cuda") + torch.cuda.reset_peak_memory_stats(device) + load_started = time.perf_counter() + pipeline = SwiftVRPipeline.from_pretrained( + args.checkpoint, + revision=_CHECKPOINT_REVISION, + device="cuda", + dtype=torch.bfloat16, + attention_window=(16, 16), + compile_blocks=args.compile_transformer, + compile_reae_encoder=args.compile_encoder, + compile_reae_decoder=args.compile_decoder, + chunk_size=args.chunk_size, + ) + model_load_seconds = time.perf_counter() - load_started + + stream = SwiftVRStream( + pipeline, + output_height=args.output_height, + output_width=args.output_width, + overlap=0, + ) + prepare_started = time.perf_counter() + for index in range(args.warmup_chunks): + stream.step(_chunk(frames, index, args.chunk_size, device)) + stream.step(frames[-1:].to(device, non_blocking=True)) + stream.flush() + torch.cuda.synchronize(device) + prepare_seconds = time.perf_counter() - prepare_started + + stream = SwiftVRStream( + pipeline, + output_height=args.output_height, + output_width=args.output_width, + overlap=0, + ) + # Populate causal state before measuring the steady path. + stream.step(_chunk(frames, 0, args.chunk_size, device)) + torch.cuda.synchronize(device) + samples = [] + for index in range(args.measured_chunks): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + stream.step(_chunk(frames, index + 1, args.chunk_size, device)) + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + + visual = _restore_visual( + pipeline, + frames, + output_height=args.output_height, + output_width=args.output_width, + chunk_size=args.chunk_size, + ) + torch.cuda.synchronize(device) + output_path = args.output_dir / f"{args.label}.pt" + torch.save(visual.to(device="cpu", dtype=torch.float16), output_path) + + median_ms = statistics.median(samples) + result: dict[str, Any] = { + "label": args.label, + "checkpoint_revision": _CHECKPOINT_REVISION, + "input_shape": list(frames.shape), + "output_shape": list(visual.shape), + "chunk_size": args.chunk_size, + "warmup_chunks": args.warmup_chunks, + "measured_chunks": args.measured_chunks, + "compile_transformer": args.compile_transformer, + "compile_encoder": args.compile_encoder, + "compile_decoder": args.compile_decoder, + "model_load_seconds": model_load_seconds, + "prepare_seconds": prepare_seconds, + "median_chunk_ms": median_ms, + "p90_chunk_ms": _percentile90(samples), + "effective_fps": args.chunk_size * 1000 / median_ms, + "peak_cuda_allocated_gib": torch.cuda.max_memory_allocated(device) / 2**30, + "samples_ms": samples, + "gpu": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "output_tensor": str(output_path), + } + result_path = args.output_dir / f"{args.label}.json" + result_path.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/integrations_v2/swiftvr/scripts/benchmark_reae_layout.py b/integrations_v2/swiftvr/scripts/benchmark_reae_layout.py new file mode 100644 index 000000000..74ddc5113 --- /dev/null +++ b/integrations_v2/swiftvr/scripts/benchmark_reae_layout.py @@ -0,0 +1,366 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Benchmark experimental SwiftVR ReAE memory layouts in a fresh process.""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +from pathlib import Path +from typing import Any, Literal + +import torch +import torch.nn.functional as F +from swiftvr.config import _resolve_checkpoint +from swiftvr.impl.decoder.network import SwiftVRTAEHV, _SwiftVRCompiledDecoder +from swiftvr.impl.encoder import SwiftVREncoder, SwiftVREncoderConfig +from torch import Tensor, nn + +from flashdreams.infra.compile import compile_module +from flashdreams.infra.cuda_graph import set_or_copy +from flashdreams.recipes.taehv.impl import Encoder, MemBlock, TPool + +_CHECKPOINT_REVISION = "743ed2530c550764905400f38eb6cc41af5abc80" +_Layout = Literal["contiguous", "channels-last", "channels-last-3d", "combined"] + + +def _arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--stage", choices=("encoder", "decoder"), required=True) + parser.add_argument("--layout", choices=_Layout.__args__, required=True) + parser.add_argument("--label", required=True) + parser.add_argument("--input-tensor", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, required=True) + parser.add_argument("--checkpoint", default="H-oliday/SwiftVR") + parser.add_argument("--output-height", type=int, default=1408) + parser.add_argument("--output-width", type=int, default=2560) + parser.add_argument("--chunk-size", type=int, default=8) + parser.add_argument("--warmup-chunks", type=int, default=5) + parser.add_argument("--measured-chunks", type=int, default=20) + parser.add_argument("--compile", action="store_true") + return parser.parse_args() + + +def _percentile90(samples: list[float]) -> float: + return sorted(samples)[max(0, int(0.9 * len(samples)) - 1)] + + +def _uses_channels_last(layout: _Layout) -> bool: + return layout in ("channels-last", "combined") + + +def _uses_channels_last_3d(layout: _Layout) -> bool: + return layout in ("channels-last-3d", "combined") + + +def _convert_weights(module: nn.Module, layout: _Layout) -> None: + temporal_convs = { + id(child) + for boundary in module.modules() + if isinstance(boundary, TPool) + for child in boundary.modules() + if isinstance(child, nn.Conv2d) + } + if _uses_channels_last(layout): + for child in module.modules(): + if isinstance(child, nn.Conv2d) and id(child) not in temporal_convs: + child.weight.data = child.weight.data.contiguous( + memory_format=torch.channels_last + ) + + +def _memblock_step( + block: MemBlock, + tensor: Tensor, + state: dict[int, Tensor], + batch: int, + *, + channels_last: bool, +) -> Tensor: + key = id(block) + bt, channels, height, width = tensor.shape + time_steps = bt // batch + video = tensor.reshape(batch, time_steps, channels, height, width) + if key not in state: + state[key] = tensor.new_zeros(batch, 1, channels, height, width) + past = torch.cat([state[key], video[:, :-1]], dim=1).reshape_as(tensor) + set_or_copy(state, key, video[:, -1:]) + if channels_last: + tensor = tensor.contiguous(memory_format=torch.channels_last) + past = past.contiguous(memory_format=torch.channels_last) + return block(tensor, past) + + +class _EncoderRunner(nn.Module): + def __init__(self, network: Encoder, *, channels_last: bool) -> None: + super().__init__() + self.network = network + self.channels_last = channels_last + + def forward(self, tensor: Tensor, state: dict[int, Tensor]) -> Tensor: + batch, time_steps, channels, height, width = tensor.shape + tensor = tensor.reshape(batch * time_steps, channels, height, width) + if self.channels_last: + tensor = tensor.contiguous(memory_format=torch.channels_last) + for block in self.network.blocks: + if isinstance(block, MemBlock): + tensor = _memblock_step( + block, + tensor, + state, + batch, + channels_last=self.channels_last, + ) + elif isinstance(block, TPool): + tensor = block(tensor.contiguous()) + if self.channels_last: + tensor = tensor.contiguous(memory_format=torch.channels_last) + else: + tensor = block(tensor) + tensor = tensor.contiguous() + _, channels, height, width = tensor.shape + return tensor.reshape(batch, -1, channels, height, width) + + +def _load_frames( + path: Path, + *, + output_height: int, + output_width: int, + chunk_size: int, + device: torch.device, +) -> Tensor: + frames = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(frames, Tensor): + raise TypeError(f"expected a tensor in {path}, got {type(frames).__name__}") + if frames.ndim != 4 or frames.shape[-1] != 3 or frames.dtype != torch.uint8: + raise ValueError( + "encoder input must be uint8 [T,H,W,3], got " + f"shape={tuple(frames.shape)} dtype={frames.dtype}" + ) + if frames.shape[0] < chunk_size or frames.shape[0] % chunk_size: + raise ValueError( + f"frame count must be a positive multiple of {chunk_size}, " + f"got {frames.shape[0]}" + ) + frames = frames.to(device=device, dtype=torch.bfloat16).permute(0, 3, 1, 2) + frames = F.interpolate( + frames, + size=(output_height, output_width), + mode="bilinear", + align_corners=False, + ).div_(255) + frames = F.pixel_unshuffle(frames, 2) + return frames.reshape(1, frames.shape[0], *frames.shape[1:]) + + +def _load_latents(path: Path, chunk_size: int, device: torch.device) -> Tensor: + latents = torch.load(path, map_location="cpu", weights_only=True) + if not isinstance(latents, Tensor) or latents.ndim != 5: + raise ValueError(f"decoder input must be a 5D tensor, got {type(latents)}") + latent_chunk = chunk_size // 4 + if latents.shape[1] < latent_chunk or latents.shape[1] % latent_chunk: + raise ValueError( + f"latent frame count must be a multiple of {latent_chunk}, " + f"got {latents.shape[1]}" + ) + return latents.to(device=device, dtype=torch.bfloat16) + + +def _measure( + step: Any, + chunks: list[Tensor], + *, + new_state: Any, + warmup_chunks: int, + measured_chunks: int, + device: torch.device, +) -> tuple[float, list[float]]: + state = new_state() + started = time.perf_counter() + for index in range(warmup_chunks): + step(chunks[index % len(chunks)], state) + torch.cuda.synchronize(device) + prepare_seconds = time.perf_counter() - started + + state = new_state() + step(chunks[0], state) + torch.cuda.synchronize(device) + samples = [] + for index in range(measured_chunks): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + step(chunks[(index + 1) % len(chunks)], state) + end.record() + end.synchronize() + samples.append(start.elapsed_time(end)) + return prepare_seconds, samples + + +@torch.inference_mode() +def _benchmark_encoder( + args: argparse.Namespace, + checkpoint_path: str, + device: torch.device, +) -> tuple[float, list[float], Tensor]: + encoder = SwiftVREncoder( + SwiftVREncoderConfig( + checkpoint_path=checkpoint_path, + dtype=torch.bfloat16, + use_compile=False, + ) + ).to(device=device, dtype=torch.bfloat16) + _convert_weights(encoder.network, args.layout) + runner: nn.Module = _EncoderRunner( + encoder.network, + channels_last=_uses_channels_last(args.layout), + ) + if args.compile: + runner = compile_module(runner, mode="default") + packed = _load_frames( + args.input_tensor, + output_height=args.output_height, + output_width=args.output_width, + chunk_size=args.chunk_size, + device=device, + ) + chunks = list(packed.split(args.chunk_size, dim=1)) + + def step(chunk: Tensor, state: dict[int, Tensor]) -> Tensor: + return runner(chunk, state) + + prepare_seconds, samples = _measure( + step, + chunks, + new_state=dict, + warmup_chunks=args.warmup_chunks, + measured_chunks=args.measured_chunks, + device=device, + ) + state: dict[int, Tensor] = {} + output = torch.cat([step(chunk, state) for chunk in chunks], dim=1) + return prepare_seconds, samples, output + + +@torch.inference_mode() +def _benchmark_decoder( + args: argparse.Namespace, + checkpoint_path: str, + device: torch.device, +) -> tuple[float, list[float], Tensor]: + network = SwiftVRTAEHV(checkpoint_path, use_compile=False).to( + device=device, dtype=torch.bfloat16 + ) + runner: nn.Module = _SwiftVRCompiledDecoder( + network.decoder, + channels_last=_uses_channels_last(args.layout), + channels_last_3d=_uses_channels_last_3d(args.layout), + ) + if args.compile: + runner = compile_module(runner) + latents = _load_latents(args.input_tensor, args.chunk_size, device) + chunks = list(latents.split(args.chunk_size // 4, dim=1)) + + def new_state() -> dict[int, Tensor]: + return {} + + def step(chunk: Tensor, state: dict[int, Tensor]) -> Tensor: + first = not state + if first: + batch, time_steps, channels, height, width = chunk.shape + network.decoder.initialize_state( + (batch, time_steps, channels, height, width), + chunk.dtype, + chunk.device, + state, + ) + output = runner(chunk, state, chunk.shape[0]).clamp_(0, 1) + n, time_steps, channels, height, width = output.shape + output = F.pixel_shuffle( + output.reshape(n * time_steps, channels, height, width), + network.patch_size, + ).reshape(n, time_steps, 3, height * 2, width * 2) + return output[:, network.frames_to_trim :] if first else output + + prepare_seconds, samples = _measure( + step, + chunks, + new_state=new_state, + warmup_chunks=args.warmup_chunks, + measured_chunks=args.measured_chunks, + device=device, + ) + state = new_state() + output = torch.cat([step(chunk, state) for chunk in chunks], dim=1) + return prepare_seconds, samples, output + + +def main() -> None: + args = _arguments() + if min(args.chunk_size, args.warmup_chunks, args.measured_chunks) <= 0: + raise ValueError("chunk and iteration counts must be positive") + if args.chunk_size % 4: + raise ValueError("chunk size must be a multiple of four") + if args.stage == "encoder" and _uses_channels_last_3d(args.layout): + raise ValueError("the encoder has no Conv3d path") + + args.output_dir.mkdir(parents=True, exist_ok=True) + checkpoint_root = _resolve_checkpoint( + args.checkpoint, + revision=_CHECKPOINT_REVISION, + ) + checkpoint_path = str(checkpoint_root / "reae.safetensors") + device = torch.device("cuda") + torch.cuda.reset_peak_memory_stats(device) + started = time.perf_counter() + if args.stage == "encoder": + prepare_seconds, samples, output = _benchmark_encoder( + args, checkpoint_path, device + ) + else: + prepare_seconds, samples, output = _benchmark_decoder( + args, checkpoint_path, device + ) + total_seconds = time.perf_counter() - started + torch.cuda.synchronize(device) + + output_path = args.output_dir / f"{args.label}.pt" + torch.save(output.to(device="cpu", dtype=torch.float16), output_path) + median_ms = statistics.median(samples) + result: dict[str, Any] = { + "label": args.label, + "stage": args.stage, + "layout": args.layout, + "compiled": args.compile, + "input_shape": list( + torch.load(args.input_tensor, map_location="cpu", weights_only=True).shape + ), + "output_shape": list(output.shape), + "chunk_size": args.chunk_size, + "warmup_chunks": args.warmup_chunks, + "measured_chunks": args.measured_chunks, + "prepare_seconds": prepare_seconds, + "total_seconds": total_seconds, + "median_chunk_ms": median_ms, + "p90_chunk_ms": _percentile90(samples), + "effective_fps": args.chunk_size * 1000 / median_ms, + "peak_cuda_allocated_gib": torch.cuda.max_memory_allocated(device) / 2**30, + "samples_ms": samples, + "gpu": torch.cuda.get_device_name(device), + "torch": torch.__version__, + "cuda": torch.version.cuda, + "cudnn": torch.backends.cudnn.version(), + "checkpoint_revision": _CHECKPOINT_REVISION, + "output_tensor": str(output_path), + } + result_path = args.output_dir / f"{args.label}.json" + result_path.write_text(json.dumps(result, indent=2) + "\n") + print(json.dumps(result, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/integrations_v2/swiftvr/tests/test_app.py b/integrations_v2/swiftvr/tests/test_app.py index 443a964e6..b29c8ebd5 100644 --- a/integrations_v2/swiftvr/tests/test_app.py +++ b/integrations_v2/swiftvr/tests/test_app.py @@ -38,9 +38,14 @@ def test_entry_point_binds_swiftvr_defaults() -> None: def test_postprocess_preset_is_discoverable() -> None: preset = resolve_postprocess_preset("swiftvr-4x") twice = resolve_postprocess_preset("swiftvr-2x") + compiled = resolve_postprocess_preset("swiftvr-2x-compiled") assert isinstance(preset, SwiftVRPostProcessorConfig) assert preset.scale == 4 assert isinstance(twice, SwiftVRPostProcessorConfig) assert twice.scale == 2 assert twice.chunk_size == 8 + assert isinstance(compiled, SwiftVRPostProcessorConfig) + assert compiled.compile_blocks + assert not compiled.compile_reae_encoder + assert compiled.compile_reae_decoder diff --git a/integrations_v2/swiftvr/tests/test_model.py b/integrations_v2/swiftvr/tests/test_model.py index f38fe6024..a4f94e4ea 100644 --- a/integrations_v2/swiftvr/tests/test_model.py +++ b/integrations_v2/swiftvr/tests/test_model.py @@ -3,6 +3,8 @@ """CPU contracts for the FlashDreams-native SwiftVR pipeline.""" +import copy +from collections.abc import Callable from pathlib import Path from types import SimpleNamespace from typing import Any, cast @@ -12,7 +14,11 @@ from swiftvr.config import build_swiftvr_pipeline from swiftvr.impl.attention import SwiftVRBlock, _axis_starts from swiftvr.impl.decoder import SwiftVRDecoderConfig -from swiftvr.impl.decoder.network import SwiftVRTAEHV, SwiftVRTemporalGrow +from swiftvr.impl.decoder.network import ( + SwiftVRTAEHV, + SwiftVRTemporalGrow, + _SwiftVRCompiledDecoder, +) from swiftvr.impl.encoder import ( SwiftVREncoder, SwiftVREncoderConfig, @@ -28,7 +34,7 @@ from flashdreams.infra.pipeline import StreamInferencePipeline from flashdreams.infra.profiler import EventProfiler from flashdreams.recipes.taehv.checkpoint import legacy_to_blocks_keys -from flashdreams.recipes.taehv.impl import TAEHV, MemBlock, TGrow, TPool +from flashdreams.recipes.taehv.impl import TAEHV, Decoder, MemBlock, TGrow, TPool from flashdreams.recipes.taehv.impl import Encoder as TAEHVEncoder from flashdreams.recipes.wan import wan_dit_state_dict_from_diffusers from flashdreams.recipes.wan.transformer.impl.network import ( @@ -52,6 +58,8 @@ def test_pipeline_config_follows_stream_inference_component_contracts( checkpoint=str(tmp_path), revision=None, attention_window=(8, 12), + compile_reae_encoder=True, + compile_reae_decoder=True, chunk_size=24, ) @@ -61,6 +69,8 @@ def test_pipeline_config_follows_stream_inference_component_contracts( assert isinstance(config.encoder, SwiftVREncoderConfig) assert isinstance(config.diffusion_model.transformer, SwiftVRTransformerConfig) assert isinstance(config.decoder, SwiftVRDecoderConfig) + assert config.encoder.use_compile + assert config.decoder.use_compile assert config.diffusion_model.transformer.network.attention_window == (8, 12) assert config.diffusion_model.transformer.latent_frames == 6 @@ -226,6 +236,79 @@ def test_reae_complete_group_adapter_preserves_stream_boundaries() -> None: torch.testing.assert_close(chunked, whole) +def test_reae_compiled_decoder_layout_preserves_stream_output() -> None: + torch.manual_seed(0) + decoder = Decoder( + n_f=(4, 4, 4, 4), + latent_channels=2, + image_channels=3, + patch_size=2, + decoder_time_upscale=(True, True), + decoder_space_upscale=(False, False, False), + act_func=torch.nn.ReLU(inplace=True), + ).eval() + for index, block in enumerate(decoder.blocks): + if isinstance(block, TGrow): + decoder.blocks[index] = SwiftVRTemporalGrow( + int(block.conv.in_channels), block.stride + ) + + baseline = copy.deepcopy(decoder) + candidate = _SwiftVRCompiledDecoder(copy.deepcopy(decoder)).eval() + baseline_state: dict[int, torch.Tensor] = {} + candidate_state: dict[int, torch.Tensor] = {} + chunks = [torch.randn(1, 2, 2, 4, 4) for _ in range(2)] + baseline.initialize_state( + (1, 2, 2, 4, 4), + chunks[0].dtype, + chunks[0].device, + baseline_state, + ) + candidate.initialize_state( + (1, 2, 2, 4, 4), + chunks[0].dtype, + chunks[0].device, + candidate_state, + ) + + expected = torch.cat( + [baseline(chunk, baseline_state, chunk.shape[0]) for chunk in chunks], + dim=1, + ) + actual = torch.cat( + [candidate(chunk, candidate_state, chunk.shape[0]) for chunk in chunks], + dim=1, + ) + + torch.testing.assert_close(actual, expected) + assert all( + convolution.weight.is_contiguous(memory_format=torch.channels_last_3d) + for convolution in candidate.modules() + if isinstance(convolution, torch.nn.Conv3d) + ) + + +def test_reae_encoder_compile_callable_is_bound_once( + monkeypatch: pytest.MonkeyPatch, +) -> None: + compiled: list[tuple[Callable[..., Any], dict[str, Any]]] = [] + + def compile_function( + function: Callable[..., Any], **kwargs: Any + ) -> Callable[..., Any]: + compiled.append((function, kwargs)) + return function + + monkeypatch.setattr(torch, "compile", compile_function) + encoder = SwiftVREncoder( + SwiftVREncoderConfig(dtype=torch.float32, use_compile=True) + ) + + assert compiled == [ + (encoder._encode_complete_groups, {"mode": "default", "fullgraph": False}) + ] + + def _diffusers_checkpoint_key(native_key: str) -> str: replacements = ( ("text_embedding.0", "condition_embedder.text_embedder.linear_1"),