[TPU][DEBUG] Provide Env Variable To Disable Sampler

[robertgshaw2-redhat]

SUMMARY:

• env variable to disable sampler
• should be merged after the triton lazy import

REVIEWER: @yaochengji

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

🚀

[mergify]

⚠️ The sha of the head commit of this PR conflicts with #15656. Mergify cannot evaluate rules on this PR. ⚠️

[robertgshaw2-redhat]

note: there is a separate PR for this that will merge separarely

[yaochengji]

Thanks for adding this support, Robert!

I think we should also change the logic in the dummy run to make it actually pre-compile for compute_logits_no_sampler?

[robertgshaw2-redhat]

Thanks for adding this support, Robert!

I think we should also change the logic in the dummy run to make it actually pre-compile for compute_logits_no_sampler?

I will double check, but I thought the dummy run calls

Thanks for adding this support, Robert!

I think we should also change the logic in the dummy run to make it actually pre-compile for compute_logits_no_sampler?

I will check in the morning, but I thought the dummy run calls execute_model so it should be covered.

[NickLucche]

Curious to know what the gain in execution time is.
There' s a path in sample_from_hidden when all_greedy=True to do just that.
This would reveal the introduced overhead.

[robertgshaw2-redhat]

Perhaps its because we wrap with torch.compile?

[bvrockwell]

it's about 5%

[NickLucche]

comparing with all_greedy=True set in request right

[yaochengji]

I will check in the morning, but I thought the dummy run calls execute_model so it should be covered.

I'm afraid not, dummy_run doesn't call execute_model, the corresponding code is at https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/tpu_model_runner.py#L803

[robertgshaw2-redhat]

I will check in the morning, but I thought the dummy run calls execute_model so it should be covered.

I'm afraid not, dummy_run doesn't call execute_model, the corresponding code is at https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/tpu_model_runner.py#L803

You're right. Ill post up a fix.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @robertgshaw2-redhat.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[yaochengji]

@robertgshaw2-redhat kindly ping about this PR.

[NickLucche]

We also need to avoid compiling sample_from_hidden altogether in capture_model

[NickLucche]

We can close this one too @robertgshaw2-redhat