# DGX Spark 上 FlashAttention 2 模型推理安装实录


## 安装背景 {#sec-5d20be071b1c}

在安装 MOSS-TTS 后，其推理速度已足以满足日常需求。然而，听说安装 FlashAttention 2 可以显著提升推理速度并减少 GPU 内存占用，这让我非常好奇实际效果究竟如何。因此，我决定着手安装 FlashAttention 2。

---

## FlashAttention 2 是什么？它如何提升推理效率？ {#sec-c62db22593c2}

据我理解，FlashAttention 2 是 Transformer 架构中用于更高效处理 Attention 运算的一种实现。我推测，它通过减少 Attention 计算过程中的内存访问和中间张量生成，或优化运算流程，从而提升速度和内存效率。不过，我判断这种效果会因模型结构、输入长度、数据类型 (torch.float16 / torch.bfloat16) 和 GPU 架构的不同而有所差异。

---

## 在 DGX Spark 上安装 FlashAttention 2 {#sec-7816f5cc040d}

参考 MOSS-TTS 的 README 文件，其中有这样一段话：

> FlashAttention 2 is only available on supported GPUs and is typically used with torch.float16 or torch.bfloat16

既然可以在 torch.float16 环境中使用，我判断 DGX Spark 应该也适用，于是尝试进行安装。


![flash-attn2-on-spark](/media/whitedec/blog_img/bbce3ba204dd4ea298238a79143832f5.webp)

---

### 1. 前期准备 {#sec-b408be4ef199}

*   DGX Spark 使用 CUDA 13.0，因此我从 `--extra-index-url https://download.pytorch.org/whl/cu130` 查找并安装了依赖包。为了确认，我还通过 `nvidia-smi` 检查了 CUDA 版本。

*   安装过程中，由于 PyTorch 已经在现有的 venv 环境中安装，我需要继续使用它。在 venv 环境中进行构建时，为了避免创建临时的隔离环境，我加入了 `--no-build-isolation` 选项。

*   在 Spark 环境中，wheel 包安装失败，安装日志中输出以下信息。这通常是由于 aarch64 架构造成的。在使用 Spark 时，这种情况屡见不鲜，我已经见怪不怪了，甚至不再感到恼火，这已是熟悉的提示。

    > `Precompiled wheel not found. Building from source...`

    别无选择，只能通过源码构建。源码构建过程中需要 `ninja`，因此将其安装到 venv 中：

    ```bash
    pip install ninja
    ```

*   主机系统需要 Python 3.12 的开发库，如果您的系统尚未安装，请执行以下命令：

    ```bash
    sudo apt update
    sudo apt install python3.12-dev -y
    ```

    `flash-attn` 在编译 C++ 和 CUDA 代码并将其与 Python 连接时，需要 `Python.h` 文件，该文件定义了 Python 的内部结构。通常的 Python 运行环境中不包含此文件，因此需要单独安装开发者包。

---

### 2. 安装命令 {#sec-8cb312ad923c}

**这是本文的核心部分。** 综合考虑上述所有事项，我使用了以下命令进行安装：

```bash
TORCH_CUDA_ARCH_LIST="12.0" MAX_JOBS=1 pip install --no-build-isolation --extra-index-url https://download.pytorch.org/whl/cu130 -e ".[flash-attn]"
```

---

### 3. 命令组合原因 (基于试错的决定) {#sec-085ee9786e8d}


最初，机器上还在运行 gpt-oss-120b。在这种情况下执行 `pip install ... -e ".[flash-attn]"` 后，CPU 使用率急剧飙升，系统随之卡死，终端也失去了响应，最终我不得不通过物理开关强制重启。此后，我暂停了所有占用资源的任务，专注于安装。

经过多次尝试，最终我通过上述命令成功完成了安装。整个过程大约耗时 1 到 2 小时。虽然没有精确测量，但安装持续了一个多小时，我离开去用餐，回来时发现安装已经完成。

安装期间，内存占用持续在 24GB 左右。主要问题在于 CPU，为了确保安装的稳定性，最好暂停其他任务。

加入这些选项的原因如下：

*   `TORCH_CUDA_ARCH_LIST="12.0"`: 明确指定以 Blackwell 架构为目标，旨在缩短安装时间。
*   `MAX_JOBS=1`: 鉴于之前系统卡死的经验，我保守地将其设置为 1。结果导致安装时间超过 60 分钟。

---

## 安装后的推理效果提升 {#sec-e872dc211143}


### 1. 速度方面 {#sec-aafe06302a64}

老实说，速度方面并没有显著的体感提升。在未安装 `flash-attn` 时，速度已经足够快，即使缩短了几秒钟，也未能给我“确实变快了”的强烈感受。

*   生成约 7 秒的结果，耗时约 8-9 秒。
*   生成约 25 秒的结果，耗时约 32 秒。
*   生成约 16 秒的结果，耗时 21 秒。

也就是说，推理时间大约是生成结果长度的 1.3 倍左右。

### 2. 内存方面 {#sec-5789cff200c8}

内存方面也难以察觉到变化。在推理过程中，我持续观察 `nvidia-smi` 的数值，但内存使用量并没有明显增加或减少。推理期间，功耗维持在 36W 左右，温度从 46 度上升到 53 度左右。

---

## 总结 {#sec-5ac1d2a40659}

*   在 DGX Spark 环境中，FlashAttention 2 的 wheel 包安装失败，最终通过源码构建完成。
*   安装虽然成功，但构建时间较长，且 CPU 负载相当高。
*   安装后，速度和内存方面的实际体验提升未达到预期。
*   如果下次再次构建，我会考虑将 `MAX_JOBS` 提升到 `4` 左右。理论上，构建时间有望缩短至原先的四分之一。


**相关文章**

- [파인튜닝 없이 고성능 TTS, MOSS-TTS NVIDIA DGX Spark 설치 가이드](/ko/whitedec/2026/3/10/moss-tts-dgx-spark-installation-guide/)