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 1 commit 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. |
1bbd72c to
2bf389d
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. |
|
@mbohlool Can you please review this? |
2bf389d to
7731417
Compare
0b094c9 to
868e94d
Compare
ccf5589 to
9c77827
Compare
3 participants
LTX2 Improvements, Optimizations & Bug Fixes
This PR introduces a comprehensive overhaul to the LTX2 implementation in MaxDiffusion, targeting architectural cleanliness, layout sharding stability, distributed attention compatibility, and massive improvements to compilation and model-loading speeds.
🌟 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.⚡ Compilation & Loading Optimizations (Ports of PRs #438 & #443)
@jax.jitfor@aot_cache.cached_jiton the main diffusion operations. Included a 2-step fast-warmup logic (with aot_cache.warmup_mode()) ingenerate_ltx2.pythat fully bypasses the prolonged tracing compilations on application startup.torch.loadand sequentialtorch2jaxroutines inltx2_utils.py. The loader now concurrently maps chunks of.safetensorsusing aThreadPoolExecutorand uses a zero-copy reinterpretation technique (tensor.view(torch.uint16).numpy().view(ml_dtypes.bfloat16)) to aggressively parallelize memory ingestion.jax.process_count() == 1. Large tensors (>=64MB) are now pushed directly to device 0 and replicated over ICI rather than enduring slow host-to-device PCIe copies._load_and_init.🌀 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.