morphottention 0.1.0__tar.gz

morphottention-0.1.0/CMakeLists.txt

@@ -0,0 +1,113 @@
+cmake_minimum_required(VERSION 3.24)
+
+if(NOT DEFINED CMAKE_CUDA_COMPILER)
+    if(DEFINED ENV{CUDA_HOME})
+        set(CMAKE_CUDA_COMPILER "$ENV{CUDA_HOME}/bin/nvcc" CACHE FILEPATH "" FORCE)
+    else()
+        find_program(_NVCC nvcc)
+        if(_NVCC)
+            set(CMAKE_CUDA_COMPILER "${_NVCC}" CACHE FILEPATH "" FORCE)
+        endif()
+    endif()
+endif()
+
+
+project(morphottention LANGUAGES CXX CUDA)
+
+#  23 host / 20 device
+set(CMAKE_CXX_STANDARD 23)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_CUDA_STANDARD 20)
+set(CMAKE_CUDA_STANDARD_REQUIRED ON)
+
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+set(CMAKE_CUDA_SEPARABLE_COMPILATION OFF)
+
+set(CMAKE_CUDA_ARCHITECTURES 90 100 120)
+set(TORCH_CUDA_ARCH_LIST "9.0;10.0;12.0")
+
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+
+set(Python_FIND_VIRTUALENV FIRST)
+
+find_package(Python REQUIRED COMPONENTS Interpreter Development.Module)
+
+execute_process(
+        COMMAND ${Python_EXECUTABLE} -c "import torch; print(torch.utils.cmake_prefix_path)"
+        OUTPUT_VARIABLE TORCH_CMAKE_PREFIX
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+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
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+)
+if(NOT EXISTS "${TORCH_PYTHON_LIBRARY}")
+    message(FATAL_ERROR "Could not locate ${_TORCH_PYTHON_LIB_NAME} at: ${TORCH_PYTHON_LIBRARY}")
+endif()
+get_filename_component(TORCH_LIB_DIR "${TORCH_PYTHON_LIBRARY}" DIRECTORY)
+
+find_package(CUDAToolkit REQUIRED)
+
+Python_add_library(_C MODULE WITH_SOABI
+        csrc/cuda/binder.cpp
+        csrc/cuda/dispatch.cpp
+
+        csrc/cuda/attention/attention.cpp
+        csrc/cuda/attention/attention.cu
+)
+
+# -lineinfo on Debug; full device debug (-G) only when explicitly sanitizing.
+option(CUDA_SANITIZE "Build device code with -G for compute-sanitizer" OFF)
+target_compile_options(_C PRIVATE
+        $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>:--threads=0>
+        $<$<AND:$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>,$<CONFIG:Debug>>:-lineinfo>
+        $<$<AND:$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>,$<CONFIG:Debug>,$<BOOL:${CUDA_SANITIZE}>>:-G>
+)
+
+
+if(MSVC)
+    target_compile_options(_C PRIVATE
+            $<$<COMPILE_LANGUAGE:CXX>:/Zc:preprocessor>
+            $<$<COMPILE_LANGUAGE:CUDA>:-Xcompiler=/Zc:preprocessor>
+    )
+endif()
+
+target_include_directories(_C PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}/csrc
+        ${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"
+    )
+endif()
+
+install(TARGETS _C LIBRARY DESTINATION morphottention)

morphottention-0.1.0/PKG-INFO

@@ -0,0 +1,62 @@
+Metadata-Version: 2.2
+Name: morphottention
+Version: 0.1.0
+Summary: CUDA attention & morphology kernels for PyTorch (sm_90/100/120)
+Keywords: attention,cuda,pytorch,transformer,morfologfy,flash-attention,ViT
+Author-Email: Vedran Hrabar <vedran.hrabar@outlook.com>
+License: MIT
+Classifier: Development Status :: 2 - Pre-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.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.14.*
+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
+```
+
+## 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.1.0/README.md

@@ -0,0 +1,27 @@
+# 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
+```
+
+## 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.1.0/csrc/cuda/attention/attention.cpp

