先寫個python腳本, test.py, 保存後運行, 或者直接在命令行裡輸入python回車後粘貼運行下方代碼:
import torch
import flashinfer
kv_len = 2048
num_kv_heads = 32
head_dim = 128
k = torch.randn(kv_len, num_kv_heads, head_dim).half().to('cuda')
v = torch.randn(kv_len, num_kv_heads, head_dim).half().to('cuda')
# CUDA Decoding for single request
q = torch.randn(num_kv_heads, head_dim).half().to('cuda')
o = flashinfer.single_decode_with_kv_cache(q, k, v)
print("FlashInfer seems ok.")如果報運行腳本時TORCH_CUDA_ARCH_LIST變量找不到, 那是因爲你編譯的時候沒限定用了哪些CUDA計算能力, 不指定CUDA架構號(CUDA計算能力/cuda compute capability)會讓編譯後的二進制包很大, 想減少二進制包大小, 見傳送門 https://www.wkwkk.com/articles/4c51566535e88f71.html
