Skip to content

Technical Reference: LLM Architecture, Embedding Math, and KV-Cache Memory Optimization

This document serves as a comprehensive reference mapping the engineering lifecycle of Transformer-based Large Language Models (LLMs) during text generation (inference), tracking token transformation from raw vector embeddings to optimized hardware physical memory footprints.

1. The Model Foundation: Embedding Dimension (\(d\))

When an LLM architecture is created, its core width—the Embedding Dimension (often denoted as \(d\), \(d_{model}\), or hidden_size)—is permanently set. This cannot change after training.

  • Vectorization: Every token (word or sub-word) entering the model is transformed into a dense vector of numbers exactly \(d\) elements long.
  • Example (Llama 3 8B): The model defines \(d = 4,096\). Therefore, every single token is processed as a unique vector containing 4,096 floating-point numbers.

2. Attention Projections: Query (Q), Key (K), and Value (V)

Inside the attention layers, the model creates three distinct "views" of the token embedding vector. To prevent a drop in semantic capacity, \(Q\), \(K\), and \(V\) are not compressed strings or simple indices; they are dense, high-dimensional vector arrays.

Core Conceptual Roles

  • Query (Q) - The "Current Search Query": Represents the immediate token being processed at the current timestep. It acts like an uncompleted puzzle piece, asking: "What contextual information do I need from the past to understand my role in this sentence?"
  • Key (K) - The "Folder Label": Represents how a past token advertises its own context and grammar to the rest of the text. It is a dynamic label, changing its numerical values based on surrounding context (e.g., distinguishing "Apple" the technology company from "Apple" the fruit).
  • Value (V) - The "Folder Content": Holds the raw semantic and informational content of a token. Once a Query successfully matches a Key, this is the payload of meaning that gets extracted and combined.

3. Multi-Head Attention (MHA) Mechanics

Instead of running a single attention calculation across the entire embedding dimension, the model processes data in parallel using independent Attention Heads to look for different linguistic patterns simultaneously (e.g., syntax on head 1, pronouns on head 2).

The Split-and-Glue Lifecycle

Step 1: Splitting into Head Dimensions (\(d_{head}\))

The fixed embedding dimension (\(d\)) is divided evenly by the total number of attention heads to determine the vector size for an individual head:

\[\text{Head Dimension } (d_{head}) = \frac{\text{Embedding Dimension } (d)}{\text{Number of Heads}}\]
  • Mathematical Example (Llama 3 8B): $\(\text{Head Dimension } = \frac{4,096 \text{ (Embedding Size)}}{32 \text{ (Heads)}} = \mathbf{128}\)$ Every individual attention head processes a Query of size 128, a Key of size 128, and a Value of size 128. Total capacity is conserved across the network (\(32 \times 128 = 4,096\)).

Step 2: Parallel Processing

All 32 heads perform their attention calculations concurrently inside their localized 128-dimensional vector spaces.

Step 3: Reconstructing (The Glue Phase)

Once the heads finish calculating, each outputs a new vector of size 128. The model concatenates (glues) these 32 individual vectors back-to-back:

\[\text{Head}_1(128) + \text{Head}_2(128) + \dots + \text{Head}_{32}(128) = \mathbf{4,096}\]

This perfectly reconstructs the original 4,096 embedding dimension to pass it forward to the next layer of the neural network. No part of the embedding space is wasted.

4. Context Length (\(N\)) and the \(O(N^2)\) Complexity Bottleneck

While the Vector Embedding Dimension (\(d\)) is a fixed architectural width, the Context Length (\(N\)) is a dynamic length representing the total number of tokens currently in the conversation history.

Why Attention Compute Scales at \(O(N^2)\) (Quadratic)

During the initial processing phase (prefill), the model forces every single token in the context window to evaluate its relationship with every other token. * If a sequence has \(N\) tokens, and each token must execute a calculation against \(N\) tokens, the total calculation density scales at \(N \times N = \mathbf{N^2}\).

Impact of Scaling Context Length

  • Context = 1,000 tokens: The model handles \(1,000 \times 1,000 = \mathbf{1,000,000}\) comparison operations.
  • Context = 2,000 tokens: The model handles \(2,000 \times 2,000 = \mathbf{4,000,000}\) comparison operations.
  • Result: Doubling the text history quadruples the computational strain on the hardware, independent of the fixed embedding size.

5. Grouped-Query Attention (GQA): Compressing the Footprint

