diff --git a/README.md b/README.md index 8be48f1..2805d9f 100644 --- a/README.md +++ b/README.md @@ -586,15 +586,22 @@ Use `--raw` to bypass weighting and fall back to plain diffusers encoding. | `vnasset upscale` | ✅ Working | | Self-contained torch wheel | ✅ Working | | `vnasset serve` (daemon/HTTP API) | 🚧 Planned | -| `torch.compile` on UNet | 🚧 Planned | -| Batch edit loop (shared VAE encode) | 🚧 Planned | +| `torch.compile` on UNet | 🔬 Needs spike (see [decisions](docs/optimization-decisions.md)) | +| Shared encode optimization | 📋 Analyzed, deferred (see [decisions](docs/optimization-decisions.md)) | ## Future Improvements -- **`torch.compile` on UNet** — the UNet forward is identical each step; ROCm's - `torch.compile` support is maturing and could cut per-step time significantly. -- **Shared encode optimization** — for N edit variants of the same input image, - run VAE encode and VL visual token encoding once, then only text-encode and - denoise per variant. - **`vnasset serve`** — lightweight daemon with Unix socket or HTTP API for integrating VNAsset into external tools. +- **`torch.compile` on UNet** — the UNet forward is identical each step, but + ROCm `torch.compile` maturity is uncertain. Needs a 30-minute spike to measure + actual speedup and check for graph breaks / crashes on RDNA 3.5. If the spike + shows 15%+ and stable, implement with a warmup step at session load. +- **Shared encode optimization** — for N edit variants of the same input image, + saves ~1.2–2.3s per variant (VAE encode + vision encoder). The dominant cost + (transformer denoising, ~45–50s) cannot be shared, so the overall gain is + 2–4%. Deferred until workload shape demands it (100+ variants). When built, + the right place is the pipeline runner, not the session API. + +See [`docs/optimization-decisions.md`](docs/optimization-decisions.md) for +full tradeoff analysis with data and cost estimates. diff --git a/docs/optimization-decisions.md b/docs/optimization-decisions.md new file mode 100644 index 0000000..9f0542f --- /dev/null +++ b/docs/optimization-decisions.md @@ -0,0 +1,169 @@ +# Optimization Decisions + +Analysis and rationale for optimization ideas considered but not yet +implemented. Each entry answers: what data says, what we'd save, what it +costs, and whether it's worth building now. + +Hardware: **AMD Strix Halo** (Ryzen AI Max 395 Pro, Radeon 8060S, 128 GB). +ROCm torch 2.11.0, Triton 3.6.0. + +--- + +## Shared encode optimization (Qwen Image Edit) + +### What it is + +When editing the same input image with N different prompts (e.g., "make her +smile", "make her angry"), the VAE encode of the input image and the vision +encoder forward pass (inside the text encoder) are identical for all variants. +Only the text prompt suffix and the denoising loop differ. + +### What the data says + +The Qwen Image Edit pipeline `__call__` performs these steps per edit: + +| Step | Cost (estimate) | Shareable? | +|------|-----------------|-----------| +| Image preprocessing (resize ×2) | negligible | Yes | +| VAE encode (1024×1024 → latent) | ~1-2s | **Yes** | +| Vision encoder (384×384 → tokens) | ~0.1-0.3s | **Yes** | +| LLM prefix (system prompt + vision tokens, ~800 tokens) | ~0.1s | Maybe (KV-cache) | +| LLM suffix (edit instruction, ~10-30 tokens) | negligible | No | +| Denoising loop (4 steps, ~9000 tokens/step, 20B params) | ~45-50s | **No** | +| VAE decode (latent → 1024×1024) | ~2-5s | No | + +The dominant cost is the transformer denoising (~45-50s of ~53s total with +flash attention). This cannot be shared — each variant needs different +conditioning, different noise trajectory, different output. + +### What we'd save + +Per variant after the first: **~1.2-2.3s** (VAE encode + vision encoder). + +For 3 expression variants: ~2.4-4.6s saved out of ~159s total (1.5-2.8%). +For 10 variants: ~10.8-20.7s saved out of ~530s total (2-4%). + +### What it costs + +- New API surface on `VnAssetsSession` (or pipeline internals touched) +- Pipeline `__call__` cannot be used as a black box — must intercept between + VAE encode/vision encode and denoising loop +- KV-caching the vision prefix adds further complexity for marginal gain + (~0.1-0.2s per variant) +- Code complexity per second saved is high compared to other options + +### Decision: Do not build yet + +The savings are modest because the bottleneck (transformer denoising) cannot +be shared. The code complexity per second saved is worse than other options. + +When the workload requires 100+ variants from a single base, the absolute +savings (100-200s) justify the complexity. At that point, the right place +for this optimization is the **pipeline runner** (`_run_edit`), not the +session API — the pipeline knows the cross-product shape and can group by +input image, encode once, denoise many times. + +--- + +## `torch.compile` on SDXL UNet + +### What it is + +Wrap `pipe.unet` with `torch.compile(mode="reduce-overhead")` so PyTorch +fuses adjacent operations into fewer GPU kernels. The UNet forward is +identical in structure every step (same shapes, same ops, different values), +which is the ideal case for compilation. + +### What the data says + +The SDXL UNet has ~20 identical forward passes per generate (1024×1024, +20 steps, ~1.0-1.3s each). The UNet contains: + +- Conv2d (MIOpen-optimized — compile won't speed these up) +- GroupNorm / LayerNorm (compile can fuse with neighbors) +- Linear (to_q, to_k, to_v, to_out) — small matmuls, compile can fuse with + reshape/transpose +- Custom attention via `simple_attention_forward` (raw matmul + softmax) — + already avoids SDPA dispatch overhead, but compile can fuse pointwise ops +- SiLU, add, multiply — pointwise ops that fuse trivially + +The ops that benefit most from `torch.compile` (pointwise fusion) are the +cheap ops. The expensive ops (conv2d, large matmuls) are already +hand-optimized by MIOpen/rocBLAS. + +### What we'd save + +Community reports on CUDA: 15-30% per-step reduction. +On ROCm with torch 2.11.0 + Triton 3.6.0: conservative hypothesis is +**10-20%**, meaning ~0.8-0.9s per steady-state step instead of ~1.0-1.3s. + +For a 20-step generate: ~16-23s total instead of ~29s (**5-13s saved**). + +| Scenario | Compilation cost | Net benefit | +|----------|-----------------|-------------| +| Single generate | 30-90s compile + ~7s saved | **Net loss** | +| 1 pipeline (1 gen + edits) | 30-90s compile + ~7s saved | **Net loss** | +| 10+ generates in one session | 30-90s compile + 50-130s saved | **Net win** | + +### ROCm-specific risks + +- `torch.compile` on ROCm has a smaller testing surface than CUDA +- `float16` already causes segfaults on RDNA 3.5; compiled graphs may + encounter more stability issues +- Custom attention processor interaction: `SimpleAttnProcessor.__call__` may + introduce graph breaks +- Numerical differences: compiled output may differ slightly from eager; + needs visual quality verification + +### What it costs + +- One `torch.compile` call in `_load_sdxl`, plus a warmup step (like the + Real-ESRGAN 256×256 tile pattern already in the codebase) +- A flag to toggle it (optional, can default off) +- A 30-minute spike to measure actual speedup and check for graph breaks + +### Decision: Spike first, decide after measurement + +The theoretical ceiling is known (10-20%), but the ROCm-specific unknowns +(graph breaks, crashes, actual speedup) dominate the risk. A 30-minute spike +wrapping the UNet, running a warmup step, and measuring 5 runs with and +without `torch.compile` answers the question definitively. + +If the spike shows <5% speedup or crashes: kill the idea. +If it shows 15%+ and stable: implement with a warmup step at session load. + +Implementation sketch for the spike: + +```python +# In _load_sdxl, after patch_unet_attention(pipe.unet): +pipe.unet = torch.compile( + pipe.unet, + mode="reduce-overhead", + fullgraph=False, +) + +# Warmup (amortize compilation into session load): +with torch.no_grad(): + dummy_latent = torch.randn(1, 4, 128, 128, device=device, dtype=dtype) + dummy_t = torch.tensor([999], device=device) + dummy_emb = torch.randn(1, 77, 2048, device=device, dtype=dtype) + pipe.unet(dummy_latent, dummy_t, encoder_hidden_states=dummy_emb).sample +``` + +--- + +## Comparison: savings per unit of complexity + +| Optimization | Savings (3-op pipeline) | Complexity | Ratio | +|-------------|------------------------|------------|-------| +| Flash attention (already done) | 34s | 1 env var | Excellent | +| `torch.compile` on UNet | 5-13s | Low (1 call + warmup) | Good (if it works) | +| Shared VAE/vision encode | 2-5s | Medium-High (pipeline internals) | Poor | +| Shared KV-cache prefix | 0.3-0.6s | High (text encoder internals) | Terrible | + +### Priority order + +1. Verify `torch.compile` with a spike +2. If `torch.compile` works: implement it +3. Shared encode only when the workload shape demands it (100+ variants) +4. `vnasset serve` (separate category — feature, not optimization)