From 8137fd85505c3fa890f5b103aefbefb9affb637c Mon Sep 17 00:00:00 2001 From: "claude[bot]" <41898282+claude[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 06:24:08 +0000 Subject: [PATCH] Add challenge 109: Fused QKV Projection (Medium) Implement the fused query/key/value projection at the entrance of a transformer attention layer: a single [3D, D] matmul over the packed weight followed by a split and reshape into the [num_heads, M, head_dim] layout used by multi-head attention. Co-Authored-By: Claude Opus 4.7 --- .../109_fused_qkv_projection/challenge.html | 162 +++++++++++++ .../109_fused_qkv_projection/challenge.py | 217 ++++++++++++++++++ .../starter/starter.cu | 5 + .../starter/starter.cute.py | 17 ++ .../starter/starter.jax.py | 15 ++ .../starter/starter.mojo | 19 ++ .../starter/starter.pytorch.py | 15 ++ .../starter/starter.triton.py | 17 ++ 8 files changed, 467 insertions(+) create mode 100644 challenges/medium/109_fused_qkv_projection/challenge.html create mode 100644 challenges/medium/109_fused_qkv_projection/challenge.py create mode 100644 challenges/medium/109_fused_qkv_projection/starter/starter.cu create mode 100644 challenges/medium/109_fused_qkv_projection/starter/starter.cute.py create mode 100644 challenges/medium/109_fused_qkv_projection/starter/starter.jax.py create mode 100644 challenges/medium/109_fused_qkv_projection/starter/starter.mojo create mode 100644 challenges/medium/109_fused_qkv_projection/starter/starter.pytorch.py create mode 100644 challenges/medium/109_fused_qkv_projection/starter/starter.triton.py diff --git a/challenges/medium/109_fused_qkv_projection/challenge.html b/challenges/medium/109_fused_qkv_projection/challenge.html new file mode 100644 index 00000000..fd3cd02e --- /dev/null +++ b/challenges/medium/109_fused_qkv_projection/challenge.html @@ -0,0 +1,162 @@ +

+ Implement the fused QKV projection used at the entrance of every transformer attention layer. + Given an input matrix x of shape [M, D] and a packed weight matrix + W_qkv of shape [3D, D] whose rows are the query, key, and value + projections stacked in that order, compute + x × W_qkvᵀ, split the result into three [M, D] matrices, and + reshape each into [num_heads, M, head_dim] — 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 float32. +

+ + + + + + + + + + + + x + [M, D] + + + + + + + x · W_qkvᵀ + [M, 3D] + fused projection + + + + + + + split + along D + + + + + reshape + transpose + [M,D] → [num_heads, M, head_dim] + + + + + reshape + transpose + [M,D] → [num_heads, M, head_dim] + + + + + reshape + transpose + [M,D] → [num_heads, M, head_dim] + + + + + Q + + + + K + + + + V + + + W_qkv rows: [W_Q; W_K; W_V] + + +

Implementation Requirements

+ + +

Example

+

+ Input: M = 2, num_heads = 2, head_dim = 2 (so + D = 4). +

+

+ \(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} + \] +

+

+ 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 num_heads and each inner \(2 \times 2\) block is one + head's \((M, \text{head\_dim})\) slice. +

+ +

Constraints

+ diff --git a/challenges/medium/109_fused_qkv_projection/challenge.py b/challenges/medium/109_fused_qkv_projection/challenge.py new file mode 100644 index 00000000..6fcf03ed --- /dev/null +++ b/challenges/medium/109_fused_qkv_projection/challenge.py @@ -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, + } diff --git a/challenges/medium/109_fused_qkv_projection/starter/starter.cu b/challenges/medium/109_fused_qkv_projection/starter/starter.cu new file mode 100644 index 00000000..8412a4a4 --- /dev/null +++ b/challenges/medium/109_fused_qkv_projection/starter/starter.cu @@ -0,0 +1,5 @@ +#include + +// 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) {} diff --git a/challenges/medium/109_fused_qkv_projection/starter/starter.cute.py b/challenges/medium/109_fused_qkv_projection/starter/starter.cute.py new file mode 100644 index 00000000..a0592b0b --- /dev/null +++ b/challenges/medium/109_fused_qkv_projection/starter/starter.cute.py @@ -0,0 +1,17 @@ +import cutlass +import cutlass.cute as cute + + +# x, W_qkv, Q, K, V are tensors on the GPU +@cute.jit +def solve( + x: cute.Tensor, + W_qkv: cute.Tensor, + Q: cute.Tensor, + K: cute.Tensor, + V: cute.Tensor, + M: cute.Int32, + num_heads: cute.Int32, + head_dim: cute.Int32, +): + pass diff --git a/challenges/medium/109_fused_qkv_projection/starter/starter.jax.py b/challenges/medium/109_fused_qkv_projection/starter/starter.jax.py new file mode 100644 index 00000000..216e76a3 --- /dev/null +++ b/challenges/medium/109_fused_qkv_projection/starter/starter.jax.py @@ -0,0 +1,15 @@ +import jax +import jax.numpy as jnp + + +# x, W_qkv are tensors on device +@jax.jit +def solve( + x: jax.Array, + W_qkv: jax.Array, + M: int, + num_heads: int, + head_dim: int, +) -> tuple[jax.Array, jax.Array, jax.Array]: + # return output tensors directly + pass diff --git a/challenges/medium/109_fused_qkv_projection/starter/starter.mojo b/challenges/medium/109_fused_qkv_projection/starter/starter.mojo new file mode 100644 index 00000000..899d020f --- /dev/null +++ b/challenges/medium/109_fused_qkv_projection/starter/starter.mojo @@ -0,0 +1,19 @@ +from gpu.host import DeviceContext +from gpu.id import block_dim, block_idx, thread_idx +from memory import UnsafePointer +from math import ceildiv + + +# x, W_qkv, Q, K, V are device pointers +@export +def solve( + x: UnsafePointer[Float32], + W_qkv: UnsafePointer[Float32], + Q: UnsafePointer[Float32], + K: UnsafePointer[Float32], + V: UnsafePointer[Float32], + M: Int32, + num_heads: Int32, + head_dim: Int32, +): + pass diff --git a/challenges/medium/109_fused_qkv_projection/starter/starter.pytorch.py b/challenges/medium/109_fused_qkv_projection/starter/starter.pytorch.py new file mode 100644 index 00000000..2978365a --- /dev/null +++ b/challenges/medium/109_fused_qkv_projection/starter/starter.pytorch.py @@ -0,0 +1,15 @@ +import torch + + +# x, W_qkv, Q, K, V are tensors on the GPU +def solve( + x: torch.Tensor, + W_qkv: torch.Tensor, + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + M: int, + num_heads: int, + head_dim: int, +): + pass diff --git a/challenges/medium/109_fused_qkv_projection/starter/starter.triton.py b/challenges/medium/109_fused_qkv_projection/starter/starter.triton.py new file mode 100644 index 00000000..2c8046ee --- /dev/null +++ b/challenges/medium/109_fused_qkv_projection/starter/starter.triton.py @@ -0,0 +1,17 @@ +import torch +import triton +import triton.language as tl + + +# x, W_qkv, Q, K, V are tensors on the GPU +def solve( + x: torch.Tensor, + W_qkv: torch.Tensor, + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + M: int, + num_heads: int, + head_dim: int, +): + pass