[core][torch.compile] not compile for profiling

[youkaichao]

we only profile once to determine the space for kv-cache, and don't run profile anymore.

compiling this run will only add guards and code cache for useless code.

removing this compilation can reduce the overhead of dynamp.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[youkaichao]

my measurement on gemma 2b shows this can reduce the Dynamo overhead by about 0.1~0.2 ms.

[WoosukKwon]

Thanks for the PR! Left minor comments.

[WoosukKwon]

@youkaichao On Gemma-2B, this PR reduced the KV cache space from 34552 to 25752, 25% drop.

[WoosukKwon]

@youkaichao Can we do this instead?

# Compiled model only used for initial memory profiling
tmp_for_profile = torch.compile(model, ...)
# Compiled model used for actual execution.
self.model = torch.compile(model, ...)

[youkaichao]

@youkaichao On Gemma-2B, this PR reduced the KV cache space from 34552 to 25752, 25% drop.

you mean you also want the compilation for profiling stage?

[WoosukKwon]

you mean you also want the compilation for profiling stage?

Yeah, while I'm not sure, there's a chance that not using torch.compile can lead to suboptimal memory allocation, overestimating the memory usage in memory profiling.

[youkaichao]

@WoosukKwon how do you like the current implementation? compilation and optimiation still happens for the profiling run, but it will be discarded and does not affect later runs.

[WoosukKwon]

@youkaichao Looks very good to me! Thanks for the quick fix!

[WoosukKwon]

For gemma 2b, I didn't see much performance difference: the latency decreased from 1.59s to 1.58s (batch size 8, input len 1024, output len 128).

[youkaichao]

I'm measuring the time spent in the decode run, and I see a clear overhead reduction.

Previously, it takes 8ms for every step, now it only takes 7.5~7.6 ms.

[WoosukKwon]

@youkaichao I see. How does it work without --enforce-eager? Will the complied model still have only 2 guards?

[youkaichao]

How does it work without --enforce-eager? Will the complied model still have only 2 guards?

yes, you already marked symbolic shapes, and all shapes match the symbolic shapes.