[torch.compile] avoid Dynamo guard evaluation overhead

[youkaichao]

We have 2 types of runtime overhead in TPU:

• Dynamo guard evaluation overhead, chooses which code to run
• torch xla overhead, convert function input to xla input

We can manage to remove the first one, via adding one layer dispatcher above Dynamo.

I did systematic measurement this time, and find that:

• pure xla execution takes 7ms for every decoding step
• combining both overhead ( the current main branch), it takes 8.2ms for every decoding step
• removing Dynamo overhead (this PR), it takes 8.0ms for every decoding step

It turns out the xla overhead is the main overhead. But I think it is still worthwhile to get rid of the Dynamo overhead before we remove the xla overhead.

[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]

NOTE: my test code is still https://github.com/vllm-project/vllm/blob/main/examples/offline_inference_tpu.py

[youkaichao]

this looks surprisingly effective.

I run python benchmarks/benchmark_throughput.py --input-len 256 --output-len 256 --model google/gemma-2b

main:

Throughput: 16.70 requests/s, 8549.39 tokens/s

this PR:

Throughput: 17.39 requests/s, 8902.73 tokens/s

it counts as 4% throughput improvement