[Proposal] Vectorize tiled VAE blending while preserving rounding and CUDA graph replay

[ShikeChen01]

Withdrawn by the author pending a complete contribution review. The previous readiness verdict was premature: local correctness and random-weight timing results do not establish that this approach and its scope satisfy the project's design and validation expectations. Repository guidance and the full related discussions are being reviewed, and the checklist will record unmet requirements explicitly. No upstream PR has been submitted.

---

I'd like to coordinate an AI-assisted contribution that vectorizes the row/frame loops used to blend overlapping VAE tiles. This is a proposal for review before opening a PR, following the AI-assisted contribution guidelines.

The prepared implementation is on this branch, at 51ef4e27cd8bb2f818da368af95e56b57eb71da3, based on main c5469b7ceb606edd7ba6570dcd17d38590a18db6.

In complete StableDiffusionPipeline calls, the local benchmark's median process latency changed from 1,128.1 to 1,008.9 ms in FP16, and 1,182.6 to 1,017.4 ms in BF16. This uses seeded random SD-v1-size UNet and VAE weights, supplied prompt embeddings, 1024×1024 output, batch 1, DDIM 2 steps, guidance scale 1, native PyTorch SDPA and VAE tiling. It measures the executable pipeline with those settings; it does not establish trained-image quality or performance for longer denoising runs or video models.

Environment: Windows 11, Python 3.11.15, PyTorch 2.9.1+cu128, CUDA 12.8, cuDNN 9.10.2, RTX 3070 Ti 8 GB, driver 591.86, Diffusers 0.41.0.dev0. Three alternating fresh-process pairs per dtype, two warmups and five synchronized measurements per process. Paired speedups ranged from 1.052–1.181× for FP16 and 1.128–1.162× for BF16. First-pair outputs matched upstream exactly for both dtypes; peak allocated memory was unchanged at about 2,449.9 MiB.

Proposed scope:

• Replace 42 blend_v, blend_h and blend_t loops across 18 existing autoencoder implementations with broadcasts over the seam, preserving clamping, in-place output and per-product rounding.
• Share coefficient construction in autoencoders/vae.py. CUDA coefficients are built on-device, with float64 tensor division to match Python-double rounding and support graph replay. Other devices retain Python-side coefficients, so MPS and other accelerators are not required to support float64.
• Regenerate copied blend methods and add numerical parity plus real tiled encode/decode CUDA graph regression tests. No API, dependency, configuration or checkpoint changes.

Validation and self-review:

• Six real AutoencoderKL CUDA graph encode/decode cases pass in FP32/FP16/BF16 on both current main and the proposed head. Each compares three replays with changed inputs against eager output exactly.
• python -m pytest tests/models/autoencoders/test_models_autoencoder_kl.py -k tiling_cuda_graph_replay -q: 6 passed.
• python -m pytest tests/models/autoencoders -k 'tile_blend_matches_reference_loop or enable_disable_tiling or enable_disable_slicing' -q: 82 passed, 50 skipped on CUDA.
• DIFFUSERS_TEST_DEVICE=cpu python -m pytest tests/models/autoencoders -k tile_blend_matches_reference_loop -q: 48 passed, 40 skipped.
• Supplementary comparisons cover every changed method: 4,032 exact eager cases, 504 exact graph replays and 84 fullgraph traces with the eager backend. This does not certify Inductor execution.
• Ran the commands behind make style and make fix-copies directly on Windows with Ruff 0.9.10; all exited 0. Copied code is consistent after normalizing Windows checkout line endings. Setup checks passed.
• Applied the repository's self-review skill to the entire diff: READY for maintainer review, with no remaining blocking findings or added dead code. The shared helper has 42 live callers. Existing usage docs remain accurate because no public behavior/configuration changes. Full video-pipeline performance, non-CUDA accelerators and Inductor remain unmeasured; very short seams may incur more setup overhead. The full test suite has not been rerun.

I reviewed similar open and closed work, including #14694, #10488 and the original size-clamping fix #2660. I did not find an existing contribution proposing this same vectorization. The LTX2 diffusion decoder overlap is limited to the existing blend methods and can be rebased if that refactor lands first.

Would maintainers be interested in a single PR with this scope and coefficient approach? I will wait for acknowledgment here before opening it.

Reproducible full-pipeline benchmark

Save the script as bench_tiled_vae.py. Select each checkout with PYTHONPATH, use the same environment, and repeat base/change then change/base in fresh processes. Repeat with --dtype bf16.