Historically, standard Multi-Head Attention (MHA) mandated a 1:1:1 ratio of Query, Key, and Value heads, causing the KV-cache to grow too large to fit into affordable GPU memory.

To solve this without sacrificing model accuracy, modern LLMs use Grouped-Query Attention (GQA).

The Structural Shift

Instead of assigning a unique Key and Value head to every single Query head, GQA groups the Query heads together. Multiple Query heads share a single, unified Key and Value head.

  • The Configuration (e.g., Llama 3 8B): Uses 32 Query heads but groups them into 8 groups.
  • The GQA Ratio: 4 Query heads share 1 Key head and 1 Value head.

1. Multi-Head Attention (MHA)

  • The Mechanism: Every single Query head has its own dedicated, completely unique Key and Value head.
  • The Ratio: 1:1 (For 32 Query heads, there are 32 Key heads and 32 Value heads).
  • The Mapping: 
    [Q1] -------> [Key 1 / Value 1]
[Q2] -------> [Key 2 / Value 2]
[Q3] -------> [Key 3 / Value 3]
[Q4] -------> [Key 4 / Value 4]
... (Repeats for all 32 heads)
  • Memory Footprint: Maximum footprint. The system must store 32 parallel channels of Key and Value vectors in the KV-cache for every layer.

2. Grouped-Query Attention (GQA)

  • The Mechanism: Query heads are divided into distinct teams or clusters. All the Query heads within a specific cluster pool their resources and share one unified Key and Value head.
  • The Ratio: Group-dependent (For 32 Query heads split into 8 groups, the ratio is 4:1).
  • The Mapping: 
    Group 1: [Q1][Q2][Q3][Q4] ----------> [Key 1 / Value 1 Head]
Group 2: [Q5][Q6][Q7][Q8] ----------> [Key 2 / Value 2 Head]
Group 3: [Q9][Q10][Q11][Q12] -------> [Key 3 / Value 3 Head]
Group 4: [Q13][Q14][Q15][Q16] -------> [Key 4 / Value 4 Head]
Group 5: [Q17][Q18][Q19][Q20] -------> [Key 5 / Value 5 Head]
Group 6: [Q21][Q22][Q23][Q24] -------> [Key 6 / Value 6 Head]
Group 7: [Q25][Q26][Q27][Q28] -------> [Key 7 / Value 7 Head]
Group 8: [Q29][Q30][Q31][Q32] -------> [Key 8 / Value 8 Head]
  • Memory Footprint: Highly compressed. The system only stores 8 channels of Key and Value vectors in the KV-cache per layer, cutting the memory storage requirement by exactly 75%.

Real-World Architectural Gains

  1. 75% Memory Reduction: By dropping from 32 Key/Value heads down to 8, the physical size of the KV-cache sitting inside GPU memory shrinks by 4x.
  2. Maintained Accuracy: Because the queries are grouped logically, the model retains nearly 100% of its reasoning capabilities compared to full MHA.
  3. Hardware Capacity Multiplier: Shrinking the cache allows cloud providers to host up to 4x more concurrent users on the exact same GPU cluster.

6. The KV-Cache and Modern Memory Optimization

The Autoregressive Inference Trick

During generation, text is written one token at a time. The past tokens never change. Therefore, their calculated Key (\(K\)) and Value (\(V\)) vectors remain completely static. * The Cache: To bypass the \(O(N^2)\) compute loop at every single token generation step, the model calculates \(K\) and \(V\) for old tokens once and saves them to the GPU's memory (The KV-Cache). * The Query Exemption: The Query (\(Q\)) vector represents only the immediate "question" being asked by the single newest token. Once that token is generated, that specific question is resolved, and \(Q\) is discarded. It is never cached.

Updated GQA KV-Cache Memory Formula

Because GQA reduces the physical number of Key and Value heads, the formula to calculate memory consumption uses KV_Heads instead of total query Heads:

\[\text{KV-Cache Size (Bytes)} = 2 \times \text{Tokens} \times \text{Layers} \times \mathbf{\text{KV-Heads}} \times \text{Head Dimension} \times \text{Bytes-per-Param}\]

(Note: The leading multiplier of 2 accounts for storing both the Key array and the Value array simultaneously).

PagedAttention Frameworks (e.g., vLLM)

Historically, serving systems had to allocate massive, continuous chunks of GPU memory matching the maximum context length upfront, creating immense internal fragmentation and wasting up to 60% of physical hardware capacity on empty space.

