morphottention 0.1.0__tar.gz → 0.2.1__tar.gz

{morphottention-0.1.0 → morphottention-0.2.1}/CMakeLists.txt

@@ -41,21 +41,14 @@ list(APPEND CMAKE_PREFIX_PATH "${TORCH_CMAKE_PREFIX}")
 find_package(Torch REQUIRED CONFIG)
 
 
-if(WIN32)
-    set(_TORCH_PYTHON_LIB_NAME "torch_python.lib")
-else()
-    set(_TORCH_PYTHON_LIB_NAME "libtorch_python.so")
-endif()
-
 execute_process(
-        COMMAND ${Python_EXECUTABLE} -c "import pathlib, torch; print(pathlib.Path(torch.__file__).resolve().parent / 'lib' / '${_TORCH_PYTHON_LIB_NAME}')"
-        OUTPUT_VARIABLE TORCH_PYTHON_LIBRARY
+        COMMAND ${Python_EXECUTABLE} -c "import pathlib, torch; print(pathlib.Path(torch.__file__).resolve().parent / 'lib')"
+        OUTPUT_VARIABLE TORCH_LIB_DIR
         OUTPUT_STRIP_TRAILING_WHITESPACE
 )
-if(NOT EXISTS "${TORCH_PYTHON_LIBRARY}")
-    message(FATAL_ERROR "Could not locate ${_TORCH_PYTHON_LIB_NAME} at: ${TORCH_PYTHON_LIBRARY}")
+if(NOT EXISTS "${TORCH_LIB_DIR}")
+    message(FATAL_ERROR "Could not locate torch lib dir at: ${TORCH_LIB_DIR}")
 endif()
-get_filename_component(TORCH_LIB_DIR "${TORCH_PYTHON_LIBRARY}" DIRECTORY)
 
 find_package(CUDAToolkit REQUIRED)
 
@@ -85,25 +78,19 @@ endif()
 
 target_include_directories(_C PRIVATE
         ${CMAKE_CURRENT_SOURCE_DIR}/csrc
+        ${CMAKE_CURRENT_SOURCE_DIR}/csrc/compat
         ${TORCH_INCLUDE_DIRS}
         ${Python_INCLUDE_DIRS}
         ${CUDAToolkit_INCLUDE_DIRS}
 )
 
-target_compile_definitions(_C PRIVATE
-        TORCH_EXTENSION_NAME=_C
-)
-
 target_link_libraries(_C PRIVATE
         ${TORCH_LIBRARIES}
-        ${TORCH_PYTHON_LIBRARY}
-        Python::Module
         CUDA::cudart
         CUDA::cublas
 )
 
 if(NOT WIN32)
