Transformers Exporters 全解：用统一 API 将任意 PreTrainedModel 导出为 PyTorch / ONNX / ExecuTorch 可执行图

Transformers Exporters 全解：用统一 API 将任意 PreTrainedModel 导出为 PyTorch / ONNX / ExecuTorch 可执行图

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

Transformers 内置的exporters模块（src/transformers/exporters）提供了一套统一的导出接口，可以把任意PreTrainedModel导出为ExportedProgram（AOT 编译 / 任意 PyTorch 运行时）、ONNXProgram（ORT / TensorRT / OpenVINO 等 ONNX 运行时）或ExecutorchProgramManager（移动端与边缘设备）。本文覆盖官方文档的全部实战路径——按后端安装依赖、单模型导出、动态形状控制、生成式模型拆分（prefill / decode / 多模态组件）、多 token decode、静态 KV cache 与零拷贝缓存更新，并结合同仓库源码解释每条导出流水线背后的五阶段实现。

三大导出器与统一接口

导出器直接内置在 Transformers 包中，而不是放在下游工具库里，这样的设计使得架构变更、新的注意力模式和自定义 cache 类型只要在 modeling 代码中落地，就能在导出时被即时支持。三个导出器及其产物、目标运行时如下：

注意：导出器目前处于实验阶段。模块内的许多补丁（patch）是针对 Torch、ONNX Script、ONNX Runtime、ExecuTorch 上游特定 bug 的临时绕行方案，一旦上游修复就会被移除。在 API 稳定之前，应把补丁视为与测试套件锁定的版本绑定：生产工具链中请锁定测试套件所用的依赖版本，并预期随着上游变化，旧补丁会消失、新补丁会出现。

最简用法只需三步——创建导出器、构造配置、调用export：

exporter = DynamoExporter() # 或 OnnxExporter、ExecutorchExporter config = DynamoConfig(dynamic=True) exported = exporter.export(model, inputs, config=config)

源码视角：HfExporter基类与版本校验

所有导出器继承自抽象基类HfExporter，它定义了export(model, sample_inputs, config)抽象方法，并内置了环境校验逻辑validate_environment：

• required_packages：必须安装的包，缺失时直接抛ImportError（例如OnnxExporter要求torch、onnx、onnxscript）；
• min_versions：硬性最低版本，低于即报错；
• tested_versions：测试套件验证过的版本。从源码结构看，校验时会剥离本地后缀（如+cu126、+cpu），版本漂移只会触发一条警告日志——这正是官方文档提示“检测到漂移时HfExporter会记录 warning”的实现来源。以 Dynamo 后端为例（exporter_dynamo.py）：

required_packages = ["torch"] min_versions = {"torch": "2.11.0"} tested_versions = {"torch": "2.12.0"}

Auto 类：在运行时选择后端

AutoHfExporter从配置中挑选导出器，AutoExportConfig从字典挑选配置类，两者遵循 Transformers 惯用的 Auto 类模式，适合后端在运行时才确定、而非在调用点硬编码的场景：

from transformers.exporters import AutoExportConfig, AutoHfExporter export_config_dict = {"export_format": "onnx", "dynamic": True} config = AutoExportConfig.from_dict(export_config_dict) exporter = AutoHfExporter.from_config(config) onnx_program = exporter.export(model, inputs, config=config)

从 auto.py 可以看到分发依据的export_format字段（对应ExportFormat枚举，取值为executorch/dynamo/onnx），以及两个注册表：

AUTO_EXPORTER_MAPPING = { "executorch": ExecutorchExporter, "dynamo": DynamoExporter, "onnx": OnnxExporter, }

模块还提供register_exporter(name)与register_export_config(name)装饰器，允许第三方注册新的后端（要求分别继承HfExporter/ExportConfigMixin）。配置类本身基于 dataclass，ExportConfigMixin提供to_dict/from_dict序列化能力，使配置可以保存后按export_format字段无损往返。

