Add FP4 matmul hard challenge

kunal-mansukhani · claude · kunal-mansukhani · commit e099a71b0a4f · 2026-04-05T15:33:14.000-07:00
Weight-only FP4 E2M1 quantized matmul (W4A16) with group-wise FP16
scales, the kernel powering low-precision LLM inference on Hopper and
Blackwell. Two FP4 values are packed per uint8 byte; each contiguous
block of group_size weights along K shares one scale.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/challenges/hard/86_fp4_matmul/challenge.html b/challenges/hard/86_fp4_matmul/challenge.html
@@ -0,0 +1,113 @@
+<p>
+  Implement an FP4 weight-only quantized matrix multiplication, the kernel at the heart of
+  modern low-precision LLM inference on Hopper and Blackwell GPUs. Given a float16 activation
+  matrix <code>x</code> of shape <code>M &times; K</code> and a weight matrix stored in packed
+  FP4 E2M1 format, compute <code>y = x &times; W<sup>T</sup></code> of shape
+  <code>M &times; N</code>, where <code>W</code> is the dequantized float16 weight matrix of
+  shape <code>N &times; K</code>.
+</p>
+
+<p>
+  <strong>FP4 E2M1 format:</strong> Each weight is encoded in 4 bits as
+  [sign | exponent (2 bits) | mantissa (1 bit)], representing one of sixteen values:
+  <code>{&plusmn;0, &plusmn;0.5, &plusmn;1, &plusmn;1.5, &plusmn;2, &plusmn;3, &plusmn;4, &plusmn;6}</code>.
+  The nibble-to-value mapping is:
+</p>
+<pre>
+  0x0 =  0.0    0x8 = -0.0
+  0x1 =  0.5    0x9 = -0.5
+  0x2 =  1.0    0xA = -1.0
+  0x3 =  1.5    0xB = -1.5
+  0x4 =  2.0    0xC = -2.0
+  0x5 =  3.0    0xD = -3.0
+  0x6 =  4.0    0xE = -4.0
+  0x7 =  6.0    0xF = -6.0
+</pre>
+
+<p>
+  <strong>Packing:</strong> Each byte of <code>w_q</code> stores two FP4 weights. The high
+  nibble (bits 7&ndash;4) holds <code>w[n, 2i]</code> and the low nibble (bits 3&ndash;0) holds
+  <code>w[n, 2i+1]</code>.
+</p>
+
+<p>
+  <strong>Dequantization:</strong> Weights are dequantized group-wise. Each contiguous block of
+  <code>group_size</code> weights along the <code>K</code> dimension shares one float16 scale:
+</p>
+<pre>
+W[n, k] = fp4_decode(w_q_nibble[n, k]) * scales[n, k // group_size]
+</pre>
+
+<h2>Implementation Requirements</h2>
+<ul>
+  <li>Use only native features (external libraries are not permitted)</li>
+  <li>The <code>solve</code> function signature must remain unchanged</li>
+  <li>The final result must be stored in <code>y</code></li>
+</ul>
+
+<h2>Example</h2>
+<p>
+  Input (<code>M</code> = 2, <code>N</code> = 4, <code>K</code> = 4, <code>group_size</code> = 2):
+</p>
+<p>
+  Activations \(x\) (float16, \(2 \times 4\)):
+  \[
+  \begin{bmatrix}
+  1.0 & 0.0 & 1.0 & 0.0 \\
+  0.0 & 1.0 & 0.0 & 1.0
+  \end{bmatrix}
+  \]
+  Packed weights \(w\_q\) (uint8, \(4 \times 2\)) decoded via the FP4 E2M1 table:
+  \[
+  \begin{bmatrix}
+  \texttt{0x22} & \texttt{0x22} \\
+  \texttt{0x44} & \texttt{0x44} \\
+  \texttt{0xAA} & \texttt{0xAA} \\
+  \texttt{0x00} & \texttt{0x00}
+  \end{bmatrix}
+  \;\Rightarrow\;
+  W_{\text{fp4}} =
+  \begin{bmatrix}
+  1.0 & 1.0 & 1.0 & 1.0 \\
+  2.0 & 2.0 & 2.0 & 2.0 \\
+  -1.0 & -1.0 & -1.0 & -1.0 \\
+  0.0 & 0.0 & 0.0 & 0.0
+  \end{bmatrix}
+  \]
+  Scales (float16, \(4 \times 2\), all entries 0.5):
+  \[
+  \begin{bmatrix}
+  0.5 & 0.5 \\
+  0.5 & 0.5 \\
+  0.5 & 0.5 \\
+  0.5 & 0.5
+  \end{bmatrix}
+  \;\Rightarrow\;
+  W_{\text{dequant}} =
+  \begin{bmatrix}
+  0.5 & 0.5 & 0.5 & 0.5 \\
+  1.0 & 1.0 & 1.0 & 1.0 \\
+  -0.5 & -0.5 & -0.5 & -0.5 \\
+  0.0 & 0.0 & 0.0 & 0.0
+  \end{bmatrix}
+  \]
+  Output \(y = x \times W^T\) (float16, \(2 \times 4\)):
+  \[
+  \begin{bmatrix}
+  1.0 & 2.0 & -1.0 & 0.0 \\
+  1.0 & 2.0 & -1.0 & 0.0
+  \end{bmatrix}
+  \]
+</p>
+
+<h2>Constraints</h2>
+<ul>
+  <li>1 &le; <code>M</code>, <code>N</code> &le; 8,192</li>
+  <li>1 &le; <code>K</code> &le; 8,192</li>
+  <li><code>K</code> is divisible by <code>2</code> and by <code>group_size</code></li>
+  <li><code>group_size</code> &isin; {2, 4, 8, 16, 32}</li>
+  <li>All tensors are stored in row-major order</li>
+  <li>Input dtype: <code>x</code> and <code>scales</code> are float16; <code>w_q</code> is uint8</li>
+  <li>Output dtype: <code>y</code> is float16</li>
+  <li>Performance is measured with <code>M</code> = 2,048, <code>N</code> = 8,192, <code>K</code> = 3,072, <code>group_size</code> = 32</li>
+</ul>
diff --git a/challenges/hard/86_fp4_matmul/challenge.py b/challenges/hard/86_fp4_matmul/challenge.py
@@ -0,0 +1,174 @@
+import ctypes
+from typing import Any, Dict, List
+
+import torch
+from core.challenge_base import ChallengeBase
+
+# OCP FP4 E2M1 lookup table: 4-bit unsigned index -> float value.
+# Bit layout: [sign | exp1 exp0 | mantissa]. Sixteen representable values.
+FP4_E2M1_TABLE = [
+    0.0,
+    0.5,
+    1.0,
+    1.5,
+    2.0,
+    3.0,
+    4.0,
+    6.0,
+    -0.0,
+    -0.5,
+    -1.0,
+    -1.5,
+    -2.0,
+    -3.0,
+    -4.0,
+    -6.0,
+]
+
+
+class Challenge(ChallengeBase):
+    def __init__(self):
+        super().__init__(
+            name="FP4 MatMul",
+            atol=5e-02,
+            rtol=5e-02,
+            num_gpus=1,
+            access_tier="free",
+        )
+
+    def reference_impl(
+        self,
+        x: torch.Tensor,
+        w_q: torch.Tensor,
+        scales: torch.Tensor,
+        y: torch.Tensor,
+        M: int,
+        N: int,
+        K: int,
+        group_size: int,
+    ):
+        assert x.shape == (M, K)
+        assert w_q.shape == (N, K // 2)
+        assert scales.shape == (N, K // group_size)
+        assert y.shape == (M, N)
+        assert x.dtype == torch.float16
+        assert w_q.dtype == torch.uint8
+        assert scales.dtype == torch.float16
+        assert y.dtype == torch.float16
+        assert x.device.type == "cuda"
+        assert w_q.device.type == "cuda"
+        assert scales.device.type == "cuda"
+        assert y.device.type == "cuda"
+
+        # Decode packed FP4 E2M1 nibbles via lookup table.
+        # w_q[n, i] holds two FP4 values: w[n, 2*i] in the high nibble (bits 7:4)
+        # and w[n, 2*i+1] in the low nibble (bits 3:0).
+        table = torch.tensor(FP4_E2M1_TABLE, device=x.device, dtype=torch.float32)
+        high = ((w_q >> 4) & 0xF).to(torch.long)  # [N, K//2]
+        low = (w_q & 0xF).to(torch.long)  # [N, K//2]
+        w_high = table[high]  # [N, K//2]
+        w_low = table[low]  # [N, K//2]
+        w_fp4 = torch.stack([w_high, w_low], dim=-1).reshape(N, K)  # [N, K]
+
+        # Apply group-wise FP16 scales: each contiguous block of `group_size`
+        # weights along K shares one scale.
+        n_groups = K // group_size
+        w_groups = w_fp4.reshape(N, n_groups, group_size)
+        scales_f = scales.float().unsqueeze(-1)  # [N, n_groups, 1]
+        w_dequant = (w_groups * scales_f).reshape(N, K)
+
+        y.copy_((x.float() @ w_dequant.T).half())
+
+    def get_solve_signature(self) -> Dict[str, tuple]:
+        return {
+            "x": (ctypes.POINTER(ctypes.c_uint16), "in"),
+            "w_q": (ctypes.POINTER(ctypes.c_uint8), "in"),
+            "scales": (ctypes.POINTER(ctypes.c_uint16), "in"),
+            "y": (ctypes.POINTER(ctypes.c_uint16), "out"),
+            "M": (ctypes.c_int, "in"),
+            "N": (ctypes.c_int, "in"),
+            "K": (ctypes.c_int, "in"),
+            "group_size": (ctypes.c_int, "in"),
+        }
+
+    def _make_test_case(self, M: int, N: int, K: int, group_size: int, zero_x: bool = False):
+        device = "cuda"
+        if zero_x:
+            x = torch.zeros(M, K, device=device, dtype=torch.float16)
+        else:
+            x = torch.randn(M, K, device=device, dtype=torch.float16) * 0.5
+        w_q = torch.randint(0, 256, (N, K // 2), dtype=torch.uint8, device=device)
+        scales = torch.rand(N, K // group_size, device=device, dtype=torch.float16) * 0.1 + 0.01
+        y = torch.empty(M, N, device=device, dtype=torch.float16)
+        return {
+            "x": x,
+            "w_q": w_q,
+            "scales": scales,
+            "y": y,
+            "M": M,
+            "N": N,
+            "K": K,
+            "group_size": group_size,
+        }
+
+    def generate_example_test(self) -> Dict[str, Any]:
+        device = "cuda"
+        M, N, K, group_size = 2, 4, 4, 2
+
+        x = torch.tensor(
+            [[1.0, 0.0, 1.0, 0.0], [0.0, 1.0, 0.0, 1.0]],
+            device=device,
+            dtype=torch.float16,
+        )
+        # Packed FP4 E2M1 weights (high nibble first).
+        # Row 0: FP4 [1.0,1.0,1.0,1.0] -> nibbles [0x2,0x2,0x2,0x2] -> bytes [0x22,0x22] = [34,34]
+        # Row 1: FP4 [2.0,2.0,2.0,2.0] -> nibbles [0x4,0x4,0x4,0x4] -> bytes [0x44,0x44] = [68,68]
+        # Row 2: FP4 [-1,-1,-1,-1] -> nibbles [0xA,0xA,0xA,0xA] -> bytes [0xAA,0xAA] = [170,170]
+        # Row 3: FP4 [0.0,0.0,0.0,0.0] -> nibbles [0x0,0x0,0x0,0x0] -> bytes [0x00,0x00] = [0,0]
+        w_q = torch.tensor(
+            [[34, 34], [68, 68], [170, 170], [0, 0]],
+            dtype=torch.uint8,
+            device=device,
+        )
+        scales = torch.full((N, K // group_size), 0.5, device=device, dtype=torch.float16)
+        y = torch.empty(M, N, device=device, dtype=torch.float16)
+
+        return {
+            "x": x,
+            "w_q": w_q,
+            "scales": scales,
+            "y": y,
+            "M": M,
+            "N": N,
+            "K": K,
+            "group_size": group_size,
+        }
+
+    def generate_functional_test(self) -> List[Dict[str, Any]]:
+        torch.manual_seed(42)
+        tests = []
+
+        # Edge cases with tiny shapes.
+        tests.append(self._make_test_case(1, 2, 4, 2, zero_x=True))
+        tests.append(self._make_test_case(2, 4, 4, 2))
+        tests.append(self._make_test_case(3, 5, 8, 4))
+
+        # Power-of-2 shapes.
+        tests.append(self._make_test_case(16, 16, 32, 16))
+        tests.append(self._make_test_case(32, 64, 64, 32))
+        tests.append(self._make_test_case(128, 128, 256, 32))
+
+        # Non-power-of-2 shapes.
+        tests.append(self._make_test_case(30, 50, 64, 32))
+        tests.append(self._make_test_case(100, 200, 128, 32))
+        tests.append(self._make_test_case(255, 100, 128, 32))
+
+        # Realistic LLM inference shape.
+        tests.append(self._make_test_case(512, 1024, 1024, 32))
+
+        return tests
+
+    def generate_performance_test(self) -> Dict[str, Any]:
+        torch.manual_seed(0)
+        # Matches the FP4 matmul shapes reported in AutoKernel community results.
+        return self._make_test_case(2048, 8192, 3072, 32)
diff --git a/challenges/hard/86_fp4_matmul/starter/starter.cu b/challenges/hard/86_fp4_matmul/starter/starter.cu
@@ -0,0 +1,7 @@
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <stdint.h>
+
+// x, w_q, scales, y are device pointers
+extern "C" void solve(const __half* x, const uint8_t* w_q, const __half* scales, __half* y, int M,
+                      int N, int K, int group_size) {}
diff --git a/challenges/hard/86_fp4_matmul/starter/starter.cute.py b/challenges/hard/86_fp4_matmul/starter/starter.cute.py
@@ -0,0 +1,17 @@
+import cutlass
+import cutlass.cute as cute
+
+
+# x, w_q, scales, y are tensors on the GPU
+@cute.jit
+def solve(
+    x: cute.Tensor,
+    w_q: cute.Tensor,
+    scales: cute.Tensor,
+    y: cute.Tensor,
+    M: cute.Int32,
+    N: cute.Int32,
+    K: cute.Int32,
+    group_size: cute.Int32,
+):
+    pass
diff --git a/challenges/hard/86_fp4_matmul/starter/starter.jax.py b/challenges/hard/86_fp4_matmul/starter/starter.jax.py
@@ -0,0 +1,17 @@
+import jax
+import jax.numpy as jnp
+
+
+# x, w_q, scales are tensors on GPU
+@jax.jit
+def solve(
+    x: jax.Array,
+    w_q: jax.Array,
+    scales: jax.Array,
+    M: int,
+    N: int,
+    K: int,
+    group_size: int,
+) -> jax.Array:
+    # return output tensor directly
+    pass
diff --git a/challenges/hard/86_fp4_matmul/starter/starter.mojo b/challenges/hard/86_fp4_matmul/starter/starter.mojo
@@ -0,0 +1,17 @@
+from std.gpu.host import DeviceContext
+from std.memory import UnsafePointer
+
+
+# x, w_q, scales, y are device pointers
+@export
+def solve(
+    x: UnsafePointer[Float16, MutExternalOrigin],
+    w_q: UnsafePointer[UInt8, MutExternalOrigin],
+    scales: UnsafePointer[Float16, MutExternalOrigin],
+    y: UnsafePointer[Float16, MutExternalOrigin],
+    M: Int32,
+    N: Int32,
+    K: Int32,
+    group_size: Int32,
+) raises:
+    pass
diff --git a/challenges/hard/86_fp4_matmul/starter/starter.pytorch.py b/challenges/hard/86_fp4_matmul/starter/starter.pytorch.py
@@ -0,0 +1,15 @@
+import torch
+
+
+# x, w_q, scales, y are tensors on the GPU
+def solve(
+    x: torch.Tensor,
+    w_q: torch.Tensor,
+    scales: torch.Tensor,
+    y: torch.Tensor,
+    M: int,
+    N: int,
+    K: int,
+    group_size: int,
+):
+    pass
diff --git a/challenges/hard/86_fp4_matmul/starter/starter.triton.py b/challenges/hard/86_fp4_matmul/starter/starter.triton.py
@@ -0,0 +1,17 @@
+import torch
+import triton
+import triton.language as tl
+
+
+# x, w_q, scales, y are tensors on the GPU
+def solve(
+    x: torch.Tensor,
+    w_q: torch.Tensor,
+    scales: torch.Tensor,
+    y: torch.Tensor,
+    M: int,
+    N: int,
+    K: int,
+    group_size: int,
+):
+    pass
