Flash Attention 3¶

Note

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

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

动机¶

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

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

训练-推理一致性：确保推理时的注意力计算与训练侧一致，这对于强化学习工作流（例如 veRL）至关重要，因为推理结果用于计算训练信号。
框架调试：一致的注意力实现消除了训练和推理之间的差异，使问题调试更加容易。
强化学习 (RL)：RL 训练通常需要确定性和一致的推演，以实现可重复性和稳定训练。

特性对比¶

下表比较了GPU FA3和Ascend NPU FA3中flash_attn_with_kvcache的功能：

特性	GPU FA3	NPU FA3
FP16 (float16)	✅	✅
BF16 (bfloat16)	✅	✅
Causal Attention	✅	✅
Sliding Window Attention	✅	-
MQA/GQA	✅	✅
Paged KV Cache	✅	✅
Rotary Position Embedding (RoPE)	✅	-
ALiBi	-	-
Softcapping	✅	-
FP8 Quantization	✅	-
Variable-length Sequences	✅	✅

与 GPU 实现的差异¶

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

NPU FA3 不支持的特性：尚不支持滑动窗口注意力、RoPE、ALiBi、Softcapping 和 FP8 量化。
图捕获：flash_attn_with_kvcache 的 tile 划分在主机端处理，目前正在优化中。它不支持 ACL 图捕获（即无法捕获到计算图中进行加速）。启用 FA3 时，请使用 compilation_config={"cudagraph_mode": "PIECEWISE"}。

硬件要求¶

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

软件要求¶

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

安装¶

安装flash_attn_npu wheel包，请参考：https://github.com/MinghuasLab/flash-attention-npu/blob/main/README.md#installation.

启用 Flash Attention 3¶

要启用FA3，您需要：

设置环境变量 export VLLM_BATCH_INVARIANT=1 以启用批次不变模式
通过 LLM 参数 FLASH_ATTN 将注意力后端指定为 attention_backend="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 的 tile 划分在主机端处理，目前不支持 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) 或稀疏闪存注意力 (SFA)。

Note

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

已测试模型¶

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

Qwen3 (密集模型)：Qwen/Qwen3-0.6B、Qwen/Qwen3-1.7B、Qwen/Qwen3-8B
Qwen3 (MoE)：Qwen/Qwen3-30B-A3B

其他模型尚未测试，未来将在测试后根据结果决定是否支持。

未来改进¶

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

开源 flash_attn_npu 包
支持 ACL 图捕获（主机端 tile 划分优化）
支持更多 NPU 系列
扩展模型覆盖范围
性能优化
更多测试与验证