[mickqian]

Branch: mick/ltx2-prompt-kv-once-per-loop (on main 02b9836), commit 8037571

Motivation

LTX-2 re-projects the prompt in every block of every denoising step. The 1024-token caption goes through caption_projection (two GEMMs, video and audio), and then each block's two prompt cross-attentions (attn2, audio_attn2) project it to K and V and RMS-normalize K. Neither the text nor the DiT weights change inside a denoising loop: the LoRA phase switch runs before the loop. On 4xB300 (CFG parallel, BCG) these are ~3.4 ms of every 55 ms stage-1 step. TensorRT-LLM caches the same projections per request.

Modifications

• models/dits/ltx_2.py
  • LTX2Attention.project_context_kv(context) returns the normed K and the V that forward would project. forward takes context_kv=(k, v) and skips to_k / to_v / k_norm.
  • LTX2VideoTransformer3DModel.prepare_prompt_kv(encoder_hidden_states, audio_encoder_hidden_states) runs the caption projections and every block's prompt K/V once. It writes them into persistent per-block buffers keyed by the text shapes and returns that key. The forward takes cached_prompt_kv=<key>, skips the caption projections and hands each block its K/V. The buffers keep their addresses across requests, so a breakable CUDA graph captured at warmup reads the current request's K/V.
  • caches_prompt_kv is false when the prompt is adaLN-modulated per step (cross_attention_adaln, LTX-2.3), whose K/V depend on the step.
• ltx_2/denoising.py: on the official CFG path, _use_cached_ltx2_prompt_kv prepares the K/V on the loop's first step (LTX2DenoisingContext.prompt_kv_key) and passes the key to every step. It keeps the per-step projection under layerwise offload (weights stream per block) and cache-dit.

Tests

• test_ltx2_prompt_kv_cache.py
  • A 2-block LTX-2 forward with cached_prompt_kv is bitwise equal to the per-step projection.
  • A second text is written into the same buffers (same data pointers), and the cached forward matches the uncached one for it.
  • caches_prompt_kv is false with cross_attention_adaln.
  • The stage prepares once per loop and leaves the kwargs untouched for an LTX-2.3 DiT, under layerwise offload and under cache-dit.
• test_ltx2_bcg_coords.py: its stub DiT reads the prompt itself, so the step test stubs the new hook.

Accuracy Tests

4xB300 (2 GPUs used), LTX-2 two-stage 768x512x121, 40 steps, CFG parallel + breakable CUDA graphs. MP4s of 4 requests are byte-identical (md5) to main + #45965 / #45977 / #46047.

Speed Tests and Profiling

Same setup, SGLANG_DIFFUSION_SYNC_STAGE_PROFILING=1, median of 3 steady requests:

	stage 1 (40 steps)	per step
base	2.205 s	55.1 ms
this PR	2.062 s	51.5 ms

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ✅ Run #39985162727
Latest PR Test (Extra): ❌ Run #39985162307
Latest PR Test (AMD ROCm 10): ❌ Run #39985162880

[mickqian]

/tag-and-rerun-ci