FasterTransformer Longformer 推理加速实践：PyTorch 集成、输入约束与端到端性能解析

• 推理引擎
• 算子库
• 大模型

Longformer 通过"局部滑动窗口注意力 + 任务驱动的全局注意力"将自注意力的复杂度从 O(n²) 降为 O(n)，是长文档建模（数千 token 甚至更长）的主流方案之一。本文以 docs/longformer_guide.md 为核心，完整梳理 FasterTransformer 中 Longformer Encoder 的实现范围、输入约束、计算流程，给出基于 PyTorch OP 的 QA 推理 Demo 全流程（环境搭建、构建、运行与参数说明），并结合 src/fastertransformer/models/longformer/LongformerEncoder.cc 等源码与官方 FP32/FP16 基准，解析其加速原理与实战要点。读完本文，你将能够在 FasterTransformer 仓库内复现 Longformer QA 加速推理，理解其与 HuggingFace 实现的替换方式及性能差异。

Longformer 模型与 FasterTransformer 的实现定位

Longformer 出自论文《Longformer: The Long-Document Transformer》。标准 Transformer 的自注意力操作复杂度随序列长度呈平方增长，导致无法处理长序列；Longformer 引入一种随序列长度线性扩展的注意力机制，易于处理数千 token 乃至更长的文档。其注意力机制是标准自注意力的"即插即用"替代品，组合了局部窗口注意力（local windowed attention）与任务驱动的全局注意力（task motivated global attention），在字符级语言建模（text8、enwik8）以及 WikiHop、TriviaQA 等长文档任务上取得了当时领先的结果。

FasterTransformer 中的 Longformer 实现与 HuggingFace 的 Longformer 对齐（longformer-large-4096-finetuned-triviaqa等模型权重可直接加载）。从仓库支持矩阵（见 README.md）看，Longformer 在 FasterTransformer 中仅在PyTorch框架下提供FP16支持（底层 CUDA/C++ 实现的精度还包括 FP32 与 BF16，见下文）。

从 src/fastertransformer/models/longformer/LongformerEncoder.h 的模板实例化可知，LongformerEncoder支持float、half以及（在ENABLE_BF16编译开关下）__nv_bfloat16三种数据类型：

template class LongformerEncoder<float>; template class LongformerEncoder<half>; #ifdef ENABLE_BF16 template class LongformerEncoder<__nv_bfloat16>; #endif

目前仓库中只实现了 Longformer 的 Encoder 部分，并配套一个真实场景的"问答（Question and Answering）"示例来同时演示正确性与加速效果。LongformerEncoder以 CUDA 和 C++ 实现，并通过 PyTorch 自定义 OP 封装（src/fastertransformer/th_op/longformer/LongformerEncoderOp.h）供用户平滑集成。

使用前提：四条硬性约束

在动手前，必须先满足原文档列出的 4 条约束，否则运行会失败或结果错误：

1. seq_len必须>= 2 * local_attn_window_size，且seq_len % local_attn_window_size == 0。这条约束在 LongformerAttentionLayer.cc 中通过FT_CHECK(seq_len % local_attn_window_size_ == 0)强制校验，在 model.py 的 Python 侧也有对应断言。
2. head_dim（即size_per_head）目前只支持64。同样在 LongformerAttentionLayer.cc 中由FT_CHECK(size_per_head_ == 64)保证。
3. max_global_token_num应< local_attn_window_size。注意源码侧的实际检查为FT_CHECK(max_global_token_num_ <= local_attn_window_size_)（见 LongformerAttentionLayer.cc），即"小于等于"也能通过校验；原文档建议严格小于，以留出安全余量。
4. 输入序列中，全局 token（global_tokens）必须连续地放在序列开头。这是 FasterTransformer 侧实现的前提假设，与 HuggingFace 中"任意位置标记全局 token"不同，使用时需要自行把全局 token 移动到序列头部。

Longformer Encoder 计算流程

原文档给出的计算流程如下：

图 1：Longformer 计算流程图（来自 docs/longformer_guide.md）

从源码层面可以还原该流程的每一步。LongformerEncoder::forward（LongformerEncoder.cc）接收三类输入：

• input：形状batch_size x seq_len x in_dim；
• local_attn_mask：形状batch_size x seq_len，0.0表示不做局部注意力，1.0表示做局部注意力；
• global_attn_mask：形状batch_size x seq_len，-10000.0表示非全局 token，0.0表示全局 token。

前向过程分两个预处理 kernel 与逐层循环：

