[Johan-de-R]

What

• --cache-dit-max-warmup-steps (env DYN_TRTLLM_CACHE_DIT_MAX_WARMUP_STEPS): passed to CacheDiTConfig(max_warmup_steps=…), the leading steps cache_dit always computes before it starts reusing. On Wan 2.2 A14B this is the high-noise expert's warm-up (the low-noise expert keeps VisualGen's own). Unset keeps VisualGen's default (4).

• Per-request guidance_scale_2: the video handler now reads it from the request's nvext (the frontend already carries the field) and falls back to --default-guidance-scale-2 as before. guidance_scale already worked this way.

• --flow-shift-by-steps (env DYN_TRTLLM_FLOW_SHIFT_BY_STEPS), e.g. 27:20,40:12: each request runs with the shift mapped to its step count. The scheduler is rebuilt only when the shift changes, and an unmapped step count is refused (ValueError) rather than run on an untested schedule. Duplicate step counts and shifts ≤ 0 fail worker init. The swap is safe because the video handler serializes generate() behind its lock.

Why the shift must follow the step count

Saturation tracks how many steps the high-noise expert gets (6 review prompts, mean HSV saturation; the reviewed look is 40 steps / shift 12):

Steps / shift	High-noise steps	Saturation	Brightness	s/clip
50 / 3 (old production)	16	0.585	0.319	225
40 / 12 (reviewed look)	26	0.263	0.362	193
27 / 12	18	0.389	0.321	160
27 / 20 (picked in review)	21	0.327	0.345	162
27 / 30	22	0.393	0.338	152
27 / 8	15	0.537	0.319	162

So one fixed shift can't serve both 27 and 40 steps. deepapi offers exactly these two step counts.

Why

Under Wan's official A14B settings (40 steps, shift 12) the high-noise expert gets 26 of 40 steps, and with the default warm-up of 4 cache_dit reuses the steps that decide layout. In review that produced an extra milk jug, extra hands and impossible legs. LPIPS against the uncached render, 1xB200, 720p/81 frames:

High-noise steps computed up front	s/clip	LPIPS vs uncached
4 (default)	163	0.33
6	183	0.22 (moving cars lost)
8	193	0.17
10	213	0.16
uncached	412	0

8 is the chosen setting; reviewed by eye. It holds at 27 steps too (160 s/clip, 0.15 vs uncached 27 steps). One server alternating 27/40/50-step requests ran each at its steady-state time, so the step count does not trigger recompiles.

Per-request dual guidance lets deepapi expose both guidance values (e.g. Artificial Analysis sends 3.5 / 4).

Tests

TestCacheDitOverrides covers the new override; test_request_guidance_overrides_config_defaults covers the handler. 84/84 pass in the …-swenc-ea1a31ba image, including the async handler tests, which need pytest-asyncio (the image doesn't ship it, so these tests didn't run before). Mutation-checked: removing the override line, or restoring the config-only guidance_scale_2, each fails exactly its test. TestFlowShiftBySteps covers parsing, switching only on change, and refusing unmapped step counts (mutation-checked: dropping the per-request call or the change guard fails it).

