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
Perseus14
force-pushed
the
ltx2-improvements
branch
3 times, most recently
from
1f64f19 to
3a13196
Compare
Perseus14
force-pushed
the
ltx2-improvements
branch
from
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. |
1 participant
Description
This PR is a comprehensive refactor and optimization sweep. It brings massive improvements to XLA compilation times, memory usage (HBM), and architectural hygiene by stripping out redundant compute, unifying duplicated logic, and optimizing JAX tracing.
🧹 Architectural Hygiene & Code Quality
LTX2VideoTransformer3DModel.__call__. The 4 separate block execution paths (scan vs. loop, perturbation vs. no-perturbation) have been consolidated using a newTransformerContextcontainer and a singleapply_blockhelper function.prepare_video_coords: Deleted the wasteful 5Dlatent_coordsblock inattention_ltx2.pythat was computing an unused, wrongly-shaped tensor only to immediately overwrite it.apply_split_rotary_emb: Cleaned up the convoluted reshape/broadcast logic for split RoPE. Removed the redundantexpand_dimsandsqueezeoperations, executing the rotation directly (first_x * cos - second_x * sin) to avoid allocating unnecessary intermediate 5D tensors.hasattr(self, "rope_type")check inLTX2Attention.max_logging.logwarning when defaulting to a zero-seedjax.random.key(0)for noise generation.⚡ XLA & JAX Compilation Optimizations
guidance_scale,stg_scale,audio_guidance_scale, etc.) fromstatic_argnamesinrun_diffusion_loop(). Tweaking these generation scales will no longer trigger expensive 10-30 minute JAX recompilations!if guidance_rescale > 0:check inside the compiled diffusion loop withjax.lax.cond. This enables the CFG rescaling logic to be fully dynamic, complementing the removal of the static scales and fixing formulation inconsistencies.nnx.scanwith standardjax.lax.scanfor the primary denoising timestep loop to ensure predictable compilation.RuntimeProgramInputMismatchforscan_layers=False: Resolved an issue where XLA would fail during warmup compilation due to unrolled layer layout mismatches. Added explicit@jax.jitwrappers withjax.lax.with_sharding_constraintto enforce layout transpositions before crossing intorun_diffusion_loop.🧠 Memory (HBM) Optimizations
target_dtypeupfront and mapping it directly, avoiding a redundant double-casting pipeline that was passing throughbfloat16.