Assignment 1: Attention Is All You Need

What is this assignment about?

This assignment is inspired by how modern Large Language Models (LLMs) such as chatbots work.

An LLM is trained to predict the next token in a sequence. Conceptually, given a prompt

the model produces a probability distribution for the next token $\texttt{token}_n$.

Think of it like autocomplete: based on what you typed so far, what is most likely to come next?

What a real LLM would do, at a high level

You do not need the exact ML details, but a real LLM roughly:

1. tokenises text into token IDs
2. looks up token embeddings
3. runs many Transformer layers
4. uses attention inside those layers

5. produces scores over a vocabulary

6. chooses the next token, appends it, and repeats

What this assignment does instead

You are not implementing a full chatbot. You are implementing a core part used inside them: a simplified single-head self-attention block plus a final KV-cache generation step.

If you want the plain-English picture of what attention is doing before you get into the loops, read the Attention explainer.

To keep the programming focused and testable:

• you begin with a small Stage 1 representation exercise that turns a token list into one-hot vectors
• the later attention stages use the provided prompt and generated embeddings directly as numeric rows
• you are given the parameter matrices Wq, Wk, and Wv directly as numbers
• you compute attention scores, weights, and outputs exactly
• you apply masking rules and stable softmax
• you do not predict English words

If the maths still feels abstract, that is fine. The implementation requirements are stated explicitly below, and the linked explainer pages fill in the intuition when you need it.

Where the numbers in this assignment come from

A real language model starts with text, not matrices. Very roughly, the pipeline is:

1. split the text into tokens
2. turn each token into an embedding vector
3. run attention over those vectors to build context-aware outputs
4. use the final vector to help predict what token should come next

This assignment still simplifies a lot, but it lets you practise the representation idea separately from the later attention maths. You do not implement a full tokenizer, embedding-table lookup, vocabulary-scoring layer, or training process.

The project therefore has two related but separate “before attention” pieces:

• a short token list used in Stage 1, where you build and print a one-hot embedding table
• the prompt and generated embedding rows used by the later attention stages

That means Stage 1 is a warm-up representation exercise on tokens, while Stages 2 to 6 operate on the provided numeric embedding matrices.

That means:

• each row of the prompt matrix is one token embedding
• n is the number of prompt tokens
• d is the number of components in each embedding vector
• the prompt positions still have an order, so “earlier token” and “later token” matter

The core job of attention is then: for one position, look back at the earlier positions that are allowed, decide which ones matter most, and blend their information into a new output row.

The whole job in one pass

Your program follows one continuous pipeline:

1. read the token list, mask, embedding rows, and parameter matrices from the input file
2. build and print the Stage 1 one-hot embeddings from the token list
3. project the provided prompt embeddings into query, key, and value matrices
4. compare prompt queries against prompt keys to get attention scores
5. turn those scores into attention weights with masking and stable softmax
6. use the weights to blend the value vectors into prompt outputs
7. extend the same logic to generated tokens while reusing cached keys and values

If you keep that pipeline in mind, the stage descriptions later on should read as one algorithm being built step by step, not as unrelated tasks.

Glossary and notation

This section collects the vocabulary and notation that help you read the assignment, before you get into the stage-by-stage implementation details.

Getting started

This section is the fast setup pass: what you have, what the scaffold already does, and what constraints matter before you start coding.

If you need more intuition while coding, use the linked explainers at the point where they become relevant: one-hot encoding for Stage 1, tokens and embeddings for the input story, projections for Stage 2, dot products and masking for Stage 3, softmax for Stage 4, attention output for Stage 5, and KV cache for Stage 6.

If reading a1.c itself is slowing you down, use the Annotated scaffold. The function pulls in the stage sections below jump straight to the matching scaffold section there.

Stage descriptions

The most useful way to read this section is as one algorithm unfolding in order. Stage 0 loads the token list and numeric snapshot. Stage 1 builds and prints one-hot embeddings. Stage 2 turns the provided prompt embeddings into Q, K, and V. Stages 3 to 5 perform attention over the prompt. Stage 6 then applies the same ideas during generation using the cache.

Each stage names the scaffold function you should implement first when there is one. Stage 1 is the small embedding-construction stage that now sits in front of the attention pipeline.

Stage 0: Reading the input

If you want to see what a real input file looks like before reading the format rules, open What an input file looks like.

Implement this scaffold function first: read_input(...) void read_input(int *n, int *d, int *g, int *text_len, char embedding_table[MAX_TEXT_SIZE][MAX_TOKEN_LENGTH + 1], int mask[MAX_TOKENS], double prompt[MAX_TOKENS][MAX_D], double gen[MAX_GEN][MAX_D], double wq[MAX_D][MAX_D], double wk[MAX_D][MAX_D], double wv[MAX_D][MAX_D]); Open in the annotated scaffold ↗ . Every later stage depends on storing the input snapshot correctly.

Read from stdin in this order:

