[TRTLLM-11707][feat] Add CUDA graph support (torch compile compatible) for LTX-2#12653
Conversation
📝 WalkthroughWalkthroughUpdates to the LTX2 model pipeline add caching for rope frequency computations, modify tensor handling in velocity calculations to preserve tensors during broadcasting, and introduce a CUDA graph runner subclass that manages cached graph capture and replay for the transformer with Modality tensor support. Changes
Sequence Diagram(s)sequenceDiagram
participant Pipeline as LTX2Pipeline
participant Runner as _LTX2CUDAGraphRunner
participant Transformer as Transformer
participant CUDA as CUDA Runtime
participant Cache as Static Buffers
Pipeline->>Runner: _setup_cuda_graphs()<br/>(wrap transformer.forward)
Note over Pipeline,Runner: Setup Phase
Pipeline->>Runner: forward(tensor_input,<br/>modality_input, ...)
Note over Runner: Check cache key<br/>(shapes + modality)
alt Cache Miss
Runner->>Runner: Clone inputs → static buffers
Runner->>CUDA: gc.collect()
Runner->>CUDA: torch.cuda.empty_cache()
Runner->>Transformer: Warmup capture
Runner->>CUDA: Synchronize
Runner->>Cache: Store graph + buffers<br/>in _graph_cache
else Cache Hit
Runner->>Cache: Retrieve cached graph
end
Note over Runner,Cache: Replay Phase
Runner->>Cache: Copy input → static buffers<br/>(in-place for Modality)
Runner->>CUDA: Launch replayed graph
Runner->>Runner: Extract output from cache
Runner->>Pipeline: Return result
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/rope.py (1)
186-189: Consider cache management for long-running processes.The module-level
_freq_grid_gpu_cachegrows unboundedly as different(generator, theta, n_dims, dim, device)combinations are used. For typical use this is fine (small number of configs), but for long-running services with varied configurations, this could accumulate GPU tensors.Consider adding a
clear_freq_cache()utility or using an LRU-bounded cache if memory pressure becomes a concern.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/rope.py` around lines 186 - 189, The module-level _freq_grid_gpu_cache stores GPU tensors keyed by (generator_id, theta, n_dims, dim, device) and can grow unbounded in long-running services; add a cache-management API such as a clear_freq_cache() function (and optionally a remove_freq_cache(key) helper) that frees GPU tensors and clears entries from _freq_grid_gpu_cache, and/or replace the dict with an LRU-bounded cache implementation to evict old entries; update any callers that rely on the cached values to use these helpers and ensure torch.cuda.synchronize() / .to('cpu') or .detach() is used appropriately when releasing GPU tensors so memory is actually freed.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/utils_ltx2.py`:
- Around line 18-26: The current implementation only checks sigma == 0 in the
non-tensor branch, so when sigma is a torch.Tensor with all-zero values the
division ((sample.to(calc_dtype) - denoised.to(calc_dtype)) / sigma) will
produce inf/nan; update the tensor branch in utils_ltx2.py to validate tensor
sigma before broadcasting (e.g., check torch.any(sigma == 0) or torch.all
depending on desired semantics) and raise a ValueError("Sigma can't be 0.0") if
a zero is detected, preserving the existing dtype conversions and the
broadcasting loop for sample.dim(); ensure the error check is done after moving
sigma to calc_dtype but before performing the division with sample and denoised.
---
Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/rope.py`:
- Around line 186-189: The module-level _freq_grid_gpu_cache stores GPU tensors
keyed by (generator_id, theta, n_dims, dim, device) and can grow unbounded in
long-running services; add a cache-management API such as a clear_freq_cache()
function (and optionally a remove_freq_cache(key) helper) that frees GPU tensors
and clears entries from _freq_grid_gpu_cache, and/or replace the dict with an
LRU-bounded cache implementation to evict old entries; update any callers that
rely on the cached values to use these helpers and ensure
torch.cuda.synchronize() / .to('cpu') or .detach() is used appropriately when
releasing GPU tensors so memory is actually freed.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 453e20c5-f035-4112-a44d-2f6e7cd935b9
📒 Files selected for processing (3)
tensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/rope.pytensorrt_llm/_torch/visual_gen/models/ltx2/ltx2_core/utils_ltx2.pytensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py
f9458db to
4762386
Compare
|
/bot run --disable-fail-fast |
luyiyun1021
changed the title
|
PR_Github #41161 [ run ] triggered by Bot. Commit: |
luyiyun1021
changed the title
luyiyun1021
changed the title
d41c719 to
51b9826
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #41198 [ run ] triggered by Bot. Commit: |
|
PR_Github #41198 [ run ] completed with state
|
51b9826 to
f4b199c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #41330 [ run ] triggered by Bot. Commit: |
|
PR_Github #41330 [ run ] completed with state
|
f4b199c to
7bedea5
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #41434 [ run ] triggered by Bot. Commit: |
|
PR_Github #41434 [ run ] completed with state
|
7bedea5 to
ec4c853
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #41619 [ run ] triggered by Bot. Commit: |
|
PR_Github #41619 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
Add Modality-aware CUDA graph capture/replay for the LTX-2 transformer, enabling monolithic graph capture of the full transformer forward including NCCL all-to-all collectives (Ulysses sequence parallelism). Key changes: - pipeline_ltx2.py: _LTX2CUDAGraphRunner subclass that handles Modality dataclass args (key derivation, static buffer clone/copy, None-tolerant output refs) and BatchedPerturbationConfig fingerprinting for STG. Override _setup_cuda_graphs to remove torch.compile mutual exclusion. - rope.py: GPU-side frequency grid cache to avoid CPU->GPU H2D transfer that blocks CUDA graph capture. - utils_ltx2.py: Keep sigma as tensor in to_velocity() to avoid .item() GPU->CPU synchronization. Compatible with torch.compile: graph capture happens during warmup after torch.compile's lazy compilation, so captured graphs contain optimized compiled kernels. Ref: NVIDIA#12603 Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com>
ec4c853 to
c21974f
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #42045 [ run ] triggered by Bot. Commit: |
|
PR_Github #42047 [ run ] triggered by Bot. Commit: |
|
PR_Github #42045 [ run ] completed with state |
|
PR_Github #42047 [ run ] completed with state
|
luyiyun1021
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #42299 [ run ] triggered by Bot. Commit: |
|
PR_Github #42299 [ run ] completed with state |
5 participants
@coderabbitai summary
Description
Add CUDA graph support for the LTX-2 diffusion transformer to reduce CPU dispatch overhead and eliminate inter-rank synchronization drift in multi-GPU Ulysses sequence parallelism.
Related PR: #12603 (alternative LTX-2 CUDA graph implementation)
Problem
In 8-GPU FP4 inference with Ulysses parallelism, each GPU processes 1/8 of the sequence. The resulting small kernels (~24us) are shorter than CPU dispatch interval (~35us), making kern/gap < 1. This exposes CPU jitter (~2us/launch) which accumulates into ~10-20ms rank drift per denoising step, degrading NCCL AllGather latency from 133us to 310us (2.3x).
Solution
Monolithic CUDA graph capture of the full
LTXModel.forward(), including NCCL all-to-all collectives. This eliminates all per-kernel CPU dispatch gaps, reducing rank drift from ~10-20ms to ~0.02ms per step.Key Changes
pipeline_ltx2.py—_LTX2CUDAGraphRunnersubclass ofCUDAGraphRunner:Modalitydataclass arguments (key derivation from nested tensor shapes, static buffer clone/copy)BatchedPerturbationConfigfingerprinting ensures different STG configurations get separate graph captures(video_vel, None)or(None, audio_vel)_setup_cuda_graphs()to remove the torch.compile mutual exclusion — both work togetherrope.py+transformer_args.py— GPU-side frequency grid cache:_generate_freq_grid_np/pytorch()returns CPU tensors vialru_cachefreq_grid_cacheparameter toprecompute_freqs_cis(): first call does H2D and caches the GPU tensor; subsequent calls (including insidetorch.cuda.graph()) reuse the cached tensorTransformerArgsPreprocessorinstance (lifecycle tied to pipeline, not global)visual_gen_ltx2.py— CLI support:--enable_cudagraphargument for the example scripttest_ltx2_transformer.py— CUDA graph correctness test:TestLTX2CUDAGraphCapture: verifies graph capture + replay produces bitwise-identical output to eager forward in AudioVideo mode (the most complex path with video/audio self-attention, bidirectional AV cross-attention, text cross-attention, and RoPE)How to Enable
Or via CLI:
python examples/visual_gen/visual_gen_ltx2.py --enable_cudagraph ...torch.compile Compatibility
The previous implementation had torch.compile and CUDA graphs as mutually exclusive. This PR removes that restriction:
torch.compileoperates at the block level (each of 48 transformer blocks compiled individually)Performance Results
8-GPU B200 FP4, 768x1280x121 frames, Ulysses=8:
Per-step synced GPU time (with
torch.cuda.synchronize(), steady-state step 20-40 of 40-step run):Rank drift (from nsys analysis):
The 4.1ms/step improvement comes from eliminating straggler CPU bubble and NCCL wait time. All 21 measured steady-state steps consistently showed improvement (-1.0 to -6.0ms).
Tests
TestLTX2CUDAGraphCapture::test_cuda_graph_audio_video_correctness— bitwise correctness for both capture and replay paths (AudioVideo mode)test_ltx2_transformer.pypass (no regression)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.