Understanding the Internal Mechanics of LLMs

An LLM response arrives as a smooth paragraph, which makes it easy to imagine the model composing the whole answer somewhere inside and then revealing it. That is not what happens. The model repeatedly scores possible next tokens, one step at a time, while a large stack of transformer blocks reworks the context it has so far.

The scale is unusual, but the path is concrete: text becomes token IDs, token IDs become vectors, attention moves information between positions, and an output layer scores what may come next. During training, many target positions can be evaluated in parallel because the correct prefix is already known. During inference, a decoder-only model remains autoregressive: token 101 cannot be sampled until tokens 1 through 100 are available.

1. Tokenization

A model does not receive words directly. A tokenizer maps text to a sequence of integer IDs drawn from a fixed vocabulary.

Modern systems commonly use subword or byte-based methods. Frequent strings may receive one token, while uncommon words split into several pieces. Byte-level fallbacks ensure that arbitrary input can still be represented.

Byte Pair Encoding (BPE) repeatedly merges frequent adjacent units according to a learned vocabulary. Unigram language-model tokenization starts from many candidate pieces and removes those that contribute least to a probabilistic objective. SentencePiece is a tokenizer toolkit that supports algorithms including BPE and unigram; it is not a third algorithm parallel to them.

Tokenization affects cost and behavior. A language or code style that splits into more pieces consumes more context and generation steps. Token boundaries can also create surprising behavior around spelling, arithmetic, and whitespace.

2. Embeddings and Position

An embedding table maps token ID $i$ to a vector:

x_i=E[i]

Token embeddings alone contain no order. Transformers therefore add or apply position information. Designs include learned absolute embeddings, relative-position biases, ALiBi, and rotary position embeddings (RoPE).

RoPE

RoPE rotates pairs of query and key coordinates by an angle determined by position. If $R_m$ is the block-diagonal rotation for position $m$ :

q'_m=R_mq_m,\qquad k'_n=R_nk_n

Their score becomes:

