[TRTLLM-12963][refactor] LTX-2 attention: drop dead k_pe parameter; require cached cross-attn - #14555
Conversation
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThe PR removes the redundant ChangesRemoving k_pe parameter and standardizing on pre_projected_kv
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/unittest/_torch/visual_gen/test_ltx2_attention.py (1)
193-199: ⚡ Quick winAdd a negative test for the newly enforced uncached cross-attention invariant.
Line 197 and Line 238 validate the new cached cross-attention path, but there’s no assertion that the old uncached pattern (
attn(x, context=ctx, ...)) now fails. Please add one failure-mode test (e.g.,assertRaisesRegex) so this contract change can’t silently regress.Proposed test addition
+ `@pytest.mark.skipif`(not torch.cuda.is_available(), reason="CUDA not available") + def test_cross_attention_uncached_context_rejected(self): + """Uncached cross-attention must be rejected after invariant tightening.""" + from tensorrt_llm._torch.visual_gen.models.ltx2.transformer_ltx2 import LTX2Attention + + batch_size = 1 + q_seq = 8 + kv_seq = 4 + query_dim = 4096 + context_dim = 4096 + heads = 32 + head_dim = 128 + dtype = torch.bfloat16 + + config = _create_config("VANILLA") + attn = ( + LTX2Attention( + query_dim=query_dim, + context_dim=context_dim, + heads=heads, + dim_head=head_dim, + config=config, + layer_idx=0, + ) + .to(self.DEVICE, dtype=dtype) + .eval() + ) + + x = torch.randn(batch_size, q_seq, query_dim, device=self.DEVICE, dtype=dtype) * 0.02 + ctx = torch.randn(batch_size, kv_seq, context_dim, device=self.DEVICE, dtype=dtype) * 0.02 + + with self.assertRaisesRegex((AssertionError, ValueError), "self-attention|cross-attention|pre_projected_kv"): + _ = attn(x, context=ctx, pe=None)Also, because this PR only touches unit tests under
tests/unittest/, QA integration test-list updates undertests/integration/test_lists/qa/are unnecessary here.As per coding guidelines: “Coverage expectations: Assess whether new/changed tests cover happy path, important edge cases, and failure modes relevant to the feature or fix.”
Also applies to: 235-240
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/visual_gen/test_ltx2_attention.py` around lines 193 - 199, The test suite adds a cached cross-attention path but misses a negative test to ensure uncached cross-attention (calling attn(x, context=ctx, ...)) now fails; add an assertion using assertRaisesRegex around calling attn(x, context=ctx, pe=None) (or attn.forward) to confirm the new invariant is enforced, e.g., call project_kv to get (k, v) and assert that calling attn with context=ctx (uncached pattern) raises the expected error message; reference the attn object, its project_kv method, and the forward/__call__ invocation to locate where to insert the failure-mode test.tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py (1)
312-348: 💤 Low valueConsider adding the same assertion to
_forward_unfusedfor consistency.The fused path (lines 285-289) asserts
context is Nonewhenpre_projected_kvis None, but_forward_unfusedlacks this check. While the docstring at lines 330-331 documents this contract ("uncached path applies pe to both Q and K (matches self-attn behavior)") and production never enters this path, adding the same assertion would make the invariant explicit and catch test/ablation misuse early.🔧 Proposed fix for consistency
if pre_projected_kv is not None: k, v = pre_projected_kv q = self.to_q(x) if self.qk_norm: q = self.norm_q(q) else: + assert context is None, ( + "uncached _forward_unfused is only valid for self-attn " + "(context=None); for cross-attn pass pre_projected_kv " + "from self.project_kv(context, pe=<k_rope>)." + ) q, k, v = self.get_qkv(x, context) q, k = self.apply_qk_norm(q, k)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py` around lines 312 - 348, Add the same invariant check present in the fused path to _forward_unfused: when pre_projected_kv is None assert that context is None to make the uncached Q/K rotation contract explicit and fail fast; place this assertion near the start of _forward_unfused (before calling self.get_qkv or any Q/K processing) so that the function enforces the intended uncached-path precondition.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py`:
- Around line 312-348: Add the same invariant check present in the fused path to
_forward_unfused: when pre_projected_kv is None assert that context is None to
make the uncached Q/K rotation contract explicit and fail fast; place this
assertion near the start of _forward_unfused (before calling self.get_qkv or any
Q/K processing) so that the function enforces the intended uncached-path
precondition.
In `@tests/unittest/_torch/visual_gen/test_ltx2_attention.py`:
- Around line 193-199: The test suite adds a cached cross-attention path but
misses a negative test to ensure uncached cross-attention (calling attn(x,
context=ctx, ...)) now fails; add an assertion using assertRaisesRegex around
calling attn(x, context=ctx, pe=None) (or attn.forward) to confirm the new
invariant is enforced, e.g., call project_kv to get (k, v) and assert that
calling attn with context=ctx (uncached pattern) raises the expected error
message; reference the attn object, its project_kv method, and the
forward/__call__ invocation to locate where to insert the failure-mode test.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c8e8bb14-96ad-4328-a460-946e01203c00
📒 Files selected for processing (2)
tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.pytests/unittest/_torch/visual_gen/test_ltx2_attention.py
|
PR_Github #50298 [ run ] triggered by Bot. Commit: |
feae178 to
7553ce1
Compare
|
/bot run --disable-fail-fast |
444d309 to
033fc8c
Compare
|
PR_Github #50301 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #50298 [ run ] completed with state |
|
PR_Github #50303 [ run ] triggered by Bot. Commit: |
|
PR_Github #50301 [ run ] completed with state |
|
PR_Github #50303 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #50331 [ run ] triggered by Bot. Commit: |
|
PR_Github #50331 [ run ] completed with state |
1c9b55e to
3575972
Compare
|
/bot run --disable-fail-fast |
3575972 to
b3b5e22
Compare
|
/bot run --disable-fail-fast |
b3b5e22 to
462f662
Compare
462f662 to
094e249
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #52932 [ run ] triggered by Bot. Commit: |
|
PR_Github #52932 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53082 [ run ] triggered by Bot. Commit: |
|
PR_Github #53082 [ run ] completed with state
|
094e249 to
11e4be3
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53097 [ run ] triggered by Bot. Commit: |
|
PR_Github #53097 [ run ] completed with state
|
11e4be3 to
b0cdceb
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53224 [ run ] triggered by Bot. Commit: |
|
PR_Github #53224 [ run ] completed with state
|
…ached cross-attn k_pe was the optional per-K rope override on LTX2Attention.forward and _forward_unfused. It became dead code after the earlier norm+rope fusion work moved K-rope into project_kv. That move had two goals: (1) eliminate the cos/sin all-gather under Ulysses (rope is per-token element-wise and commutes with seq-dim concat, so rotating K on the local shard before all-gather is bit-identical and saves a collective); (2) enable K-norm + K-rope to go through a single fused kernel (fused_dit_split_norm_rope) inside project_kv. After that, every production cross-attn caller pre-rotates K via project_kv and calls forward with pre_projected_kv=(k, v) and k_pe=None. Cleanup: - Drop k_pe parameter from forward + _forward_unfused, drop k_pe=None kwargs from BasicAVTransformerBlock's a2v / v2a calls. - Replace the k_pe-ternary 'k_pe if k_pe is not None else pe' in the uncached SEPARATE_QKV branch with just pe (it always was pe in practice; non-None k_pe was never wired). - Enforce a clean cached-cross-attn invariant: cross-attn (context != None) requires pre_projected_kv. The uncached SEPARATE_QKV branch is now self-attn-only (context must be None; Q and K share pe by construction). Cross-attn through the uncached branch raises ValueError directing the caller to project_kv + pre_projected_kv (ValueError survives python -O). In LTX-2 production every cross-attn caller already uses the cached pattern, so this is a no-op there; it just kills a silent-mis-rotation footgun and the documented-but-unreachable inline fallback. - Update BasicAVTransformerBlock.forward docstring (text_kv_video / text_kv_audio are required when their stream runs cross-attn, not 'falls back to inline computation'). - Update class + project_kv docstrings to drop k_pe references. - Update TestLTX2CrossAttention unit tests to use the production cached pattern (project_kv → pre_projected_kv) instead of the now-rejected uncached cross-attn shape. Tests: test_ltx2_attention.py (6/6) + test_ltx2_transformer.py (14/14) — 20/20 PASS. No production behavior change. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
b0cdceb to
6df8849
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53335 [ ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #53471 [ run ] triggered by Bot. Commit: |
|
PR_Github #53471 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53766 [ run ] triggered by Bot. Commit: |
|
PR_Github #53766 [ run ] completed with state |
@coderabbitai summary
Description
Drop the now-dead
k_peparameter fromLTX2Attention.forwardand_forward_unfused, and enforce a clean cached-cross-attn invariant: cross-attn requirespre_projected_kvfromproject_kv. The uncached SEPARATE_QKV branch becomes self-attn-only.Why
k_pewas the optional per-K rope override onLTX2Attention.forward— it let cross-attn pass a different rope for K than for Q. It became dead code after the earlier norm+rope fusion work moved K-rope intoproject_kv. That move had two goals:fused_dit_split_norm_rope) insideproject_kv.After that, every production cross-attn caller pre-rotates K via
project_kvand then callsforwardwithpre_projected_kv=(k, v)andk_pe=None. Grep confirms zero non-Nonek_pecall sites. The "inline fallback whentext_kv_video/text_kv_audiois None" path documented onBasicAVTransformerBlock.forwardis also unreachable in production:prepare_text_cachealways populates the KV list when its modality is enabled, and the block only invokes the cross-attn when the corresponding modality is enabled — so the cached KV is always present.The old uncached SEPARATE_QKV cross-attn branch additionally had a latent silent footgun:
k_pe if k_pe is not None else peapplied Q's rope to K when both came from different sources with potentially different sequence lengths.What changed (+40 / −44, 2 files)
tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py:k_pefromLTX2Attention.forwardand_forward_unfused.k_pe=Nonekwargs fromBasicAVTransformerBlock's a2v / v2a call sites.k_pe if k_pe is not None else pein the uncached SEPARATE_QKV branch with justpe.if context is not None: raise ValueError(...)(survivespython -O, unlikeassert). Cross-attn must usepre_projected_kvfromproject_kv. Same invariant in_forward_unfused.BasicAVTransformerBlock.forwarddocstring:text_kv_video/text_kv_audioare required when the corresponding stream runs cross-attn (built byLTXModel.prepare_text_cache) — drop the "falls back to inline computation" wording.project_kvdocstrings to dropk_pereferences.tests/unittest/_torch/visual_gen/test_ltx2_attention.py:TestLTX2CrossAttention.test_cross_attention_sanityandtest_cross_attention_different_dimsnow use the production cached pattern:k, v = attn.project_kv(ctx, pe=...)→attn(x, pre_projected_kv=(k, v), pe=...). The previousattn(x, context=ctx, pe=None)shape was exercising the unreachable inline-fallback path and is now rejected.Test Coverage
Local on B200, wt-3 container:
tests/unittest/_torch/visual_gen/test_ltx2_attention.py— 6/6 PASStests/unittest/_torch/visual_gen/test_ltx2_transformer.py— 14/14 PASSNo production behavior change.
PR Checklist
PR description clearly explains what and why.
PR Follows TRT-LLM CODING GUIDELINES.
Test cases are provided for new code paths.
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if significant design change.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.