Tag: llm architecture

  • Multi-Head Latent Attention: Scaling LLM Serving Throughput via Low-Rank KV Compression

    Multi-Head Latent Attention: Scaling LLM Serving Throughput via Low-Rank KV Compression

    TL;DR

    • Multi-Head Latent Attention (MLA) overcomes serving memory bottlenecks by compressing Key-Value states into compact latent vectors via low-rank projections rather than relying on coarse head-sharing.

    • During inference, MLA absorbs key and value up-projection matrices into query and output projections; it pairs this with Decoupled RoPE so that position embeddings do not interfere with matrix absorption.

    • In DeepSeek-V2, MLA enables a 93.3% reduction in KV cache footprint compared to DeepSeek 67B. Together with sparse MoE computation (activating 21B out of 236B parameters), this contributes to a 5.76-fold boost in maximum generation throughput, while separate architectural ablations show MLA outperforms standard Multi-Head Attention.

    The Memory Wall: Why the KV Cache Constrains LLM Serving

    During Autoregressive Generation in Large Language Models (LLMs), Inference Throughput is fundamentally bottlenecked by the Key-Value (KV) Cache. In standard Multi-Head Attention (MHA), the model must retain the Key and Value tensors of all previously generated tokens across every layer. These cached states reside in High-Bandwidth Memory (HBM) and must be repeatedly transferred into on-chip memory at every decoding step to compute attention for each new token. This persistent, high-volume memory traffic creates severe HBM bandwidth contention, capping decoding throughput and strictly limiting both the serving batch size and the practical context window length.

    Prior attempts to alleviate this memory wall rely primarily on Multi-Query Attention (MQA) and Grouped-Query Attention (GQA). Both techniques reduce memory requirements by forcing multiple query heads to share a reduced number of KV heads. While this head-sharing structure lowers the absolute volume of data stored in HBM and transferred to on-chip memory, it introduces an architectural trade-off: sharing KV projections across disparate query heads compromises per-head representational diversity and impairs retrieval expressiveness.

    Consequently, LLM serving faces a structural tension between memory bandwidth efficiency and model capacity. Standard MHA preserves distinct per-head projections at the cost of high HBM bandwidth consumption, whereas MQA and GQA reduce cache overhead at the expense of representational richness. Overcoming this bottleneck requires moving beyond coarse head-sharing, addressing the memory wall via Low-Rank Joint Compression of the stored representations into a Compressed KV Latent Vector rather than sacrificing individual per-head projection capacity.

    Sources: [2405.04434] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model, arxiv.org/pdf/2405.04434, vertexaisearch.cloud.google.com/grounding-api-redirect/AUZIYQF73W_INb…, Not Acceptable!, Paper page – TransMLA: Multi-head Latent Attention Is All You Need

    Empirical Gains: KV Footprint Reduction and Serving Speedups in Practice

    The practical efficacy of this low-rank compression strategy is clearly demonstrated in modern architectures. Compared to the dense DeepSeek 67B baseline, DeepSeek-V2 reduces the Key-Value (KV) Cache footprint by 93.3% and boosts maximum generation throughput up to 5.76 times. This throughput gain is a system-wide result: the KV cache reduction stems directly from MLA’s low-rank compression, while overall serving efficiency is further multiplied by the DeepSeekMoE architecture, which activates only 21B parameters out of 236B total parameters per token.

    Standard Multi-Head Attention (MHA) caches all uncompressed keys and values across every attention head, requiring 2 * n_h * d_h * l elements per token, where n_h represents the head count, d_h is the per-head dimension, and l denotes the layer count. Alternative mechanisms like Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) scale memory down to 2 * n_g * d_h * l elements (for n_g groups) and 2 * d_h * l elements, respectively, but sacrifice modeling capability.

    In contrast, MLA projects keys and values into a single low-rank latent vector c_t^KV of dimension d_c. However, standard Rotary Position Embedding (RoPE) is position-sensitive and couples directly with keys and queries. If RoPE were applied directly to the content keys, the key up-projection matrix W^UK could not be absorbed into the query projection during inference because a token-dependent rotation matrix would sit between them, destroying associativity and forcing the system to recompute prefix keys. MLA resolves this by introducing Decoupled RoPE: the content keys k_t^C remain unrotated (enabling matrix absorption), while a separate, low-dimensional shared key k_t^R of dimension d_h^R carries the positional embeddings.

    During inference, MLA caches only the compressed latent vector c_t^KV alongside the decoupled key k_t^R, requiring (d_c + d_h^R) * l elements per token. In the DeepSeek-V2 configuration (n_h = 128, d_h = 128, d_c = 512 = 4 * d_h, and d_h^R = 64 = d_h / 2), the cache requirement is (d_c + d_h^R) * l = 4.5 * d_h * l elements. In terms of per-token cache storage footprint, this is equivalent to GQA configured with only 2.25 groups. In architectural evaluations, MLA nevertheless achieves stronger performance than full MHA.

    Furthermore, during inference, the key up-projection matrix W^UK is absorbed directly into the query projection W^Q, and the value up-projection matrix W^UV is absorbed into the output projection W^O. As a result, the full per-head keys and values never need to be materialized in memory for attention computation, greatly reducing HBM bandwidth traffic.

    Sources: DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model

    Closing thoughts

    These architectural efficiencies highlight a practical approach to addressing memory bottlenecks in LLM serving. Rather than relying entirely on coarse head-sharing mechanisms like GQA or MQA that compromise representational diversity, Multi-Head Latent Attention demonstrates that memory bandwidth contention can be significantly mitigated through low-rank joint compression.

    A key design element is decoupling positional embeddings from content projections. Because standard RoPE operations are position-dependent and prevent key up-projection matrices from being absorbed into queries, Decoupled RoPE isolates the positional stream. This allows the content stream to remain purely linear, enabling inference engines to absorb W^UK into W^Q and W^UV into W^O without materializing uncompressed per-head Key-Value states in memory.

    MLA does not eliminate all architectural trade-offs: practitioners must still select appropriate compression dimensions d_c, allocate cache for decoupled positional keys, and navigate specialized CUDA kernel implementations. Nevertheless, its empirical impact in DeepSeek-V2—where MLA’s 93.3% KV cache reduction relative to DeepSeek 67B combines with sparse MoE execution to deliver up to a 5.76-fold throughput improvement, alongside ablation results outperforming standard MHA—makes latent compression a compelling pattern for high-throughput serving systems.

    Frequently Asked Questions

    Why does the standard Key-Value (KV) cache create a bottleneck in LLM serving?

    In standard Multi-Head Attention, retaining uncompressed Key and Value tensors across all heads and layers requires repeatedly transferring massive volumes of data from High-Bandwidth Memory to on-chip memory at each decoding step. This creates severe memory bandwidth contention that restricts decoding throughput, serving batch sizes, and context window lengths.

    How does Multi-Head Latent Attention (MLA) differ from MQA and GQA?

    MQA and GQA reduce memory by forcing multiple query heads to share a reduced count of KV heads, which compromises representational diversity and retrieval expressiveness. In contrast, MLA compresses stored representations into a compact latent vector via low-rank joint compression without relying on coarse head-sharing.

    What does MLA store in the KV cache per token?

    MLA caches only the Compressed KV Latent Vector alongside a decoupled key carrying Decoupled Rotary Position Embedding (Decoupled RoPE). This requires (d_c + d_h^R) * l elements per token rather than materializing full uncompressed keys and values across every head.

    How does MLA optimize computation during inference?

    During inference, the key up-projection matrix is absorbed directly into the query projection, and the value up-projection matrix is absorbed into the output projection. Because Decoupled RoPE keeps the content stream position-free, this linear absorption avoids having to compute or materialize full per-head Key and Value tensors in memory during attention.

    What empirical throughput and memory gains were achieved by DeepSeek-V2 using MLA?

    Compared to DeepSeek 67B, DeepSeek-V2 reduced the KV cache footprint by 93.3% and increased maximum generation throughput up to 5.76 times. The KV cache reduction is driven by MLA, while the throughput gains reflect both the compressed cache and the sparse DeepSeekMoE architecture activating 21B out of 236B parameters. In terms of cache memory footprint per token, MLA is equivalent to GQA with just 2.25 groups, while separate paper ablations show it delivers stronger modeling performance than full Multi-Head Attention.


    Keep reading: LLM Inference & Serving

    All guides in this series: LLM Inference & Serving Guides: vLLM, Ollama & Hugging Face

  • Byte Latent Transformer: Replacing Tokenizers with Dynamic Entropy-Based Patching

    Byte Latent Transformer: Replacing Tokenizers with Dynamic Entropy-Based Patching

    TL;DR

    • The Byte Latent Transformer introduces a tokenizer-free architecture that overcomes the noise brittleness and vocabulary biases of traditional subword tokenization.
    • It processes raw byte streams efficiently by dynamically segmenting bytes into patches based on next-byte entropy.
    • BLT matches or exceeds the performance of traditional tokenizer-based models like LLaMA at scales up to 8 billion parameters.

    Introduction: The Shift from Fixed Tokenizers to Byte-Level Modeling

    Traditional large language models depend heavily on heuristic Subword Tokenization algorithms, such as Byte-Pair Encoding (BPE), to segment text into static subword vocabularies. While widely adopted, fixed Subword Tokenization introduces fundamental structural vulnerabilities:

    • Brittleness to Noise and Manipulation: Models exhibit high sensitivity to typos, spelling errors, and character-level perturbations, while struggling on character-level manipulation and arithmetic tasks.
    • Vocabulary and Representation Issues: Fixed vocabularies lead to out-of-vocabulary fragmentation and create structural biases favoring English and high-resource scripts over low-resource languages.

    Operating directly on raw byte streams eliminates the need for tokenizers, but naive Byte-Level Language Model design introduces severe computational bottlenecks. Expanding text directly into raw byte sequences dramatically inflates sequence lengths, imposing extreme quadratic computational overhead on standard Transformer self-attention mechanisms.

    To overcome both the rigidity of subword tokenizers and the computational expense of naive byte sequences, researchers from Meta FAIR, the University of Washington, and the University of Chicago (Artidoro Pagnoni, Ram Pasunuru, Pedro Rodriguez, John Nguyen, Benjamin Muller, Margaret Li, Chunting Zhou, Lili Yu, Jason Weston, Luke Zettlemoyer, Gargi Ghosh, Mike Lewis, Ari Holtzman, and Srinivasan Iyer) introduced the Byte Latent Transformer (BLT). BLT represents the first compute- and FLOP-controlled scaling study demonstrating that a tokenizer-free Byte-Level Language Model can match or exceed the performance of standard tokenizer-based models, such as LLaMA baselines, across scales up to 8 billion parameters and trillions of training bytes.

    Sources: youtube.com/watch, Byte Latent Transformer: Patches Scale Better Than Tokens, [2412.09871] Byte Latent Transformer: Patches Scale Better Than Tokens, OpenReview (paper discussion page), Byte Latent Transformer: Patches Scale Better Than Tokens – ACL Anthology

    Mechanism: Dynamically Segmenting Patches via Next-Byte Entropy

    At the core of this breakthrough is the Byte Latent Transformer (BLT) architecture’s approach to sequence structure: raw bytes are segmented into dynamically sized patches through Dynamic Entropy-Based Patching, serving as the primary units of computation and bypassing the need for Subword Tokenization. Patch segmentation is driven by Next-Byte Entropy predicted by an Entropy Model, which enables the model to dynamically allocate more compute and model capacity where increased data complexity demands it.

    When the underlying data is predictable, Dynamic Entropy-Based Patching constructs longer patches. This entropy-based segmentation enhances both training and inference efficiency, allowing BLT to simultaneously scale model size and patch length under fixed inference costs.

    Sources: Byte Latent Transformer: Patches Scale Better Than Tokens

    Closing thoughts

    Ultimately, by tying sequence segmentation directly to informational entropy rather than static vocabulary heuristics, the Byte Latent Transformer successfully resolves the long-standing computational bottleneck of naive byte-level processing. In my view, the fact that BLT matches or exceeds standard baselines like LLaMA at scales up to 8 billion parameters demonstrates that the structural brittleness and vocabulary biases of BPE are no longer necessary compromises for scalable architectures. Dynamically scaling patch lengths according to data predictability proves to be a well-reasoned solution, allowing byte-level modeling to achieve character-level robustness while strictly maintaining compute and inference efficiency.

    Frequently Asked Questions

    What are the main drawbacks of traditional subword tokenization methods like BPE?

    Fixed subword tokenization introduces brittleness to noise, typos, and character-level perturbations, while causing out-of-vocabulary fragmentation. It also struggles with character manipulation tasks and creates structural biases favoring English and high-resource scripts over low-resource languages.

    Why is naive byte-level modeling computationally inefficient?

    Directly expanding text into raw byte sequences dramatically inflates sequence lengths. This imposes extreme quadratic computational overhead on standard Transformer self-attention mechanisms.

    How does the Byte Latent Transformer (BLT) process byte sequences efficiently?

    BLT uses Dynamic Entropy-Based Patching driven by next-byte entropy predicted by an Entropy Model. It constructs longer patches when data is predictable, dynamically allocating more compute and capacity where data complexity demands it.

    How does the performance of BLT compare to standard tokenizer-based models?

    BLT matches or exceeds the performance of standard tokenizer-based models, such as LLaMA baselines, across scales up to 8 billion parameters and trillions of training bytes.


    Keep reading: Training & Architectures

    All guides in this series: LLM Training & Architecture Guides: Fine-Tuning, RL & New Models

  • Demystifying Mixture of Experts Routing: Core Mechanisms, Optimization, and Practical Challenges

    Demystifying Mixture of Experts Routing: Core Mechanisms, Optimization, and Practical Challenges

    TL;DR

    • Mixture of Experts architectures scale model capacity while keeping active computational costs low by using a gating network to route tokens to specific expert subnetworks.
    • Empirical analysis demonstrates that expert activation patterns are highly task-conditioned, with prompts from the same category producing highly similar routing signatures.
    • To resolve single-GPU memory bottlenecks, deployment strategies utilize post-training expert pruning and predictive caching systems like ExpertFlow to reduce memory usage and improve throughput.

    An Introduction to Mixture of Experts (MoE) Routing

    MoE routing represents a core architectural component in modern large language models (LLMs), serving as the foundational mechanism for conditional computation. This approach allows a model to scale its overall parameter capacity to hundreds of billions or even trillions of parameters. Crucially, it achieves this scaling while keeping the active computational cost—the number of parameters actually activated to process a given input—equivalent to that of a much smaller model.

    In a standard Transformer architecture, every input token in a sequence is processed by the exact same Feed-Forward Network (FFN) at each layer. A SparseMoE model restructures this setup by replacing that single FFN with multiple parallel subnetworks known as “experts.” To coordinate this architecture, a gating network is positioned before the expert layer. The gating network is a lightweight, trainable neural network that functions as a traffic director. For each incoming token, it determines which specific experts are most mathematically suited to handle the processing and directs the token to them accordingly.

    The standard mathematical execution of MoE routing operates through three primary steps:

    1. Compute Logits: For an incoming token representation x, the gating network computes a dot product with a learnable weight matrix W_g to produce affinity scores (logits) for each expert:
      Logits = x * W_g

    2. Top-k Routing: Rather than sending the token to all experts, which would be computationally expensive, the gating network selects only the top-k experts that yielded the highest affinity scores. In standard setups, k is typically set to 1 or 2.

    3. Softmax Normalization and Weighted Combine: The logits of the selected experts are normalized using a softmax function. Once these chosen experts process the token, their respective outputs are multiplied by these normalized gating network probabilities and summed together to produce the final representation of the token.

    Sources: Mixture of Experts (MoEs) in Transformers, Mixture-of-Experts with Expert Choice Routing, A Visual Guide to Mixture of Experts (MoE) – Maarten Grootendorst, Artificial Intelligence Articles, MoE(Mixture of Experts) by hugging face

    Routing Paradigms: Top-k vs. Soft Routing Mechanisms

    Building on these standard routing mechanics, alternative paradigms have emerged to optimize performance and stability. When evaluating SoftMoE and SparseMoE classifier heads under comparable model capacity, both routing paradigms achieve slightly higher validation accuracy than a dense baseline on the CIFAR10 image classification dataset. To avoid routing collapse, both variants prevent load imbalance through auxiliary loss.

    The primary distinctions between these paradigms emerge in their generalization and loss surface curvature. Hessian-based sharpness metrics at convergence—specifically the largest eigenvalue and the trace of the loss Hessian evaluated on both training and test data—show that SoftMoE exhibits higher sharpness. Meanwhile, SparseMoE and dense models lie in a similar curvature regime, despite all of these models achieving comparable generalization performance. Furthermore, loss surface perturbation analyses show qualitative differences in non-local behavior under finite parameter perturbations between dense and Mixture of Experts models.

    In terms of operational performance, there is a pronounced gap between theoretical and realized efficiency in SparseMoE models. Naively implemented MoE routing fails to yield empirical inference speedups on modern hardware at this scale, highlighting the practical challenges of deploying sparse routing mechanisms.

    Sources: Mixture-of-Experts Models in Vision: Routing, Optimization, and Generalization

    Analyzing Real-World Routing: Task-Conditioned Signatures

    To address these practical deployment challenges and gain deeper insights into how experts actually behave, researchers have introduced routing signatures as a tool to analyze whether expert selection in Mixture of Experts architectures exhibits task-conditioned structure. A routing signature is a vector representation that summarizes expert activation patterns across layers for a given prompt.

    Empirical analysis using the OLMoE-1B-7B-0125-Instruct model as a testbed demonstrates that prompts from the same task category induce highly similar routing signatures, whereas prompts from different categories have significantly lower similarity. Specifically, within-category routing similarity reaches 0.8435 +/- 0.0879, which substantially exceeds the across-category similarity of 0.6225 +/- 0.1687, representing a Cohen’s d of 1.44. The distinctiveness of these signatures is strong enough that a simple logistic regression classifier trained solely on routing signatures can achieve 92.5% +/- 6.1% cross-validated accuracy on four-way task classification.

    To confirm that this separation is a genuine reflection of task-specific routing rather than an artifact of architectural constraints, the analysis incorporated permutation and load imbalance baselines. These baselines demonstrate that the observed task-based clustering cannot be explained by sparsity or load imbalance constraints alone. Furthermore, examining layer-wise signal strength and low-dimensional projections of routing signatures reveals that this task-conditioned structure becomes increasingly apparent in the deeper layers of the model. These findings indicate that routing in sparse transformers operates as a measurable, task-sensitive component of conditional computation rather than serving merely as a load imbalance prevention mechanism. To facilitate further routing telemetry and analysis, researchers have released a lightweight toolkit named MOE-XRAY.

    Sources: Task-Conditioned Routing Signatures in Sparse Mixture-of-Experts Transformers

    Addressing Memory Bottlenecks: Expert Pruning and Predictive Caching

    Beyond analyzing routing signatures, another major hurdle in utilizing these architectures is managing their immense parameter sizes. To mitigate this heavy memory footprint, particularly in memory-constrained environments like single-GPU deployments, researchers have developed post-training sparsification techniques and predictive offloading systems.

    One primary direction to optimize deployment efficiency is pruning and skipping inactive or redundant components. Recent research has introduced post-training approaches that provide plug-and-play, expert-level sparsification through task-agnostic and task-specific expert pruning and expert skipping. Unlike historical weight pruning methods that rely on specialized hardware, these techniques prune or skip entire experts to simultaneously reduce overall model sizes and increase inference speeds while maintaining satisfactory performance across a wide range of tasks.

    Another approach to overcoming memory bottlenecks on a single GPU is dynamic offloading, where inactive experts are stored in CPU memory and loaded on demand. Because static caches disregard input-dependent routing and training separate predictor models is often inaccurate or computationally expensive, the ExpertFlow inference system coordinates three distinct components to manage routing dependencies efficiently:

    • Transformer-based routing path predictor: This component estimates expert usage across all Mixture of Experts layers in a single forward pass.
    • Token scheduler: This groups tokens with similar predicted routes to maximize expert utilization.
    • Predictive expert cache: This system loads only the necessary experts into memory while dynamically correcting any routing mispredictions at runtime.

    By coordinating these three components, ExpertFlow enables highly efficient expert loading and execution. On a single GPU, this predictive caching and scheduling pipeline reduces GPU memory usage by up to 93.72% and improves inference throughput by up to 10x compared to strong offloading baselines.

    Sources: ExpertFlow: Efficient Mixture-of-Experts Inference via Predictive Expert Caching and Token Scheduling, Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models

    Closing thoughts

    Ultimately, as these optimization techniques demonstrate, the viability of Mixture of Experts architectures hinges on closing the pronounced gap between their theoretical computational elegance and the physical realities of modern hardware deployment. While naive routing struggles with empirical inference speeds, the discovery of highly structured, task-conditioned routing signatures in deeper layers demonstrates that these gating networks are executing genuine, meaningful computation rather than arbitrary load balancing. Consequently, by leveraging this predictable behavior through targeted expert pruning and dynamic predictive caching systems like ExpertFlow, developers can successfully resolve severe single-GPU memory bottlenecks. This suggests that the true maturity of MoE technology will be defined not merely by expanding model sizes, but by co-designing hardware-aware routing algorithms with robust diagnostic telemetry.

    Frequently Asked Questions

    What is a routing signature in Mixture of Experts (MoE) architectures?

    A routing signature is a vector representation that summarizes expert activation patterns across layers for a given prompt. It is used by researchers to analyze whether expert selection in MoE models exhibits task-conditioned structures.

    How does the ExpertFlow system optimize single-GPU memory and inference performance?

    ExpertFlow coordinates a Transformer-based routing path predictor, a token scheduler, and a predictive expert cache to manage routing dependencies. This pipeline reduces GPU memory usage by up to 93.72% and improves inference throughput by up to 10x compared to strong offloading baselines.

    What are the three steps involved in standard MoE routing?

    First, the gating network computes affinity scores (logits) for each expert using a dot product. Second, top-k routing selects only the experts with the highest scores, typically setting k to 1 or 2. Finally, the chosen experts’ outputs are normalized using a softmax function and summed together to produce the final representation.


    Keep reading: Training & Architectures

    All guides in this series: LLM Training & Architecture Guides: Fine-Tuning, RL & New Models