perf(hi3): implement the batched-step replay hook for HunyuanImage3

[leviking98z-rgb]

Stacked on Tencent-Hunyuan#299. Base that PR's branch, not main. The diff shown against main will include Tencent-Hunyuan#299's refactor; the changeset owned by this PR is the 5 files below. Merge Tencent-Hunyuan#299 first, then this.

Summary

Tencent-Hunyuan#299 introduces SingleStreamDiffusionRunner._replay_batched as an optional one-forward replay fast path, and the runner's serial loop is the default when it returns None. SD3 implements the hook; HI3 does not. So HunyuanImage3 still issues one transformer forward per selected SDE step.

That is worth fixing for HI3 specifically: its DiT is a large MoE, so every extra replay forward pays its own kernel launches, FSDP all-gather and activation recompute at a batch of only B — a shape that does not saturate the device.

This implements the hook. Replay is stateless (_step_kwargs returns no state outside mode="sample"), so the S selected transitions carry no cross-step dependency and can be stacked on the batch dim:

• sample / prev_sample → [S*B, C, H, W], step-major (rows [k*B:(k+1)*B] are all B samples at target[k])
• sigma / sigma_next → [S*B] vectors, since one batch now mixes different steps
• one step_with_logp → one transformer forward + one vectorized SDE transition
• log-probs reshape back to the public [B, S]

With num_sde_steps=2 that removes one of two serial forwards; the win scales with num_sde_steps.

The tiling is the non-obvious part, and the reason this isn't a copy of SD3's. The runner CFG-expands before replay (_prepare_conditions → _expand_cfg_for_forward), so the incoming fused batch is branch-major [cond_B, uncond_B] while the latents are step-major. predict_noise splits the fused batch at its midpoint into the cond/uncond halves, so the tiled layout must stay branch-major with each branch internally step-major:

[cond_step0_B, cond_step1_B, ..., uncond_step0_B, uncond_step1_B, ...]

Getting that nesting backwards is silent — wrong log-probs, no crash. See the mutation check in the Test Plan.

Two further details specific to HI3:

• rope_cache tiles like any other row-aligned field. In this stack it is a per-sample [B, 2, L, D] CONCAT tensor that predict_noise genuinely consumes (unbound to (cos, sin) at the model boundary), not the base's shared (cos, sin) tuple. Passing it through verbatim would feed the wrong rows.
• Conditional-image payloads are lists in the it2i path even though their generic annotations are tensors, so _tile handles tensor- and sequence-backed fields alike — mirroring what _expand_cfg_for_forward already does.

Guards fall back to the serial loop (return None) when: the flag is off; S <= 1; the strategy is not a stateless SDEStrategy (a stateful solver may consume step_index or carry cross-step state); or the guidance scale varies across the target steps — a per-step schedule cannot be folded into one forward, because the CFG combine is per-row and would need a per-row scale.

Because the π_old anchor is replayed through this same method, the on-policy ratio stays exactly 1.

Default is False on HunyuanImage3PipelineConfig; the two HI3 unified recipes opt in.

Related Issue

Depends on Tencent-Hunyuan#299. Extends the batched-step replay work in Tencent-Hunyuan#144 / Tencent-Hunyuan#156 to HunyuanImage3.

Test Plan

