训练百万 Token 的模型,像让 AI 一次读完一座小型图书馆:如果每段文字都要和其他段落逐一比较,计算量和显存占用会迅速膨胀。Flash-MSA 的思路是做“稀疏注意力”——只挑出值得关注的部分关系来计算,并用 CuTeDSL(一种编写高性能 GPU 运算的专用语言)为 NVIDIA Hopper、Blackwell 架构定制训练内核。
最值得注意的是它对“挑选结果”的复用。MSA 每次选择的是 128 个 Token 组成的信息块,而非单个 Token;内核可以把这些块的编号一直保留到反向传播——也就是模型根据误差调整参数的阶段。作者称,因此整个训练步骤里,只有负责初次筛选的代理注意力前向计算仍会随上下文长度产生两两比较,其余环节都复用已缓存的稀疏块。
作者将其称为首个高性能开源 MSA 训练内核,但未提供完整实测数字;项目也不是 MiniMax 官方实现。现阶段更适合作为长上下文工程团队值得跟踪的底层方案。