-    target_link_options(_C PRIVATE "-Wl,--no-as-needed")
     set_target_properties(_C PROPERTIES
             BUILD_RPATH   "${TORCH_LIB_DIR}"
             INSTALL_RPATH "$ORIGIN/../torch/lib"

morphottention-0.2.1/PKG-INFO

@@ -0,0 +1,130 @@
+Metadata-Version: 2.2
+Name: morphottention
+Version: 0.2.1
+Summary: Mathematical Morphology-based self-attention module for PyTorch (CUDA) using Flash-style kernel fusion.
+Keywords: attention,cuda,pytorch,transformer,morphology,flash-attention,ViT
+Author-Email: Vedran Hrabar <vedran.hrabar@outlook.com>
+License: MIT
+Classifier: Development Status :: 3 - Alpha
+Classifier: Environment :: GPU
+Classifier: Environment :: GPU :: NVIDIA CUDA
+Classifier: Environment :: GPU :: NVIDIA CUDA :: 13
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Programming Language :: C++
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Image Recognition
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Typing :: Typed
+Project-URL: repository, https://github.com/vhrabar/morphottention
+Project-URL: documentation, https://github.com/vhrabar/morphottention/wiki
+Project-URL: Bug Tracker, https://github.com/vhrabar/morphottention/issues
+Requires-Python: >=3.12
+Requires-Dist: torch>=2.12
+Description-Content-Type: text/markdown
+
+# Morphottention
+Mathematical Morphology-based self-attention module for PyTorch using Flash-style kernel fusion.
+
+## Install
+
+Prebuilt wheels are published for CPython 3.14 on Linux (x86_64, aarch64) and
+Windows (x86_64). A working CUDA-enabled PyTorch (`torch >= 2.12`) must already
+be installed in the environment.
+
+```bash
+pip install morphottention
+```
+
+## Usage
+
+The package exposes an `nn.Module` (`MorphoAttention`), a functional entry point
+(`morpho_attention`), and the raw autograd bridge (`MorphoAttentionFunction`).
+All inputs must be CUDA tensors; the module defaults to `float16`.
+
+### As an `nn.Module`
+
+```python
+import torch
+from morphottention import MorphoAttention
+
+attn = MorphoAttention(
+    dim=256,       # model dimension D
+    num_heads=8,   # number of attention heads H
+    cube_m=16,     # hypercube width per head
+    scale=1.0,     # softmax temperature
+    causal=False,  # casual masking flag
+    device="cuda",
+)
+
+x = torch.randn(2, 128, 256, dtype=torch.float16, device="cuda")  # (B, N, D)
+out = attn(x)          # (B, N, D)
+out.sum().backward()
+```
+
+### Functional form
+
+```python
+from morphottention import morpho_attention
+
+out = morpho_attention(
+    x, 
+    W_phi, 
+    gate_q, 
+    gate_k, 
+    W_V,
+    num_heads=8, cube_m=16, scale=1.0, 
+    causal=False,
+)
+```
+
+### Raw autograd bridge
+
+```python
+import torch
+from morphottention import MorphoAttentionFunction
+
+B, N, D, H, cube_m = 2, 128, 256, 8, 16
+
+x = torch.randn(B, N, D, dtype=torch.float16, device="cuda", requires_grad=True)
+W_phi = torch.randn(D, H * cube_m, dtype=torch.float16, device="cuda", requires_grad=True)
+gate_q = torch.ones(H, cube_m, dtype=torch.float16, device="cuda", requires_grad=True)
+gate_k = torch.ones(H, cube_m, dtype=torch.float16, device="cuda", requires_grad=True)
+W_V = torch.randn(D, D, dtype=torch.float16, device="cuda", requires_grad=True)
+
+out = MorphoAttentionFunction.apply(
+    x, W_phi, gate_q, gate_k, W_V,
+    H, cube_m, 1.0, False,  # num_heads, cube_m, scale, causal
+)  # (B, N, D)
+out.sum().backward()
+```
+
+`W_phi` has shape `(D, H * cube_m)`, `W_V` has shape `(D, D)`, and `gate_q` /
+`gate_k` each have shape `(H, cube_m)`.
+
+## Building from source
+
+Requires the CUDA 13.X toolkit (`nvcc`) and a matching `torch` build:
+
+```bash
+uv sync --package morphottention --no-dev --group build
+uv build --package morphottention --wheel --no-build-isolation
+```
+
+## License
+
+MIT
+
+Copyright © 2026 Vedran Hrabar.

morphottention-0.2.1/README.md

@@ -0,0 +1,93 @@
+# Morphottention
+Mathematical Morphology-based self-attention module for PyTorch using Flash-style kernel fusion.
+
+## Install
+
+Prebuilt wheels are published for CPython 3.14 on Linux (x86_64, aarch64) and
+Windows (x86_64). A working CUDA-enabled PyTorch (`torch >= 2.12`) must already
+be installed in the environment.
+
+```bash
+pip install morphottention
+```
+
+## Usage
+
+The package exposes an `nn.Module` (`MorphoAttention`), a functional entry point
+(`morpho_attention`), and the raw autograd bridge (`MorphoAttentionFunction`).
+All inputs must be CUDA tensors; the module defaults to `float16`.
+
+### As an `nn.Module`
+
+```python
+import torch
+from morphottention import MorphoAttention
+
+attn = MorphoAttention(
+    dim=256,       # model dimension D
+    num_heads=8,   # number of attention heads H
+    cube_m=16,     # hypercube width per head
+    scale=1.0,     # softmax temperature
+    causal=False,  # casual masking flag
+    device="cuda",
+)
+
+x = torch.randn(2, 128, 256, dtype=torch.float16, device="cuda")  # (B, N, D)
+out = attn(x)          # (B, N, D)
+out.sum().backward()
+```
+
+### Functional form
+
+```python
+from morphottention import morpho_attention
+
+out = morpho_attention(
+    x, 
+    W_phi, 
+    gate_q, 
+    gate_k, 
+    W_V,
+    num_heads=8, cube_m=16, scale=1.0, 
+    causal=False,
+)
+```
+
+### Raw autograd bridge
+
+```python
+import torch
+from morphottention import MorphoAttentionFunction
+
+B, N, D, H, cube_m = 2, 128, 256, 8, 16
+
+x = torch.randn(B, N, D, dtype=torch.float16, device="cuda", requires_grad=True)
+W_phi = torch.randn(D, H * cube_m, dtype=torch.float16, device="cuda", requires_grad=True)
+gate_q = torch.ones(H, cube_m, dtype=torch.float16, device="cuda", requires_grad=True)
+gate_k = torch.ones(H, cube_m, dtype=torch.float16, device="cuda", requires_grad=True)
+W_V = torch.randn(D, D, dtype=torch.float16, device="cuda", requires_grad=True)
+
+out = MorphoAttentionFunction.apply(
+    x, W_phi, gate_q, gate_k, W_V,
+    H, cube_m, 1.0, False,  # num_heads, cube_m, scale, causal
+)  # (B, N, D)
+out.sum().backward()
+```
+
+`W_phi` has shape `(D, H * cube_m)`, `W_V` has shape `(D, D)`, and `gate_q` /
+`gate_k` each have shape `(H, cube_m)`.
+
+## Building from source
+
+Requires the CUDA 13.X toolkit (`nvcc`) and a matching `torch` build:
+
+```bash
+uv sync --package morphottention --no-dev --group build
+uv build --package morphottention --wheel --no-build-isolation
+```
+
+## License
+
+MIT
+
+Copyright © 2026 Vedran Hrabar.

morphottention-0.2.1/build.toml

@@ -0,0 +1,37 @@
+[general]
+name = "morphottention"
+backends = ["cuda"]
+license = "MIT"
+version = 0
+
+[general.hub]
+repo-id = "vhrabar/morphottention"
+
+[torch]
+include = ["csrc", "torch-ext"]
+src = [
+  "csrc/cuda/binder.cpp",
+  "csrc/cuda/dispatch.cpp",
+  "csrc/cuda/dispatch.h",
+]
+
+[kernel.morpho]
+backend = "cuda"
+cuda-capabilities = ["8.9", "9.0", "10.0", "12.0", "12.0a"]
+depends = ["torch"]
+include = ["csrc"]
+src = [
+  "csrc/cuda/attention/attention.cpp",
+  "csrc/cuda/attention/attention.cu",
+  "csrc/cuda/attention/attention.cuh",
+  "csrc/cuda/dispatch.h",
+  "csrc/cuda/morfology/cube.cuh",
+  "csrc/cuda/morfology/soft_morph.cuh",
+  "csrc/cuda/sm120/matmul.cuh",
+  "csrc/cuda/sm120/project.cuh",
+  "csrc/cuda/sm120/smem.cuh",
+  "csrc/cuda/utils/declarations.cuh",
+  "csrc/cuda/utils/reductions.cuh",
+  "csrc/cuda/utils/smem.cuh",
+  "csrc/cuda/utils/utils.cuh",
+]

morphottention-0.2.1/csrc/compat/registration.h

@@ -0,0 +1,14 @@
+#ifndef MORPHOTTENTION_COMPAT_REGISTRATION_H
+#define MORPHOTTENTION_COMPAT_REGISTRATION_H
+
+// Compatibility shim for the local CMake / PyPI build.
+
+#include <torch/library.h>
+
+#define TORCH_LIBRARY_EXPAND(NAME, MODULE) TORCH_LIBRARY(NAME, MODULE)
+
+#define TORCH_EXTENSION_NAME morphottention
+
+#define REGISTER_EXTENSION(NAME)
+
+#endif // MORPHOTTENTION_COMPAT_REGISTRATION_H

morphottention-0.2.1/csrc/cuda/attention/attention.cpp

@@ -0,0 +1,118 @@
+#include "attention.cuh"
+
+#include <cuda/dispatch.h>
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/torch.h>
+
+auto check = [](const torch::Tensor& t, const char* name) {
+    TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+    TORCH_CHECK(t.scalar_type() == torch::kHalf, name, " must be float16");
+    TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+};
+
+std::vector<torch::Tensor> morpho_forward(const torch::Tensor& X, const torch::Tensor& W_phi,
+                                          const torch::Tensor& gate_q, const torch::Tensor& gate_k,
+                                          const torch::Tensor& W_V, const int64_t H, const int64_t cube_m,
+                                          const double scale, const bool causal) {
+    // checkers
+    check(X, "X");
+    check(W_phi, "W_phi");
+    check(gate_q, "gate_q");
+    check(gate_k, "gate_k");
+    check(W_V, "W_V");
+
+    // shape managment
+    TORCH_CHECK(X.dim() == 3, "X must be [B, N, D]");
+    const int64_t B = X.size(0);
+    const int64_t N = X.size(1);
+    const int64_t D = X.size(2);
+
+    TORCH_CHECK(H > 0 && D % H == 0, "D must be divisible by H");
+    const int64_t head_dim_v = D / H;
+
+    TORCH_CHECK(W_phi.dim() == 2 && W_phi.size(0) == D && W_phi.size(1) == H * cube_m, "W_phi must be [D, H*cube_m]");
+    TORCH_CHECK(W_V.dim() == 2 && W_V.size(0) == D && W_V.size(1) == H * head_dim_v, "W_V must be [D, H*head_dim_v]");
+    TORCH_CHECK(gate_q.dim() == 2 && gate_q.size(0) == H && gate_q.size(1) == cube_m, "gate_q must be [H, cube_m]");
+    TORCH_CHECK(gate_k.dim() == 2 && gate_k.size(0) == H && gate_k.size(1) == cube_m, "gate_k must be [H, cube_m]");
+
+    auto out = torch::empty_like(X);
+    auto lse = torch::empty({B * H, N}, X.options().dtype(torch::kFloat32));
+
+    // launcher
+    const c10::cuda::CUDAStreamGuard guard(c10::cuda::getCurrentCUDAStream());
+
+    attention_forward_kernel_launcher(reinterpret_cast<const __half*>(X.data_ptr<at::Half>()),
+                                      reinterpret_cast<const __half*>(W_phi.data_ptr<at::Half>()),
+                                      reinterpret_cast<const __half*>(gate_q.data_ptr<at::Half>()),
+                                      reinterpret_cast<const __half*>(gate_k.data_ptr<at::Half>()),
+                                      reinterpret_cast<const __half*>(W_V.data_ptr<at::Half>()),
+                                      reinterpret_cast<__half*>(out.data_ptr<at::Half>()), lse.data_ptr<float>(),
+                                      static_cast<int>(B), static_cast<int>(N), static_cast<int>(D),
+                                      static_cast<int>(H), static_cast<int>(cube_m), static_cast<int>(head_dim_v),
+                                      static_cast<float>(scale), c10::cuda::getCurrentCUDAStream());
+
+    return {out, lse};
+}
+
+std::vector<torch::Tensor> morpho_backward(const torch::Tensor& grad_out, const torch::Tensor& X,
+                                           const torch::Tensor& W_phi, const torch::Tensor& gate_q,
+                                           const torch::Tensor& gate_k, const torch::Tensor& W_V,
+                                           const torch::Tensor& out, const torch::Tensor& lse, const int64_t H,
+                                           const int64_t cube_m, const double scale, const bool causal) {
+    // checkers
+    check(grad_out, "grad_out");
+    check(X, "X");
+    check(W_phi, "W_phi");
+    check(gate_q, "gate_q");
+    check(gate_k, "gate_k");
+    check(W_V, "W_V");
+    check(out, "out");
+    TORCH_CHECK(lse.is_cuda(), "lse must be a CUDA tensor");
+    TORCH_CHECK(lse.scalar_type() == torch::kFloat32, "lse must be float32");
+    TORCH_CHECK(lse.is_contiguous(), "lse must be contiguous");
+
+    // shape managment
+    TORCH_CHECK(X.dim() == 3, "X must be [B, N, D]");
+    const int64_t B = X.size(0);
+    const int64_t N = X.size(1);
+    const int64_t D = X.size(2);
+
+    TORCH_CHECK(H > 0 && D % H == 0, "D must be divisible by H");
+    const int64_t head_dim_v = D / H;
+
+    TORCH_CHECK(grad_out.sizes() == X.sizes(), "grad_out must match X shape [B, N, D]");
+    TORCH_CHECK(out.sizes() == X.sizes(), "out must match X shape [B, N, D]");
+    TORCH_CHECK(W_phi.dim() == 2 && W_phi.size(0) == D && W_phi.size(1) == H * cube_m, "W_phi must be [D, H*cube_m]");
+    TORCH_CHECK(W_V.dim() == 2 && W_V.size(0) == D && W_V.size(1) == H * head_dim_v, "W_V must be [D, H*head_dim_v]");
+    TORCH_CHECK(gate_q.dim() == 2 && gate_q.size(0) == H && gate_q.size(1) == cube_m, "gate_q must be [H, cube_m]");
+    TORCH_CHECK(gate_k.dim() == 2 && gate_k.size(0) == H && gate_k.size(1) == cube_m, "gate_k must be [H, cube_m]");
+    TORCH_CHECK(lse.dim() == 2 && lse.size(0) == B * H && lse.size(1) == N, "lse must be [B*H, N]");
+
+    // grad outputs (zero-initialised: the kernel accumulates into these via atomics)
+    auto dX = torch::zeros_like(X);
+    auto dW_phi = torch::zeros_like(W_phi);
+    auto d_gate_q = torch::zeros_like(gate_q);
+    auto d_gate_k = torch::zeros_like(gate_k);
+    auto dW_V = torch::zeros_like(W_V);
+
+    // launcher
+    const c10::cuda::CUDAStreamGuard guard(c10::cuda::getCurrentCUDAStream());
+
+    attention_backward_kernel_launcher(
+        reinterpret_cast<const __half*>(grad_out.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(X.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(W_phi.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(gate_q.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(gate_k.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(W_V.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(out.data_ptr<at::Half>()), lse.data_ptr<float>(),
+        reinterpret_cast<__half*>(dX.data_ptr<at::Half>()), reinterpret_cast<__half*>(dW_phi.data_ptr<at::Half>()),
+        reinterpret_cast<__half*>(d_gate_q.data_ptr<at::Half>()),
+        reinterpret_cast<__half*>(d_gate_k.data_ptr<at::Half>()), reinterpret_cast<__half*>(dW_V.data_ptr<at::Half>()),
+        static_cast<int>(B), static_cast<int>(N), static_cast<int>(D), static_cast<int>(H), static_cast<int>(cube_m),
+        static_cast<int>(head_dim_v), static_cast<float>(scale), c10::cuda::getCurrentCUDAStream());
+
+    return {dX, dW_phi, d_gate_q, d_gate_k, dW_V};
+}