initial softmax attention kernel

.gitignore

@@ -0,0 +1,10 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv

.python-version

@@ -0,0 +1 @@
+3.14

README.md

@@ -0,0 +1,56 @@
+---
+tags:
+  - kernel
+license: apache-2.0
+---
+
+# my_softmax_function
+
+Fused **softmax-attention** kernel written in the [CUTLASS Python DSL](https://github.com/NVIDIA/cutlass)
+(`cutlass.cute`), packaged for the Hugging Face [`kernels`](https://github.com/huggingface/kernels) Hub.
+
+Computes, in a single online pass (max-shifted for numerical stability):
+
+```
+out = softmax(scale * Q @ Kᵀ) @ V
+```
+
+## Usage
+
+```python
+import torch
+from kernels import get_kernel
+
+k = get_kernel("monteiro-t/my_softmax_function")   # replace with your repo id
+
+Q = torch.randn(128, 64, device="cuda", dtype=torch.float32)
+K = torch.randn(256, 64, device="cuda", dtype=torch.float32)
+V = torch.randn(256, 64, device="cuda", dtype=torch.float32)
+
+out = k.attention(Q, K, V)          # scale defaults to 1/sqrt(d)
+out = k.attention(Q, K, V, scale=0.5)
+```
+
+## API
+
+| Function | Signature | Notes |
+|---|---|---|
+| `attention(Q, K, V, scale=None)` | `(M,d),(N,d),(N,d) -> (M,d)` | Convenience wrapper; `scale` defaults to `1/sqrt(d)`. |
+| `softmax_attention(Q, K, V, scale)` | registered `torch.ops` op | `scale` required; `torch.compile`-friendly (has a fake/meta impl). |
+
+## Supported shapes / dtypes / hardware
+
+| Property | Support |
+|---|---|
+| Dtype | `float32` only |
+| Layout | 2-D, contiguous (inputs are forced contiguous) |
+| Shapes | `Q: (M, d)`, `K: (N, d)`, `V: (N, d)` → `out: (M, d)` |
+| Batch / heads | **Not supported** (single 2-D attention; no leading batch/head dims) |
+| Masking | **None** (dense attention only) |
+| Backward pass | **None** (inference / forward only) |
+| Device | CUDA only |
+
+## Limitations
+
+Naive one-row-per-thread implementation intended as a reference/demo kernel. No
+batching, masking, dropout, or gradients. Not performance-tuned.

build.toml

@@ -0,0 +1,14 @@
+[general]
+name = "my_softmax_function"
+backends = [
+    "cpu",
+    "cuda",
+    "metal",
+    "rocm",
+    "xpu",
+]
+
+[torch]
+src = []
+
+[kernel]

flake.lock

@@ -0,0 +1,95 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1765121682,
+        "narHash": "sha256-4VBOP18BFeiPkyhy9o4ssBNQEvfvv1kXkasAYd0+rrA=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "65f23138d8d09a92e30f1e5c87611b23ef451bf3",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs"
+      },
+      "locked": {
+        "lastModified": 1775482375,
+        "narHash": "sha256-RUbxfJGs96jwnwSci3+8h08GB2/katOI67yZTCaV+aE=",
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "rev": "dffbce5a048648febb96bbed79f3811ab6d7a577",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1766341660,
+        "narHash": "sha256-4yG6vx7Dddk9/zh45Y2KM82OaRD4jO3HA9r98ORzysA=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "26861f5606e3e4d1400771b513cc63e5f70151a6",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "my_softmax_function — fused softmax-attention kernel";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernel-builder";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

main.py

@@ -0,0 +1,6 @@
+def main():
+    print("Hello from my-softmax-function!")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "my-softmax-function"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.14"
+dependencies = [
+    "nvidia-cutlass-dsl>=4.6.1",
+    "torch>=2.13.0",
+]
+
+[dependency-groups]
+dev = [
+    "kernels",
+    "pytest",
+]

tests/test_attention.py

@@ -0,0 +1,58 @@
+"""Test the built kernel through the `kernels` loader against a torch reference.
+
+Run AFTER building (so build/torch-universal/my_softmax_function exists):
+
+    pytest tests/ -v
+"""
+
+import math
+from pathlib import Path
+
+import pytest
+import torch
+
+from kernels import get_local_kernel
+
+REPO_ROOT = Path(__file__).resolve().parent.parent
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="kernel requires CUDA"
+)
+
+
+def _load():
+    # Loads build/torch-universal/my_softmax_function from the local repo,
+    # exactly the way kernels.get_kernel(...) would load it from the Hub.
+    return get_local_kernel(REPO_ROOT, "my_softmax_function")
+
+
+def _reference(Q, K, V, scale):
+    scores = (Q @ K.transpose(-1, -2)) * scale
+    probs = torch.softmax(scores, dim=-1)
+    return probs @ V
+
+
+@pytest.mark.parametrize("M,N,d", [(1, 1, 1), (17, 33, 8), (128, 256, 64)])
+def test_matches_reference(M, N, d):
+    kernel = _load()
+    torch.manual_seed(0)
+    Q = torch.randn(M, d, device="cuda", dtype=torch.float32)
+    K = torch.randn(N, d, device="cuda", dtype=torch.float32)
+    V = torch.randn(N, d, device="cuda", dtype=torch.float32)
+    scale = 1.0 / math.sqrt(d)
+
+    out = kernel.attention(Q, K, V)
+    ref = _reference(Q, K, V, scale)
+
+    torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
+
+
+def test_custom_scale():
+    kernel = _load()
+    Q = torch.randn(8, 16, device="cuda", dtype=torch.float32)
+    K = torch.randn(12, 16, device="cuda", dtype=torch.float32)
+    V = torch.randn(12, 16, device="cuda", dtype=torch.float32)
+
+    out = kernel.attention(Q, K, V, scale=0.5)
+    ref = _reference(Q, K, V, 0.5)
+    torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)

