morphottention 0.2.0__cp313-cp313-win_amd64.whl

morphottention/_C.pyi

@@ -0,0 +1,27 @@
+import torch
+
+def forward(
+    X: torch.Tensor,
+    W_phi: torch.Tensor,
+    gate_q: torch.Tensor,
+    gate_k: torch.Tensor,
+    W_V: torch.Tensor,
+    H: int,
+    cube_m: int,
+    scale: float,
+    causal: bool,
+) -> list[torch.Tensor]: ...
+def backward(
+    grad_out: torch.Tensor,
+    X: torch.Tensor,
+    W_phi: torch.Tensor,
+    gate_q: torch.Tensor,
+    gate_k: torch.Tensor,
+    W_V: torch.Tensor,
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    H: int,
+    cube_m: int,
+    scale: float,
+    causal: bool,
+) -> list[torch.Tensor]: ...

morphottention/__init__.py

@@ -0,0 +1,3 @@
+from .autograd import MorphoAttention, MorphoAttentionFunction, morpho_attention
+
+__all__ = ["MorphoAttention", "MorphoAttentionFunction", "morpho_attention"]

morphottention/autograd.py

@@ -0,0 +1,141 @@
+"""
+Autograd wrapper and nn.Module around the compiled Morphottention CUDA kernels.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from . import _C
+
+
+class MorphoAttentionFunction(torch.autograd.Function):
+    """
+    Bridges the CUDA attention kernels into autograd.
+    """
+
+    @staticmethod
+    def forward(
+        ctx: torch.autograd.function.FunctionCtx,
+        x: torch.Tensor,
+        W_phi: torch.Tensor,
+        gate_q: torch.Tensor,
+        gate_k: torch.Tensor,
+        W_V: torch.Tensor,
+        H: int,
+        cube_m: int,
+        scale: float,
+        causal: bool,
+    ) -> torch.Tensor:
+        if not x.is_cuda:
+            raise ValueError("MorphoAttention expects a CUDA tensor")
+
+        x = x.contiguous()
+        out, lse = _C.forward(x, W_phi, gate_q, gate_k, W_V, H, cube_m, scale, causal)
+
+        ctx.save_for_backward(x, W_phi, gate_q, gate_k, W_V, out, lse)
+        ctx.H = H  # type: ignore[attr-defined]
+        ctx.cube_m = cube_m  # type: ignore[attr-defined]
+        ctx.scale = scale  # type: ignore[attr-defined]
+        ctx.causal = causal  # type: ignore[attr-defined]
+        return out
+
+    @staticmethod
+    def backward(
+        ctx: torch.autograd.function.FunctionCtx,
+        grad_out: torch.Tensor,
+    ) -> tuple[torch.Tensor | None, ...]:
+        x, W_phi, gate_q, gate_k, W_V, out, lse = ctx.saved_tensors  # type: ignore[attr-defined]
+
+        grad_out = grad_out.contiguous()
+        dX, dW_phi, d_gate_q, d_gate_k, dW_V = _C.backward(
+            grad_out,
+            x,
+            W_phi,
+            gate_q,
+            gate_k,
+            W_V,
+            out,
+            lse,
+            ctx.H,  # type: ignore[attr-defined]
+            ctx.cube_m,  # type: ignore[attr-defined]
+            ctx.scale,  # type: ignore[attr-defined]
+            ctx.causal,  # type: ignore[attr-defined]
+        )
+
+        return dX, dW_phi, d_gate_q, d_gate_k, dW_V, None, None, None, None
+
+
+class MorphoAttention(nn.Module):
+    """
+    Morphological hypercube attention
+    :param dim: model dimension ``D`` (must be divisible by ``num_heads``).
+    :param num_heads: number of attention heads ``H``.
+    :param cube_m: hypercube width per head ``m`` (the score-GEMM contraction dim).
+    :param scale: softmax temperature.
+    :param causal: whether to apply causal masking.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        cube_m: int,
+        *,
+        scale: float = 1.0,
+        causal: bool = False,
+        dtype: torch.dtype = torch.float16,
+        device: torch.device | str | None = None,
+    ) -> None:
+        super().__init__()
+        if dim % num_heads != 0:
+            raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})")
+
+        self.dim = dim
+        self.num_heads = num_heads
+        self.cube_m = cube_m
+        self.scale = scale
+        self.causal = causal
+
+        head_dim_v = dim // num_heads
+        self.W_phi = nn.Parameter(torch.empty(dim, num_heads * cube_m, dtype=dtype, device=device))
+        self.gate_q = nn.Parameter(torch.empty(num_heads, cube_m, dtype=dtype, device=device))
+        self.gate_k = nn.Parameter(torch.empty(num_heads, cube_m, dtype=dtype, device=device))
+        self.W_V = nn.Parameter(torch.empty(dim, num_heads * head_dim_v, dtype=dtype, device=device))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        std = self.dim**-0.5
+        with torch.no_grad():
+            for w in (self.W_phi, self.W_V):
+                w.copy_(torch.empty(w.shape, dtype=torch.float32, device=w.device).normal_(0.0, std))
+            self.gate_q.fill_(1.0)
+            self.gate_k.fill_(1.0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return MorphoAttentionFunction.apply(  # type: ignore[no-any-return, no-untyped-call]
+            x, self.W_phi, self.gate_q, self.gate_k, self.W_V, self.num_heads, self.cube_m, self.scale, self.causal
+        )
+
+
+def morpho_attention(
+    x: torch.Tensor,
+    W_phi: torch.Tensor,
+    gate_q: torch.Tensor,
+    gate_k: torch.Tensor,
+    W_V: torch.Tensor,
+    num_heads: int,
+    cube_m: int,
+    *,
+    scale: float = 1.0,
+    causal: bool = False,
+) -> torch.Tensor:
+    """
+    Apply the Morphottention attention mechanism with autograd support.
+    :param: x: input activations of shape (B, N, D) on a CUDA device.
+    :returns: the attention output, differentiable w.r.t. ``x``.
+    """
+    return MorphoAttentionFunction.apply(  # type: ignore[no-any-return, no-untyped-call]
+        x, W_phi, gate_q, gate_k, W_V, num_heads, cube_m, scale, causal
+    )

morphottention-0.2.0.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,8 @@
+morphottention/__init__.py,sha256=geq-PCIo1BkEd77w7_3-vTUKGHBFFHs-CTzT3tlMwGQ,159
+morphottention/_C.cp313-win_amd64.pyd,sha256=UZvNR1ewZ5Zsjd15H6Ry9TCmo-wCTw5dsOKpjspDXZM,786944
+morphottention/_C.pyi,sha256=jFKwfXOxsS7Y9FjwKBGaVDpi2YqwN_x0cP1mJJCJO_I,572
+morphottention/autograd.py,sha256=w07iFr_y8D7qMBlCbbdjZDtObCP5DbKXRceVe-jwyDY,4633
+morphottention/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+morphottention-0.2.0.dist-info/METADATA,sha256=tCP9s_kBbh8fzl1kLw1a5zUVozbvBxw598XksuFgn-w,4256
+morphottention-0.2.0.dist-info/WHEEL,sha256=UZrbbE4r80xj7Ncfa6JoeTVe-77bdXLkKUA63V8pKWQ,106
+morphottention-0.2.0.dist-info/RECORD,,

morphottention-0.2.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: scikit-build-core 0.12.2
+Root-Is-Purelib: false
+Tag: cp313-cp313-win_amd64
+