[Tech Review] Core AI: 온디바이스 생성형 모델을 위한 변환 스택의 재설계

안녕하세요, 옵트에이아이 박성재입니다.

아래 내용은 저희 팀이 온디바이스 LLM 배포 파이프라인을 분석하는 과정에서 정리한 기술 리뷰입니다. LLM이나 Diffusion 같은 생성형 모델을 디바이스 위에서 실행하는 일은 기존 Vision/CNN 모델 배포와는 성격이 다릅니다. KV 캐시라는 가변 상태를 매 토큰마다 갱신해야 하고, prefill과 decode라는 서로 다른 실행 모드를 오가야 하며, 대형 vocabulary에 대한 출력까지 처리해야 합니다.

Core AI는 Apple이 WWDC 2026에서 발표한 온디바이스 생성형 AI 프레임워크로, 신경망·트랜스포머 워크로드에서 Core ML을 잇는 공식 후속 스택입니다. Apple Intelligence를 구동하는 것과 동일한 추론 기술을 서드파티 개발자에게 개방한 것으로, CPU/GPU/Neural Engine을 하나의 Swift API로 아우르고 AOT 컴파일과 디바이스별 specialization을 지원합니다. PyTorch 모델은 coreai-torch 변환기를 통해 .aimodel로 변환되며, Apple이 공개한 coreai-models 저장소에서 Qwen, Mistral 등 주요 모델의 export 레시피를 제공합니다. 본 리뷰에서는 Qwen3를 예시로, 기존 coremltools 기반 파이프라인(ANEMLL)과의 구조적 차이를 세 가지 축에서 살펴봅니다.

1. 서론: 두 파이프라인의 전체 구조

ANEMLL은 HF 모델을 ANE(Apple Neural Engine)용으로 재구현한 뒤 torch.jit.trace와 coremltools를 거쳐 .mlmodelc를 산출하는 방식입니다.

HF checkpoint
   │  가중치 로드
   ▼
QwenForCausalLM  (anemll/models/qwen_model.py)   ← HF 모델을 ANE용으로 재구현
   │    - 모든 Linear → nn.Conv2d(kernel=1), 입력 [B, hidden, 1, seq]
   │    - KV 캐시 = register_buffer("kv_cache_0", ...)  (K/V 통합 단일 버퍼)
   │    - forward 입력: input_ids, position_ids, causal_mask, current_pos, update_mask
   ▼
torch.jit.trace(wrapper, sample_inputs)          ← 정적 그래프 캡처
   ▼
ct.convert(traced,
           inputs=[...], outputs=[logits1..16],   ← 큰 vocab → 16분할 출력
           states=[ct.StateType(... kv_cache_0 ...)],
           compute_units=CPU_AND_NE,
           minimum_deployment_target=iOS18,
           convert_to="mlprogram")                ← coremltools MIL로 변환
   ▼
cto.coreml.palettize_weights(...)                ← LUT4/8 양자화 (변환 후)
   ▼
.mlpackage ──compile──► .mlmodelc  (+ dedup, chunk 결합)
   ▼
chat.py / Swift:  state = model.make_state();  model.predict(inputs, state)

Core AI는 torch.export로 그래프를 추출한 뒤 자체 MLIR 기반 컴파일러로 .aimodel을 만들어 냅니다. 압축(양자화/팔레타이제이션)이 파이프라인에 내장되어 있으며, 토크나이저와 메타데이터까지 하나의 에셋으로 번들링됩니다.

HF checkpoint
   │  AutoConfig → model_type 판별 → registry에서 모델 클래스 resolve
   ▼
Qwen3 (coreai_models/models/ios/qwen3.py)        ← HF 가중치만 차용, 자체 구현
   │
   ▼ 압축이 파이프라인에 내장:
   │    macOS: QuantizerConfig + C4 캘리브레이션
   │    iOS  : KMeansPalettizerConfig
   ▼
torch.export.export(model.extend / .prompt_opt / .gather_embeddings / .load_embeddings,
                    dynamic_shapes={... torch.export.Dim ...})  ← 4개의 ExportedProgram
   ▼
TorchConverter()
   .add_exported_program(... entrypoint_name=..., state_names=["key_cache","value_cache"])
   .to_coreai()
   .set_static_shape_config({(cache_len, q_len) 조합})           ← enumerated 특수화
   .set_hardware_constraints(IOSurface, interleave=8, BC1S)     ← 메모리 배치를 선언
   .optimize()                                                  ← MLIR 최적화 패스
   ▼
AIProgram → .aimodel
   ▼
