xLSTM：长短期记忆扩展架构使用指南

NX-AI/xlstm 是面向学术界与工业界的新一代循环神经网络架构实现，当前 Star 数约 2.2k，采用 Apache-2.0 协议，主要语言为 Python。它在经典 LSTM 的思想上做了两个关键升级：一是通过"指数门控"配合适当的归一化与稳定化技巧，二是引入全新的矩阵记忆（Matrix Memory），从而克服了原始 LSTM 难以扩展到大规模语言建模的缺陷，并在与 Transformer、状态空间模型（SSM）的对比中展现出有竞争力的性能。项目最重磅的成果是：团队已经在 2.3 万亿 token 上训练出了一个 70 亿参数的 xLSTM 大语言模型，即优化后的"xLSTM Large"。相关论文有两篇：架构原始论文 xLSTM: Extended Long Short-Term Memory（NeurIPS 2024），以及 7B 模型论文 xLSTM 7B: A Recurrent LLM for Fast and Efficient Inference（ICML 2025），作者包括 LSTM 发明人 Sepp Hochreiter 本人领衔的林茨大学 NX-AI 团队。对关注"后 Transformer 架构"的研究者和工程师来说，xLSTM 是目前为数不多把循环结构真正推到百亿参数级、并拿出完整权重与推理方案的开源工作。它的出现，打破了过去几年"Scaling 只能靠 Transformer"的共识：证明只要把训练稳定性和并行内核做扎实，传统循环结构也能吃到大数据红利。对工程师而言，这意味着长序列场景下可以用固定大小的状态替代随长度膨胀的 KV-cache，部署成本显著下降；对研究者而言，仓库里同时给出了小模型教学版（NeurIPS 架构）与大模型工业版（xlstm_large），可以对照着读。换句话说，这个仓库既是论文配套代码，也是一份可落地的"循环架构大模型"参考实现。

一、项目官方信息

二、官方图片

README 顶部附有 res/desc_xlstm_overview.svg 架构总览图，展示 sLSTM 与 mLSTM 两种变体的门控结构与矩阵记忆如何层层组合；7B 章节另附 res/xlstm_7b_poster.svg 论文海报图，汇总了模型规模、训练数据与推理性能对比。本系列封面使用仓库官方图，两张图分别对应架构论文与 7B 模型发布两个重要节点。

三、环境要求

本包基于 PyTorch 开发，官方测试要求 PyTorch 1.8 及以上；官方最充分验证的环境写在 environment_pt240cu124.yaml 里，直接用 conda 一键创建：conda env create -n xlstm -f environment_pt240cu124.yaml && conda activate xlstm，这份环境固定了 PyTorch 2.4.0 与 CUDA 12.4 的组合，能避免版本不匹配带来的编译问题。要跑 xLSTM Large 7B 模型，还必须单独安装官方的 mlstm_kernels 包，它提供 xLSTM 专用的高性能 Triton/CUDA 算子；没有这个包，7B 模型只能跑慢路径。sLSTM 的 CUDA 内核要求显卡 Compute Capability ≥ 8.0（即 Ampere 及以上，见 NVIDIA 官方列表），老于 30 系的 N 卡无法使用加速内核。建议第一次接触的同学直接用官方 yaml 建环境，不要自己拼版本，否则极容易在编译 Triton 内核时踩坑；conda 环境建好后再装 mlstm_kernels 与 xlstm 两个包，顺序不能反。如果机器没有 GPU，也可以先纯 CPU 跑通 demo，只是前向速度会慢到只能做正确性验证。

四、完整部署步骤

1. 安装

xLSTM 的安装分两步：先装 mlstm_kernels（7B 模型的高性能内核），再装 xlstm 包本身。用 pip 一行即可，不建议跳步，因为 xLSTM 包在 import 时会去检查内核是否可用。也可以从源码安装，便于阅读与魔改：

git clone https://github.com/NX-AI/xlstm.git
cd xlstm
pip install -e .

源码安装后，仓库里的 experiments 目录、notebooks 目录都能直接跑，适合跟着论文复现实验。对想深入读源码的同学，建议先从 xlstm_large/model.py 这个单文件实现看起，它比 NeurIPS 版更紧凑，集中体现了 7B 模型的工程取舍。

2. 运行 xLSTM Large 7B 快速开始

