Together AI 解析如何高效服务 MiniMax M3:1M token 上下文与多模态推理优化
Serving MiniMax-M3 for efficient inference: Unlocking 1M-Token Context and Multimodality Without Regrets
Together AI 作为 MiniMax M3 的首选云合作伙伴,介绍了为该模型提供生产级推理服务的工程方案。M3 采用 MiniMax Sparse Attention 块稀疏注意力,使 prefilling 提速超 9x、decoding 提速超 15x,支持 1M token 上下文与原生多模态。
原文由服务方工程团队拆解 MSA 的内核优化与多模态网关预处理方案,给出可参考的长上下文推理工程路径和实测吞吐提升数字。
MiniMax 发布了其最新一代顶尖模型 M3,Together AI 很荣幸成为其首选云合作伙伴,助力 MiniMax 高效地大规模生产化部署 M3。一旦 MiniMax M3 在未来几天内作为开放权重模型发布,Together AI 也将直接为开发者提供该模型的端点托管服务。在这一规模背后,是我们推理团队和内核团队的卓越工作,他们推动了深度的性能优化,并确保了这一前沿模型的生产级可靠性:1M token 上下文窗口、原生多模态,以及一个需要严肃工程投入才能高效服务的架构。在这篇文章中,我们将介绍我们是如何做到的。祝贺 MiniMax 团队完成这一里程碑式的模型发布以及持续不断的创新。
MiniMax M3 是一款一体化模型,集最先进的编码性能、智能体工作流支持和原生多模态推理于一身。在这些能力之上,它还设计为支持 1M 上下文,同时在服务成本上极具经济性。这使其非常适合真实世界任务,在这些任务中,长文档、代码库、工具使用、图像和迭代推理往往同时出现,且上下文负载很重。与前一代相比,服务 M3 带来了更多挑战,因为新能力需要在更多维度上进行优化,包括稀疏注意力计算、更大的 KV 缓存管理、多模态处理等。
架构 / 特性
M3 中最具创新性的架构变更是 MiniMax 稀疏注意力(MSA),它旨在解决 MiniMax M2.7 中出现的注意力计算瓶颈。其块稀疏注意力机制限制了每个 query 可以关注的最大 token 数量,降低了长上下文处理的成本,使更长的上下文窗口变得切实可行。这在预填充阶段带来了超过 9 倍的加速,在解码阶段带来了超过 15 倍的加速。

本质上,MSA 的计算由两部分组成:一部分是分数计算,用于确定每个 KV 组最相关的 K 个块进行关注;另一部分是 query token 与这些块之间的稠密注意力。这种设计在 KV 组维度上保留了表达能力,同时对 query token 关注的 KV token 最大数量设定了上限。注意力计算本身不再随上下文长度以 N^2 扩展,因此非常适合长上下文工作负载。
我们在 B200 上、并发度为 8 的智能体式流量形态(60k 前缀缓存)下测量了内核执行时间分解。MSA 显著降低了每次迭代中实际注意力计算所占的墙钟时间百分比。

除了注意力架构的变更外,M3 还配备了多模态支持,包含视觉组件以及新的图像和视频预处理功能。
鉴于这些根本性变化,Together AI 与 MiniMax 的工程团队紧密合作,应对新出现的挑战。一些主要挑战包括:
- 尽管 MiniMax 稀疏注意力计算本身非常高效,但从工程角度来看,支持 1M 上下文长度仍然具有挑战性。
- 视频和图像处理天生就比文本分词更复杂。
优化
KV 块主序稀疏注意力
在预填充阶段,对于长上下文输入,注意力计算仍然可能是一个重要因素,因为对于每个 token,我们需要计算 Selected Block * KV Head Group * Tokens。块稀疏注意力的特性允许多个查询关注相同的键值块。因此,如果我们逐个查询计算与键值块的注意力,就会在 GPU 上重复从 HBM 到 SRAM 的 KV 移动。将键值组放在外层循环中迭代,并在内层循环中计算查询 token 之间的注意力,可以实现更好的算术强度,因为 KV 缓存只移动一次。
为了实现这一点,我们需要将 {q, kv block} 的映射重新组织为 {kv block, q},并重新实现注意力内核。由于我们只计算 kv 块的部分 O 输出,我们需要基于 Log-Sum-Exp 进行最终的“归约”,以重新缩放输出 O 并求和。过程如下:

