[TPU] Fix the test_sampler

[bythew3i]

TLDR: this PR is temporary quick fix for broken CI for tests/v1/tpu/test_sampler.py. Kernel fix is WIP.

After #16458, for the test_sampler_different in tests/v1/tpu/test_sampler.py, we will use (32, 32) block shape instead of (128, 32) block shape. However, this triggers some correctness issue. The root cause is kernel assumes kv_cache has no NAN value. We should fix that in kernel! Meanwhile, just to unblock current CI for TPU, we can reduce max_num_batched_tokens to make the test pass and this should not affect correctness testing. (I have no idea why when max_num_batched_tokens==512, the kv cache will have uninitialized NAN value.)

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

🚀

[bythew3i]

CC: @yaochengji @yarongmu-google @lsy323 @bvrockwell

[bythew3i]

Can you PTAL? @mgoin Thanks!

[mgoin]

Cc @NickLucche

[mgoin]

LGTM thanks for fix, let's see if CI is green!

[yaochengji]

Yeah, it's strange when theres's NaN as kv cache is initialized by using torch.zeros, see https://github.com/vllm-project/vllm/blob/main/vllm/v1/worker/tpu_model_runner.py#L1281.

The NaN value should be introduced during the model execution, then it is written to kv cache.

[yaochengji]

LGTM, thanks for the fix!

[NickLucche]

kernel assumes kv_cache has no NAN value. We should fix that in kernel!

The NaN part is a bit concerning, can we dig into the reason as of why we're producing nan values at runtime? I think the kernel assumption is fine as we can't recover from that situation anyway.