Tokens, Embeddings, Tensors, and Transformer Mechanics

A language model does not receive a sentence and return a sentence. A common decoder-only path is:

Unicode text
  → tokenizer pipeline
  → token IDs [sequence]
  → embedding lookup + position information [sequence, model_width]
  → repeated transformer blocks [sequence, model_width]
  → vocabulary projection [sequence, vocabulary]
  → logits for the next position
  → sampling policy chooses one token ID
  → append token and repeat
  → token decoder produces bytes/text

Each arrow is an application contract. A wrong tokenizer changes IDs. A wrong axis makes matrix multiplication invalid or, worse, broadcasts into plausible nonsense. A missing causal mask leaks future tokens during training/evaluation. An unstable softmax creates NaNs. An underestimated sequence length can turn an apparently small request into quadratic attention work and a large KV cache.

This chapter implements one transparent, single-head causal block in mosaic-ml. It is small enough to inspect line by line and strict enough to expose malformed shapes, non-finite numbers, out-of-vocabulary IDs and capacity violations as typed errors. It is not a production tensor runtime and its fixed numbers are not trained weights. Once the data flow is clear, later chapters replace the loops with Candle, mistral.rs or an inference service.

Learning objectives

By the end you will be able to:

• describe normalization, pre-tokenization, subword modeling, post-processing and decoding as separate tokenizer stages;
• explain why tokens are neither words nor characters and why token counts vary by model and input;
• distinguish a token string, token ID, embedding vector, hidden state, logit and probability;
• read two- and three-dimensional tensor shapes without losing the meaning of each axis;
• reason about dtype, device, layout, stride, contiguity and broadcasting at an application boundary;
• derive the shapes in scaled dot-product attention;
• explain query, key and value as learned projections rather than database concepts;
• apply a causal mask and numerically stable softmax;
• distinguish multi-head attention from the chapter’s single-head implementation;
• follow residual, normalization and feed-forward paths in a modern decoder block at a useful level;
• identify the last-position logits used for autoregressive generation;
• connect image patches, audio frames and video/media features to transformer sequences without pretending all modalities are text tokens;
• estimate why full attention grows quadratically with sequence length;
• test the causality invariant: changing a future token cannot change an earlier output;
• know which mechanics belong to a model runtime and which limits remain the harness’s job.

Dependencies and completion artifact

This chapter assumes:

• ownership and slices;
• checked matrix dimensions and fallible constructors;
• Unicode/text boundary knowledge;
• the previous model-bundle chapter;
• basic algebra: vectors, dot products and matrix multiplication.

Run the artifact:

cd rust/rust-ai-engineering
cargo run -q -p mosaic-ml --example causal_mechanics

It reports:

{
  "embedding_shape": [3, 3],
  "attention_shape": [3, 3],
  "logits_shape": [3, 4],
  "argmax_token": { "id": 2, "text": "builds" },
  "warning": "didactic fixed weights; not a trained language model"
}

The exact report also includes the lower-triangular attention weights and final logits. The result does not demonstrate intelligence. It demonstrates a verifiable shape and causality contract.

Notation: name every axis

We use:

Symbol	Meaning
B	batch size
S	input sequence length
T	generated/decoded tokens so far
D	model width (d_model)
H	attention head count
Dk	query/key width per head
Dv	value width per head
V	vocabulary size
L	transformer layer count

For one unbatched sequence, hidden state is [S, D]. Production runtimes usually add batch: [B, S, D]. Some kernels use different physical/order conventions. Never infer axis meaning from two equal numbers. A [32, 32] tensor could be sequence-by-vocabulary, head-by-head, or an image tile; the type/API and model contract supply semantics.

Write shapes beside equations:

token_ids                     [S]
token_embedding_table         [V, D]
embedded sequence X           [S, D]
Wq                            [D, Dk]
Q = X × Wq                    [S, Dk]
K = X × Wk                    [S, Dk]
Vvalues = X × Wv              [S, Dv]
scores = Q × transpose(K)     [S, S]
attention = softmax(scores)   [S, S]
context = attention × Vvalues [S, Dv]
lm_head                       [D, V]
logits                        [S, V]

