CUDA error: no kernel image is available for execution on the device

[NigelNelson]

I am attempting to use Exllama on a unique device. My platform is aarch64 and I have a NVIDIA A6000 dGPU. As such, the only compatible torch 2.1.0 build I can find is one for Python 3.8. The only issues that arose from this were Python 3.10 type checks failing, which I fixed by importing the relevant types from typing. My issue is when the torch extension is built, all I see in ~/.cache/torch_extensions/ is py38_cu114 not py38_cu118 which is what I was expecting. Does anyone know how these extensions are built so I can potentially hack it for my use case? Torch.cuda.is_available() returns True.
python==3.8.17
torch==2.1.0a0+41361538.nv23.06
cuda==11.8

Stacktrace:

 ## Warning: this project requires Python 3.9 or higher.

 -- Perplexity:
 -- - Dataset: datasets/wikitext2_val_sample.jsonl
 -- - Chunks: 100
 -- - Chunk size: 2048 -> 2048
 -- - Chunk overlap: 0
 -- - Min. chunk size: 50
 -- - Key: text
 -- Tokenizer: /media/m2/models/Llama-2-7B-GPTQ/tokenizer.model
 -- Model config: /media/m2/models/Llama-2-7B-GPTQ/config.json
 -- Model: /media/m2/models/Llama-2-7B-GPTQ/gptq_model-4bit-128g.safetensors
 -- Sequence length: 2048
 -- Tuning:
 -- --sdp_thd: 8
 -- --matmul_recons_thd: 8
 -- --fused_mlp_thd: 2
 -- Options: ['perf', 'perplexity']
Traceback (most recent call last):
  File "exllama_inf.py", line 125, in <module>
    model = timer("Load model", lambda: ExLlama(config))
  File "exllama_inf.py", line 56, in timer
    ret = func()
  File "exllama_inf.py", line 125, in <lambda>
    model = timer("Load model", lambda: ExLlama(config))
  File "/home/ubuntu/miniconda3/envs/new_tina/lib/python3.8/site-packages/exllama/model.py", line 862, in __init__
    inv_freq = 1.0 / (self.config.rotary_embedding_base ** (torch.arange(0, self.config.head_dim, 2, device = device).float() / self.config.head_dim))
RuntimeError: CUDA error: no kernel image is available for execution on the device
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.```

[turboderp]

I could suggest setting verbose = True at the top of cuda_ext.py to get the output from nvcc. Maybe that'll highlight why it's building the extension for cu114 instead of cu118. PyTorch 2.1.0 should be fine, I've tested on that, but your specific build of it, I don't know.

[NigelNelson]

Thank you for the suggestion. Unfortunately this only allowed me to confirm that the cuda version used by the NVIDIA Developer Torch download is 11.4. So, I believe the issue is that I have Cuda Toolkit 11.8, but the Cuda version used by torch is 11.4. I guess I'm stuck waiting for the proper wheel to be released. Has anybody had success building Pytorch from source on aarch64?

[turboderp]

Maybe you could try installing CUDA 11.4. I haven't tested any of the CUDA code on that version, but it might work?