Jetson AGX Orin 64GB : flash-attention 설치

파란크리스마스 2026. 7. 8. 09:10

2026. 7. 8. 09:10

출처

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/Dao-AILab/flash-attention
(.venv) bluesanta@localhost:~/llm$ cd flash-attention
(.venv) bluesanta@localhost:~/llm/flash-attention$ git checkout v2.8.3
(.venv) bluesanta@localhost:~/llm/pytorch$ git submodule update --init --recursive

환경 설정

(.venv) bluesanta@localhost:~/llm/flash-attention$ export MAX_JOBS=4
(.venv) bluesanta@localhost:~/llm/flash-attention$ export FLASH_ATTN_CUDA_ARCHS=87

빌드

(.venv) bluesanta@localhost:~/llm/flash-attention$ python setup.py bdist_wheel
/home/bluesanta/llm/.venv/lib/python3.12/site-packages/wheel/bdist_wheel.py:4: FutureWarning: The 'wheel' package is no longer the canonical location of the 'bdist_wheel' command, and will be removed in a future release. Please update to setuptools v70.1 or later which contains an integrated version of this command.
  warn(
 
torch.__version__  = 2.12.0a0+git0d62256
 
running bdist_wheel
Guessing wheel URL:  https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3+cu12torch2.12cxx11abiTRUE-cp312-cp312-linux_aarch64.whl

설치

(.venv) bluesanta@localhost:~/llm/flash-attention$ ls dist
flash_attn-2.8.3-cp312-cp312-linux_aarch64.whl
(.venv) bluesanta@localhost:~/llm/flash-attention$ pip install dist/flash_attn-2.8.3-cp312-cp312-linux_aarch64.whl 
Processing ./dist/flash_attn-2.8.3-cp312-cp312-linux_aarch64.whl
Requirement already satisfied: torch in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from flash-attn==2.8.3) (2.12.0a0+git0d62256)
Requirement already satisfied: einops in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from flash-attn==2.8.3) (0.8.2)
Requirement already satisfied: filelock in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from torch->flash-attn==2.8.3) (3.29.6)
Requirement already satisfied: typing-extensions>=4.10.0 in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from torch->flash-attn==2.8.3) (4.16.0)
Requirement already satisfied: setuptools<82 in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from torch->flash-attn==2.8.3) (80.10.2)
Requirement already satisfied: sympy>=1.13.3 in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from torch->flash-attn==2.8.3) (1.14.0)
Requirement already satisfied: networkx>=2.5.1 in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from torch->flash-attn==2.8.3) (3.6.1)
Requirement already satisfied: jinja2 in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from torch->flash-attn==2.8.3) (3.1.6)
Requirement already satisfied: fsspec>=0.8.5 in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from torch->flash-attn==2.8.3) (2026.6.0)
Requirement already satisfied: mpmath<1.4,>=1.1.0 in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from sympy>=1.13.3->torch->flash-attn==2.8.3) (1.3.0)
Requirement already satisfied: MarkupSafe>=2.0 in /home/bluesanta/llm/.venv/lib/python3.12/site-packages (from jinja2->torch->flash-attn==2.8.3) (3.0.3)
Installing collected packages: flash-attn
Successfully installed flash-attn-2.8.3

설치 확인

test_flash_attn.py

import torch
from flash_attn import flash_attn_func

# 检查 PyTorch 的 CUDA 架构支持，这是关键的一步
print(f"PyTorch 编译时支持的 CUDA 架构: {torch.cuda.get_arch_list()}")

# 创建一个简单的测试用例
batch_size, seq_len, num_heads, head_dim = 2, 128, 8, 64
q = torch.randn(batch_size, seq_len, num_heads, head_dim, dtype=torch.float16, device='cuda')
k = torch.randn(batch_size, seq_len, num_heads, head_dim, dtype=torch.float16, device='cuda')
v = torch.randn(batch_size, seq_len, num_heads, head_dim, dtype=torch.float16, device='cuda')

# 执行 Flash Attention 函数
output = flash_attn_func(q, k, v)
print(f"测试成功！输出张量形状: {output.shape}")

실행

(.venv) bluesanta@localhost:~/llm$ python test_flash_attn.py
PyTorch 编译时支持的 CUDA 架构: ['sm_87']
测试成功！输出张量形状: torch.Size([2, 128, 8, 64])

저작자표시 (새창열림)

파란크리스마스

Jetson AGX Orin 64GB : flash-attention 설치

+ Recent posts

티스토리툴바