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.32024-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为输入张量的每个元素加 1add_one内核mandelbrot计算 Mandelbrot 集合vector_addition使用手写 GPU 函数进行向量加法eager_vector_addition与vector_addition相同的算子但使用 Eager Tensor APItop_kTop-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_pipelinepixi 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, forwardlambda x: ( ops.custom( nameadd_one, # Mojo 内核注册名 deviceDeviceRef.from_device(device), values[x], out_types[ TensorType( dtypex.dtype, shapex.tensor.shape, deviceDeviceRef.from_device(device), ) ], )[0].tensor ), input_types[ TensorType(dtype, shape[rows, columns], deviceDeviceRef.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 上——若系统存在受支持的加速器则优先调度到加速器否则回退到 CPUvector_addition示例展示了这一机制通过编译期特化compile-time specializationMAX 会为给定硬件架构选择最优代码路径。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, targettarget关键点拆解extensibility.register(add_one)将结构体注册为名为add_one的自定义算子Python 图通过ops.custom(nameadd_one, ...)与之对应execute的参数化参数target: StaticString值为cpu或gpu决定内核调度到哪种设备ctx: DeviceContext在执行某些 GPU 调用时需要foreachelementwise_add_one, targettarget把逐元素内核分发到输出张量的每个坐标文件底部还演示了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 APITorchInputSpec支持命名动态维度可告知 MAX 哪些动态维度要求相同尺寸Mojo API新增ShapeElement类型用于在InputSpec中表达命名动态维度C APIM_newTorchInputSpec()新增dimNames参数以支持命名动态维度。MAX Graph API 变更类型系统收敛与直接实例化 Graph发布说明指出max.graphAPI 仍处于快速演进但开始趋于稳定本版本的改动集中在类型重命名、冗余类型移除与建图入口简化三方面旧名称新名称 / 替代方案说明AnyMoTypeType类型统一重命名MOTensorTensorType张量类型统一重命名MOListListType列表类型统一重命名ElementTypeDType移除ElementType改用DTypeTypeTupleList[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/ShapeElementMojo 侧的 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/ 是现成的上手模板优化手段命名动态维度通过TorchInputSpecPython、ShapeElement/InputSpecMojo、dimNamesC API表达让编译器获得等长约束这一额外信息API 迁移清单Python/Mojo 的InferenceSession.load()参数收敛为custom_ops_path/input_specsload_model()改名load()LoadOptions系列全部移除max.engine.engine改名max.engine.infomax.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），仅供参考