추론 파이프라인¶

목차¶

텍스트 생성의 동작 방식
전체 파이프라인 시퀀스 다이어그램
1단계: 토크나이제이션
토크나이저 백엔드
채팅 템플릿 적용
2단계: 프리필
순전파
프리필 계획
MLX 연산 그래프 실행
3단계: 디코드 루프
룩어헤드 파이프라이닝
EOS 감지와 정지
토큰 스트리밍
4단계: 디토크나이제이션
프로파일링과 통계
VLM 변형: 비전-언어 생성
VLM 파이프라인 시퀀스
이미지 전처리
시각 인코딩과 프로젝션
임베딩 병합과 프리필
after_prefill 훅

범위

이 페이지는 표준 autoregressive text/VLM 경로를 설명합니다. Diffusion 생성, audio STT/TTS, Florence2 계열 sequence-to-sequence 실행, RT-DETR 검출은 전용 runtime을 사용하므로 아래의 모든 KV-cache/decode 단계를 따르지 않습니다.

텍스트 생성의 동작 방식¶

대형 언어 모델은 한 번에 하나의 토큰을 예측하여 텍스트를 생성합니다. 각 단계에서 모델은 지금까지 본 토큰의 전체 이력(프롬프트와 이전에 생성된 모든 토큰)을 받고 신경망 레이어를 거쳐 순전파를 실행하며 어휘에 대한 확률 분포를 생성합니다. 그런 다음 샘플링 전략이 이 분포에서 다음 토큰을 선택합니다.

이를 자기회귀 생성이라고 합니다. 생성된 각 토큰이 다음 단계의 입력 일부가 됩니다.

이 프로세스에는 두 가지 뚜렷한 단계가 있습니다.

프리필: 모델이 전체 입력 프롬프트를 하나의 순전파로 처리합니다. 프롬프트가 사전에 알려져 있기 때문에 모든 토큰을 병렬로 처리할 수 있습니다. 이것이 가장 연산 집약적인 단계입니다.
디코드: 모델이 KV 캐시를 사용하여 전체 이력에 대한 어텐션을 재계산하지 않고 단계마다 새로 생성된 하나의 토큰을 처리합니다. 이 단계는 연산 집약적이라기보다 메모리 대역폭 집약적입니다.

src/lib/mlxcel-core/src/engine/의 mlxcel 디코딩 엔진이 CLI와 서버 양쪽에서 이 과정을 구현하며, 디코드 단계 사이에 GPU가 쉬지 않도록 룩어헤드 파이프라이닝을 사용합니다.

전체 파이프라인 시퀀스 다이어그램¶

sequenceDiagram
    participant User
    participant CLI/Server
    participant Tokenizer
    participant ChatTemplate
    participant DirectEngine
    participant Model
    participant MLX

    User->>CLI/Server: prompt text (+ optional images)
    CLI/Server->>ChatTemplate: apply_chat_template(messages)
    ChatTemplate-->>CLI/Server: formatted prompt string
    CLI/Server->>Tokenizer: encode(formatted_prompt)
    Tokenizer-->>CLI/Server: [token_id_0, token_id_1, ...]

    CLI/Server->>DirectEngine: generate(request, on_token)

    Note over DirectEngine: PREFILL PHASE
    DirectEngine->>Model: forward(PrefillPlan pieces, caches, mask=None)
    Model->>MLX: execute computation graph
    MLX-->>Model: logits [1, seq_len, vocab_size]
    Model-->>DirectEngine: logits
    DirectEngine->>DirectEngine: RowSampler draw → first_token
    DirectEngine->>MLX: async_eval(first_token)

    Note over DirectEngine: DECODE LOOP
    loop until EOS or max_tokens
        DirectEngine->>Model: forward([last_token], caches, mask=None)
        Model->>MLX: execute computation graph (single token)
        MLX-->>Model: logits [1, 1, vocab_size]
        Model-->>DirectEngine: logits
        DirectEngine->>DirectEngine: sample next token (async)
        DirectEngine->>DirectEngine: ffi::item_i32(current_token): sync GPU
        DirectEngine->>CLI/Server: on_token(token_id)
        CLI/Server->>Tokenizer: decode([token_id]) → text fragment
        CLI/Server->>User: stream text fragment
    end

    DirectEngine-->>CLI/Server: all generated token IDs

