Tensor Parallelism with CUDA - Standard Attention

Tensor Parallelism from Scratch

Intro:

In the last few posts, I started building up from the most basic ideas: what a tensor is, why GPUs are a good fit for tensor operations, how matrix multiplication works, and how we can start splitting that work across multiple GPUs. If you missed those, I’d recommend starting with them, because this article is where the math starts to look a little more like the stuff people actually talk about when they talk about LLMs. We buildin’ baby! For reference, this is the first article:

Tensor Parallelism with CUDA - Intro

For those just jumping into the series, I came to machine learning from a bioinformatics / HPC background. I don’t have decades of deep learning research experience, but I do know about GPUs, CUDA, C++, and over a decade of experience operating and writing applications for HPC clusters at scale. So the point of this series is not “let’s use PyTorch and call .attention().” The point is to peel back the abstraction and determine what is actually happening.

Standard Attention:

The 2017 paper Attention Is All You Need introduced the Transformer architecture and replaced recurrence / convolution with attention as the primary sequence-modeling mechanism. The paper describes attention as mapping a query plus key-value pairs to an output, where the output is a weighted sum of values and the weights come from query-key compatibility. In this article, we will implement the plain, naive, single-GPU version of scaled dot-product attention:

\(O = softmax ( QK^T / \sqrt{d} ) V\)

This may look complicated, but it is just:

1. A matrix multiplication

2. a softmax

3. another matrix multiplication

That’s it. Well, “that’s it” in the same way a GPU is “just a bunch of cores”. Technically true, but a lot hiding in the details.

The implementation I will show is intentionally naive. I decided not to use any CUDA core libraries such as cuBLAS, I avoided kernel fusions, memory tiling, etc. The whole point was to write the textbook version first so we can see exactly how the textbook implementation maps to the GPU (what is slow, what is memory hungry, where we could fuse kernels), to show why later optimizations like FlashAttention are such a big deal.

Buckle up and get ready to dive in!

Why Attention?

Let’s start with the problem attention is solving.

The starting step of LLM inference is the “tokenization” of a sequence. From this, we generate a list of tokens from an input prompt. A token might be a word, part of a word, punctuation, whitespace, whatever the “tokenizer” decides. Each token gets represented as a vector of numbers.

So a sentence like, “The cat slept” becomes token IDs: [791, 8415, 22341], which then becomes a list of vectors:

• token 0 → [0.12, -0.44, …]

• token 1 → [0.03. 0.81, …]

• token 2 → [-0.2, 0.19, …]

Attention provides a way for tokens to learn from other tokens, something commonly referred to as context.

For example, in the sentence:

“The animal didn’t cross the street because it was tired.”

What does, “it” refer to?

A model needs some way for the token “it” to look around the sequence and decide which previous tokens are relevant. Attention is the mechanism that lets every token ask:

“Which other tokens should I care about, and how much should I care about them?”

The “how much” part is important. Attention does not usually make a hard decision like “only look at token 2.” Instead, it creates a weighted average over all tokens. That means a token can attend:

• 70% to “animal”

• 20% to “tired”

• 10% to everything else

This is hand-wavy, but it’s the right mental model to take forward as we appreciate the math. Essentially, attention is a learned information-routing system.

Q, K, and V

In attention, each token carries three vectors:

• Q: query

• K: key

• V: value

The easiest way I’ve found to think about these:

• Q: What am I looking for?

• K: What do I currently contain?

• V: What do I pass along if someone pays attention to me?

If token ‘i’ wants information, it uses its query vector. If token ‘j’ might contain useful information, it exposes a key vector. The model compares query ‘i’ against key ‘j’. If they match well, token ‘i’ pays more attention to token ‘j’. Then, token ‘i’ pulls information from token ‘j’s value vector.

So:

• query asks

• key matches

• value contributes

We’ve discussed attention in plain English, so now we move to math. Sorry.

Scary Math Notation Like I’m 15