将 MSA 与分页注意力集成
在现代推理引擎中,分页注意力通常用于管理请求的 KV 缓存上下文。大多数高度优化的注意力内核都是针对固定页面大小支持编写的。阻止我们使用这些内核的障碍在于,所选块在不同 KV 组之间是不同的。
在 Together AI,我们提出了一种将 MiniMax 稀疏注意力集成到引擎中的新方法。在解码期间,我们首先基于所选块构建页表,将 KV 组维度展平到批次维度,并利用 KV 缓存张量的跨步视图,为注意力内核提供检索 KV 页所需的指针。技巧在于步长:页地址按 D 前进以选择虚拟页起始,而 token 按 Hkv * D 前进。这将一个物理张量解交织为每个头的页,因此每个展平的行现在可以使用不同的页表。

这一设计使我们能够使用支持 GQA 的现有注意力内核,而无需从头重写一个支持稀疏注意力的新内核。由于每个查询的所选块是有限的,用于查找块到页映射的内核开销非常低。这一设计使我们的解码吞吐量提升了 5%。
解码索引评分内核优化
对于解码,MSA 将大部分开销从稠密注意力中移出,转移到 score/top-k 索引器中。对于每个解码查询,引擎将查询侧索引向量与候选键侧索引向量进行比较,将每个 128 token 的 KV 块归约为单个分数,并仅保留 top 块用于真正的注意力内核。这个扫描位于每个生成 token 的关键路径上,而在长上下文长度下,候选块的数量会随上下文长度增长。解码评分具有小查询索引、长键索引的形态。人们很容易将一批解码查询视为一个更大的 GEMM,但 score/indexing 步骤并不仅仅是稠密矩阵乘法:每个请求和 K 组都有自己的候选块范围、掩码、逐块归约和 top-k 边界。将查询拼接在一起仍然会在 GEMM 周围留下参差不齐的 gather-and-reduce 问题,同时还会将填充和额外簿记强加到关键路径上。因此,我们优化的路径使用 AB 交换的 HMMA 布局:128 token 的键索引块成为 MMA 的 M 维度,而查询侧仅填充到较小的 N 维度。该内核通过异步拷贝暂存 128 token 的 K 索引,预取下一页,使用 HMMA 以 bfloat16 计算点积,并将每页归约为一个块分数。

网关处的多模态预处理
SMG(Serving Model Gateway)是一个基于 Rust 的模型网关,位于 OpenAI 兼容 API 和推理引擎之间。除了路由和分词之外,SMG 还承担了对多模态模型尤为重要的一项职责:它在请求到达 GPU worker 之前,在 CPU 上执行所有视觉预处理。
图像和视频输入在能被视觉编码器使用之前,需要相当多的 CPU 工作:下载、解码、帧采样、缩放以及转换为 patch 张量。在推理引擎内部执行这些操作会占用本应用于生成的资源。SMG 改为在网关中处理所有这些工作,因此当请求到达 GPU 时,张量已经准备就绪。

对于 M3,这意味着:获取视频,用 FFmpeg 提取帧,根据 FPS(每秒帧数)选择子集,缩放并归一化,然后进行 patchify,并将时间维度融入其中。输出的是一个扁平的 patch 张量和一个小的网格元数据张量,打包进 gRPC 消息中。worker 直接运行视觉编码器——其侧无需预处理。
此外,SMG 的多模态流水线围绕 Rust trait 构建,将模型特定的预处理逻辑与流水线管道分离。添加 M3 多模态支持意味着使用 M3 特定的常量来实现这些 trait;流水线本身没有改变。相同的架构适用于大多数具有视觉能力的开源模型,并可跨推理引擎运行时通用化。
性能结果
自从收到 MiniMax M3 权重和模型架构以来,我们一直致力于提升推理性能。在常见的 agentic 形态流量上,我们在各种并发级别下实现了 81% - 125% 的提升。

未来工作
新架构带来了新的基础设施和工程挑战。在 Together AI,我们的目标是提供最佳的推理性能。我们正在积极围绕 M3 开展的一些工作包括:
1. 稀疏注意力架构引入了更多更小的内核,例如对 kv 块进行 topk、将 q-kv 映射重映射为 kv-q 等。这里有更多的内核融合机会。我们的 Kernel Agent Research 团队正在积极开发能够编写生产级内核的智能体。
2. k-index 和实际 kv 缓存的 CPU 缓存卸载现在可以解耦。我们正在研究加载完整的 k-index,并基于 topk 选择按需加载 kv 缓存。
下一步
加入等待列表,以便在 M3 开放权重在 Together AI 上可用时收到通知。
查看来自 MiniMax 模型设计者和 Together 推理团队的深度解析。
来源:Together AI 研究与产品博客 · together.ai