[TPU] Increase block size and reset block shapes

[bythew3i]

Increase kv cache block size and reset kernel block shapes based on autotuned results from kernel.
But still need to retune the kernel block shapes in kernel.

Note: we should wait for pytorch/xla#9041 to be checkin and update new torch_xla version in requirements.txt

Benchmarked without cache:

v6e-1 (single chip): 7.87 -> 8.37 req / sec

Benchmarking script:

VLLM_USE_V1=1 vllm serve meta-llama/Llama-3.1-8B-Instruct  --disable-log-requests --port 8003 --gpu-memory-utilization 0.95 --max-num-batched-tokens 512 --tensor-parallel-size 1 --max-model-len 2048 --max_num_seqs 512 &> /tmp/serve.log &

python benchmarks/benchmark_serving.py --backend vllm --model meta-llama/Llama-3.1-8B-Instruct --dataset-name random --port=8003 --random-input-len 1800 --random-output-len 128

Before:

============ Serving Benchmark Result ============
Successful requests:                     987       
Benchmark duration (s):                  125.47    
Total input tokens:                      1776600   
Total generated tokens:                  118669    
Request throughput (req/s):              7.87      
Output token throughput (tok/s):         945.80    
Total Token throughput (tok/s):          15105.47  
---------------Time to First Token----------------
Mean TTFT (ms):                          61168.93  
Median TTFT (ms):                        60913.25  
P99 TTFT (ms):                           121404.06 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          31.70     
Median TPOT (ms):                        31.98     
P99 TPOT (ms):                           32.56     
---------------Inter-token Latency----------------
Mean ITL (ms):                           31.69     
Median ITL (ms):                         31.93     
P99 ITL (ms):                            33.08     
==================================================

After:

============ Serving Benchmark Result ============
Successful requests:                     987       
Benchmark duration (s):                  117.96    
Total input tokens:                      1776600   
Total generated tokens:                  118669    
Request throughput (req/s):              8.37      
Output token throughput (tok/s):         1006.00   
Total Token throughput (tok/s):          16066.94  
---------------Time to First Token----------------
Mean TTFT (ms):                          57649.43  
Median TTFT (ms):                        57545.37  
P99 TTFT (ms):                           113943.35 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          29.78     
Median TPOT (ms):                        29.96     
P99 TPOT (ms):                           30.49     
---------------Inter-token Latency----------------
Mean ITL (ms):                           29.76     
Median ITL (ms):                         29.97     
P99 ITL (ms):                            30.93     
==================================================

v6e-8 (multi chip): 4.92 -> 5.42 req / sec

VLLM_USE_V1=1 vllm serve "meta-llama/Llama-3.1-70B" --download_dir "/root/.cache" --disable-log-requests --tensor_parallel_size=8 --max-model-len=2048 --gpu-memory-utilization 0.95 --max-num-batched-tokens 512 --max_num_seqs 512  &> /tmp/serve.log &


python benchmarks/benchmark_serving.py --backend vllm --model meta-llama/Llama-3.1-70B --dataset-name random --port=8003 --random-input-len 1800 --random-output-len 128

Before:

============ Serving Benchmark Result ============
Successful requests:                     987       
Benchmark duration (s):                  200.42    
Total input tokens:                      1776600   
Total generated tokens:                  111817    
Request throughput (req/s):              4.92      
Output token throughput (tok/s):         557.91    
Total Token throughput (tok/s):          9422.22   
---------------Time to First Token----------------
Mean TTFT (ms):                          98990.46  
Median TTFT (ms):                        99069.74  
P99 TTFT (ms):                           195396.73 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          51.12     
Median TPOT (ms):                        51.44     
P99 TPOT (ms):                           52.66     
---------------Inter-token Latency----------------
Mean ITL (ms):                           51.27     
Median ITL (ms):                         51.38     
P99 ITL (ms):                            53.50     
==================================================

After

============ Serving Benchmark Result ============
Successful requests:                     987       
Benchmark duration (s):                  182.04    
Total input tokens:                      1776600   
Total generated tokens:                  111445    
Request throughput (req/s):              5.42      
Output token throughput (tok/s):         612.19    
Total Token throughput (tok/s):          10371.49  
---------------Time to First Token----------------
Mean TTFT (ms):                          89566.91  
Median TTFT (ms):                        89369.13  
P99 TTFT (ms):                           177012.56 
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          46.24     
Median TPOT (ms):                        46.64     
P99 TPOT (ms):                           47.39     
---------------Inter-token Latency----------------
Mean ITL (ms):                           46.45     
Median ITL (ms):                         46.59     
P99 ITL (ms):                            48.17     
==================================================

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