The letter V is unfortunately used for both vocabulary size and the value matrix in common literature. This chapter writes Vvalues when ambiguity matters.

Step 1: text becomes token IDs

Tokens are learned units, not linguistic truth

A tokenizer may emit:

• a whole frequent word;
• a word fragment;
• punctuation or whitespace;
• one or several Unicode characters;
• byte-level representations;
• special control units that never appear as ordinary source text.

The spelling token often means three different things:

1. token surface/piece, such as "rust" or a whitespace-prefixed piece;
2. integer token ID, such as 18421;
3. position in a particular encoded sequence.

Keep them separate in APIs and logs. A token ID is meaningful only under the exact tokenizer identity. ID 42 from tokenizer A is not portable to tokenizer B.

The tokenizer pipeline

The Hugging Face Tokenizers documentation separates four encoding stages:

raw text
  → normalization
  → pre-tokenization
  → tokenization model
  → post-processing
  → IDs + offsets + masks/type IDs

Normalization may perform Unicode normalization, case folding, accent handling or other transformations. Destructive normalization can make exact reconstruction impossible unless alignment metadata is retained.

Pre-tokenization proposes boundaries—for example whitespace/punctuation or byte-level splits. It constrains where the trained subword algorithm operates.

Model applies WordPiece, BPE, Unigram, WordLevel or another learned mapping and assigns IDs. SentencePiece is notable for training subword models directly from raw sentences rather than requiring a language-specific word tokenizer.

Post-processing adds model-required special tokens, sequence separators, segment/type IDs and templates. Chat templating happens conceptually before or around this pipeline: messages are serialized to the exact control-token structure expected by the model, then encoded.

Decoding maps IDs back through the tokenizer’s decoder rules. decode(encode(text)) == text is not universal because normalization and unknown-token behavior may be lossy.

What the application must retain

For provenance and citation, keep:

• original bytes or trusted decoded text artifact;
• tokenizer identity;
• token IDs only when needed and protected from excessive logging;
• token-to-source byte/character offsets where supported;
• special-token mask;
• attention/padding mask;
• truncation window and original length;
• rendered template identity.

Offsets are essential for mapping generated citations or token classifications back to the original text. Unicode byte, scalar-value, grapheme and display-column offsets are different coordinate systems; label the one you store.

Counting is model-specific

Characters, words and bytes are not safe substitutes for tokens. The same string can have different counts under different vocabularies, normalization rules and templates. Code, JSON, uncommon scripts, long numbers, whitespace and binary-to-text encodings can expand unexpectedly.

Admission needs the selected model’s tokenizer or a provider’s authoritative counting contract:

rendered system + tools + retrieved evidence + history + user input
  → exact tokenizer
  → input token count
  → reserve room for output and runtime overhead

If a remote provider’s private tokenizer is unavailable, use its count endpoint/response usage when available, maintain a conservative measured bound, and still handle provider-side rejection.

Step 2: IDs select embedding rows

An embedding table is a learned matrix:

E [V, D]

For each token ID, lookup selects one row. For IDs [7, 12, 7], output is:

X_tokens [3, D] = [E[7], E[12], E[7]]

This is table lookup, not a semantic search. During training, gradients adjust rows so they become useful coordinates for the model’s objective. Similar meanings may occupy useful geometric relationships, but no individual dimension has to mean “plural” or “positive sentiment.”

The companion makes bounds explicit:

pub fn embed(&self, token_ids: &[usize]) -> Result<Matrix, MlError> {
    if token_ids.is_empty() {
        return Err(MlError::EmptySequence);
    }
    for (position, &token_id) in token_ids.iter().enumerate() {
        let row = self.values.row(token_id).ok_or(
            MlError::TokenOutOfVocabulary {
                position,
                token_id,
                vocabulary_size: self.vocabulary_size(),
            },
        )?;
        // append exactly D values
    }
    // result shape: [S, D]
}