torch-ext/my_softmax_function/__init__.py

@@ -0,0 +1,8 @@
+"""my_softmax_function — fused softmax-attention kernel (CUTLASS Python DSL).
+
+Public API exposed to `kernels` consumers.
+"""
+
+from ._ops import attention, softmax_attention
+
+__all__ = ["attention", "softmax_attention"]

torch-ext/my_softmax_function/_ops.py

@@ -0,0 +1,68 @@
+"""Public torch custom-op wrapper around the CUTLASS-DSL softmax-attention kernel.
+
+This is the API a Hugging Face `kernels` consumer sees: a normal, registered
+`torch.ops` op that works with autograd tracing / `torch.compile`, not the raw
+`@cute.kernel`.
+"""
+
+import math
+
+import torch
+from cutlass.cute.runtime import from_dlpack
+
+from .attention_v3 import solve
+
+_OP_NAME = "my_softmax_function::softmax_attention"
+
+
+@torch.library.custom_op(_OP_NAME, mutates_args=())
+def softmax_attention(
+    Q: torch.Tensor,        # (M, d)
+    K: torch.Tensor,        # (N, d)
+    V: torch.Tensor,        # (N, d)
+    scale: float,
+) -> torch.Tensor:          # (M, d)
+    """Fused, max-shifted softmax attention: softmax(scale * Q @ K^T) @ V.
+
+    Q, K, V must be contiguous, 2-D, float32, and on the same CUDA device.
+    """
+    if not (Q.is_cuda and K.is_cuda and V.is_cuda):
+        raise ValueError("Q, K, V must be CUDA tensors")
+    if not (Q.dim() == K.dim() == V.dim() == 2):
+        raise ValueError("Q, K, V must be 2-D (M,d)/(N,d)/(N,d)")
+
+    M, d = Q.shape
+    N = K.shape[0]
+    if K.shape[1] != d or V.shape[1] != d or V.shape[0] != N:
+        raise ValueError("shape mismatch between Q/K/V")
+
+    Q = Q.contiguous()
+    K = K.contiguous()
+    V = V.contiguous()
+    output = torch.empty((M, d), dtype=torch.float32, device=Q.device)
+
+    solve(
+        from_dlpack(Q),
+        from_dlpack(K),
+        from_dlpack(V),
+        from_dlpack(output),
+        M,
+        N,
+        d,
+        float(scale),
+    )
+    return output
+
+
+@softmax_attention.register_fake
+def _(Q, K, V, scale):
+    # Shape/dtype/device metadata only — no compute. Lets torch.compile trace.
+    M, d = Q.shape
+    return Q.new_empty((M, d))
+
+
+def attention(Q, K, V, scale=None):
+    """Convenience entry point with a default 1/sqrt(d) scale."""
+    if scale is None:
+        scale = 1.0 / math.sqrt(Q.shape[-1])
+    return torch.ops.my_softmax_function.softmax_attention(Q, K, V, scale)

