Skip to content

Add challenge 109: Fused QKV Projection (Medium) #291

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
claude wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add challenge 109: Fused QKV Projection (Medium) #291

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
162 changes: 162 additions & 0 deletions challenges/medium/109_fused_qkv_projection/challenge.html
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,162 @@
<p>
Implement the fused QKV projection used at the entrance of every transformer attention layer.
Given an input matrix <code>x</code> of shape <code>[M, D]</code> and a packed weight matrix
<code>W_qkv</code> of shape <code>[3D, D]</code> whose rows are the query, key, and value
projections stacked in that order, compute
<code>x &times; W_qkv&#x1D40;</code>, split the result into three <code>[M, D]</code> matrices, and
reshape each into <code>[num_heads, M, head_dim]</code> — the layout expected by multi-head
attention. Fusing the three projections into a single matmul reuses the input reads and cuts
launch overhead compared with three separate GEMMs. All tensors are <code>float32</code>.
</p>

<svg width="680" height="260" viewBox="0 0 680 260" xmlns="http://www.w3.org/2000/svg"
style="display:block; margin:20px auto; font-family:monospace;">
<rect width="680" height="260" fill="#222" rx="8"/>
<defs>
<marker id="arr" markerWidth="8" markerHeight="8" refX="6" refY="3" orient="auto">
<path d="M0,0 L0,6 L8,3 z" fill="#888"/>
</marker>
</defs>

<!-- x box -->
<rect x="16" y="110" width="70" height="40" rx="4" fill="#2a4a7f" stroke="#5588cc" stroke-width="1.5"/>
<text x="51" y="134" fill="#ccc" font-size="12" text-anchor="middle">x</text>
<text x="51" y="164" fill="#666" font-size="9" text-anchor="middle">[M, D]</text>

<!-- Arrow x -> matmul -->
<line x1="86" y1="130" x2="128" y2="130" stroke="#888" stroke-width="1.5" marker-end="url(#arr)"/>

<!-- W_qkv matmul box -->
<rect x="130" y="100" width="120" height="60" rx="4" fill="#5a3a1a" stroke="#cc8844" stroke-width="1.5"/>
<text x="190" y="126" fill="#ccc" font-size="12" text-anchor="middle">x &#xb7; W_qkv&#x1D40;</text>
<text x="190" y="146" fill="#666" font-size="9" text-anchor="middle">[M, 3D]</text>
<text x="190" y="94" fill="#cc8844" font-size="9" text-anchor="middle">fused projection</text>

<!-- Arrow matmul -> split -->
<line x1="250" y1="130" x2="290" y2="130" stroke="#888" stroke-width="1.5" marker-end="url(#arr)"/>

<!-- Split box -->
<rect x="292" y="100" width="70" height="60" rx="4" fill="#3a1a5a" stroke="#8844cc" stroke-width="1.5"/>
<text x="327" y="134" fill="#ccc" font-size="11" text-anchor="middle">split</text>
<text x="327" y="152" fill="#666" font-size="8" text-anchor="middle">along D</text>

<!-- Q branch -->
<line x1="362" y1="115" x2="404" y2="52" stroke="#888" stroke-width="1.5" marker-end="url(#arr)"/>
<rect x="406" y="32" width="150" height="40" rx="4" fill="#2a4a7f" stroke="#5588cc" stroke-width="1.5"/>
<text x="481" y="52" fill="#ccc" font-size="11" text-anchor="middle">reshape + transpose</text>
<text x="481" y="66" fill="#666" font-size="8" text-anchor="middle">[M,D] &#x2192; [num_heads, M, head_dim]</text>

<!-- K branch -->
<line x1="362" y1="130" x2="404" y2="130" stroke="#888" stroke-width="1.5" marker-end="url(#arr)"/>
<rect x="406" y="110" width="150" height="40" rx="4" fill="#2a4a7f" stroke="#5588cc" stroke-width="1.5"/>
<text x="481" y="130" fill="#ccc" font-size="11" text-anchor="middle">reshape + transpose</text>
<text x="481" y="144" fill="#666" font-size="8" text-anchor="middle">[M,D] &#x2192; [num_heads, M, head_dim]</text>

<!-- V branch -->
<line x1="362" y1="145" x2="404" y2="208" stroke="#888" stroke-width="1.5" marker-end="url(#arr)"/>
<rect x="406" y="188" width="150" height="40" rx="4" fill="#2a4a7f" stroke="#5588cc" stroke-width="1.5"/>
<text x="481" y="208" fill="#ccc" font-size="11" text-anchor="middle">reshape + transpose</text>
<text x="481" y="222" fill="#666" font-size="8" text-anchor="middle">[M,D] &#x2192; [num_heads, M, head_dim]</text>