Real tokenizers normally emit their unknown-token ID rather than an invalid ID for unknown text. An out-of-range ID instead suggests mismatched tokenizer/weights, corrupted input, unsafe deserialization or a runtime bug. It must fail before indexing memory.

Position information

Self-attention alone has no inherent order. Models add or incorporate position using learned absolute embeddings, sinusoidal encodings, rotary position embeddings, relative biases or other schemes. These are not freely interchangeable.

The teaching block adds a learned-style position table:

X = token_embeddings + position_embeddings

Both operands are [S, D]. The block rejects S beyond the position table instead of wrapping or silently reusing positions. Real rotary/relative systems have different limits and scaling behavior; a runtime advertising a larger context does not prove the model remains accurate there.

Input embeddings vs retrieval embeddings

Two uses of “embedding” must remain distinct:

• input token embeddings are internal rows consumed by a generative model;
• retrieval embeddings are output vectors intended for similarity/search across chunks or media.

They may have different dimensions, normalization, pooling, training objective and model identity. Do not put internal token hidden states directly into a vector database unless the model contract explicitly defines that representation.

Step 3: understand tensor contracts

A tensor is a typed, multidimensional view over storage. Its useful contract includes:

• shape: length of every axis;
• dtype: f32, f16, bf16, integer/quantized form, and so on;
• device: host CPU, CUDA device, Metal device or another accelerator;
• layout/strides: how axis indices map to storage offsets;
• contiguity/alignment;
• ownership/lifetime;
• gradients or inference-only status;
• logical semantics for every axis.

The teaching Matrix supports only owned contiguous row-major f32 with two nonzero dimensions. Its constructor uses checked multiplication, exact element count and finite-value validation:

pub fn try_new(rows: usize, cols: usize, values: Vec<f32>) -> Result<Self, MlError> {
    if rows == 0 || cols == 0 {
        return Err(MlError::ZeroDimension);
    }
    let expected = rows.checked_mul(cols).ok_or(MlError::ShapeOverflow)?;
    if values.len() != expected {
        return Err(MlError::ElementCountMismatch { /* exact facts */ });
    }
    if values.iter().any(|value| !value.is_finite()) {
        return Err(MlError::NonFiniteValue);
    }
    Ok(Self { rows, cols, values })
}

The simplicity is educational, not performant. Production tensor libraries must handle views, higher ranks, non-contiguous transposes, vectorized kernels, device memory, asynchronous execution and many dtypes.

Matrix multiplication

For:

A [M, K] × B [K, N] = C [M, N]

the inner K dimensions must match. The companion rejects:

[2, 3] × [2, 2]

with both shapes and operation name. A deep runtime error like “mat1 and mat2 shapes cannot be multiplied” is more actionable when the harness also records the model identity and semantic axes—without logging sensitive tensor content.

Broadcasting

Tensor runtimes often expand dimensions logically without copying. Adding bias [D] to [B,S,D] usually broadcasts across batch and sequence. This is powerful and dangerous:

• a missing axis may broadcast instead of fail;
• a padding mask [B,S] may need reshaping to [B,1,1,S];
• a per-head mask may accidentally share across heads;
• an image normalization vector must align with channel axis, not width.

At boundaries, assert the intended shape rather than relying on an operation to accept something.

Dtype and accumulation

Lower-precision storage reduces bandwidth and memory, but operations often accumulate in a wider type for stability. Quantized formats represent blocks using integers plus scales/metadata rather than a naive fixed number of bits per independent scalar.

Changing dtype can alter:

• memory and transfer size;
• supported kernels;
• throughput/latency;
• overflow/underflow behavior;
• output logits and sampling decisions;
• model quality.

Treat dtype/quantization as part of model release identity and run the same eval suite.

Device execution can be asynchronous

Launching an accelerator kernel may return before work completes. A host timer without synchronization can measure enqueue time rather than inference time. Data transfers, compilation, memory allocation and warm-up can dominate early requests. The performance chapter later defines correct timing boundaries.