bundle_llm_asset(): tokenizer/ + metadata.json 자동 동봉
   ▼
Swift CoreAILanguageModels 런타임에서 로드·generate

ANEMLL은 컴파일러에 넘기기 전에 사람이 그래프를 하드웨어에 맞춰 직접 다듬고, Core AI는 모델 코드를 평범하게 유지한 채 변환기와 컴파일러가 하드웨어 특화를 담당합니다. 이 철학 차이가 이후의 모든 설계 차이로 이어집니다.

2. 핵심 차이 1: 그래프 캡처 — torch.jit.trace vs torch.export

두 스택은 PyTorch 그래프를 캡처하는 방식부터 다릅니다.

trace가 생성하는 TorchScript IR은 C++ 측 torch::jit::Graph 객체입니다. SSA 형태의 Value/Node/Block 구조로 구성되며, aten op 외에 prim op가 혼재합니다.

graph(%self : __torch__.Model,
      %x : Float(1, 10)):
  %w  : Tensor = prim::GetAttr[name="weight"](%self.linear)   # 가중치를 모듈 속성에서 조회
  %b  : Tensor = prim::GetAttr[name="bias"](%self.linear)
  %1  : Tensor = aten::linear(%x, %w, %b)                     # 고수준 aten op 유지
  %d  : int    = prim::Constant[value=-1]()                   # 상수도 prim 노드
  %2  : Tensor = aten::softmax(%1, %d, ...)
  return (%2)

반면 torch.export가 생성하는 ExportedProgram은 ATen 수준의 FX 그래프 기반입니다. functionalize 과정을 거쳐 in-place 연산과 mutation이 제거된 함수형 그래프로 정규화되며, 가중치는 그래프의 입력으로 승격됩니다.

class GraphModule(torch.nn.Module):
    def forward(self, p_linear_weight, p_linear_bias, x):   # 가중치가 입력으로 승격
        permute  = torch.ops.aten.permute.default(p_linear_weight, [1, 0])
        addmm    = torch.ops.aten.addmm.default(p_linear_bias, x, permute)
        softmax  = torch.ops.aten._softmax.default(addmm, -1, False)
        return (softmax,)

# Graph signature:
#   inputs  = [p_linear_weight (PARAMETER), p_linear_bias (PARAMETER), x (USER_INPUT)]
#   outputs = [softmax (USER_OUTPUT)]

제어 흐름 처리도 다릅니다. trace는 실제 지나간 경로만 기록하므로, 데이터 의존적 분기가 있으면 예시 입력이 통과한 분기만 그래프에 남습니다.

def forward(self, x):
    if x.sum() > 0:      # trace는 예시 입력이 통과한 분기만 기록
        return x * 2
    else:
        return x * 3

export는 이러한 분기를 감지하여 오류를 발생시키거나 torch.cond 같은 명시적 제어 흐름 연산자를 요구합니다. 잘못된 그래프를 조용히 생성하는 대신, 변환 시점에 문제를 드러내는 방식입니다.

ExportedProgram은 그래프 하나가 아니라 graph_module + graph_signature + state_dict + range_constraints를 묶은 컨테이너입니다. 분해된 순수 ATen op만 존재하고, 함수화되어 있으며, 가중치가 입력으로 승격되어 있어 — 변환기가 노드 단위로 순회하며 IR로 매핑하기에 규칙적이고 예측 가능한 입력이 됩니다.

이 규칙성 위에서 TorchConverter는 ATen op를 Core AI op로 lowering합니다. lowering rule은 coreai_torch/_aten_to_core.py에 정의되어 있습니다.

def replace_softmax(
    values_map: dict[str, Value], node: fx.Node, loc: Location
) -> Value:
    x = _get_operand(values_map, node, 0)
    dim = node.args[1]
    return coreai.softmax(x, dim + x.type.rank if dim < 0 else dim)

전체 변환 흐름은 다음과 같이 정리됩니다.

to_coreai()
  ↓
_get_graph_op()
  ↓
for node in exported_program.graph_module.graph.nodes
  ↓
_get_operation(node)
  ↓
_handle_call_function_op(node)
  ↓
namespace == "aten"
  ↓
_aten_to_core_resolver[target](values_map, node, loc)
  ↓
coreai.*Op 생성
  ↓
_values_map[node.name] = op_result

3. 핵심 차이 2: KV 캐시 — 모델이 소유 vs 변환기가 승격