torch-ext/my_softmax_function/attention_v3.py

@@ -0,0 +1,98 @@
+import cutlass
+import cutlass.cute as cute
+from cutlass.cute.runtime import from_dlpack
+
+NEG_BIG = -3.0e38  # finite stand-in for -inf
+
+
+def transpose_matrix(T: cute.Tensor, i, j):
+    """Element (i, j) of T-transpose: a pure index swap, no data movement."""
+    return T[j, i]
+
+
+def matrix_multiplication(acc, A: cute.Tensor, B: cute.Tensor, row, col, k,
+                          transpose_b=False):
+    """One MAC step of C[row, col] = sum_k A[row, k] * B[k, col].
+    transpose_b is a Python bool, resolved at trace time."""
+    b = transpose_matrix(B, k, col) if transpose_b else B[k, col]
+    return acc + A[row, k] * b
+
+
+def softmax_computation(score, row_max, denom):
+    """One softmax probability, max-shifted for stability. Pass denom=1.0 for the
+    unnormalized numerator (online softmax normalizes once at the end)."""
+    return cute.math.exp(score - row_max) / denom
+
+
+@cute.kernel
+def softmax_attention(
+    Q: cute.Tensor,        # (M, d)
+    K: cute.Tensor,        # (N, d)
+    V: cute.Tensor,        # (N, d)
+    output: cute.Tensor,   # (M, d)
+    M: cutlass.Int32,
+    N: cutlass.Int32,
+    d: cutlass.Int32,
+    host_scale: cutlass.Float32,
+):
+    bx, _, _ = cute.arch.block_idx()
+    bdx, _, _ = cute.arch.block_dim()
+    tx, _, _ = cute.arch.thread_idx()
+
+    row = bx * bdx + tx
+
+    if row < M:
+        scale = host_scale
+
+        # output[row, :] is the running accumulator
+        for k in cutlass.range(d):
+            output[row, k] = cutlass.Float32(0.0)
+
+        running_max = cutlass.Float32(NEG_BIG)
+        running_sum = cutlass.Float32(0.0)
+
+        # single pass: score, max, denominator and V-accumulation together
+        for n in cutlass.range(N):
+            s = cutlass.Float32(0.0)
+            for k in cutlass.range(d):
+                s = matrix_multiplication(s, Q, K, row, n, k, transpose_b=True)
+            s = s * scale
+
+            new_max = running_max
+            if s > new_max:
+                new_max = s
+
+            # rescales everything accumulated under the old max
+            correction = softmax_computation(running_max, new_max, 1.0)
+            p = softmax_computation(s, new_max, 1.0)
+
+            running_sum = running_sum * correction + p
+            running_max = new_max
+
+            for k in cutlass.range(d):
+                output[row, k] = output[row, k] * correction + p * V[n, k]
+
+        # normalize once, with a reciprocal instead of d divisions
+        inv_sum = 1.0 / running_sum
+        for k in cutlass.range(d):
+            output[row, k] = output[row, k] * inv_sum
+
+
+@cute.jit
+def solve(
+    Q: cute.Tensor,
+    K: cute.Tensor,
+    V: cute.Tensor,
+    output: cute.Tensor,
+    M: cutlass.Int32,
+    N: cutlass.Int32,
+    d: cutlass.Int32,
+    host_scale: cutlass.Float32,
+):
+    block_size = 256
+    grid_size = (M + block_size - 1) // block_size
+
+    softmax_attention(Q, K, V, output, M, N, d, host_scale).launch(
+        grid=(grid_size, 1, 1),
+        block=(block_size, 1, 1),
+    )