# Lab 01 — Attention and masks

A CPU-only companion lab for lessons 4–6 of the LLM Internals track. It implements
single-head and multi-head attention from scratch, checks tensor shapes, stable softmax,
causal future-token isolation, and numerical parity with PyTorch's reference attention.
It also contains one deliberately wrong causal mask, a diagnostic that catches it, and
the repair.

No model download, API key, GPU, dataset, or network access is required.

## Prerequisites

Read lessons 4–6 first:

- `self-attention-in-pytorch` — Q/K/V projections, scaled scores, stable softmax.
- `causal-attention-masks` — lower-triangular causal mask, key-padding mask, AND combination.
- `multi-head-attention` — head split, parallel attention, merge, output projection.

You need basic Python and enough PyTorch to create tensors and call `matmul`.

## Tested environment

Recorded from the machine that produced `HANDOFF.md`:

- macOS (Apple silicon), CPU only; `torch.cuda.is_available()` is `False`
- Python 3.12.10
- PyTorch 2.15.0 (CPU build)
- pytest 9.0.2

The code uses only stable PyTorch APIs and should work on other recent CPU builds, but the
exact tolerances below were measured on this environment.

## Exact setup

From this directory:

```bash
python3.12 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r requirements.txt
```

`requirements.txt` pins:

```
torch==2.15.0
pytest==9.0.2
```

On Linux or Windows, to guarantee a CPU-only wheel and avoid CUDA packages:

```bash
python -m pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/cpu
python -m pip install pytest==9.0.2
```

On macOS the default PyPI wheel is CPU-only.

## Run the tests

```bash
python -m pytest -q
```

Expected: `28 passed`. The test suite is deterministic: all random tensors use explicit
seeds or `torch.Generator` objects.

Run the exercise on its own:

```bash
python exercises/wrong_mask.py
```

## Files

| Path | Purpose |
| --- | --- |
| `attention.py` | `stable_softmax`, `scaled_dot_product_attention`, `causal_mask`, `combine_causal_and_padding`, `SelfAttention`, `MultiHeadSelfAttention` |
| `exercises/wrong_mask.py` | Upper-triangular "causal" mask, future-leak diagnostic, and the correct lower-triangular replacement |
| `tests/test_shapes.py` | Shape contracts for single-head and multi-head paths, mask broadcasting, explicit shape errors |
| `tests/test_stable_softmax.py` | Parity with `torch.softmax`, overflow resistance, fully masked rows |
| `tests/test_causal_isolation.py` | Zero future weights and behavioural isolation under value/key perturbation; padding masks |
| `tests/test_parity_torch.py` | Parity with `F.scaled_dot_product_attention` and `nn.MultiheadAttention` |
| `tests/test_wrong_mask.py` | Proves the wrong mask passes shape checks, fails the leak test, and that the fix passes both |

## Implementation conventions

- Activations are `[..., T, d_head]`; the leading dimensions are `[B]` for single-head and
  `[B, h]` for multi-head.
- Masks are boolean and broadcastable to `[..., T, T]`. `True` means **keep / may attend**.
  This matches lessons 4–6 and PyTorch's boolean `attn_mask` semantics.
- Scaling defaults to `1 / sqrt(d_head)`, not `1 / sqrt(d_model)`.
- Multi-head uses one combined `nn.Linear(d_model, 3 * d_model)` in-projection and one
  `nn.Linear(d_model, d_model)` output projection, both `bias=False` by default.
- A query row with **no allowed keys** produces all-zero attention weights and an all-zero
  output. It never produces NaN.

  The lab's `stable_softmax` deliberately maps fully‑masked rows to zero — a design choice
  so that a query with no visible keys contributes nothing to the output, rather than
  poisoning the batch with NaN. Raw `torch.softmax` does return NaN when a row is entirely
  ``-inf`` (0/0), and PyTorch's SDPA reference kernel can also return NaN for fully masked
  rows depending on the backend. On CPU with PyTorch 2.15.0 and the MATH backend
  (`torch.nn.attention.SDPBackend.MATH`), the attention weights are stabilized via max‑subtraction
  before ``exp``, so NaN does not arise. Fused CPU kernels may accumulate differently; the
  parity tests use the MATH backend explicitly via ``sdpa_kernel(backends=[SDPBackend.MATH])``
  to ensure deterministic comparison.

  That the lab returns zero (not NaN) for fully masked rows is a deliberate pedagogical choice
  and is tested separately in ``test_causal_isolation.py``::``test_fully_masked_row_is_zero_not_nan``.