安装：按目标后端锁定依赖

只安装你计划导出的目标后端对应的依赖。以下版本是导出器测试套件锁定的版本；更新或更旧的版本通常也能工作，但导出器补丁针对特定 API 面编写，生产工具链建议照此锁定，并留意漂移警告。

• Dynamo：

pip install transformers "torch==2.12.0"

• ONNX：

pip install transformers "torch==2.12.0" "onnx==1.21.0" "onnxscript==0.7.0" onnxruntime

• ExecuTorch：

pip install transformers "torch==2.12.0" "executorch==1.3.1"

导出一个模型

所有导出器共享同一套接口：用配置创建导出器，调用HfExporter.export。切换运行时只需替换导出器类。

Dynamo 导出

from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.exporters import DynamoExporter, DynamoConfig model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B") tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") inputs = tokenizer("Hello, world!", return_tensors="pt") exporter = DynamoExporter() config = DynamoConfig(dynamic=True) exported = exporter.export(model, inputs, config=config) # 直接运行导出后的图 outputs = exported.module()(**inputs)

ONNX 导出

from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.exporters import OnnxExporter, OnnxConfig model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B") tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") inputs = tokenizer("Hello, world!", return_tensors="pt") exporter = OnnxExporter() config = OnnxConfig(dynamic=True) onnx_program = exporter.export(model, inputs, config=config) # 保存并用 ONNX Runtime 加载 onnx_program.save("model.onnx") import onnxruntime as ort session = ort.InferenceSession("model.onnx") ort_inputs = {k: v.numpy() for k, v in inputs.items()} outputs = session.run(None, ort_inputs)

OnnxConfig在DynamoConfig字段之外还提供 ONNX 专属开关（见 configs.py）：output_path（为None时模型仅驻留在内存中的ONNXProgram，不落盘）、opset_version（默认取已安装 onnxscript 支持的最新 opset）、external_data（默认True，大权重写入独立的.onnx_data旁车文件，规避 protobuf 2 GB 上限）、optimize（默认True，运行 onnxscript 优化通道；遇到上游优化器 bug 时可关闭）、export_params（默认True，嵌入权重；置False则导出无权重图）、keep_initializers_as_inputs（旧运行时 opset < 9 需要）。

ExecuTorch 导出

ExecutorchConfig的backend参数默认xnnpack，面向 CPU，CPU-only 环境即可运行；cuda面向 GPU，需要 CUDA 环境，在非 CUDA 环境请求它会抛出RuntimeError。

from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.exporters import ExecutorchExporter, ExecutorchConfig model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B") tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") inputs = tokenizer("Hello, world!", return_tensors="pt") exporter = ExecutorchExporter() config = ExecutorchConfig(backend="xnnpack", dynamic=True) et_program = exporter.export(model, inputs, config=config) # 保存用于端侧部署 et_program.save("model.pte") # 通过 ExecuTorch Python 运行时加载并运行 from executorch.runtime import Runtime program = Runtime.get().load_program("model.pte") method = program.load_method("forward") outputs = method.execute(list(inputs.values()))

ExecutorchConfig还有三个内存规划开关：alloc_graph_input/alloc_graph_output/alloc_mutable_buffers（均默认True）。它们的含义在“零拷贝原地更新”一节会再次出现——静态 cache 的 in-place 写入正要求关闭它们。

动态形状：dynamic=True与显式dynamic_shapes

传dynamic=True会把每个张量的每个维度都标记为动态，导出的图在运行时即可接受任意大小的输入而无需重新追踪。若需精细控制哪些维度动态，直接传显式dynamic_shapes——它会原样转交给torch.export.export。