Modern frameworks like vLLM utilize PagedAttention to eliminate this: 1. Dynamic Pages: The KV-cache for a conversation is sliced into small, discrete pages (typically 16 tokens per page). 2. Virtual Allocation: These pages are scattered dynamically across any open, non-contiguous physical addresses in the GPU's memory. 3. The Page Table: A virtual routing table maps the logical order of the conversation to its scattered physical addresses, cutting memory waste to near 0% and allowing 2x to 4x more concurrent users to fit on the same hardware cluster.

Advanced Inference Optimizations: MoE and Speculative Decoding

As models scale to hundreds of billions of parameters, hardware boundaries require structural changes to how models are organized (MoE) and how text is generated (Speculative Decoding).

1. Mixture-of-Experts (MoE) — Sparse Activation

Traditional "dense" models force 100% of their parameters to calculate every single token. MoE shifts this to a "sparse" architecture, breaking neural network layers into specialized sub-networks called Experts.

The Mechanism

  1. The Router: A lightweight gating network sits at the entrance of the layer. It inspects the incoming token vector.
  2. Dynamic Routing: The router determines which experts are best suited for the token (e.g., routing a math-heavy token to a calculation expert, or a coding token to a syntax expert).
  3. Sparse Execution: Only a fraction of the total parameters (typically 1 or 2 experts) activate to process the token. The remaining experts stay completely idle.

Architectural Impact

  • Total vs. Active Parameters: A model can have 120 Billion total parameters (the capacity to know many things) but only use 22 Billion active parameters per token (the speed of a lightweight model).
  • Efficiency: Drastically reduces the computational cost (\(O\) operations) required per token, allowing massive foundational models to run on standard enterprise hardware clusters.

2. Speculative Decoding — Parallel Verification

LLM inference is traditionally execution-bound by a step-by-step autoregressive loop, generating exactly one token per forward pass. Speculative Decoding breaks this bottleneck by using a dual-model system.

The Mechanism

  1. The Draft Model: A tiny, highly efficient model (e.g., a 1B model) aggressively guesses a short sequence of upcoming tokens (e.g., the next 5 tokens) one-by-one at extreme speed.
  2. The Target Model: The primary, massive LLM (e.g., a 70B model) takes the entire 5-token guessed sequence and processes it simultaneously in a single parallel batch step.
  3. Verification Matrix:
  4. Tokens Accepted: If the Target model agrees with the Draft model's logic, it validates all 5 tokens at once. The system prints 5 tokens in the span of a single execution cycle.
  5. Tokens Rejected: If the Draft model guesses incorrectly at token 3, the Target model discards tokens 3-5, outputs the correct token 3, and resets the drafting loop.

Architectural Impact

  • Latency Reduction: Because modern GPUs can verify a batch of tokens almost as fast as they can compute a single token, Speculative Decoding typically yields a 2x to 3x wall-clock speedup in text generation without changing the model's final mathematical output weights.

Real-World Production Adoption of Speculative Decoding

Speculative decoding is a standard optimization utilized across commercial AI endpoints. Because it guarantees mathematically identical outputs to a model operating alone, it is a primary driver for increasing tokens-per-second throughput in user-facing applications.

Frontier Production Examples

1. Google Gemini Lineup

  • Architecture: Utilizes small, highly aligned "drafter" models running in parallel with multi-trillion parameter "verifier" models.
  • Open-Source Integration: Google's Gemma 4 open-weight models natively integrate speculative decoding setups to achieve up to a 3x speed increase out of the box.

2. Meta Llama 3.1 & Llama 4 Systems

  • Architecture: Meta relies on distributed speculative decoding to make massive dense/MoE models viable for consumer chat latency.
  • Performance: Running Llama 3.1 405B or Llama 4 clusters using NVIDIA's TensorRT-LLM with a tailored draft model yields over a 3x latency reduction.

3. DeepSeek V3/V4 & Coder Series

  • Architecture: Replaces the traditional "two-model" (Draft + Target) approach with a native Multi-Token Prediction (MTP) architecture.
  • Mechanism: Auxiliary speculator heads are built directly onto the final hidden layer of the primary model. These layers are co-trained to predict multiple tokens ahead natively, eliminating the memory overhead of maintaining a separate secondary draft model in VRAM.

4. NVIDIA Nemotron-3 Ultra

  • Architecture: A hybrid Mixture-of-Experts (MoE) and Mamba-Attention architecture designed specifically for modern enterprise inference.
  • Mechanism: Integrates native MTP layers built directly into the silicon optimization pipeline, maximizing token generation speeds on Blackwell and Hopper architectures.