Run from repo root at 84506f9 (on top of Tencent-Hunyuan#299's f717a5c2); venv torch 2.13.0+cu130.

1. The property that matters — batched output equals serial output. A CPU fake-kernel harness whose per-row result deliberately mixes in the conditioning rows, so a wrong tiling permutation changes the numbers (uncommitted, per the no-tests/ policy):

PASS  cfg=1: batched returns a ReplayResult
PASS  cfg=1: log_probs shape [B,S]
PASS  cfg=1: log_probs == serial
PASS  cfg=1: prev_sample_means == serial
PASS  cfg=2 (branch-major): batched returns a ReplayResult
PASS  cfg=2 (branch-major): log_probs shape [B,S]
PASS  cfg=2 (branch-major): log_probs == serial
PASS  cfg=2 (branch-major): prev_sample_means == serial
PASS  flag off -> None (serial)
PASS  S==1 -> None (serial)
PASS  non-SDE strategy -> None (serial)
PASS  step-varying guidance -> None (serial)
PASS  tiled batch is S*cfg*B
PASS  tiling is branch-major with step-major inside
PASS  repeats=1 is identity
PASS  rope_cache is tiled (not shared verbatim)
PASS  fused_uncond cleared after tiling
PASS  repeats=0 rejected
PASS  repeats=-1 rejected

all checks passed

2. Mutation check — the harness has teeth. Swapping the tiling nesting order (step-major outer instead of branch-major outer, i.e. exactly the bug this layout guards against):

FAIL  cfg=2 (branch-major): log_probs == serial          max|d|=1.928e+04
FAIL  cfg=2 (branch-major): prev_sample_means == serial  max|d|=3.373e+04
FAIL  tiling is branch-major with step-major inside

So the 19 passes above are not vacuous. Reverted after the check.

3. Hook conforms to the base contract:

config has batch_replay_steps: True
stage __init__ has flag: True
_replay_batched overridden: True
_tile_conditions present: True
hook signature matches base: True -> ['self','conditions','segment','params','target','sigmas','sigma_max','device']

4. Recipe keys audited against the config dataclass:

hi3_vllmomni.yaml: batch_replay_steps=True unknown=[]
hi3_vllmomni_veomni_ep.yaml: batch_replay_steps=True unknown=[]
all recipe keys map to config fields

5. Repo gates:

$ ruff check unirl/models/hunyuan_image3/ examples/
All checks passed!
$ ruff format --check unirl/models/hunyuan_image3/
19 files already formatted
$ python lint/check_recipe_targets.py
check-recipe-targets: 2434 recipe _target_ paths resolve.
$ python lint/check_experimental_boundaries.py
check-experimental-boundaries: ok

(ruff check unirl/ repo-wide reports 4 pre-existing F841 in unirl/models/boogu_image/vendor/attention_processor.py, untouched by this PR.)

Not run: GPU real-checkpoint parity, and any throughput number on this base. See Reviewer Notes — this is the main gap and I would rather flag it than paper over it.

Compatibility / Risk

• No behavior change by default. batch_replay_steps defaults to False on the config; with it off, _replay_batched returns None and the runner's serial loop runs exactly as today.
• The two HI3 unified recipes flip it on. If you would rather land the mechanism inert and enable it in a follow-up after a GPU parity run, say so and I will set both to false.
• No config-schema break: the new field is additive with a False default.
• Only unirl/models/hunyuan_image3/* and the two HI3 recipes are touched. SD3's own implementation and the shared runner are untouched.
• The fallback guards mean an unsupported strategy or a step-varying guidance schedule silently uses the serial path rather than producing wrong numbers.
• Memory: the batched forward's activation peak is S× the serial one's, since the batch is S*B instead of B. At num_sde_steps=2 that is 2×. This is the real cost and I have not measured it on HI3 — a recipe that is already near its memory ceiling should keep the flag off until it has.

Verification

I reviewed the full diff and ran every command quoted above.

Duplicate-work check: scanned the open PRs. Tencent-Hunyuan#299 defines the hook but leaves HI3 unimplemented (_replay_batched appears 4× in sd3/diffusion.py, 0× in hunyuan_image3/diffusion.py on its head). Tencent-Hunyuan#156 extends batched replay to qwen_image / z_image / flux2_klein and does not touch HI3. Tencent-Hunyuan#144 is the original SD3 work. Nothing else touches unirl/models/hunyuan_image3/.

Reviewer Notes

The evidence gap, stated plainly. What I have verified is equivalence — batched replay produces the same log-probs and prev-sample means as the serial path, on a fake kernel, including a mutation check that the test would catch a wrong layout. What I have not verified on this base is (a) parity with a real HI3 checkpoint on GPU, and (b) any speedup number.

I do have a measurement of this mechanism from an earlier, pre-Tencent-Hunyuan#299 HI3 stack: training-phase time −24.2%, with per-phase timing attributing effectively all of it to the image backward (118.07s → 86.85s, −26.4%) while AR backward stayed flat (19.04s → 19.60s). I am deliberately not presenting that as evidence for this diff — different base, different conditions plumbing (fused_uncond was not split out then, and rope_cache was a shared tuple rather than a per-sample tensor), so the code paths differ in exactly the places this PR had to get right. Treat it as motivation for why the hook is worth implementing, not as a result.

Happy to hold this as draft until Tencent-Hunyuan#299 lands and I can run a real-checkpoint parity check (ratio=1.0000, log-prob max_abs_diff) plus a controlled A/B on the same 8×H20 shape Tencent-Hunyuan#299's own test plan uses.

Where I'd look first: _tile_conditions, specifically the branch-major/step-major nesting and the cfg_factor branch that decides whether a B-sized payload is duplicated per branch or split. That is where a subtle error would be silent. Second: whether the SDEStrategy + constant-guidance guard pair is sufficient for every HI3 strategy you expect in practice.