🚀

[alexm-redhat]

@bythew3i which model did you test on single and multi-chip setups?

[alexm-redhat]

256 block size seems a bit aggressive. Maybe you can try the sharegpt (it has average short prompts) benchmark and see if you don't see a regression.

[yaochengji]

And usually the page_size should not be larger than max_model_len.

[bythew3i]

re: @alexm-redhat I modified the code to calculate the block size based on the max-model-len. PTAL.

BTW, can you please share the cmds used for sharegpt benchmarking?

[bythew3i]

re: @yaochengji now the way how we choose page size should handle this. PTAL get_page_size in pallas.py

[bythew3i]

@bythew3i which model did you test on single and multi-chip setups?

Hi @alexm-redhat, sorry for the late reply. I benchmarked meta-llama/Llama-3.1-8B-Instruct and meta-llama/Llama-3.1-70B. I also updated benchmarking cmds in the PR description.

[bythew3i]

CC: @yarongmu-google @bvrockwell

[yaochengji]

LGTM, when CI is green. Thanks for the contribution.

[bythew3i]

QQ: what is the cmd to format the code?

[yaochengji]

QQ: what is the cmd to format the code?

You can use pre-commit run --all-files, from https://docs.vllm.ai/en/stable/contributing/overview.html#testing

Sometimes there's still some lines not formatted, you can install a ruff plungin in vscode and format selected

[bythew3i]

@WoosukKwon @alexm-redhat PTAL! Thanks!

[lsy323]

cc @mgoin

[Akshat-Tripathi]

Hi, not a review, but this is interesting information, is there anywhere I can find it online?

[bythew3i]

Hi Akshat, thanks for asking! I can not find any TPU's SREGs number documented anywhere publicly. So I think it is better to not mention this in the comments.

[Akshat-Tripathi]

Oh ok thanks!

[mgoin]

In trying to run a quick test I accidentally ran V0. It seems this PR breaks V0 by not specifying the block_size in that flow

  File "/home/mgoin/code/vllm/vllm/executor/uniproc_executor.py", line 56, in collective_rpc
    answer = run_method(self.driver_worker, method, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgoin/code/vllm/vllm/utils.py", line 2463, in run_method
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgoin/code/vllm/vllm/worker/worker_base.py", line 594, in init_worker
    self.worker = worker_class(**kwargs)
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/mgoin/code/vllm/vllm/worker/tpu_worker.py", line 51, in __init__
    self.model_runner: TPUModelRunner = TPUModelRunner(
                                        ^^^^^^^^^^^^^^^
  File "/home/mgoin/code/vllm/vllm/worker/tpu_model_runner.py", line 111, in __init__
    self.max_num_blocks_per_seq = (self.model_config.max_model_len //
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: unsupported operand type(s) for //: 'int' and 'NoneType'

        self.max_num_blocks_per_seq = (self.model_config.max_model_len //
                                       self.block_size)

[bythew3i]

@mgoin PTAL! Thanks!

[mgoin]

LGTM thank you!

[bythew3i]

@mgoin Can you please help merge this PR?

[mgoin]

@bythew3i the TPU V1 sampler test is failing https://buildkite.com/vllm/ci/builds/19212#01969083-3aed-45d6-9738-f4d601113fd5/6-1707
Can you see if you can reproduce locally? It is concerning that the outputs are all "!"

[bythew3i]

@bythew3i the TPU V1 sampler test is failing https://buildkite.com/vllm/ci/builds/19212#01969083-3aed-45d6-9738-f4d601113fd5/6-1707 Can you see if you can reproduce locally? It is concerning that the outputs are all "!"

@mgoin Is this the right cmd to test?

VLLM_USE_V1=1 pytest tests/v1/tpu/test_sampler.py

I tested on local... it also failed at main branch... let me pull the latest change to see if the failure still exist

[bythew3i]

@bythew3i the TPU V1 sampler test is failing https://buildkite.com/vllm/ci/builds/19212#01969083-3aed-45d6-9738-f4d601113fd5/6-1707 Can you see if you can reproduce locally? It is concerning that the outputs are all "!"

The error seems not related to this PR... It fails at HEAD on main branch @mgoin