[Johan-de-R]

What

• --flow-shift (env DYN_TRTLLM_FLOW_SHIFT): after the pipeline loads, rebuild its scheduler from the same config with the shift replaced, under whichever of flow_shift (UniPC) or shift (FlowMatchEulerDiscreteScheduler) the scheduler's __init__ accepts. Choosing by __init__ rather than by config presence matters because diffusers keeps unknown keys, and Wan's FlowMatchEuler configs carry an ignored flow_shift beside shift. Unset keeps the checkpoint's scheduler config. A scheduler with neither key raises.
• feat: expose cache_dit threshold / max-continuous / Wan 2.2 low-noise budget #56 review nits: --cache-dit-low-noise-max-cached-steps rejects values below -1 on every cache backend (-1 = unlimited, as in CacheDiTConfig), and logs a warning when it cannot apply (backend is not cache_dit, or the pipeline has no transformer_2).

Why

Wan2.2-T2V-A14B-Diffusers ships a UniPC scheduler with flow_shift: 3.0, which VisualGen uses as-is. Wan's official A14B T2V config (wan_t2v_A14B.py) uses sample_steps = 40, sample_shift = 12.0. The shift decides how the steps split across the two experts at boundary 0.875: production gives the high-noise expert 16 of 50 steps; the official settings give it 26 of 40.

Measured on 1x B200 (1280x720, 81 frames, guidance 4.0/3.0, cache_dit, image from #55), 12 public prompts:

	Steps	Shift	s/clip
Production now	50	3	225
Wan official	40	12	163

Side-by-side review: official looked better in 10 of 12 clips, generally much less oversaturated; two clips regressed (latte: an extra milk jug and a third hand; ballerina: an impossible leg movement). An uncached official render of those two is running to tell the settings from cache_dit's reuse (17 of 26 high-noise steps reused at shift 12).

Not enabled by this PR. Shipping the official settings also needs deepapi to send 40 steps (Wan22T2VIn.num_inference_steps is hard-coded to 50), and dropping the budget-24 / threshold-0.32 cache args.

Tests

TestFlowShift: the shift key is replaced and the rest of the config kept (UniPC flow_shift and FlowMatchEuler shift); a FlowMatchEuler config with a stray flow_shift still gets shift changed; a scheduler without either raises; budget below -1 raises with cache_dit and with no cache backend (mutation-checked). 49/49 non-handler tests pass in the #56 image; against real diffusers, UniPC ends at flow_shift 12 and FlowMatchEuler {shift: 5, flow_shift: 3} at shift 12.

🤖 Generated with Claude Code

[temirulan]

Approving. The production target, UniPC on Wan2.2-T2V-A14B, works correctly. There is one latent bug on another scheduler type that's worth a quick follow-up (inline).

What I verified against TRT-LLM v1.3.0rc14

• Swapping the scheduler after load is safe. The Wan T2V and I2V pipelines only read self.scheduler per request: set_timesteps, config.num_train_timesteps for the expert boundary, and the denoise loop. cache_dit's split_wan22_inference_steps() also reads pipeline.scheduler at refresh time, so the per-expert step split follows the new shift. Warmup runs with the old scheduler, but the timesteps don't change compiled shapes, so that's fine.
• The UniPC path rebuilds correctly. type(s).from_config(dict(s.config), flow_shift=x) is what VisualGen itself uses to build the scheduler, and every other config key is kept.
• The two #56 follow-ups are fixed: the budget must be >= -1, and there's now a warning when the flag can't apply (wrong backend, or no transformer_2).

Finding (inline, should fix in a follow-up)

1. On a FlowMatchEulerDiscreteScheduler whose config also has flow_shift, the PR overrides flow_shift. FlowMatchEuler ignores that key, so shift stays the same and the log falsely reports the change. I reproduced it with diffusers. VisualGen's Wan loader explicitly expects such configs. The fix is to pick the key the scheduler's __init__ accepts, or to check the result after the rebuild.

Nit

• The low-noise range check only runs on the cache_dit path. With another backend, -5 passes and only produces the "ignored" warning. That's harmless, so this is just a note.

On rollout: as the description says, the 10/12 quality result also depends on deepapi sending 40 steps and on retuning the cache args. Shift 12 with 50 steps and the budget-24 settings was not measured, so I wouldn't flip --flow-shift alone.

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

[temirulan]

--flow-shift can silently do nothing on a FlowMatchEulerDiscreteScheduler whose config also has a flow_shift key.

diffusers keeps unexpected config attributes in scheduler.config; it logs "will be ignored" but still stores them. VisualGen's Wan loader expects FlowMatchEuler configs to carry a stray flow_shift: pipeline_wan.py:249-250 in rc14 copies it into shift. With that config, keys[0] is flow_shift, from_config ignores it again, and shift never changes. The info log still prints flow_shift 3.0 -> 12.0.

Repro with diffusers:

s = FlowMatchEulerDiscreteScheduler.from_config({"shift": 5.0, "flow_shift": 3.0})
config = dict(s.config)  # still contains flow_shift
n = type(s).from_config(config, flow_shift=12.0)
n.config.shift  # 5.0, unchanged

This doesn't affect the Wan2.2-T2V-A14B rollout, because UniPC takes flow_shift as a real init arg. Suggested fix: choose the key the scheduler's __init__ accepts ("flow_shift" in inspect.signature(type(scheduler).__init__).parameters, else "shift"). Alternatively, after the rebuild, assert that the new scheduler's config holds shift under the chosen key and raise if not. A test with both keys on a FlowMatchEuler-like fake would cover it.

[temirulan]

Fixed in 41cfbaf. I checked it with real diffusers. inspect.signature sees through @register_to_config, so FlowMatchEulerDiscreteScheduler exposes shift but not flow_shift. For a config with both keys, the code now picks shift, and the rebuilt scheduler has shift == 12.0. Upstream UniPCMultistepScheduler.__init__ takes flow_shift and no shift, so the A14B path picks flow_shift as before. The new test_stray_flow_shift_key_does_not_hide_the_real_shift covers this case.

[temirulan]

Round 2: still approving at 450251e. Both follow-ups look correct.

• 41cfbaf (flow-shift key choice): I checked it against real diffusers. inspect.signature(type(s).__init__) sees through @register_to_config, so the code now finds the parameter each scheduler actually accepts. For FlowMatchEuler that is shift; for current upstream UniPC it is flow_shift. My repro, a FlowMatchEuler with both shift and a stray flow_shift, now ends up with shift == 12.0. config.get() in the log line also stops a KeyError when the key is accepted but missing from the config.
• 450251e (budget check): the >= -1 check now runs at the top of _build_diffusion_args() for every cache backend. The parametrized test covers cache_dit and none.

No new findings. I couldn't run the tests locally because they need the dynamo.llm Rust extension.