diff --git a/challenges/medium/108_cross_attention/challenge.html b/challenges/medium/108_cross_attention/challenge.html new file mode 100644 index 00000000..a6235c35 --- /dev/null +++ b/challenges/medium/108_cross_attention/challenge.html @@ -0,0 +1,154 @@ +

+Implement multi-head cross-attention, the building block used in encoder-decoder transformers such +as T5, BART, and Whisper, as well as the text-conditioning layers of diffusion models like Stable +Diffusion. Given decoder queries Q of shape (M, H, D) and encoder keys +K / values V of shape (N, H, D), compute, for each head +h, scaled dot-product attention where every one of the M queries attends +to all N encoder positions: +

+

+\[ \text{head}_h = \text{softmax}\!\left(\frac{Q_h K_h^{\top}}{\sqrt{D}}\right) V_h \] +

+

+There is no causal mask — encoder positions are fully visible to every decoder query — and +the query length M may differ from the key/value length N. All tensors +use float32. +

+ + + + Cross-Attention: decoder queries attend to encoder keys/values + + + Encoder K, V + (N positions) + + K[0], V[0] + + K[1], V[1] + + ... + + K[N-1], V[N-1] + + + Decoder Q + (M queries, may differ from N) + + Q[0] + + Q[1] + + Q[M-1] + + + + + + + + + + + + + + for each head h: + S = Q_h K_h^T / sqrt(D) + A = softmax(S, dim=-1) + O_h = A V_h + no causal mask + M may differ from N + + + + + + + + +

Implementation Requirements

+ + +

Example

+

+ With M = 2, N = 3, H = 2, D = 2. Tensors are + laid out as (position, head, dim). Per-head views: +

+

+ \(Q_0\) (2×2): + \[ + \begin{bmatrix} + 1 & 0 \\ + 0 & 1 + \end{bmatrix} + \qquad + Q_1\text{ (2}\times\text{2):} + \begin{bmatrix} + 0 & 1 \\ + 1 & 0 + \end{bmatrix} + \] + \(K_0\) (3×2): + \[ + \begin{bmatrix} + 1 & 0 \\ + 0 & 1 \\ + 1 & 1 + \end{bmatrix} + \qquad + K_1\text{ (3}\times\text{2):} + \begin{bmatrix} + 0 & 1 \\ + 1 & 0 \\ + 1 & 1 + \end{bmatrix} + \] + \(V_0\) (3×2): + \[ + \begin{bmatrix} + 1 & 2 \\ + 3 & 4 \\ + 5 & 6 + \end{bmatrix} + \qquad + V_1\text{ (3}\times\text{2):} + \begin{bmatrix} + 7 & 8 \\ + 9 & 10 \\ + 11 & 12 + \end{bmatrix} + \] +

+

+ Output (values rounded to 4 decimals): + \[ + \text{output}_0\text{ (head 0, 2}\times\text{2):} + \begin{bmatrix} + 3.0000 & 4.0000 \\ + 3.4067 & 4.4067 + \end{bmatrix} + \qquad + \text{output}_1\text{ (head 1, 2}\times\text{2):} + \begin{bmatrix} + 9.0000 & 10.0000 \\ + 9.4067 & 10.4067 + \end{bmatrix} + \] +

+ +

Constraints

