add built softmax attention kernel

.gitignore

@@ -1,10 +1,13 @@
 # Python-generated files
 __pycache__/
 *.py[oc]
-build/
 dist/
 wheels/
 *.egg-info
 
 # Virtual environments
 .venv
+
+# Nix build symlink (do NOT commit; build/ is the real artifact and IS tracked)
+result
+result-*

build/torch-cuda/__init__.py

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

build/torch-cuda/_kernel.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)

build/torch-cuda/_ops.py

@@ -0,0 +1,8 @@
+import torch
+ops = torch.ops._my_softmax_function_2af76fc
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_my_softmax_function_2af76fc::{op_name}"

build/torch-cuda/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),
+    )

build/torch-cuda/metadata.json

@@ -0,0 +1,5 @@
+{
+  "python-depends": [
+    "nvidia-cutlass-dsl"
+  ]
+}

build/torch-cuda/my_softmax_function/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+from pathlib import Path
+from types import ModuleType
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

result

@@ -1 +0,0 @@
-/nix/store/4mkpzj01sp4ixa9wbig2hiqfaavhxqck-torch-ext-bundle