Exporters

Exporters

어떤 PreTrainedModel이든 타깃 런타임과 무관하게 ONNX, ExecuTorch, 또는 독립형 PyTorch 프로그램으로 내보낼 수 있어요.

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

exporters는 다운스트림 라이브러리 대신 Transformers 안에 살고 있어서, 아키텍처 변경, 새 attention 패턴, 커스텀 캐시 타입이 모델링 코드에 들어오는 즉시 export 시점에 지원돼요.

[!WARNING] exporters는 실험적입니다. 이 모듈의 많은 patch들은 특정 업스트림 버그(Torch, ONNX Script, ONNX Runtime, ExecuTorch)를 우회하며, 수정이 업스트림에 반영되는 즉시 제거될 것입니다. API가 안정화될 때까지 patch들이 테스트 스위트에서 사용된 버전에 묶여 있다고 간주하세요. 프로덕션 툴링에서는 그 버전들을 고정하고, 업스트림 변경이 반영됨에 따라 새 patch가 나타나고 옛 patch가 사라질 것을 기대하세요.

Exporter	Output	Runtime
DynamoExporter	ExportedProgram	Any PyTorch runtime, AOT compilation
OnnxExporter	ONNXProgram	Any ONNX runtime (ORT, TensorRT, OpenVINO)
ExecutorchExporter	ExecutorchProgramManager	Mobile and edge devices (ExecuTorch)

AutoHfExporter는 config에서 올바른 exporter를 선택하고, AutoExportConfig는 dict에서 올바른 config 클래스를 선택해요. 둘 다 Transformers의 동일한 auto-class 패턴을 따르는데, 이는 백엔드가 호출 지점에서 하드코딩되는 대신 런타임에 선택될 때 유용해요.

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)

출처: 문서

본문

설치

내보낼 백엔드의 의존성을 설치해요.

[!TIP] 아래 버전들은 exporter 테스트 스위트가 고정된 버전입니다. 더 새롭거나 오래된 릴리스도 보통 동작하지만, exporter patch들은 특정 API 표면을 타겟으로 하므로 프로덕션 툴링에서는 이것들을 고정하고 HfExporter가 드리프트(drift)를 감지하면 경고를 로그하도록 기대하세요.

pip install transformers "torch==2.12.0"

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

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

모델 내보내기

모든 exporters는 동일한 인터페이스를 공유해요. config로 exporter를 만들고 export()를 호출해요.

exporter 클래스를 바꿔서 런타임을 전환해요.

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)

# run the exported graph directly
outputs = exported.module()(**inputs)

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)

# save and load with 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)

backend는 기본값이 xnnpack이며 CPU를 타겟팅해서 CPU 전용 설치에서도 작동해요. 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)

# save for on-device deployment
et_program.save("model.pte")

# load and run via the ExecuTorch Python runtime
from executorch.runtime import Runtime

program = Runtime.get().load_program("model.pte")
method = program.load_method("forward")
outputs = method.execute(list(inputs.values()))

동적 형태 (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}},
    # Emit data-dependent shape guards as runtime asserts instead of failing the export when a
    # guard wouldn't hold across the explicit symbolic range. Most LLMs need this under fine-grained
    # ``Dim(min=, max=)`` bounds. Not needed with ``dynamic=True`` / ``Dim.AUTO``, where torch.export
    # infers shape relations instead of verifying them against user-stated bounds.
    prefer_deferred_runtime_asserts_over_guards=True,
)
exported = exporter.export(model, inputs, config=config)

import torch
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!", "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 = OnnxExporter()
config = OnnxConfig(
    dynamic_shapes={"input_ids": {0: batch, 1: seq}, "attention_mask": {0: batch, 1: seq}},
    # Emit data-dependent shape guards as runtime asserts instead of failing the export when a
    # guard wouldn't hold across the explicit symbolic range. Most LLMs need this under fine-grained
    # ``Dim(min=, max=)`` bounds. Not needed with ``dynamic=True`` / ``Dim.AUTO``, where torch.export
    # infers shape relations instead of verifying them against user-stated bounds.
    prefer_deferred_runtime_asserts_over_guards=True,
)
onnx_program = exporter.export(model, inputs, config=config)