Architectural Alternatives: State Space Models & Mamba

While Transformers rely on Attention and the KV-cache, alternative architectures seek to bypass the \(O(N^2)\) complexity bottleneck by re-engineering sequence processing.

1. The Core Concept of Mamba (State Space Models / SSMs)

Mamba is a non-Transformer architecture based on Selective State Space Models. Instead of cross-referencing a new token against every individual past token, Mamba processes information sequentially through a continuously updated, fixed-size internal "state."

  • Transformer Ingestion: Retains raw history matrices (Filing Cabinet analogy).
  • Mamba Ingestion: Compresses history dynamically into a single mathematical summary vector that updates token-by-token (Mental Summary analogy).

2. Structural Breakdown: Transformer vs. Mamba

Architectural Feature Transformer Architecture Mamba Architecture (SSM)
Compute Complexity Quadratic \(O(N^2)\) Linear \(O(N)\)
Memory Scaling Scales linearly with context length (KV-Cache grows) Completely flat (Fixed internal state size)
Data Retrival Perfect long-range needle-in-a-haystack recall Lossy long-range compression over massive contexts
Hardware Constraint Highly memory-bandwidth bound Highly compute-bound (Matrix updates)

3. The Shift to Hybrid Systems (e.g., NVIDIA Nemotron)

To exploit the speed of Mamba without sacrificing the absolute pinpoint recall accuracy of Transformers, frontier enterprise architectures frequently deploy Hybrid Attention-Mamba networks.

The Hybrid Topology

Instead of utilizing a single paradigm, the layers alternate within the model: 1. Mamba Layers (80%): Handle the heavy lifting of sequence processing, context compression, and rapid text generation at linear \(O(N)\) speed with zero KV-cache overhead. 2. Transformer Attention Layers (20%): Injected periodically to act as pinpoint memory anchors, preserving the model's ability to execute complex cross-referencing and exact "needle-in-a-haystack" data retrieval.

Hardware Deployment Epilogue: NVIDIA H200 vs. Blackwell

The algorithmic concepts of embeddings, attention heads, KV-caches, and model architectures map directly onto physical silicon constraints. Choosing between Hopper (H200) and Blackwell (B200/GB200) depends entirely on cluster scale and precision strategy.

1. NVIDIA H200 (Hopper Architecture)

The H200 serves as the ultimate maturation of the standard datacenter GPU. It was engineered specifically to alleviate memory bandwidth bottlenecks in standard AI inference and classic scientific computing clusters.

  • Primary Workload Target: Enterprise LLM Inference, Serving, Retrieval-Augmented Generation (RAG), and high-double-precision (FP64) classic HPC physics simulations.
  • Pros:
    • Massive, ultra-fast memory footprint (141 GB HBM3e running at 4.8 TB/s).
    • Excellent standard-precision handling for large KV-caches and long conversation contexts.
    • Drop-in compatibility; operates flawlessly in standard, air-cooled PCIe or HGX datacenter racks without infrastructure redesigns.
  • Cons:
    • Lacks the ultra-dense low-precision compute matrix (FP4) of the next generation.
    • Interconnect bandwidth (NVLink 4 at 900 GB/s) is outpaced by Blackwell when scaling to thousands of continuous nodes for massive foundational training.

2. NVIDIA Blackwell (B200 / GB200 Architecture)

Blackwell represents a generational architectural leap utilizing a dual-die chiplet design packing 208 billion transistors. It is built to act as the primary engine for massive "AI Factories" scaling multi-trillion parameter models.

  • Primary Workload Target: Frontier Foundation Model Training, Large-Scale Fine-Tuning, and high-concurrency multi-modal real-time reasoning.
  • Pros:
    • Massive training leap with 5th-Gen NVLink providing double the interconnect speed (1.8 TB/s) to eliminate GPU-to-GPU data transmission lag.
    • Introduces native FP4 precision and a 2nd-Gen Transformer Engine, yielding a 15x inference throughput boost for dense networks.
    • Higher memory capacity (192 GB HBM3e) and faster bandwidth (8.0 TB/s) to house giant parallel batch KV-caches.
  • Cons:
    • Imposes extreme structural demands on datacenters, frequently requiring custom, liquid-cooled, whole-rack infrastructure setups (like the NVL72) due to massive power targets (up to 14 kW per system).