1. n d g text_len. The three attention sizes, followed by the number of tokens in the Stage 1 text.
2. text_len whitespace-separated tokens. This is the Stage 1 text that you turn into the one-hot embedding table.
3. n integers for mask. These say which prompt positions are real and which ones are padding.
4. n prompt embeddings of length d. These rows are the numeric prompt matrix used by the later attention stages.
5. g generated embeddings of length d. These rows are used later in the generation stage.
6. d lines for Wqthe query projection matrix.
7. d lines for Wkthe key projection matrix.
8. d lines for Wvthe value projection matrix.

Bounds

Use scanf. Do not open files in your C program.

Whitespace-separated input means spaces and newlines are equivalent for scanf.

Stage 1: Creating embeddings

This stage is the small representation exercise before the attention maths begins. You take the input text, treat each whitespace-separated piece as a token, build the unique-token list, sort it lexicographically, and then print one one-hot vector per token in that sorted list.

The scaffold function for this stage is create_embeddings(...) int compare2Dchar(const void *a, const void *b); void create_embeddings(char embedding_table[MAX_TEXT_SIZE][MAX_TOKEN_LENGTH + 1], int text_len); Open in the annotated scaffold ↗ . If you use qsort(...), the comparison helper sits right beside it in the scaffold.

The core job in this stage is:

1. sort the token array stored by Stage 0 lexicographically

2. remove duplicate tokens from the sorted array
3. give each token a vector that is all zero except for one 1 based on the index in the sorted array
4. print the token together with its one-hot embedding as per the Output Format

That is the one-hot encoding idea. It is a simple bridge from categories to numbers, and the one-hot explainer shows it visually.

If this stage makes sense in principle but you have not really worked with C strings yet, use Strings in C. It is meant as a short crash course on the string operations this stage relies on.

This stage is the simple practice version of token representation.

You can use either qsort(...) or your own sorting code for the lexicographic ordering step.

algorithms, best, foundations, of, subject, the, to, welcome

Stage 1: Create Embeddings
"algorithms" -> (1 0 0)
"subject" -> (0 1 0)
"the" -> (0 0 1)

Stage 2: Projections

In an LLM once the prompt text has been tokenised, it will then be represented as a sequence of embedding vectors.

In the stage, we provide you a sequence of embedding vectors corresponding to a tokenised prompt, and not rely on the previous stage to keep it seperate. If we did use the previous stage, for a given prompt, you would look up the embedding for each token, and pass an array of those embeddings to this stage. Notably the embedding for every token is the same regardless of where in the prompt it is, despite the location of words, in relation to others, altering the meaning. This is where attention comes in!

Once the input is stored, our first computation to implement attention is to build the prompt query, key, and value matrices .

Before reading on make sure you are familar with the notation

The input prompt embeddings are the rows of a matrix $\mathbf{X}$. To get queries, keys, and values, you project the same prompt matrix three times with three different matrices:

What Q, K, and V are, and why are there three of them.

Implement the scaffold function for this stage: compute_projection(...) void compute_projection(int count, int d, double src[MAX_TOKENS][MAX_D], double w[MAX_D][MAX_D], double dest[MAX_TOKENS][MAX_D]); Open in the annotated scaffold ↗ .

If $\mathbf{X}$ is the input matrix, $\mathbf{W}$ is one of the projection matrices, and $\mathbf{Y}$ is the projected output matrix, then:

Be careful to reset the sum to 0 for each output component.

This function is called three times to compute prompt $\mathbf{Q}$, $\mathbf{K}$, and $\mathbf{V}$

Stage 3: Attention scores (prompt)

With $\mathbf{Q}$ and $\mathbf{K}$ available, the next step is to measure how strongly each key vector corresponding to prompt token $j$ should attend to each query vector corresponding prompt token $i$. For this task you will be computing the $\mathbf{score}$ matrix, using the projections from the previous stage.

You should implement the scaffold function: compute_attention_scores_or_weights_prompt(...) void compute_attention_scores_or_weights_prompt( int n, int d, const int mask[MAX_TOKENS], double q[MAX_TOKENS][MAX_D], double k[MAX_TOKENS][MAX_D], double scores_or_weights[MAX_TOKENS][MAX_TOKENS], int apply_softmax); Open in the annotated scaffold ↗

For each query vector corresponding to prompt token $i$ ($\vec{\mathbf{Q}_i}$) and key vector corresponding to prompt token $j$ ($\vec{\mathbf{K}_j}$), we compute an attention score. The score is a scaled dot product:

When computing scores, you should apply two different masks:

• causal mask: if j > i, mask that position by setting the score to $-\infty$
• padding mask:
  • if mask[i] == 0, the whole score row becomes $-\infty$.
  • if mask[j] == 0, mask that position by setting the score to $-\infty$

In C, the score of $-\infty$ will be stored as -INFINITY using the INFINITY constant defined in <math.h>.

Stage 4: Attention weights (prompt)

The Stage 3 scores are not the final answer. They must be normalised into weights.