Step 4: project query, key and value

For one attention head in this chapter:

X  [S,D]
Wq [D,D]
Wk [D,D]
Wv [D,D]

Q = XWq [S,D]
K = XWk [S,D]
Vvalues = XWv [S,D]

The projection matrices are learned weights. An intuition—not a formal semantic guarantee—is:

• query represents what this position is looking for;
• key represents how each source position can be matched;
• value represents information to mix if that source receives weight.

They are not cached documents or ordinary lookup keys/values. The network learns useful projection spaces from its objective.

The constructor requires every projection to be exactly [D,D] for this single-head design:

for (name, projection) in [
    ("query", &query),
    ("key", &key),
    ("value", &value),
    ("output", &output),
] {
    if projection.rows() != width || projection.cols() != width {
        return Err(MlError::InvalidProjection { /* ... */ });
    }
}

A general implementation can use Dk and Dv, then split/reshape into heads. Keeping one head lets us see the core arithmetic without hiding it behind four-dimensional layout transforms.

Step 5: compute scaled dot-product attention

The original Transformer defines:

Attention(Q,K,V) = softmax(QKᵀ / sqrt(Dk)) V

For every query position i and key position j, compute:

score[i,j] = dot(Q[i], K[j]) / sqrt(Dk)

The dot product produces one compatibility score. Dividing by sqrt(Dk) controls scale as the projection width grows; without it, larger dot-product magnitudes can push softmax toward extremely saturated distributions.

The score matrix is [S,S]. That is the origin of a major sequence-cost term: doubling S makes four times as many pairwise scores in full attention. Optimized kernels can avoid materializing all intermediates and improve memory traffic, but they do not make ordinary full attention linear in sequence length.

Apply the causal mask

At position i, a decoder may attend only to 0..=i. Future keys j > i must have zero probability:

row 0: [allowed, blocked, blocked]
row 1: [allowed, allowed, blocked]
row 2: [allowed, allowed, allowed]

Many implementations add negative infinity (or an appropriate very negative representable value) before softmax. The teaching code avoids creating blocked scores at all, computes softmax only over the allowed prefix, and leaves future weights initialized to exact zero.

This test is stronger than checking a triangular picture:

let prefix = block.forward(&[0, 1])?;
let extended = block.forward(&[0, 1, 3])?;

assert_eq!(prefix.logits.row(0), extended.logits.row(0));
assert_eq!(prefix.logits.row(1), extended.logits.row(1));

Changing a future token cannot alter earlier logits. That causality invariant catches mask direction, off-by-one and broadcasting errors.

Use numerically stable softmax

Naive softmax computes exp(x_i) / sum(exp(x_j)). exp(10_000) overflows in f32. Softmax is unchanged when subtracting the same constant from every logit, so subtract the maximum:

let maximum = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
for value in values.iter_mut() {
    *value = (*value - maximum).exp();
}
let denominator = values.iter().sum::<f32>();
for value in values.iter_mut() {
    *value /= denominator;
}

The companion checks inputs and denominator for finiteness. Its test with [10_000, 9_999] produces finite probabilities summing to one. In training or log-probability code, log-softmax is often preferable for numerical properties.

Mix values

After softmax, each attention row is a distribution over allowed source positions:

context [S,Dv] = attention_weights [S,S] × Vvalues [S,Dv]

The context row is a weighted sum of value vectors. An output projection maps/combines head output back to model width.

The example’s flattened attention matrix is:

[
  1.0000000, 0.0000000, 0.0000000,
  0.3608598, 0.6391402, 0.0000000,
  0.2640470, 0.2655759, 0.4703772
]

The numbers are consequences of fixed identity projections and tiny invented embeddings. Do not interpret them as explanations of language behavior. Attention weights are internal routing signals, not automatically faithful human explanations.

Step 6: multi-head attention and the rest of the block

Production transformer blocks commonly use several heads:

project X
  → reshape [B,S,H,Dk]
  → transpose/layout for kernel
  → attention per head
  → concatenate heads [B,S,H×Dv]
  → output projection [B,S,D]

Heads can learn different useful relationships, but individual head meaning is not guaranteed. Grouped-query and multi-query attention share some key/value projections across query heads to reduce KV-cache cost. Architecture details belong to the exact model config/runtime.

A decoder block also typically contains:

• normalization (LayerNorm, RMSNorm or another exact variant);
• residual connections;
• a positionwise feed-forward network/MLP, often with gated activation;
• architecture-specific positional operation such as RoPE;
• optional biases, dropout during training, mixture-of-experts routing or other features.

Pre-norm and post-norm ordering differ. Activation functions and epsilon values matter. A loader cannot replace “close enough” operations because tensor names and shapes fit.

The teaching block intentionally implements only:

X = token + position
A = causal_attention(X)
hidden = X + A
logits = hidden × lm_head

It omits normalization, MLP, multiple layers, biases, batches, training and cache. The omission is stated in code and output so a runnable example is not mistaken for a miniature accurate LLM.

Step 7: hidden states become logits

The language-model head projects model width to vocabulary:

hidden [S,D] × lm_head [D,V] = logits [S,V]

A logit is an unnormalized real score. It is not a probability and does not have to lie in [0,1]. Softmax turns one position’s logits into a distribution when needed.

For a prompt of length S, generation normally uses the final position’s [V] logits to choose the next token. The example’s last row is:

[0.2667027, 0.2802796, 1.4703772, 0.5469823]

Argmax selects ID 2 ("builds"). Argmax is deterministic greedy choice, not the general model generation algorithm. Temperature, top-k/top-p, seeds and reproducibility are the next chapter.

Some models tie the LM head to the transpose of input embeddings; others do not. The bundle/config defines this. Classification, embedding and reranker models expose different heads and pooling contracts even when their internal transformer blocks look related.

Step 8: autoregressive generation

Conceptually:

ids = encode(render(messages))
repeat:
    logits = model(ids)
    next = sampling_policy(logits[last])
    ids.push(next)
    if stop(next, ids) or budget_exhausted:
        break
return decode(ids generated after prompt)

Recomputing the entire prefix each step is wasteful. During prefill, the runtime processes the prompt and builds key/value state. During decode, it processes a new token and reuses a KV cache. The next chapter analyzes those phases, cache shapes and context costs.

The application still owns:

• maximum input/output/total token budgets;
• cancellation and deadline;
• stop token/sequences;
• structured-output validation;
• tool-call state;
• streaming semantics;
• usage accounting;
• content/safety policy;
• trace/eval identity.

The tensor runtime should not be asked to infer product policy from prompt text.

From text tokens to multimodal sequences

Transformers can operate on sequences derived from other modalities, but “token” becomes broader:

Images may be resized/cropped/normalized, divided into patches, linearly projected, and supplied with position information. A vision encoder may produce features that a learned projector maps into the language model’s width.

Audio may become waveform frames, spectrogram/mel features, codec tokens, or encoder states. Sampling rate, hop/window and padding define time alignment.

Video adds frame/clip sampling and space-time position. Naively turning every pixel/frame into attention positions is prohibitively expensive, so systems sample, compress, pool or use specialized architectures.

Generated images/audio/video often use latent, diffusion, autoregressive codec or flow-based representations with different schedules and heads. The word “token” does not imply the same tokenizer or loss as text.

At a multimodal boundary record:

source coordinates
→ deterministic decoder/normalizer
→ feature/patch/frame coordinates
→ model sequence positions
→ output/citation coordinates

Without that mapping, a model may describe evidence but the product cannot prove where it came from.

Read the Rust implementation

Owned row-major matrix

Matrix stores:

pub struct Matrix {
    rows: usize,
    cols: usize,
    values: Vec<f32>,
}

It exposes immutable rows and values, checked matmul, exact-shape add, and private positional prefix extraction. There is no unchecked indexing at public boundaries and workspace lint policy forbids unsafe code.

