tekBlueprint

Part 2 covered quantization and KV cache optimization. This is Part 3 of 4 — it focuses on the techniques that accelerate inference: multi-token prediction and distributed parallelism strategies.

Part 1: Architecture and Metrics → Part 2: Quantization and KV Cache ← Part 3: Parallelism and Speed → Part 4: Frameworks and Operations

Multi-Token Prediction (MTP)

Multi-token prediction represents a structural change to how autoregressive models generate text. Instead of producing one token per forward pass, an MTP model predicts multiple tokens ahead in a single pass, then verifies them. When the predictions are correct, you skip intermediate forward passes and generate several tokens at the cost of one.

What Is MTP

The traditional autoregressive pattern follows a fixed cadence: call forward, get logits, sample token, append to sequence, repeat. Each step requires a complete forward pass through the model. For a 70B model, this means streaming 70B+ weights from HBM for every single generated token.

Speculative decoding introduced the idea of drafting tokens with a smaller model and verifying them with the large model. A small draft model generates n tokens independently. The large model then verifies all n tokens in one forward pass, because its logits at each position are independent given the prompt and draft tokens. If the draft tokens match what the large model would have generated, you accept them all and advance by n positions. If a draft token is rejected, you accept the tokens up to that point and fall back to standard decoding.

MTP takes this idea and bakes it into the model architecture. Instead of using a separate draft model, the main model includes additional heads that predict multiple future tokens. These heads are trained jointly with the main generation head, so they learn the same distribution rather than approximating it from a smaller model. The verification step still happens, but the draft and target come from the same parameter space.

The paradigm shift is elimination of the draft quality gap. External speculative decoding suffers when the small draft model makes confidently wrong predictions. The large model must still compute its full forward pass to verify, gaining nothing. With MTP, the drafting heads are part of the trained model, producing higher acceptance rates – typically 60-85% vs 40-60% for external drafters.

How MTP Works

An MTP generation step has two phases. In the draft phase, the model runs a forward pass through the main generation path, producing logits for the next token. Simultaneously, it passes hidden states to a series of MTP heads. Each MTP head predicts one additional token ahead. For a 3-step MTP head, the model produces logits for positions t+1, t+2, and t+3 in a single forward pass.

Draft token generation takes two forms depending on the architecture. In Medusa-style MTP, multiple heads sit on top of the main hidden states and predict future tokens in parallel – there is no sequential dependency between the heads themselves, but acceptance is sequential: a token at offset k is only accepted if all prior offsets are accepted. In Gemma 4-style MTP, tokens are generated sequentially within the MTP drafter: each token is produced autoregressively, with a separate embedder, attending to prior MTP tokens and the main model’s KV cache. Both approaches produce a draft sequence of n candidate tokens.

In the verify phase, the model runs a single forward pass through the main generation path on the accumulated context (original prompt plus all draft tokens). This produces logits for every draft position. The model compares its own logits against the draft tokens at each position. This comparison is the acceptance test.

The accept/reject logic follows a specific pattern. Starting from position t+1, the model checks if the draft token matches the argmax (or a high-probability match) of its own logits. If it matches, the token is accepted and the check continues to t+2. If it does not match, the model resamples from its own logits at that position to produce the correct token. All subsequent draft tokens are discarded.

This means MTP provides a variable acceleration speedup. On easy continuation paths, all draft tokens are accepted and you skip n forward passes. On uncertain paths, only the first token is accepted and MTP provides no speedup beyond the single verification pass. The average speedup depends on the acceptance rate.

A concrete example illustrates the savings. Suppose the MTP head drafts 3 ahead tokens, and the per-position acceptance probability is 70% (assuming independence across positions). Each MTP step consists of two forward passes: one draft pass and one verification pass. The expected tokens generated per MTP step is:

E[tokens per step] = 1 + p + p^2 + p^3
                   = 1.0 + 0.7 + 0.49 + 0.343 ≈ 2.52

The “1” accounts for the mandatory first token from the main generation path; the remaining terms are expected accepted draft tokens. Per forward pass, this yields approximately 1.26 tokens on average. In practice, the speedup is lower – around 1.3-1.5x effective tokens per forward pass – because of KV cache management during verification, the overhead of the MTP drafter, and the fact that draft token acceptance probabilities are not fully independent across positions.

MTP Architecture