<!-- Output Q, K, V labels -->
<line x1="556" y1="52" x2="596" y2="52" stroke="#888" stroke-width="1.5" marker-end="url(#arr)"/>
<rect x="600" y="32" width="60" height="40" rx="4" fill="#3a1a3a" stroke="#cc44cc" stroke-width="1.5"/>
<text x="630" y="56" fill="#ccc" font-size="14" text-anchor="middle">Q</text>

<line x1="556" y1="130" x2="596" y2="130" stroke="#888" stroke-width="1.5" marker-end="url(#arr)"/>
<rect x="600" y="110" width="60" height="40" rx="4" fill="#3a1a3a" stroke="#cc44cc" stroke-width="1.5"/>
<text x="630" y="134" fill="#ccc" font-size="14" text-anchor="middle">K</text>

<line x1="556" y1="208" x2="596" y2="208" stroke="#888" stroke-width="1.5" marker-end="url(#arr)"/>
<rect x="600" y="188" width="60" height="40" rx="4" fill="#3a1a3a" stroke="#cc44cc" stroke-width="1.5"/>
<text x="630" y="212" fill="#ccc" font-size="14" text-anchor="middle">V</text>

<!-- Weight layout label -->
<text x="190" y="182" fill="#666" font-size="9" text-anchor="middle">W_qkv rows: [W_Q; W_K; W_V]</text>
</svg>

<h2>Implementation Requirements</h2>
<ul>
<li>Implement the <code>solve</code> function with the signature unchanged.</li>
<li>Do not use external libraries beyond the framework provided.</li>
<li>Write the results into <code>Q</code>, <code>K</code>, and <code>V</code> in-place.</li>
<li><code>D = num_heads &times; head_dim</code>.</li>
<li>
Row ranges of <code>W_qkv</code>: rows <code>[0, D)</code> project to <code>Q</code>, rows
<code>[D, 2D)</code> to <code>K</code>, rows <code>[2D, 3D)</code> to <code>V</code>.
</li>
</ul>

<h2>Example</h2>
<p>
Input: <code>M</code> = 2, <code>num_heads</code> = 2, <code>head_dim</code> = 2 (so
<code>D</code> = 4).
</p>
<p>
\(x\) (float32, \(2 \times 4\)):
\[
x = \begin{bmatrix} 1.0 & 0.0 & 0.0 & 0.0 \\ 0.0 & 1.0 & 0.0 & 0.0 \end{bmatrix}
\]
\(W_{qkv}\) (\(12 \times 4\), rows stacked as \(W_Q\); \(W_K\); \(W_V\)):
\[
W_{qkv} = \begin{bmatrix}
1.0 & 0.0 & 0.0 & 0.0 \\
0.0 & 1.0 & 0.0 & 0.0 \\
0.0 & 0.0 & 1.0 & 0.0 \\
0.0 & 0.0 & 0.0 & 1.0 \\
\hline
0.0 & 1.0 & 0.0 & 0.0 \\
1.0 & 0.0 & 0.0 & 0.0 \\
0.0 & 0.0 & 0.0 & 1.0 \\
0.0 & 0.0 & 1.0 & 0.0 \\
\hline
2.0 & 0.0 & 0.0 & 0.0 \\
0.0 & 2.0 & 0.0 & 0.0 \\
0.0 & 0.0 & 2.0 & 0.0 \\
0.0 & 0.0 & 0.0 & 2.0
\end{bmatrix}
\]
</p>
<p>
The packed projection gives \(x \cdot W_{qkv}^{\top}\) (\(2 \times 12\)):
\[
\begin{bmatrix} 1.0 & 0.0 & 0.0 & 0.0 & 0.0 & 1.0 & 0.0 & 0.0 & 2.0 & 0.0 & 0.0 & 0.0 \\
0.0 & 1.0 & 0.0 & 0.0 & 1.0 & 0.0 & 0.0 & 0.0 & 0.0 & 2.0 & 0.0 & 0.0 \end{bmatrix}
\]
After splitting along the last axis into three \(2 \times 4\) blocks, reshaping each to
\((M, \text{num\_heads}, \text{head\_dim})\) = \((2, 2, 2)\), and transposing to
\((\text{num\_heads}, M, \text{head\_dim})\):
\[
Q =
\begin{bmatrix}
\begin{bmatrix} 1.0 & 0.0 \\ 0.0 & 1.0 \end{bmatrix} \\
\begin{bmatrix} 0.0 & 0.0 \\ 0.0 & 0.0 \end{bmatrix}
\end{bmatrix}, \quad
K =
\begin{bmatrix}
\begin{bmatrix} 0.0 & 1.0 \\ 1.0 & 0.0 \end{bmatrix} \\
\begin{bmatrix} 0.0 & 0.0 \\ 0.0 & 0.0 \end{bmatrix}
\end{bmatrix}, \quad
V =
\begin{bmatrix}
\begin{bmatrix} 2.0 & 0.0 \\ 0.0 & 2.0 \end{bmatrix} \\
\begin{bmatrix} 0.0 & 0.0 \\ 0.0 & 0.0 \end{bmatrix}
\end{bmatrix}
\]
where the outer bracket indexes <code>num_heads</code> and each inner \(2 \times 2\) block is one
head's \((M, \text{head\_dim})\) slice.
</p>

