[TPU] Support sliding window and logit soft capping in the paged attention kernel for TPU.

[vanbasten23]

After I added the sliding window and logit soft-capping support to the paged attention Pallas kernel, this PR is intended to added the sliding window and logit soft-capping at the vLLM level, so that we can support gemma model.

Test plans:

• pytest -s -v /workspace/vllm/tests/v1/tpu/test_pallas.py
• pytest -v -s tests/entrypoints/llm/test_accuracy.py::test_lm_eval_accuracy_v1_engine 2>&1 | tee out.txt

The Gemma model that we care about are google/gemma-3-27b-it and google/gemma-3-1b-it.

cc @miladm

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

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

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

[brittrock]

This is great @vanbasten23 thanks for piping this in! Would we be able to add a Gemma 3 [text-only] test in please?

Cc @robertgshaw2-redhat @yarongmu-google @yaochengji

[yaochengji]

Have you compared this score with GPU?

[vanbasten23]

Thanks for the suggestion! I ran on GPU and got the same accuracy (0.25) for gemma-3-1b (test output on GPU: https://gist.github.com/vanbasten23/eda147d28db5f11187178771372ee39f)

[yaochengji]

LGTM, thanks!

Note that the CI is not green.

[DarkLight1337]

You shuold merge from main to fix docker build

[NickLucche]

Great job!
I think some of the commits need signing, please take a look

[bvrockwell]

@yaochengji we shouldn't merge this until @vanbasten23 validates 27B as well. Let's wait for these results
.

[bvrockwell]

waiting for multi chip validation - thank you!

[vanbasten23]

@yaochengji we shouldn't merge this until @vanbasten23 validates 27B as well. Let's wait for these results .

I tested the 27b and it fails:

MODEL=google/gemma-3-27b-it
VLLM_USE_V1=1 vllm serve $MODEL --seed 42 --disable-log-requests --gpu-memory-utilization 0.95 --max-num-batched-tokens 1024 --max-num-seqs 512 --tensor-parallel-size 4 --max-model-len 2048 &
VLLM_USE_V1=1 python3 benchmarks/benchmark_serving.py --model $MODEL  --dataset-name random --random-input-len 1800 --random-output-len 128

the failure is in https://gist.github.com/vanbasten23/637bd3af3c3946c00e6fb1ce27892e46 which fails at tpu_woker.py, outside the kernel.

But I don't think it should block this PR. This PR intends to add sliding window/logit soft-capping support. @bvrockwell The gemma-3 27B failure is a separate issue and can be fixed in a later PR.

[bvrockwell]

@yaochengji we shouldn't merge this until @vanbasten23 validates 27B as well. Let's wait for these results .

I tested the 27b and it fails:

MODEL=google/gemma-3-27b-it
VLLM_USE_V1=1 vllm serve $MODEL --seed 42 --disable-log-requests --gpu-memory-utilization 0.95 --max-num-batched-tokens 1024 --max-num-seqs 512 --tensor-parallel-size 4 --max-model-len 2048 &
VLLM_USE_V1=1 python3 benchmarks/benchmark_serving.py --model $MODEL  --dataset-name random --random-input-len 1800 --random-output-len 128

the failure is in https://gist.github.com/vanbasten23/637bd3af3c3946c00e6fb1ce27892e46 which fails at tpu_woker.py, outside the kernel.

But I don't think it should block this PR. This PR intends to add sliding window/logit soft-capping support. @bvrockwell The gemma-3 27B failure is a separate issue and can be fixed in a later PR.

ok understood @vanbasten23 , looking forward to the multichip gemma 3 PR! Thanks for adding sliding window and logits soft capping.

[mergify]

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

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

[vanbasten23]

hi @robertgshaw2-redhat , I resolved the merge conflicts but the auto-merge was canceled. Could you enable the auto-merge again? Thanks!

[mgoin]

FYI the TPU V1 test was failing when this was merged

[vanbasten23]

FYI the TPU V1 test was failing when this was merged

Good catch. This test fails even before this PR. For example, in the PR before #15586, the TPU CI soft fails with the same error.

I think we should hard fail the test failure in TPU CI.

[vanbasten23]

The failing test seems to have been fixed later in #16041 for example:

======================================================================= warnings summary =======================================================================
--
  | tests/tpu/test_compilation.py::test_tpu_compilation
  | :0: UserWarning: /usr/local/lib/python3.10/site-packages/depyf/explain/enable_debugging.py:163: You are trying to debug `torch.compile`. Please make sure the code runs multiple times to cover all the possible branches.
  |  
  | -- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
  | =========================================================== 1 passed, 1 warning in 128.19s (0:02:08) ===========================================================