xLSTM Large 的独立实现放在 xlstm/xlstm_large/model.py，只依赖 mlstm_kernels，与 NeurIPS 版架构代码相互独立，方便单独研究大模型形态。官方提供 demo.ipynb 笔记本，演示如何导入配置与模型类、随机初始化、跑一次前向。关键配置项包括：嵌入维度 embedding_dim、注意力头数 num_heads、块数 num_blocks、词表大小 vocab_size、运行模式 mode，以及三类内核选择——分块内核 chunkwise_kernel、序列内核 sequence_kernel、单步内核 step_kernel；想跑 TFLA Triton 加速就分别选 chunkwise--triton_xl_chunk、native_sequence__triton、triton。mode="inference" 用于部署，训练时换对应模式即可。权重在 HuggingFace 开放下载，按常规 from_pretrained 方式加载即可。建议第一次跑通时先用官方 demo 里的随机小模型确认环境没问题，再去拉 7B 权重，避免几百 GB 下载完才发现内核没装上。推理时若显存吃紧，可以减小 num_blocks 或换用更小的配置先验证流程，再逐步放大。

3. 其他硬件适配

模型主要在 NVIDIA GPU 上测试，Triton 内核理论上也能跑 AMD GPU，因为 Triton 本身是跨厂商的中间表示层。Apple Metal 等平台暂时建议改用原生 PyTorch 实现：把 chunkwise_kernel 换成 chunkwise--native_autograd、sequence_kernel 换成 native_sequence__native、step_kernel 换成 native，代价是速度慢一些但能跑通，适合先验证算法正确性。如果你在 Apple MLX 生态里工作，可以看社区维护的 xLSTM-metal 移植版，它提供 MLX 原生、面向 Apple Silicon 的实现，统一内存架构下推理效率会明显更好。

4. 使用 NeurIPS 版架构

非语言类应用或想把它当骨干网络嵌入其他项目，用 xLSTMBlockStack；语言建模或其他 token 级任务用 xLSTMLMModel（在 block stack 外面包了 token embedding 与 lm head）。block stack 通过配置 dataclass 分别描述 mLSTM 块（卷积核大小、QKV 投影块大小、头数）与 sLSTM 块（后端、头数、偏置初始化、前馈投影因子与激活函数），再指定上下文长度、总块数、嵌入维度和 sLSTM 在哪些位置出现。输入形状是 (batch, seq_len, embedding_dim)，输出形状与输入一致，和 Transformer block 栈完全对称。如果习惯 YAML 配置，可以用 dacite 把 yaml 字典直接转成配置 dataclass，方便团队把超参数沉淀进配置文件，也便于做超参扫描。整段代码读完可以发现：xLSTM 的接口设计刻意模仿 Transformer，老用户几乎零学习成本就能替换骨干网络。这种设计选择对工业落地非常友好——现有的训练框架、日志、监控、推理部署脚本大多是按 Transformer 形状写的，xLSTM 保持了对称的输入输出形状，迁移时主要改模型定义，业务层基本不动。

5. CUDA 编译排错

编译 sLSTM CUDA 内核报错时，先试 export TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0" 显式指定目标算力；torch 与 CUDA 版本必须严格匹配，否则编译出的内核无法加载；需要自定义头文件路径时用环境变量 XLSTM_EXTRA_INCLUDE_PATHS='/usr/local/include/cuda/:/usr/include/cuda/'（shell 或 Python 里设置均可）。这套排错思路来自社区贡献者在 issue 区的总结，覆盖了大部分"编译能过但运行报找不到符号"的情况。想要更极致的单步速度，可以试官方的 FlashRNN 独立内核库，它把 sLSTM 的循环步进一步做了融合优化，适合对延迟敏感的在线服务。对个人研究者来说，遇到编译问题先看 issue 搜索，大部分坑都有人踩过并写在讨论里。如果你的需求只是理解算法而不是跑 7B，完全可以先用 native 内核在 CPU 上跑小模型，把流程走通。

五、核心特性与使用方法

两种变体互补。 sLSTM（标量 LSTM）通过记忆混合提供状态追踪能力，mLSTM（矩阵 LSTM）通过矩阵记忆与状态扩展提供记忆能力。打个比方：sLSTM 像一个会做"当前状态判断"的小控制器，擅长追踪奇偶、计数这类需要持续更新的规则；mLSTM 像一块大容量的联想记忆板，擅长记住早期出现过的 key 并在后期被 query 时召回。合成实验里，Parity 任务只有 sLSTM 能解，Multi-Query Associative Recall 任务则靠 mLSTM 的矩阵记忆占优，两者按层组合后两类任务都表现最好。复现实验：进入仓库后跑 PYTHONPATH=. python experiments/main.py --config experiments/parity_xlstm01.yaml（纯 sLSTM）、parity_xlstm10.yaml（纯 mLSTM）、parity_xlstm11.yaml（两者混合），注意训练循环不含早停与测试评估，需要自己外接。

