xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-10 14:25:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

Prefix Caching- fix t4 triton error (#2517)

Browse Source
This commit is contained in:
shiyi.c_98 2024-02-16 14:17:55 -08:00 committed by GitHub
parent 5255d99dc5
commit 64da65b322
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/model_executor/layers/triton_kernel/prefix_prefill.py
View File

@ -618,7 +618,9 @@ if triton.__version__ >= "2.1.0":
b_ctx_len, b_ctx_len,
max_input_len, max_input_len,
alibi_slopes=None): alibi_slopes=None):
BLOCK = 128
cap = torch.cuda.get_device_capability()
BLOCK = 128 if cap[0] >= 8 else 64
# shape constraints # shape constraints
Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1] Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
assert Lq == Lk and Lk == Lv assert Lq == Lk and Lk == Lv