1단계: 토크나이제이션¶

토크나이저 백엔드¶

mlxcel은 src/tokenizer/mod.rs의 MlxcelTokenizer 열거형으로 세 가지 토크나이저 형식을 지원합니다.

Variant	파일	라이브러리
`HuggingFace`	`tokenizer.json`	HuggingFace tokenizers 크레이트
`SentencePiece`	`tokenizer.model`	sentencepiece 크레이트
`Tiktoken`	`*.tiktoken`	`src/tokenizer/tiktoken.rs`의 커스텀 구현

load_tokenizer() 함수는 모델 디렉터리에서 어떤 형식이 사용 가능한지 감지하고 적절한 variant를 구성합니다. MlxcelTokenizer 열거형에서는 기반 형식이 무엇이든 encode() 및 decode() 인터페이스가 동일합니다.

SentencePiece 백엔드는 특수 토큰 처리를 포함합니다. 특수 토큰은 길이 내림차순으로 정렬되고 나머지 텍스트를 기반 SentencePiece 프로세서에 전달하기 전에 탐욕적으로 매칭됩니다. 덕분에 기본 SentencePiece 어휘에 <|im_start|> 또는 <|begin_of_text|> 같은 특수 토큰을 추가하는 모델을 올바르게 처리합니다.

채팅 템플릿 적용¶

대부분의 현대 모델은 특정 대화 형식으로 파인튜닝되어 있습니다. 단순한 프롬프트 문자열을 전송하면 모델이 학습 형식을 기대하기 때문에 좋지 않은 결과를 낼 수 있습니다.

mlxcel은 Jinja2 템플릿 엔진의 Rust 구현인 minijinja를 사용하여 모델의 채팅 템플릿을 자동으로 적용합니다. 채팅 템플릿은 모델의 tokenizer_config.json 파일에 Jinja2 템플릿 문자열로 저장됩니다.

사용자 템플릿에서 알아 둘 minijinja 동작은 다음과 같습니다. 리스트나 딕셔너리를 바로 출력하면 Python 방식 따옴표를 쓰고, //와 %는 내림 기준으로 계산하며, round는 짝수 쪽으로 반올림하고, 0으로 나누면 오류가 납니다.

템플릿은 메시지 목록(role과 content 필드 포함)을 받아 토크나이저가 인코딩할 형식화된 문자열을 생성합니다. 모델마다 형식이 다릅니다.

Llama 3은 <|begin_of_text|><|start_header_id|>user<|end_header_id|>...를 사용합니다.
Qwen은 <|im_start|>user\n...<|im_end|>를 사용합니다.
Gemma는 <start_of_turn>user\n...<end_of_turn>을 사용합니다.

채팅 템플릿 적용은 CLI에서 --no-chat-template로 비활성화하거나, 서버에서 --chat-template 또는 --chat-template-file로 재정의할 수 있습니다.

2단계: 프리필¶

순전파¶

LanguageModel::forward() 메서드는 모델의 연산으로 들어가는 단일 진입점입니다. 프리필 중에 다음을 받습니다.

input_ids: 전체 프롬프트의 토큰 ID로 이루어진 [1, seq_len] 텐서.
caches: 트랜스포머 레이어당 하나씩 있는 KVCache 구조체의 가변 슬라이스.
mask: 프리필 중에는 None(모델이 내부적으로 인과 마스크를 생성함).

모델은 모든 프롬프트 토큰을 병렬로 처리하고 모든 토큰에 대한 키-값 프로젝션을 캐시에 기록하며 [1, seq_len, vocab_size] 형태의 로짓 텐서를 반환합니다. 마지막 시퀀스 위치의 로짓만 첫 번째 생성 토큰을 샘플링하는 데 사용됩니다.

프리필 계획¶

CLI와 서버는 하나의 PrefillPlan(src/lib/mlxcel-core/src/prefill_plan.rs)으로 프롬프트를 같은 방식으로 나눕니다.

