docs: add optimization-decisions analysis, update README future plans
Detail the shared-encode and torch.compile tradeoffs with data-driven estimates from the pipeline code. Shared encode saves ~1.2-2.3s per variant (2-4% total) and is deferred; torch.compile needs a 30-minute ROCm spike before deciding. Update Current State and Future Improvements to reflect the analysis.
This commit is contained in:
21
README.md
21
README.md
@@ -586,15 +586,22 @@ Use `--raw` to bypass weighting and fall back to plain diffusers encoding.
|
|||||||
| `vnasset upscale` | ✅ Working |
|
| `vnasset upscale` | ✅ Working |
|
||||||
| Self-contained torch wheel | ✅ Working |
|
| Self-contained torch wheel | ✅ Working |
|
||||||
| `vnasset serve` (daemon/HTTP API) | 🚧 Planned |
|
| `vnasset serve` (daemon/HTTP API) | 🚧 Planned |
|
||||||
| `torch.compile` on UNet | 🚧 Planned |
|
| `torch.compile` on UNet | 🔬 Needs spike (see [decisions](docs/optimization-decisions.md)) |
|
||||||
| Batch edit loop (shared VAE encode) | 🚧 Planned |
|
| Shared encode optimization | 📋 Analyzed, deferred (see [decisions](docs/optimization-decisions.md)) |
|
||||||
|
|
||||||
## Future Improvements
|
## 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
|
- **`vnasset serve`** — lightweight daemon with Unix socket or HTTP API for
|
||||||
integrating VNAsset into external tools.
|
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.
|
||||||
|
|||||||
169
docs/optimization-decisions.md
Normal file
169
docs/optimization-decisions.md
Normal file
@@ -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)
|
||||||
Reference in New Issue
Block a user