The MTP heads are not independent modules. They share the KV cache with the main model, meaning they read from the same key-value storage that the main attention heads write to. Each accepted draft token is embedded and processed through the drafter’s attention layers, which write new entries to the shared KV cache at their respective positions.

A typical MTP implementation uses a separate, smaller Transformer block stacked on top of the main model. This block receives the final hidden state from the main model and processes it through additional Transformer layers. Each layer in the MTP block produces logits for one ahead position through a separate output projection.

The cross-attention mechanism is critical. The MTP block’s attention layers attend to the KV cache shared with the main model. This means MTP tokens can see the full context – they are not generated in isolation. The MTP block also writes to the KV cache as it generates, so deeper MTP heads can attend to earlier MTP predictions.

The design keeps MTP heads smaller than the main model. A common configuration uses the same hidden dimension but fewer attention heads. For example, if the main model uses 64 attention heads, the MTP block might use 16. This reduces the overhead of the MTP computation while maintaining prediction quality.

Memory overhead is manageable. The KV cache grows by n additional token positions during the draft phase, but this memory is freed when draft tokens are rejected. Accepted tokens reuse the draft KV cache entries, so there is no doubling of cache size for accepted tokens.

Gemma 4 Case Study

Gemma 4, released by Google in early May 2026, supports Multi-Token Prediction through dedicated MTP drafters. The MTP drafters were released weeks after the initial model launch as separate checkpoint assets – each target model pairs with a corresponding assistant model (e.g., gemma-4-12B-it with gemma-4-12B-it-assistant).

The MTP drafter for Gemma 4 is a 4-layer Transformer model loaded alongside the main model. It generates ahead tokens sequentially — each token is produced using a separate embedder and attends to the KV cache of prior MTP tokens and the main model. The drafter shares the KV cache with the main model and uses a reduced attention head count relative to the main layers. Commonly configured for 3 ahead tokens, the architecture supports arbitrary draft length. The drafter weights add approximately 10-25% to total parameter loading, depending on the model size.

Google’s benchmarks and the technical report (arXiv:2607.02770, published July 2026) showed up to 3x generation speedup on hardware with strong memory bandwidth. The speedup is most pronounced on H200 and H100 GPUs where memory bandwidth is the primary bottleneck. On memory-constrained hardware, the speedup drops to 1.3-1.8x because loading the additional MTP weights partially offsets the savings from skipped forward passes.

The quality impact was measured at zero across standard benchmarks. Since the MTP heads are trained jointly with the main model, they do not alter the generation distribution for accepted tokens. Rejected tokens are resampled from the main model’s logits, so they follow the same distribution as standard autoregressive generation. The only potential quality impact is on token-level perplexity during the verify pass, which Gemma 4’s design does not introduce.

EAGLE and Medusa: Advanced External Drafting

Beyond simple token-level MTP, two advanced drafting architectures have emerged that push acceptance rates to 70-85% on predictable workloads without significantly increasing model size.

Medusa adds multiple lightweight prediction heads on top of the model’s final hidden states, each head predicting the token at a different future offset (t+1, t+2, t+3, …). These heads share the main model’s representations and are trained jointly with the main model, each optimizing its respective offset prediction through a combined loss.

EAGLE (Extrapolation Algorithm for Greater Language-model Efficiency) takes a more sophisticated approach: instead of predicting future tokens, it predicts future hidden states – the feature vectors that the main model would produce at subsequent layers. A small autoregressive head then maps these predicted features into token logits. Because feature distributions are smoother and more predictable than discrete token distributions, EAGLE achieves higher acceptance rates with a smaller drafting overhead than pure token-level MTP.

Both techniques integrate cleanly with the verify-pass acceptance logic described above, and both are supported by modern serving frameworks like vLLM and SGLang.

MTP vs External Speculative Decoding

External speculative decoding uses a separate small model (often 1/10th the size) as the drafter. MTP bakes the drafter into the model architecture. Both achieve the same end – verifying draft tokens to skip forward passes – but with different trade-offs.

External speculative decoding has one advantage: flexibility. You can combine any draft model with any target model. This is useful when working with older architectures that lack built-in MTP. You can also swap draft models to tune the acceptance rate versus memory overhead trade-off.

MTP has the advantage of a tighter optimization surface. The drafter learns to predict what the main model will produce, not what it would produce if it were a different architecture. The acceptance rate converges to the theoretical maximum for the given distribution, limited only by inherent uncertainty in the next-token distribution.

