feat: add support for pytorch compile of the codec model

[no2chem]

This PR adds support for using torch.compile on the codec model. It results in a minor speedup.

I tried to add support for the other models, but it seems there are some issues - in the coarse model I run into an issue with past_kv (setting Dynamic=True), and in the fine model it seems to slow things down a bit.

You can turn off compile by setting SUNO_DISABLE_COMPILE to true.

[tongbaojia]

Hi @no2chem thanks for the PR.

I tested this but I think the codec model isn't the bottleneck for computation in general. It is only used in the codec_decode step and compiling it or not doesn't make a big difference for me...

Compiled:

324 ms ± 3.2 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Original:

324 ms ‡ 11.6 ms per loop (mean + std. dev. of 7 runs, 1 loop each)

[no2chem]

It's a very small change. Yes, the main issue is the coarse model and the semantic model. I have pytorch compilation of those models working locally, it yields a small speedup. I probably can push it in a bit, just have been busy.

[tongbaojia]

@no2chem yep, coarse and semantic models if compiled would be very beneficial. Take your time and let me know if you have updates.