PYTHONPATH=/path/to/base/src python bench_tiled_vae.py --label base --dtype fp16 --out base.json --save-output base.pt
PYTHONPATH=/path/to/change/src python bench_tiled_vae.py --label change --dtype fp16 --out change.json --save-output change.pt

"""Full StableDiffusionPipeline timing with seeded random SD-v1-size components.

No checkpoint download or trained-image quality claim. Uses supplied prompt embeddings,
1024x1024 output, batch 1, DDIM with 2 steps, guidance_scale=1, and VAE tiling.
Choose the Diffusers checkout with PYTHONPATH. Run base/change in alternating fresh
processes; use --save-output on the first pair to compare output tensors exactly.
"""
import argparse
import json
from pathlib import Path
import platform
import statistics
import subprocess
import time

import torch
import diffusers
from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNet2DConditionModel

parser = argparse.ArgumentParser()
parser.add_argument("--label", required=True)
parser.add_argument("--dtype", choices=["fp16", "bf16"], default="fp16")
parser.add_argument("--out", type=Path, required=True)
parser.add_argument("--reps", type=int, default=5)
parser.add_argument("--save-output", type=Path)
args = parser.parse_args()
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16}[args.dtype]
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.manual_seed(0)
vae = AutoencoderKL(
 in_channels=3, out_channels=3, down_block_types=("DownEncoderBlock2D",) * 4,
 up_block_types=("UpDecoderBlock2D",) * 4, block_out_channels=(128, 256, 512, 512),
 layers_per_block=2, latent_channels=4, norm_num_groups=32, sample_size=512,
).to(device="cuda", dtype=dtype).eval()
vae.enable_tiling()
torch.manual_seed(0)
unet = UNet2DConditionModel(
 sample_size=64, in_channels=4, out_channels=4, layers_per_block=2,
 block_out_channels=(320, 640, 1280, 1280),
 down_block_types=("CrossAttnDownBlock2D",) * 3 + ("DownBlock2D",),
 up_block_types=("UpBlock2D",) + ("CrossAttnUpBlock2D",) * 3,
 cross_attention_dim=768, attention_head_dim=8,
).to(device="cuda", dtype=dtype).eval()
scheduler = DDIMScheduler(
 beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear",
 clip_sample=False, set_alpha_to_one=False, steps_offset=1,
)
pipe = StableDiffusionPipeline(
 vae=vae, unet=unet, scheduler=scheduler, text_encoder=None, tokenizer=None,
 safety_checker=None, feature_extractor=None, requires_safety_checker=False,
)
pipe.set_progress_bar_config(disable=True)
generator = torch.Generator(device="cuda").manual_seed(1234)
embeds = torch.randn((1, 77, 768), generator=generator, device="cuda", dtype=dtype)
latents = torch.randn((1, 4, 128, 128), generator=generator, device="cuda", dtype=dtype)


def run():
 return pipe(prompt_embeds=embeds, latents=latents, height=1024, width=1024,
 num_inference_steps=2, guidance_scale=1.0, output_type="pt").images


for _ in range(2):
 output = run()
torch.cuda.synchronize()
assert torch.isfinite(output).all().item()
samples = []
for _ in range(args.reps):
 torch.cuda.synchronize()
 start = time.perf_counter()
 output = run()
 torch.cuda.synchronize()
 samples.append((time.perf_counter() - start) * 1000)
torch.cuda.reset_peak_memory_stats()
output = run()
torch.cuda.synchronize()
peak = torch.cuda.max_memory_allocated() / 2**20
assert torch.isfinite(output).all().item()
if args.save_output:
 torch.save(output.cpu(), args.save_output)
result = {
 "label": args.label, "dtype": args.dtype, "pipeline_ms": samples,
 "median_ms": statistics.median(samples), "peak_allocated_mib": peak,
 "finite": True, "output_shape": list(output.shape), "random_weights": True,
 "config": "SD-v1-size UNet and VAE; DDIM 2 steps; guidance 1; supplied embeddings; 1024x1024",
 "python": platform.python_version(), "platform": platform.platform(),
 "torch": torch.__version__, "cuda": torch.version.cuda,
 "diffusers": diffusers.__version__, "diffusers_file": diffusers.__file__,
 "gpu": torch.cuda.get_device_name(), "cudnn": torch.backends.cudnn.version(),
 "driver": subprocess.check_output(["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader"], text=True).strip(),
 "attention": "native PyTorch SDPA (default backend selection)",
}
args.out.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
print(json.dumps(result), flush=True)