import torch
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!", "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 = ExecutorchExporter()
config = ExecutorchConfig(
    backend="xnnpack",
    dynamic_shapes={"input_ids": {0: batch, 1: seq}, "attention_mask": {0: batch, 1: seq}},
    # Emit data-dependent shape guards as runtime asserts instead of failing the export when a
    # guard wouldn't hold across the explicit symbolic range. Most LLMs need this under fine-grained
    # ``Dim(min=, max=)`` bounds. Not needed with ``dynamic=True`` / ``Dim.AUTO``, where torch.export
    # infers shape relations instead of verifying them against user-stated bounds.
    prefer_deferred_runtime_asserts_over_guards=True,
)
et_program = exporter.export(model, inputs, config=config)

생성 모델 (Generative models)

자기회귀 생성의 경우 모델의 forward는 프리필(prefill) 단계(전체 프롬프트, KV 캐시 없음)와 디코드 단계(단일 토큰, 채워진 KV 캐시)에서 다른 형태를 가져요. Exporters는 두 단계를 모두 분해해서 각각 내보내는 ~HfExporter.export_for_generation을 노출해요.

멀티모달 생성 모델의 경우 프리필은 추가로 이미지 또는 오디오 인코더, 언어 모델, lm_head로 나뉘어요. 인코더와 언어 모델 발견은 get_encoder() (modality="image" 또는 "audio") 및 get_decoder() 접근자를 사용하므로, 이들을 사용하는 새 아키텍처는 수정 없이 바로 동작해요.

프로젝터 구성 요소는 모델이 속성 이름(multi_modal_projector, connector, embed_vision, embed_audio)으로 하나를 노출할 때만 나타나요. 아래 Qwen2-VL은 프로젝터를 비전 타워에 접어 넣으므로 그 구성 요소 dict에는 별도의 multi_modal_projector 키가 없어요. 새 아키텍처는 목록을 늘리는 대신 프로젝터 속성을 이 이름 중 하나에 맞춰야 해요.

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}

from transformers import AutoModelForImageTextToText, AutoProcessor
from transformers.exporters import OnnxExporter, OnnxConfig

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 = OnnxExporter()
config = OnnxConfig(dynamic=True)
components = exporter.export_for_generation(model, inputs, config=config)
# components = {"image_encoder": ONNXProgram, "language_model": ONNXProgram, "lm_head": ONNXProgram, "decode": ONNXProgram}

from transformers import AutoModelForImageTextToText, AutoProcessor
from transformers.exporters import ExecutorchExporter, ExecutorchConfig

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 = ExecutorchExporter()
config = ExecutorchConfig(backend="xnnpack", dynamic=True)
components = exporter.export_for_generation(model, inputs, config=config)
# components = {"image_encoder": ExecutorchProgramManager, "language_model": ..., "lm_head": ..., "decode": ...}

[!WARNING] 내보낸 구성 요소들은 독립적인 그래프이며 즉시 실행 가능한 추론 파이프라인이 아닙니다. 호출자는 각 인코더를 실행하고, 임베딩을 프로젝션하며, 생성 루프를 오케스트레이션할 책임이 있습니다.

export_for_generation이 동작하는 방식

decompose_for_generation()은 model.generate(**inputs, max_new_tokens=2)을 한 번 실행하고 model.forward에 훅을 걸어 실제 프리필과 디코드 kwargs를(모델이 멀티모달이면 각 인코더/프로젝터/언어 모델의 훅을 통해 서브모듈별 kwargs도) 캡처해요. 그래서 디코더 전용, SSM, 인코더-디코더, 멀티모달을 포함한 모든 아키텍처에서 모델별 글루(glue) 없이 동작하는 거예요. export_for_generation은 이것 위의 원라이너(one-liner)예요.

캡처는 inputs로 모델을 즉시(eager) 실행하므로, 짧은 프롬프트, 작은 이미지 하나, 몇 개의 오디오 프레임 같은 작지만 대표적인 값들을 전달해요. 내보낸 프로그램은 그 크기에 묶이지 않지만(동적 형태가 여전히 흐름), 더 작은 캡처 입력은 decompose_for_generation을 더 저렴하게 만들고 symbolic-shape 추론을 다루기 쉽게 유지해요.

