Mixer Benchmarks

.github/workflows/install-matrix.yml

@@ -0,0 +1,127 @@
+name: Install Matrix
+
+on:
+  pull_request:
+    paths:
+      - "nvsubquadratic/**"
+      - "pyproject.toml"
+      - "VERSION"
+      - ".github/workflows/install-matrix.yml"
+  push:
+    branches: [main]
+    paths:
+      - "nvsubquadratic/**"
+      - "pyproject.toml"
+      - "VERSION"
+      - ".github/workflows/install-matrix.yml"
+  workflow_dispatch:
+
+# Cancel in-flight runs for the same PR/branch when a new commit is pushed
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
+jobs:
+  # Build the pure-Python wheel once; all matrix legs share the same artifact.
+  # This catches sdist/wheel-build regressions independently of the install check.
+  build-wheel:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Build distribution
+        run: |
+          python -m pip install --upgrade pip build
+          python -m build
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/*.whl
+          retention-days: 1
+
+  # Stronger than compileall: tests a real wheel install into a fresh venv,
+  # not just byte-compilation. Catches broken core deps, missing extras, and
+  # import-graph regressions (e.g. CUDA kernel leaking into core).
+  install-clean:
+    name: Install clean (py${{ matrix.python-version }})
+    needs: build-wheel
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        # Keep in sync with `requires-python` and the classifiers in pyproject.toml.
+        python-version: ["3.10", "3.11", "3.12"]
+    steps:
+      - uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist
+
+      # ── Step 2 ── install the built wheel (not `pip install .`) into a fresh venv.
+      # torch/torchvision are fetched from the PyTorch CPU index so the install
+      # succeeds on a CPU-only runner without the CUDA toolkit.
+      - name: Install wheel into clean venv (core deps, CPU torch)
+        run: |
+          python -m pip install --upgrade pip
+          python -m venv /tmp/clean
+          /tmp/clean/bin/pip install --upgrade pip
+          /tmp/clean/bin/pip install \
+            --extra-index-url https://download.pytorch.org/whl/cpu \
+            dist/*.whl
+
+      # ── Step 3 ── import the public operator surface from the clean venv.
+      - name: Import public surface
+        run: |
+          /tmp/clean/bin/python -c "import nvsubquadratic; \
+            import nvsubquadratic.lazy_config, nvsubquadratic.modules.hyena_nd, \
+            nvsubquadratic.modules.ckconv_nd, nvsubquadratic.modules.kernels_nd, \
+            nvsubquadratic.modules.masks_nd; print(nvsubquadratic.__version__)"
+
+      # ── Step 4 ── assert the CUDA kernel was NOT pulled in by the core install.
+      # subquadratic-ops-torch-cu12 lives in the [cuda] extra (requires nvcc);
+      # finding it in a core install means the extra boundary was broken.
+      - name: Assert CUDA kernel absent from core install
+        run: |
+          /tmp/clean/bin/python -c "import importlib.util, sys; \
+            sys.exit(1) if importlib.util.find_spec('subquadratic_ops_torch') else None"
+
+      # ── Step 5 ── lean-path check: only torch + einops + omegaconf + numpy.
+      # Asserts the importable operator surface genuinely needs nothing heavier
+      # than these four runtime deps (no datasets, pytorch-lightning, wandb, …).
+      # This is exactly the dep footprint that downstream projects (e.g. MONAI) rely on.
+      - name: Lean-path import (torch + einops + omegaconf + numpy only)
+        run: |
+          python -m venv /tmp/lean
+          /tmp/lean/bin/pip install --upgrade pip
+          /tmp/lean/bin/pip install --no-deps dist/*.whl
+          /tmp/lean/bin/pip install \
+            --index-url https://download.pytorch.org/whl/cpu \
+            "torch>=2.10.0,<2.11.0" "torchvision>=0.25.0,<0.26.0"
+          /tmp/lean/bin/pip install einops omegaconf numpy
+          /tmp/lean/bin/python -c "import nvsubquadratic; \
+            import nvsubquadratic.lazy_config, nvsubquadratic.modules.hyena_nd, \
+            nvsubquadratic.modules.ckconv_nd, nvsubquadratic.modules.kernels_nd, \
+            nvsubquadratic.modules.masks_nd; print(nvsubquadratic.__version__)"
+
+      # ── Step 6 ── [cuda] extra: metadata/resolution smoke (informational).
+      # We do NOT build the kernel (no nvcc). The goal is to catch metadata
+      # breakage (e.g. the extra disappearing from wheel METADATA), not to
+      # compile the extension. Resolution failure because the sdist isn't on a
+      # reachable index is expected and logged, not treated as a job failure.
+      - name: "[cuda] extra metadata smoke (informational, no nvcc)"
+        run: |
+          /tmp/clean/bin/pip install --upgrade "pip>=23.1"
+          /tmp/clean/bin/pip install --dry-run --find-links dist/ "nvsubquadratic[cuda]" \
+            || echo "NOTE: [cuda] resolution incomplete — expected if subquadratic-ops-torch-cu12 is not on a reachable index or requires nvcc to build"

Dockerfile

@@ -74,6 +74,83 @@ RUN MAX_JOBS="${MAX_JOBS}" pip install -v --disable-pip-version-check --no-cache
     --config-settings "--build-option=--cuda_ext" \
     git+https://github.com/NVIDIA/apex.git
 
+# ── Mamba baseline: mamba-ssm + causal-conv1d from source ─────────────────────
+# Optional comparison baseline (e.g. the 2D forward-time benchmark's Mamba2
+# series); NOT a project dependency, so it is installed explicitly here rather
+# than via an extra. Placed before COPY of the full tree so unrelated code
+# changes do not trigger a rebuild. The tiny patch script is copied first so we
+# can rewrite upstream setup.py before compiling.
+#
+# Upstream hardcodes -gencode for sm_75..sm_120 and ignores TORCH_CUDA_ARCH_LIST,
+# which OOMs / gcc-ICEs under QEMU arm64. We clone, patch gencodes to match
+# TORCH_CUDA_ARCH_LIST, and honor NVCC_THREADS (arm64 build sets this to 1).
+#
+# mamba-ssm depends on an unpinned `torch`, so a plain install lets pip swap
+# torch / nvidia-cudnn-cu12 out from under the 2.10 stack — which breaks cuDNN
+# (CUDNN_STATUS_NOT_INITIALIZED at runtime). Freeze the current torch/nvidia/
+# triton versions into a constraints file so the mamba install cannot touch them.
+# Gated by INSTALL_MAMBA (default false — benchmark-only baseline, not a project dep).
+# Build with --build-arg INSTALL_MAMBA=true for the benchmark image; the default lean
+# build skips the (slow, QEMU-heavy) source compile and the benchmark's Mamba2 series
+# then reports 'unavailable' and is omitted. The patch script is COPYed unconditionally
+# (tiny, harmless if unused; COPY cannot be gated).
+ARG NVCC_THREADS=4
+ARG INSTALL_MAMBA=false
+COPY scripts/docker/patch_mamba_cuda_arches.py /tmp/patch_mamba_cuda_arches.py
+RUN if [ "${INSTALL_MAMBA}" != "true" ]; then \
+        echo "INSTALL_MAMBA=${INSTALL_MAMBA} — skipping Mamba baseline (mamba-ssm / causal-conv1d)"; \
+    else \
+        pip freeze | grep -iE '^(torch|torchvision|torchaudio|nvidia-|triton|pytorch-triton)' > /tmp/mamba-constraints.txt && \
+        cat /tmp/mamba-constraints.txt && \
+        git clone --depth 1 https://github.com/Dao-AILab/causal-conv1d.git /tmp/causal-conv1d && \
+        git clone --depth 1 https://github.com/state-spaces/mamba.git /tmp/mamba && \
+        python /tmp/patch_mamba_cuda_arches.py /tmp/causal-conv1d/setup.py /tmp/mamba/setup.py && \
+        MAX_JOBS="${MAX_JOBS}" NVCC_THREADS="${NVCC_THREADS}" \
+        CAUSAL_CONV1D_FORCE_BUILD=TRUE MAMBA_FORCE_BUILD=TRUE \
+        pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \
+        -c /tmp/mamba-constraints.txt \
+        /tmp/causal-conv1d /tmp/mamba && \
+        rm -rf /tmp/causal-conv1d /tmp/mamba ; \
+    fi
+
+# ── FlashAttention-4 baseline (Blackwell / Hopper) ────────────────────────────
+# Optional baseline for the forward-time benchmark's FlashAttention-4 series
+# (attn_impl="fa4"); NOT a project dependency. FA4 is a pure-Python CuTe-DSL wheel
+# that JIT-compiles its kernel at runtime on the target GPU, so — unlike apex /
+# mamba above — there is NO CUDA source build here: this layer is cheap and
+# QEMU-safe. It needs a Hopper/Blackwell GPU + CUDA >= 12.3 at *run* time (the
+# build host needs no GPU; the JIT fires on the first forward). Alpha release, so
+# --pre; on a CUDA 13 base use the `[cu13]` extra on both pins below.
+#
+# VERSION LOCK: flash-attn-4 pins an EXACT matching CuTe-DSL dev build (b23 ->
+# nvidia-cutlass-dsl==4.6.0.dev0). A skew between the two crashes the JIT with
+# "fmax() takes 2 positional arguments but 3 given" in flash_attn/cute/softmax.py.
+# So (1) install the matched pair explicitly, and (2) EXCLUDE nvidia-cutlass-dsl
+# from the torch/CUDA constraints pin — otherwise a DSL already present in the base
+# image gets frozen to the wrong version and strands FA4 on a mismatched API. The
+# rest of the nvidia/torch/triton stack is still pinned (the cuDNN-clobber guard).
+# Gated by INSTALL_FA4 (default false — benchmark-only baseline, not a project dep).
+# Build with --build-arg INSTALL_FA4=true for the benchmark image; the default lean
+# build skips the alpha flash-attn-4 wheel and the benchmark's FA4 series then reports
+# 'unavailable' and is omitted. When enabled, the install IS fatal on failure (so a bad
+# version pin fails the build loudly); only the post-install import probe is best-effort
+# (braced so its `|| echo` cannot mask an install failure) since importing the CuTe DSL
+# may touch the CUDA driver on a GPU-less build host.
+ARG FLASH_ATTN4_VERSION=4.0.0b23
+ARG CUTLASS_DSL_VERSION=4.6.0.dev0
+ARG INSTALL_FA4=false
+RUN if [ "${INSTALL_FA4}" != "true" ]; then \
+        echo "INSTALL_FA4=${INSTALL_FA4} — skipping FlashAttention-4 baseline"; \
+    else \
+        pip freeze | grep -iE '^(torch|torchvision|torchaudio|nvidia-|triton|pytorch-triton)' \
+          | grep -viE '^nvidia-cutlass-dsl' > /tmp/fa4-constraints.txt && \
+        pip install --no-cache-dir --pre -c /tmp/fa4-constraints.txt \
+          "nvidia-cutlass-dsl==${CUTLASS_DSL_VERSION}" "flash-attn-4==${FLASH_ATTN4_VERSION}" && \
+        { python -c "import nvidia_cutlass_dsl as c, flash_attn.cute; from flash_attn.cute import flash_attn_func; \
+print('flash-attn-4 import OK / cutlass-dsl', c.__version__)" \
+          || echo "WARNING: flash-attn-4 import check failed at build (JIT may probe CUDA); verify on-GPU." ; } ; \
+    fi
+
 # ── Dev deps: cached until requirements-dev.txt changes ──────────────────────
 COPY requirements-dev.txt .
 RUN pip install --no-cache-dir -r requirements-dev.txt

README.md

@@ -79,10 +79,11 @@ docker build -t nvsubquadratic:dev .
 docker run --gpus all -p 8888:8888 -v $(pwd):/workspaces/nvSubquadratic nvsubquadratic:dev
 ```
 
-The Dockerfile builds NVIDIA Apex from source for a broad set of NVIDIA archs by default (`7.0;7.5;8.0;8.6;8.9;9.0;10.0;12.0` — Volta through Blackwell). Two build-args let you tune the compile:
+The Dockerfile builds NVIDIA Apex from source for a broad set of NVIDIA archs by default (`7.0;7.5;8.0;8.6;8.9;9.0;10.0;12.0` — Volta through Blackwell). Build-args let you tune the compile:
 
-- `TORCH_CUDA_ARCH_LIST` — narrow to your GPU(s) to speed up the build (e.g. `9.0` for H100, `8.6` for A6000, `8.9` for L4).
-- `MAX_JOBS` — number of parallel nvcc jobs. Defaults to unconstrained. Set to a small number (e.g. `2`) if the build OOMs (typical under qemu emulation).
+- `TORCH_CUDA_ARCH_LIST` — narrow to your GPU(s) to speed up the build (e.g. `9.0` for H100, `8.6` for A6000, `8.9` for L4). Also applied to the patched `mamba-ssm` / `causal-conv1d` source builds (upstream otherwise hardcodes sm_75..sm_120).
+- `MAX_JOBS` — number of parallel nvcc/ninja jobs. Defaults to unconstrained. Set to `1` if the build OOMs or gcc ICEs (typical under qemu emulation for arm64).
+- `NVCC_THREADS` — nvcc `--threads` for the mamba stack (default `4`; use `1` under QEMU).
 
 ```bash
 docker build \
@@ -92,13 +93,13 @@ docker build \
 
 ### Enroot (SLURM clusters)
 
-For SLURM deployments that use enroot/pyxis, [`scripts/slurm/enroot/build_sqsh.sh`](scripts/slurm/enroot/build_sqsh.sh) builds the Docker image and converts it to an enroot `.sqsh` in one step. It selects the right `TORCH_CUDA_ARCH_LIST` and `MAX_JOBS` per platform:
+For SLURM deployments that use enroot/pyxis, [`scripts/slurm/enroot/build_sqsh.sh`](scripts/slurm/enroot/build_sqsh.sh) builds the Docker image and converts it to an enroot `.sqsh` in one step. It selects the right `TORCH_CUDA_ARCH_LIST`, `MAX_JOBS`, and `NVCC_THREADS` per platform:
 
 ```bash
 # H100 (x86-64, default)
 scripts/slurm/enroot/build_sqsh.sh
 
-# GB200 (ARM64) — uses qemu emulation on an x86 build host
+# GB200 (ARM64) — QEMU on x86: keep ≥64GB free RAM+swap; defaults MAX_JOBS=1 NVCC_THREADS=1
 PLATFORM=arm64 scripts/slurm/enroot/build_sqsh.sh
 ```