VAE decode → GPU handoff (FLUX/LTX render services)

Companion to PREFILL_ATTENTION_GPU_HANDOFF.md. Same class of fix ("un-GEMM'd CPU loop → GEMM-ify on GPU"), same constraint (needs a CUDA box — no local nvcc), same route (lands via the gnosis GitHub remote a CUDA box is already pulling).

The measurement (Cloud Run L4, this session)

A FLUX 256px / 4-step render = 681s end-to-end, broken down from live logs:

LTX is worse: its VAE is a native 3D causal decode (model_ltx_vae.rs, LtxVaeDecoder, reference-verified cos≈1.0 vs diffusers AutoencoderKLLTXVideo.decode); the bin comment says "keep clips SMALL". The DiT already runs on GPU; the VAE does not — it's the analog of the un-GEMM'd attention loop.

The CPU VAE code (read these for exact shapes)

• FLUX: src/model_flux_vae.rs — conv2d_nchw (:42, the per-layer CPU convolution), decode_flux_vae_boundary_rgb (:128). Walk the decoder: latent → conv_in → mid (resnet/attn) → up-blocks (resnet + upsample/nearest to full res) → conv_out → RGB. The late up-blocks at full spatial resolution dominate (largest H×W × Cin×kh×kw), as in any VAE decoder.
• LTX: src/model_ltx_vae.rs — LtxVaeDecoder::decode (called from src/bin/ltx-txt2video.rs:337, Vol volume type). 3D causal conv (depth=time + H×W), causal padding on the time axis. Heavier (extra dim).

GPU plan — im2col + cuBLAS SGEMM (pragmatic, reuses existing wiring)

For each conv layer, on the GPU:

1. im2col: gather input patches → matrix A = [out_h·out_w × (Cin·kh·kw)] (3D/LTX: [out_t·out_h·out_w × (Cin·kt·kh·kw)] with causal time padding). A small CUDA kernel (one thread per output element column-group) or a cublas- friendly layout.
2. GEMM: Out[out_spatial × Cout] = A · Wᵀ where W = [Cout × (Cin·kh·kw)]. Reuse the raw-pointer cublas SGEMM body already in cuda_matmul.rs (CudaBlas, Gemm, GemmConfig, the gemm_with_weight* raw *const c_void path). f32 is fine (VAE is small vs the DiT; keep weights f32-resident).
3. bias + activation (SiLU/GELU per layer) fused into the epilogue or a tiny kernel; upsample (nearest/conv-transpose) as a kernel or as a strided im2col. Keep activations resident on-device across layers (only the final RGB comes back to host for PNG/ffmpeg).

Correctness gate (non-negotiable — Sardis discipline)

The GPU result MUST match the CPU reference to cos≈1.0 (the repo's bar; the CPU path stays source of truth). Harnesses already exist: FLUX --dump-vae-decode / flux-vae-dump, LTX ltx-vae-dump / --dump-vae-decode (the bins that produced the cos≈1.0-vs-diffusers verification). Gate the GPU path behind an env flag (default safe), mirror the gpucheck_* / try_* return-bool-then-CPU-fallback discipline in cuda_matmul.rs. Any mismatch → fall back to conv2d_nchw / LtxVaeDecoder.

Insertion points

• FLUX: where flux-txt2img.rs calls decode_flux_vae_boundary_rgb (the seeder spawns /opt/flux/bin/flux-txt2img). Add try_vae_decode_gpu(...) -> bool.
• LTX: src/bin/ltx-txt2video.rs:337 dec.decode(&latent); add a GPU path on LtxVaeDecoder with the same fallback.

Build/verify loop (CUDA box)

No local nvcc here → must iterate on a CUDA-equipped box: edit → cargo build --release -p distributed-inference --bin flux-txt2img --features cuda → run the *-vae-dump parity harness (cos≈1.0) → only then ship the image.

Expected win

VAE 321s → single-digit seconds (massively parallel GPU conv vs the CPU loop). Combined with the baked-knot cold-start fix (already shipped: Dockerfile.{flux,ltx}-gpu knot-baker stage → :baked-v1 images, ~7.6min download → ~seconds mmap) and a DiT weight-residency fix (secondary: logs show Q8 weight upload ... MISS ... q8 resident 9.34 GB / budget 19.0 GB — the Q8 DiT weights re-dequant+upload each of the 4 steps instead of staying resident; making them resident across steps cuts into the 273s DiT), a FLUX render drops from ~681s+cold toward ~tens of seconds warm.

Priority order

1. ✅ Bake knots (DONE — kills the ~7.6min cold start; non-CUDA).
2. GPU VAE (this doc — the ~321s, biggest single render chunk).
3. DiT weight residency across denoise steps (the 273s; check cuda_matmul.rs residency keying for the DiT weights vs the per-step re-upload).