[PR 2/5] LTX2: fix block_kv_compute_in fallback, framewise VAE tiling, v6e tile-search VMEM cap - #495
Closed
jitendra-jalwaniya wants to merge 1 commit into
Closed
jitendra-jalwaniya wants to merge 1 commit into
jitendra-jalwaniya wants to merge 1 commit into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces several configuration and logic updates for LTX2 video generation. Specifically, it adds VAE tiling and slicing options to the configuration files, caps the VMEM limit to 32MB on v6 JAX devices, refactors block size extraction in attention_flax.py to handle bkv_compute_in dynamically, and updates the VAE pipeline to toggle use_framewise_decoding alongside VAE tiling. There are no review comments, and I have no feedback to provide.
jitendra-jalwaniya
force-pushed
the
ltx2_tpu_fixes
branch
from
September 29, 2026 07:35
54af45c to
5d55cc9
Compare
jitendra-jalwaniya
changed the title
jitendra-jalwaniya
requested review from
Perseus14
and removed request for
entrpn
September 29, 2026 07:54
Perseus14
reviewed
Sep 29, 2026
| keys.get("tile_search_vmem_limit_bytes") or config.flash_block_sizes.get("vmem_limit_bytes") or 64 * 1024 * 1024 | ||
| ) | ||
| if "v6" in jax.devices()[0].device_kind.lower(): | ||
| vmem_limit_bytes = min(vmem_limit_bytes, 32 * 1024 * 1024) |
Collaborator
There was a problem hiding this comment.
v6e has ~128 MiB, why are we scoping it to 32 MiB
Collaborator
Author
There was a problem hiding this comment.
Thanks for highlighting this issue.
I had some confusion around XLA's scoped VMEM vs full VMEM.
Remove these changes.
jitendra-jalwaniya
force-pushed
the
ltx2_tpu_fixes
branch
from
September 29, 2026 10:36
5d55cc9 to
278ed61
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Small, independent LTX-2 fixes for TPU, split out of #489.
block_kv_compute_inwas set explicitly. It now falls back toblock_kv_computewhen unset, and is only clamped when it is given.-
ltx2_pipeline.py:enable_vae_tiling()/disable_vae_tiling()now also togglevae.use_framewise_decoding, so tiled decoding also decodes frame by frame.-
ltx2_video.yml/ltx2_3_video.yml: addenable_vae_tiling: Falseandenable_vae_slicing: Falseso the options are declared and can be overridden from the command line.-
generate_ltx2.py: on v6e, cap the tile-search VMEM budget at 32 MiB, down from the 64 MiB default. This keeps the grid search from picking tiles that don't fit v6e's scoped VMEM.