[TPU][V1] Fix Sampler recompilation

[NickLucche]

Fix XLA recompilations introduced .
Namely it factors out the on-device slicing that is happening inside InputBatch._make_sampling_metadata as well as an issue with XLA not pre-compiling sample_from_hidden when its output isn't moved to cpu.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[yaochengji]

Hi @NickLucche , one thing I noticed is that the host overhead becomes larger.

Previously the average is 3.5ms, now the average is about 5ms.

You can set the environemnt variable VLLM_TORCH_PROFILER_DIR and add --profile in benchmark_script.py to enable profiling.

[lsy323]

Hi @NickLucche I found sampler test is not added to TPU CI. Also it's running in enforce_eager, could we turn off enforce_eager so that recompilation can be checked in our tests?