ANEMLL은 모델이 캐시를 소유합니다. 캐시 텐서가 register_buffer로 모델의 일부가 되며, K와 V가 (num_layers*2, n_kv_heads, state_length, head_dim) shape의 단일 버퍼에 통합됩니다.

# anemll/models/qwen_model.py — K/V 통합 단일 버퍼 등록
cache_size = (
    2 * config.num_hidden_layers,
    config.num_key_value_heads,
    config.state_length,
    self.head_dim
)
self.register_buffer("kv_cache_0", torch.zeros(cache_size, dtype=MODEL_DTYPE, device=TEST_DEVICE))

in-place 슬라이스 대입이 trace에 포착되어 ct.StateType으로 변환되고, 추론 시에는 make_state() → predict(inputs, state) 형태로 런타임이 캐시를 암묵적으로 관리합니다. state_length(물리 길이)와 context_length(논리 길이)를 분리해 슬라이딩/회전을 처리합니다.

Core AI는 변환기가 입력을 state로 승격합니다. 모델 코드는 캐시를 평범한 입출력 텐서(key_cache → new_k_cache)로 다루고, state 선언과 메모리 배치는 변환 단계에서 선언적으로 지정됩니다.

# coreai_models/export/ios.py
state_names = ["key_cache", "value_cache"]
converter.add_exported_program(extend_program, input_names=..., state_names=state_names, ...)

cache_constraints = HardwareConstraints(
    AllocationType.IOSurface,
    interleave=[1, 1, KV_CACHE_INTERLEAVE_FACTOR, 1, 1],   # ANE 친화 정렬
    alignments=[1, 1, 1, 1, KV_CACHE_INTERLEAVE_FACTOR * max_context_length, 1],
)

캐시 shape은 (num_layers, 1, kv_cached_embed_size, 1, cache_len) — 시퀀스가 마지막 축인 BC1S 레이아웃입니다.

iOS와 macOS의 분리. Core AI는 두 플랫폼의 경로를 분리해 관리합니다. ANE는 본래 CNN/Vision 워크로드를 위해 설계된 가속기로, static shape과 특정 메모리 레이아웃(B, C, 1, S)을 선호하며 op 집합이 제한적이고 디스패치 오버헤드가 큽니다. 반면 macOS의 주력 타겟은 GPU(Metal)로, 범용 컴퓨트이므로 dynamic shape과 유연한 레이아웃이 가능합니다. 같은 LLM이라도 ANE 친화와 GPU 친화를 동시에 만족하는 코드는 사실상 불가능하므로, iOS는 캐시 길이 고정 + enumerated 특수화, macOS는 torch.export.Dim("k_seq_len", ...)을 통한 동적 시퀀스 길이 방식을 택합니다.

여기서 한 가지 문제가 발생합니다. torch.export는 그래프를 functionalize하여 in-place 연산을 제거하지만, KV 캐시 갱신은 본질적으로 버퍼의 특정 슬라이스를 제자리에서 덮어쓰는 mutation입니다. 이를 일반 ATen op 조합(slice + copy)으로 분해하면 컴파일러가 "state 업데이트"라는 의도를 인식할 수 없게 됩니다.

Core AI는 이를 custom op로 해결합니다.

@torch.library.custom_op("coreai::mutable_slice_update", mutates_args=["x"])
def mutable_slice_update(
    x: Tensor,
    update: Tensor,
    begin: Tensor,
    end: Tensor,
) -> Tensor:
    """
    Mutable slice update operation for cache updates.
    Updates a slice of tensor x with the update tensor using dynamic begin/end indices.
    """
    begin = torch.split(begin, 1, dim=0)
    end = torch.split(end, 1, dim=0)
    slices = tuple(slice(b.item(), e.item()) for b, e in zip(begin, end, strict=False))
    x[slices] = update
    # Note: Not actually in-place for torch
    return x.clone()

등록 과정은 다음과 같습니다. mutates_args=["x"]는 이 op가 입력 x를 변경하는 side effect를 가진다는 선언이며, register_fake는 실제 계산 없이 출력 shape/dtype을 알려주는 fake(meta) 구현입니다.

@torch.library.custom_op 로 PyTorch dispatcher에 op 등록
  ↓
eager 실행용 구현 등록
  ↓
@register_fake 로 fake/meta 구현 등록
  ↓
torch.export / tracing 시 shape 추론 가능
  ↓
Core AI converter가 이 op를 state update primitive / custom lowering으로 인식