프롬프트 캐시에서 채택한 접두부는 다시 순전파하지 않습니다.
채팅 프롬프트라면 이력 경계(생성 프롬프트를 뺀 렌더링)까지의 구간을 패딩 없는 조각 하나로 처리하므로, 프롬프트 캐시가 정확히 그 토큰들에 해당하는 스냅샷을 저장할 수 있습니다.
나머지는 모델이 여러 번에 걸친 프리필을 지원하고 입력이 미리 병합한 임베딩이 아닌 토큰일 때 프리필 청크 크기 이하의 조각으로 자릅니다.
프리필 타일 정렬을 쓰는 하드웨어에서는 모델이 패딩 프리필을 지원하면 조각을 Neural Accelerator 타일 크기로 패딩합니다. 이력 경계에서 끝나는 조각은 패딩하지 않습니다.

mlxcel generate는 2048 토큰 청크를 씁니다. 서버는 다른 시퀀스가 디코딩 중이 아닐 때 2048, 디코딩 배치가 돌고 있을 때 512를 쓰며, --prefill-chunk-size나 MLXCEL_PREFILL_CHUNK를 명시하면 두 상태 모두 그 값으로 고정됩니다.

MLX 연산 그래프 실행¶

MLX는 지연 평가를 사용합니다. Rust 코드가 MLX 연산을 (ffi 모듈을 통해) 호출하면, 그 연산들은 연산 그래프에 기록되지만 즉시 GPU에서 실행되지 않습니다. 실행은 특정 텐서를 구체화하기 위해 ffi::eval()이 호출될 때 발생합니다.

그래서 MLX는 GPU 작업을 스케줄링하기 전에 연산 그래프를 최적화할 수 있습니다. 연산을 융합하고 효율적인 커널을 선택하는 식입니다. 생성 루프는 ffi::async_eval()을 사용하여 Rust 스레드를 차단하지 않고 연산을 스케줄링한 다음, 토큰 값이 실제로 필요할 때 ffi::item_i32()로 동기화합니다.

체크포인트 가중치는 지연 평가의 예외입니다. 모델 로드가 가중치를 평가해 업로드하므로(모델 로딩 참조) 첫 프리필은 가중치 I/O를 기다리지 않습니다. 비전·오디오 타워 가중치는 요청이 해당 모달리티를 담을 때까지 지연 상태로 남습니다.

프리필 이후, ffi::clear_memory_cache()는 더 이상 필요하지 않은 중간 텐서를 해제하여 디코드 루프가 시작하기 전에 GPU 메모리를 회수합니다.

3단계: 디코드 루프¶

룩어헤드 파이프라이닝¶

모든 디코드 단계는 엔진의 분할 단계(Engine::submit과 Engine::finish_rows)를 거칩니다. mlxcel generate가 쓰는 단일 시퀀스 루프 DirectEngine::generate()와 서버의 BatchScheduler는 모두 이를 룩어헤드 파이프라인으로 구동해 GPU 연산과 CPU 측 처리를 겹칩니다. 각 반복에서:

다음 토큰의 순전파와 샘플링을 시작합니다(GPU 작업이 비동기적으로 제출됨).
이전 반복에서 현재 토큰을 구체화합니다(GPU와 동기화).
현재 토큰을 EOS 토큰 ID에 대해 확인합니다.
현재 토큰 ID와 함께 스트리밍 콜백을 호출합니다.
다음 반복을 위해 "next"를 "current"로 이동합니다.

이 방식 덕에 CPU가 현재 토큰을 처리하는 동안에도 GPU는 유휴 상태에 빠지지 않습니다. EOS 확인, 콜백 디스패치, 이력 유지의 CPU 오버헤드가 GPU 연산 뒤에 숨겨집니다.

융합 추출(fused draw)로 처리할 수 없는 샘플러(이력 페널티, DRY, mirostat, adaptive-p)는 대신 행별 파이프라인으로 순전파를 한 단계 앞서 실행합니다. mirostat과 adaptive-p는 추출 결과를 호스트에서 읽으므로 종료 판정만 다음 순전파와 겹칩니다. MLXCEL_FORCE_SYNC를 설정하면 CLI와 서버 모두에서 파이프라인이 꺼지고 모든 단계가 동기식 Engine::step으로 실행됩니다. 디버깅에 유용합니다.

