Inference Pipeline¶

Table of Contents¶

• How Text Generation Works
• Full Pipeline Sequence Diagram
• Step 1: Tokenization
• Tokenizer Backends
• Chat Template Application
• Step 2: Prefill
• The Forward Pass
• The Prefill Plan
• MLX Computation Graph Execution
• Step 3: The Decode Loop
• Lookahead Pipelining
• EOS Detection and Stopping
• Token Streaming
• Step 4: Detokenization
• Profiling and Statistics
• VLM Variant: Vision-Language Generation
• VLM Pipeline Sequence
• Image Preprocessing
• Visual Encoding and Projection
• Embedding Merge and Prefill
• The after_prefill Hook

How Text Generation Works¶

Large language models generate text by predicting one token at a time. At each step, the model receives the entire history of tokens seen so far (the prompt plus all previously generated tokens), runs a forward pass through its neural network layers, and produces a probability distribution over its vocabulary. A sampling strategy then selects the next token from that distribution.

This is called autoregressive generation: each generated token becomes part of the input for the next step.

There are two distinct steps in this process:

• Prefill: The model processes the entire input prompt in one forward pass. Because the prompt is known in advance, all tokens can be processed in parallel. This is the most computationally intensive step.
• Decode: The model processes one newly generated token per step, using the KV cache to avoid recomputing attention over the entire history. This is memory-bandwidth intensive rather than compute intensive.

mlxcel's decode engine in src/lib/mlxcel-core/src/engine/ implements this process for both the CLI and the server, with lookahead pipelining to keep the GPU busy between decode steps.

Full Pipeline Sequence Diagram¶

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

Step 1: Tokenization¶

Tokenizer Backends¶

mlxcel supports three tokenizer formats through the MlxcelTokenizer enum in src/tokenizer/mod.rs:

The load_tokenizer() function detects which format is available in the model directory and constructs the appropriate variant. The MlxcelTokenizer enum provides a uniform encode() and decode() interface regardless of the underlying format.

The SentencePiece backend includes special token handling: special tokens are sorted by length descending and matched greedily before passing remaining text to the underlying SentencePiece processor. This correctly handles models that add special tokens like <|im_start|> or <|begin_of_text|> to a base SentencePiece vocabulary.

Chat Template Application¶

Most modern models are fine-tuned with a specific conversation format. Sending a bare prompt string may give poor results because the model expects its training format.

mlxcel applies the model's chat template automatically using minijinja, a Rust implementation of the Jinja2 templating engine. The chat template is stored in the model's tokenizer_config.json file as a Jinja2 template string.

minijinja differs from Python Jinja2 in a few cases that a custom template can reach: a list or dict rendered directly uses Python quoting, // and % floor, round rounds half to even, and division by zero is an error.

The template receives a list of messages (with role and content fields) and produces the formatted string that the tokenizer will encode. Different models have different formats:

• Llama 3 uses <|begin_of_text|><|start_header_id|>user<|end_header_id|>...
• Qwen uses <|im_start|>user\n...<|im_end|>
• Gemma uses <start_of_turn>user\n...<end_of_turn>

Chat template application can be disabled with --no-chat-template in the CLI, or overridden with --chat-template or --chat-template-file in the server.

Step 2: Prefill¶

The Forward Pass¶

The LanguageModel::forward() method is the single entry point into the model's computation. During prefill, it receives:

• input_ids: A [1, seq_len] tensor of token IDs for the entire prompt.
• caches: A mutable slice of KVCache structs, one per transformer layer.
• mask: None during prefill (the model generates its own causal mask internally).

The model processes all prompt tokens in parallel, writes key-value projections for every token into the caches, and returns a logits tensor of shape [1, seq_len, vocab_size]. Only the logits at the last sequence position are used to sample the first generated token.

The Prefill Plan¶

The CLI and the server partition a prompt the same way through one PrefillPlan (src/lib/mlxcel-core/src/prefill_plan.rs):

1. A prefix adopted from the prompt cache is never forwarded again.
2. For a chat prompt, the segment up to the history boundary (the rendering without its generation prompt) is one unpadded piece, so the prompt cache can store a snapshot that describes exactly those tokens.
3. The rest is cut into pieces of at most the prefill chunk size when the model supports a multi-call prefill and the input is tokens rather than pre-merged embeddings.
4. On hardware with prefill tile alignment, a piece is padded to the Neural Accelerator tile when the model supports padded prefill; a piece that ends at the history boundary is never padded.

mlxcel generate uses a 2048-token chunk. The server uses 2048 while no other sequence is decoding and 512 while a decode batch is live; an explicit --prefill-chunk-size or MLXCEL_PREFILL_CHUNK pins both.

MLX Computation Graph Execution¶

MLX uses lazy evaluation. When Rust code calls MLX operations (through the ffi module), those operations are recorded into a computation graph but not immediately executed on the GPU. Execution happens when ffi::eval() is called to materialize a specific tensor.

This allows MLX to optimize the computation graph (fusing operations, selecting efficient kernels) before scheduling GPU work. The generation loop uses ffi::async_eval() to schedule computation without blocking the Rust thread, then syncs with ffi::item_i32() when the token value is actually needed.

Checkpoint weights are the exception to lazy evaluation: model load evaluates and uploads them (see Model Loading), so the first prefill does not wait on weight I/O. Vision and audio tower weights stay lazy until a request carries that modality.