<h2>Constraints</h2>
<ul>
<li>1 &le; <code>M</code> &le; 8,192</li>
<li>1 &le; <code>num_heads</code> &le; 64</li>
<li>1 &le; <code>head_dim</code> &le; 256</li>
<li><code>D</code> = <code>num_heads</code> &times; <code>head_dim</code> &le; 8,192</li>
<li>All tensors are <code>float32</code> on the GPU.</li>
<li>Input values are in the range [-10, 10].</li>
<li>
Performance is measured with <code>M</code> = 512, <code>num_heads</code> = 32,
<code>head_dim</code> = 128
</li>
</ul>
217 changes: 217 additions & 0 deletions challenges/medium/109_fused_qkv_projection/challenge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,217 @@
import ctypes
from typing import Any, Dict, List

import torch
from core.challenge_base import ChallengeBase, OutTensor, RandnTensor


class Challenge(ChallengeBase):
name = "Fused QKV Projection"
atol = 0.0001
rtol = 0.0001
num_gpus = 1
access_tier = "free"

def reference_impl(
self,
x: torch.Tensor,
W_qkv: torch.Tensor,
Q: torch.Tensor,
K: torch.Tensor,
V: torch.Tensor,
M: int,
num_heads: int,
head_dim: int,
):
D = num_heads * head_dim
assert x.shape == (M, D)
assert W_qkv.shape == (3 * D, D)
assert Q.shape == (num_heads, M, head_dim)
assert K.shape == (num_heads, M, head_dim)
assert V.shape == (num_heads, M, head_dim)
assert x.dtype == W_qkv.dtype == Q.dtype == K.dtype == V.dtype == torch.float32

# Fused projection: single matmul producing packed Q, K, V.
qkv = x @ W_qkv.t() # [M, 3*D]

# Split into Q, K, V along the last dimension.
q_flat, k_flat, v_flat = qkv.split(D, dim=-1) # each [M, D]

# Reshape to (M, num_heads, head_dim) and transpose to (num_heads, M, head_dim).
Q.copy_(q_flat.reshape(M, num_heads, head_dim).transpose(0, 1))
K.copy_(k_flat.reshape(M, num_heads, head_dim).transpose(0, 1))
V.copy_(v_flat.reshape(M, num_heads, head_dim).transpose(0, 1))

def reference_impl_jax(self, x, W_qkv, M, num_heads, head_dim):
import jax.numpy as jnp

D = num_heads * head_dim
qkv = x @ W_qkv.T # [M, 3*D]
q_flat = qkv[:, :D]
k_flat = qkv[:, D : 2 * D]
v_flat = qkv[:, 2 * D :]

def to_heads(t):
return jnp.transpose(jnp.reshape(t, (M, num_heads, head_dim)), (1, 0, 2))

return to_heads(q_flat), to_heads(k_flat), to_heads(v_flat)

def get_solve_signature(self) -> Dict[str, tuple]:
return {
"x": (ctypes.POINTER(ctypes.c_float), "in"),
"W_qkv": (ctypes.POINTER(ctypes.c_float), "in"),
"Q": (ctypes.POINTER(ctypes.c_float), "out"),
"K": (ctypes.POINTER(ctypes.c_float), "out"),
"V": (ctypes.POINTER(ctypes.c_float), "out"),
"M": (ctypes.c_int, "in"),
"num_heads": (ctypes.c_int, "in"),
"head_dim": (ctypes.c_int, "in"),
}

def _make_test_case(self, M, num_heads, head_dim, zero_x=False):
device = self.device
dtype = torch.float32
D = num_heads * head_dim
if zero_x:
x = torch.zeros(M, D, device=device, dtype=dtype)
else:
x = torch.randn(M, D, device=device, dtype=dtype) * 0.1
W_qkv = torch.randn(3 * D, D, device=device, dtype=dtype) * 0.02
Q = torch.empty(num_heads, M, head_dim, device=device, dtype=dtype)
K = torch.empty(num_heads, M, head_dim, device=device, dtype=dtype)
V = torch.empty(num_heads, M, head_dim, device=device, dtype=dtype)
return {
"x": x,
"W_qkv": W_qkv,
"Q": Q,
"K": K,
"V": V,
"M": M,
"num_heads": num_heads,
"head_dim": head_dim,
}

