使用 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 块等模块。