Advanced Inference and Serving¶
Running a model is only the beginning of serving it well. The important production questions are where time is spent, which work can safely be avoided or parallelized, what quality contract must be preserved, and whether the workload is generation, retrieval, or ranking. This chapter develops those ideas and then maps them to mlxcel's actual behavior.
Reading Map¶
| Page | Central question | Current mlxcel status |
|---|---|---|
| Speculative Decoding | How can a target model validate several possible next tokens in one pass without changing the requested decoding contract? | Prompt lookup needs no drafter and runs offline and on the server; a smaller full-model drafter that shares the target's tokenizer runs on the batch-native engine on every surface that takes --draft-model; MTP and DFlash use specialized paths. Temperature 0 output of block verification is not guaranteed to match plain decode (#2778). |
| MTP and DFlash | How do learned multi-token and diffusion drafters differ, and how are target and drafter checkpoints paired? | External Gemma 4 assistant, Qwen 3.5 MTP, and GLM-4.7-Flash MTP drafters, plus Qwen 3.5 DFlash, LFM2 DSpark, Muse Glimmer assistant, and Laguna DFlash drafters, are implemented with documented scope limits. |
| Block Diffusion | What changes when a model generates by iteratively revealing a block instead of appending one token at a time? | DiffusionGemma and LLaDA-2 MoE have dedicated generation and serial-server workers. |
| Runtime Model Surgery | How can weights be transformed reproducibly at load time without rewriting the checkpoint? | YAML-defined scale, add, prune, replace, and interpolate operations are available on supported common-loader paths. |
| Embeddings and Reranking | How do retrieval systems first find candidates cheaply and then score them accurately? | Offline commands, dedicated workers, dense and multi-vector embeddings, multimodal inputs, and cross-encoder or generative rerankers are implemented. |
One Latency Budget, Several Different Optimizations¶
For an interactive service, separate at least four quantities:
- Model load time: time before the process can accept work.
- Time to first token (TTFT): queueing plus prompt preprocessing and prefill.
- Inter-token latency (ITL): time between visible generated tokens after prefill.
- End-to-end latency and throughput: completion time for one request and useful work completed across all requests.
The mechanisms in this chapter attack different terms. Prompt caching and chunked prefill reduce repeated or disruptive prefill work. Continuous batching raises aggregate throughput. Tensor parallelism distributes a single model. Speculative decoding primarily reduces ITL by amortizing target-model verification over multiple accepted tokens. Block diffusion changes the decoding process itself and can revise many positions per iteration. Embedding and reranking models do not generate a response at all; they transform retrieval and relevance scoring into separate serving stages.
An optimization can improve one metric while hurting another. A second drafter consumes memory; a larger speculative block can waste work after an early rejection; a cross-encoder improves relevance while adding a second inference stage; a larger server batch can raise throughput while increasing queueing latency. Always identify the service objective before choosing a mechanism.
A Practical Selection Guide¶
| Workload or constraint | Start with | Why |
|---|---|---|
| One interactive autoregressive request whose reply copies the prompt (edits, rewrites, reformatting) | Prompt lookup | No drafter checkpoint or extra model memory; target verification retains authority. |
| A larger target with a smaller model of the same family and tokenizer | Full-model drafter (--draft-model) | No target-specific training is needed, but the drafter's cost can exceed its gain; measure the real pair first. |
| Gemma 4 target with its dedicated assistant | MTP | The assistant is trained for the target and shares target-side context interfaces. |
| Supported Qwen 3.5-family text target with a published specialized drafter | MTP or DFlash | Both can propose blocks; compare acceptance, drafter cost, and exactness engagement on the actual checkpoint and hardware. |
| GLM-4.7-Flash, LFM2/LFM2.5, Muse Glimmer, or Laguna target | The family's own MTP or DFlash drafter | Each family accepts only its published drafter; see the support matrix in MTP and DFlash. |
| DiffusionGemma or LLaDA-2 checkpoint | Native block diffusion | These are diffusion models, not autoregressive targets to which ordinary speculative decoding should be bolted on. |
| Reproducible ablation, task-vector composition, or temporary weight edits | Runtime Model Surgery | Applies declarative transformations while loading and leaves the checkpoint on disk unchanged. |
| Semantic search over a large corpus | Embedding retrieval | Document vectors can be precomputed and searched approximately. |
| Better ordering of a short retrieved candidate list | Reranking after retrieval | A joint query-document scorer spends more computation only on the shortlisted documents. |
| High-concurrency autoregressive serving | Continuous batching first, then measure speculation | Speculative B=1 wins do not automatically survive queueing and larger batch shapes. |
The Optimization Contract¶
Before enabling any acceleration, write down four things:
- Correctness contract: byte-identical greedy output, preservation of the target sampling distribution, statistically comparable quality, or task-level quality only.
- Measurement unit: TTFT, p50/p95 ITL, request latency, accepted tokens per verification, tokens/s, queries/s, or retrieval quality such as recall and nDCG.
- Resource ceiling: resident memory, memory bandwidth, GPU occupancy, power, and acceptable model-start time.
- Fallback behavior: whether an unsupported pairing must fail startup, decline per request, or transparently return to the baseline path.
This prevents a common mistake: reporting that a method is “faster” because one internal kernel took less time, even though drafter synchronization, verification, rollback, queuing, or a second ranking stage made the full request slower.
Benchmarking Checklist¶
- Warm the Metal kernels and model caches before recording steady-state measurements.
- Keep target checkpoint, quantization, prompt set, output budget, sampler, seed, and penalties fixed.
- Record target-only results before enabling an acceleration path.
- Measure real prompt families rather than a single easy completion; acceptance and diffusion convergence are content-dependent.
- Report resident and peak memory as well as speed.
- For speculative decoding, record proposal width, acceptance rate, mean accepted prefix, decline/fallback counts, and whether an exactness gate engaged.
- For block diffusion, record canvas or block length, denoising steps, early-stop behavior, and visible output tokens rather than counting masked internal positions as completed work.
- For retrieval, evaluate both system latency and task quality. A fast vector search with poor recall cannot be repaired by a reranker that never sees the missing document.
Implementation Boundary¶
This chapter distinguishes three kinds of statements:
- a general inference concept,
- a capability published by an upstream model or research project,
- and a path that the documented mlxcel source tree actually exposes.
Those are not interchangeable. For example, Qwen checkpoints can contain MTP weights in their ecosystem format, but the documented acceleration path uses a supported external qwen3_5_mtp drafter and does not execute an in-checkpoint mtp.layers.* head. Generation models also use internal embedding tensors, but the public search-vector subsystem is a separate implementation with explicit detection, pooling, normalization, and API contracts. Each page calls out such boundaries explicitly.