## What the tests establish

| Test group | Check | What a failure means |
| --- | --- | --- |
| Shapes | `Q/K/V`, scores, weights, output have the documented shapes; bad widths raise `ValueError` | Projection, transpose, split, or merge is wrong |
| Stable softmax | Matches `torch.softmax` on ordinary logits; finite on `[1000, 1001, 999, 1002]`; zeros for an all-`-inf` row | Naive `exp` overflowed, or fully masked rows were not handled |
| Causal isolation | Upper triangle is exactly zero; changing a future value or key leaves earlier outputs bit-identical | The mask leaks future information |
| Padding | Padded key columns get zero weight; padded queries still attend to real past keys; fully masked row is zero, not NaN | Padding mask is missing, inverted, or OR-ed instead of AND-ed |
| PyTorch parity | Max difference against `F.scaled_dot_product_attention` (MATH backend) and `nn.MultiheadAttention` | Ordering, scaling, head split/merge, or projection layout is wrong |

Measured maximum absolute differences on the tested environment (float32, CPU):

| Comparison | Max abs difference |
| --- | --- |
| Single-head, non-causal vs SDPA | `2.38e-07` |
| Single-head, causal vs SDPA `is_causal=True` | `1.19e-07` |
| Boolean padding mask vs SDPA `attn_mask` | `1.19e-07` |
| Multi-head vs `nn.MultiheadAttention` (copied weights) | `8.94e-08` |
| Stable softmax vs `torch.softmax` on the large-logit example | `0.0` |

These are numerical-agreement measurements, not performance measurements.

## Exercise: the mask that looks right

`exercises/wrong_mask.py` builds the causal mask with `torch.triu` instead of `torch.tril`.
Every tensor still has shape `[T, T]`, so every shape assertion passes. The model can see
future tokens.

1. Run `python exercises/wrong_mask.py`. It prints that both masks have shape `(4, 4)`,
   then reports 6 future-token weights for the buggy mask (largest `0.740962` in the tested
   run) and 0 for the fixed mask.
2. Run `python -m pytest tests/test_wrong_mask.py -q`. The diagnostic
   `assert_no_future_leak` raises `AssertionError: buggy mask: ... future-token weights ...`.
3. Open `exercises/wrong_mask.py`, change `buggy_causal_mask` to return
   `torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))`, and rerun both commands.
   The leak count becomes 0 and the repaired mask also matches PyTorch's `is_causal=True`
   reference.

Why the fix works: a decoder query at position `t` may attend only to keys `j <= t`.
`torch.tril` keeps exactly that region and masks every `j > t` with `-inf` before softmax,
so those weights become exactly zero.

Why the diagnostic is necessary: shape checks and row-sum checks both pass for the wrong
mask. Only a behavioural check — "is any weight placed on a key after the query?" — detects
it. The same diagnostic also catches a missing mask, an inverted mask, and an OR-combined
padding mask.

## Assumptions and limits

- The lab teaches the mechanism, not a production kernel. It uses explicit matmuls, not
  FlashAttention, and makes no speed or memory claims.
- Parity is checked against the deterministic MATH backend. Fused CPU kernels can differ
  in accumulation order; the `1e-6` and `2e-5` thresholds are for this reference path.
- The fully masked row behaviour is a deliberate safety choice. It is not a claim about how
  every framework handles such rows.
- No training loop, dropout, positional encoding, or KV cache is included. Those belong to
  other labs.
- Padding semantics follow lessons 4–6: `True` = real token. Some libraries use the
  opposite convention.

## Lesson 5 correction

Lesson 5 (`causal-attention-masks`) previously printed an inverted causal-mask table in
§1.3 (upper triangle) while defining the correct lower-triangular formula in §1.1.
That table has now been corrected to the lower-triangular form, matching this lab and
the lesson's own formula. The lab was written against the formula, so no lab code
changed.

## Source support

- Vaswani et al., "Attention Is All You Need" (arXiv:1706.03762), §3.2.1–3.2.2:
  scaled dot-product attention and multi-head concatenation.
- PyTorch documentation for `torch.nn.functional.scaled_dot_product_attention`:
  query/key/value shapes, boolean `attn_mask` semantics (`True` = keep), `is_causal`,
  and `scale`.
- PyTorch documentation for `torch.nn.MultiheadAttention`: `in_proj_weight` layout and
  additive masks.