rotary-embedding-torch

收录理由

旋转位置编码如今已是现代Transformer模型处理序列顺序的主流做法，而这个库把这项技术打包成一个轻巧、专注的PyTorch工具，你可以直接把它嵌进自己的注意力代码里。用的时候只需实例化一个RotaryEmbedding，对查询和键做旋转，原有的注意力循环照旧，不必为了加RoPE而引入整套框架。它把实践中容易踩坑的边界情况都考虑到了：推理时遇到键值缓存需要偏移位置、视频或图像Transformer要用的轴向嵌入，以及通过XPos做长度外推。另外还提供了一条融合的Flash Attention路径，能在单个内核里完成旋转，没有Triton时也有参考实现兜底。如果你在写自定义注意力，或者想试验不同的位置编码方案，这个库是个边界清晰的基础构件，不会拖进笨重的依赖。