import torch from transformers import AutoModelForCausalLM, AutoTokenizer from transformers.exporters import DynamoExporter, DynamoConfig model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-0.6B") tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") inputs = tokenizer(["Hello, world!", "Hi"], padding=True, return_tensors="pt") batch = torch.export.Dim("batch", min=1, max=32) seq = torch.export.Dim("seq", min=1, max=2048) exporter = DynamoExporter() config = DynamoConfig( dynamic_shapes={"input_ids": {0: batch, 1: seq}, "attention_mask": {0: batch, 1: seq}}, # 当显式 Dim(min=, max=) 区间内某守卫无法在整个符号区间成立时， # 把数据依赖的形状守卫发射为运行时 assert，而不是让导出失败。 # 大多数 LLM 在细粒度 Dim(min=, max=) 边界下需要开启此项； # 使用 dynamic=True / Dim.AUTO 时则不需要（torch.export 会自行推断形状关系）。 prefer_deferred_runtime_asserts_over_guards=True, ) exported = exporter.export(model, inputs, config=config)

同样的dynamic_shapes写法对OnnxConfig和ExecutorchConfig（后者额外带backend="xnnpack"）完全一致。

从源码看，两个开关的行为边界在DynamoExporter.export中清晰可见：dynamic=True且未提供显式dynamic_shapes时，会调用get_auto_dynamic_shapes为所有张量/缓存输入的每个维度生成Dim.AUTO，并打印一次提示——只标记真正变化的维度可以绕开对固定轴（batch、长度为 1 的 decode 步、num_heads/head_dim）的符号形状解析，导出速度显著更快。DynamoConfig的全部字段与默认值（dynamic=False、strict=False、dynamic_shapes=None、prefer_deferred_runtime_asserts_over_guards=False）定义在 configs.py。

源码视角：Dynamo 导出的五阶段流水线

DynamoExporter并非简单包装torch.export.export，模块头注释（exporter_dynamo.py）列出了按执行顺序的五段流水线，理解它们是排查导出失败的关键：

1. 模型签名补丁（patch_forward_signature）：把model.forward替换为由sample_inputs推导出的扁平显式签名。Transformers 模型普遍使用**kwargs签名，若不处理，torch.export会把它展开成与dynamic_shapes字典不匹配的combined_args大包；
2. 模型补丁（_PATCHES["dynamo"]，经apply_patches("dynamo")安装）：可逆的类属性替换，把不可导出的模式（数据依赖循环、in-place 操作、掩码检查）换成可导出等价物。例如_patch_sdpa在 CPU 上把 SDPA 路由到 MATH 后端以绕开 batch=1 的调度守卫，_patch_chunked_vision_attention则把 VLM 视觉编码器中split → 逐段 SDPA → cat的 Python 循环改写成“reshape 成批 + 单次 SDPA”的可追踪形式（安装了约 19 个视觉/音频注意力类）；
3. Pytree 注册（register_cache_pytrees_for_model）：为所有Cache子类注册 flatten/unflatten 钩子，使torch.export能穿透 cache 对象做追踪，也让后续torch.export.load能把StaticCache这类输入反序列化；
4. 动态形状：即上文get_auto_dynamic_shapes的Dim.AUTO推断；
5. 模型状态清理（reset_model_state）：把非 Cache 的有状态属性（如cached_rotary_positional_embedding）在追踪期间置空、追踪后恢复，防止 eager 前向残留状态泄漏进符号追踪，或 FakeTensor 泄漏回 eager 执行。

OnnxExporter继承DynamoExporter，在其之上再加五段 ONNX 专属处理（见 exporter_onnx.py 模块注释）：Torch 补丁（追踪期 monkey-patch torch op，使发射出可降级到 ONNX 的模式）、torch.onnx内部钩子（让 FX 节点修复在run_decompositions之后再跑一遍）、FX 节点修复（就地改写 ONNX 无法降级的节点：别名、in-place、死比较、_assert_*）、onnxscript 自定义翻译表（覆盖 torchlib 默认降级有 bug 的 aten op）、以及针对 ORT 兼容性的 ONNX IR 后置修复。patch_model_outputs会把输出拍平成具名张量字典，并对同名的输入/输出加input./output.前缀去歧义——这解释了后文 ONNX Runtime 静态 cache 循环里input.<name>/output.<name>配对命名的由来。

生成式模型：export_for_generation

对自回归生成，模型的forward在 prefill 步（完整提示词、无 KV cache）与 decode 步（单 token、已填充的 KV cache）具有不同形状。导出器因此提供HfExporter.export_for_generation，把两个阶段拆开并分别导出。

对多模态生成模型，prefill 还会进一步拆成图像/音频编码器、语言模型与lm_head三个组件。编码器与语言模型的发现走PreTrainedModel.get_encoder（modality="image"或"audio"）与get_decoder()访问器，因此任何使用这些访问器的新架构都能开箱即用。投影层（projector）组件仅在模型以约定属性名暴露时才出现——multi_modal_projector、connector、embed_vision、embed_audio（该名单固化在 utils.py 的_MULTIMODAL_PROJECTOR_NAMES中）。例如 Qwen2-VL 把投影器折叠进了视觉塔，其组件字典里就没有单独的multi_modal_projector键；新架构应让自己的投影器属性对齐到上述名单之一，而不是让名单继续增长。

以多模态视觉语言模型为例（Dynamo / ONNX / ExecuTorch 三版代码，仅导出器与配置类不同）：

from transformers import AutoModelForImageTextToText, AutoProcessor from transformers.exporters import DynamoExporter, DynamoConfig model = AutoModelForImageTextToText.from_pretrained("Qwen/Qwen2-VL-2B-Instruct") processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct") messages = [{"role": "user", "content": [{"type": "image", "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/pipeline-cat-chonk.jpeg"}, {"type": "text", "text": "Describe this image."}]}] text = processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) inputs = processor(text=text, images=messages[0]["content"][0]["url"], return_tensors="pt").to(model.device) exporter = DynamoExporter() config = DynamoConfig(dynamic=True) components = exporter.export_for_generation(model, inputs, config=config) # components = {"image_encoder": ExportedProgram, "language_model": ExportedProgram, "lm_head": ExportedProgram, "decode": ExportedProgram}

ONNX 版把导出器换成OnnxExporter/OnnxConfig，返回ONNXProgram组件；ExecuTorch 版换成ExecutorchExporter/ExecutorchConfig(backend="xnnpack", dynamic=True)，返回ExecutorchProgramManager组件，结构完全一致。

注意：导出的组件是彼此独立的图，不是拿来即用的推理管线。按序运行各编码器、投影嵌入、编排生成循环都是调用方的责任。

export_for_generation的内部机制

底层函数decompose_for_generation只运行一次model.generate(**inputs, max_new_tokens=2)，并 hookmodel.forward捕获真实的 prefill 与 decode kwargs（多模态时还会给每个编码器/投影器/语言模型挂 hook 捕获各自的子模块 kwargs）。正因为复用完整的生成机制，它对 decoder-only、SSM、encoder-decoder 与多模态架构都无需逐模型胶水代码即可工作；export_for_generation只是在它之上的一行封装——拿到{name: (submodel, fwd_kwargs)}后对每个组件调用self.export（见 base.py 中export_for_generation的实现，其中config既可以是单一配置，也可以是按组件名键控的字典——后者要求覆盖全部组件名）。

由于捕获过程是 eager 运行inputs，应传入小而具代表性的值：短提示词、单张小图或几帧音频。导出的图并不绑定这些尺寸（动态形状照常生效），但更小的捕获输入让decompose_for_generation更廉价，也便于符号形状推断保持可控。若想在“拆分”和“导出”之间插入自定义动作——如先跑一次 eager 前向做校验、替换某子模块输入、跳过某个阶段——可以直接调用decompose_for_generation：

from transformers.exporters.utils import decompose_for_generation components = decompose_for_generation(model, inputs) # {"image_encoder": (submodel, fwd_kwargs), "language_model": (...), ..., "decode": (...)} exported = {} for name, (submodel, subinputs) in components.items(): eager_outputs = submodel(**subinputs) # 导出前先 sanity-check eager 前向 exported[name] = exporter.export(submodel, subinputs, config=config)

多 token decode

默认情况下decode组件是单 token步——针对 KV cache 的一个查询 token——torch.export会把查询序列轴特化为 1。传multi_token_decode=True则把decode捕获为多 tokendecode：decompose_for_generation用max_new_tokens=3捕获，再通过_merge_decode_calls把两个连续 decode 步合并成一次前向，使该轴保持符号化。这样一张图就能服务所有查询长度：单 token（普通解码）、多 token（如投机 token 的整块接受）、以及 cache 为空时的普通 prefill。

from transformers.exporters import DynamoExporter, DynamoConfig exporter = DynamoExporter() config = DynamoConfig(dynamic=True) components = exporter.export_for_generation(model, inputs, config=config, multi_token_decode=True) # components["decode"] 现在接受可变数量的查询 token

ONNX / ExecuTorch 版写法相同（换成各自导出器与配置类）。注意：查询轴只有在动态形状导出（dynamic=True）下才保持符号化；静态导出会把它冻结在捕获长度上，得到一张定长的多 token 图。该特性与下文的静态 KV cache 正交可组合——合并后的 decode 会把每步的 token 原地写入定长 cache，写入位置由 cache 内部跟踪。

静态 KV cache

generate()默认使用DynamicCache，随序列增长不断重新分配——对导出的图来说是移动目标。静态cache 则是分配一次的定长缓冲区，每步在当前位置原地写入。配合多 token decode，可以把整个生成过程收敛为单张导出图：decode图接收定长 cache 与可变数量的查询 token，从而同时承担提示词（空 cache → prefill）与每个生成 token（已填充 cache → decode）。导出方式是在multi_token_decode=True的基础上，传一个cache_implementation="static"（并指定max_cache_len）的GenerationConfig：

from transformers import GenerationConfig from transformers.exporters import DynamoExporter, DynamoConfig exporter = DynamoExporter() gen_config = GenerationConfig(cache_implementation="static", max_cache_len=2048) components = exporter.export_for_generation( model, inputs, config=DynamoConfig(dynamic=True), generation_config=gen_config, multi_token_decode=True )

ONNX / ExecuTorch 版仅需替换导出器与配置类。此时decode图带有两个符号轴：查询长度（你喂入多少 token）与 cache 长度（max_cache_len，加载时可调整）。dynamic=True会把它们（以及所有其他轴）标记为Dim.AUTO，导出的图在加载时接受任意提示词长度与 cache 尺寸。

零拷贝原地更新

静态 cache 作为输入传入并被原地修改，因此一块缓冲区就能携带跨越 decode 步的状态、无需主机拷贝——前提是运行时绑定调用方提供的缓冲区，而不是复制到自有 arena。这是各后端仅剩的差异化部分：

• Dynamo：导出后的程序把 cache 写入建模为USER_INPUT_MUTATION，直接调用components["decode"].module()(...)就会更新你传入的 cache 张量。每步复用同一批张量即可，无需任何配置。

• ONNX Runtime：decode 图把 cache 暴露为成对的input.<name>/output.<name>。ORT 的CudaSession.set_buffer_sharing（onnxruntime.transformers.io_binding_helper）把每对绑定到同一设备缓冲区，cache 在循环中就地读写，无主机往返。

• ExecuTorch：关闭ExecutorchConfig中的内存规划分配，使 in-place 写入落在调用方自己的张量上：

  config = ExecutorchConfig( backend="xnnpack", dynamic=True, alloc_graph_input=False, alloc_graph_output=False, alloc_mutable_buffers=False, )

  说明：零拷贝原地写入还需要调用方在运行时通过Method::set_output_data_ptr绑定输出缓冲区——Python 运行时没有暴露该 API（executorch.runtime.Method只提供execute/set_inputs/get_outputs）。上述标志只是完成了设置，真正的 in-place 写入是C++ 专属路径（见下文 ExecuTorch decode 循环示例）。从 Python 侧只能在每步从方法输出读回更新后的 cache。

decode 循环推理示例

三个后端的循环形态完全相同——全程使用同一张图：从空的定长 cache 出发，一次性喂入整个提示词（空 cache → prefill），随后逐 token 解码（已填充 cache → decode）。每次调用传入input_ids、因果attention_mask和position_ids（每步推进新增 token 数）以及 cache，返回每个查询位置的 logits；token 落在 cache 何处由静态 cache 内部跟踪，调用方无需额外传参。差异只在于 cache 的搭建方式：Dynamo 用StaticCache对象，ONNX Runtime 用原始设备缓冲区，ExecuTorch 在 C++ 中用调用方数组。

Dynamo：torch.export把静态 cache 写入记录为USER_INPUT_MUTATION，因此加载后的module()会直接更新你传入的StaticCache——一个 cache 携带整轮循环的状态，无需绑定或回传。register_pytree_node(StaticCache)让torch.export.load能够展开StaticCache输入。cache 必须预先初始化（torch.export 把已分配的 K/V 烙进输入规格，惰性空 cache 无法匹配）——但保存的程序自带example_inputs，复用其中已初始化的StaticCache模板并重置为空即可：

import copy import torch from transformers import StaticCache from transformers.exporters.exporter_dynamo import register_pytree_node register_pytree_node(StaticCache) exported = torch.export.load("decode.pt2") decode = exported.module() # 在其输入 / cache 所在设备上运行（此处为 CUDA） # 产物携带一个已初始化的 StaticCache 模板——复用它（重置为空） _, example_kwargs = exported.example_inputs past_key_values = copy.deepcopy(example_kwargs["past_key_values"]) past_key_values.reset() def causal_mask(positions, cache_len): # [1, 1, len(positions), cache_len] return (torch.arange(cache_len, device="cuda")[None, :] <= positions[:, None])[None, None] # prefill：一次调用喂入整个提示词 positions = torch.arange(prompt_len, device="cuda") logits = decode(input_ids=prompt_ids, attention_mask=causal_mask(positions, max_cache_len), position_ids=positions[None], past_key_values=past_key_values).logits next_token = logits[:, -1:].argmax(-1) # decode：query=1 缓冲区原地复用 input_ids = torch.empty((1, 1), dtype=torch.long, device="cuda") position_ids = torch.empty((1, 1), dtype=torch.long, device="cuda") attention_mask = torch.empty((1, 1, 1, max_cache_len), dtype=torch.bool, device="cuda") slots = torch.arange(max_cache_len, device="cuda") for position in range(prompt_len, max_cache_len): input_ids.copy_(next_token) position_ids.fill_(position) attention_mask[0, 0, 0].copy_(slots <= position) logits = decode(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values).logits next_token = logits[:, -1:].argmax(-1)

ONNX Runtime：ORT 原样执行该图，原地 cache 更新借助CudaSession（onnxruntime.transformers.io_binding_helper，ORT io-binding 的薄封装）完成。set_buffer_sharing把 cache 的input.<name>与对应output.<name>绑定到同一个设备缓冲区，修改后的 K/V/计数器直接写回输入；allocate_buffers只分配其余（非共享）输出——这里只有logits；infer(feed_dict)按指针绑定你的 CUDA 张量并执行。cache 缓冲区直接来自图自身的输入元数据（get_inputs()的 shape 与 type），无需模型配置——唯一的符号轴（cache 长度）取值为max_cache_len：

import torch import onnxruntime as ort from onnxruntime.transformers.io_binding_helper import CudaSession, TypeHelper def causal_mask(positions, cache_len): return (torch.arange(cache_len, device="cuda")[None, :] <= positions[:, None])[None, None] session = ort.InferenceSession("decode.onnx", providers=["CUDAExecutionProvider"]) cuda = CudaSession(session, torch.device("cuda")) # 由每个 cache 输入自身的 shape/dtype 构建全新设备 cache 缓冲区； # 把每个 input.<name>/output.<name> 对共享到同一缓冲区，更新即可落地 cache = {} for i in session.get_inputs(): if not i.name.startswith("input."): continue name = i.name[len("input.") :] dims = [max_cache_len if isinstance(d, str) and not d.isdigit() else int(d) for d in i.shape] cache[name] = torch.zeros(dims, dtype=TypeHelper.ort_type_to_torch_type(i.type), device="cuda") cuda.set_buffer_sharing(f"input.{name}", f"output.{name}") cache_feed = {f"input.{name}": buf for name, buf in cache.items()} vocab_size = next(o.shape[-1] for o in session.get_outputs() if o.name.endswith("logits")) # prefill：一次调用喂入整个提示词 positions = torch.arange(prompt_len, device="cuda") cuda.allocate_buffers({"logits": (1, prompt_len, vocab_size)}) out = cuda.infer({"input_ids": prompt_ids, "attention_mask": causal_mask(positions, max_cache_len), "position_ids": positions[None], **cache_feed}) next_token = out["logits"][:, -1:].argmax(-1) # decode：query=1 缓冲区原地复用 cuda.allocate_buffers({"logits": (1, 1, vocab_size)}) input_ids = torch.empty((1, 1), dtype=torch.long, device="cuda") position_ids = torch.empty((1, 1), dtype=torch.long, device="cuda") attention_mask = torch.empty((1, 1, 1, max_cache_len), dtype=torch.bool, device="cuda") slots = torch.arange(max_cache_len, device="cuda") for position in range(prompt_len, max_cache_len): input_ids.copy_(next_token) position_ids.fill_(position) attention_mask[0, 0, 0].copy_(slots <= position) out = cuda.infer({"input_ids": input_ids, "attention_mask": attention_mask, "position_ids": position_ids, **cache_feed}) next_token = out["logits"][:, -1:].argmax(-1)

ExecuTorch：端侧运行时是 C++，原地 cache 更新依赖Method::set_output_data_ptr——如前所述，Python 运行时未暴露它，因此零拷贝 decode 是 C++ 专属路径。把每个被修改 cache 的输出绑定到对应的 cache输入缓冲区，新的 K/V/计数器就会落入调用方的StaticCache缓冲区且零拷贝。与其他示例一样，形状与大小来自产物本身——这里是程序的method_meta（input_tensor_meta/output_tensor_meta→TensorInfo::nbytes()），C++ 侧没有 Python 运行时可供查询（示例省略了 Result 错误检查）：

#include <executorch/extension/data_loader/file_data_loader.h> #include <executorch/extension/tensor/tensor_ptr.h> #include <executorch/runtime/executor/method.h> #include <executorch/runtime/executor/program.h> using namespace executorch::runtime; using executorch::extension::FileDataLoader; using executorch::extension::make_tensor_ptr; // 加载导出的 decode 及其 `forward` 方法 auto loader = FileDataLoader::from("decode.pte"); auto program = Program::load(&loader.get()); // 方法执行内存：一个固定 arena 用于簿记 + 按方法内存规划每块一个缓冲区 std::array<uint8_t, 4 * 1024 * 1024> arena; MemoryAllocator method_allocator(arena.size(), arena.data()); auto meta = program->method_meta("forward"); std::vector<std::vector<uint8_t>> planned(meta->num_memory_planned_buffers()); std::vector<Span<uint8_t>> planned_spans; for (size_t i = 0; i < planned.size(); ++i) { planned[i].resize(meta->memory_planned_buffer_size(i).get()); planned_spans.push_back({planned[i].data(), planned[i].size()}); } HierarchicalAllocator planned_allocator({planned_spans.data(), planned_spans.size()}); MemoryManager memory_manager(&method_allocator, &planned_allocator); auto decode = std::move(program->load_method("forward", &memory_manager).get()); // 形状与大小来自 method_meta——无需外部配置。 // 输入为 [ids, mask, position_ids, cache×N]；输出为 [返回 cache×N, logits, 修改后 cache×N]。 const size_t num_cache_tensors = meta->num_inputs() - 3; const size_t logits_out_idx = num_cache_tensors; std::vector<size_t> cache_nbytes(num_cache_tensors), cache_out_idx(num_cache_tensors); for (size_t i = 0; i < num_cache_tensors; ++i) { cache_nbytes[i] = meta->input_tensor_meta(3 + i)->nbytes(); cache_out_idx[i] = num_cache_tensors + 1 + i; // 修改后的 cache = 最后 N 个输出 } const size_t logits_nbytes = meta->output_tensor_meta(logits_out_idx)->nbytes(); // 一组 StaticCache 缓冲区（K/V + 每层计数器）跨步复用，每次调用原地绑定 auto forward = & { decode.set_input(EValue(*input_ids), 0); decode.set_input(EValue(*mask), 1); decode.set_input(EValue(*position_ids), 2); for (int i = 0; i < num_cache_tensors; ++i) decode.set_input(EValue(*cache_tensor[i]), 3 + i); // 把每个修改后 cache 输出绑定到同一输入张量的数据 → 写入原地落地，零拷贝 for (int i = 0; i < num_cache_tensors; ++i) decode.set_output_data_ptr(cache_tensor[i]->mutable_data_ptr(), cache_nbytes[i], cache_out_idx[i]); decode.set_output_data_ptr(logits_data, logits_nbytes, logits_out_idx); decode.execute(); // cache 原地更新；logits 写入 logits_data return argmax_last(logits_data); // 贪心选择 }; // 先 prefill 整个提示词，再每步解码一个 token——cache 在所有调用间原地传递 int64_t next_token = forward(prompt_ids, prompt_mask, prompt_positions); for (int64_t position = prompt_len; position < max_cache_len; ++position) { next_token = forward(make_tensor_ptr({1, 1}, &next_token, ScalarType::Long), causal_mask(position), // [1, 1, 1, max_cache_len] bool make_tensor_ptr({1, 1}, &position, ScalarType::Long)); }

限制与绕行方案

torch.export、torch.onnx.export和 ExecuTorch 各自在特定 PyTorch 模式上存在毛边。导出器用一组可逆补丁与 FX 级修复，在导出流程的明确节点处理这些问题。公共exportAPI 看不到这些细节，但最常见、值得知道的有两点：

• FlashAttention 与 FlexAttention 在任何后端都不可导出。应优先使用sdpa，eager也可用（更慢）。若模型当前配置的不是其中之一，在调用export前先设置好。
• grouped_mm在DynamoExporter下可以正常追踪，并被OnnxExporter自动翻译；对使用 XNNPACK 后端的ExecutorchExporter，由于 XNNPACK 没有_grouped_mm.out内核，导出器会把 MoE 专家替换为batched_mm。

从源码看，这类绕行的组织方式也值得留意：补丁与修复通过@register_patch(backend, *paths)/@register_fx_node_fix/@register_fx_program_fix三个注册表按后端键控（utils.py），在torch.export追踪块内安装、块退出时还原，因此对 modeling 代码零侵入。

延伸阅读

• 想为新架构或新后端添加导出支持，可参考 扩展导出器文档 中介绍的 patch 与 fix 注册表机制；
• 相关测试见 tests/exporters/test_export.py 与 tests/test_executorch.py，它们以锁定版本验证了上文各后端流水线的行为；
• 核心源码入口：src/transformers/exporters/init.py（公共 API 清单）、src/transformers/exporters/base.py（基类与export_for_generation）、src/transformers/exporters/configs.py（三个配置 dataclass）、src/transformers/exporters/utils.py（decompose_for_generation等共享工具）。

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers