MAX v24.3 发布说明深度解读：Mojo 自定义算子、命名动态维度与引擎 API 重构-迪斯科星球

MAX v24.3 发布说明深度解读：Mojo 自定义算子、命名动态维度与引擎 API 重构

【免费下载链接】mojoThe Modular Platform (includes MAX & Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo

本篇文章基于 Modular 平台 MAX v24.3（2024-05-02 发布）的官方发布说明，系统梳理该版本的核心变更：包括用 Mojo 为模型编写自定义算子这一里程碑能力、MAX Engine 支持"命名动态维度"的优化手段，以及 MAX Graph / MAX Engine 在 Python、Mojo、C 三层 API 上的重构细节与 TensorFlow 支持的移除。阅读完本文，你将能够：使用ops.custom与extensibility在 MAX Graph 中接入 Mojo 内核、通过命名动态维度（TorchInputSpec/ShapeElement/dimNames）让编译器识别需要等长的动态维度，并理解新旧 API 的对应关系以便迁移代码。

版本概览：MAX v24.3 的核心主题

MAX v24.3 是一次以"扩展性"与"API 一致性"为主线的版本更新，主要变更集中在四个方向：

自定义算子（Custom Ops）：开发者首次可以直接用 Mojo 为模型编写自定义算子，这是该版本最突出的新能力；
命名动态维度（Named Dynamic Dimensions）：允许为动态维度命名，使多个动态维度在运行时强制等长，帮助 MAX Engine 编译器做额外优化；
API 简化与重构：统一模型输入规格（Input Spec）的加载接口，Python、Mojo、C 三层 API 均围绕InferenceSession.load()与自定义算子加载路径做了重新设计；
移除 TensorFlow 支持：MAX SDK 不再支持加载 TensorFlow SavedModel 进行推理。

这些变更的落地代码与示例均可在当前仓库中找到对应实现，下文逐一展开。

🔥 里程碑特性：用 Mojo 编写自定义算子

MAX v24.3 起，你可以直接用 Mojo 为模型编写自定义算子（Custom Ops），将高性能内核以原生 Mojo 的形式接入 MAX Graph，并在 GPU（若可用）或 CPU 上编译、调度与执行。

示例总览与运行方式

仓库中 max/examples/custom_ops/ 目录提供了从入门到复杂的全套示例，据其 README.md 说明，各示例主题如下：

示例	内容
`addition`	为输入张量的每个元素加 1（`add_one`内核）
`mandelbrot`	计算 Mandelbrot 集合
`vector_addition`	使用手写 GPU 函数进行向量加法
`eager_vector_addition`	与`vector_addition`相同的算子，但使用 Eager Tensor API
`top_k`	Top-K token 采样器，展示大语言模型处理流水线中的真实算子场景
`matrix_multiplication`	多种矩阵乘法算法，使用内存布局抽象
`fused_attention`	融合注意力算子，综合运用 MAX GPU 编程特性
`image_pipeline`	串联灰度、提亮、模糊多个自定义算子的图像流水线，数据全程留在 GPU

每个示例由kernels/目录下的 Mojo 内核实现与基目录下的 Python 图构建代码组成，使用单一 Pixi 命令即可运行，例如：

pixi run addition pixi run mandelbrot pixi run vector_addition pixi run top_k pixi run matrix_multiplication pixi run fused_attention pixi run image_pipeline

pixi run <example>负责保证max包依赖可见，Python 端构建图与推理会话状态，Mojo 内核代码在需要时即时（重新）编译，确保执行的图始终使用最新的 Mojo 代码。仓库还提供了基准测试入口：

pixi run benchmark

在 Python 图中接入 Mojo 自定义算子

以 max/examples/custom_ops/addition.py 为例，图中通过ops.custom按字符串名字引用 Mojo 算子，并声明输入张量列表与预期输出类型：

from max.driver import CPU, Accelerator, Buffer, accelerator_count from max.dtype import DType from max.engine import InferenceSession from max.graph import DeviceRef, Graph, TensorType, ops device = CPU() if accelerator_count() == 0 else Accelerator() graph = Graph( "addition", forward=lambda x: ( ops.custom( name="add_one", # Mojo 内核注册名 device=DeviceRef.from_device(device), values=[x], out_types=[ TensorType( dtype=x.dtype, shape=x.tensor.shape, device=DeviceRef.from_device(device), ) ], )[0].tensor ), input_types=[ TensorType(dtype, shape=[rows, columns], device=DeviceRef.from_device(device)), ], custom_extensions=[mojo_kernels], # Mojo 内核源码目录 ) session = InferenceSession(devices=[device]) compiled = session.compile(graph) model = session.init(compiled)

其中custom_extensions=[mojo_kernels]指定了 Mojo 内核所在的源码目录（Path(__file__).parent / "kernels"）。一个值得注意的细节是：同一份 Mojo 内核代码无需改动即可同时运行在 CPU 与 GPU 上——若系统存在受支持的加速器则优先调度到加速器，否则回退到 CPU；vector_addition示例展示了这一机制：通过编译期特化（compile-time specialization），MAX 会为给定硬件架构选择最优代码路径。

Mojo 侧的内核实现与注册

自定义算子在 Mojo 端通过extensibility模块注册。以 max/examples/custom_ops/kernels/add_one.mojo 为例：

import extensibility from max.gpu.host import DeviceContext from extensibility import InputTensor, OutputTensor, foreach from std.utils.coord import Coord from std.utils.index import IndexList @extensibility.register("add_one") struct AddOne: @staticmethod def execute # The kind of device this will be run on: "cpu" or "gpu" target: StaticString, raises: @__parameter @always_inline def elementwise_add_onewidth: Int -> SIMD[x.dtype, width]: return x.loadwidth + 1 foreachelementwise_add_one, target=target

关键点拆解：

@extensibility.register("add_one")：将结构体注册为名为add_one的自定义算子，Python 图通过ops.custom(name="add_one", ...)与之对应；
execute的参数化参数target: StaticString：值为"cpu"或"gpu"，决定内核调度到哪种设备；ctx: DeviceContext在执行某些 GPU 调用时需要；
foreachelementwise_add_one, target=target：把逐元素内核分发到输出张量的每个坐标；
文件底部还演示了@extensibility.register_shape_function("add_one")用于注册形状推导函数——只有当你在图中未手动标注输出形状时才需要实现它。

同目录下的 kernels/add_constant.mojo 展示了带编译期参数[value: Int]的变体（AddConstant），对应parametric_addition.py中的参数化使用方式，说明自定义算子可以像 Mojo 泛型一样携带静态参数。

命名动态维度：让编译器理解"必须等长"的动态尺寸

MAX v24.3 新增对命名动态维度的支持：此前模型中动态维度的尺寸只能用None表示，无法表达"两个或多个动态维度在运行时尺寸必须相等"这一约束；现在可以为这些维度赋予名字，MAX Engine 编译器即可据此执行额外优化。

典型应用场景是带有 batch 维与序列维的模型——例如若干输入张量共享同一个动态 batch 大小。通过为这些维度指定相同名称（而非各自None），开发者向编译器明确"这些尺寸在运行时保持一致"，从而获得更优的代码生成。

三种语言的 API 落地

命名动态维度贯穿 MAX Engine 的三种编程接口：

Python API：TorchInputSpec支持命名动态维度，可告知 MAX 哪些动态维度要求相同尺寸；
Mojo API：新增ShapeElement类型，用于在InputSpec中表达命名动态维度；
C API：M_newTorchInputSpec()新增dimNames参数，以支持命名动态维度。

MAX Graph API 变更：类型系统收敛与直接实例化 Graph

发布说明指出max.graphAPI 仍处于快速演进但开始趋于稳定，本版本的改动集中在类型重命名、冗余类型移除与建图入口简化三方面：

旧名称	新名称 / 替代方案	说明
`AnyMoType`	`Type`	类型统一重命名
`MOTensor`	`TensorType`	张量类型统一重命名
`MOList`	`ListType`	列表类型统一重命名
`ElementType`	`DType`	移除`ElementType`，改用`DType`
`TypeTuple`	`List[Type]`	移除`TypeTuple`，改用`List[Type]`
`Module`	直接实例化`Graph`	移除`Module`类型，建图时直接创建`Graph`

同时max.ops中新增了一批算子，包括对自定义算子的支持——即上文ops.custom的用法，详见 max/examples/custom_ops/addition.py 中的实际调用。

MAX Engine Python API 变更：`load()`重新设计

用`custom_ops_path`取代混乱的`options`参数

MAX v24.3 重新设计了InferenceSession.load()：原本令人困惑的options参数被替换为custom_ops_path（用于加载自定义算子）。作为结果，CommonLoadOptions、TorchLoadOptions、TensorFlowLoadOptions三个类被全部移除。

从当前仓库源码看，这一思路在后续演进中被进一步延伸：如今的 api.py 中load()接收的是custom_extensions（支持指向编译产物.mojoc或.mojo源文件的路径，.mojo源文件会被自动编译为包后再加载），并且官方推荐将load()拆分为compile()（api.py 第 923 行起，只编译不绑定权重，支持跨编译场景）与init()（api.py 第 985 行起，绑定权重生成可执行Model）两步，以复用编译产物。多模型产物（如视觉编码器与语言模型一起编译）则使用load_all()/init_all()按sym_name分别获取。

`TorchInputSpec`支持命名动态维度

TorchInputSpec现在支持命名动态维度：此前动态维度尺寸只能写None，现在可以为需要等长的动态维度赋予名字，帮助 MAX 更好地优化模型（详见上文"命名动态维度"一节）。

MAX Engine Mojo API 变更：`load()`与新的`InputSpec`/`ShapeElement`

Mojo 侧的 API 变更与 Python 侧对齐：

InferenceSession.load_model()重命名为load()；
InferenceSession.load()重新设计：原先的config参数被替换为两个职责明确的参数——custom_ops_path（加载自定义算子时使用）与input_specs（加载 TorchScript 模型时使用）；
随之移除了LoadOptions类型，并引入新的InputSpec类型来定义模型的输入形状/类型；
新增ShapeElement类型，用于在InputSpec中表达命名动态维度；
max.engine.engine模块重命名为max.engine.info。

迁移要点：如果你此前用LoadOptions声明输入规格，v24.3 起需要改用InputSpec（配合ShapeElement表达动态/命名维度）；调用入口从load_model()改为load()。

MAX Engine C API 变更

C 层 API 的变更聚焦于命名动态维度的打通：M_newTorchInputSpec()新增dimNames参数，使 C 语言接口同样能够指定动态维度名称，与 Python 的TorchInputSpec、Mojo 的ShapeElement保持能力一致。

❌ 移除项：TensorFlow 支持下线

MAX v24.3 移除了 MAX SDK 中的 TensorFlow 支持，具体影响与原因如下：

影响：无法再加载 TensorFlow SavedModel 进行推理；TensorFlow 仍面向企业客户提供；
原因（据发布说明）：行业范围内 TensorFlow 使用量显著下滑，尤其在最新 AI 创新领域；移除 TensorFlow 还可将包体积削减超过 50%，并加速其他客户需求特性的开发；
连带移除：Python 的CommonLoadOptions、TorchLoadOptions、TensorFlowLoadOptions三个类，以及 Mojo 的LoadOptions类型——它们均随InferenceSession.load()的参数重设计一并删除。

如果你有 TensorFlow 模型的生产级需求，发布说明建议联系 Modular 官方洽谈企业方案；在 MAX SDK 路线中，推理入口统一收敛为 Mojo/Python/C 三层中与自定义算子、命名动态维度能力对齐的新接口。

版本性能表现（发布说明口径）

发布说明给出了与上一版本 v24.2 的对比数据：MAX Engine v24.3 在 PyTorch 模型上平均加速 10%，在动态量化（dynamically quantized）的 ONNX Transformer 模型上平均加速 20%。该数据为官方发布说明口径，具体收益会因模型结构、量化方式与硬件环境而异，可作为性能回归测试的参考基线。

总结与迁移建议

MAX v24.3 的核心脉络可以概括为"扩展 + 收敛"：

新能力：Mojo 自定义算子打通了从内核编写（extensibility.register）到图接入（ops.custom）再到推理会话（custom_ops_path/custom_extensions）的完整链路，仓库 max/examples/custom_ops/ 是现成的上手模板；
优化手段：命名动态维度通过TorchInputSpec（Python）、ShapeElement/InputSpec（Mojo）、dimNames（C API）表达，让编译器获得"等长约束"这一额外信息；
API 迁移清单：Python/Mojo 的InferenceSession.load()参数收敛为custom_ops_path/input_specs；load_model()改名load()；LoadOptions系列全部移除；max.engine.engine改名max.engine.info；max.graph中AnyMoType→Type、MOTensor→TensorType、MOList→ListType、ElementType→DType、TypeTuple→List[Type]，建图直接实例化Graph；
注意移除项：TensorFlow SavedModel 推理在 MAX SDK 中不再可用。

对照上述清单即可完成从 v24.2 及更早版本到 v24.3 的平滑迁移，并第一时间用上 Mojo 自定义算子与命名动态维度带来的编译优化能力。

【免费下载链接】mojoThe Modular Platform (includes MAX & Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考

企业官网建设流程全解析