# 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)