预处理（longformer_kernels.cu）

• invokeInitLongformerIdx：基于全局注意力掩码，通过cub::DevicePartition::Flagged提取每个 batch 中全局 token 的索引global_idx与数量global_token_nums，供后续全局注意力按索引取值；
• invokeLocalAttnMaskShift：将局部注意力掩码经fma(mask, 10000.f, -10000.f)变换为注意力打分可加的偏移量（对应 Python 侧extended_mask_b = (global_attn_mask_b + dummy_local_attn_mask_b) * 10000. - 10000.的做法）。

每个 Transformer 层（forwardLayer，LongformerEncoder.cc）

1. 用cublasMMWrapper::Gemm计算 Q、K、V，以及全局路径的kg（global K）、vg（global V）五个投影；全局 Q（qg）因只对max_global_token_num_个 token 计算，用stridedBatchedGemm单独处理；
2. invokeAddBiasTransposeToMultiHead把 Q/K/V/global-K/global-V 五组结果加 bias 并转置为多头的[batch, head, seq, dim]布局；
3. LongformerAttentionLayer::forward完成局部滑动窗口注意力（序列被切成 head/tail 与 middle 三段，对应不同 stride 的窗口 GEMM，见 LongformerAttentionLayer.cc 中的buf_sizes/buf_strides设计）与全局注意力；
4. invokeLongformerMHASoftmax（longformer_kernels.cu 中的longformerMHASoftmaxKernel）对局部 + 全局注意力打分做带 mask 的 softmax：局部 token 只 softmax 窗口内及其前方的全局 token，全局 token 则对全序列 softmax，且当mask_blk[row_idx] != 0时按 Python 实现语义整体置零；
5. invokeTransposeMultiHeadToSingle把多头注意力输出还原，再经attention_output_weights线性层、AddBiasResidualLayerNorm；
6. GeluFfnLayer完成 FFN（intermediate GEMM + bias + GELU + output GEMM），最后再一次AddBiasResidualLayerNorm输出该层结果。

整体呈现出"局部窗口 GEMM 复用共享内存缓冲 + 全局 token 索引化 gather + 融合 softmax"的 CUDA 优化形态，这也是相对 HuggingFace 朴素实现获得加速的核心所在。

Demo：以 PyTorch OP 运行 Longformer QA

仓库提供了问答示例脚本 examples/pytorch/longformer/longformer_qa.py。它会先用随机输入对比 HuggingFace 与 FasterTransformer 的 Longformer Encoder 延迟，再分别解码输出答案并打印，用于同时验证性能与正确性。

环境要求

• CMake >= 3.13（用于编译 PyTorch 扩展）；
• CUDA 11.0 或更高版本；
• Python 3（仅验证过 Python 3）；
• PyTorch：验证版本 1.8.0，>= 1.5.0应可工作。

推荐直接使用 NVIDIA NGC 的 PyTorch Docker 镜像（原文档示例使用nvcr.io/nvidia/pytorch:22.09-py3），并确保宿主机具备基于 Pascal、Volta、Turing 或 Ampere 架构的 GPU。

环境搭建

1. 启动 Docker 容器并挂载项目目录：

docker run \ -it \ --shm-size 5g \ --rm \ --ipc=host \ --gpus=all \ -v {YOUR_FASTER_TRANSFORMER_PROJECT_DIR_ON_HOST}:/workspace/FasterTransformer \ --workdir /workspace/FasterTransformer \ nvcr.io/nvidia/pytorch:22.09-py3 bash

也可以换成其他带 CUDA 的 PyTorch NGC 容器，但需满足上述版本要求。

2. 安装 HuggingFace Transformers（指定版本以对齐依赖）：

pip install transformers==4.9.2

3. 以 PyTorch 模式构建 FasterTransformer：

mkdir build && cd build cmake -DCMAKE_BUILD_TYPE=Release -DBUILD_PYT=ON .. make -j12

4. 下载 HuggingFace Longformer QA 模型（需 git-lfs）：

apt-get update && apt-get install git-lfs git lfs install git clone https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa ../examples/pytorch/longformer/longformer-large-4096-finetuned-triviaqa

运行

cd /workspace/FasterTransformer python3 examples/pytorch/longformer/longformer_qa.py \ --ft-longformer-lib build/lib/libth_transformer.so \ --model-dir examples/pytorch/longformer/longformer-large-4096-finetuned-triviaqa \ --passage "Jim Henson was a nice puppet" \ --question "Who was Jim Henson?" \ --repeat-test-num 50

