[Johan-de-R]

Follows #55 (merged). Also carries #55 review nit 1: a malformed --warmup-resolutions / --warmup-num-frames now raises an error naming the flag (925595e).

What

Three optional flags, unset = today's behaviour (CacheDiTConfig() with VisualGen defaults):

• --cache-dit-residual-diff-threshold (env DYN_TRTLLM_CACHE_DIT_RESIDUAL_DIFF_THRESHOLD)
• --cache-dit-max-continuous-cached-steps (env DYN_TRTLLM_CACHE_DIT_MAX_CONTINUOUS_CACHED_STEPS)
• --cache-dit-low-noise-max-cached-steps (env DYN_TRTLLM_CACHE_DIT_LOW_NOISE_MAX_CACHED_STEPS): Wan 2.2's low-noise expert budget, hard-coded to 20 in VisualGen. Set on VisualGen's module constant before load; a TRT-LLM without it raises.

Why

In production the low-noise expert reuses exactly 20 of its 34 steps: the budget, not the similarity test, forces the last ~4 full computes. Measured on 1x B200 with the #55 image (6 prompts, fixed seeds):

Setting	s/clip	LPIPS vs #55 output	LPIPS vs full 50-step output
#55 (budget 20, threshold 0.24)	223-229	—	0.160
budget 24, threshold 0.32	188.8	0.066	0.165
budget 26, threshold 0.32, max-continuous 4	168.6	0.088	0.177

Two identical runs with separate compile caches differ by ~0.08 LPIPS. 12-prompt side-by-side review on the shipped image: budget 24 = 185.6 s/clip, LPIPS 0.018 vs production (11/12 clips <= 0.02); Johan saw one neutral difference. Rolling out to Wan2.2-T2V-A14B with --cache-dit-low-noise-max-cached-steps=24 --cache-dit-residual-diff-threshold=0.32 (image ...-swenc-abe8c4ee, built from abe8c4e).

Tests

TestCacheDitOverrides: unset flags call CacheDiTConfig() with no kwargs and leave the budget at 20; set flags reach CacheDiTConfig and the module constant; unknown TRT-LLM raises. 38/38 non-handler tests pass in the #55 image (handler tests need pytest_asyncio, absent from the runtime image).

🤖 Generated with Claude Code

[temirulan]

Approving. I checked the change against upstream TRT-LLM v1.3.0rc14 source, and it does what the description says.

What I verified

• The low-noise budget override takes effect. _WAN22_LOW_NOISE_MAX_CACHED_STEPS is a plain module global in cache/cache_dit_enablers.py. enable_cache_dit_for_wan() reads it when it is called, in the second ParamsModifier. It is not bound as a default argument or imported by name elsewhere. _build_diffusion_args() sets it before PipelineLoader runs post_load_weights(), so the dual-expert BlockAdapter picks up the new value. The hasattr guard means a TRT-LLM that renames the constant fails loudly instead of silently ignoring the flag.
• The threshold and max-continuous flags map to real CacheDiTConfig fields. These are residual_diff_threshold (ge=0, default 0.24) and max_continuous_cached_steps (ge=-1, default 3). Pydantic still validates both. They feed the shared DBCacheConfig, so they apply to both experts.
• Unset flags keep today's behaviour. Each override is dropped when it is None, so the default path is still CacheDiTConfig() with no kwargs, and the module constant stays at 20.
• The #55 nit is fixed. Malformed warmup shapes now raise an error that names the flag, and a trailing comma (81,) is rejected too.

Non-blocking nits

1. --cache-dit-low-noise-max-cached-steps is only applied inside the cache_backend == "cache_dit" branch. It is also only read by the Wan 2.2 dual-transformer enabler. Setting it with another backend, or on Wan 2.1, FLUX or LTX2, is silently ignored. A warning would keep a misconfigured model-CD arg from going unnoticed.
2. Unlike the other two flags, the budget has no range check: -1 means unlimited in cache_dit, and other negative values pass through as they are. A >= -1 check would match CacheDiTConfig's convention.

On quality: an LPIPS of 0.066 against #55 output is inside the ~0.08 run-to-run noise you measured, and the 12-clip review at budget 24 backs that up. The rollout values look reasonable.

I couldn't run the tests locally because they need the dynamo.llm Rust extension. I'm relying on your 38/38 result from the #55 image.