After prefill, ffi::clear_memory_cache() releases intermediate tensors that are no longer needed, reclaiming GPU memory before the decode loop begins.

Step 3: The Decode Loop¶

Lookahead Pipelining¶

Every decode step runs through the engine's split step (Engine::submit and Engine::finish_rows). DirectEngine::generate(), the one-sequence loop mlxcel generate decodes on, and the server's BatchScheduler both drive it as a lookahead pipeline that overlaps GPU computation with CPU-side bookkeeping. At each iteration:

1. Start computing the next token's forward pass and sampling (GPU work submitted asynchronously).
2. Materialize the current token from the previous iteration (sync with GPU).
3. Check the current token against EOS token IDs.
4. Call the streaming callback with the current token ID.
5. Move "next" to "current" for the next iteration.

The GPU never goes idle waiting for the CPU: the cost of EOS checking, callback dispatch, and history maintenance is hidden behind GPU compute.

A sampler that the fused draw cannot run (history penalties, DRY, mirostat, adaptive-p) keeps the forward ahead on a per-row pipeline instead; mirostat and adaptive-p read their draw on the host, so only their finish step overlaps the next forward. Setting MLXCEL_FORCE_SYNC disables the pipeline in both the CLI and the server, and every step then runs the synchronous Engine::step. This is useful for debugging.

EOS Detection and Stopping¶

The generation loop merges two sources of EOS token IDs:

1. The model's built-in EOS tokens from LanguageModel::eos_token_ids(), which are set by the model implementation based on the vocabulary.
2. Any additional stop token IDs provided in SamplingConfig::stop_token_ids, which come from the generation_config.json file or the API request.

When the sampled token ID is in this merged set, the loop terminates without including the stop token in the output.

Generation also terminates when max_tokens is reached, or when the streaming callback returns false (used by the server to cancel a request).

Token Streaming¶

DirectEngine::generate() takes an on_token callback F: FnMut(i32) -> bool that is called for each emitted token. The CLI uses this to decode and print each token immediately rather than waiting for the complete generation. The server streams each row's tokens to its HTTP client as the scheduler finishes them.

The callback returns true to continue or false to abort. This allows the server to stop generation if the client disconnects.

Step 4: Detokenization¶

Token IDs are decoded back to text using the same MlxcelTokenizer that encoded the prompt. In streaming mode, each individual token is decoded and emitted immediately. Some tokenizers (especially SentencePiece-based ones) may not produce a valid UTF-8 string for a single token in isolation; the server accumulates partial bytes until a complete codepoint is formed before sending the response fragment to the client.

Profiling and Statistics¶

The inference session's generate_with_stats() and the CLI's generate path report a GenerationStats struct containing:

The CLI's --profile flag prints these statistics after generation.

VLM Variant: Vision-Language Generation¶

Vision-language models extend the standard text generation pipeline with an additional multimodal preprocessing step that runs before the prefill step.

VLM Pipeline Sequence¶

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...

Image Preprocessing¶

Each VLM family has specific image preprocessing requirements. Image loading and resizing is handled in src/vision/processors/. The preprocessing pipeline:

1. Loads the image from disk using the image crate.
2. Resizes and crops the image to the dimensions expected by the visual encoder.
3. Normalizes pixel values using the mean and standard deviation specified in the model's preprocessor_config.json.
4. Returns a float tensor of shape [n_images, channels, height, width].

Some VLMs support dynamic tiling (splitting a high-resolution image into multiple tiles) to handle images larger than the encoder's fixed input size. Qwen2.5-VL's windowed ViT and Phi-4's NaFlex encoder are examples of this approach.

Visual Encoding and Projection¶

The visual encoder runs the preprocessed pixel tensor through an image backbone, typically a Vision Transformer (ViT), SigLIP, or CLIP model, and produces a sequence of patch feature vectors.

A projection module (sometimes called the connector, adapter, or resampler depending on the architecture) maps these visual features from the encoder's hidden dimension into the language model's embedding dimension. Different families use different projector designs:

• Linear projection: Qwen2-VL, most LLaVA variants
• MLP projection: Gemma3 VLM, Phi-3 Vision
• Resampler / cross-attention: MiniCPM-o
• Patch merger: Mistral3 VLM, Pixtral

Embedding Merge and Prefill¶

The chat template for a VLM request contains image placeholder tokens at the positions where visual information should be inserted. After projection, the EmbedMerger (in src/vision/merge.rs) replaces these placeholder positions in the token embedding sequence with the projected visual embeddings.

The result is a merged embedding tensor of shape [1, total_seq_len, hidden_dim] where total_seq_len = text_tokens + visual_tokens. This tensor is passed to forward_with_embeddings() instead of the standard forward().

An attention mask is also constructed to handle the different positional relationships between text tokens and visual tokens, and between visual tokens from different images.

The after_prefill Hook¶

Some VLM architectures modify their internal state between the prefill and decode steps. For example, Phi-4 Multimodal fuses vision LoRA weights during prefill for better image understanding, then unfuses them for the decode step so the base language model weights are used.

The LanguageModel::after_prefill() hook is called once after the prefill forward pass completes and before the decode loop begins. The default implementation is a no-op; only models that need it override it.

Because MLX uses lazy evaluation, the prefill graph must be explicitly evaluated (via ffi::eval()) before after_prefill() is called. This ensures that the computation graph has materialized its results before any weight modifications take effect.