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

Update:

• Persistent sampling metadata tensors, re-using the ones in input_batch to reduce waste
• (hopefully) substantially simplified code. TPUSupportedSamplingMetadata is now a simpler wrapper around tensors managed in input_batch.

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

[NickLucche]

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

Thanks for looking into it. I had to move the sampling pre-processing (slicing) from on device to host. I'll look into optimizing it, but on GPU we just keep everything on device and do the slicing there, which may cause recompilation depending on the number of reqs being scheduled. TL;DR It's a trade-off in integrating existing logic.

Also it's running in enforce_eager, could we turn off enforce_eager so that recompilation can be checked in our tests?

Yes but I'd rather change the logic then, probably merge the two tests into one.

[NickLucche]

@lsy323 there's still a very small graph being re-compiled at runtime that I haven't been able to track down. I have to hold on that test until it is figured out.

[mgoin]

Looks good to me, just needs some cleanup. Have you validated a performance win smoke test?

[mgoin]

Worth leaving a comment why this isn't at the end of the loop

[yaochengji]

@lsy323 there's still a very small graph being re-compiled at runtime that I haven't been able to track down. I have to hold on that test until it is figured out.

@NickLucche , we can use PT_XLA_DEBUG_LEVEL=2 to get where the recompilation is triggered. And please remember to clear the xla compilation cache before the execution.