Let’s jump into the actual math behind the kernels by beginning to define some terms.

Let:

• N = sequence length

• d = head dimension

For this example:

• N = 4096

• d = 128

That means we have 4096 tokens, and each token has query/key/value vectors of width 128. So our matrices look like this:

\(\begin{align} Q∈R^{N×d}\\ K∈R^{N×d}\\ V∈R^{N×d} \end{align}\)

Let’s revisit these notations in case you missed previous posts:

Q ∈ R ^ (N x d)

Q is a matrix of real numbers with N rows and d columns. Each row is one token, each column is one feature in that token’s query vector.

The same applies for K & V.

These yield an output of:

\(O∈R^{N×d}\)

One output vector per token.

Then, the full attention equation is:

\(O = softmax (QK^T / \sqrt{d})V\)

This looks scarier than it is. We’ll break it apart:

Step 1: QKᵀ

The first thing we compute is:

\(S = QK^T / \sqrt{d}\)

Ignore the divide by sqrt(d) for one second.

What is QK^T?

Well, Q is:

[N, d]

K is also:

[N, d]

But we transpose (turn on its side) K, giving:

K^T = [d, N]

So:

Q      x  K^T    = S
[N,d]  x [d,N]   = [N,N]

The result is an N x N matrix.

This is our score matrix.

Each element:

\(S_{ij}\)

means:

how much query token i matches key token j

So row i of S contains token i’s attention scores against every token in the sequence.

That means:

S[0, :] = how much token 0 cares about every token
S[1, :] = how much token 1 cares about every token
S[2, :] = how much token 2 cares about every token
...

This is why attention gets expensive.

If N = 4096, then S has:

4096 x 4096 = 16,777,216 values

And that’s for one attention head.

Have fun at long context lengths (we see 1M in modern models).

What is a Dot Product?

Each score in S is a dot product.

This kernel is here.

Given two vectors:

q = [1, 2, 3]
k = [4, 5, 6]

The dot product is:

q · k = 1(4) + 2(5) + 3(6)
      = 4 + 10 + 18
      = 32

In attention:

q = query vector for one token
k = key vector for another token

If the two vectors point in a similar “direction,” the dot product tends to be larger. If they do not align, the dot product tends to be smaller.

That makes dot product a simple compatibility score.

This is powerful because the model learns how to create Q and K. We are not manually saying “this word relates to that word.” The model learns projections that produce query and key vectors where useful relationships hopefully line up.

So when we calculate:

\(QK^T\)

we are calculating every query-key match in the sequence.

Every token asks:

how much do I care about token 0?
how much do I care about token 1?
how much do I care about token 2?
...

And the result is the N x N score matrix.

Why Divide by sqrt(d)?

Now let’s bring back the scale:

\(S = QK^T / \sqrt{d}\)

Why are we dividing by sqrt(d​)?

The original Transformer paper makes this point:

dot products can grow large in magnitude as the key/query dimension gets larger. When those values get too large, the softmax can get pushed into regions with tiny gradients, which makes learning harder. Their fix is to scale the dot products by sqrt​(dk)​.

Let’s say d = 128.

Then:

sqrt(128) ≈ 11.31

So instead of feeding raw dot products into softmax, we calm them down a bit.

This matters because softmax is very sensitive to big differences.

For example:

softmax([1, 2, 3])

is one thing.

But:

softmax([10, 20, 30])

is basically “the last value wins.”

Scaling helps keep the scores in a range where the softmax is useful instead of immediately becoming nearly one-hot.

The short version:

dot product tells us similarity
sqrt(d) keeps the similarity values numerically sane

Step 2: Softmax

After QK^T / sqrt(d), we have raw scores, but raw scores are not weights yet. We need each row to become a probability distribution.

That means each row should:

1. contain non-negative numbers
2. sum to 1

Softmax does exactly that.

Follow the naive kernel here.

For a vector xx, softmax is:

\(softmax(x_i) = e^{x_i} / \sum_j e^{x_j}\)

