[TPU][V1] Refine tpu_model_runner to mitigate future recompilation issues

[yaochengji]

• Wrap all the TPU computation into torch.compile, currently they're
  • forward: total_bucket_size = token_bucket_num
  • select_hidden_states: total_bucket_size = token_bucket_num * req_bucket_num
  • sample_from_hidden: total_bucket_size = req_bucket_num * 2
• Remove pytorch operations in TPUSupportedSamplingMetadata.from_input_batch
• Change the implementation of torch.where to if else branches

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

🚀

[NickLucche]

I can't review rn sorry. Is it faster than main?

[yaochengji]

I can't review rn sorry. Is it faster than main?

Yes, both compilation and execution

[NickLucche]

Thanks for the PR, overall this lgtm, especially the from_input_batch change.

I don't fully grasp where the optimization lies with the if/else branch as we're just forking two different graphs now, but that's about it.

[yaochengji]

I don't fully grasp where the optimization lies with the if/else branch as we're just forking two different graphs now, but that's about it.

If we use torch.where(pred, A, B), we only need a compile one graph for both pred is True and False, but both A and B will be computed no matter what the value of pred is.

If we use

if pred:
    A
else:
    B

We need to compile two graphs for pred is True or False. The benefit is that we don't need to compute both A and B in execution.

also cc @robertgshaw2-redhat , I remember you had similar questions.

[NickLucche]

but both A and B will be computed no matter what the value of pred is

Ok I see the issue now. The simpler if/else solution fits our use-case surely better. Thanks for explaining!

[vanbasten23]

nit: it seems that the impl directly uses the cpu tensor and move them to the xla device

[yaochengji]

I think it aligns with the description?

[mgoin]

Looks like the TPU V1 test is failing

ERROR collecting tests/v1/tpu/worker/test_tpu_model_runner.py 
  | ImportError while importing test module '/workspace/vllm/tests/v1/tpu/worker/test_tpu_model_runner.py'.
  | Hint: make sure your test modules/packages have valid Python names.
  | Traceback:
  | /usr/local/lib/python3.10/importlib/__init__.py:126: in import_module
  | return _bootstrap._gcd_import(name[level:], package, level)
  | tests/v1/tpu/worker/test_tpu_model_runner.py:10: in <module>
  | from vllm.v1.worker.tpu_model_runner import (TPUModelRunner,
  | E   ImportError: cannot import name '_get_paddings' from 'vllm.v1.worker.tpu_model_runner' (/workspace/vllm/vllm/v1/worker/tpu_model_runner.py)

[yaochengji]

Looks like the TPU V1 test is failing

@mgoin thanks for reminding, it should be fixed now. BTW, can we make the TPU CI test not a soft-fail now? cc @robertgshaw2-redhat

[mgoin]

Nice work, I feel this is a good path to go down! Thank you