MTP in Practice

MTP provides predictable speedup when the generation path has low per-token entropy. Code completion, documentation generation, structured output, and translation – these tasks have high next-token predictability. MTP acceptance rates climb to 80-90% on these workloads, delivering 2-3x effective speedup.

The speedup drops for creative writing, open-ended reasoning, and dialogue. These tasks have high token uncertainty, meaning the model’s own logits have flatter distributions. When the main model cannot be confident about the next token, neither can its MTP heads. Acceptance rates drop to 40-50%, and MTP provides marginal or zero acceleration.

Batch size considerations are critical. MTP and speculative decoding are most effective at small batch sizes (BS=1-4), where inference is memory-bound and the GPU compute units are underutilized. The MTP heads fill the idle compute cycles with useful work. At batch sizes above 8-16, inference becomes compute-bound, and the MTP overhead competes with the main model for Tensor Cores, potentially reducing throughput. For high-throughput workloads with large batches, consider disabling MTP entirely and relying on other optimizations.

MTP and other optimization techniques stack. You can combine MTP with quantization, continuous batching, and PagedAttention. The combined effect is multiplicative on throughput. A quantized model with MTP and PagedAttention can achieve 3-5x effective throughput compared to the unquantized autoregressive baseline. The key constraint remains VRAM: all optimizations must fit within the available memory budget.

Parallelism Strategies

A single GPU rarely holds a modern large model at useful batch sizes. Serving a 70B parameter model at batch 32 with fp16 weights and activations demands tens of gigabytes of VRAM and exceeds what even an A100 80G can provide without some form of distributed computation. Parallelism strategies solve this problem by splitting model or work across multiple devices. No single strategy is universally optimal, and production systems combine several strategies to match their hardware topology and workload profile.

Tensor Parallelism (TP)

Tensor parallelism splits the weight matrices themselves across GPUs within a single node. The key insight is that many linear layers in a transformer – especially the large feed-forward projection and attention output projection – can be partitioned without fundamentally changing the computation graph. This approach belongs to one of the oldest and most mature distributed training techniques, dating back to the Megatron paper.

