Flash Attention 3¶

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

动机¶

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

FA3 对以下场景至关重要：

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

功能对比¶

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

与 GPU 实现的差异¶

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

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

硬件要求¶

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

软件要求¶

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

安装¶

要安装 flash_attn_npu wheel 包，请参阅：https://github.com/MinghuasLab/flash-attention-npu/blob/trainInferConsist/README.md#installation（trainInferConsist 分支）。

启用 Flash Attention 3¶

通过 RL 配置启用 FA3，将 rl_config.enabled 和 rl_config.enable_training_consistency 设置为 true。

批不变模式独立于 FA3。如果也需要该模式，请将 rl_config.enable_batch_invariant 设置为 true。使用此配置路径时，无需手动设置 VLLM_BATCH_INVARIANT、HCCL_DETERMINISTIC 或 LCCL_DETERMINISTIC。

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

要启动启用 FA3 的 vLLM 服务器：

vllm serve Qwen/Qwen3-8B \
  --additional-config '{"rl_config": {"enabled": true, "enable_training_consistency": true}}' \
  --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 的离线批量推理：

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,
    additional_config={
        "rl_config": {
            "enabled": True,
            "enable_training_consistency": True,
        },
    },
    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 不支持滑动窗口注意力。需要滑动窗口的模型必须使用默认 FIA 后端。
• 不支持 ACL 图捕获：flash_attn_with_kvcache 的分块处理在主机侧进行，目前不支持 ACL 图捕获。启用 FA3 时请使用 compilation_config={"cudagraph_mode": "PIECEWISE"}。
• 不支持 RoPE：FA3 不支持在注意力内核中使用旋转位置编码。vLLM-Ascend 通过改用 PyTorch 原生 RoPE 回退方案来修补此问题。
• 不支持 ALiBi：FA3 不支持 ALiBi（带线性偏置的注意力）。
• 不支持 Softcapping：FA3 不支持注意力 logit 软截断。
• 不支持 FP8 量化：FA3 不支持 FP8 量化注意力。
• 不支持 MLA 和 SFA：FA3 不支持多头潜在注意力（MLA）或稀疏 Flash Attention（SFA）。

已测试模型¶

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 系列
• 扩大模型覆盖范围
• 性能优化
• 额外的测试和验证