Let’s take:

x = [2, 1, 0]

Exponentiate each value:

e^2 ≈ 7.39
e^1 ≈ 2.72
e^0 = 1

Sum:

7.39 + 2.72 + 1 = 11.11

Divide each by the sum:

[7.39 / 11.11, 2.72 / 11.11, 1 / 11.11]
≈ [0.665, 0.245, 0.090]

Now we have probabilities.

The biggest input score became the biggest probability, but the other scores still contribute.

In attention, we apply softmax row-wise:

\(P = softmax(S)\)

So every token gets its own probability distribution over all tokens.

P[i, j] = how much token i attends to token j

One important detail: we do a numerically stable softmax.

Instead of:

\(e^{x_i}/\sum_j e^{x_j}\)

we compute:

\(e^{x_i - max(x)} / \sum_j e^{x_j - max(x)}\)

Subtracting the row max does not change the result, but it avoids exploding exponentials.

In the CUDA code, the softmax kernel does exactly this: one block handles one row, first reducing to find the row max, then reducing again to find the sum of exp(score - max), then normalizing the row.

Step 3: PV

Now we have:

P = attention probabilities
V = values

The last step is:

\(O = PV\)

The kernel is here.

Shapes:

P      x V      = O
[N,N]  x [N,d]  = [N,d]

Each output token is a weighted average of value vectors.

For token i:

\(O_i = \sum_j P_{ij}V_j\)

In English:

for token i,
look at every token j,
take token j's value vector,
multiply it by how much token i attends to token j,
sum everything up

That gives token i a new representation containing information from the rest of the sequence.

That’s the whole algorithm.

S = QK^T / sqrt(d)
P = softmax(S)
O = PV

Two matmuls with a softmax wedged in the middle.

Why This Ordering?

The CUDA implementation follows the math directly:

qkKernel      -> compute S = QK^T / sqrt(d)
softmaxKernel -> compute P = softmax(S)
pvKernel      -> compute O = PV

This ordering is not arbitrary, we cannot run softmax first because softmax needs the scores.

We cannot run PV first, because P does not exist yet.

So the dependency chain is:

Q, K
 ↓
scores S
 ↓
probabilities P
 ↓
output O

In qkKernel, one CUDA thread computes one element of the score matrix:

S[row, col] = scale * sum_k Q[row, k] * K[col, k]

That is just the dot product between one query row and one key row.

In softmaxKernel, one CUDA block owns one row of S.

That matters because softmax couples the entire row together. To normalize one value in the row, we need the row max and the row sum. So unlike the first matmul, where each output element can be computed independently, the softmax needs coordination across the row.

Then in pvKernel, one CUDA thread computes one element of the output:

O[row, k] = sum_j P[row, j] * V[j, k]

Again, this looks like plain matrix multiplication, and that’s the interesting part. The first and third steps are matmuls, while the middle step is the weird one.

The softmax is what makes attention more annoying than a regular matmul. Each row needs a reduction, normalization, and then the probabilities are used by the next matmul.

The Naive HBM Algorithm

The implementation in this article intentionally materializes both intermediate matrices:

S = score matrix
P = probability matrix

The code comments describe this as the textbook three-pass algorithm:

1. Load Q and K, compute S, write S to HBM
2. Read S, compute P = softmax(S), write P to HBM
3. Read P and V, compute O, write O to HBM

That is exactly what we want for learning. It is not what we want for maximum performance.

For this example:

N = 4096
d = 128
dtype = FP16 storage, FP32 accumulation

S is:

4096 x 4096 x 2 bytes = 32 MiB

P is also:

4096 x 4096 x 2 bytes = 32 MiB

So just those two intermediate matrices occupy:

64 MiB

And remember: we do not merely allocate them

• We write S.

• Then read S.

• Then write P.

• Then read P.

That is a lot of memory traffic for data that only exists to bridge the three steps.

