[Johan-de-R]

What

1. Compile the transformer blocks in place when cache_dit is on and TRT-LLM < 1.3.0rc23. On those versions PipelineLoader attaches cache_dit (in post_load_weights) before torch_compile, and torch_compile swaps a new ModuleList of compiled blocks onto the transformer. cache_dit's BlockAdapter keeps the original list, so every computed denoising step runs eager. Upstream fixed the order in NVIDIA/TensorRT-LLM#16508 (first in rc23); our video image is on rc14. The engine now replaces BasePipeline.torch_compile with one that compiles each block into the existing list, only under those three conditions.
2. --warmup-resolutions / --warmup-num-frames (env DYN_TRTLLM_WARMUP_RESOLUTIONS / DYN_TRTLLM_WARMUP_NUM_FRAMES), mapped to VisualGen's CompilationConfig. Unset keeps the pipeline defaults. Wan's default warmup compiles {480x832, 720x1280} x {33, 81}; our API only serves 720x1280x81.

Measured (1x B200, Wan2.2-T2V-A14B, 1280x720, 81 frames, 50 steps, guidance 4.0/3.0 @ 0.875, cache_dit, prod image trtllm-visualgen:1.2.1-wan22t2v-cachedit-swenc-a9f16fe9)

	Production	This change
Computed step	12.04 s	10.40 s (dense compiled: 10.38 s)
Cached step	0.32 s	0.28 s
Per clip, 12 prompts	265.2 s	229.0 s (1.16x)
Steps skipped by cache_dit	9 + 20	identical schedule
Load + warmup	139.8 s	99.8 s with --warmup-resolutions 720x1280 --warmup-num-frames 81

Quality, 12 public prompts with fixed seeds: LPIPS vs the full 50-step uncached output 0.160 (production: 0.169), CLIP -0.3%. LPIPS vs production 0.124 mean: compiled kernels round differently from eager ones. Side-by-side review of all 12 pairs: no visible difference except one production artifact (convertible clip) that this output does not have.

Tests

• TestWarmupShapes: string parsing; configured shapes become CompilationConfig(resolutions=..., num_frames=...); unset keeps defaults.
• TestCompileBlocksInPlace: version gate (rc22 patched, rc23+/1.4 not); compiled blocks land in the list object cache_dit holds (fails if compile swaps in a new list, checked by mutation); fallback to the original method when a component has no block list.
• Ran in the prod image: 11/11 targeted tests pass. The 8-9 async handler tests fail there before and after this change (pytest_asyncio is not installed in the runtime image).

After review

• Speedup: 1.16x is the rc14 in-place patch measured in the harness (265.2 -> 229.0 s, 12 prompts); the built image trtllm-visualgen:1.2.1-wan22t2v-cachedit-swenc-43e67c0f measured 223 s (1.19x). In production since Oct 21, 4:06 pm PT.
• With warmup narrowed to 720x1280x81, any other shape recompiles on its first request; the API (Wan22T2VIn) only serves 720p landscape 81 frames.
• A clearer error for malformed --warmup-resolutions goes into feat: expose cache_dit threshold / max-continuous / Wan 2.2 low-noise budget #56 to keep this approved PR unchanged.

Rollout

New image = current prod image + these 3 files, then Wan2.2-T2V-A14B via model-CD with --warmup-resolutions=720x1280 --warmup-num-frames=81 added.

🤖 Generated with Claude Code

[temirulan]

Approving. I checked the change against upstream TRT-LLM v1.3.0rc14 source, and the fix is correct.

What I verified

• The bug is real, and rc23 is the right gate. In rc14, WanPipeline.post_load_weights() sets up cache_dit. cache_dit_enablers.py builds BlockAdapter(blocks=[pipeline.transformer.blocks, transformer_2.blocks]) from the original ModuleList objects. PipelineLoader then calls pipeline.torch_compile(), which does setattr(model, block_name, nn.ModuleList(compiled_blocks)), so cache_dit is left holding the uncompiled list. pipeline_loader.py in rc22 vs rc23 confirms _setup_cache_acceleration() moves to after torch_compile() starting in rc23.
• The replacement torch_compile behaves like upstream. It uses the same transformer_components (transformer and transformer_2 on Wan 2.2), the same _find_transformer_blocks, mode="default", dynamic=None and enable_fullgraph. The only difference is that blocks are assigned into the existing list. Any component without a block list (or missing) falls back to the original method. The gate can only misfire on dev builds like 1.3.0rc23.dev0, and that is harmless because compiling in place is correct whichever order the loader uses.
• The warmup flags match rc14's CompilationConfig. Both fields are Optional[...] = None, and resolve_warmup_plan() falls back to each model default separately. Setting only one of the two flags therefore behaves as expected, and leaving both unset passes no compilation kwarg.
• The measurements fit the mechanism. Computed steps take 10.40 s with cache_dit vs 10.38 s for dense compiled, so the cached path now runs the compiled blocks.

Non-blocking nits

1. A malformed --warmup-resolutions (e.g. 720x1280x81 or 720*1280) fails at startup with a bare ValueError: too many values to unpack. A short error message naming the flag would make a bad model-CD arg quicker to diagnose.
2. The title says 1.18x but the table shows 1.16x per clip; 1.16x is the per-step speedup on computed steps.
3. When rolling out, remember that any request outside 720x1280x81 will trigger a recompile on first hit once warmup is narrowed to that one shape. That's fine as long as the API keeps pinning the shape.

I couldn't run the test file locally because it needs the dynamo.llm Rust extension. I'm relying on your 11/11 result from the prod image.