使用 FlagGems 模块构造自己的模型#
在某些使用场景中，用户可能希望从头构建自己的 AI 模型， 或者对现有的模型进行适配，以更好地满足自己的特定使用场景。 为了支持这种需求，FlagGems 提供一个不断增长的高性能模块集合， 这些模块在大语言模型（LLM）中使用很普遍。
这些组件是使用 FlagGems 加速过的算子实现的，可以像你使用标准的 torch.nn.Module
一样使用。你可以将它们无缝集成到自己的系统重，在不需要编写定制的 CUDA
代码或者 Triton 代码的前提下，从内核级的加速中获益。
FlagGems 所支持的模块代码位于源码仓库的 flag_gems/modules. 目录下。
可用的模块#
| 模块 | 描述 | 支持的特性 |
|---|---|---|
GemsRMSNorm | RMS LayerNorm | 对残差求和进行融合，支持inplace outplace 模式 |
GemsRope | 标准的旋转位置编码 | inplace 和 outplace 模式 |
GemsDeepseekYarnRoPE | 带外推的旋转位置编码，用于 DeepSeek 风格的 LLM | inplace 和 outplace 模式 |
GemsSiluAndMul | SiLU 激活函数与逐元素乘法的融合 | 仅支持 outplace 模式 |
我们鼓励用户将这些模块作为等价 PyTorch 层的替换方案。 团队正在开发融合的注意力机制、MoE 层以及 Transformer 块等模块。