분해와 내보내기 사이에 행동하려면(검증을 위한 eager forward 실행, 서브모듈 입력 교체, 단계 건너뛰기 등) 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 the eager forward before exporting
    exported[name] = exporter.export(submodel, subinputs, config=config)

멀티 토큰 디코드

기본적으로 decode 구성 요소는 단일 토큰 단계예요. KV 캐시에 대한 쿼리 토큰 하나이므로 torch.export는 쿼리-시퀀스 축을 1로 특수화해요. multi_token_decode=True를 전달하면 decode를 멀티 토큰 디코드로 캡처해요: decompose_for_generation()이 연속된 두 디코드 단계(max_new_tokens=3으로 캡처)를 하나의 forward로 병합해서 그 축이 symbolic으로 유지되게 해요. 그러면 단일 그래프가 모든 쿼리 길이를 처리해요 — 토큰 하나(일반 디코딩), 한 번에 여러 토큰(과거로부터의 연속, 예: 추측(speculative) 토큰 청크 수락), 캐시가 비어 있을 때의 일반 프리필.

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"] now accepts a variable number of query tokens

from transformers.exporters import OnnxExporter, OnnxConfig

exporter = OnnxExporter()
config = OnnxConfig(dynamic=True)
components = exporter.export_for_generation(model, inputs, config=config, multi_token_decode=True)
# components["decode"] now accepts a variable number of query tokens

from transformers.exporters import ExecutorchExporter, ExecutorchConfig

exporter = ExecutorchExporter()
config = ExecutorchConfig(backend="xnnpack", dynamic=True)
components = exporter.export_for_generation(model, inputs, config=config, multi_token_decode=True)
# components["decode"] now accepts a variable number of query tokens

쿼리 축은 동적 형태 export(dynamic=True)에서만 symbolic으로 유지돼요. 정적 export는 이를 캡처된 길이로 고정해서 고정된 멀티 토큰 그래프를 만들어요. 아래 정적 KV 캐시와 결합돼요 — 병합된 디코드가 각 단계의 토큰을 고정 크기 캐시에 제자리에서 쓰고, 캐시가 어디에 놓일지 내부적으로 처리해요.

정적 KV 캐시

generate()는 기본적으로 DynamicCache를 키우며 시퀀스가 늘어남에 따라 재할당해요 — 내보낸 그래프에겐 움직이는 목표죠. 정적 캐시는 고정 크기 버퍼로, 한 번 할당되고 각 단계에서 현재 위치에 제자리에서 쓰여요. 멀티 토큰 디코드와 결합하면 생성을 단일 내보낸 그래프로 압축해요: decode 그래프가 고정 크기 캐시와 가변 개수의 쿼리 토큰을 받으므로, 한 그래프가 프롬프트(빈 캐시 → 프리필)와 각 생성 토큰(채워진 캐시 → 디코드)을 모두 처리해요. cache_implementation="static"(그리고 max_cache_len)이 있는 GenerationConfig를 multi_token_decode=True와 함께 전달해서 내보내요:

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
)

from transformers import GenerationConfig
from transformers.exporters import OnnxExporter, OnnxConfig

exporter = OnnxExporter()
gen_config = GenerationConfig(cache_implementation="static", max_cache_len=2048)
components = exporter.export_for_generation(
    model, inputs, config=OnnxConfig(dynamic=True), generation_config=gen_config, multi_token_decode=True
)

from transformers import GenerationConfig
from transformers.exporters import ExecutorchExporter, ExecutorchConfig

exporter = ExecutorchExporter()
gen_config = GenerationConfig(cache_implementation="static", max_cache_len=2048)
components = exporter.export_for_generation(
    model, inputs, config=ExecutorchConfig(backend="xnnpack", dynamic=True), generation_config=gen_config, multi_token_decode=True
)

decode 그래프는 이제 두 개의 symbolic 축을 가져요 — 쿼리 길이(공급하는 토큰 수)와 캐시 길이(max_cache_len, 로드 시 조정 가능). dynamic=True는 이것들(및 다른 모든 축)을 Dim.AUTO로 표시하므로 내보낸 그래프가 로드 시 어떤 프롬프트 길이와 캐시 크기도 받아들여요.

