RWKV cuda kernel loading
Author: Hannibal046
Pull request: no

### System Info

- `transformers` version: 4.31.0.dev0
- Platform: Linux-5.4.0-147-generic-x86_64-with-glibc2.17
- Python version: 3.8.16
- Huggingface_hub version: 0.14.1
- Safetensors version: 0.3.1
- PyTorch version (GPU?): 2.0.0+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no


### Who can help?

cc @younesbelkada  @ArthurZucker @sgugger 

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

my `demo.py`
```
from transformers import AutoTokenizer, RwkvModel
import torch
device = torch.device("cuda:5")

tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")
model = RwkvModel.from_pretrained("RWKV/rwkv-4-169m-pile").to(device)

inputs = tokenizer("Hello, my dog is cute", return_tensors="pt").to(device)
outputs = model(**inputs)

last_hidden_states = outputs.last_hidden_state
```

I commented out try except in `modeling_rwkv.py` to force it to load cuda kernel for RWKV attention. And i got this:
```
python demo.py
Traceback (most recent call last):
  File "/data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1893, in _run_ninja_build
    subprocess.run(
  File "/data/chengxin/anaconda3/envs/rwkv/lib/python3.8/subprocess.py", line 516, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['ninja', '-v']' returned non-zero exit status 1.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "demo.py", line 6, in <module>
    model = RwkvModel.from_pretrained("RWKV/rwkv-4-169m-pile").to(device)
  File "/data/chengxin/rwkv/transformers/src/transformers/modeling_utils.py", line 2675, in from_pretrained
    model = cls(config, *model_args, **model_kwargs)
  File "/data/chengxin/rwkv/transformers/src/transformers/models/rwkv/modeling_rwkv.py", line 604, in __init__
    self.blocks = nn.ModuleList([RwkvBlock(config, layer_id=idx) for idx in range(config.num_hidden_layers)])
  File "/data/chengxin/rwkv/transformers/src/transformers/models/rwkv/modeling_rwkv.py", line 604, in <listcomp>
    self.blocks = nn.ModuleList([RwkvBlock(config, layer_id=idx) for idx in range(config.num_hidden_layers)])
  File "/data/chengxin/rwkv/transformers/src/transformers/models/rwkv/modeling_rwkv.py", line 378, in __init__
    self.attention = RwkvSelfAttention(config, layer_id)
  File "/data/chengxin/rwkv/transformers/src/transformers/models/rwkv/modeling_rwkv.py", line 263, in __init__
    load_wkv_cuda_kernel(config.context_length)
  File "/data/chengxin/rwkv/transformers/src/transformers/models/rwkv/modeling_rwkv.py", line 87, in load_wkv_cuda_kernel
    rwkv_cuda_kernel = load_kernel(
  File "/data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1284, in load
    return _jit_compile(
  File "/data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1509, in _jit_compile
    _write_ninja_file_and_build_library(
  File "/data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1624, in _write_ninja_file_and_build_library
    _run_ninja_build(
  File "/data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/utils/cpp_extension.py", line 1909, in _run_ninja_build
    raise RuntimeError(message) from e
RuntimeError: Error building extension 'wkv_1024': [1/3] /usr/bin/nvcc  -DTORCH_EXTENSION_NAME=wkv_1024 -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/TH -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/THC -isystem /data/chengxin/anaconda3/envs/rwkv/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 --compiler-options '-fPIC' -res-usage --maxrregcount 60 --use_fast_math -O3 -Xptxas -O3 --extra-device-vectorization -DTmax=1024 -std=c++17 -c /data/chengxin/rwkv/transformers/src/transformers/kernels/rwkv/wkv_cuda.cu -o wkv_cuda.cuda.o
FAILED: wkv_cuda.cuda.o
/usr/bin/nvcc  -DTORCH_EXTENSION_NAME=wkv_1024 -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/TH -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/THC -isystem /data/chengxin/anaconda3/envs/rwkv/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 --compiler-options '-fPIC' -res-usage --maxrregcount 60 --use_fast_math -O3 -Xptxas -O3 --extra-device-vectorization -DTmax=1024 -std=c++17 -c /data/chengxin/rwkv/transformers/src/transformers/kernels/rwkv/wkv_cuda.cu -o wkv_cuda.cuda.o
nvcc fatal   : Unknown option '-extra-device-vectorization'
[2/3] /usr/bin/nvcc  -DTORCH_EXTENSION_NAME=wkv_1024 -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/TH -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/THC -isystem /data/chengxin/anaconda3/envs/rwkv/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 --compiler-options '-fPIC' -res-usage --maxrregcount 60 --use_fast_math -O3 -Xptxas -O3 --extra-device-vectorization -DTmax=1024 -std=c++17 -c /data/chengxin/rwkv/transformers/src/transformers/kernels/rwkv/wkv_cuda_bf16.cu -o wkv_cuda_bf16.cuda.o
FAILED: wkv_cuda_bf16.cuda.o
/usr/bin/nvcc  -DTORCH_EXTENSION_NAME=wkv_1024 -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1011\" -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/TH -isystem /data/chengxin/anaconda3/envs/rwkv/lib/python3.8/site-packages/torch/include/THC -isystem /data/chengxin/anaconda3/envs/rwkv/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_86,code=compute_86 -gencode=arch=compute_86,code=sm_86 --compiler-options '-fPIC' -res-usage --maxrregcount 60 --use_fast_math -O3 -Xptxas -O3 --extra-device-vectorization -DTmax=1024 -std=c++17 -c /data/chengxin/rwkv/transformers/src/transformers/kernels/rwkv/wkv_cuda_bf16.cu -o wkv_cuda_bf16.cuda.o
nvcc fatal   : Unknown option '-extra-device-vectorization'
ninja: build stopped: subcommand failed.

```

### Expected behavior

Load cuda kernel successfully.

Comments:
I found the problem was about the `nvcc` and `cuda`  mismatch. The installation of `cudatoolkit` from `conda` doesn't necessarily download all relevant cuda things(e.g. `nvcc`).

So I solve it on my side by `conda install cuda -c nvidia`

Awesome thanks for sharing the solution!