[None][feat] LTX-2: run v2a cross-attention under distributed Attention2D - #16301
Draft
luyiyun1021 wants to merge 2 commits into
Draft
[None][feat] LTX-2: run v2a cross-attention under distributed Attention2D#16301luyiyun1021 wants to merge 2 commits into
luyiyun1021 wants to merge 2 commits into
Conversation
…on2D When Attention2D (attn2d) is active, let the v2a cross-attention (the only sequence-parallel-enabled cross-attn direction) run under distributed Attention2D + Ulysses instead of all-gathering the full video K/V. An `attn2d_active` flag gates `enable_sp`, the Ulysses dual-attn wiring, and `use_ulysses`; a2v and text cross-attn keep sequence parallelism off (unchanged). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
…fallback Ulysses (head-sharding) is orthogonal to CP (sequence-sharding), so it composes with pure-Ulysses, Attention2D and async for SEPARATE_QKV cross-attn too; Ring + SEPARATE_QKV is the only unsupported combination and is already hard-errored below. The (qkv != SEPARATE_QKV or async_ulysses or cp_size == 1 or attn2d_active) inner condition was therefore a tautology for every config that survives that check, and the not-use_ulysses all-gather fallback was unreachable. Collapse the gate to use_ulysses = ulysses_size > 1 and enable_sequence_parallel and drop the redundant cp_size/attn2d_active locals, the dead fallback + debug log, and the now-unused logger import. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
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.
Description
When Attention2D (
attn2d) is active, let the v2a cross-attention — the only sequence-parallel-enabled cross-attn direction — run under distributed Attention2D + Ulysses instead of all-gathering the full video K/V. Anattn2d_activeflag opens three gates inLTX2Attention/Attention:enable_spfor cross-attn, the Ulysses dual-attn wiring, anduse_ulysses.a2vand text cross-attn keep sequence parallelism off (unchanged). Extends the v2a-Ulysses path (#15303) to the Attention2D mesh.Test Coverage
Structurally exercised on 16-GPU LTX-2 (cfg2 × attn2d). End-to-end numerical validation is still pending — opening as draft.
Notes
Orthogonal to
async_ulysses(a separate flag, off here). The config validator only forbidsasync_ulysses+ring_size>1, notattn2d; whether the async fast-path composes correctly with the Attention2D wrapper is a separate question and not exercised by this change.