@@ -0,0 +1,71 @@
+#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& X, const torch::Tensor& grad_out) {
+    TORCH_CHECK(X.is_cuda(), "Input tensor must be a CUDA tensor");
+    TORCH_CHECK(grad_out.is_cuda(), "Gradient output tensor must be a CUDA tensor");
+
+    auto X_contig = X.contiguous();
+    auto grad_out_contig = grad_out.contiguous();
+
+    const int B = static_cast<int>(X_contig.size(0));
+    const int N = static_cast<int>(X_contig.size(1));
+    const int D = static_cast<int>(X_contig.size(2));
+
+    return {X};
+}

morphottention-0.1.0/csrc/cuda/attention/attention.cu

@@ -0,0 +1,235 @@
+#include "attention.cuh"
+
+#include <cuda_pipeline.h>
+
+#include <cfloat>
+
+template <int HEAD_DIM_V, int CUBE_M, int BR, int BC>
+__global__ void morpho_attention_forward_kernel(const __half* __restrict__ X, const __half* __restrict__ W_phi,
+                                                const __half* __restrict__ gate_q, const __half* __restrict__ gate_k,
+                                                const __half* __restrict__ W_V, __half* __restrict__ out, float* lse,
+                                                int B, int N, int D, int H, float scale) {
+
+    static_assert(HEAD_DIM_V <= BR, "HEAD_DIM_V must be <= BR");
+    static_assert(HEAD_DIM_V <= BC, "HEAD_DIM_V must be <= BC");
+
+    // NOLINTNEXTLINE(readability-suspicious-call-argument)
+    auto [t, b, tid] = get_coords();
+    const unsigned int warp = tid / 32;
+    const unsigned int lane = tid & 31;
+
+    const unsigned int bh = t;
+    const unsigned int batch = bh / static_cast<unsigned int>(H);
+    const unsigned int head = bh % static_cast<unsigned int>(H);
+
+    const int q_row_start = static_cast<int>(b) * BR;
+    const int q_row_end = min(q_row_start + BR, N);
+    const int q_rows = q_row_end - q_row_start;
+    if (q_rows <= 0) {
+        return;
+    }
+
+    const int ldphi = H * CUBE_M;
+    const int ldv = H * HEAD_DIM_V;
+
+    const __half* X_b = X + batch * N * D;
+    const __half* W_phi_h = W_phi + head * CUBE_M;
+    const __half* W_V_h = W_V + head * HEAD_DIM_V;
+    const __half* gate_q_h = gate_q + head * CUBE_M;
+    const __half* gate_k_h = gate_k + head * CUBE_M;
+    __half* out_s = out + batch * N * D + head * HEAD_DIM_V;
+    float* lse_s = lse + bh * N;
+
+    // smem carve
+    constexpr int XT_ROWS = (BR > BC) ? BR : BC;
+    extern __shared__ __align__(16) unsigned char smem[];
+    auto* q_mem = reinterpret_cast<__half*>(smem); // q codes [BR, CUBE_M]
+    auto* k_mem = q_mem + BR * CUBE_M;             // k codes [BC, CUBE_M]
+    auto* v_mem = k_mem + BC * CUBE_M;             // V tile  [BC, HEAD_DIM_V]
+    auto* xt_mem = v_mem + BC * HEAD_DIM_V;        // token staging [max(BR,BC), D]
+
+    auto* s_mem = reinterpret_cast<float*>(xt_mem + XT_ROWS * D); // S/P fp32 tile [BR, BC]
+    auto* p_h_mem = reinterpret_cast<__half*>(s_mem + BR * BC);   // P fp16 tile [BR, BC]
+    auto* state_mem = reinterpret_cast<float*>(p_h_mem + BR * BC);
+
+    float* max_mem = state_mem;            // [BR]
+    float* lse_mem = state_mem + BR;       // [BR]
+    float* corr_mem = state_mem + 2 * BR;  // [BR]
+    float* qbias_mem = state_mem + 3 * BR; // [BR]
+    float* cbias_mem = state_mem + 4 * BR; // [BC]
+    float* o_acc = cbias_mem + BC;         // [BR, HEAD_DIM_V]
+    float* pv_mem = s_mem;                 // PV scratch
+
+    // runing vars in SMEM
+    for (unsigned int i = tid; i < BR * HEAD_DIM_V; i += blockDim.x) {
+        o_acc[i] = 0.0f;
+    }
+    for (unsigned int row = tid; row < BR; row += blockDim.x) {
+        max_mem[row] = -FLT_MAX;
+        lse_mem[row] = 0.0f;
+    }
+
+    // move X to SRAM
+    arch::impl::smem_load(xt_mem, X_b + q_row_start * D, static_cast<unsigned int>(q_rows),
+                          static_cast<unsigned int>(D), tid, BR);
+    __syncthreads();
+    // project Q to the gated unit-hypercube
+    // q = gate_q @ sigma(W_phi.t, x_q) -> q_mem, q_bias = gate_k x q_i
+    arch::impl::project_phi<BR, CUBE_M, WARPS>(xt_mem, W_phi_h, gate_q_h, gate_k_h, q_mem, qbias_mem, s_mem, D, ldphi,
+                                               warp, lane);
+
+    const float scale_log2 = scale * LOG2E;
+
+    // loop over KV vlocks
+    unsigned int n_tiles = (static_cast<unsigned int>(N) + BC - 1) / BC;
+    const int q_global_max = q_row_start + q_rows - 1;
+    if (CAUSAL) {
+        // full mask after last row-q
+        n_tiles = min(n_tiles, static_cast<unsigned int>(q_global_max) / BC + 1u);
+    }
+
+    for (unsigned int tile = 0; tile < n_tiles; tile++) {
+        const int kv_block_start = static_cast<int>(tile) * BC;
+        const int kv_valid = min(BC, N - kv_block_start);
+
+        // stage K/V tokens
+        arch::impl::smem_load(xt_mem, X_b + kv_block_start * D, static_cast<unsigned int>(kv_valid),
+                              static_cast<unsigned int>(D), tid, BC);
+        __syncthreads();
+
+        // project K to the gated hypercube
+        // k = gate_k @ sigma(W_phi.t, x_j) -> k_mem, c_bias = gate_q x k_j
+        arch::impl::project_phi<BC, CUBE_M, WARPS>(xt_mem, W_phi_h, gate_k_h, gate_q_h, k_mem, cbias_mem, s_mem, D,
+                                                   ldphi, warp, lane);
+
+        // v = W_v x x_j -> s_mem -> v_curr
+        arch::impl::matmul<BC, HEAD_DIM_V, WARPS, false>(s_mem, xt_mem, W_V_h, D, D, ldv, HEAD_DIM_V, 1.0f);
+        __syncthreads();
+        for (unsigned int i = tid; i < BC * HEAD_DIM_V; i += blockDim.x) {
+            v_mem[i] = __float2half(s_mem[i]);
+        }
+        __syncthreads();
+
+        // S_raw = QK.T -> s_mem
+        arch::impl::matmul<BR, BC, WARPS, true>(s_mem, q_mem, k_mem, CUBE_M, CUBE_M, CUBE_M, BC, 1.0f);
+        __syncthreads();
+
+        // S = (2S_raw − qbias[i] − cbias[j])·scale_log2e
+        // symetry -> -inf mask
+        for (unsigned int row = 0; row < BR / WARPS; row++) {
+            const unsigned int row_cor = row + warp * (BR / WARPS);
+            const int q_global = q_row_start + static_cast<int>(row_cor);
+
+            float tile_max = -FLT_MAX;
+            for (unsigned int c = lane; c < BC; c += 32) {
+                float s = (2.0f * s_mem[row_cor * BC + c] - qbias_mem[row_cor] - cbias_mem[c]) * scale_log2;
+                const int k_global = kv_block_start + static_cast<int>(c);
+
+                bool drop = (static_cast<int>(c) >= kv_valid) || (static_cast<int>(row_cor) >= q_rows);
+                if (CAUSAL && k_global > q_global)
+                    drop = true;
+                if (MASK_DIAG && k_global == q_global)
+                    drop = true;
+
+                s = drop ? -FLT_MAX : s;
+                s_mem[row_cor * BC + c] = s;
+                tile_max = fmaxf(tile_max, s);
+            }
+            tile_max = warpMax(tile_max);
+
+            const float max_prev = max_mem[row_cor];
+            const float lse_prev = lse_mem[row_cor];
+            const float max_new = fmaxf(max_prev, tile_max);
+            const float corr = exp2f(max_prev - max_new);
+
+            float partial = 0.0f;
+            for (unsigned int c = lane; c < BC; c += 32) {
+                const float s = s_mem[row_cor * BC + c];
+                const float p = (s == -FLT_MAX) ? 0.0f : exp2f(s - max_new);
+                p_h_mem[row_cor * BC + c] = __float2half(p);
+                partial += p;
+            }
+            const float lse_new = lse_prev * corr + warpAllReduceSum(partial);
+
+            if (lane == 0) {
+                corr_mem[row_cor] = corr;
+                lse_mem[row_cor] = lse_new;
+                max_mem[row_cor] = max_new;
+            }
+        }
+        __syncthreads();
+
+        // PV -> s_mem -> pv_mem
+        arch::impl::matmul<BR, HEAD_DIM_V, WARPS, false>(pv_mem, p_h_mem, v_mem, BC, BC, HEAD_DIM_V, HEAD_DIM_V, 1.0f);
+        __syncthreads();
+
+        // O = ALpha O + PV
+        for (unsigned int i = tid; i < BR * HEAD_DIM_V; i += blockDim.x) {
+            o_acc[i] = o_acc[i] * corr_mem[i / HEAD_DIM_V] + pv_mem[i];
+        }
+        __syncthreads();
+    }
+
+    // store to GMEM
+    for (unsigned int row = 0; row < BR / WARPS; row++) {
+        const unsigned int row_cor = row + warp * (BR / WARPS);
+        const unsigned int g_row = q_row_start + row_cor;
+        if (row_cor >= static_cast<unsigned int>(q_rows) || g_row >= static_cast<unsigned int>(N)) {
+            continue;
+        }
+        const float denom = lse_mem[row_cor];
+        const float inv = (denom > 0.0f) ? 1.0f / denom : 0.0f;
+        for (unsigned int d = lane; d < HEAD_DIM_V; d += 32) {
+            out_s[g_row * D + d] = __float2half(o_acc[row_cor * HEAD_DIM_V + d] * inv);
+        }
+        if (lane == 0) {
+            lse_s[g_row] = (denom > 0.0f) ? max_mem[row_cor] + log2f(denom) : -FLT_MAX;
+        }
+    }
+}
+
+void attention_forward_kernel_launcher(const __half* X, const __half* W_phi, const __half* gate_q, const __half* gate_k,
+                                       const __half* W_V, __half* out, float* lse, const int B, const int N,
+                                       const int D, const int H, const int cube_m, const int head_dim_v,
+                                       const float scale, cudaStream_t stream) {
+
+    // checkers
+    TORCH_CHECK(B > 0 && N > 0 && D > 0 && H > 0, "Invalid dimensions");
+    TORCH_CHECK(D % H == 0, "D must be divisible by H");
+    TORCH_CHECK(H * head_dim_v == D, "H * head_dim_v must equal D");
+    TORCH_CHECK(cube_m == CUBE_M_FWD && head_dim_v == HEAD_DIM_V_FWD, "kernel built for fixed (cube_m, head_dim_v)");
+    TORCH_CHECK(X && W_phi && gate_q && gate_k && W_V && out && lse, "Null pointer");
+
+    // smem carve
+    const size_t smem = sizeof(__half) * (BR_FWD * CUBE_M_FWD +      // q codes
+                                          BC_FWD * CUBE_M_FWD +      // k codes
+                                          BC_FWD * HEAD_DIM_V_FWD +  // V tile
+                                          BR_FWD * BC_FWD)           // P fp16 tile
+                        + sizeof(__half) * (XT_ROWS * D)             // shared raw-token staging tile
+                        + sizeof(float) * (BR_FWD * BC_FWD +         // S/raw fp32 tile
+                                           4 * BR_FWD +              // max + lse + corr + qbias
+                                           BC_FWD +                  // column bias
+                                           BR_FWD * HEAD_DIM_V_FWD); // O accumulator
+
+    // kernel instance
+    const auto kernel = morpho_attention_forward_kernel<HEAD_DIM_V_FWD, CUBE_M_FWD, BR_FWD, BC_FWD>;
+
+    // smem alloc check
+    static int fwd_cached = -1;
+    configure_kernel_smem(kernel, smem, fwd_cached, "forward_kernel");
+
+    // launch
+    dim3 grid(B * H, (N + BR_FWD - 1) / BR_FWD);
+    dim3 block(BLOCK_SIZE);
+
+    kernel<<<grid, block, smem, stream>>>(X, W_phi, gate_q, gate_k, W_V, out, lse, B, N, D, H, scale);
+
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+void attention_backward_kernel_launcher(const __half* grad_out, const __half* X, const __half* dX, const __half* d_se,
+                                        int B, int N, int D, cudaStream_t stream) {
+    TORCH_CHECK(B > 0 && N > 0 && D > 0, "Invalid dimensions");
+
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+}

morphottention-0.1.0/csrc/cuda/attention/attention.cuh

@@ -0,0 +1,36 @@
+#ifndef MORPHOTTENTION_ATTENTION_CUH
+#define MORPHOTTENTION_ATTENTION_CUH
+
+#include <cuda_runtime.h>
+
+#include <cuda_fp16.h>
+
+#ifdef __CUDACC__
+#include <cuda/sm120/matmul.cuh>
+#include <cuda/sm120/project.cuh>
+#include <cuda/sm120/smem.cuh>
+#include <cuda/utils/declarations.cuh>
+#include <cuda/utils/smem.cuh>
+#include <cuda/utils/utils.cuh>
+
+#include <c10/cuda/CUDAException.h>
+#include <c10/util/Exception.h>
+
+namespace arch {
+#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 120)
+namespace impl = sm120; // Consumer Blackwell
+#else
+namespace impl = sm120; // fallback
+#endif
+} // namespace arch
+#endif // __CUDACC__
+
+void attention_forward_kernel_launcher(const __half* X, const __half* W_phi, const __half* gate_q, const __half* gate_k,
+                                       const __half* W_V, __half* out, float* lse, const int B, const int N,
+                                       const int D, const int H, const int cube_m, const int head_dim_v,
+                                       const float scale, cudaStream_t stream);
+
+void attention_backward_kernel_launcher(const __half* grad_out, const __half* X, const __half* dX, const __half* d_se,
+                                        int B, int N, int D, cudaStream_t stream);
+
+#endif // MORPHOTTENTION_ATTENTION_CUH

morphottention-0.1.0/csrc/cuda/binder.cpp

@@ -0,0 +1,12 @@
+#include "dispatch.h"
+
+#include <torch/extension.h>
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "Morphottention CUDA attention kernels";
+
+    m.def("forward", &forward, "Attention forward dispatcher", py::arg("X"), py::arg("W_phi"), py::arg("gate_q"),
+          py::arg("gate_k"), py::arg("W_V"), py::arg("H"), py::arg("cube_m"), py::arg("scale"), py::arg("causal"));
+
+    m.def("backward", &backward, "Attention backward dispatcher", py::arg("grad_out"), py::arg("X"));
+}