(q'_m)^Tk'_n=q_m^TR_m^TR_nk_n =q_m^TR_{n-m}k_n

This gives the dot product a relative-position structure. The exact frequency schedule, scaling method, and extrapolation behavior depend on the model.

3. A Decoder Transformer Block

A current decoder-only transformer usually contains:

normalization;
causal self-attention;
a residual connection;
another normalization;
a feed-forward or MoE sublayer;
another residual connection.

Architectures differ in whether normalization is placed before or after a sublayer, which normalization is used, and how the feed-forward block is gated.

Residual connections are essential. They give information and gradients a direct route around each sublayer:

h_{l+1}=h_l+F(\operatorname{Norm}(h_l))

4. Causal Self-Attention

For a matrix of hidden states $X$ :

Q=XW_Q,\qquad K=XW_K,\qquad V=XW_V

Scaled dot-product attention is:

\operatorname{Attention}(Q,K,V) =\operatorname{softmax} \left(\frac{QK^T}{\sqrt{d_k}}+M\right)V

$M$ is a causal mask. Entries that would let position $i$ attend to a later position $j>i$ receive negative infinity before softmax. A training example can therefore calculate losses at many positions in parallel without allowing a target token to leak into its own prediction.

Multi-head, multi-query, and grouped-query attention

Multi-head attention gives each head separate query, key, and value projections. Multi-query attention shares one key/value head across several query heads. Grouped-query attention uses a smaller number of key/value heads than query heads.

Sharing key/value heads reduces the KV-cache size and memory bandwidth during generation. It changes model capacity and kernel shape, so the best design depends on the quality and serving target.

A Worked Forward Pass, With Shapes

The equations become easier to reason about when the dimensions stay visible. Consider a decoder with:

hidden width $d=4096$ ;
32 query heads;
8 key/value heads;
head dimension $d_h=128$ ;
batch size $B$ and sequence length $S$ .

The residual stream enters a block with shape $[B,S,4096]$ . After projection and reshaping, grouped-query attention produces:

Q:[B,32,S,128],\qquad K,V:[B,8,S,128].

Four query heads share each key/value head. An implementation may broadcast that relationship logically rather than materializing four copies of every key and value. The conceptual score tensor is still $[B,32,S,S]$ during full causal attention, but an efficient kernel does not need to write that entire tensor to high-bandwidth memory.

After attention, the heads are concatenated back to $[B,S,4096]$ and projected into the residual stream. The feed-forward block expands each token to its intermediate width, applies its gate and nonlinearity, then projects back to 4096. The same weights are used at every sequence position; positions differ because their residual vectors and attention context differ.

Writing shapes beside equations catches common implementation errors: transposing the wrong axes, applying softmax over heads instead of keys, broadcasting a mask incorrectly, or assuming grouped-query attention reduces the number of query heads.

5. The Feed-Forward Sublayer

Attention moves information between positions. The feed-forward network transforms each position independently with shared parameters. A gated form such as SwiGLU can be written schematically as:

\operatorname{SwiGLU}(x) =W_o\left(\operatorname{SiLU}(W_gx)\odot W_ux\right)

In many transformers, these matrices contain a large share of the model’s parameters. Sparse Mixture-of-Experts models replace some dense feed-forward blocks with routed experts, but attention and other shared components remain active.

6. From Hidden State to Next-Token Scores

After the final block, the model projects the hidden state to one logit per vocabulary token:

z_t=W_oh_t+b

Softmax converts the logits into a categorical distribution:

p(x_{t+1}=v\mid x_{\le t}) =\frac{e^{z_{t,v}}}{\sum_u e^{z_{t,u}}}

The model has not selected a token yet. A decoding rule does that. Greedy decoding takes the largest probability; sampling may use temperature, top- $k$ , top- $p$ , repetition controls, or task-specific constraints.

A lower sampling temperature sharpens the distribution, while a higher temperature flattens it:

p_T(v)=\operatorname{softmax}(z/T)_v

This inference temperature is unrelated to the temperature used in knowledge-distillation training, even though the equation looks similar.

7. Pretraining

For tokens $x_1,\ldots,x_n$ , decoder pretraining minimizes next-token negative log-likelihood:

\mathcal{L}=-\sum_{t=1}^{n-1} \log p_\theta(x_{t+1}\mid x_{\le t})

Training commonly uses teacher forcing: the model sees the actual previous tokens rather than its own sampled output. All positions can be evaluated in one forward pass under the causal mask.

Data preparation matters as much as the formula. Deduplication, filtering, mixture weights, licensing, privacy controls, and contamination checks all affect the resulting model and its evaluation.

8. Post-Training

Pretraining produces a next-token model, not automatically a reliable assistant. Post-training may include:

supervised fine-tuning on demonstrations;
preference data and reward modeling;
direct preference objectives;
tool-use examples;
safety training and targeted evaluations.

These stages change behavior, not the fundamental fact that inference produces a probability distribution over the next token. They also introduce their own data biases and reward-hacking risks.

9. Prefill and Autoregressive Decode

Inference has two visibly different phases.

Prefill

The model processes the prompt and builds attention keys and values for every layer. Prompt positions can be handled in parallel, so prefill resembles a large matrix workload. Its cost grows strongly with prompt length, though attention variants may alter the exact scaling.

Decode

The model generates one token, appends it, and repeats. Cached keys and values avoid recomputing the full prefix at every step. Decode is often limited by memory bandwidth because each new token must read a large amount of model and cache state.

The KV cache grows with layers, cached sequence length, key/value heads, head dimension, batch size, and bytes per element. Grouped-query attention, lower cache precision, sliding windows, and paged cache management can reduce the pressure.

A KV-Cache Calculation You Can Audit

For a decoder with $L$ layers, batch size $B$ , cached length $S$ , $H_{kv}$ key/value heads, head dimension $d_h$ , and $b$ bytes per value, the unsharded cache is approximately:

M_{KV}=2LB S H_{kv}d_hb.

The factor of two accounts for keys and values. With 32 layers, one sequence of 32,768 tokens, 8 key/value heads, head dimension 128, and BF16 values:

M_{KV} =2\cdot32\cdot1\cdot32768\cdot8\cdot128\cdot2 =4{,}294{,}967{,}296\text{ bytes},

or 4 GiB for one sequence. This excludes model weights, temporary buffers, allocator fragmentation, and cache metadata. With 32 key/value heads instead of 8, the same calculation reaches 16 GiB. GQA is therefore not a cosmetic architectural detail; it changes how many concurrent long requests a server can hold.

Paged cache management addresses a different problem. It reduces waste caused by reserving contiguous space for sequences whose lengths grow unpredictably. It does not change the bytes required by a live key or value element.

The calculation is small enough to keep in an experiment notebook:

def kv_cache_gib(
    layers: int,
    batch: int,
    tokens: int,
    kv_heads: int,
    head_dim: int,
    bytes_per_value: int,
) -> float:
    total = 2 * layers * batch * tokens * kv_heads * head_dim * bytes_per_value
    return total / (1024 ** 3)


print(kv_cache_gib(32, 1, 32_768, 8, 128, 2))  # 4.0

This formula is also a useful sanity check for tensor-parallel serving. If key/value heads are partitioned across devices, calculate the local share and then add any replicated or communication buffers separately.

10. FlashAttention

The basic attention equation suggests materializing an $n\times n$ score matrix. FlashAttention calculates attention in tiles and carefully maintains the softmax statistics, reducing traffic between high-bandwidth memory and on-chip memory. It is an exact attention algorithm up to normal floating-point differences, not an approximation to a different objective.

Later variants improve work partitioning and hardware use. Backward passes can recompute selected intermediate values instead of storing the full attention matrix. The main gain is I/O efficiency; the mathematical dense-attention workload still has quadratic pairwise interactions in sequence length.

FlashAttention-3 goes further on Hopper GPUs by overlapping tensor-core work with data movement and by supporting an FP8 path with block quantization. Its reported speedups are properties of named hardware, precisions, shapes, and kernels—not a promise that every attention call becomes twice as fast.

Sparse and linear-attention methods have different complexity. Sparse attention is not universally $O(n\log n)$ ; the result depends on its pattern and number of attended positions.

11. Numeric Precision

Training frequently combines formats:

FP32 for selected accumulations and optimizer state;
BF16 or FP16 for many matrix operations;
sometimes lower-precision formats on supported hardware.

FP16 has a narrow exponent range, so loss scaling is often needed to prevent small gradients from underflowing. BF16 has roughly FP32’s exponent range and normally does not require FP16-style loss scaling, though it has fewer fraction bits. “Mixed precision” should not be summarized as one universal recipe.

12. Quantization

Quantization maps real values to a finite integer range. For an affine scheme:

q=\operatorname{clip} \left(\operatorname{round}\left(\frac{x}{s}\right)+z, q_{\min},q_{\max}\right)

and dequantization is:

\hat x=s(q-z)

$s$ is the scale and $z$ the zero point. Symmetric quantization often uses $z=0$ . Scales may be per tensor, per channel, or per group. Weight-only quantization reduces model storage and weight bandwidth; weight-and-activation quantization can save more compute but is harder to calibrate.

The bit width alone does not predict quality or speed. Group size, outlier handling, calibration data, kernel availability, cache precision, and hardware support matter.

13. Scaling Laws, Carefully Stated

Empirical language-model studies have found power-law relationships between reducible validation loss, parameters, data, and compute over particular ranges. A schematic form is:

L(N,D)\approx L_\infty+A N^{-\alpha}+B D^{-\beta}

The irreducible term and fitted constants matter. The exponents come from measured model families and should not be treated as laws of nature.

Kaplan et al. emphasized model scaling under their setup. The later Chinchilla study found that, for its compute-optimal regime, model size and training tokens should grow roughly together. These results cannot be reduced to “data matters slightly more” by comparing two exponents without their definitions and constraints.

14. Emergent Abilities

Some benchmark curves appear flat for smaller models and then jump at a larger scale. That observation does not justify a universal sigmoid equation for capability emergence. Discrete metrics such as exact match can turn smooth changes in token probabilities into an apparent threshold. Test-set size and model selection can add more discontinuity.

Abrupt changes may still be important, but they should be reported as observations under a named metric and model family, not as a settled physical law of LLMs.

15. What the Model “Knows”

Training compresses statistical structure from data into parameters. A model can recall facts, combine patterns, follow instructions, and sometimes solve unfamiliar tasks. It can also produce fluent falsehoods because next-token likelihood is not a database consistency check.

Internal activations are distributed and context-dependent. Attention maps, probes, and feature visualizations can reveal useful patterns, but none provides a complete transcript of the model’s reasoning.

Reading an Inference Trace

When a serving run is slow, “the model is large” is not yet a diagnosis. Separate at least four quantities:

time to first token, which includes queueing and prefill;
inter-token latency, which describes decode responsiveness;
request throughput, which depends heavily on batching and scheduling;
goodput, which counts requests that meet a stated latency objective.

For batch-one decode, every new token uses the full stack of model weights while doing relatively little matrix work per byte loaded. That often makes decode memory-bandwidth-bound. Batching lets several tokens reuse weights loaded for the same step and raises arithmetic intensity, but it can also increase latency and KV-cache pressure.

A rough lower-bound exercise makes the constraint tangible. Seven billion BF16 parameters occupy about 14 GB. If a decode step had to stream those weights from a device sustaining 2 TB/s, weight traffic alone would take roughly 7 ms. That is an optimistic bound: it ignores the KV cache, synchronization, kernel launch overhead, incomplete bandwidth utilization, and every non-weight operation. It should be used to reject impossible throughput claims, not to predict an exact benchmark.

A credible serving report therefore names the model, precision, accelerator, tensor-parallel degree, prompt and output-length distributions, concurrency, scheduler, cache policy, and latency percentiles. “Tokens per second” without those details is difficult to compare with anything.

Conclusion

The transformer itself is only part of an LLM system. Tokenization shapes the input, causal attention and feed-forward blocks perform the main computation, training data shapes the parameters, post-training shapes behavior, and the serving stack determines whether generation is affordable. Keeping training, prefill, and decode separate clears up many of the myths that collect around these models.

References

Vaswani et al., Attention Is All You Need (2017).
Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding (2021).
Kaplan et al., Scaling Laws for Neural Language Models (2020).
Hoffmann et al., Training Compute-Optimal Large Language Models (2022).
Dao et al., FlashAttention (2022).
Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention (2023).
Shah et al., FlashAttention-3 (2024).
Schaeffer, Miranda, and Koyejo, Are Emergent Abilities of Large Language Models a Mirage? (2023).