The naive multiplication is three loops. That makes shape semantics visible but is not suitable for large inference: it lacks SIMD-aware blocking, parallelism, accelerator kernels, lower precision and cache optimization. The benchmark lesson applies: never compare this teaching loop to a production backend as though they solve the same performance problem.

Attention owns its projections

CausalAttention validates weights once at construction, then forward:

1. validates input width;
2. computes Q/K/V projections;
3. calculates only causal-prefix scores;
4. applies stable softmax per row;
5. constructs [S,S] weights;
6. mixes values;
7. applies output projection.

Every intermediate rejects non-finite arithmetic. This is intentionally fail-fast. A production model may naturally produce very large finite values; NaN/Inf handling policy should include model identity, dtype, device and layer/stage diagnostics, then fail the request/model instance rather than emit corrupted output.

Tiny block exposes intermediates

CausalBlockOutput returns:

pub struct CausalBlockOutput {
    pub input_embeddings: Matrix,
    pub attention_weights: Matrix,
    pub hidden: Matrix,
    pub logits: Matrix,
}

This is appropriate for a teaching/debug artifact and inappropriate for a normal public inference API. Exposing every attention matrix in production can allocate O(S²) output, leak sensitive relationships, and couple the product to architecture internals. Production observability should record bounded shapes, dtypes, timings and health counters; capture tensor contents only in an explicitly protected research workflow.

Performance and capacity implications

Complexity

For full self-attention, score work/storage is associated with S×S per head. Projection and MLP work scale differently, often involving S×D². Which dominates depends on sequence, width, architecture, kernel, device and batch.

Useful application-level levers:

• reject or truncate before tokenization/model allocation according to policy;
• retrieve compact relevant evidence instead of pasting an entire corpus;
• cache static rendered/tokenized prefixes when semantics allow;
• batch compatible requests under latency and memory bounds;
• use prefill/decode-aware scheduling;
• select an architecture/context strategy fit for the workload;
• compress images/video/audio before model sequence construction while preserving eval quality;
• measure length distributions by modality and route.

Padding and batching

Batch tensors often pad sequences to a shared length. An attention mask prevents padded positions from contributing. Padding to the longest outlier wastes compute; length bucketing can help. But batching changes queue delay, device utilization, cache layout and per-request cancellation.

Never mix tenants or sensitive requests in a shared batch unless the runtime’s isolation, memory clearing and side-channel threat model permits it.

Shape telemetry

Safe bounded fields include:

• model/harness release ID (bounded known values);
• batch bucket;
• input/output token bucket;
• modality and route enum;
• dtype/device enum;
• prefill/decode duration histogram;
• cache hit/eviction/bytes;
• rejection reason enum.

Do not use raw prompt, token IDs, user ID, model output or arbitrary dimensions as metric labels.

Failure investigations

Failure 1: valid IDs, wrong vocabulary semantics

Symptoms: model runs with correct shapes but output is incoherent. The tokenizer vocabulary length matches the embedding row count.

The length check cannot detect permuted ID assignments. Compare exact tokenizer digest, template, special token IDs and golden encoding vectors against the evaluated bundle. Repair the bundle composition; do not add a prompt instruction to compensate for broken tokenization.

Failure 2: training eval is impossibly good

Symptoms: next-token accuracy is suspiciously high, then deployment fails.

Inspect the causal mask. A transposed or inverted mask may let position i read future ground-truth tokens. Add the prefix-invariance test and a direct assertion that every weight above the diagonal is zero. Also inspect label shifting: inputs and targets must be aligned according to the training objective.

Failure 3: attention contains NaN

Work backward:

1. identify first stage/layer producing non-finite values;
2. check input normalization and corrupt media;
3. verify dtype/quantization/kernel support;
4. inspect mask values and rows where every key is masked;
5. confirm max-subtracted softmax/log-softmax;
6. check model weights and runtime version;
7. compare CPU/higher-precision reference on a minimal frozen case.