The basic partitioning scheme uses two complementary operations. Column parallel splits a weight matrix W along its output dimensions. If W has shape (d_model, d_ff), GPU 0 computes W[:, :d_ff//2] @ x and GPU 1 computes W[:, d_ff//2:] @ x, each producing half the output vector. Row parallel splits along the input dimensions. GPU 0 computes W[:d_model//2, :] @ y, GPU 1 computes W[d_model//2:, :] @ y, and their partial results must be summed.

These two operations naturally appear together inside a transformer layer. The query, key, and value projections use column parallelism: each GPU produces one-quarter of the outputs for each Q, K, and V projection respectively. After the attention mechanism completes on a per-head basis, the output projection uses row parallelism to recombine the split heads back into the full d_model representation.

Both column and row parallelism require a post-layer all-reduce to synchronize partial results across GPUs. Every layer performs this collective operation, meaning the communication volume scales with the number of layers. For a 70B model with 80 layers, that is 80 all-reduce operations per forward pass. This communication overhead is the primary constraint on TP.

The TP degree equals the number of GPUs per model replica. A TP degree of 4 means four GPUs jointly compute one copy of the model. Increasing the total throughput requires multiple independent replicas, each with its own TP group. This separation matters for deployment: TP degree is bounded by the GPU count within a single node because inter-node communication is an order of magnitude slower than intra-node NVLink.

In practice, TP works best when the goal is single-request latency reduction for a large model. Splitting a 70B model across 8 GPUs with NVLink often yields near-linear FLOP scaling, bringing time-to-first-token down to interactive levels. The communication overhead remains bounded to the node fabric, keeping per-layer synchronization in the microsecond to low-millisecond range. The trade-off is that TP degree cannot grow beyond the physical GPU count per node.

Pipeline Parallelism (PP)

Pipeline parallelism takes a different axis: instead of splitting tensors, it splits the model’s layer sequence across GPUs. If a model has 40 layers and two GPUs are available, GPU 0 handles layers 0-19 and GPU 1 handles layers 20-39. Each GPU runs a complete forward computation for its slice before passing intermediate activations to the next stage.

This approach works across nodes as well as within a single node. The intermediate activations between stages travel over PCIe or standard network links, and their bandwidth requirements are modest compared to TP’s all-reduce traffic. The communication per step is point-to-point: each stage sends one tensor to its neighbor and receives one tensor from its predecessor.

The main inefficiency in pipeline parallelism is the pipeline bubble. During a forward pass with G stages, the first stage finishes its work and passes data to stage 2, while stage 2 is idle until it receives that data. If a single batch processes through the pipeline, the last stage begins working only after the first stage finishes, leaving (G-1)/G of the total GPU-time idle.

Micro-batching solves this inefficiency. Instead of sending one large batch through the pipeline, the batch splits into M smaller micro-batches. While stage 1 processes micro-batch 2, stage 0 feeds micro-batch 1 forward. Piping micro-batches through the stages keeps all GPUs busy, though residual bubble waste persists. The 1F1B scheduling technique (primarily used in training) optimizes further: after a warmup phase fills the pipeline with forward micro-batches, it interleaves one backward pass step with one forward step. This backpressure strategy minimizes peak memory for gradients while keeping GPU utilization near-optimal. For inference-only deployments, PP with micro-batching typically uses a simpler forward-only interleaved pipeline schedule.

Interleaved pipeline parallelism adds another degree of freedom. Instead of assigning one contiguous chunk of layers per GPU, it assigns multiple non-contiguous chunks in a round-robin order. GPU 0 handles layers [0, k, 2k, …], GPU 1 handles [k, 2k, 3k, …]. This interleaving reduces the bubble fraction because each GPU participates more frequently in each pipeline step, improving both memory balancing and compute utilization.

PP is the right choice for very deep models or multi-node deployments where NVLink bandwidth is insufficient. Llama 70B with PP across 4 nodes still runs without constant all-reduce chatter between racks. The scheduling complexity is the penalty: PP requires careful micro-batch sizing, memory tracking for activations and KV cache state, and often a custom communication library.

Expert Parallelism (EP)

Expert parallelism exists only for Mixture-of-Experts (MoE) models. In an MoE architecture, the feed-forward network inside each transformer layer contains K independent “expert” sub-networks, and a routing gate selects which experts process each token. At inference time, a token typically activates 2-4 out of 64 or 128 experts.

EP assigns each GPU a subset of the total expert pool. If there are 64 experts across 8 GPUs, each GPU holds 8 experts. During forward computation, tokens route to whichever GPU owns their assigned experts. This routing creates an all-to-all communication pattern: every GPU sends tokens to every other GPU and receives tokens from every other GPU.

All-to-all communication is significantly more expensive than all-reduce or point-to-point transfers. The communication volume depends on both the number of experts and the token distribution across them. If routing is perfectly uniform, each GPU sends roughly 1/K of its tokens to each peer GPU. Under skewed routing, hotspots emerge and some GPUs process far more tokens than others, creating stragglers.

EP almost always combines with TP. The typical pattern is EP across nodes – distributing expert subsets to different nodes – and TP within each node – splitting the individual expert weights across GPUs on shared NVLink. This two-level decomposition balances communication cost against memory requirements and is the strategy used by DeepSeek V3 and Mixtral serving stacks.

Data and Sequence Parallelism (DP/SP)

Data parallelism is the simplest and most widely understood strategy. Each GPU holds a complete model replica, and the input batch divides across GPUs. GPU 0 processes tokens from batch slice 0, GPU 1 processes slice 1, and so on. During training, gradients synchronize across replicas via all-reduce. During inference, each replica operates independently without inter-replica communication. Data parallelism scales throughput linearly with GPU count and requires no model splitting.

The limitation of plain data parallelism is that each GPU still needs enough VRAM for one full model replica. Models that exceed single-GPU memory cannot use DP alone, regardless of batch size.

Sequence parallelism (SP) addresses memory bottlenecks in non-attention layers (like LayerNorm and Dropout) by splitting activations along the sequence dimension, reducing peak memory usage during forward and backward passes. This requires all-gather and reduce-scatter operations to maintain consistency across GPUs.

Context parallelism (CP), on the other hand, is designed specifically for long context windows. For a 128K context, the attention computation and KV cache footprint grow quadratically or linearly with sequence length. CP splits the attention matrix and the KV cache along the token dimension. GPU 0 computes attention for tokens 0-64K, GPU 1 for tokens 64K-128K. During attention, GPUs exchange their local KV caches via collective communication (like Ring Attention). CP is particularly valuable for streaming long documents or multi-turn conversations where context extends beyond a single GPU’s VRAM budget.

Disaggregated Serving: Separating Prefill from Decode

The most significant architectural shift in LLM serving over the past few years is disaggregation – the separation of prefill and decode phases onto different GPU pools. Systems like DistServe, Splitwise, and Mooncake pioneered this approach, and it has become the standard for large-scale production deployments.

The core insight is that prefill and decode have fundamentally different resource profiles and interfere with each other when colocated. A single long-context prefill (compute-bound) sharing a GPU with multiple decode sequences (memory-bound) causes “bubble” effects: the decode sequences experience ITL spikes while waiting for the prefill’s large matmuls to complete, and the prefill underutilizes memory bandwidth while the decode sequences compete for it.

Disaggregation solves this by dedicating separate GPU pools to each phase. A prefill pool is optimized for compute throughput: large batch sizes of prompt tokens, aggressive matrix fusion, and Tensor Core saturation. A decode pool is optimized for memory bandwidth: small batches, KV cache locality, and continuous batching. When a new request arrives, it is processed by the prefill pool, which computes the initial hidden states and KV cache, then transfers them over high-speed interconnect (NVLink, InfiniBand, or RoCE) to the decode pool for token generation.

The benefits are substantial. Disaggregated systems achieve 2-4x better hardware utilization than colocated systems, because each pool operates at its optimal regime without interference. TTFT typically improves for long-context requests where prefill compute dominates queue time, though the added KV transfer overhead can offset gains for short prompts. ITL improves because decode GPUs never stall on long prefills. The trade-off is additional network overhead for transferring KV cache state, which becomes significant at very long context lengths (64K+ tokens).

Modern serving frameworks like vLLM and TensorRT-LLM now support disaggregated configurations natively, and purpose-built systems like Mooncake (used by Moonshot AI) have demonstrated its viability at planetary scale.

3D Parallelism

Production serving systems rarely rely on a single parallelism strategy. 3D parallelism combines TP, PP, and DP to optimize compute utilization, communication overhead, and deployment flexibility simultaneously.

The typical decomposition works as follows. TP splits tensor operations within each node, taking advantage of NVLink bandwidth to minimize intra-node communication latency. PP chains node groups across the rack, distributing model depth to avoid memory concentration. DP replicates the combined TP+PP pipeline across GPU clusters, achieving throughput scaling at the workload level.

Megatron-LM popularized the TP+DP combination for training and adapted it for serving. DeepSpeed added ZeRO optimizer states and PP for extreme-scale models. On the inference side, vLLM relies on TP as its primary parallelism primitive, deploying one TP-grouped replica per serving instance and scaling horizontally with DP. TensorRT-LLM supports both TP and PP, enabling compilation-aware placement of parallel stages for maximal throughput across single and multi-node deployments.

Choosing the right 3D configuration depends on three factors: model size relative to per-GPU memory, desired latency versus throughput profile, and the available hardware topology.

Parallelism Comparison Table

Glossary

References

1. Parallelism — Tensor, Pipeline, and Data
   1. Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Tensor + Pipeline Parallelism). arXiv:1909.08053
   2. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models (Data Parallelism + ZeRO). arXiv:1910.02054
2. Parallelism — Sequence and Context
   1. Ring Attention with Blockwise Transformers for Near-Infinite Context (Context Parallelism). arXiv:2310.01889
3. Disaggregated Serving
   1. DistServe: Disaggregating Prefill and Decoding for Goodput-optimized LLM Serving. arXiv:2401.09670
   2. Splitwise: Efficient Generative LLM Inference Using Phase Splitting. arXiv:2311.18677
4. Multi-Token Prediction (MTP) and Speculative Decoding
   1. Fast Inference from Transformers via Speculative Decoding. arXiv:2211.17192
   2. Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. arXiv:2401.10774
   3. EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. arXiv:2401.15077
   4. Gemma 4 Technical Report (MTP heads, multimodal architecture). arXiv:2607.02770
5. Models and Architecture References
   1. DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (MLA + EP). arXiv:2405.04434
   2. DeepSeek-V3 Technical Report (MoE with 256 experts, expert parallelism at scale). arXiv:2412.19437

Previous: Part 2 — Quantization and KV Cache

Next: Part 4 — Frameworks and Operations covers serving frameworks, production operations, and a decision matrix.