EOS 감지와 정지¶

생성 루프는 두 가지 EOS 토큰 ID 소스를 병합합니다.

LanguageModel::eos_token_ids()의 모델 내장 EOS 토큰들. 어휘를 기반으로 모델 구현이 설정합니다.
SamplingConfig::stop_token_ids에 제공된 추가 정지 토큰 ID들. generation_config.json 파일 또는 API 요청에서 옵니다.

샘플링된 토큰 ID가 이 병합된 세트에 있으면, 출력에 정지 토큰을 포함하지 않고 루프가 종료됩니다.

max_tokens에 도달하거나 스트리밍 콜백이 false를 반환할 때도 생성이 종료됩니다(서버가 요청을 취소하는 데 사용).

토큰 스트리밍¶

DirectEngine::generate()는 출력 토큰마다 호출되는 on_token 콜백 F: FnMut(i32) -> bool을 받습니다. CLI는 이 콜백으로 생성이 끝나기를 기다리지 않고 토큰을 바로 디코딩해 출력합니다. 서버는 스케줄러가 각 행의 토큰을 마무리하는 대로 해당 HTTP 클라이언트에 스트리밍합니다.

콜백은 계속하려면 true를 반환하고 중단하려면 false를 반환합니다. 클라이언트가 연결을 끊으면 서버가 생성을 중지하는 데 사용됩니다.

4단계: 디토크나이제이션¶

토큰 ID는 프롬프트를 인코딩한 것과 동일한 MlxcelTokenizer를 사용하여 텍스트로 다시 디코딩됩니다. 스트리밍 모드에서 각 개별 토큰은 즉시 디코딩되고 방출됩니다. 일부 토크나이저(특히 SentencePiece 기반)는 단독으로 격리된 하나의 토큰에 대해 유효한 UTF-8 문자열을 생성하지 못할 수 있습니다. 서버는 완전한 코드포인트가 형성될 때까지 부분 바이트를 누적한 후 클라이언트에 응답 조각을 전송합니다.

프로파일링과 통계¶

추론 세션의 generate_with_stats()와 CLI의 generate 경로는 다음 필드를 담은 GenerationStats 구조체를 보고합니다.

필드	설명
`prompt_tokens`	입력 프롬프트의 토큰 수
`generated_tokens`	생성된 토큰 수
`prefill_time_ms`	프롬프트 처리 시간
`decode_time_ms`	디코드 루프 시간
`prefill_tok_per_sec`	프리필 처리량
`decode_tok_per_sec`	디코드 처리량

CLI의 --profile 플래그는 생성 이후 이 통계를 출력합니다.

VLM 변형: 비전-언어 생성¶

비전-언어 모델은 프리필 단계 이전에 실행되는 추가적인 멀티모달 전처리 단계로 표준 텍스트 생성 파이프라인을 확장합니다.

VLM 파이프라인 시퀀스¶

sequenceDiagram
    participant User
    participant CLI/Server
    participant ImagePreproc
    participant VisionEncoder
    participant Projector
    participant EmbedMerger
    participant Generator
    participant TextModel
    participant MLX

    User->>CLI/Server: prompt text + image files
    CLI/Server->>CLI/Server: apply chat template\ninsert image placeholder tokens

    loop for each image
        CLI/Server->>ImagePreproc: load_and_resize(image_path)
        ImagePreproc-->>CLI/Server: pixel_values tensor
    end

    CLI/Server->>VisionEncoder: encode(pixel_values)
    VisionEncoder->>MLX: ViT/SigLIP/CLIP forward pass
    MLX-->>VisionEncoder: visual features [n_patches, hidden_dim]
    VisionEncoder-->>CLI/Server: visual features

    CLI/Server->>Projector: project(visual_features)
    Projector->>MLX: linear/MLP projection
    MLX-->>Projector: projected embeddings [n_tokens, lang_hidden_dim]
    Projector-->>CLI/Server: projected embeddings

    CLI/Server->>EmbedMerger: merge(text_tokens, projected_embeddings, placeholder_positions)
    EmbedMerger->>TextModel: embed_tokens(text_token_ids)
    TextModel-->>EmbedMerger: text embeddings
    EmbedMerger-->>CLI/Server: merged input embeddings + attention mask

    Note over CLI/Server,Generator: VLM PREFILL
    CLI/Server->>Generator: generate_streaming_with_embeddings(model, tokens, embeddings, mask)
    Generator->>TextModel: forward_with_embeddings(tokens, embeddings, caches, mask)
    TextModel->>MLX: transformer forward with merged embeddings
    MLX-->>TextModel: logits
    TextModel-->>Generator: logits

    Note over Generator: DECODE LOOP (identical to text-only)
    Generator->>Generator: sample, stream, repeat...