Do not merely replace NaN with zero. That hides a broken numerical contract and can produce plausible but untrustworthy output.

Failure 4: changing a future token changes earlier logits

Possible causes:

• mask direction reversed;
• off-by-one allowed i+1;
• mask broadcast across wrong axis;
• reused output buffer/state from a later position;
• non-causal architecture used under a causal assumption.

Reduce to a three-token deterministic case like the companion, print only small test matrices, then compare row by row. Keep the invariant test independent of model quality.

Failure 5: context fits advertised limit but OOMs

The advertised maximum may assume a particular batch, KV dtype, attention implementation, offloading plan or device. The request may also include hidden template/tool/media tokens not counted by raw user text.

Record exact rendered count, batch, model/runtime/device, cache dtype, concurrency and memory peaks. Admission must use a measured capacity envelope, not only maximum_context_tokens.

Failure 6: GPU timing looks impossibly fast

The host timer ended after enqueue while kernels were still running. Synchronize at the measurement boundary or use runtime/device events correctly. Separate cold load, warm-up, prefill, time-to-first token, per-token decode and end-to-end latency.

Failure 7: quoted source span points into the wrong characters

Token offsets were measured against normalized text but applied to original UTF-8 bytes, or character indices were treated as byte offsets. Retain alignment metadata and explicitly convert coordinate systems. Test combining characters, emoji sequences, non-Latin scripts and normalized forms.

Security boundaries

• Tokenization is attacker-controlled CPU/memory work; bound bytes, tokens and batch before expensive allocation where possible.
• Template expansion and tool schemas can multiply tokens; count the complete rendered input.
• Extremely long repeated strings, unusual Unicode and media dimensions are denial-of-service cases.
• Token IDs/embeddings can still encode sensitive content; do not classify them as anonymous.
• Shape/dtype parsers use checked arithmetic before allocation.
• Custom tokenizer/model code is code execution; pin, review and isolate it.
• Avoid returning raw logits/hidden states unless the product explicitly requires and protects them.
• Cross-request KV/prefix caches need tenant/policy-aware keys and lifecycle clearing.
• A causal mask protects model semantics, not authorization. The model can still repeat earlier unauthorized context if the harness retrieved it.

Exercises

Exercise 1: add batched matrix semantics

Introduce a rank-3 Tensor3 with [B,S,D], checked offset arithmetic and a batched projection. Do not add implicit broadcasting. Implement explicit bias addition and tests where [D], [1,1,D] and [B,S,D] are deliberately distinguished.

Exercise 2: padding mask plus causal mask

Batch sequences of lengths two and four into [2,4,D]. Combine padding and causal rules so:

• queries cannot attend future positions;
• no query attends padded keys;
• policy for padded query rows is explicit and never creates all-negative softmax NaNs;
• unpadded outputs match individually evaluated sequences within a stated tolerance.

Exercise 3: two-head attention

Use D=4, H=2, Dk=2. Implement reshape/transpose/concatenate explicitly. Write every intermediate shape and a test that swapping head and sequence axes fails a semantic shape wrapper even if raw element counts match.

Exercise 4: tokenizer golden suite

Using the Rust tokenizers crate and a pinned fixture, store pieces, IDs, offsets, special mask and decoded output for:

• ASCII plus punctuation/whitespace;
• composed/decomposed Unicode;
• emoji with joiners;
• code indentation;
• an unknown or byte-fallback case;
• rendered system/user/tool messages.

Changing tokenizer/template identity must require an intentional golden review.

Exercise 5: quantify sequence growth

For S = 512, 2,048, 8,192, calculate score element counts per layer/head and raw bytes for f32, f16 and bf16 as if materialized. Then explain why an optimized attention kernel’s memory behavior can differ while arithmetic/sequence capacity still demands measurement.

Exercise 6: add RMS normalization and gated MLP

Implement a clearly specified RMSNorm with epsilon and a small SwiGLU-like feed-forward path. Test zero/large inputs, non-finite values, exact shapes and residual ordering. Compare your equations to one pinned architecture config; do not label the result “universal Transformer.”

