Flash Attention 3

Flash Attention 3#

备注

昇腾上的Flash Attention 3目前处于测试阶段。FA3所需的flash_attn_npu包已在GitHub上开源。更多详情请参考flash-attention-npu仓库。

本文档介绍如何在vLLM-Ascend中启用Flash Attention 3 (FA3)。FA3为昇腾NPU提供了训练-推理一致的注意力实现。

动机#

在veRL等RL训练框架中，训练期间的注意力计算使用Flash Attention。当vLLM-Ascend作为推理后端时，默认的融合推理注意力（FIA）实现与训练侧的Flash Attention不同，可能导致训练-推理不一致。为解决此问题，vLLM-Ascend引入了FA3注意力后端以保持与训练侧的一致性。

FA3在以下场景中至关重要：

  • 训练-推理一致性：确保推理期间的注意力计算与训练侧一致，这对RL工作流（如veRL）至关重要，因为推理结果用于计算训练信号。

  • 框架调试：一致的注意力实现通过消除训练和推理之间的差异，使问题调试更加容易。

  • 强化学习（RL）：RL训练通常需要确定性和一致的推演，以确保可复现性和稳定训练。

特性对比#

下表比较了GPU FA3和昇腾NPU FA3中flash_attn_with_kvcache的特性：

特性

GPU FA3

NPU FA3

FP16 (float16)

✅

✅

BF16 (bfloat16)

✅

✅

因果注意力

✅

✅

滑动窗口注意力

✅

-

MQA/GQA

✅

✅

分页KV缓存

✅

✅

旋转位置编码（RoPE）

✅

-

ALiBi

-

-

Softcapping

✅

-

FP8量化

✅

-

变长序列

✅

✅

与GPU实现的差异#

NPU上的flash_attn_with_kvcache接口在API参数方面与GPU FA3版本语义一致。主要差异如下：

  1. NPU FA3不支持的特性：滑动窗口注意力、RoPE、ALiBi、Softcapping和FP8量化暂不支持。

  2. 图捕获：flash_attn_with_kvcache的tiling在主机侧处理，目前正在优化中。不支持ACL图捕获（即无法捕获到计算图中进行加速）。启用FA3时请使用compilation_config={"cudagraph_mode": "PIECEWISE"}。

硬件要求#

FA3 目前需要昇腾 Atlas A2 和 A3 推理 NPU。未来我们将支持其他 NPU。

软件要求#

FA3需要flash_attn_npu包，该包提供包含flash_attn_with_kvcache算子的flash_attn_npu_v3模块。

安装#

要安装 flash_attn_npu wheel 包，请参阅：MinghuasLab/flash-attention-npu。

启用Flash Attention 3#

要启用FA3，您需要：

  1. 设置环境变量export VLLM_BATCH_INVARIANT=1以启用批处理不变模式

  2. 通过LLM参数attention_backend="FLASH_ATTN"将注意力后端指定为FLASH_ATTN

在线推理（服务器模式）#

启动启用FA3的vLLM服务器：

VLLM_BATCH_INVARIANT=1 vllm serve Qwen/Qwen3-8B \
  --attention-backend FLASH_ATTN \
  --compilation-config '{"cudagraph_mode": "PIECEWISE"}'

然后使用兼容OpenAI的客户端：

from openai import OpenAI

client = OpenAI(
    api_key="EMPTY",
    base_url="http://localhost:8000/v1",
)

response = client.completions.create(
    model="Qwen/Qwen3-8B",
    prompt="The future of AI is",
    max_tokens=100,
    temperature=0.7,
    seed=42,
)

print(response.choices[0].text)

离线推理#

使用FA3进行离线批量推理：

import os
os.environ["VLLM_BATCH_INVARIANT"] = "1"

from vllm import LLM, SamplingParams

prompts = [
    "The future of AI is",
    "Machine learning enables",
    "Deep learning models can"
]

sampling_params = SamplingParams(
    temperature=0.7,
    max_tokens=100,
    seed=42,
)

llm = LLM(
    model="Qwen/Qwen3-8B",
    tensor_parallel_size=1,
    attention_backend="FLASH_ATTN",
    compilation_config={"cudagraph_mode": "PIECEWISE"},
)

outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt: {prompt!r}")
    print(f"Generated: {generated_text!r}\n")

限制#

  • 包尚未开源：FA3所需的flash_attn_npu包尚未发布。在包可用之前，外部用户无法使用FA3。

  • 不支持滑动窗口：FA3不支持滑动窗口注意力。需要滑动窗口的模型需使用默认的FIA后端。

  • 不支持ACL图捕获：flash_attn_with_kvcache的tiling在主机侧处理，目前不支持ACL图捕获。启用FA3时请使用compilation_config={"cudagraph_mode": "PIECEWISE"}。

  • 不支持RoPE：FA3不支持注意力内核内的旋转位置编码。vLLM-Ascend通过使用PyTorch原生RoPE回退来修补此问题。

  • 不支持ALiBi：FA3不支持ALiBi（线性偏置注意力）。

  • 不支持Softcapping：FA3不支持注意力logit softcapping。

  • 不支持FP8量化：FA3不支持FP8量化注意力。

  • 不支持MLA和SFA：FA3不支持多头潜在注意力（MLA）或稀疏Flash注意力（SFA）。

备注

启用FA3可能导致性能相比默认的FIA后端有所下降。这种权衡是有意为之，旨在保证训练与推理的一致性。

已测试模型#

FA3已在以下模型上完成测试和验证：

  • Qwen3（密集模型）：Qwen/Qwen3-0.6B、Qwen/Qwen3-1.7B、Qwen/Qwen3-8B

  • Qwen3（MoE）：Qwen/Qwen3-30B-A3B

其他模型尚未测试，若测试后不支持，将在未来版本中提供支持。

未来改进#

FA3功能正在积极开发中。计划中的改进包括：

  • 开源flash_attn_npu包

  • 支持ACL图捕获（主机端分块优化）

  • 支持更多NPU系列

  • 扩展模型覆盖范围

  • 性能优化

  • 更多测试与验证