AlphaGPU · claude · Jul 3, 2026
@@ -0,0 +1,154 @@
+<p>
+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 <code>Q</code> of shape <code>(M, H, D)</code> and encoder keys
+<code>K</code> / values <code>V</code> of shape <code>(N, H, D)</code>, compute, for each head
+<code>h</code>, scaled dot-product attention where every one of the <code>M</code> queries attends
+to all <code>N</code> encoder positions:
+</p>
+<p>
+\[ \text{head}_h = \text{softmax}\!\left(\frac{Q_h K_h^{\top}}{\sqrt{D}}\right) V_h \]
+</p>
+<p>
+There is no causal mask &mdash; encoder positions are fully visible to every decoder query &mdash; and
+the query length <code>M</code> may differ from the key/value length <code>N</code>. All tensors
+use <code>float32</code>.
+</p>
+
+<svg width="700" height="270" viewBox="0 0 700 270" xmlns="http://www.w3.org/2000/svg" style="display:block; margin:20px auto;">
+  <rect width="700" height="270" fill="#222" rx="10"/>
+  <text x="350" y="28" fill="#ccc" font-family="monospace" font-size="14" text-anchor="middle">Cross-Attention: decoder queries attend to encoder keys/values</text>
+
+  <!-- Encoder KV column -->
+  <text x="110" y="60" fill="#aaa" font-family="monospace" font-size="12" text-anchor="middle">Encoder K, V</text>
+  <text x="110" y="76" fill="#888" font-family="monospace" font-size="11" text-anchor="middle">(N positions)</text>
+  <rect x="50" y="88" width="120" height="28" fill="#1d4ed8" rx="4"/>
+  <text x="110" y="106" fill="#fff" font-family="monospace" font-size="11" text-anchor="middle">K[0], V[0]</text>
+  <rect x="50" y="122" width="120" height="28" fill="#1d4ed8" rx="4"/>
+  <text x="110" y="140" fill="#fff" font-family="monospace" font-size="11" text-anchor="middle">K[1], V[1]</text>
+  <rect x="50" y="156" width="120" height="28" fill="#1d4ed8" rx="4"/>
+  <text x="110" y="174" fill="#fff" font-family="monospace" font-size="11" text-anchor="middle">...</text>
+  <rect x="50" y="190" width="120" height="28" fill="#1d4ed8" rx="4"/>
+  <text x="110" y="208" fill="#fff" font-family="monospace" font-size="11" text-anchor="middle">K[N-1], V[N-1]</text>
+
+  <!-- Decoder Q column -->
+  <text x="400" y="60" fill="#aaa" font-family="monospace" font-size="12" text-anchor="middle">Decoder Q</text>
+  <text x="400" y="76" fill="#888" font-family="monospace" font-size="11" text-anchor="middle">(M queries, may differ from N)</text>
+  <rect x="340" y="100" width="120" height="28" fill="#7c3aed" rx="4"/>
+  <text x="400" y="118" fill="#fff" font-family="monospace" font-size="11" text-anchor="middle">Q[0]</text>
+  <rect x="340" y="140" width="120" height="28" fill="#7c3aed" rx="4"/>
+  <text x="400" y="158" fill="#fff" font-family="monospace" font-size="11" text-anchor="middle">Q[1]</text>
+  <rect x="340" y="180" width="120" height="28" fill="#7c3aed" rx="4"/>
+  <text x="400" y="198" fill="#fff" font-family="monospace" font-size="11" text-anchor="middle">Q[M-1]</text>
+
+  <!-- Arrows: each query attends to every K/V -->
+  <line x1="170" y1="102" x2="340" y2="112" stroke="#60a5fa" stroke-width="1.2" marker-end="url(#arr)"/>
+  <line x1="170" y1="136" x2="340" y2="114" stroke="#60a5fa" stroke-width="1.2" marker-end="url(#arr)"/>
+  <line x1="170" y1="170" x2="340" y2="116" stroke="#60a5fa" stroke-width="1.2" marker-end="url(#arr)"/>
+  <line x1="170" y1="204" x2="340" y2="118" stroke="#60a5fa" stroke-width="1.2" marker-end="url(#arr)"/>
+
+  <line x1="170" y1="102" x2="340" y2="192" stroke="#c4b5fd" stroke-width="1.2" marker-end="url(#arr)"/>
+  <line x1="170" y1="136" x2="340" y2="194" stroke="#c4b5fd" stroke-width="1.2" marker-end="url(#arr)"/>
+  <line x1="170" y1="170" x2="340" y2="196" stroke="#c4b5fd" stroke-width="1.2" marker-end="url(#arr)"/>
+  <line x1="170" y1="204" x2="340" y2="198" stroke="#c4b5fd" stroke-width="1.2" marker-end="url(#arr)"/>
+
+  <!-- Formula panel -->
+  <text x="510" y="110" fill="#4ade80" font-family="monospace" font-size="12">for each head h:</text>
+  <text x="510" y="132" fill="#4ade80" font-family="monospace" font-size="12">S = Q_h K_h^T / sqrt(D)</text>
+  <text x="510" y="154" fill="#4ade80" font-family="monospace" font-size="12">A = softmax(S, dim=-1)</text>
+  <text x="510" y="176" fill="#4ade80" font-family="monospace" font-size="12">O_h = A V_h</text>
+  <text x="510" y="204" fill="#aaa" font-family="monospace" font-size="11">no causal mask</text>
+  <text x="510" y="222" fill="#aaa" font-family="monospace" font-size="11">M may differ from N</text>
+
+  <defs>
+    <marker id="arr" markerWidth="6" markerHeight="6" refX="3" refY="3" orient="auto">
+      <path d="M0,0 L0,6 L6,3 z" fill="#888"/>
+    </marker>
+  </defs>
+</svg>
+
+<h2>Implementation Requirements</h2>
+<ul>
+  <li>Implement the function <code>solve(Q, K, V, output, M, N, H, D)</code>.</li>
+  <li>Do not change the function signature or use external libraries beyond the standard GPU framework.</li>
+  <li>Write the result into the provided <code>output</code> buffer of shape <code>(M, H, D)</code>.</li>
+  <li>Use scaled dot-product attention with scale factor <code>1 / sqrt(D)</code> and a softmax over the key dimension.</li>
+  <li>No causal mask is applied &mdash; every query attends to all <code>N</code> encoder positions.</li>
+</ul>
+
+<h2>Example</h2>
+<p>
+  With <code>M</code> = 2, <code>N</code> = 3, <code>H</code> = 2, <code>D</code> = 2. Tensors are
+  laid out as <code>(position, head, dim)</code>. Per-head views:
+</p>
+<p>
+  \(Q_0\) (2&times;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&times;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&times;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}
+  \]
+</p>
+<p>
+  <strong>Output</strong> (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}
+  \]
+</p>
+
+<h2>Constraints</h2>
+<ul>
+  <li>1 &le; <code>M</code> &le; 4,096</li>
+  <li>1 &le; <code>N</code> &le; 4,096</li>
+  <li>1 &le; <code>H</code> &le; 64</li>
+  <li>8 &le; <code>D</code> &le; 256; <code>D</code> is a multiple of 8</li>
+  <li>All tensor values are <code>float32</code></li>
+  <li>Performance is measured with <code>M</code> = 1,024, <code>N</code> = 2,048, <code>H</code> = 16, <code>D</code> = 128</li>
+</ul>
@@ -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,
+        }
@@ -0,0 +1,5 @@
+#include <cuda_runtime.h>
+
+// 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) {}
@@ -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
@@ -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