제로-카피 제자리 업데이트

정적 캐시는 전달되어 제자리에서 변경(mutate)되므로, 런타임이 자체 아레나를 통해 복사하는 대신 호출자의 버퍼를 바인딩하기만 하면 스텝 간에 호스트 복사 없이 하나의 버퍼가 상태를 옮겨요. 그게 유일하게 남은 백엔드별 부분이에요:

• Dynamo — 내보낸 프로그램은 캐시 쓰기를 USER_INPUT_MUTATION으로 모델링하므로, components["decode"].module()(...)을 호출하면 전달하는 캐시 텐서를 직접 업데이트해요. 매 단계마다 동일한 텐서를 재사용하면 돼요. 구성할 게 없어요.

• ONNX Runtime — 디코드 그래프는 캐시를 일치하는 input.<name> / output.<name> 쌍으로 노출해요. ORT의 CudaSession.set_buffer_sharing(onnxruntime.transformers.io_binding_helper)은 각 쌍을 하나의 디바이스 버퍼에 바인딩하므로, 캐시가 호스트 왕복 없이 루프 전체에 걸쳐 제자리에서 읽고 업데이트돼요.

• ExecuTorch — ExecutorchConfig에서 메모리 계획 할당을 꺼서 제자리 쓰기가 호출자 자신의 텐서에 들어가게 해요 (각 플래그가 하는 일은 참조 확인):

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

  [!NOTE] 제로-카피 제자리 쓰기는 또한 호출자가 런타임에서 Method::set_output_data_ptr로 출력 버퍼를 바인딩해야 합니다 — Python 런타임에서는 노출되지 않습니다 (executorch.runtime.Method는 execute/set_inputs/get_outputs만 노출). 위 플래그들이 설정하지만, 제자리 쓰기는 C++ 전용 경로입니다 (아래 ExecuTorch 디코드 루프 예시 참조). Python에서는 매 단계마다 메서드 출력에서 업데이트된 캐시를 다시 읽으세요.

디코드-루프 추론 예시

루프는 모든 백엔드에서 동일한 형태예요 — 전체에 걸쳐 동일한 그래프죠. 빈 고정 크기 캐시로 시작해서, 전체 프롬프트를 한 번 넣고(빈 캐시 → 프리필), 그 다음 한 번에 한 토큰씩(채워진 캐시 → 디코드) 넣어요. 각 호출은 input_ids, 인과(causal) attention_mask, position_ids(매 단계 새 토큰 수만큼 진행), 그리고 캐시를 전달하고 모든 쿼리 위치의 로짓을 돌려받아요. 각 토큰이 캐시의 어디에 놓일지는 정적 캐시가 내부적으로 추적하므로, 호출을 통해 지나갈 추가 작업은 없어요. 캐시가 어떻게 설정되는지는 런타임마다 달라요(Dynamo용 StaticCache 객체, ONNX Runtime용 원시 디바이스 버퍼, ExecuTorch용 C++ 호출자 배열). 그래서 각 탭이 아래에 자신의 것을 만듭니다. Dynamo와 ONNX Runtime 탭은 캐시를 제자리에서 업데이트하고, ExecuTorch의 제자리 경로는 C++입니다 (Python 런타임은 위에서 언급했듯 불가능).

torch.export는 정적-캐시 쓰기를 USER_INPUT_MUTATION으로 기록하므로, 로드된 그래프의 module()은 전달하는 StaticCache를 직접 업데이트해요 — 하나의 캐시가 바인딩하거나 다시 내보낼 것 없이 루프 전체에서 상태를 옮겨요. register_pytree_node(StaticCache)는 torch.export.load가 StaticCache 입력을 언플래튼할 수 있게 해요. 캐시는 사전에 초기화되어야 해요 (torch.export는 할당된 K/V를 입력 스펙에 구워 넣으므로, 지연된 빈 캐시는 일치하지 않음) — 하지만 저장된 프로그램은 자체 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()   # runs on the device its inputs / cache live on (CUDA here)

