Perf(LTX2): Comprehensive XLA, Memory, and Transformer Code Quality Optimizations - #422
Perf(LTX2): Comprehensive XLA, Memory, and Transformer Code Quality Optimizations#422Perseus14 wants to merge 30 commits into
Conversation
Perseus14
changed the title
3a13196 to
36e0f5d
Compare
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
36e0f5d to
1bbd72c
Compare
|
🤖 Hi @Perseus14, I've received your request, and I'm working on it now! You can track my progress in the logs for more details. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
This pull request introduces a highly comprehensive set of optimizations and refactorings for the LTX2 model family, delivering significant enhancements to memory usage, model modularity, and JAX compilation behavior. By transforming static continuous hyperparameters into dynamic inputs, the sweep successfully eliminates expensive compilation timeouts while resolving underlying layout mismatches and unneeded computational overhead. The architectural quality of the codebase is substantially improved, particularly in the block consolidation and cleaner RoPE logic.
🔍 General Feedback
- Exceptional Block Refactoring: The transition of
LTX2VideoTransformer3DModelblock application to a single unifiedapply_blockhelper and an immutableLTX2BlockContextis a masterclass in reducing hazard-prone code duplication. - RoPE & Coordinates Cleanup: Removing the duplicated and dead coordinate generation block in
prepare_video_coordsand simplifying split-RoPE dimension manipulation are excellent and highly effective hygiene improvements. - Robustness in Fallbacks & Dtypes: Incorporating proper warning logs for the noise generator fallback and upfront dtype casting for prompt embeddings ensures both runtime reliability and optimal HBM performance.
|
🤖 I'm sorry @Perseus14, but I was unable to process your request. Please see the logs for more details. |
jax.lax.scan bakes iteration count into array shapes, so a 2-step warmup compiles a different XLA program than the N-step real run, causing a wasteful double compilation. Only reduce warmup steps when scan_diffusion_loop=False (Python loop), where per-step JIT compilation is independent of total iteration count.
This reverts commit aa1c697.
… of nnx.List indices
48884c1 to
fedbc0d
Compare
…r on Pipeline getattr
3 participants
LTX2 Improvements, Optimizations & Bug Fixes
This PR introduces a massive overhaul to the LTX2 and LTX2.3 implementations in MaxDiffusion. The primary focus spans architectural cleanliness, persistent memory caching, AOT loop optimizations, layout sharding stability, distributed attention compatibility, and dramatic improvements to end-to-end latency (both compilation and post-processing).
🌟 Architectural & JAX Optimizations
flax.struct.dataclass(LTX2BlockContext), significantly improving code readability and type safety across attention mechanisms.guidance_scale,stg_scale) out ofstatic_argnamesand introduced boolean triggers (do_cfg,do_stg). This drastically minimizes recompilation overhead and tracing stalls._apply_ltx2_guidanceto heavily deduplicate and unify the computational paths for standard CFG, STG, and audio guidance logic, resulting in a significantly cleaner XLA graph.⚡ Caching, Loading, and Compilation (Ports of PRs #438 & #443)
try_load_converted_weights) for LTX2. This safely caches converted weights to disk, drastically accelerating model loading on subsequent runs by bypassing redundant conversions.torch.loadand sequentialtorch2jaxroutines in favor of mapping.safetensorsconcurrently with aThreadPoolExecutor. Memory ingestion is aggressively parallelized using zero-copy reinterpretation (tensor.view(torch.uint16).numpy().view(ml_dtypes.bfloat16)).@jax.jitfor@aot_cache.cached_jiton the main diffusion loop. Integrated 2-step fast-warmup logic (with aot_cache.warmup_mode()) which ensures the full step count is traced whenscan_diffusion_loop=True, bypassing prolonged compilations on application startup.jax.process_count() == 1), rather than enduring slow host-to-device PCIe copies for all devices.🚀 Post-Processing Latency
np_uint8pipeline output type and updated the YAML configs to default to it. This entirely bypasses the massive CPU bottleneck of sequentially instantiatingPIL.Imageobjects across 121 frames.uint8orfloat32) as vectorized JAX operations natively on the TPU prior to host transfer, we eliminate massive host-side memory thrashing and CPU iteration overhead.🌀 Distributed Attention
ring_attentionandulysses_attentionparameters down through allLTX2VideoTransformerBlockinstantiations (affecting bothnnx.scanand un-scanned eager block initializations).NNXAttentionOpin LTX2 cross-attention mechanisms, ensuring layout rules bind properly when Ulysses sharding configurations are enabled.📊 Results
DiT Execution Time:
11.2s➔9.6s7.8s➔5.6sVideo Post-Processing Time:
0.200s➔0.040s0.500s➔0.052s