Attention is implemented as standard scaled dot-product attention on one GPU, with full S and P materialized in HBM, deliberately avoiding shared-memory tiling, tensor cores, fused kernels, or CUDA libraries. This is the attention version of our naive matrix multiplication from the previous blog.

Not optimal, but very useful for learning.

Let’s Check the Code Shape

The three kernels map nicely to the three math operations:

Pass 1:
S = QK^T / sqrt(d)

Pass 2:
P = softmax(S)

Pass 3:
O = PV

The README summarizes the pass structure like this:

qkKernel      reads Q, K      writes S
softmaxKernel reads S         writes P
pvKernel      reads P, V      writes O

At N = 4096 and d = 128, the example output on a single H100 shows the three passes taking about 5.75 ms total, with QK^T dominating the runtime in this naive implementation.

That timing result is not the point by itself, the point is the breakdown.

We can now see attention as real GPU work:

matmul
row-wise reduction + normalization
matmul

And that helps us understand both the parallelism and the pain.

The matmuls are easy to reason about because each output element can be assigned to a thread.

The softmax is still parallel, but it needs cooperation across each row.

And the memory traffic is the monster hiding in the corner.

Why This Was Such a Big Deal in 2017

Before Transformers took over, sequence models often relied heavily on recurrence. Recurrent models process tokens in order, which creates an inherent sequential dependency. The Transformer paper argued for replacing that recurrence with attention, making the model significantly more parallelizable.

That is the key systems insight.

Attention lets every token look at every other token in a more parallel-friendly way.

Yes, the N x N score matrix is expensive.

But the work maps beautifully onto GPUs:

many tokens
many pairwise scores
many independent dot products
many output elements

This is why attention became such a core building block for modern LLMs. It is also why GPU kernels for attention became such a competitive optimization target. Once attention is everywhere, shaving time off attention matters a lot.

Forward Looking: Tensor Parallel Attention

In the next article, we’ll take this same standard attention calculation and shard it across multiple GPUs.

The trick is that production transformers do not usually have one giant attention head. They have multiple heads.

So instead of thinking:

one attention head
one GPU

we can think:

many attention heads
split heads across GPUs

Each GPU owns some subset of heads.

That part is nice because heads are mostly independent during the attention calculation itself:

GPU 0 computes heads 0..3
GPU 1 computes heads 4..7
...

But after the heads are computed, the model applies an output projection that mixes head outputs back into the model dimension. That output projection creates the communication point. Each GPU can compute a partial result, but the final result requires summing the partials across GPUs — an all-reduce.

That will be the next post.

The important part for today:

we first needed a correct single-GPU attention baseline

Now we have one.

Forward Looking: Optimization Opportunities

We are also leaving a ton of performance on the table.

On purpose.

Some obvious future optimization directions:

1. shared-memory tiling
2. kernel fusion
3. online softmax
4. tensor cores
5. vectorized memory loads
6. better memory coalescing
7. avoiding materialized S and P
8. FlashAttention-style IO-aware scheduling

A real production attention kernel is not going to launch three totally separate naive kernels and repeatedly bounce giant intermediates through HBM. But that is exactly why this version is valuable.

• It gives us the baseline.

• It gives us the story.

• It shows the math clearly.

And it makes the bottleneck obvious enough that optimizations covered in future posts can actually tie back to pain points in this naive approach.

Final Remarks

If you made it here, nice job. This was a dense one.

We started with tokens and vectors, introduced Q, K, and V, broke down scaled dot-product attention, walked through softmax, and then mapped the math directly to three CUDA kernels:

qkKernel
softmaxKernel
pvKernel

The big takeaway is that standard attention is not conceptually that wild. It is pairwise matching, normalization, and weighted averaging. The hard part is making it fast and memory-efficient at scale.

That is where the next few posts are going:

1. tensor-parallel attention across multiple GPUs
2. optimization opportunities like FlashAttention

As usual, the full code is in the repo..

Thanks for reading, and I hope you found this useful.