Repository: vllm-project/vllm
Pull request: Prefix Caching- fix t4 triton error
Fix #2513, need a smaller block size for Turing GPUs

Problem statement:
prefix caching error with baichuan model
## Machine Info
- Ubuntu
- cuda driver 12.1
- T4 GPU x4

## Reproduce Step

Only change llm line in `examples/offline_inference_with_prefix.py` with
```
llm = LLM(model="baichuan-inc/Baichuan2-13B-Chat", tensor_parallel_size=4, enforce_eager=True, dtype="half", trust_remote_code=True)
``` 

## Error message:
```
triton.runtime.autotuner.OutOfResources: out of resource: shared memory, Required: 65538, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.
```

When i use a common llama arch like `TInyllama` model, it's ok. Might related `Alibi`? 

@DouHappy @caoshiyi @zhuohan123  Please take a look!

Hints from issue discussion:

Commit message: minor

Patch to review:
diff --git a/vllm/model_executor/layers/triton_kernel/prefix_prefill.py b/vllm/model_executor/layers/triton_kernel/prefix_prefill.py
index 8fa70054f02c..10e5f2c76d89 100644
--- a/vllm/model_executor/layers/triton_kernel/prefix_prefill.py
+++ b/vllm/model_executor/layers/triton_kernel/prefix_prefill.py
@@ -5,6 +5,8 @@
 import triton
 import triton.language as tl
 
+TESLA = 'Tesla' in torch.cuda.get_device_name(0)
+
 if triton.__version__ >= "2.1.0":
 
     @triton.jit
@@ -618,7 +620,8 @@ def context_attention_fwd(q,
                               b_ctx_len,
                               max_input_len,
                               alibi_slopes=None):
-        BLOCK = 128
+
+        BLOCK = 128 if not TESLA else 64
         # shape constraints
         Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
         assert Lq == Lk and Lk == Lv


Code review comments:
diff_hunk: @@ -5,6 +5,8 @@
 import triton
 import triton.language as tl
 
+TESLA = 'Tesla' in torch.cuda.get_device_name(0)
path: vllm/model_executor/layers/triton_kernel/prefix_prefill.py
text: @user1:
would it be possible to check for compute capability instead? also, we should do this inside `context_attention_fwd`, as calling CUDA APIs before we set `CUDA_VISIBLE_DEVICES` will lead to errors.

@user2:
Maybe we can set `prefix_block_size` as a parameter in `CacheConfig` and allow user configure in `LLM`? 

@user1:
this sort of a thing should be ideally derived automatically.

@user2:
@user1 @author Does the block size affect the memory utilization or prefix speed?

@author:
@user2 The block size is mainly dependent on the shared mem size for different GPU architectures. It will affect the prefix-prefill kernel speed a little bit but has nothing to do with the GPU memory utilization.

diff_hunk: @@ -5,6 +5,8 @@
 import triton
 import triton.language as tl
 
+TESLA = 'Tesla' in torch.cuda.get_device_name(0)
path: vllm/model_executor/layers/triton_kernel/prefix_prefill.py
text: @user1:
Can we set this variable in a function instead of a global variable? Setting it in global variable may lead to issues in distributed setting.