部分代码说明
| 文件名称 | 实现功能 | 文件地址 |
|---|---|---|
MHA | 多头注意力模块,支持 flash_attn,输入数据格式为：x:(B,T,C),atten_mask:(B,T) | 🔗 |
GQA | 分组注意力模块,支持 flash_attn,输入数据格式为：x:(B,T,C),atten_mask:(B,T) | 🔗 |
MQA | 多查询注意力模块,支持 flash_attn,输入数据格式为：x:(B,T,C),atten_mask:(B,T) | 🔗 |
SWA | 滑动窗口注意力模块,支持 flash_attn,输入数据格式为：x:(B,T,C),atten_mask:(B,T) | 🔗 |
MoBA | Kimi MoBA论文稀疏注意力计算,输入数据格式为：x:(B,T,C),atten_mask:(B,T) | 🔗 |
PosEncoding | 位置编码,RotaryPositionalEncoding,AbsolutePositionEmbedding,LearnedPositionEmbedding。输入：x:(B,T,C) | 🔗 |
Norm | 归一化操作,LayerNorm,BatchNorm,RMSNorm,InstanceNorm,GlobalResponseNorm。输入：(B,T,C) 或者 (B,C,H,W) | 🔗 |
ResNet | 视觉编码器,ResNet50, ResNet101, ResNet152系列 | 🔗 |
ConvNeXt | 视觉编码器,ConvNeXt v1系列 | 🔗 |
Vit | 视觉编码器,Vit | 🔗 |
SwinTransformer | 视觉编码器,SwinTransformer | 🔗 |