为什么关注 7B 循环模型。 xLSTM 7B 是少有的把循环架构推到百亿参数级并跑通完整预训练的工作，权重已在 HuggingFace 开放；循环结构在推理时的优势是序列越长、相对 Transformer 的显存与算力优势越明显，因为它不像注意力那样随长度二次增长。这让它特别适合长文档理解、代码生成、长对话记忆、端侧部署等场景——在这些场景里 Transformer 的 KV-cache 会迅速吃掉显存，而 xLSTM 只需维护固定大小的循环状态。作者团队还专门针对 7B 做了训练吞吐与稳定性优化，代码放在 xlstm/xlstm_large 目录，与 NeurIPS 教学版并列，读者可以对比学习"研究版"和"工程版"的差异。对小团队来说，这类架构的现实意义是：在有限显存的消费级显卡上跑长上下文对话，循环模型往往比 Transformer 更划算。你不需要为一个 32k 上下文准备几十 GB 的 KV-cache，只需要一份固定大小的循环状态随对话滚动更新。

六、常见问题 FAQ

1. xLSTM 和原来的 LSTM 是什么关系？ 它是经典 LSTM 的"扩展版"：保留了循环门控思想，但用指数门控、归一化稳定化和矩阵记忆解决了原版难以扩展到大规模模型的问题，由 LSTM 发明人团队亲自操刀，相当于给这个老架构装上了现代训练稳定性。从这个意义上说，它是一次"老树新枝"式的升级，而不是另起炉灶。

2. 7B 模型权重许可怎么看？ 代码是 Apache-2.0，7B 权重采用 NXAI 社区许可（见仓库 LICENSE），商用前请到 HuggingFace 模型页确认条款，避免踩许可红线。个人研究与学习使用则完全免费，商用前再仔细确认许可条款细节内容之后即可使用。

3. 只有 AMD 卡或 Apple Silicon 能跑吗？ AMD 可以直接用 Triton 内核；Apple 暂时把三个 kernel 都换成 native 实现，或用社区的 xLSTM-metal（MLX 原生）。AMD 性能虽未经官方充分验证，但 Triton 后端理论上兼容，遇到问题建议先在 native 路径跑通再切加速。

4. sLSTM 编译报错怎么办？ 先设 TORCH_CUDA_ARCH_LIST="8.0;8.6;9.0"，确认 Compute Capability ≥ 8.0，再用 XLSTM_EXTRA_INCLUDE_PATHS 指到正确的 CUDA 头文件目录；还不行就去仓库 issue 区搜关键词。这一步是新手最容易卡住的地方，因为 sLSTM 的 CUDA 内核需要在本机编译，而不是 pip 装好就能用。

5. 想把它嵌进自己的项目当骨干网络？ 用 xLSTMBlockStack，输入输出形状与 Transformer block 栈一致，直接替换即可；想做语言建模用 xLSTMLMModel，词表大小在配置里指定。如果已有 Transformer 训练 pipeline，迁移成本主要在内核环境上，模型本身几乎不用改业务代码。

6. 引用格式？ 学术使用请引用两篇论文：Beck et al. 2024（arXiv:2405.04517）与 Beck et al. 2025（arXiv:2503.13427），BibTeX 见仓库 Citation 章节，分别对应架构提出与 7B 模型发布两个节点。

7. 它能直接当聊天模型用吗？ 仓库主要提供架构代码与 7B 权重，开箱聊天体验不如成熟对话框架；建议先读 demo 与论文，把它当作"研究/二次开发基座"，再按需接入自己的推理与对话层。如果你只是想要一个能直接用的本地聊天助手，现在的成熟选择更多；如果你在评估下一代架构，xLSTM 是绕不开的参考。它代表了"循环结构能否在大模型时代复兴"这一问题的最有力回答，也是理解后 Transformer 路线必读的开源实现。

8. 学习路线建议？ 先读 NeurIPS 论文理解 sLSTM/mLSTM 的数学动机，再跑 experiments 里的 parity 与 associative recall 合成任务，最后上 7B 权重做前向推理；不要一上来就微调，否则很容易在工程问题里卡住。仓库自带的 notebook 是最快的上手路径，建议跟着敲一遍代码再回头看论文，理解会深刻得多。

七、跳转链接

• GitHub 仓库：NX-AI/xlstm
• 7B 模型权重：HuggingFace: xLSTM-7b
• 快速内核：mlstm_kernels
• 极致加速：FlashRNN
• Apple MLX 移植：xLSTM-metal
• 架构论文：arXiv:2405.04517
• 7B 论文：arXiv:2503.13427