9c4474b93c
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.
170 lines
6.6 KiB
Markdown
170 lines
6.6 KiB
Markdown
# 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)
|