Ollama 支持的 flash attention 能提升推理速度吗？我们一起测测看吧

这是 Ollama 支持的 flash attention 能提升推理速度吗？我们一起测测看吧 的笔记哦，查看更详尽的内容，请观看视频，谢谢。

ollama 最近的更新还是蛮频繁的。继上次更新了并发请求之后，最新的版本 0.1.39 则是支持了 flash attention 。

flash attention 可以显著减少注意力机制计算的时间，使得Transformer模型在训练和推理时能够更快地处理大量数据，而且它通过优化内存访问模式，减少了计算过程中的内存占用。这不仅有助于降低内存开销，还减少了计算过程中的内存带宽瓶颈。

这里可以理解为在 attention 计算比较密集的情况下，flash attention 能够显著减少计算量，在推理时能够更快地处理大量数据。那么对应的，上下文越长，flash attention 所带来的优势就越明显。

这里我就对比下开关 flash attention 对推理时间的影响。

环境准备

这里还是使用一块 4090 云 GPU 进行测试。我还是用 llama3 模型。通过命令获取。

ollama pull llama3

然后通过下面的命令可以启动 ollama：

OLLAMA_HOST=0.0.0.0:6006 OLLAMA_NUM_PARALLEL=16 OLLAMA_MODELS=/root/autodl-tmp/models ollama serve

其中：

1. OLLAMA_HOST=0.0.0.0:6006：指定 Ollama 服务监听的主机和端口，0.0.0.0 表示监听所有网络接口上的连接，端口为 6006。
2. OLLAMA_NUM_PARALLEL=16：设置 Ollama 服务的并行处理数量为 16，表示可以同时处理 16 个请求。
3. OLLAMA_MODELS=/root/autodl-tmp/models：指定 Ollama 模型的存储路径为 /root/autodl-tmp/models。

默认 ollama 并没有打开 flash attention 需要我们在启动的时候，增加环境变量 OLLAMA_FLASH_ATTENTION=1 来启动:

OLLAMA_FLASH_ATTENTION=1 \
OLLAMA_NUM_PARALLEL=16 \
OLLAMA_HOST=0.0.0.0:6006 \
OLLAMA_MODELS=/root/autodl-tmp/models \
ollama serve

修改 open webui 的配置，直接使用云端的 ollama

首先，我在本地通过 docker 已经运行了一个 open webui。然后点击左下角 User -> Settings -> Connections ，修改 Ollama Base URL 为云端 ollama 的地址。

然后准备一个特别长的提示词：

下面是一系列的新闻：

---

<此处略过 1000 字>

请你帮我进行分析和处理，做如下事情：

1. 帮我做总结
2. 帮我依据以上内容准备 5 个问题

这里我先插入了一段从 csdn 摘录的新闻，然后让 llama3 帮我进行总结，并准备 5 个问题。这也是目前非常流行的 RAG 的场景：先提供一个上下文，然后基于上下文所提供的信息提问。

然后就可以通过 OLLAMA_FLASH_ATTENTION=1 的有无去对比下实际效果了。

脚本准备

下一步我还是使用之前的脚本对 flash attention 进行测试，但是这次我需要一个更长的上下文。

import aiohttp
import asyncio
import time
from tqdm import tqdm

question = """
下面是一系列的新闻：

---

<此处省略 1000 字，如果你想要使用这个脚本，务必把这里替换为一个很长的新闻稿>

---

请你帮我进行分析和处理，做如下事情：

1. 帮我做总结
2. 帮我依据以上内容准