def generate_example_test(self) -> Dict[str, Any]:
device = self.device
dtype = torch.float32
# M = 2 tokens, D = 4, num_heads = 2, head_dim = 2.
M, num_heads, head_dim = 2, 2, 2

# Input rows: two basis-like vectors.
x = torch.tensor(
[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]],
device=device,
dtype=dtype,
)
# W_qkv is [3*D=12, D=4]. Rows 0..3 are Q weights, 4..7 are K, 8..11 are V.
# We build each 4x4 block to give a small, hand-checkable result.
w_q = torch.tensor(
[
[1.0, 0.0, 0.0, 0.0],
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
],
device=device,
dtype=dtype,
)
w_k = torch.tensor(
[
[0.0, 1.0, 0.0, 0.0],
[1.0, 0.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 1.0],
[0.0, 0.0, 1.0, 0.0],
],
device=device,
dtype=dtype,
)
w_v = torch.tensor(
[
[2.0, 0.0, 0.0, 0.0],
[0.0, 2.0, 0.0, 0.0],
[0.0, 0.0, 2.0, 0.0],
[0.0, 0.0, 0.0, 2.0],
],
device=device,
dtype=dtype,
)
W_qkv = torch.cat([w_q, w_k, w_v], dim=0) # [12, 4]

Q = torch.empty(num_heads, M, head_dim, device=device, dtype=dtype)
K = torch.empty(num_heads, M, head_dim, device=device, dtype=dtype)
V = torch.empty(num_heads, M, head_dim, device=device, dtype=dtype)
return {
"x": x,
"W_qkv": W_qkv,
"Q": Q,
"K": K,
"V": V,
"M": M,
"num_heads": num_heads,
"head_dim": head_dim,
}

def generate_functional_test(self) -> List[Dict[str, Any]]:
torch.manual_seed(42)
tests = []

# Edge: single token, single head, tiny head_dim.
tests.append(self._make_test_case(1, 1, 4))

# Zero input.
tests.append(self._make_test_case(4, 2, 8, zero_x=True))

# Power-of-2 sizes, small.
tests.append(self._make_test_case(16, 4, 8))

# Power-of-2 sizes, moderate.
tests.append(self._make_test_case(64, 8, 16))

# Non-power-of-2 M.
tests.append(self._make_test_case(30, 4, 16))

# Non-power-of-2 M, larger head layout.
tests.append(self._make_test_case(100, 8, 32))

# Non-power-of-2 M, medium size.
tests.append(self._make_test_case(255, 4, 32))

# Negative-only inputs to exercise sign handling.
M, num_heads, head_dim = 32, 4, 16
D = num_heads * head_dim
tests.append(
{
"x": torch.full((M, D), -0.5, device=self.device, dtype=torch.float32),
"W_qkv": torch.randn(3 * D, D, device=self.device, dtype=torch.float32) * 0.02,
"Q": torch.empty(num_heads, M, head_dim, device=self.device, dtype=torch.float32),
"K": torch.empty(num_heads, M, head_dim, device=self.device, dtype=torch.float32),
"V": torch.empty(num_heads, M, head_dim, device=self.device, dtype=torch.float32),
"M": M,
"num_heads": num_heads,
"head_dim": head_dim,
}
)

# Realistic small inference batch (GPT-2-small-style: D=768, 12 heads, dh=64).
tests.append(self._make_test_case(128, 12, 64))

# Realistic medium inference batch (LLaMA-2-7B-style: D=4096, 32 heads, dh=128).
tests.append(self._make_test_case(256, 32, 128))

return tests

def generate_performance_test(self) -> Dict[str, Any]:
# LLaMA-2-7B-style attention projection: D=4096, num_heads=32, head_dim=128.
# M=512 = batch 4 x seq_len 128 (a realistic prefill workload).
M, num_heads, head_dim = 512, 32, 128
D = num_heads * head_dim # 4096
return {
"x": RandnTensor((M, D), std=0.1),
"W_qkv": RandnTensor((3 * D, D), std=0.02),
"Q": OutTensor((num_heads, M, head_dim)),
"K": OutTensor((num_heads, M, head_dim)),
"V": OutTensor((num_heads, M, head_dim)),
"M": M,
"num_heads": num_heads,
"head_dim": head_dim,
}
5 changes: 5 additions & 0 deletions challenges/medium/109_fused_qkv_projection/starter/starter.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
#include <cuda_runtime.h>

// x, W_qkv, Q, K, V are device pointers
extern "C" void solve(const float* x, const float* W_qkv, float* Q, float* K, float* V, int M,
int num_heads, int head_dim) {}
Loading
Loading