핵심은 KV 캐시 갱신을 coreai::mutable_slice_update(x, update, begin, end)라는 단일 노드로 그래프에 보존한다는 점입니다. 컴파일러는 이 노드를 단순 slice/copy가 아닌 mutable state buffer의 슬라이스 갱신 연산으로 이해하고, lowering 단계에서 state update primitive 또는 타겟 특화 커널로 낮춥니다. Core AI 문서에서도 기존 PyTorch 모델 변환 외에 composite op, custom op lowering, inline Metal GPU kernel을 통한 모델 authoring을 지원한다고 설명합니다.

동일한 KV 캐시 갱신을 두 스택이 어떻게 표현하는지 Qwen3 변환 코드로 나란히 놓고 보면 차이가 선명하게 드러납니다.

ANEMLL은 앞서 본 것처럼 모델 안의 통합 버퍼에 in-place 슬라이스 대입을 수행하고, 변환 시 qwen_converter.py의 GetTransformerStates가 해당 버퍼를 StateType으로 선언하여 coremltools가 state로 인식하게 합니다.

# anemll/models/qwen_model.py — prefill 시 in-place 슬라이스 갱신
key_idx = layer_in_group_idx
value_idx = layer_in_group_idx + layers_per_group

# Store the full sequence length in prefill mode
seq_length = key_states.shape[2]  # Get actual sequence length
kv_cache[key_idx:key_idx + 1, :, current_pos:current_pos + seq_length, :] = key_states
kv_cache[value_idx:value_idx + 1, :, current_pos:current_pos + seq_length, :] = value_states

# anemll/ane_converter/qwen_converter.py — GetTransformerStates
num_layers_this_part = num_layers * 2   # K+V 통합이므로 레이어 수의 2배

states = [
    ct.StateType(
        wrapped_type=ct.TensorType(
            shape=(
                num_layers_this_part,
                model.config.num_key_value_heads,
                model.config.state_length,
                head_dim,
            ),
            dtype=np.float16,
        ),
        name=f"{prefix}kv_cache_0",  # Only one group for unified cache
    )
]

Core AI(coreai_models/models/ios/qwen3.py)는 custom op로 mutation을 표현하고, torch.export가 functionalize한 뒤 변환 단계에서 state로 선언합니다. 동적 begin/end 텐서로 슬라이스 위치를 지정하므로 그래프 재생성 없이 임의 위치 갱신이 가능합니다.

begin, end = self.gen_slice_args(layer_idx, offset, num_token_updates)
# update k — iOS는 dimension 4(마지막 차원)에서 갱신
mutable_slice_update(x=self._k_cache, update=k.unsqueeze(0), begin=begin, end=end)
# update v
mutable_slice_update(x=self._v_cache, update=v.unsqueeze(0), begin=begin, end=end)

4. 핵심 차이 3: 컴파일러 스택 — MIL vs MLIR

[ANEMLL]
TorchScript → coremltools MIL(자체 IR, ≠MLIR) → MIL 최적화 패스 → mlprogram(.mlmodelc)
                                                                → Apple 비공개 ANE 컴파일러(espresso)

[Core AI]
ExportedProgram → MLIR(coreai dialect) → MLIR 최적화 패스(optimize) → AIProgram(.aimodel)
                                                                → Apple 비공개 ANE 런타임/컴파일러

Core AI의 IR은 MLIR입니다. export/mlir_ops.py가 from coreai._compiler.dialects import coreai, coreaix로 자체 MLIR dialect를 사용하며, .to_mlir() lowering과 post-MLIR INT4 quantization까지 수행합니다.

두 스택의 철학 차이가 이 지점에서 가장 분명하게 드러납니다. ANEMLL은 컴파일러에 넘기기 전에 PyTorch 그래프를 사람이 ANE 형태로 미리 다듬어 둡니다(1×1 Conv2d, 헤드별 SDPA, logits 16분할, register_buffer 캐시). 컴파일러의 자유도를 줄이는 대신 결과를 통제하고 예측할 수 있으며, anemll-profile로 검증합니다. Core AI는 RoPE/RMSNorm/SDPA 같은 패턴을 named composite로 보존한 채 MLIR에 전달하고, MLIR 최적화 패스와 HardwareConstraints를 통해 fusion과 레이아웃 결정을 컴파일러에 위임합니다. 자동화 수준과 최적화 상한을 높이는 방향입니다.

5. 그래프 분할과 prefill/decode

LLM 추론은 프롬프트 전체를 한 번에 처리하는 prefill과 토큰을 하나씩 생성하는 decode로 나뉘며, 이를 담아내는 방식도 두 스택이 다릅니다.

