Jetson AGX Orin : flashinfer(v0.6.14) 설치

CUDA 버전 확인

bluesanta@localhost:~$ uname -a
Linux localhost.localdomain 6.8.12-1021-tegra #1 SMP PREEMPT Mon Jun  1 13:25:46 PDT 2026 aarch64 aarch64 aarch64 GNU/Linux
bluesanta@localhost:~$ cat /etc/nv_tegra_release
## R39 (release), REVISION: 2.0, GCID: 45755727, BOARD: generic, EABI: aarch64, DATE: Mon Jun  1 09:28:48 PM UTC 2026
## KERNEL_VARIANT: oot
TARGET_USERSPACE_LIB_DIR=nvidia
TARGET_USERSPACE_LIB_DIR_PATH=usr/lib/aarch64-linux-gnu/nvidia
bluesanta@localhost:~$ nvidia-smi --query-gpu=name,compute_cap,driver_version --format=csv
name, compute_cap, driver_version
Orin (nvgpu), 8.7, 595.78
bluesanta@localhost:~$ nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2026 NVIDIA Corporation
Built on Thu_Mar_19_11:11:41_PM_PDT_2026
Cuda compilation tools, release 13.2, V13.2.78
Build cuda_13.2.r13.2/compiler.37668154_0

빌드 도구 설치

bluesanta@localhost:~$ sudo apt update
bluesanta@localhost:~$ sudo apt install -y cmake ninja-build gcc g++ git build-essential git cmake ninja-build libopenblas-dev libopenmpi-dev openmpi-bin libatlas-base-dev libprotobuf-dev protobuf-compiler libssl-dev zlib1g-dev libffi-dev python3-pip libopenblas-dev ccache git-lfs libjpeg-dev libpng-dev libtiff-dev
bluesanta@localhost:~$ sudo apt install -y sox libsox-dev pkg-config cmake ninja-build

NCCL 설치

bluesanta@localhost:~$ wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/sbsa/cuda-keyring_1.1-1_all.deb
bluesanta@localhost:~$ sudo dpkg -i cuda-keyring_1.1-1_all.deb
bluesanta@localhost:~$ sudo apt update
bluesanta@localhost:~$ sudo apt install libnccl2 libnccl-dev

Python 가상 환경 생성 및 활성화

bluesanta@localhost:~$ cd llm
bluesanta@localhost:~/llm$ python -m venv .venv
bluesanta@localhost:~/llm$ source .venv/bin/activate
(.venv) bluesanta@localhost:~/llm$ python --version
Python 3.12.3
(.venv) bluesanta@localhost:~/llm$ pip install --upgrade pip

flash-attention 설치

(.venv) bluesanta@localhost:~/llm$ git clone https://github.com/flashinfer-ai/flashinfer.git
(.venv) bluesanta@localhost:~/llm$ cd flashinfer
(.venv) bluesanta@localhost:~/llm/flashinfer$ git checkout v0.6.14
(.venv) bluesanta@localhost:~/llm/flashinfer$ git submodule update --init --recursive

환경 설정

(.venv) bluesanta@localhost:~/llm/flashinfer$ export TORCH_CUDA_ARCH_LIST="8.7;8.9;9.0"
(.venv) bluesanta@localhost:~/llm/flashinfer$ export FLASHINFER_ENABLE_AOT=1

빌드

(.venv) bluesanta@localhost:~/llm/flashinfer$ python -m build --wheel

설치

(.venv) bluesanta@localhost:~/llm/flashinfer$ ls dist/
flashinfer_python-0.6.14-py3-none-any.whl
(.venv) bluesanta@localhost:~/llm/flashinfer$ pip install dist/flashinfer_python-0.6.14-py3-none-any.whl

설치 확인

test_flashinfer.py

import torch
import flashinfer

print('='*50)
print(f'✅  PyTorch 버전: {torch.__version__}')
print(f'✅  CUDA 사용 가능 여부: {torch.cuda.is_available()}')
if torch.cuda.is_available():
    print(f'✅  현재 GPU 장치: {torch.cuda.get_device_name(0)}')
    print(f'✅  PyTorch 인식 CUDA 버전: {torch.version.cuda}')

print(f'✅  FlashInfer 버전: {flashinfer.__version__}')

try:
    # 아주 간단한 FlashInfer 모듈을 호출하여 런타임 에러가 없는지 테스트
    workspace = torch.empty(32 * 1024 * 1024, dtype=torch.uint8, device='cuda:0')
    print('✅  FlashInfer CUDA 워크스페이스 할당 테스트 성공! (CUDA 완벽 지원)')
except Exception as e:
    print(f'❌  FlashInfer 동작 테스트 실패: {e}')
print('='*50)

실행

(.venv) bluesanta@localhost:~/llm$ python test_flashinfer.py 
==================================================
✅ PyTorch 버전: 2.12.0a0+git0d62256
✅ CUDA 사용 가능 여부: True
✅ 현재 GPU 장치: Orin
✅ PyTorch 인식 CUDA 버전: 13.2
✅ FlashInfer 버전: 0.6.14
✅ FlashInfer CUDA 워크스페이스 할당 테스트 성공! (CUDA 완벽 지원)
==================================================