morphottention-0.1.0/csrc/cuda/dispatch.cpp

@@ -0,0 +1,13 @@
+#include "dispatch.h"
+
+std::vector<torch::Tensor> 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) {
+    return morpho_forward(X, W_phi, gate_q, gate_k, W_V, H, cube_m, scale, causal);
+}
+
+std::vector<torch::Tensor> backward(const torch::Tensor& grad_out, const torch::Tensor& X) {
+    TORCH_CHECK(grad_out.is_cuda() && X.is_cuda(), "Gradient output and X must be a CUDA tensors");
+
+    return morpho_backward(X, grad_out);
+}

morphottention-0.1.0/csrc/cuda/dispatch.h

@@ -0,0 +1,22 @@
+#ifndef MORPHOTTENTION_DISPATCH_H
+#define MORPHOTTENTION_DISPATCH_H
+
+#include <torch/extension.h>
+
+// py-facing dispatchers
+std::vector<torch::Tensor> 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, int64_t H, int64_t cube_m,
+                                   double scale, bool causal);
+
+std::vector<torch::Tensor> backward(const torch::Tensor& grad_out, const torch::Tensor& X);
+
+// CUDA-facing  dispatchers
+
+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, int64_t H, int64_t cube_m, double scale,
+                                          bool causal);
+
+std::vector<torch::Tensor> morpho_backward(const torch::Tensor& grad_out, const torch::Tensor& X);
+
+#endif // MORPHOTTENTION_DISPATCH_H