脚本会先对比 HuggingFace PyTorch 与 FasterTransformer 的 Longformer Encoder 性能，随后打印问答答案。更多命令行参数以python3 longformer_qa.py -h查看，核心参数如下：

脚本其余参数还包括--passage/--passage-file（正文文本二选一）与--model-dir、--ft-longformer-lib等（见 longformer_qa.py）。

FasterTransformer Longformer Encoder PyTorch OP

PyTorch 侧的封装类FTLongformerEncoder定义在 model.py，它通过torch.classes.load_library(ft_longformer_lib)加载libth_transformer.so，再实例化torch.classes.FasterTransformer.LongformerEncoder，其底层对应 LongformerEncoderOp.h 中的FasterTransformerLongformerEncoder。模型权重由from_hf_longformer_weight_to_ft从 HuggingFace 的pytorch_model.bin中按层抽取 q/k/v、global q/k/v、attention output、两处 LayerNorm、intermediate/output 等全部参数并转置、展平、拼接（FP32/FP16/BF16 按需转换），形成连续权重张量——这也正是 LongformerEncoderOp.h 中"q k v kg vg 权重与 bias 必须连续存放"的要求。

构建 FT Longformer Encoder

# 传入必要配置与参数 weights_file = os.path.join(hf_model_dir, 'pytorch_model.bin') ft_encoder = FTLongformerEncoder(weights_file, layer_num, head_num, size_per_head, intermediate_size, local_attn_window_size, max_global_token_num, batch_size, seq_len, attn_scaler, ft_longformer_lib) # 如需无缝替换原 HuggingFace Longformer Encoder： ft_longformer = build_hf_longformer(hf_model_dir) ft_encoder.set_hf_plugin_mode(True) ft_longformer.longformer.encoder = ft_encoder

推理

# 方式一：直接向 FT Encoder 传准备好的输入 ft_encoder.set_hf_plugin_mode(False) output = ft_encoder.forward(input_b, local_attn_mask_b, global_attn_mask_b) # local_attn_mask_b 与 global_attn_mask_b 需与输入等长：1.0 表示有效注意力，0.0 表示无注意力 # 方式二：按 HuggingFace Longformer 的常规用法直接调用 outputs = ft_longformer(input_ids_b, attention_mask=local_attn_mask_b, global_attention_mask=global_attn_mask_b)

hf_plugin_mode为True时，forward会把 HuggingFace 传入的扩展注意力掩码（extended_attention_mask）转换回 FT 需要的local_attn_mask与global_attn_mask，并返回LongformerBaseModelOutput结构，从而对 HuggingFace 代码保持透明（见 model.py）。

性能基准：FP32 与 FP16

官方在 A100 上对比了 HuggingFace Longformer Encoder（不含 embedding 与最终qa_output线性层）与 FT Longformer Encoder 的延迟。硬件环境为 8×A100-80GB（mclk 1593MHz，pclk 1410MHz）+ AMD EPYC 7742 64 核处理器。模型超参数：head_num=16、size_per_head=64、intermediate_size=4096、local_attn_window_size=512、num_layers=24。

FP32

FP16

数据解读（基于上表事实）：FP16 在 batch=1、seq=1024 时加速比最高达 5.71×；FP32 在相同配置下为 3.54×。整体上小 batch、短序列的加速更明显，大 batch 下受 GEMM 效率与访存带宽影响加速比有所回落，但所有测试配置均快于 HuggingFace 基线。该数据为官方原文档在特定软硬件环境下的测试结果，实际部署中应根据自身 GPU 型号与序列长度重新基准。

小结

FasterTransformer 的 Longformer 实现把长文档 Transformer 中最昂贵的自注意力拆解为"局部滑动窗口 GEMM + 全局 token 索引化注意力"，并通过融合 softmax、bias-transpose 等 CUDA kernel 与 PyTorch OP 封装，实现了对 HuggingFace Longformer Encoder 的透明替换与显著加速。使用时的核心要点是遵守四条输入约束（序列长度与窗口的整除关系、head_dim=64、全局 token 数量上限、全局 token 置于序列开头），并按照"构建 lib → 权重转换 → 替换 encoder"的流程集成。需要进一步探索时，可继续阅读 LongformerAttentionLayer.cc 的窗口 GEMM 缓冲布局，或参考仓库通用的模型接入模板 templates/adding_a_new_model/README.md。

• 推理引擎
• 算子库
• 大模型