3. The Deployment Matrix Summary

Core Metric H200 (Hopper) Blackwell (B200)
Optimal Use Case Model Serving / LLM Inference / HPC Frontier Model Training / Mass Scale AI
Interconnect (NVLink) 900 GB/s 1,800 GB/s (1.8 TB/s)
VRAM Capacity 141 GB HBM3e 192 GB HBM3e
Cooling Strategy Standard Air-Cooled / Standard Racks Liquid-Cooled / Whole-Rack Solutions

Hardware Economics and Inference Concurrency Reference

1. Approximate Market Costs (Enterprise Datacenter Tiers)

  • NVIDIA H200 (Hopper): Individual chip MSRP ~\(35k–\)40k | Complete 8-GPU Server Node ~\(400k–\)500k. Cloud rental runs ~$3.50 to $10.00/hr.
  • NVIDIA Blackwell B200: Individual chip MSRP ~\(45k–\)50k | Complete Server Node/Rack setups frequently scale from $500k to \(1M+. Cloud rental runs ~\)5.00 to $18.00/hr.

2. Single-GPU Concurrency Benchmarks

  • Lightweight Networks (~8B Parameters): Supports 64 to 256+ concurrent user streams per GPU using modern PagedAttention block systems.
  • Medium Networks (~70B Parameters): Supports 16–32 concurrent users on H200 (VRAM limited) and 32–64 concurrent users on Blackwell B200 (benefiting from 192 GB VRAM capacity).

3. Production Deployment Architecture (Frontier Models: Claude / Gemini)

  • Distributed Architecture: Ultra-large frontier models do not execute on single cards; they are partitioned across multi-GPU cluster nodes (Tensor Parallelism) to store massive weight matrices.
  • Throughput Scaling: An 8-GPU cluster node using H200 hardware can process roughly 400–800 continuous parallel streams before hitting memory saturation. Due to FP4 optimization and 1.8 TB/s NVLink 5 data streaming buses, an 8-GPU Blackwell B200 cluster node expands this ceiling to support 2,000 to 4,000+ simultaneous concurrent users on the exact same structural network footprint.

Infrastructure & Serving Framework Reference: vLLM vs. TensorRT-LLM

Deploying modern architectures requires an inference serving engine that orchestrates contiguous batching, memory allocation (PagedAttention), and multi-GPU tensor communication.

1. vLLM (The Ecosystem Default)

Developed by UC Berkeley and backed by a massive open-source community (including direct contributions from Meta, Google, and DeepSeek), vLLM is the most common serving layer in the AI ecosystem.

  • Primary Advantage: Flexibility and Velocity. It features hardware-agnostic architecture (running on NVIDIA, AMD, Google TPUs, and AWS Neuron) and offers instant, plug-and-play support for newly released models without preprocessing.
  • Optimal Use Case: Fast prototyping, multi-vendor hardware environments, dynamic model switching, and organizations prioritizing rapid engineering iteration.

2. TensorRT-LLM (The Silicon Squeezer)

Developed natively by NVIDIA, TensorRT-LLM is a highly specialized framework designed to squeeze the absolute maximum performance boundaries out of NVIDIA hardware (Hopper, Blackwell).

  • Primary Advantage: Raw Throughput. By utilizing Ahead-Of-Time (AOT) compilation, it fuses mathematical layers natively at the CUDA kernel level, yielding a 15% to 30% performance advantage over vLLM.
  • The Trade-Off: High development friction. It requires a lengthy compilation step (often 30+ minutes) every time a model or configuration changes, and it locks your deployment pipeline strictly into the NVIDIA hardware ecosystem.
  • Optimal Use Case: Massive, high-volume production APIs where a model is frozen for months and every millisecond saved translates directly to millions of dollars in reduced server costs.

3. SGLang (The Agentic Specialization)

An emerging, highly competitive open-source runtime engine that focuses heavily on speeding up complex, multi-step LLM behaviors (like AI agents and complex chain-of-thought workflows).

  • Primary Advantage: Prefill Cache Sharing (RadixAttention). While vLLM manages token pages for individual chats, SGLang automatically recognizes and caches massive system prompts or codebase structures across completely different users. If 1,000 developers ask Copilot questions about the same code repository, SGLang reads the context from GPU memory once and instantly serves all 1,000 queries.
  • Optimal Use Case: Multi-turn AI Agent workflows, heavy structured generation (JSON routing), and coding assistants with massive system prompts.