morphottention 0.1.0__tar.gz → 0.2.0__tar.gz

morphottention-0.2.0/PKG-INFO

@@ -0,0 +1,130 @@
+Metadata-Version: 2.2
+Name: morphottention
+Version: 0.2.0
+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.0/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.0/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/extension.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};
+}