torch.compile ae.decode

[yorickvP]

It takes about 80 seconds on my machine to compile this. Makes the encoding step about 50% faster on A5000 (0.3 -> 0.2s), haven't tried H100.

[daanelson]

this is great! you can push to an internal H100 model (just don't leave it running 😄) on Replicate to test perf in prod, good to have solid metrics on that before we merge

[daanelson]

nit - since these flags are just simple little flags we set setup for dev/schnell predictor, I don't mind adding a separate compile_ae flag

[daanelson]

any way we can compile once with craftier use of dynamo.mark_dynamic - add a max=192 on dims 2 & 3? I assume you've tried this, curious how it breaks

[yorickvP]

I tried max=192, but it didn't have any effect. Setting torch.compile(dynamic=True) makes for one fewer recompile, but I should check the runtime performance of that.

[yorickvP]

Did some H100 benchmarks.

flux-schnell 1 image, VAE not compiled

• 30ms prepare
• 355ms denoise-single-item
• 117ms vae-decode
• total: 505ms

flux-schnell 4 images, VAE not compiled

• 30 ms prepare
• 4x 355 ms denoise-single-item
• 3.21s vae-decode
• total: 4.69s

flux-schnell 4 images, VAE compiled

• 30ms prepare
• 4x 355 ms denoise-single-item
• 152ms vae-decode
• total: 1.62s

---

The VAE speed seems reproducible, where the uncompiled VAE spends a lot of time in nchwToNhwcKernel while the compiled version manages to avoid it.

At the same time, I had a cog bug saying output streams failed to drain, crashing the pod instantly, but this seems unrelated to my PR.

[jonluca]

Did you figure out what the output streams failed to drain issue was? I'm seeing that in prod with our cog deploy as well

[yorickvP]

@jonluca as I understand it, it was a regression in cog and should be fixed when building with 0.9.25 and later.
It was caused by cog replacing stdout/stderr during predictions, but not during setup, causing forked processes to attempt to write to the original stdout/stderr. Should be fixed in replicate/cog#1969 but let me know if it's not!