vnassets/docs/optimization-decisions.md

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:

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