# the artifact carries an initialized StaticCache template — reuse it (reset to empty)
_, 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: the whole prompt in one call
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 buffers reused in place
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 io-binding의 얇은 래퍼인 ORT의 CudaSession(onnxruntime.transformers.io_binding_helper)으로 수행돼요. set_buffer_sharing은 캐시 input.<name>과 그에 맞는 output.<name>을 하나의 디바이스 버퍼에 바인딩하므로, 변경된 K/V/카운터가 입력에 그대로 다시 쓰여요. allocate_buffers는 남은(공유되지 않은) 출력 — 여기서는 logits만 — 을 할당하고, infer(feed_dict)는 CUDA 텐서를 포인터로 바인딩하고 실행해요. 캐시 버퍼는 그래프 자체의 입력 메타데이터(get_inputs() 형태와 타입)에서 직접 오므로 모델 config가 필요 없어요 — 유일한 symbolic 축(캐시 길이)이 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"))

# fresh device cache buffers built from each cache input's own shape/dtype; share each
# input.<name>/output.<name> pair on one buffer so the update lands in place
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: the whole prompt in one call
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 buffers reused in place
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++이고, 제자리 캐시 업데이트는 Method::set_output_data_ptr에 의존해요 — Python 런타임에서는 노출되지 않으므로 (executorch.runtime.Method는 execute/set_inputs/get_outputs만 노출) 제로-카피 디코드는 C++ 전용 경로예요. 변경된 각 캐시 출력을 일치하는 캐시 입력 버퍼에 바인딩하면 새 K/V/카운터가 복사 없이 호출자의 StaticCache 버퍼에 들어가요. 다른 탭들처럼 형태와 크기는 아티팩트 자체에서 오는데, 여기서는 프로그램의 method_meta(input_tensor_meta/output_tensor_meta → TensorInfo::nbytes())에서 오고 C++에는 질의할 Python 런타임이 없어요:

#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;

// load the exported decode and its `forward` method (Result error-checks elided for brevity)
auto loader = FileDataLoader::from("decode.pte");
auto program = Program::load(&loader.get());

// method-execution memory: a fixed arena for bookkeeping + one buffer per the method's memory plan
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());

// shapes and sizes come from method_meta — no external config. Inputs are [ids, mask, position_ids,
// cache×N]; outputs are [returned cache×N, logits, mutated 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;   // mutated cache = the last N outputs
}
const size_t logits_nbytes = meta->output_tensor_meta(logits_out_idx)->nbytes();

// one set of StaticCache buffers (K/V + per-layer counters) is reused across steps, bound in place each call
auto forward = [&](const TensorPtr& input_ids, const TensorPtr& mask, const TensorPtr& position_ids) {
    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);
    // bind each mutated-cache output onto that same input tensor's data → the write lands in place, zero copies
    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 updated in place; logits written to logits_data
    return argmax_last(logits_data);  // greedy pick
};

// prefill the whole prompt, then decode one token per step — the cache carries in place across all calls
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 패턴 주위에 거친 가장자리(rough edge)가 있어요. exporters는 export 흐름에서 잘 정의된 지점에 적용되는 작고 되돌릴 수 있는 patch 세트와 FX 수준 수정으로 이를 우회해요. 이 중 어떤 것도 공개 export API에서 보이지 않지만, 가장 흔히 알아야 할 것들:

• FlashAttention과 FlexAttention은 어떤 백엔드에서도 내보낼 수 없어요. sdpa가 선호 설정이고 eager도 작동해요(더 느림). 모델이 다른 것을 사용하고 있다면 export를 호출하기 전에 그 중 하나를 설정해요.
• grouped_mm은 DynamoExporter를 통해 잘 추적되고 OnnxExporter에서 자동 번역돼요. XNNPACK 백엔드의 ExecutorchExporter에서는 XNNPACK에 _grouped_mm.out kernel이 없으므로 exporter가 MoE 전문가를 batched_mm으로 교체해요.

더 알아보기 (Learn more)

• Extending the exporters의 patch 및 fix 레지스트리로 새 아키텍처 또는 백엔드의 export 지원을 추가해 보세요.