[WIP] SwinUNETR + HyenaND: subquadratic alternative to windowed self-attention

.github/workflows/cicd_tests.yml

@@ -237,6 +237,55 @@ jobs:
         python -m unittest -v
       shell: bash
 
+  hyena-dep:  # Optional HyenaND dependency + the no-CUDA Hyena tests.
+    # nvsubquadratic >= 0.1.1 supports Python >= 3.10 and keeps its CUDA-kernel sdist
+    # (subquadratic-ops-torch-cu12) plus the megatron / dali / timm packages in opt-in
+    # extras, so it installs on a CPU runner.  We still pass ``--no-deps`` deliberately:
+    #   (1) the HyenaND operators import only torch + einops + omegaconf at runtime, so
+    #       skipping the (still batteries-included: datasets/lightning/wandb) core deps
+    #       keeps this job lean; and
+    #   (2) nvsubquadratic pins torch>=2.10,<2.11, which would otherwise upgrade/clash
+    #       with the torch this job (and MONAI's matrix) installs.
+    # CUDA-required Hyena tests skip cleanly here; the GPU surface is covered by
+    # ``.github/workflows/pythonapp-hyena-gpu.yml``.
+    runs-on: ubuntu-latest
+    timeout-minutes: 30
+    steps:
+    - name: Clean unused tools
+      run: |
+        find  /opt/hostedtoolcache/* -maxdepth 0 ! -name 'Python' -exec rm -rf {} \;
+        sudo rm -rf /usr/share/dotnet
+        sudo rm -rf /usr/local/lib/android
+        sudo rm -rf /opt/ghc /usr/local/.ghcup
+        sudo docker system prune -f
+    - uses: actions/checkout@v6
+    - name: Set up Python ${{ env.PYTHON_VER1 }}
+      uses: actions/setup-python@v6
+      with:
+        python-version: ${{ env.PYTHON_VER1 }}
+        cache: 'pip'
+    - name: Install dependencies + nvsubquadratic (no-deps)
+      run: |
+        python -m pip install --upgrade pip wheel
+        python -m pip install torch==${PYTORCH_VER1} torchvision==${TORCHVISION_VER1}
+        python -m pip install --no-build-isolation -r requirements-dev.txt
+        python -m pip install -e .
+        # nvsubquadratic runtime imports need only torch + einops + omegaconf; install
+        # the package itself without its core dependency tree (see job comment above).
+        python -m pip install omegaconf
+        python -m pip install --no-deps 'nvsubquadratic>=0.1.1'
+        python -m pip list
+      shell: bash
+    - name: Run Hyena tests (CUDA-required cases skip cleanly)
+      run: |
+        python -c "from monai.networks.blocks.hyena import is_nvsubquadratic_available; \
+                   assert is_nvsubquadratic_available(), 'nvsubquadratic must be importable'"
+        python -m pytest -v \
+          tests/networks/blocks/test_hyena_block.py \
+          tests/networks/nets/test_hyena_nd_unetr.py \
+          tests/networks/nets/test_swin_unetr.py
+      shell: bash
+
   packaging:  # Test package generation
     runs-on: ubuntu-latest
     env:

.github/workflows/pythonapp-hyena-gpu.yml

@@ -0,0 +1,74 @@
+# Optional self-hosted GPU CI for the HyenaND test surface.
+#
+# This workflow exercises the CUDA-required Hyena tests
+# (tests/networks/blocks/test_hyena_block.py CUDA cases, the four-paper-variant
+# forward and gradient cases in tests/networks/nets/test_swin_unetr.py and
+# tests/networks/nets/test_hyena_nd_unetr.py, the SwinUNETR(use_hyena=False)
+# golden-hash backward-compat regression, and sliding-window inference).
+#
+# Disabled by default (``if: false``).  To enable:
+#   1. Ensure a self-hosted runner with the labels below is available, AND
+#   2. Ensure the runner has CUDA-capable hardware visible (the existing
+#      ``pythonapp-gpu.yml`` uses ``--gpus all`` against ``[self-hosted, linux,
+#      x64, common]``).  Reuse that pool if possible.
+#   3. Flip ``if: false`` to ``if: github.event.pull_request.merged != true``
+#      (mirroring ``pythonapp-gpu.yml``'s gating pattern).
+#
+# nvsubquadratic (Hyena's optional dep) requires Python >= 3.10; any NGC base with
+# Python >= 3.10 works. The accelerated [cuda] kernels build against the container nvcc.
+
+name: hyena-gpu
+
+on:
+  workflow_dispatch:
+
+concurrency:
+  group: hyena-gpu-${{ github.event.pull_request.number || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  GPU-Hyena:
+    if: ${{ false }}  # See header for enable instructions.
+    strategy:
+      matrix:
+        environment:
+          # NGC PyTorch 25.05 ships Python 3.12 and CUDA 12.5.  Bump as needed.
+          - "NGC25.05+PY312"
+        include:
+          - environment: NGC25.05+PY312
+            base: "nvcr.io/nvidia/pytorch:25.05-py3"
+    container:
+      image: ${{ matrix.base }}
+      options: --gpus all --env NVIDIA_DISABLE_REQUIRE=true
+    runs-on: [self-hosted, linux, x64, common]
+    steps:
+    - uses: actions/checkout@v6
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip wheel
+        python -c "import sys; assert sys.version_info >= (3, 10), f'Python >= 3.10 required for nvsubquadratic, got {sys.version}'"
+        python -m pip install -r requirements-dev.txt
+        python -m pip install -e .
+        # Install nvsubquadratic with --no-deps: the default torch_fft path needs only
+        # torch + einops + omegaconf, and nvsubquadratic pins torch>=2.10,<2.11 which can
+        # clash with the container's torch.  To exercise the accelerated fused CUDA
+        # kernels instead, install the [cuda] extra (subquadratic-ops-torch-cu12, builds
+        # against the container's nvcc) and set fft_backend="subq_ops" in the tests.
+        python -m pip install omegaconf
+        python -m pip install --no-deps 'nvsubquadratic>=0.1.1'
+        python -m pip list
+      shell: bash
+    - name: Verify CUDA + nvsubquadratic
+      run: |
+        nvidia-smi
+        python -c "import torch; assert torch.cuda.is_available(); print('CUDA OK:', torch.cuda.get_device_name(0))"
+        python -c "from monai.networks.blocks.hyena import is_nvsubquadratic_available; \
+                   assert is_nvsubquadratic_available(), 'nvsubquadratic must be importable'"
+      shell: bash
+    - name: Run Hyena test suite (CUDA + no-CUDA)
+      run: |
+        python -m pytest -v \
+          tests/networks/blocks/test_hyena_block.py \
+          tests/networks/nets/test_hyena_nd_unetr.py \
+          tests/networks/nets/test_swin_unetr.py
+      shell: bash

CHANGELOG.md

@@ -4,6 +4,11 @@ All notable changes to MONAI are documented in this file.
 The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
+### Added
+* `HyenaMixer`, `HyenaTransformerBlock`, and `DepthwiseFFTConv{2,3}d` in `monai.networks.blocks`: subquadratic O(N log N) alternatives to windowed self-attention, backed by the HyenaND operator from the optional `nvsubquadratic` package.
+* `HyenaNDUNETR` (`monai.networks.nets.HyenaNDUNETR`): thin `SwinUNETR` subclass with a `get_variant(name)` classmethod for the three Hyena variants (`HHHH`, `HAHA`, `HHAA`) from the NeurIPS 2026 paper "Native Multi-Dimensional Subquadratic Operators via Input Dependent Long Convolutions" (paper id 26539).
+* `SwinUNETR.use_hyena` and `SwinUNETR.hyena_stages` kwargs to thread HyenaND blocks through any subset of Swin stages. Default `use_hyena=False` preserves bit-identical forward behavior of the existing code path.
+* New `[hyena]` extras_require in setup.cfg (`pip install monai[hyena]`).
 
 ## [1.6.0] - 2026-06-12
 

docs/source/installation.md

@@ -254,10 +254,15 @@ Since MONAI v0.2.0, the extras syntax such as `pip install 'monai[nibabel]'` is
 - The options are
 
 ```
-[nibabel, skimage, scipy, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, cucim, openslide, pandas, einops, transformers, mlflow, clearml, matplotlib, tensorboardX, tifffile, imagecodecs, pyyaml, fire, jsonschema, ninja, pynrrd, pydicom, h5py, nni, optuna, onnx, onnxruntime, zarr, lpips, pynvml, huggingface_hub]
+[nibabel, skimage, scipy, pillow, tensorboard, gdown, ignite, torchvision, itk, tqdm, lmdb, psutil, cucim, openslide, pandas, einops, transformers, mlflow, clearml, matplotlib, tensorboardX, tifffile, imagecodecs, pyyaml, fire, jsonschema, ninja, pynrrd, pydicom, h5py, nni, optuna, onnx, onnxruntime, zarr, lpips, pynvml, huggingface_hub, hyena]
 ```
 
 which correspond to `nibabel`, `scikit-image`,`scipy`, `pillow`, `tensorboard`,
-`gdown`, `pytorch-ignite`, `torchvision`, `itk`, `tqdm`, `lmdb`, `psutil`, `cucim`, `openslide-python`, `pandas`, `einops`, `transformers`, `mlflow`, `clearml`, `matplotlib`, `tensorboardX`, `tifffile`, `imagecodecs`, `pyyaml`, `fire`, `jsonschema`, `ninja`, `pynrrd`, `pydicom`, `h5py`, `nni`, `optuna`, `onnx`, `onnxruntime`, `zarr`, `lpips`, `nvidia-ml-py`, `huggingface_hub` and `pyamg` respectively.
+`gdown`, `pytorch-ignite`, `torchvision`, `itk`, `tqdm`, `lmdb`, `psutil`, `cucim`, `openslide-python`, `pandas`, `einops`, `transformers`, `mlflow`, `clearml`, `matplotlib`, `tensorboardX`, `tifffile`, `imagecodecs`, `pyyaml`, `fire`, `jsonschema`, `ninja`, `pynrrd`, `pydicom`, `h5py`, `nni`, `optuna`, `onnx`, `onnxruntime`, `zarr`, `lpips`, `nvidia-ml-py`, `huggingface_hub`, `pyamg`, and `nvsubquadratic` respectively.
+
+The `hyena` extra pulls in [`nvsubquadratic`](https://github.com/NVIDIA-BioNeMo/nvSubquadratic),
+required by `HyenaNDUNETR` / `HyenaMixer` / `HyenaTransformerBlock` (subquadratic
+O(N log N) alternatives to windowed self-attention). Install with
+`pip install 'monai[hyena]'`.
 
 - `pip install 'monai[all]'` installs all the optional dependencies.

docs/source/networks.rst

@@ -129,6 +129,23 @@ Blocks
 .. autoclass:: TransformerBlock
     :members:
 
+`Hyena Mixer`
+~~~~~~~~~~~~~
+.. autoclass:: HyenaMixer
+    :members:
+
+`Hyena Transformer Block`
+~~~~~~~~~~~~~~~~~~~~~~~~~
+.. autoclass:: HyenaTransformerBlock
+    :members:
+
+`Depthwise FFT Convolution`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+.. autoclass:: DepthwiseFFTConv2d
+    :members:
+.. autoclass:: DepthwiseFFTConv3d
+    :members:
+
 `UNETR Block`
 ~~~~~~~~~~~~~
 .. autoclass:: UnetrBasicBlock
@@ -591,6 +608,11 @@ Nets
 .. autoclass:: SwinUNETR
     :members:
 
+`HyenaNDUNETR`
+~~~~~~~~~~~~~~
+.. autoclass:: HyenaNDUNETR
+    :members:
+
 `BasicUNet`
 ~~~~~~~~~~~
 .. autoclass:: BasicUNet

monai/networks/blocks/__init__.py

@@ -26,6 +26,13 @@
 from .encoder import BaseEncoder
 from .fcn import FCN, GCN, MCFCN, Refine
 from .feature_pyramid_network import ExtraFPNBlock, FeaturePyramidNetwork, LastLevelMaxPool, LastLevelP6P7
+from .hyena import (
+    DepthwiseFFTConv2d,
+    DepthwiseFFTConv3d,
+    HyenaMixer,
+    HyenaTransformerBlock,
+    is_nvsubquadratic_available,
+)
 from .localnet_block import LocalNetDownSampleBlock, LocalNetFeatureExtractorBlock, LocalNetUpSampleBlock
 from .mednext_block import MedNeXtBlock, MedNeXtDownBlock, MedNeXtOutBlock, MedNeXtUpBlock
 from .mlp import MLPBlock