For this stage, you should modify the scaffold function from stage 3 compute_attention_scores_or_weights_prompt(...) void compute_attention_scores_or_weights_prompt( int n, int d, const int mask[MAX_TOKENS], double q[MAX_TOKENS][MAX_D], double k[MAX_TOKENS][MAX_D], double scores_or_weights[MAX_TOKENS][MAX_TOKENS], int apply_softmax); Open in the annotated scaffold ↗ so that it can apply softmax.

Hence, when apply_softmax is 1 (this stage), you should take the scores and apply stable softmax for each row in the output matrix.

If a row has no unmasked columns, output 0.0 for each weight in that row, do not divide by zero!

(2, 0, -1)

(0.84, 0.11, 0.05)

Stage 5: Attention output (prompt)

Once you have the attention weights, you can compute the actual prompt outputs. Implement this scaffold function: compute_attention_output_prompt(...) void compute_attention_output_prompt( int n, int d, double weights[MAX_TOKENS][MAX_TOKENS], double v[MAX_TOKENS][MAX_D], double out[MAX_TOKENS][MAX_D]); Open in the annotated scaffold ↗ .

Compute the attention output vectors as a matrix $\mathbf{output}$ for each prompt token using:

Note that $\overrightarrow{(\mathbf{V}^T)}_j$ just means, take the vector corresponding to column $j$ of matrix $\mathbf{V}$. Writing $\overrightarrow{\mathbf{V}}_j$ would refer to row $j$. Doing the transpose first, turns the columns into rows, so we can extract row $j$ to get our vector corresponding to the column.

This is a weighted sum of value vectors using the attention weights you just computed.

Stage 6: Generated outputs with a KV cache

In this task you will be using the prompt and the generated tokens (provided as part of the input) and a KV cache to compute attention outputs. You will not be generating the tokens as these are provided.

The final stage applies the same attention logic during generation, but now the program reuses past keys and values through the cache.

Implement this scaffold function: compute_generation_with_cache(...) long compute_generation_with_cache( int n, int d, int g, int t, const int mask[MAX_TOKENS], double gen[MAX_GEN][MAX_D], double wq[MAX_D][MAX_D], double wk[MAX_D][MAX_D], double wv[MAX_D][MAX_D], double k_cache[MAX_TOKENS + MAX_GEN][MAX_D], double v_cache[MAX_TOKENS + MAX_GEN][MAX_D], double output[MAX_D]); Open in the annotated scaffold ↗ .

For each generated token $t$ which follows $n$ prompt tokens, you compute:

1. $\vec{\mathbf{Q}}_{n+t}$: the query vector for the token numbered $n+t$.
2. $\vec{\mathbf{K}}_{n+t}$ and $\vec{\mathbf{V}}_{n+t}$: the key and value vectors for the token numbered $n+t$, which should be appended to the KV cache.
3. score the new query $\vec{\mathbf{Q}}_{n+t}$ against all cached keys $\vec{\mathbf{K}}_{i}$ for $0 \le i \le n+t$, remembering to mask any prompt cache entries ($0 \leq i < n$) where mask[i] == 0 (same as one row of stage 3)
4. apply stable softmax over the cached scores to turn them into weights (same as one row of stage 4)
5. compute the output as a weighted sum of cached values (same as one row of stage 5)

There is no separate future-token causal mask in Stage 6 because the cache only contains past items plus the current generated token. The current generated token is included, so it can attend to itself.

Dot Products Computed

You must also report the exact number of dot products computed for each generated token.

In this stage, define “a dot product computed”, as computing the dot product between any two vectors (which includes any matrix multiplication between a given row and column of a matrix), for any weight matrix, and any unmasked prompt or any generated token embedding or projection vectors.

That means:

• For each generated token, you should compute each unmasked score exactly once (store it in an array, reuse it for both max and denominator).
• The total count you return should equal the total number of unmasked cached key positions processed across all generated tokens.
• If a prompt cache entry is masked (mask[j] == 0), it contributes 0 to the count.

Testing, Marking, and Output Rules

This section collects the parts of the spec that are directly useful while checking your work: testing, common failure modes, marking, and the exact output contract.

clang -Wall -Wextra -Wno-newline-eof -Wno-unused-parameter -pedantic -std=c17 -lm -o a1 a1.c

./a1 < test_data/test0.txt

./a1 < test_data/test0.txt > myout.txt
diff -u test_data/test0-out.txt myout.txt

Stage 1: Create Embeddings
Stage 2: Projections
Stage 3: Attention Scores (Prompt)
Stage 4: Attention Weights (Prompt)
Stage 5: Attention Output (Prompt)
Stage 6: Generated Outputs

Q Projection:
<n lines of d numbers>
K Projection:
<n lines of d numbers>
V Projection:
<n lines of d numbers>

Gen 0: <d numbers>
Dot products computed: <an integer>
Gen 1: <d numbers>
Dot products computed: <an integer>
...