+ diff --git a/challenges/medium/108_cross_attention/challenge.py b/challenges/medium/108_cross_attention/challenge.py new file mode 100644 index 00000000..b00da98f --- /dev/null +++ b/challenges/medium/108_cross_attention/challenge.py @@ -0,0 +1,194 @@ +import ctypes +import math +from typing import Any, Dict, List + +import torch +from core.challenge_base import ChallengeBase, OutTensor, RandnTensor + + +class Challenge(ChallengeBase): + name = "Cross-Attention" + atol = 1e-04 + rtol = 1e-04 + num_gpus = 1 + access_tier = "free" + + def reference_impl( + self, + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + output: torch.Tensor, + M: int, + N: int, + H: int, + D: int, + ): + assert Q.shape == (M, H, D) + assert K.shape == (N, H, D) + assert V.shape == (N, H, D) + assert output.shape == (M, H, D) + assert Q.dtype == K.dtype == V.dtype == output.dtype + + # (M, H, D) -> (H, M, D); (N, H, D) -> (H, N, D) + Qt = Q.transpose(0, 1) + Kt = K.transpose(0, 1) + Vt = V.transpose(0, 1) + + scale = 1.0 / math.sqrt(D) + scores = torch.matmul(Qt, Kt.transpose(-2, -1)) * scale # (H, M, N) + attn = torch.softmax(scores, dim=-1) # (H, M, N) + out = torch.matmul(attn, Vt) # (H, M, D) + # (H, M, D) -> (M, H, D) + output.copy_(out.transpose(0, 1)) + + def get_solve_signature(self) -> Dict[str, tuple]: + return { + "Q": (ctypes.POINTER(ctypes.c_float), "in"), + "K": (ctypes.POINTER(ctypes.c_float), "in"), + "V": (ctypes.POINTER(ctypes.c_float), "in"), + "output": (ctypes.POINTER(ctypes.c_float), "out"), + "M": (ctypes.c_int, "in"), + "N": (ctypes.c_int, "in"), + "H": (ctypes.c_int, "in"), + "D": (ctypes.c_int, "in"), + } + + def generate_example_test(self) -> Dict[str, Any]: + dtype = torch.float32 + M, N, H, D = 2, 3, 2, 2 + Q = torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[0.0, 1.0], [1.0, 0.0]], + ], + device=self.device, + dtype=dtype, + ) + K = torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[0.0, 1.0], [1.0, 0.0]], + [[1.0, 1.0], [1.0, 1.0]], + ], + device=self.device, + dtype=dtype, + ) + V = torch.tensor( + [ + [[1.0, 2.0], [7.0, 8.0]], + [[3.0, 4.0], [9.0, 10.0]], + [[5.0, 6.0], [11.0, 12.0]], + ], + device=self.device, + dtype=dtype, + ) + output = torch.empty((M, H, D), device=self.device, dtype=dtype) + return {"Q": Q, "K": K, "V": V, "output": output, "M": M, "N": N, "H": H, "D": D} + + def _make_case(self, M, N, H, D, kind="randn"): + dtype = torch.float32 + device = self.device + if kind == "zeros": + Q = torch.zeros((M, H, D), device=device, dtype=dtype) + K = torch.zeros((N, H, D), device=device, dtype=dtype) + V = torch.zeros((N, H, D), device=device, dtype=dtype) + elif kind == "uniform": + Q = torch.empty((M, H, D), device=device, dtype=dtype).uniform_(-1.0, 1.0) + K = torch.empty((N, H, D), device=device, dtype=dtype).uniform_(-1.0, 1.0) + V = torch.empty((N, H, D), device=device, dtype=dtype).uniform_(-1.0, 1.0) + else: + Q = torch.randn(M, H, D, device=device, dtype=dtype) + K = torch.randn(N, H, D, device=device, dtype=dtype) + V = torch.randn(N, H, D, device=device, dtype=dtype) + output = torch.empty((M, H, D), device=device, dtype=dtype) + return {"Q": Q, "K": K, "V": V, "output": output, "M": M, "N": N, "H": H, "D": D} + + def generate_functional_test(self) -> List[Dict[str, Any]]: + torch.manual_seed(42) + dtype = torch.float32 + tests = [] + + # Basic example (matches generate_example_test) + tests.append( + { + "Q": torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[0.0, 1.0], [1.0, 0.0]], + ], + device=self.device, + dtype=dtype, + ), + "K": torch.tensor( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[0.0, 1.0], [1.0, 0.0]], + [[1.0, 1.0], [1.0, 1.0]], + ], + device=self.device, + dtype=dtype, + ), + "V": torch.tensor( + [ + [[1.0, 2.0], [7.0, 8.0]], + [[3.0, 4.0], [9.0, 10.0]], + [[5.0, 6.0], [11.0, 12.0]], + ], + device=self.device, + dtype=dtype, + ), + "output": torch.empty((2, 2, 2), device=self.device, dtype=dtype), + "M": 2, + "N": 3, + "H": 2, + "D": 2, + } + ) + + # Edge case: single query, single key, single head + tests.append(self._make_case(1, 1, 1, 8)) + + # Decode-like: single query vs many keys + tests.append(self._make_case(1, 16, 4, 8)) + + # Prefill-like: many queries vs single key (attention collapses to V) + tests.append(self._make_case(4, 1, 2, 8)) + + # Zero inputs (softmax should be uniform 1/N) + tests.append(self._make_case(3, 5, 2, 4, kind="zeros")) + + # Negative + mixed values via uniform + tests.append(self._make_case(4, 6, 2, 8, kind="uniform")) + + # Power-of-2 sizes + tests.append(self._make_case(16, 32, 4, 32)) + + # Non-power-of-2: M != N with odd dims + tests.append(self._make_case(30, 45, 6, 32)) + + # Larger non-power-of-2 + tests.append(self._make_case(100, 200, 8, 64)) + + # Realistic Whisper-encoder-decoder-like sizes + tests.append(self._make_case(64, 256, 8, 64)) + + # Realistic BART/T5-like sizes + tests.append(self._make_case(128, 512, 16, 64)) + + return tests + + def generate_performance_test(self) -> Dict[str, Any]: + # BART-large-style cross-attention: 1024 decoder queries attending to + # 2048 encoder tokens, 16 heads, head_dim=128. + M, N, H, D = 1024, 2048, 16, 128 + return { + "Q": RandnTensor((M, H, D)), + "K": RandnTensor((N, H, D)), + "V": RandnTensor((N, H, D)), + "output": OutTensor((M, H, D)), + "M": M, + "N": N, + "H": H, + "D": D, + } diff --git a/challenges/medium/108_cross_attention/starter/starter.cu b/challenges/medium/108_cross_attention/starter/starter.cu new file mode 100644 index 00000000..c5d05b52 --- /dev/null +++ b/challenges/medium/108_cross_attention/starter/starter.cu @@ -0,0 +1,5 @@ +#include + +// Q, K, V, output are device pointers +extern "C" void solve(const float* Q, const float* K, const float* V, float* output, int M, int N, + int H, int D) {} diff --git a/challenges/medium/108_cross_attention/starter/starter.cute.py b/challenges/medium/108_cross_attention/starter/starter.cute.py new file mode 100644 index 00000000..a036e178 --- /dev/null +++ b/challenges/medium/108_cross_attention/starter/starter.cute.py @@ -0,0 +1,17 @@ +import cutlass +import cutlass.cute as cute + + +# Q, K, V, output are tensors on the GPU +@cute.jit +def solve( + Q: cute.Tensor, + K: cute.Tensor, + V: cute.Tensor, + output: cute.Tensor, + M: cute.Int32, + N: cute.Int32, + H: cute.Int32, + D: cute.Int32, +): + pass diff --git a/challenges/medium/108_cross_attention/starter/starter.jax.py b/challenges/medium/108_cross_attention/starter/starter.jax.py new file mode 100644 index 00000000..ddb4266d --- /dev/null +++ b/challenges/medium/108_cross_attention/starter/starter.jax.py @@ -0,0 +1,17 @@ +import jax +import jax.numpy as jnp + + +# Q, K, V are tensors on device +@jax.jit +def solve( + Q: jax.Array, + K: jax.Array, + V: jax.Array, + M: int, + N: int, + H: int, + D: int, +) -> jax.Array: + # return output tensor directly + pass diff --git a/challenges/medium/108_cross_attention/starter/starter.mojo b/challenges/medium/108_cross_attention/starter/starter.mojo new file mode 100644 index 00000000..ff829202 --- /dev/null +++ b/challenges/medium/108_cross_attention/starter/starter.mojo @@ -0,0 +1,17 @@ +from std.gpu.host import DeviceContext +from std.memory import UnsafePointer + + +# Q, K, V, output are device pointers +@export +def solve( + Q: UnsafePointer[Float32, MutExternalOrigin], + K: UnsafePointer[Float32, MutExternalOrigin], + V: UnsafePointer[Float32, MutExternalOrigin], + output: UnsafePointer[Float32, MutExternalOrigin], + M: Int32, + N: Int32, + H: Int32, + D: Int32, +) raises: + pass diff --git a/challenges/medium/108_cross_attention/starter/starter.pytorch.py b/challenges/medium/108_cross_attention/starter/starter.pytorch.py new file mode 100644 index 00000000..9179ab73 --- /dev/null +++ b/challenges/medium/108_cross_attention/starter/starter.pytorch.py @@ -0,0 +1,15 @@ +import torch + + +# Q, K, V, output are tensors on the GPU +def solve( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + output: torch.Tensor, + M: int, + N: int, + H: int, + D: int, +): + pass diff --git a/challenges/medium/108_cross_attention/starter/starter.triton.py b/challenges/medium/108_cross_attention/starter/starter.triton.py new file mode 100644 index 00000000..12017739 --- /dev/null +++ b/challenges/medium/108_cross_attention/starter/starter.triton.py @@ -0,0 +1,17 @@ +import torch +import triton +import triton.language as tl + + +# Q, K, V, output are tensors on the GPU +def solve( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + output: torch.Tensor, + M: int, + N: int, + H: int, + D: int, +): + pass