morphottention-0.1.0/csrc/cuda/morfology/cube.cuh

@@ -0,0 +1,81 @@
+#ifndef MORPHOTTENTION_CUBE_CUH
+#define MORPHOTTENTION_CUBE_CUH
+
+#include <cuda_runtime.h>
+
+#include <math.h>
+
+namespace morph {
+
+enum class CubeProjection { Sigmoid, HardTanh01, Identity };
+
+__device__ __forceinline__ float cube_sigmoid(const float x) {
+    return 1.0f / (1.0f + __expf(-x));
+}
+__device__ __forceinline__ float cube_hardtanh01(const float x) {
+    return fminf(1.0f, fmaxf(0.0f, 0.5f * (x + 1.0f)));
+}
+__device__ __forceinline__ float cube_project(const float x, CubeProjection p) {
+    switch (p) {
+    case CubeProjection::Sigmoid:
+        return cube_sigmoid(x);
+    case CubeProjection::HardTanh01:
+        return cube_hardtanh01(x);
+    default:
+        return x;
+    }
+}
+
+__device__ __forceinline__ float cube_sigmoid_grad(float x) {
+    const float s = cube_sigmoid(x);
+    return s * (1.0f - s);
+}
+__device__ __forceinline__ float cube_hardtanh01_grad(float x) {
+    return (x > -1.0f && x < 1.0f) ? 0.5f : 0.0f;
+}
+__device__ __forceinline__ float cube_project_grad(float x, CubeProjection p) {
+    switch (p) {
+    case CubeProjection::Sigmoid:
+        return cube_sigmoid_grad(x);
+    case CubeProjection::HardTanh01:
+        return cube_hardtanh01_grad(x);
+    default:
+        return 1.0f;
+    }
+}
+
+template <CubeProjection P>
+struct Cube;
+template <>
+struct Cube<CubeProjection::Sigmoid> {
+    static __device__ __forceinline__ float project(const float x) {
+        return cube_sigmoid(x);
+    }
+    static __device__ __forceinline__ float grad(const float x) {
+        return cube_sigmoid_grad(x);
+    }
+};
+
+template <>
+struct Cube<CubeProjection::HardTanh01> {
+    static __device__ __forceinline__ float project(const float x) {
+        return cube_hardtanh01(x);
+    }
+    static __device__ __forceinline__ float grad(const float x) {
+        return cube_hardtanh01_grad(x);
+    }
+};
+
+template <>
+struct Cube<CubeProjection::Identity> {
+    static __device__ __forceinline__ float project(const float x) {
+        return x;
+    }
+    static __device__ __forceinline__ float grad(float) {
+        return 1.0f;
+    }
+};
+
+} // namespace morph
+
+#endif // MORPHOTTENTION_CUBE_CUH