이미지 전처리¶

VLM 패밀리마다 이미지 전처리 요구 사항이 다릅니다. 이미지 로딩과 크기 조정은 src/vision/processors/에서 처리됩니다. 전처리 파이프라인은 다음을 수행합니다.

image 크레이트를 사용하여 디스크에서 이미지를 로드합니다.
시각 인코더가 예상하는 크기로 이미지를 조정하고 자릅니다.
모델의 preprocessor_config.json에 지정된 평균과 표준 편차를 사용하여 픽셀 값을 정규화합니다.
[n_images, channels, height, width] 형태의 float 텐서를 반환합니다.

일부 VLM은 인코더의 고정 입력 크기보다 큰 이미지를 처리하기 위해 동적 타일링(고해상도 이미지를 여러 타일로 분할)을 지원합니다. Qwen2.5-VL의 윈도우 ViT와 Phi-4의 NaFlex 인코더가 이 접근 방식의 예입니다.

시각 인코딩과 프로젝션¶

시각 인코더는 전처리된 픽셀 텐서를 이미지 백본(일반적으로 Vision Transformer(ViT), SigLIP, 또는 CLIP 모델)을 통해 실행하고 패치 특성 벡터의 시퀀스를 생성합니다.

프로젝션 모듈(아키텍처에 따라 커넥터, 어댑터, 또는 리샘플러라고도 불림)은 이러한 시각적 특성을 인코더의 히든 차원에서 언어 모델의 임베딩 차원으로 매핑합니다. 패밀리마다 다른 프로젝터 설계를 사용합니다.

Linear projection: Qwen2-VL, 대부분의 LLaVA 변형
MLP projection: Gemma3 VLM, Phi-3 Vision
Resampler / cross-attention: MiniCPM-o
Patch merger: Mistral3 VLM, Pixtral

임베딩 병합과 프리필¶

VLM 요청의 채팅 템플릿에는 시각 정보가 삽입되어야 할 위치에 이미지 플레이스홀더 토큰이 포함됩니다. 프로젝션 이후, src/vision/merge.rs의 EmbedMerger는 토큰 임베딩 시퀀스에서 이러한 플레이스홀더 위치를 프로젝션된 시각 임베딩으로 교체합니다.

결과는 total_seq_len = text_tokens + visual_tokens인 [1, total_seq_len, hidden_dim] 형태의 병합된 임베딩 텐서입니다. 이 텐서는 표준 forward() 대신 forward_with_embeddings()에 전달됩니다.

텍스트 토큰과 시각 토큰 사이, 그리고 서로 다른 이미지의 시각 토큰 사이의 다른 위치 관계를 처리하기 위해 어텐션 마스크도 구성됩니다.

after_prefill 훅¶

일부 VLM 아키텍처는 프리필과 디코드 단계 사이에 내부 상태를 수정합니다. 예를 들어, Phi-4 Multimodal은 더 나은 이미지 이해를 위해 프리필 중에 vision LoRA 가중치를 융합한 다음, 베이스 언어 모델 가중치가 사용되도록 디코드 단계에서 이를 분리합니다.

LanguageModel::after_prefill() 훅은 프리필 순전파가 완료된 후, 디코드 루프가 시작되기 전에 한 번 호출됩니다. 기본 구현은 no-op이며 필요한 모델만 이를 재정의합니다.

MLX는 지연 평가를 사용하기 때문에, after_prefill()이 호출되기 전에 프리필 그래프를 명시적으로 평가(ffi::eval()을 통해)해야 합니다. 그래야 가중치 수정이 적용되기 전에 연산 그래프가 결과를 구체화했음이 보장됩니다.