Not enabled by this PR; the Wan config gets --cache-dit-max-warmup-steps=8 --flow-shift-by-steps=27:20,40:12,50:12 with the rollout; 50:12 covers the window before deepapi stops sending 50 steps and is dropped after (see the merge gate in deepinfra/backend#6590).

🤖 Generated with Claude Code

[Thachnh]

Sauron · Review of #58 together with deepinfra/backend#6590

Verdict: the code looks right. The rollout order in backend#6590 is broken because of how this PR rejects step counts. Fix the order before step 2 runs.

Blocking (rollout, not code): unmapped step counts fail every request during the config-switch window

_use_flow_shift_for raises ValueError for any step count missing from the map (diffusion_engine.py, in generate()). The merge gate in backend#6590 says:

2. Model config: new image, --flow-shift-by-steps=27:20,40:12 … Until deepapi deploys, the worker still gets 50 steps (226 s at the new schedule).
3. Then merge this PR and deploy deepapi.

But deepapi on main hardcodes Wan22T2VIn.num_inference_steps = 50 and always sends it in nvext. Between step 2 and step 3, every Wan2.2 request becomes --flow-shift-by-steps has no shift for 50 steps. It does not run at 226 s. The worker's own fallback, --default-num-inference-steps=50 (baked in redis_model_manager.py), hits the same refusal.

Fix: in step 2, map 50 as well, e.g. --flow-shift-by-steps=27:20,40:12,50:12 (or 50:3 if you want the old look until cut-over). Drop 50: in a step 4, after deepapi deploys. Separately, the baked --default-num-inference-steps should move to 27 (I left a comment on backend#6590).

Checked and fine

• Scheduler swap is serialized. engine.generate runs only inside async with self._generate_lock → asyncio.to_thread, so apply_flow_shift replacing pipeline.scheduler can't race another denoise loop.
• Warm-up is unaffected. VisualGen's warm-up runs inside loader.load(), not through DiffusionEngine.generate, so the step-count map can't fail worker init. With no --flow-shift, _scheduler_shift starts as None and the first request always applies its mapped shift.
• guidance_scale_2 is carried end to end. It's in VideoNvExt (common/protocols/video_protocol.py:46) and in the Rust frontend nvext (lib/llm/src/protocols/openai/videos/nvext.rs:62), so the field survives the frontend. Falling back to --default-guidance-scale-2 keeps today's behaviour for callers that don't send it.
• max_warmup_steps is passed through like the other cache_dit overrides, and None keeps VisualGen's default.

Nits (non-blocking)

• parse_flow_shift_by_steps("27:20,27:12") silently keeps the last value, and it accepts 27:0 or negative shifts. A one-line duplicate check and a shift > 0 check would turn an operator typo into an init error instead of a quiet schedule change.
• An unmapped step count raises a ValueError from inside the engine. That's fine while deepapi is the only caller and enforces Literal[27, 40], but the handler returns it as a generic generation failure. If a non-deepapi caller ever appears, a 4xx-shaped error would be kinder.

Not done: I did not run the test suite (no TRT-LLM image here). The review is from reading the diff plus the surrounding handler/engine/protocol code at cf702bbb.

[temirulan]

Sauron · Second pass at ea9bd1fe. Not approving yet.

The only new commit since my last review is ea9bd1fe (black formatting, no behaviour change), and it looks fine. The one blocking item from the first pass is still open. It is a rollout issue, not a code defect in this PR:

• Blocking (rollout): _use_flow_shift_for() raises ValueError for any step count missing from the map. Deepapi on main still hardcodes 50 steps for Wan2.2 (Wan22T2VIn.num_inference_steps returns 50 in deepinfra/deepapi/inf_models/ttv.py). Step 2 of backend#6590's merge gate sets --flow-shift-by-steps=27:20,40:12, so from then until deepapi deploys, every Wan2.2 request fails. The PR text assumes those requests keep running at 50 steps / 226 s. Two ways to fix it: add 50:12 (or 50:3) to the map for that window and drop it once deepapi is live, or reorder the gate.
• Nit (unchanged, non-blocking): parse_flow_shift_by_steps silently accepts duplicate step keys (last one wins) and non-positive shifts.

I'll approve once the rollout plan in backend#6590 (or the map used during the transition) no longer breaks the 50-step window.

[Johan-de-R]

Thanks. Both items are addressed:

• Blocking (rollout): agreed, it was a real gap. The rollout map is now 27:20,40:12,50:12 until deepapi stops sending 50 steps; then 50:12 is dropped. That's in the merge gate of deepinfra/backend#6590 and in this PR's description. The worker's baked --default-num-inference-steps also moved to 27 in backend#6590 (4cc189d).
• Nit, now fixed (26d503b): duplicate step counts and shifts ≤ 0 fail worker init. Covered by test_parse_rejects_malformed (3 new cases). Mutation-checked: removing the check fails exactly those 3. 84/84 pass in the …-ea1a31ba image.
• Unmapped steps returning a generic failure rather than a 4xx: left as is for now, since deepapi is the only caller and enforces 27/40. Worth revisiting if another caller appears.

[temirulan]

Sauron · Third pass at 26d503b9. Approving. Everything from the earlier rounds is resolved:

• Rollout: the transition map is now 27:20,40:12,50:12. The worker keeps serving the 50 steps that deepapi on main sends, and 50:12 is dropped after backend#6590 deploys. I traced each phase of the merge gate and every step count deepapi sends is mapped in each one. That includes the baked worker default (50 before backend#6590 merges, 27 after).
• Parser: duplicate step counts and shifts ≤ 0 now fail at init. I checked this against the edge cases (27:20,27:12, 27:0, 27:-3 and the malformed forms all raise).
• Tiny non-blocking nit: 27:nan and 27:inf still parse, because float() accepts them and nan <= 0 is False. A math.isfinite(shift) check would close that off if you want it, but it's an operator-typo edge case and not worth another round on its own.

Agreed on leaving unmapped step counts as a generic failure while deepapi is the only caller.