Core AI iOS export가 산출하는 4개의 entrypoint는 각각 load_embeddings(임베딩 테이블 반환), gather_embeddings(토큰 ID → 임베딩), extend(decode 모드 forward), prompt_opt(prefill 모드 forward)를 담당합니다(export/ios.py). ANEMLL이 모델을 여러 파일로 분할하고 후처리 스크립트로 가중치 중복을 제거하는 반면, Core AI는 하나의 모듈에서 네 entrypoint를 각각 export하여 단일 멀티펑션 에셋으로 묶습니다. 가중치 공유가 구조적으로 보장되며 배포 단위도 단순해집니다.

대형 vocabulary 처리 방식도 대비됩니다. ANEMLL이 logits 출력을 16개로 분할해 ANE 제약을 우회했다면, Core AI는 lm_head를 nn.Linear(hidden_size, vocab_size) 단일 레이어로 유지하는 대신, vocab_size × hidden_size 크기의 거대한 임베딩 테이블을 트랜스포머 그래프에서 분리합니다. load_embeddings가 int8로 저장된 테이블을 반환하고, gather_embeddings가 coreai::fused_dequant_gather_reshape custom op로 dequantization·gather·reshape을 하나의 연산으로 융합해 필요한 토큰 임베딩만 추출합니다. tie_word_embeddings 설정 시 이 테이블은 lm_head와도 공유됩니다.

# coreai_models/primitives/ios/embedding.py — 임베딩 분리를 위한 fused custom op
@torch.library.custom_op("coreai::fused_dequant_gather_reshape", mutates_args=[])
def fused_dequant_gather_reshape(
    embedding_table: torch.Tensor,
    input_ids: torch.Tensor,
    scale: torch.Tensor,
    final_shape: list[int],
) -> torch.Tensor:
    return (embedding_table[input_ids].to(scale.dtype) * scale).reshape(final_shape)

6. 결론: 직접 다듬기에서 위임으로

두 스택의 차이는 한 문장으로 요약됩니다. ANEMLL은 사람이 그래프를 하드웨어에 맞추고, Core AI는 모델을 평범하게 유지한 채 변환기와 컴파일러가 하드웨어에 맞춥니다. 이 철학이 그래프 캡처(trace vs export), KV 캐시(모델 소유 vs 변환기 승격), 컴파일러(MIL 수동 다듬기 vs MLIR 위임), 배포 단위(다중 파일 vs 단일 멀티펑션 에셋)까지 일관되게 이어집니다.

torch.export 기반의 규칙적인 그래프, custom op를 통한 state 의도 보존, MLIR dialect와 선언적 하드웨어 제약 — Core AI는 생성형 모델의 온디바이스 실행을 위해 변환 파이프라인 전체를 다시 설계한 스택이라 할 수 있습니다. 다음 글에서는 실제 모델을 Core AI로 변환하여 디바이스에서 측정한 벤치마크 결과를 다루겠습니다. 많은 관심 부탁드립니다.

torch.export가 산출하는 컨테이너로, functionalize된 ATen 수준 FX 그래프와 graph_signature, state_dict, range_constraints를 함께 담습니다. 변환기 입장에서 규칙적이고 예측 가능한 입력이 됩니다.

모델 코드에서는 평범한 입출력 텐서로 다루던 KV 캐시를, 변환 단계에서 런타임이 관리하는 mutable state로 선언하는 방식입니다. 모델 코드와 하드웨어 특화를 분리하는 핵심 장치입니다.

시퀀스가 마지막 축에 오는 (Batch, Channel, 1, Sequence) 형태의 메모리 배치로, ANE가 선호하는 레이아웃입니다.

동적 shape을 지원하지 않는 하드웨어를 위해, 허용되는 (cache_len, q_len) 조합을 열거하여 하나의 모델 안에 특수화된 그래프들을 담는 기법입니다.

참고 자료

1. Apple, "Integrate on-device AI models into your app using Core AI" — WWDC26 Session 326 (developer.apple.com/videos/play/wwdc2026/326/)
2. InfoQ, "Apple Launches Core AI for Apple-Silicon Optimized On-Device Generative AI" (2026.06)
3. apple/coreai-models — Core AI 공식 모델·export 레시피 저장소 (github.com/apple/coreai-models)
4. ANEMLL — Apple Neural Engine용 LLM 변환 오픈소스 프로젝트 (github.com/anemll/anemll)
5. PyTorch, torch.export documentation — ExportedProgram, Dim, torch.cond
6. PyTorch, torch.library documentation — custom_op, register_fake
7. coremltools documentation — MIL, StateType, palettize_weights