[LeslieWylie]

What does this PR do?

Add optional process-group tile decoding to AutoencoderKLMiniMaxH3. Pass an initialized group to the existing enable_tiling(process_group=...) API; each rank decodes a share of the original spatial tiles and receives the full result. Default decoding and encoding stay local, and temporal chunks, tile geometry, and blending stay unchanged.

This is a draft pending real CUDA/NCCL and performance validation. No isolated VAE speedup or full-pipeline quality improvement is claimed.

vae.enable_tiling(process_group=dist.group.WORLD)  # an initialized group or subgroup
with torch.inference_mode():
    decoded = vae.decode(latents).sample

Validation

On current main d42ebf9c6e0b46103b74283f5324357a67acaa04, using PyTorch 2.14 CPU and Gloo:

• The added real tiny-VAE multiprocess test passes for WORLD and a subgroup excluding global rank zero. It checks exact serial/parallel outputs, batch and temporal overlap, BF16 autocast output dtype, fewer tiles than ranks, and total decoder work. Non-members and gradient-enabled distributed decoding are rejected.
• The existing temporal-geometry test passes.
• Affected-file Ruff, formatting, copy consistency, forward docstrings/return annotations, agent-guide consistency, and git diff --check pass.

PYTHONPATH=src DIFFUSERS_TEST_DEVICE=cpu PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 \
  .venv/bin/python -m pytest -q \
  tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py::TestAutoencoderKLMiniMaxH3SlicingTiling::test_parallel_tiled_decoding \
  tests/models/autoencoders/test_models_autoencoder_kl_minimax_h3.py::TestAutoencoderKLMiniMaxH3::test_encode_decode_temporal_geometry

These are CPU correctness checks, not current-main GPU or pretrained-checkpoint validation. Real NCCL execution, standalone decode timing, and full-resolution output comparison remain pending. Distributed decoding is inference-only and outside a full-graph compiled region.

Related work and self-review

#12590 proposes parallel tiling for AutoencoderKL/Wan; it does not modify the H3 VAE. This draft stays within H3's existing tile/chunk/blend implementation and adds no generic parallel framework. Searches for H3 VAE and AutoencoderKLMiniMaxH3 parallel decoding found no matching open PR. No maintainer acknowledgment for this API has been obtained.

Self-review: READY for draft review. No blocking correctness or dead-code findings remain after tracing decode → _decode → _decode_clip → _stitch_tiles and checking the public opt-in call in the added test. Subgroup source ranks are mapped to global ranks and allocation metadata uses the actual decoder output dtype. Remaining review items are API alignment with related parallel-tiling work, real NCCL validation, and performance evidence. AI assistance was used for implementation, checks, and this description; no submitting-human test execution is claimed.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc? (important for complex PRs)
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline (only applicable for model/pipeline related PRs)?

Who can review?

Community review of the optional API and its relationship to existing parallel-tiling work is welcome.