Solution guidance

1. A rank-specific type should store dimensions and checked contiguous storage, expose semantic methods, and require explicit conversion/reshape. Broadcasting should be a named operation with exact accepted shapes.
2. Construct an allowed boolean mask first. For rows representing padding, either skip output or define a safe sentinel path before softmax. Test the combined mask directly and compare real token rows to unbatched evaluation.
3. Project [S,4], reshape [S,2,2], transpose to [2,S,2], attend per head, reverse transpose and concatenate to [S,4]. Semantic wrapper types prevent an accidental [S,H,Dk]/[H,S,Dk] confusion.
4. Pin tokenizer bytes in the bundle and assert arrays plus coordinate convention. A decoded string alone misses ID/special/offset regressions.
5. S² is 262,144, 4,194,304, and 67,108,864 elements respectively, per conceptual head/layer score matrix. Multiply by dtype bytes and then by relevant heads/layers only for the deliberately materialized thought experiment; real kernels may tile/recompute/fuse.
6. State epsilon, dtype/accumulation, width and ordering. Compare to reference runtime outputs on a tiny deterministic tensor within justified tolerance.

Check your understanding

1. Why can two tokenizers with the same vocabulary size still be incompatible?
2. What does row 4 of an input embedding matrix represent?
3. Why does QKᵀ have shape [S,S]?
4. What invariant distinguishes causal from bidirectional self-attention?
5. Why subtract the maximum before exponentiating logits?
6. Is an attention weight a probability that a source statement is true?
7. Which logits are normally used to select the first generated token after a prompt?
8. Why can a quantized checkpoint occupy more runtime memory than its file size?
9. How do padding mask and causal mask answer different questions?
10. What must be retained to map a token-level citation back to original Unicode bytes?

Answers:

1. Token strings and ID assignments can differ; the same ID then selects a different embedding row.
2. The hidden vector at sequence position four—not necessarily token ID four.
3. Each of S queries is compared with each of S keys.
4. Earlier outputs are invariant to changes in later input positions.
5. Softmax is shift-invariant and max subtraction prevents exponential overflow.
6. No. It is a normalized internal routing coefficient, not calibrated truth or faithful explanation.
7. The final prompt position’s vocabulary logits.
8. Dequantization, host/device copies, KV cache, activations, workspaces and allocator overhead.
9. Causal mask blocks future positions; padding mask blocks non-content positions.
10. Tokenizer identity plus aligned offsets and an explicit byte/scalar/grapheme coordinate system, along with original or normalized source as appropriate.

Completion checklist

• tokenizer and template bytes are pinned with the model;
• golden encodings include IDs, offsets, special tokens and Unicode boundaries;
• token count covers complete rendered context, not raw user text;
• every tensor boundary names shape axes, dtype, device and layout;
• dimension arithmetic and allocations use checked integers;
• invalid IDs, shapes, non-finite values and context overflow fail explicitly;
• causal and padding masks have direct invariant tests;
• softmax/log-softmax is numerically stable for the chosen dtype/path;
• application distinguishes input embeddings, hidden states and retrieval embeddings;
• LM-head output shape and selected generation position are verified;
• runtime memory is measured beyond file/parameter size;
• async device timing synchronizes correctly;
• logs/metrics contain bounded metadata, not raw tensors or sensitive tokens;
• the production backend is compared to a small reference case;
• the limitations of any teaching implementation are visible in code and output.

References

• Vaswani et al., Attention Is All You Need
• Kudo and Richardson, SentencePiece
• Hugging Face Tokenizers pipeline
• Hugging Face Tokenizer API, including Rust examples
• PyTorch Softmax reference
• Candle core tensor/runtime documentation

Next

We now understand one full forward pass and the vocabulary logits it emits. The next chapter makes greedy, temperature, top-k and nucleus selection explicit, then separates a seeded sampler replay from the much stronger—and often impossible—claim of universal model determinism.