diff --git a/.github/workflows/install-matrix.yml b/.github/workflows/install-matrix.yml new file mode 100644 index 00000000..f7e3f992 --- /dev/null +++ b/.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" diff --git a/CHANGELOG.md b/CHANGELOG.md index 7e951ccd..22fed49f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,18 @@ All notable changes to nvSubquadratic are documented here. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/) and the project follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## \[Unreleased\] + +### Changed + +- **Bumped the CUDA runtime from 12.9 to 13.0.** The Docker image now builds on + `nvcr.io/nvidia/cuda:13.0.3-devel-ubuntu22.04` with PyTorch `cu130` wheels, and + the CUDA extras target CUDA 13: `[cuda]` → `subquadratic-ops-torch-cu13` + (`>=0.2.1`, resolved from public PyPI — wheels for both aarch64 and x86_64), + `[dali]` → `nvidia-dali-cuda130`. The FlashAttention-4 benchmark baseline uses + the `[cu13]` extra. The default `fft_backend="torch_fft"` path is unaffected and + still needs no CUDA kernel. + ## \[0.1.1\] ### Changed diff --git a/Dockerfile b/Dockerfile index 4297d96a..68b33b3a 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,7 +9,12 @@ # 3. requirements-dev.txt (changes when dev deps change) # 4. COPY . . + pip install (changes on every code push — fast) -FROM nvcr.io/nvidia/cuda:12.9.0-devel-ubuntu22.04 +# CUDA 13.0 toolkit — matched to PyTorch's cu130 (CUDA 13.0) wheels. apex's (and +# mamba's) build-time check requires the base nvcc CUDA to match torch's CUDA +# EXACTLY: a 13.2 base fails apex with "Cuda extensions ... compiled with Cuda 13.0" +# vs nvcc 13.2. torch ships only cu130 (there is no cu132), so CUDA 13.0 is the +# runtime regardless — pin the base to 13.0.x to build the extensions cleanly. +FROM nvcr.io/nvidia/cuda:13.0.3-devel-ubuntu22.04 ARG MINIFORGE_NAME=Miniforge3 ARG MINIFORGE_VERSION=25.3.0-3 @@ -41,10 +46,10 @@ RUN conda install --yes \ && conda clean --all --yes RUN pip install --no-cache-dir \ - torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu129 \ - && pip install --no-cache-dir nvidia-dali-cuda120 \ + torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu130 \ + && pip install --no-cache-dir nvidia-dali-cuda130 \ && pip install --no-cache-dir \ - torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu129 \ + torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu130 \ && conda clean --all --yes # Create ubuntu user with sudo privileges @@ -74,6 +79,84 @@ 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-cu13 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. This CUDA-13 image uses the `[cu13]` extra on both pins below; on a CUDA-12 +# base drop `[cu13]` (the default wheel is cu12). +# +# 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[cu13]==${CUTLASS_DSL_VERSION}" "flash-attn-4[cu13]==${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 @@ -87,10 +170,12 @@ RUN git config --global --add safe.directory /workspaces/nvSubquadratic # (distributed/Megatron CP tests, timm baselines, DALI, subq_ops CUDA kernels) # can run. After the 0.1.1 dependency restructure, megatron-core/timm/etc. are # optional extras ([distributed]/[baselines]/...), so a bare install no longer -# pulls them — [all] restores the complete pre-restructure dependency set. +# pulls them — [all] restores the complete pre-restructure dependency set. The +# [cuda] extra resolves subquadratic-ops-torch-cu13 via the normal pip index chain +# (wheel-stub sdist → prebuilt wheel). RUN pip install --no-cache-dir wheel-stub \ && pip install --no-cache-dir --no-build-isolation ".[all]" \ - --extra-index-url https://download.pytorch.org/whl/cu129 + --extra-index-url https://download.pytorch.org/whl/cu130 # Set up ubuntu user's home directory and permissions RUN chown -R ubuntu:ubuntu /workspaces && \ diff --git a/README.md b/README.md index 4a300031..0580cd5e 100644 --- a/README.md +++ b/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 ``` @@ -125,7 +126,7 @@ bash setup_conda_env.sh conda activate nvsubquadratic ``` -This script creates the `nvsubquadratic` conda environment with Python 3.12 and PyTorch 2.10 (CUDA 12.9), installs all dev dependencies, builds NVIDIA Apex from source, and installs `quack-kernels`. +This script creates the `nvsubquadratic` conda environment with Python 3.12 and PyTorch 2.10 (CUDA 13.0), installs all dev dependencies, builds NVIDIA Apex from source, and installs `quack-kernels`. ### Local Installation (venv) @@ -135,7 +136,7 @@ python3 -m venv venv source venv/bin/activate # Install PyTorch with CUDA support first (before package dependencies) -pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu129 +pip install torch==2.10.0 torchvision==0.25.0 --index-url https://download.pytorch.org/whl/cu130 # Install development dependencies pip install -r requirements-dev.txt diff --git a/benchmarks/README.md b/benchmarks/README.md index b5355f08..c9743ebc 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -123,8 +123,8 @@ Logs go to `logs/`. ## Environment - GPU: NVIDIA H100 SXM 80GB -- PyTorch 2.6.0+cu129 -- CUDA 12.9 +- PyTorch 2.10.0+cu130 +- CUDA 13.0 - Apex 0.1 (FusedLAMB) - QuACK 0.2.10 (fused RMSNorm) - NVIDIA DALI 1.53.0 diff --git a/benchmarks/benchmark_forward_time_nd_resolution.py b/benchmarks/benchmark_forward_time_nd_resolution.py new file mode 100644 index 00000000..01b05b52 --- /dev/null +++ b/benchmarks/benchmark_forward_time_nd_resolution.py @@ -0,0 +1,636 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Forward-time vs resolution benchmark in 1D / 2D / 3D (the ND analogue of Figure 1, right). + +Select the dimensionality with ``--data-dim {1,2,3}``; the input is +``[B, R, ...(N axes)..., C]`` and the token count is ``L = R^N``. Hyena is causal +in 1D (the fused genomics kernel) and non-causal in 2D/3D; 3D falls back to +torch_fft (subq_ops has no 3D kernel). The 2D case (the default) is described below. + +The paper's Figure 1 (right) plots single-operator forward time against **1D +sequence length** (4K->1M) for HyenaND (nSubQ) / attention / Mamba2, showing +attention's O(L^2) wall while HyenaND scales as O(L log L). This script draws +the same comparison with a **2D context** on the x-axis: it sweeps a square +spatial resolution ``R = H = W`` so the token count ``L = R^2`` grows, and times +the forward pass of a single mixer *layer* at each resolution. + + R 64 128 256 512 1024 + L = R^2 4096 16384 65536 262144 1048576 (== 4K 16K 64K 256K 1M) + +Three mixers are compared at a shared ``hidden_dim`` (channels-last +``[B, R, R, hidden_dim]`` input): + + * ``hyena`` -- ``QKVSequenceMixer`` wrapping ``Hyena`` with a 2D + ``CKConvND`` on the ``subq_ops`` CUDA backend (the nSubQ fused FFT-conv). + * ``attention`` -- ``QKVSequenceMixer`` wrapping ``Attention`` with the default + SDPA kernel (PyTorch auto-selects flash / cuDNN / fallback), 2D axial RoPE. + * ``flex`` -- same ``Attention`` layer with the compiled FlexAttention + kernel (``torch.nn.attention.flex_attention``). Requires head_dim >= 16. + * ``fa4`` -- same ``Attention`` layer with FlashAttention-4 (the external + ``flash-attn-4`` CuTe-DSL wheel). Requires head_dim >= 16 and a Hopper/ + Blackwell GPU; if ``flash_attn`` is absent the points show ``unavailable``. + * ``mamba`` -- a **bare** ``Mamba`` (bidirectional Mamba2) that rasterizes + the 2D grid into a 1D scan. Not wrapped in ``QKVSequenceMixer`` (its + ``forward`` takes a single tensor, not q/k/v). + +``attention``/``flex``/``fa4`` are three interchangeable attention *kernels* on +one shared q/k/v + RoPE path; run them together for an apples-to-apples flash +comparison (needs head_dim >= 16, so a wider ``hidden_dim`` than the tiny-width +16M-reach sweep). + +The mixer config builders are reused from ``benchmark_patch_size_2d.py``. We +time a single layer (not the 4-block ``ResidualNetwork`` that script builds): +it is faithful to Figure 1's operator-level timing, uses far less memory (so it +reaches R=1024), and removes the confounds of the residual net's projections, +MLPs and extra norms. + +Output is a JSONL file (one row per ``(mixer, resolution)``); plotting is a +separate step (``scripts/visualization/visualize_forward_time_nd.py``) that +reads the JSONL, so the GPU run needs no matplotlib. + +Local smoke test (any CUDA GPU, no ``subquadratic_ops_torch`` needed):: + + PYTHONPATH=. python benchmarks/benchmark_forward_time_nd_resolution.py \\ + --fft-backend torch_fft --no-compile --batch-size 1 \\ + --resolutions 8 16 32 --mixers hyena attention \\ + --num-warmup 2 --num-iters 3 --output /tmp/smoke_2d.jsonl + +GB200 production run (fused nSubQ kernels, all three mixers):: + + PYTHONPATH=. python benchmarks/benchmark_forward_time_nd_resolution.py \\ + --fft-backend subq_ops --dtype bf16 --batch-size 1 --hidden-dim 256 \\ + --resolutions 64 128 256 512 1024 \\ + --output benchmarks/results/forward_time.jsonl +""" + +from __future__ import annotations + +import argparse +import gc +import json +import math +import sys +import time +import traceback +from pathlib import Path +from typing import Any + +import torch + + +_BENCH_DIR = Path(__file__).resolve().parent +_PROJECT_ROOT = _BENCH_DIR.parent +for _p in (str(_BENCH_DIR), str(_PROJECT_ROOT)): + if _p not in sys.path: + sys.path.insert(0, _p) + +# Reuse the concrete mixer-config builders (backward-compatibly extended to take +# per-resolution kwargs). ``benchmarks/`` has no __init__.py, hence the direct +# module import after the sys.path insert above. +from benchmark_patch_size_2d import ( + _attention_mixer_cfg, + _hyena_mixer_cfg, + _mamba_mixer_cfg, +) + +from nvsubquadratic.lazy_config import instantiate + + +MIXER_CHOICES = ("attention", "flex", "fa4", "hyena", "mamba") +# Attention kernel per mixer key: SDPA (auto cuDNN/flash), compiled FlexAttention, +# or FlashAttention-4 (external flash_attn). All share the same q/k/v + RoPE path. +_ATTN_IMPL = {"attention": "sdpa", "flex": "flex", "fa4": "fa4"} +DTYPE_MAP = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} +# 2D axial RoPE needs head_dim divisible by 4; 1D by 2; 3D by 6. +_ROPE_DIVISOR = {1: 2, 2: 4, 3: 6} + + +# ─── Module construction (fresh per resolution) ─────────────────────────────── + + +def build_module( + name: str, + *, + hidden_dim: int, + resolution: int, + fft_backend: str, + grid_type: str, + num_heads: int, + attn_rope: bool, + mamba_headdim: int, + mamba_expand: int, + data_dim: int, +) -> torch.nn.Module: + """Instantiate a single mixer layer sized for a ``resolution`` grid in ``data_dim`` dims. + + Only the resolution-dependent construction args change with ``resolution``: + the SIREN kernel cache (``L_cache``) for Hyena and the RoPE tables for + Attention. Mamba is resolution-independent (it rasterizes at forward time). + Hyena is causal in 1D (the fused genomics path) and non-causal in 2D/3D. + """ + if name == "hyena": + cfg = _hyena_mixer_cfg( + hidden_dim, + fft_backend, + canvas_size=resolution, + grid_type=grid_type, + data_dim=data_dim, + is_causal=(data_dim == 1), + ) + elif name in _ATTN_IMPL: # attention (sdpa) / flex / fa4 — shared q/k/v + RoPE path + head_dim = hidden_dim // num_heads + div = _ROPE_DIVISOR[data_dim] # 1D:2, 2D:4, 3D:6 + use_rope = attn_rope and (head_dim % div == 0) + if attn_rope and not use_rope: + print( + f" [warn] disabling {data_dim}D RoPE: head_dim={head_dim} " + f"(hidden_dim={hidden_dim} / num_heads={num_heads}) not divisible by {div}.", + flush=True, + ) + cfg = _attention_mixer_cfg( + hidden_dim, + num_heads=num_heads, + use_rope=use_rope, + rope_spatial_dims=(resolution,) * data_dim if use_rope else None, + attn_impl=_ATTN_IMPL[name], + ) + elif name == "mamba": + cfg = _mamba_mixer_cfg( + hidden_dim, + headdim=mamba_headdim, + expand=mamba_expand, + bidirectional=True, + ) + else: # pragma: no cover - guarded by argparse choices + raise ValueError(f"unknown mixer '{name}'") + + return instantiate(cfg) + + +# ─── Timing ─────────────────────────────────────────────────────────────────── + + +def _is_oom(exc: BaseException) -> bool: + if isinstance(exc, torch.cuda.OutOfMemoryError): + return True + return isinstance(exc, RuntimeError) and "out of memory" in str(exc).lower() + + +def time_forward( + name: str, + resolution: int, + *, + hidden_dim: int, + batch_size: int, + dtype: torch.dtype, + num_warmup: int, + num_iters: int, + compile_mode: str | None, + fft_backend: str, + grid_type: str, + num_heads: int, + attn_rope: bool, + mamba_headdim: int, + mamba_expand: int, + data_dim: int, + max_seconds: float, + device: torch.device, +) -> dict[str, Any]: + """Build the layer, run a forward, and return a result dict. + + Returns keys ``status`` (``ok``/``oom``/``error``/``timeout``), ``ms`` and + ``mem_gb`` (``None`` unless the point produced a usable timing). A single + warmup forward is wall-timed first; if it already exceeds ``max_seconds`` the + point is marked ``timeout`` and the (expensive) timed loop is skipped, so the + worst case is one slow forward rather than ``num_iters`` of them. + """ + if name == "hyena" and compile_mode is not None and fft_backend == "torch_fft": + # Only the torch.fft path needs this flag; the subq_ops custom op is + # compile-safe on its own. + import nvsubquadratic.ops.fftconv as _fftconv + + _fftconv.COMPILE_COMPATIBLE = True + + module = None + x = None + try: + module = ( + build_module( + name, + hidden_dim=hidden_dim, + resolution=resolution, + fft_backend=fft_backend, + grid_type=grid_type, + num_heads=num_heads, + attn_rope=attn_rope, + mamba_headdim=mamba_headdim, + mamba_expand=mamba_expand, + data_dim=data_dim, + ) + .to(device) + .eval() + ) + + if compile_mode is not None: + module = torch.compile(module, mode=compile_mode) + + x = torch.randn(batch_size, *([resolution] * data_dim), hidden_dim, device=device, dtype=torch.float32) + torch.cuda.reset_peak_memory_stats(device) + + # ── Single wall-timed forward (compile + timeout guard) ────────────── + torch.cuda.synchronize(device) + t0 = time.perf_counter() + with torch.inference_mode(), torch.autocast("cuda", dtype=dtype): + _ = module(x) + torch.cuda.synchronize(device) + first_forward_s = time.perf_counter() - t0 + + if first_forward_s > max_seconds: + return { + "status": "timeout", + "ms": first_forward_s * 1000.0, + "mem_gb": torch.cuda.max_memory_allocated(device) / (1024**3), + } + + # ── Adaptive iteration counts ───────────────────────────────────────── + # Per-forward cost spans ~1e-3 s (small hyena) to minutes (attention at + # multi-M tokens). A fixed count would be noisy for fast ops and run for + # hours on slow ones (30 iters x a 5-min forward = 2.5 h for one point). + # Scale both warmup and timed counts to the measured cost, keeping each + # point within the per-point time budget. num_warmup/num_iters are caps. + target_timed_s = 5.0 + min_timed_iters = 3 + n_timed = round(target_timed_s / max(first_forward_s, 1e-9)) + n_timed = max(min_timed_iters, min(num_iters, n_timed)) + # Never let the timed loop exceed the budget for slow ops. + n_timed = min(n_timed, max(1, int(max_seconds / max(first_forward_s, 1e-9)))) + # Extra warmup only pays off for cheap forwards; skip it for slow ones. + extra_warmup = max(0, num_warmup - 1) if first_forward_s < 1.0 else 0 + + # ── Remaining warmup ────────────────────────────────────────────────── + with torch.inference_mode(), torch.autocast("cuda", dtype=dtype): + for _ in range(extra_warmup): + _ = module(x) + torch.cuda.synchronize(device) + + # ── Timed iterations (CUDA events) ──────────────────────────────────── + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + torch.cuda.synchronize(device) + start.record() + with torch.inference_mode(), torch.autocast("cuda", dtype=dtype): + for _ in range(n_timed): + _ = module(x) + end.record() + torch.cuda.synchronize(device) + + return { + "status": "ok", + "ms": start.elapsed_time(end) / n_timed, + "mem_gb": torch.cuda.max_memory_allocated(device) / (1024**3), + "iters": n_timed, + } + except Exception as exc: + if isinstance(exc, ImportError): # e.g. mamba_ssm not installed — not a wall + status = "unavailable" + elif _is_oom(exc): + status = "oom" + else: + status = "error" + # Print the full traceback to the log for genuine errors (not clean + # OOM/unavailable) so failures like Mamba's construction error are + # diagnosable without a rerun. Only the repr goes into the JSONL. + if status == "error": + traceback.print_exc(file=sys.stdout) + return {"status": status, "ms": None, "mem_gb": None, "detail": repr(exc)} + finally: + del module, x + gc.collect() + torch.cuda.empty_cache() + if compile_mode is not None and hasattr(torch, "_dynamo"): + torch._dynamo.reset() + + +def _predicted_ms(name: str, seq_len: int, last_seq_len: int, last_ms: float) -> float: + """Extrapolate step time to ``seq_len`` from the last successful point. + + Uses each mixer's asymptotic law so we can *skip* (rather than launch and + block on) points that will clearly exceed the time budget: attention grows + as O(L^2), Hyena as O(L log L), Mamba as O(L). + """ + r = seq_len / last_seq_len + if name == "attention": + return last_ms * r * r + if name == "hyena": + return last_ms * r * (math.log2(seq_len) / math.log2(last_seq_len)) + return last_ms * r # mamba (linear) and any other operator + + +# ─── Main ───────────────────────────────────────────────────────────────────── + + +def _validate(args: argparse.Namespace) -> None: + if "attention" in args.mixers: + if args.hidden_dim % args.num_heads != 0: + raise SystemExit(f"hidden_dim={args.hidden_dim} not divisible by num_heads={args.num_heads}.") + if "mamba" in args.mixers: + d_inner = args.hidden_dim * args.mamba_expand + if d_inner % args.mamba_headdim != 0: + raise SystemExit( + f"Mamba2 d_inner=hidden_dim*expand={d_inner} not divisible by mamba_headdim={args.mamba_headdim}." + ) + # subq_ops (1D/2D only) needs even resolutions; 3D uses torch_fft so it is exempt. + if args.fft_backend == "subq_ops" and args.data_dim in (1, 2) and "hyena" in args.mixers: + odd = [r for r in args.resolutions if r % 2 != 0] + if odd: + raise SystemExit(f"subq_ops requires even resolutions; got odd {odd}.") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--resolutions", + nargs="+", + type=int, + default=[64, 128, 256, 512, 1024], + help="Per-axis resolution R. Token count L=R^data_dim. Swept ascending.", + ) + parser.add_argument( + "--data-dim", + type=int, + choices=[1, 2, 3], + default=2, + help="Spatial dimensionality N. Input is [B, R,...(N)..., C], L=R^N. Hyena is causal in 1D " + "(the fused genomics kernel) and non-causal in 2D/3D; 3D auto-falls back to torch_fft " + "(subq_ops has no 3D kernel).", + ) + parser.add_argument( + "--mixers", + nargs="+", + choices=MIXER_CHOICES, + default=list(MIXER_CHOICES), + help="Subset of mixers to benchmark.", + ) + parser.add_argument("--batch-size", type=int, default=1, help="Batch size for the timed forward pass.") + parser.add_argument("--hidden-dim", type=int, default=256, help="Shared channel dim across all mixers.") + parser.add_argument("--num-heads", type=int, default=8, help="Attention heads (head_dim=hidden_dim/num_heads).") + parser.add_argument("--mamba-headdim", type=int, default=64, help="Mamba2 head dim.") + parser.add_argument("--mamba-expand", type=int, default=2, help="Mamba2 expansion factor.") + parser.add_argument("--num-warmup", type=int, default=10, help="Warmup iterations (also covers compile).") + parser.add_argument("--num-iters", type=int, default=30, help="Timed iterations.") + parser.add_argument("--dtype", choices=list(DTYPE_MAP), default="bf16", help="Autocast dtype.") + parser.add_argument( + "--fft-backend", + choices=["subq_ops", "torch_fft"], + default="subq_ops", + help="FFT-conv backend for Hyena. Use 'torch_fft' on hosts without the subq_ops CUDA kernels.", + ) + parser.add_argument( + "--grid-type", + choices=["double", "single"], + default="double", + help="Hyena SIREN kernel grid. 'double' = non-circular / zero-padded (kernel size ~2R per axis); " + "'single' = circular (kernel size ~R). The materialized kernel is ~4x smaller with 'single', which " + "helps memory at high resolution. Both size the 2D FFT to ~2*next_pow2(R) from the input, so R=8192 " + "(~16K-point FFT) sits near the cuFFTDx size ceiling either way.", + ) + rope = parser.add_mutually_exclusive_group() + rope.add_argument("--attn-rope", dest="attn_rope", action="store_true", help="Enable 2D axial RoPE (default).") + rope.add_argument("--no-attn-rope", dest="attn_rope", action="store_false", help="Disable RoPE.") + parser.set_defaults(attn_rope=True) + compile_grp = parser.add_mutually_exclusive_group() + compile_grp.add_argument( + "--compile-mode", + type=str, + default=None, + help="torch.compile mode (e.g. 'max-autotune-no-cudagraphs'). Default: eager.", + ) + compile_grp.add_argument("--no-compile", action="store_true", help="Disable torch.compile (default).") + parser.add_argument( + "--max-seconds-per-point", + type=float, + default=120.0, + help="If a single forward exceeds this, mark the point 'timeout' and skip larger resolutions " + "for that mixer (attention's O(L^2) explosion).", + ) + parser.add_argument( + "--disable-cudnn", + action="store_true", + help="Turn off cuDNN (Conv2d native fallback; SDPA avoids the cuDNN backend). " + "Workaround for images where cuDNN fails to initialize (CUDNN_STATUS_NOT_INITIALIZED).", + ) + parser.add_argument( + "--output", + type=str, + default="benchmarks/results/forward_time.jsonl", + help="Output JSONL path (one row per mixer/resolution).", + ) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA device required for the forward-time benchmark.") + _validate(args) + + device = torch.device("cuda") + torch.backends.cudnn.benchmark = True + torch.set_float32_matmul_precision("high") + if args.disable_cudnn: + torch.backends.cudnn.enabled = False + if hasattr(torch.backends.cuda, "enable_cudnn_sdp"): + torch.backends.cuda.enable_cudnn_sdp(False) + print("[cudnn] disabled (Conv2d native fallback; SDPA flash/mem-efficient).", flush=True) + + dtype = DTYPE_MAP[args.dtype] + compile_mode = None if args.no_compile else args.compile_mode + resolutions = sorted(args.resolutions) + device_name = torch.cuda.get_device_name(device) + data_dim = args.data_dim + # subq_ops has no 3D kernel — Hyena falls back to torch_fft in 3D. + hyena_backend = "torch_fft" if data_dim == 3 and args.fft_backend == "subq_ops" else args.fft_backend + if hyena_backend != args.fft_backend: + print("[note] data_dim=3: Hyena fft_backend forced to torch_fft (subq_ops is 1D/2D only).") + + print(f"Device: {device_name}") + print( + f"Settings: data_dim={data_dim} batch_size={args.batch_size} hidden_dim={args.hidden_dim} dtype={args.dtype} " + f"compile={compile_mode} fft_backend={hyena_backend} grid_type={args.grid_type} rope={args.attn_rope} " + f"warmup={args.num_warmup} timed={args.num_iters} max_s={args.max_seconds_per_point}" + ) + print(f"Mixers: {args.mixers} Resolutions: {resolutions}") + + # Attention-kernel eligibility. SDPA (auto cuDNN / flash) needs head_dim % 8 == 0 + # and is *optimized* for 64/128; below a multiple of 8 it silently falls back to a + # weak math / mem-efficient path. FlexAttention and FA4 additionally *require* + # head_dim >= 16 — below that they raise, so those points error out (rather than + # degrade) at a too-small head_dim. + attn_mixers = [m for m in args.mixers if m in _ATTN_IMPL] + if attn_mixers: + head_dim = args.hidden_dim // args.num_heads + print( + f"[attn] kernels={attn_mixers} head_dim={head_dim} (hidden {args.hidden_dim} / heads {args.num_heads})", + flush=True, + ) + if head_dim % 8 != 0: + print( + f"[attn] head_dim={head_dim} is NOT a multiple of 8 → SDPA flash/cuDNN INELIGIBLE " + f"(falls back to math/mem-efficient, a weak non-flash baseline).", + flush=True, + ) + elif head_dim < 64: + print( + f"[attn] head_dim={head_dim}: flash eligible but BELOW its optimized regime " + f"(flash/FA4 tuned for head_dim 64/128).", + flush=True, + ) + else: + print( + f"[attn] head_dim={head_dim} → flash/cuDNN eligible; SDPA auto-selects the fused kernel.", flush=True + ) + needs16 = sorted({"flex", "fa4"} & set(attn_mixers)) + if needs16 and head_dim < 16: + print( + f"[attn] head_dim={head_dim} < 16 → {needs16} will ERROR " + f"(FlexAttention/FA4 require head_dim >= 16); those points are marked 'error'. " + f"Raise hidden_dim or lower num_heads.", + flush=True, + ) + + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + + # rows[mixer][seq_len] = result dict, for the summary table + rows: dict[str, dict[int, dict[str, Any]]] = {m: {} for m in args.mixers} + + with out_path.open("w") as fh: + for mixer in args.mixers: + last_ok: tuple[int, float] | None = None # (seq_len, ms) of last ok/timeout point + for R in resolutions: + seq_len = R**data_dim + label = f"[{mixer:>9s} R={R:>4d} L={seq_len:>9d}]" + + # Predictive skip: never launch a point that will clearly blow the budget. + if last_ok is not None: + pred_ms = _predicted_ms(mixer, seq_len, last_ok[0], last_ok[1]) + if pred_ms > args.max_seconds_per_point * 1000.0: + print( + f"\n{label}\n [skip] predicted ~{pred_ms / 1000.0:.0f}s > budget; marking timeout.", + flush=True, + ) + result = {"status": "timeout", "ms": None, "mem_gb": None} + rows[mixer][seq_len] = result + record = { + "mixer": mixer, + "resolution": R, + "seq_len": seq_len, + "data_dim": data_dim, + "backend": hyena_backend if mixer == "hyena" else None, + "batch_size": args.batch_size, + "hidden_dim": args.hidden_dim, + "num_heads": args.num_heads, + "dtype": args.dtype, + "device": device_name, + **result, + } + fh.write(json.dumps(record) + "\n") + fh.flush() + continue + + print(f"\n{label}", flush=True) + t0 = time.perf_counter() + result = time_forward( + mixer, + R, + hidden_dim=args.hidden_dim, + batch_size=args.batch_size, + dtype=dtype, + num_warmup=args.num_warmup, + num_iters=args.num_iters, + compile_mode=compile_mode, + fft_backend=hyena_backend, + grid_type=args.grid_type, + num_heads=args.num_heads, + attn_rope=args.attn_rope, + mamba_headdim=args.mamba_headdim, + mamba_expand=args.mamba_expand, + data_dim=data_dim, + max_seconds=args.max_seconds_per_point, + device=device, + ) + wall = time.perf_counter() - t0 + + ms, mem = result.get("ms"), result.get("mem_gb") + if result["status"] == "ok": + n = result.get("iters", "?") + print( + f" ms/fwd = {ms:9.3f} (n={n}) | peak mem = {mem:6.2f} GB | wall = {wall:5.1f}s", + flush=True, + ) + last_ok = (seq_len, ms) + elif result["status"] == "timeout": + shown = f"{ms / 1000.0:.1f}s" if ms is not None else "n/a" + print(f" [timeout] single forward = {shown} | wall = {wall:5.1f}s", flush=True) + last_ok = (seq_len, ms) if ms is not None else last_ok + else: + print(f" [{result['status']}] {result.get('detail', '')} | wall = {wall:5.1f}s", flush=True) + + rows[mixer][seq_len] = result + record = { + "mixer": mixer, + "resolution": R, + "seq_len": seq_len, + "data_dim": data_dim, + "backend": hyena_backend if mixer == "hyena" else None, + "batch_size": args.batch_size, + "hidden_dim": args.hidden_dim, + "num_heads": args.num_heads, + "dtype": args.dtype, + "device": device_name, + "status": result["status"], + "ms": ms, + "mem_gb": mem, + "iters": result.get("iters"), + } + if "detail" in result: + record["detail"] = result["detail"] + fh.write(json.dumps(record) + "\n") + fh.flush() + + # ── Summary table ───────────────────────────────────────────────────────── + print(f"\n{'=' * 78}") + print("Forward-time summary (ms/fwd, lower is better; 'oom'/'timeout'/'error' otherwise)") + print(f"{'=' * 78}") + header = f"{'R':>6s} {f'L=R^{data_dim}':>10s} " + " ".join(f"{m:>14s}" for m in args.mixers) + print(header) + print("-" * len(header)) + for R in resolutions: + seq_len = R**data_dim + row = f"{R:>6d} {seq_len:>10d} " + for m in args.mixers: + res = rows[m].get(seq_len) + if res is None: + cell = "-" + elif res["status"] == "ok": + cell = f"{res['ms']:.3f}" + else: + cell = res["status"] + row += f"{cell:>14s} " + print(row.rstrip()) + + print(f"\n[done] wrote {out_path.resolve()}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/benchmark_patch_size_2d.py b/benchmarks/benchmark_patch_size_2d.py index 4526f67b..97886c19 100644 --- a/benchmarks/benchmark_patch_size_2d.py +++ b/benchmarks/benchmark_patch_size_2d.py @@ -106,34 +106,46 @@ MAMBA_EXPAND = 2 MAMBA_BIDIRECTIONAL = True +# Per-dimensionality short depthwise conv (1D/2D/3D). +_CONV_ND = {1: torch.nn.Conv1d, 2: torch.nn.Conv2d, 3: torch.nn.Conv3d} + # ─── Mixer config builders (concrete — no OmegaConf interpolation) ──────────── -def _hyena_mixer_cfg(hidden_dim: int, fft_backend: str) -> LazyConfig: +def _hyena_mixer_cfg( + hidden_dim: int, + fft_backend: str, + *, + canvas_size: int = CANVAS_SIZE, + grid_type: str = "double", + data_dim: int = DATA_DIM, + is_causal: bool = False, +) -> LazyConfig: return LazyConfig(QKVSequenceMixer)( hidden_dim=hidden_dim, mixer_cfg=LazyConfig(Hyena)( global_conv_cfg=LazyConfig(CKConvND)( - data_dim=DATA_DIM, + data_dim=data_dim, hidden_dim=hidden_dim, kernel_cfg=LazyConfig(SIRENKernelND)( - data_dim=DATA_DIM, + data_dim=data_dim, out_dim=hidden_dim, mlp_hidden_dim=KERNEL_MLP_HIDDEN_DIM, num_layers=KERNEL_NUM_LAYERS, embedding_dim=KERNEL_EMBEDDING_DIM, omega_0=KERNEL_OMEGA_0, - L_cache=CANVAS_SIZE, + L_cache=canvas_size, use_bias=True, hidden_omega_0=KERNEL_HIDDEN_OMEGA_0, ), mask_cfg=LazyConfig(torch.nn.Identity)(), - grid_type="double", + grid_type=grid_type, fft_padding="zero", fft_backend=fft_backend, + is_causal=is_causal, ), - short_conv_cfg=LazyConfig(torch.nn.Conv2d)( + short_conv_cfg=LazyConfig(_CONV_ND[data_dim])( in_channels=3 * hidden_dim, out_channels=3 * hidden_dim, kernel_size=3, @@ -144,32 +156,50 @@ def _hyena_mixer_cfg(hidden_dim: int, fft_backend: str) -> LazyConfig: gate_nonlinear_cfg=LazyConfig(torch.nn.Identity)(), pixelhyena_norm_cfg=LazyConfig(torch.nn.LayerNorm)(normalized_shape=hidden_dim), qk_norm_cfg=None, - use_rope=False, - rope_base=10000.0, ), init_method_in=small_init, init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers=NUM_BLOCKS), ) -def _attention_mixer_cfg(hidden_dim: int) -> LazyConfig: +def _attention_mixer_cfg( + hidden_dim: int, + *, + num_heads: int = ATTN_NUM_HEADS, + use_rope: bool = ATTN_USE_ROPE, + rope_spatial_dims: tuple[int, ...] | None = None, + attn_impl: str = "sdpa", +) -> LazyConfig: return LazyConfig(QKVSequenceMixer)( hidden_dim=hidden_dim, mixer_cfg=LazyConfig(Attention)( hidden_dim=hidden_dim, - num_heads=ATTN_NUM_HEADS, + num_heads=num_heads, apply_qk_norm=True, - use_rope=ATTN_USE_ROPE, + use_rope=use_rope, is_causal=False, rope_base=10000.0, attn_dropout=0.0, + rope_spatial_dims=rope_spatial_dims, + attn_impl=attn_impl, ), init_method_in=small_init, init_method_out=LazyConfig(partial_wang_init_fn_with_num_layers)(num_layers=NUM_BLOCKS), ) -def _mamba_mixer_cfg(hidden_dim: int) -> LazyConfig: +def _mamba_mixer_cfg( + hidden_dim: int, + *, + headdim: int = MAMBA_HEADDIM, + expand: int = MAMBA_EXPAND, + bidirectional: bool = MAMBA_BIDIRECTIONAL, +) -> LazyConfig: + # mamba-ssm >= 2.3 eagerly imports Mamba3 in its package __init__, which pulls + # in tilelang -> tvm and crashes on this stack (tvm_ffi AttributeError under + # py3.12). We only use Mamba2, so make that optional `import tilelang` fail as a + # clean ImportError (which mamba_ssm's __init__ skips) by neutering it here. + sys.modules.setdefault("tilelang", None) from mamba_ssm import Mamba2 from nvsubquadratic.modules.mamba_nd import Mamba as MambaNDMixer @@ -177,10 +207,10 @@ def _mamba_mixer_cfg(hidden_dim: int) -> LazyConfig: return LazyConfig(MambaNDMixer)( mamba_layer_cfg=LazyConfig(Mamba2)( d_model=hidden_dim, - headdim=MAMBA_HEADDIM, - expand=MAMBA_EXPAND, + headdim=headdim, + expand=expand, ), - bidirectional=MAMBA_BIDIRECTIONAL, + bidirectional=bidirectional, ) diff --git a/docs/getting_started.md b/docs/getting_started.md index f2abbbff..9b313e4f 100644 --- a/docs/getting_started.md +++ b/docs/getting_started.md @@ -46,7 +46,7 @@ For an alternative venv-based install: python3 -m venv venv source venv/bin/activate pip install torch==2.10.0 torchvision==0.25.0 \ - --index-url https://download.pytorch.org/whl/cu129 + --index-url https://download.pytorch.org/whl/cu130 pip install -r requirements-dev.txt pip install --no-build-isolation -e . ``` diff --git a/examples/vit5_imagenet/_deprecated/vit5_small_pretrain_apex_dali.py b/examples/vit5_imagenet/_deprecated/vit5_small_pretrain_apex_dali.py index a48797a9..23690ab2 100644 --- a/examples/vit5_imagenet/_deprecated/vit5_small_pretrain_apex_dali.py +++ b/examples/vit5_imagenet/_deprecated/vit5_small_pretrain_apex_dali.py @@ -17,7 +17,7 @@ Identical to vit5_small_pretrain_apex.py except: - Uses DALIImageNetDataModule for GPU-pipelined data loading -- Requires: pip install nvidia-dali-cuda120 +- Requires: pip install nvidia-dali-cuda130 """ import os diff --git a/examples/vit5_imagenet/_deprecated/vit5_small_pretrain_apex_dali_optimized_plus.py b/examples/vit5_imagenet/_deprecated/vit5_small_pretrain_apex_dali_optimized_plus.py index 7d3a97d6..d756049b 100644 --- a/examples/vit5_imagenet/_deprecated/vit5_small_pretrain_apex_dali_optimized_plus.py +++ b/examples/vit5_imagenet/_deprecated/vit5_small_pretrain_apex_dali_optimized_plus.py @@ -18,7 +18,7 @@ Based on optimized_v2 with additional training-loop improvements: - check_val_every_n_epoch=4 (validate every 4 epochs → fewer DALI pipeline disruptions) - local_staging_dir=/scratch/imagenet_folder (stage data to node-local NVMe) -- Requires: pip install nvidia-dali-cuda120 +- Requires: pip install nvidia-dali-cuda130 """ import os diff --git a/examples/vit5_imagenet/v2/vit5_small_pretrain_apex_dali_fused.py b/examples/vit5_imagenet/v2/vit5_small_pretrain_apex_dali_fused.py index a05e0bae..d188b260 100644 --- a/examples/vit5_imagenet/v2/vit5_small_pretrain_apex_dali_fused.py +++ b/examples/vit5_imagenet/v2/vit5_small_pretrain_apex_dali_fused.py @@ -19,7 +19,7 @@ DALI pipeline, eliminating ~25ms of serial GPU augmentation overhead per step. Based on optimized_plus with the datamodule swapped to DALIImageNetFusedDataModule. -Requires: pip install nvidia-dali-cuda120 +Requires: pip install nvidia-dali-cuda130 """ import os diff --git a/examples/vit5_imagenet/v5_patch/README.md b/examples/vit5_imagenet/v5_patch/README.md index 5a7b1646..e9785ab6 100644 --- a/examples/vit5_imagenet/v5_patch/README.md +++ b/examples/vit5_imagenet/v5_patch/README.md @@ -33,7 +33,7 @@ Attention does not use prepend_registers, so T = 1 + 4 + `num_patches` (no paddi ```bash git clone && cd nvSubquadratic -# PyTorch with CUDA 12.8 +# PyTorch with CUDA 13.0 pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128 # Core dependencies @@ -47,7 +47,7 @@ DALI provides fused decode + crop + augmentations on the GPU, significantly speeding up data loading for ImageNet. ```bash -pip install nvidia-dali-cuda120 +pip install nvidia-dali-cuda130 ``` ### 3. Install Apex (FusedLAMB optimizer) diff --git a/experiments/datamodules/_deprecated/dali_imagenet.py b/experiments/datamodules/_deprecated/dali_imagenet.py index f91f8535..0914d0b7 100644 --- a/experiments/datamodules/_deprecated/dali_imagenet.py +++ b/experiments/datamodules/_deprecated/dali_imagenet.py @@ -20,7 +20,7 @@ Remaining augmentations (ThreeAugment, ColorJitter) run as batched PyTorch ops on uniform-sized GPU tensors. -Requires: ``pip install nvidia-dali-cuda120`` +Requires: ``pip install nvidia-dali-cuda130`` """ import os diff --git a/experiments/datamodules/_deprecated/dali_imagenet_optimized.py b/experiments/datamodules/_deprecated/dali_imagenet_optimized.py index aa25250b..c74d3fe3 100644 --- a/experiments/datamodules/_deprecated/dali_imagenet_optimized.py +++ b/experiments/datamodules/_deprecated/dali_imagenet_optimized.py @@ -26,7 +26,7 @@ 6. ``torch.compile``-friendly augmentation modules 7. Optional NCHW output (``channels_first=True``) to skip redundant permute -Requires: ``pip install nvidia-dali-cuda120`` +Requires: ``pip install nvidia-dali-cuda130`` """ import os diff --git a/experiments/datamodules/dali_imagenet_fused.py b/experiments/datamodules/dali_imagenet_fused.py index b4f81573..c574bdb8 100644 --- a/experiments/datamodules/dali_imagenet_fused.py +++ b/experiments/datamodules/dali_imagenet_fused.py @@ -22,7 +22,7 @@ This eliminates the ~25-33 ms of serial GPU augmentation that blocks the training loop in ``dali_imagenet_optimized.py``. -Requires: ``pip install nvidia-dali-cuda120`` +Requires: ``pip install nvidia-dali-cuda130`` """ import math diff --git a/nvsubquadratic/modules/attention.py b/nvsubquadratic/modules/attention.py index addc59f3..f077975f 100644 --- a/nvsubquadratic/modules/attention.py +++ b/nvsubquadratic/modules/attention.py @@ -121,6 +121,31 @@ from nvsubquadratic.utils import qk_norm, rope +def _resolve_fa4_func(): + """Return the ``flash_attn_func`` callable for the installed FlashAttention. + + FlashAttention-4 (Blackwell/Hopper, ``pip install --pre flash-attn-4``) is a + pure-Python CuTe-DSL JIT wheel exposing ``flash_attn.cute.flash_attn_func``. + We prefer it, then fall back to the Hopper FA3 build (``flash_attn_interface``) + and finally the mainline FA2 wheel (``flash_attn``). All expose a + ``flash_attn_func(q, k, v, softmax_scale=, causal=)`` on ``[B, S, H, D]`` + tensors. Raise a clear error if none is importable. + """ + for mod_name in ("flash_attn.cute", "flash_attn_interface", "flash_attn"): + try: + mod = __import__(mod_name, fromlist=["flash_attn_func"]) + except ImportError: + continue + func = getattr(mod, "flash_attn_func", None) + if func is not None: + return func + raise ImportError( + "attn_impl='fa4' needs FlashAttention. Install FA4 (Blackwell/Hopper): " + "`pip install --pre flash-attn-4` (exposes flash_attn.cute), or a " + "flash_attn_interface / flash_attn build for your GPU arch." + ) + + class Attention(torch.nn.Module): r"""Multi-head scaled dot-product self-attention for 1D/2D/3D spatial inputs. @@ -227,6 +252,11 @@ class Attention(torch.nn.Module): Examples: ``(4096,)`` for 1D, ``(64, 64)`` for 2D, ``(8, 64, 64)`` for 3D. Must match the spatial shape seen during ``forward``. + attn_impl (str): Attention kernel — ``"sdpa"`` (default; PyTorch + ``scaled_dot_product_attention`` auto-selecting cuDNN / flash / + fallback), ``"flex"`` (compiled ``torch.nn.attention.flex_attention``), + or ``"fa4"`` (FlashAttention-4 via the external ``flash_attn`` package). + ``flex``/``fa4`` require ``head_dim >= 16``. Example:: @@ -256,6 +286,7 @@ def __init__( attn_dropout: float = 0.0, rope_base: float = 10000.0, rope_spatial_dims: tuple[int, ...] | None = None, + attn_impl: str = "sdpa", ): """Initialise the Attention module and precompute RoPE buffers. @@ -278,6 +309,9 @@ def __init__( or ``extra_repr``). The corresponding cos/sin buffers are stored as non-persistent registered buffers (``rope_cos``, ``rope_sin``, etc.). + attn_impl (str): Attention kernel — ``"sdpa"`` (default), + ``"flex"`` (compiled FlexAttention), or ``"fa4"`` + (FlashAttention-4). ``flex``/``fa4`` require ``head_dim >= 16``. Raises: AssertionError: If ``hidden_dim % num_heads != 0``. @@ -299,6 +333,19 @@ def __init__( self.rope_base = rope_base self.is_causal = is_causal self.attn_dropout = attn_dropout + # Attention kernel: "sdpa" (PyTorch auto-select: cuDNN / flash / fallback), + # "flex" (compiled FlexAttention), or "fa4" (FlashAttention-4 via flash_attn). + # flex/fa4 back-ends are resolved eagerly here so a missing dependency fails + # at construction, not mid-forward. + self.attn_impl = attn_impl + if attn_impl == "flex": + from torch.nn.attention.flex_attention import flex_attention + + self._flex_attention = torch.compile(flex_attention) + elif attn_impl == "fa4": + self._fa4_func = _resolve_fa4_func() + elif attn_impl != "sdpa": + raise ValueError(f"Unknown attn_impl={attn_impl!r}; use 'sdpa', 'flex', or 'fa4'.") # ── Precomputed RoPE cos/sin buffers ────────────────────────────── # @@ -550,19 +597,44 @@ def forward( key = rearrange(key, "(b h) t d -> b h t d", h=local_num_heads) value = rearrange(value, "(b h) t d -> b h t d", h=local_num_heads) - # Scaled dot-product attention — let PyTorch auto-select the best - # backend (CuDNN on H100, FlashAttention on A100, etc.). - # No manual dtype cast: autocast handles precision. - out = F.scaled_dot_product_attention( - query, - key, - value, - dropout_p=self.attn_dropout if self.training else 0.0, - is_causal=self.is_causal, - # When QK-norm is applied (cosine attention), disable the default - # 1/sqrt(d) scaling — it would flatten the normalised logits. - scale=self.scale if not self.apply_qk_norm else 1.0, - ) + # When QK-norm is applied (cosine attention), disable the default + # 1/sqrt(d) scaling — it would flatten the normalised logits. + scale = self.scale if not self.apply_qk_norm else 1.0 + dropout_p = self.attn_dropout if self.training else 0.0 + if self.attn_impl in ("flex", "fa4"): + # RoPE / qk_norm can upcast q,k to fp32 while v stays in the autocast + # dtype. SDPA's math fallback tolerates the mismatch, but flash-class + # kernels (FlexAttention, FA4) require q,k,v to share one dtype. + query = query.to(value.dtype) + key = key.to(value.dtype) + if self.attn_impl == "flex": + # Compiled FlexAttention. Non-causal here (no block_mask); is_causal + # would need a causal BlockMask, unused by the ND benchmark. + out = self._flex_attention(query, key, value, scale=scale) + elif self.attn_impl == "fa4": + # FlashAttention-4. flash_attn expects [B, S, H, D], not SDPA's [B, H, S, D]. + # softmax_scale/causal are the args common to FA2/FA3/FA4; dropout_p is + # omitted (dropped in FA3+, and 0 here anyway). Some builds return + # (out, softmax_lse) — take the first element. + out = self._fa4_func( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + softmax_scale=scale, + causal=self.is_causal, + ) + if isinstance(out, tuple): + out = out[0] + out = out.transpose(1, 2) + else: # sdpa — PyTorch auto-selects cuDNN / flash / fallback. + out = F.scaled_dot_product_attention( + query, + key, + value, + dropout_p=dropout_p, + is_causal=self.is_causal, + scale=scale, + ) # Merge heads: [B, H, T, D] -> [B, T, H*D] out = rearrange(out, "b h t d -> b t (h d)", h=local_num_heads) diff --git a/pyproject.toml b/pyproject.toml index 7691a084..0970a0e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,13 +73,15 @@ Repository = "https://github.com/NVIDIA-BioNeMo/nvSubquadratic" Issues = "https://github.com/NVIDIA-BioNeMo/nvSubquadratic/issues" [project.optional-dependencies] -# Accelerated CUDA kernels (fused FFT-conv / direct causal-conv). Source-only -# sdist — building requires nvcc, so this is intentionally NOT a core dep: that -# is what lets `pip install nvsubquadratic` succeed without the CUDA toolkit -# (e.g. a downstream project's CPU CI). The operators import it lazily and the -# default fft_backend="torch_fft" path runs without it; install for the fast path: -# pip install nvsubquadratic[cuda] -cuda = ["subquadratic-ops-torch-cu12>=0.2.0"] +# Accelerated CUDA-13 kernels (fused FFT-conv / direct causal-conv). An opt-in +# extra, NOT a core dep, so `pip install nvsubquadratic` still succeeds without the +# CUDA toolkit (e.g. a downstream project's CPU CI): the operators import it lazily +# and the default fft_backend="torch_fft" path runs without it. Install for the +# fast path: pip install nvsubquadratic[cuda] +# Pinned to >=0.2.1, the newest cu13 on public PyPI, so it resolves with a plain +# `pip install`. (0.2.2 exists but only on the internal NVIDIA GitLab registry — +# bump this pin once it is published to public PyPI.) +cuda = ["subquadratic-ops-torch-cu13>=0.2.1"] # QuACK fused RMSNorm kernel (Hopper/Blackwell only: H100, B200, B300). # On Ampere and older GPUs RMSNorm falls back to pure PyTorch automatically. @@ -87,7 +89,7 @@ quack = ["quack-kernels"] # NVIDIA DALI GPU data-loading pipeline, used by the ImageNet / The Well # examples. Not required by the importable library; ~400 MB download. -dali = ["nvidia-dali-cuda120==1.53.0"] +dali = ["nvidia-dali-cuda130==1.53.0"] # Megatron parallel state, needed only for context-parallel / distributed # training via nvsubquadratic.parallel.utils.init_parallel_state. The operators diff --git a/scripts/docker/patch_mamba_cuda_arches.py b/scripts/docker/patch_mamba_cuda_arches.py new file mode 100644 index 00000000..3b24ac2f --- /dev/null +++ b/scripts/docker/patch_mamba_cuda_arches.py @@ -0,0 +1,126 @@ +#!/usr/bin/env python3 + +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Patch mamba-ssm / causal-conv1d setup.py to honor TORCH_CUDA_ARCH_LIST. + +Upstream hardcodes -gencode for sm_75..sm_120 (and ignores TORCH_CUDA_ARCH_LIST), +which OOMs / gcc-ICEs when building under QEMU arm64. This rewrites the gencode +block from TORCH_CUDA_ARCH_LIST and makes append_nvcc_threads honor NVCC_THREADS. +""" + +from __future__ import annotations + +import os +import re +import sys +from pathlib import Path + + +# TORCH_CUDA_ARCH_LIST token -> sm number used in -gencode +ARCH_TO_SM = { + "7.5": 75, + "8.0": 80, + "8.6": 86, + "8.7": 87, + "8.9": 89, + "9.0": 90, + "10.0": 100, + "12.0": 120, +} + + +def parse_arch_list(raw: str) -> list[tuple[str, int]]: + arches: list[tuple[str, int]] = [] + for part in raw.replace(",", ";").split(";"): + token = part.strip() + if not token: + continue + # Drop optional "a" suffix (e.g. 9.0a / 10.0a). + base = token.rstrip("aA") + if base not in ARCH_TO_SM: + raise SystemExit( + f"Unsupported arch {token!r} in TORCH_CUDA_ARCH_LIST={raw!r}. Known: {', '.join(ARCH_TO_SM)}" + ) + arches.append((base, ARCH_TO_SM[base])) + if not arches: + raise SystemExit("TORCH_CUDA_ARCH_LIST is empty; nothing to compile") + return arches + + +def gencode_block(arches: list[tuple[str, int]], indent: str) -> str: + lines: list[str] = [f"{indent}# Patched by patch_mamba_cuda_arches.py from TORCH_CUDA_ARCH_LIST"] + for base, sm in arches: + lines.append(f'{indent}cc_flag.append("-gencode")') + lines.append(f'{indent}cc_flag.append("arch=compute_{sm},code=sm_{sm}") # {base}') + return "\n".join(lines) + "\n" + + +# Matches the hardcoded compute_75 start through the version-gated arch block, +# stopping before the CXX11 ABI HACK comment present in both setup.py files. +GENCODE_BLOCK_RE = re.compile( + r"(?P[ \t]*)cc_flag\.append\(\"-gencode\"\)\n" + r"[ \t]*cc_flag\.append\(\"arch=compute_75,code=sm_75\"\)\n" + r"(?:.*\n)*?" + r"(?=[ \t]*# HACK: The compiler flag)", + re.MULTILINE, +) + +NVCC_THREADS_RE = re.compile( + r"def append_nvcc_threads\(nvcc_extra_args\):\n" + r"[ \t]*return nvcc_extra_args \+ \[\"--threads\", \"4\"\]\n" +) + + +def patch_file(path: Path, arches: list[tuple[str, int]]) -> None: + text = path.read_text() + match = GENCODE_BLOCK_RE.search(text) + if match is None: + raise SystemExit(f"{path}: could not find hardcoded gencode block to patch") + + indent = match.group("indent") + text = GENCODE_BLOCK_RE.sub(gencode_block(arches, indent), text, count=1) + + new_threads, n = NVCC_THREADS_RE.subn( + "def append_nvcc_threads(nvcc_extra_args):\n" + ' return nvcc_extra_args + ["--threads", os.getenv("NVCC_THREADS", "4")]\n', + text, + count=1, + ) + if n != 1: + raise SystemExit(f"{path}: could not patch append_nvcc_threads") + text = new_threads + + path.write_text(text) + sm_list = ", ".join(f"sm_{sm}" for _, sm in arches) + print(f"Patched {path}: gencodes -> {sm_list}; NVCC_THREADS via env") + + +def main(argv: list[str]) -> None: + if len(argv) < 2: + raise SystemExit(f"usage: {argv[0]} setup.py [setup.py ...]") + + raw = os.environ.get("TORCH_CUDA_ARCH_LIST", "").strip() + if not raw: + raise SystemExit("TORCH_CUDA_ARCH_LIST must be set when patching") + + arches = parse_arch_list(raw) + for arg in argv[1:]: + patch_file(Path(arg), arches) + + +if __name__ == "__main__": + main(sys.argv) diff --git a/scripts/slurm/enroot/README.md b/scripts/slurm/enroot/README.md index 5b192b1d..5528790e 100644 --- a/scripts/slurm/enroot/README.md +++ b/scripts/slurm/enroot/README.md @@ -1,6 +1,6 @@ # slurm/enroot — Container Image Build -Builds the top-level [`Dockerfile`](../../Dockerfile) and converts the result to an enroot `.sqsh` for use with `srun --container-image=...` / `pyxis` on SLURM clusters. +Builds the top-level [`Dockerfile`](../../../Dockerfile) and converts the result to an enroot `.sqsh` for use with `srun --container-image=...` / `pyxis` on SLURM clusters. ## Build @@ -11,17 +11,21 @@ PLATFORM=arm64 bash build_sqsh.sh # GB200 (ARM64, built via qemu emulation) The script selects per-platform `--build-arg` values: -| `PLATFORM` | `TORCH_CUDA_ARCH_LIST` | `MAX_JOBS` | Target HW | -| ---------- | ---------------------- | ---------- | ------------------- | -| `x86_64` | `9.0` | unset | H100 | -| `arm64` | `10.0;12.0` | `2` | GB200 (B200 / 5090) | +| `PLATFORM` | `TORCH_CUDA_ARCH_LIST` | `MAX_JOBS` | `NVCC_THREADS` | Target HW | +| ---------- | ---------------------- | ---------- | -------------- | ------------------- | +| `x86_64` | `9.0` | unset | `4` | H100 | +| `arm64` | `10.0;12.0` | `1` | `1` | GB200 (B200 / 5090) | -`MAX_JOBS=2` on arm64 caps parallel nvcc jobs to avoid OOM under qemu emulation. On x86_64 it stays unset (parallel) for fastest builds. +`MAX_JOBS=1` / `NVCC_THREADS=1` on arm64 serializes work to reduce OOM/gcc-ICE failures under QEMU. Upstream `mamba-ssm` / `causal-conv1d` ignore `TORCH_CUDA_ARCH_LIST`; the Dockerfile patches their `setup.py` so only the arches above are compiled. On x86_64, `MAX_JOBS` stays unset (parallel) for fastest builds. + +**ARM64 on x86 hosts:** `PLATFORM=arm64` cross-builds via QEMU. Apex/mamba compilation can take many hours and may fail with `gcc: internal compiler error: Segmentation fault` if the host is memory-constrained. Keep ≥64GB combined free RAM+swap (`free -h`); Docker Engine on Linux uses host memory (no separate VM slider). Prefer a native `aarch64` host when possible. ## Override -| Env var | Default | -| ------------- | --------------------------------- | -| `PLATFORM` | `x86_64` | -| `DOCKER_TAG` | `nvsubquadratic:${PLATFORM}` | -| `OUTPUT_SQSH` | `nvsubquadratic-${PLATFORM}.sqsh` | +| Env var | Default | +| -------------- | --------------------------------- | +| `PLATFORM` | `x86_64` | +| `DOCKER_TAG` | `nvsubquadratic:${PLATFORM}` | +| `OUTPUT_SQSH` | `nvsubquadratic-${PLATFORM}.sqsh` | +| `MAX_JOBS` | platform default (see table) | +| `NVCC_THREADS` | platform default (see table) | diff --git a/scripts/slurm/enroot/build_sqsh.sh b/scripts/slurm/enroot/build_sqsh.sh index e5a97566..a30e64ef 100755 --- a/scripts/slurm/enroot/build_sqsh.sh +++ b/scripts/slurm/enroot/build_sqsh.sh @@ -9,31 +9,80 @@ # PLATFORM x86_64 (default, H100) | arm64 (GB200) # DOCKER_TAG image tag (default: nvsubquadratic:) # OUTPUT_SQSH output file (default: nvsubquadratic-.sqsh) +# MAX_JOBS parallel nvcc/ninja jobs (arm64 default: 1) +# NVCC_THREADS nvcc --threads (arm64 default: 1) +# INSTALL_BASELINES turn ON both benchmark baselines below (default: false). +# Use for the benchmark image: INSTALL_BASELINES=true ./build_sqsh.sh +# INSTALL_MAMBA build the Mamba2 baseline (default: INSTALL_BASELINES; true = slow +# QEMU source build; false → benchmark's Mamba2 series 'unavailable') +# INSTALL_FA4 build the FlashAttention-4 baseline (default: INSTALL_BASELINES; +# the alpha flash-attn-4 wheel) set -euo pipefail PLATFORM="${PLATFORM:-x86_64}" case "${PLATFORM}" in - x86_64) DOCKER_PLATFORM="linux/amd64"; TARGET_HW="H100 (x86-64)"; CUDA_ARCHS="9.0"; MAX_JOBS="" ;; - arm64) DOCKER_PLATFORM="linux/arm64"; TARGET_HW="GB200 (ARM64)"; CUDA_ARCHS="10.0;12.0"; MAX_JOBS="2" ;; + x86_64) DOCKER_PLATFORM="linux/amd64"; TARGET_HW="H100 (x86-64)"; CUDA_ARCHS_DEFAULT="9.0"; MAX_JOBS_DEFAULT=""; NVCC_THREADS_DEFAULT="4" ;; + arm64) DOCKER_PLATFORM="linux/arm64"; TARGET_HW="GB200 (ARM64)"; CUDA_ARCHS_DEFAULT="10.0;12.0"; MAX_JOBS_DEFAULT="1"; NVCC_THREADS_DEFAULT="1" ;; *) echo "Error: unknown PLATFORM=${PLATFORM}. Use x86_64 or arm64."; exit 1 ;; esac +# Override with CUDA_ARCHS=... (e.g. CUDA_ARCHS=10.0 to ease QEMU builds). +CUDA_ARCHS="${CUDA_ARCHS:-${CUDA_ARCHS_DEFAULT}}" +MAX_JOBS="${MAX_JOBS:-${MAX_JOBS_DEFAULT}}" +NVCC_THREADS="${NVCC_THREADS:-${NVCC_THREADS_DEFAULT}}" + +if [[ "${PLATFORM}" == "arm64" && "$(uname -m)" != "aarch64" ]]; then + echo "Warning: building linux/arm64 on $(uname -m) uses QEMU emulation." + echo " Apex/mamba CUDA compiles are slow and may ICE/OOM (gcc segfault)." + echo " Defaults: MAX_JOBS=1 NVCC_THREADS=1; mamba gencodes narrowed to ${CUDA_ARCHS}." + echo " Keep plenty of free host RAM+swap (64GB+ combined recommended)." + + # An OUTDATED qemu-user emulator mis-emulates the arm64 nvcc and crashes it with + # SIGSEGV (even on `nvcc -V`), failing the apex/mamba builds. Register a current + # qemu-aarch64 via tonistiigi/binfmt so cross-compilation works. One-time per host + # boot; skip with SKIP_BINFMT_SETUP=1 if your host QEMU is already current. + if [[ "${SKIP_BINFMT_SETUP:-0}" != "1" ]]; then + echo "Refreshing QEMU arm64 emulation (tonistiigi/binfmt --install arm64)..." + docker run --privileged --rm tonistiigi/binfmt --install arm64 + fi +fi + +# Preflight: free RAM matters more than MAX_JOBS under QEMU (single TUs still spike). +if command -v free >/dev/null 2>&1; then + avail_gb=$(free -g | awk '/^Mem:/{print $7}') + swap_gb=$(free -g | awk '/^Swap:/{print $2}') + if [[ "${PLATFORM}" == "arm64" && "${avail_gb}" =~ ^[0-9]+$ && "${avail_gb}" -lt 24 ]]; then + echo "Warning: only ~${avail_gb}GiB MemAvailable (swap=${swap_gb}GiB)." + echo " Free RAM or add swap before retrying; gcc ICE usually means OOM." + fi +fi + SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" DOCKER_TAG="${DOCKER_TAG:-nvsubquadratic:${PLATFORM}}" OUTPUT_SQSH="${OUTPUT_SQSH:-${SCRIPT_DIR}/nvsubquadratic-${PLATFORM}.sqsh}" +# Benchmark-only baselines, OFF by default (leaner/faster image). Turn on BOTH with +# INSTALL_BASELINES=true (the benchmark-image shortcut), or each individually. +INSTALL_BASELINES="${INSTALL_BASELINES:-false}" +INSTALL_MAMBA="${INSTALL_MAMBA:-${INSTALL_BASELINES}}" +INSTALL_FA4="${INSTALL_FA4:-${INSTALL_BASELINES}}" echo "Platform: ${DOCKER_PLATFORM} (${TARGET_HW})" echo "Image: ${DOCKER_TAG}" echo "Output: ${OUTPUT_SQSH}" +echo "Arches: ${CUDA_ARCHS} MAX_JOBS=${MAX_JOBS:-unset} NVCC_THREADS=${NVCC_THREADS}" +echo "Baselines: mamba=${INSTALL_MAMBA} fa4=${INSTALL_FA4}" docker buildx build \ --platform "${DOCKER_PLATFORM}" \ --build-arg TORCH_CUDA_ARCH_LIST="${CUDA_ARCHS}" \ --build-arg MAX_JOBS="${MAX_JOBS}" \ + --build-arg NVCC_THREADS="${NVCC_THREADS}" \ + --build-arg INSTALL_MAMBA="${INSTALL_MAMBA}" \ + --build-arg INSTALL_FA4="${INSTALL_FA4}" \ -t "${DOCKER_TAG}" \ -f "${REPO_ROOT}/Dockerfile" \ --load \ @@ -42,5 +91,6 @@ docker buildx build \ enroot import -o "${OUTPUT_SQSH}" "dockerd://${DOCKER_TAG}" echo "Done: ${OUTPUT_SQSH}" -echo " PLATFORM=arm64 ./build_sqsh.sh # GB200" -echo " PLATFORM=x86_64 ./build_sqsh.sh # H100 (default)" +echo " PLATFORM=arm64 ./build_sqsh.sh # GB200, lean (no baselines)" +echo " PLATFORM=x86_64 ./build_sqsh.sh # H100 (default), lean" +echo " INSTALL_BASELINES=true PLATFORM=arm64 ./build_sqsh.sh # + Mamba2 & FA4 (benchmark image)" diff --git a/scripts/slurm/setup_env.sh b/scripts/slurm/setup_env.sh index 5c3cfb80..cd0e2c8f 100644 --- a/scripts/slurm/setup_env.sh +++ b/scripts/slurm/setup_env.sh @@ -38,10 +38,10 @@ python3 -c "import urllib.request; urllib.request.urlretrieve('https://bootstrap python3 /tmp/get-pip.py pip --version -# Install torch cu129 -echo "=== Installing torch cu129 ===" +# Install torch cu130 +echo "=== Installing torch cu130 ===" pip install torch==2.10.0 torchvision==0.25.0 torchaudio==2.10.0 \ - --index-url https://download.pytorch.org/whl/cu129 + --index-url https://download.pytorch.org/whl/cu130 # Verify torch sees the GPU echo "=== Verifying torch+CUDA ===" diff --git a/scripts/slurm/submit_forward_time_flash_kernels.sh b/scripts/slurm/submit_forward_time_flash_kernels.sh new file mode 100755 index 00000000..f615aa01 --- /dev/null +++ b/scripts/slurm/submit_forward_time_flash_kernels.sh @@ -0,0 +1,67 @@ +#!/bin/bash +# ============================================================================= +# Attention-KERNEL comparison: SDPA vs FlexAttention vs FlashAttention-4 vs +# HyenaND — a ready-to-submit wrapper over submit_forward_time_nd.sh. +# +# This is the SECOND of the two forward-time runs, and deliberately a different +# config from the "reach-to-16M" sweep: +# +# * reach sweep (submit_forward_time_nd.sh defaults): hidden 8, head_dim 4 — +# the largest shared width that keeps the qkv tensor under torch's 2^31 index +# limit out to 16M tokens. flex/fa4 CANNOT run there (they require head_dim +# >= 16), and SDPA there is a weak non-flash fallback. +# * this sweep: hidden 512 / 4 heads => head_dim 128 — the regime FlashAttention +# (2/3/4), cuDNN SDPA, and FlexAttention are actually OPTIMIZED for (FA3/FA4 +# on Hopper/Blackwell are tuned for head_dim 128). head_dim 16 is only the +# eligibility floor; 64/128 is where these kernels — and real models — live. +# At hidden 512 the 32-bit-index wall lands at ~1M tokens (2D), which is far +# past where attention time-walls, so the comparison is complete. +# +# Runs HyenaND and bidirectional Mamba2 alongside the three attention kernels, so +# the plot keeps the Figure-1 punchline (subquadratic scaling past every attention +# kernel's O(L^2) wall, with Mamba2 as the O(L) SSM reference that walls on its +# causal_conv1d grid limit ~1-2M). The fa4 series needs flash-attn-4 in the image +# (rebuild the sqsh with the version-locked layer); absent or mismatched, it +# records 'unavailable'/'error' and the plotter simply omits it. Mamba2 needs +# mamba-ssm (already in the image); absent, it is likewise omitted. +# +# Usage (defaults to 2D; override DATA_DIM / any submit_forward_time_nd.sh var): +# scripts/slurm/submit_forward_time_flash_kernels.sh +# DATA_DIM=1 scripts/slurm/submit_forward_time_flash_kernels.sh +# # lighter head_dim 64 instead of 128: +# HIDDEN_DIM=256 scripts/slurm/submit_forward_time_flash_kernels.sh +# ============================================================================= +set -euo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +DATA_DIM="${DATA_DIM:-2}" +HIDDEN_DIM="${HIDDEN_DIM:-512}" # /4 heads => head_dim 128 (flash-optimized) +NUM_HEADS="${NUM_HEADS:-4}" +MAMBA_HEADDIM="${MAMBA_HEADDIM:-64}" # d_inner = 512*expand(2) = 1024 / 64 => 16 heads +MIXERS="${MIXERS:-hyena attention flex fa4 mamba}" + +# Per-dim reach: largest R keeping 3*HIDDEN_DIM*R^DATA_DIM under 2^31 at hidden 512. +case "${DATA_DIM}" in + 1) RESOLUTIONS="${RESOLUTIONS:-4096 16384 65536 262144 1048576}" ;; + 2) RESOLUTIONS="${RESOLUTIONS:-64 128 256 512 1024}" ;; + 3) RESOLUTIONS="${RESOLUTIONS:-16 32 64}" ;; + *) echo "DATA_DIM must be 1, 2, or 3 (got '${DATA_DIM}')"; exit 1 ;; +esac + +echo "[flash-kernels] DATA_DIM=${DATA_DIM} hidden=${HIDDEN_DIM} heads=${NUM_HEADS} " \ + "(head_dim=$((HIDDEN_DIM / NUM_HEADS))) mamba_headdim=${MAMBA_HEADDIM} mixers='${MIXERS}'" +echo "[flash-kernels] resolutions='${RESOLUTIONS}'" + +# Distinct output stem so results/plots don't clobber the reach sweep's +# forward_time_${DATA_DIM}d.* files. +OUT="forward_time_flash_${DATA_DIM}d" + +# Hand off to the main sbatch script. Export the knobs into the environment and +# propagate the whole env with --export=ALL, rather than packing space-containing +# values (MIXERS, RESOLUTIONS) into an --export=KEY=VAL,... string — Slurm splits +# that list on commas and mishandles embedded spaces on some versions. +export DATA_DIM MIXERS HIDDEN_DIM NUM_HEADS MAMBA_HEADDIM RESOLUTIONS OUT +exec sbatch \ + --job-name="nvsubq-fwd-flash-${DATA_DIM}d" \ + --export=ALL \ + "${HERE}/submit_forward_time_nd.sh" diff --git a/scripts/slurm/submit_forward_time_nd.sh b/scripts/slurm/submit_forward_time_nd.sh new file mode 100755 index 00000000..124fce04 --- /dev/null +++ b/scripts/slurm/submit_forward_time_nd.sh @@ -0,0 +1,160 @@ +#!/bin/bash +#SBATCH --account=healthcareeng_bionemo +#SBATCH --nodes=1 +#SBATCH --partition=36x2-a01r +#SBATCH --ntasks-per-node=1 +#SBATCH --time=03:00:00 +#SBATCH --mem=0 +#SBATCH --job-name=nvsubq-fwdtime-2d +#SBATCH --mail-type=FAIL +#SBATCH --exclusive + +set -x + +# ============================================================================= +# 2D forward-time-vs-resolution benchmark (the 2D analogue of Figure 1, right). +# +# Times a single HyenaND / Attention / Mamba2 layer at growing square +# resolutions on ONE GPU and writes a JSONL, then renders the paper-style +# log-log plot. Single-GPU microbenchmark (no torchrun / CP): we pin +# CUDA_VISIBLE_DEVICES=0 even on an exclusive multi-GPU node. +# +# Usage (defaults target the GB200 arm64 image + partition below): +# sbatch scripts/slurm/submit_forward_time_nd.sh +# +# # H100 instead: use the x86_64 image and a Hopper partition (the defaults +# # sweep unchanged — the ~15 GB peak at 2048^2 fits an 80 GB H100, and the +# # ceiling is the 4096^2 torch 32-bit-index wall, not memory): +# SQSH_PATH=/lustre/.../enroot/nvsubquadratic-x86_64.sqsh \ +# sbatch --partition= scripts/slurm/submit_forward_time_nd.sh +# +# The defaults sweep 64x64 .. 8192x8192 (4K .. 64M tokens) at hidden_dim=64 with +# the circular (single-grid) kernel, so a *single* run produces the whole story: +# the HyenaND/Attention (+Mamba2 if installed) comparison where they overlap, and +# HyenaND scaling far past attention's O(L^2) wall. Observed reach at hidden_dim=64 +# is ~2048^2 (4M tokens, ~1200x faster than attention there); at 4096^2+ every +# operator hits torch's 2^31-element 32-bit-indexing limit and errors (that top-end +# 'x' is a torch limit, not subq_ops or memory). Adaptive timing keeps the slow +# high-R points cheap; MAX_SECONDS marks where attention becomes an 'x'. +# +# All mixers ALWAYS share one hidden_dim (one comparable plot). Peak memory at +# 2048^2 is ~15 GB, so hidden_dim=64 fits an 80 GB H100 and a GB200 alike; the +# ceiling is the 32-bit-index wall above, not memory. Every parameter is an +# environment override (one script, not modes); these change the shared value / +# range for the WHOLE run, they do not mix dims: +# # all three at hidden 256, non-circular (whole run) — caps ~4096^2: +# HIDDEN_DIM=256 GRID_TYPE=double RESOLUTIONS="64 128 256 512 1024 2048 4096" \ +# sbatch scripts/slurm/submit_forward_time_nd.sh +# # HyenaND-reach-to-8K, apple-to-apple: all three at hidden 8 (the largest +# # shared width whose qkv tensor 3*hidden*R^2 stays under 2^31 at 8192^2), so +# # HyenaND continues to 64M while Attention (~16M) and Mamba (~4M) x out at the +# # SAME config. num_heads/mamba_headdim reduced to divide hidden 8: +# HIDDEN_DIM=8 NUM_HEADS=2 MAMBA_HEADDIM=8 \ +# sbatch scripts/slurm/submit_forward_time_nd.sh +# # Attention-KERNEL comparison (SDPA vs FlexAttention vs FlashAttention-4 vs +# # HyenaND) — use the ready-made wrapper, which sets head_dim 128 (hidden 512 / +# # 4 heads), the regime these flash kernels are OPTIMIZED for. flex/fa4 only +# # *accept* head_dim >= 16, but 64/128 is where they (and real models) run fast; +# # 16/32 sits on a slow small-tile path. Needs flash-attn-4 in the image for the +# # fa4 series (absent -> 'unavailable', omitted from the plot): +# scripts/slurm/submit_forward_time_flash_kernels.sh # 2D, head_dim 128 +# DATA_DIM=1 scripts/slurm/submit_forward_time_flash_kernels.sh +# # HyenaND only: +# MIXERS=hyena sbatch scripts/slurm/submit_forward_time_nd.sh +# +# Prerequisites: +# * The repo (with these benchmark scripts) is on the cluster at CODE_PATH; it +# is mounted and run with PYTHONPATH=. so the latest code is always used. +# * The image (SQSH_PATH) has the subq_ops v2 CUDA kernels ([cuda] extra). +# Optional: mamba-ssm for the Mamba2 series — absent, it shows as 'x' and +# the other two still render. +# ============================================================================= + +# ── Paths (adjust to your cluster layout) ──────────────────────────────────── +SQSH_PATH="${SQSH_PATH:-/lustre/fsw/healthcareeng_bionemo/farhadr/enroot/nvsubquadratic-arm64.sqsh}" +CODE_PATH="${CODE_PATH:-/lustre/fsw/healthcareeng_bionemo/farhadr/nvsubquadratic_workdir/nvSubquadratic}" +CODE_MOUNT=/workspace/nvsubq +RESULTS_HOST="${CODE_PATH}/benchmarks/results" +MOUNTS="${CODE_PATH}:${CODE_MOUNT}" + +mkdir -p "${RESULTS_HOST}" + +# ── Benchmark parameters (override any from the environment) ────────────────── +DATA_DIM="${DATA_DIM:-2}" # 1 | 2 | 3 (L = R^DATA_DIM) +MIXERS="${MIXERS:-attention hyena mamba}" +# Shared width. hidden 8 (num_heads 2, mamba_headdim 8) keeps the qkv tensor +# 3*hidden*R^N under torch's 2^31 index limit through ~16M tokens in every dim. +HIDDEN_DIM="${HIDDEN_DIM:-8}" +NUM_HEADS="${NUM_HEADS:-2}" +MAMBA_HEADDIM="${MAMBA_HEADDIM:-8}" +GRID_TYPE="${GRID_TYPE:-single}" +# Per-dim resolution sweep (R), each topping out near 16M tokens. Override with RESOLUTIONS=. +# 1D: L = R 2D: L = R^2 3D: L = R^3 +case "${DATA_DIM}" in + 1) RESOLUTIONS="${RESOLUTIONS:-4096 16384 65536 262144 1048576 4194304 16777216}" ;; + 2) RESOLUTIONS="${RESOLUTIONS:-64 128 256 512 1024 2048 4096}" ;; + 3) RESOLUTIONS="${RESOLUTIONS:-16 32 64 128 256}" ;; + *) echo "DATA_DIM must be 1, 2, or 3 (got '${DATA_DIM}')"; exit 1 ;; +esac +FFT_BACKEND="${FFT_BACKEND:-subq_ops}" # 1D=causal fused, 2D=fused; 3D auto-falls back to torch_fft +DTYPE="${DTYPE:-bf16}" +BATCH_SIZE="${BATCH_SIZE:-1}" +NUM_WARMUP="${NUM_WARMUP:-10}" +NUM_ITERS="${NUM_ITERS:-30}" +MAX_SECONDS="${MAX_SECONDS:-300}" +OUT="${OUT:-forward_time_${DATA_DIM}d}" + +JSONL="${CODE_MOUNT}/benchmarks/results/${OUT}.jsonl" +PNG="${CODE_MOUNT}/benchmarks/results/${OUT}.png" +MEM_PNG="${CODE_MOUNT}/benchmarks/results/${OUT}_memory.png" + +# ── In-container command ───────────────────────────────────────────────────── +read -r -d '' COMMAND </dev/null; then + PYTHONPATH=. python scripts/visualization/visualize_forward_time_nd.py \ + --input ${JSONL} --out ${PNG} --no-fail-markers + PYTHONPATH=. python scripts/visualization/visualize_forward_time_nd.py \ + --input ${JSONL} --out ${MEM_PNG} --metric memory --no-fail-markers +else + echo "[plot] matplotlib unavailable — plot on the login node from ${OUT}.jsonl." +fi +EOF + +srun \ + --output "${RESULTS_HOST}/${OUT}-%j.out" \ + --error "${RESULTS_HOST}/${OUT}-error-%j.out" \ + --container-image="${SQSH_PATH}" \ + --container-mounts "${MOUNTS}" \ + bash -c "${COMMAND}" + +set +x diff --git a/scripts/visualization/visualize_forward_time_nd.py b/scripts/visualization/visualize_forward_time_nd.py new file mode 100644 index 00000000..dcf38b4c --- /dev/null +++ b/scripts/visualization/visualize_forward_time_nd.py @@ -0,0 +1,417 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Plot the forward-time-vs-resolution sweep (1D / 2D / 3D) in paper (Figure 1) style. + +Dimensionality, hardware, and per-operator gaps are read from the JSONL, so the +title (``{N}D Forward-Time Scaling``), x-axis (``L = R^N``), tick exponents, and +device subtitle adapt automatically. + +Reads the JSONL written by +``benchmarks/benchmark_forward_time_nd_resolution.py`` and renders forward time +(log y) against 2D token count ``L = R^2`` (log x) for HyenaND / Attention / +Mamba2 — the 2D analogue of Figure 1 (right). Points that ran out of memory or +exceeded the time budget are drawn as a green "OOM" ``x`` (as in Figure 1), and +the largest HyenaND-vs-Attention gap is annotated (the "339x"-style label). + +The plotter is decoupled from the GPU run: it needs only the JSONL and +matplotlib, so it can run on a login node / laptop. + +Usage:: + + python scripts/visualization/visualize_forward_time_nd.py \\ + --input benchmarks/results/forward_time.jsonl \\ + --out benchmarks/results/forward_time.png +""" + +from __future__ import annotations + +import argparse +import json +import math +from collections import defaultdict +from pathlib import Path + +import matplotlib +import matplotlib.pyplot as plt + + +matplotlib.rcParams["pdf.fonttype"] = 42 +matplotlib.rcParams["ps.fonttype"] = 42 + + +# mixer key -> (legend label, colour, marker) — colours match Figure 1. The three +# attention kernels (SDPA / FlexAttention / FA4) share a cool-toned family so the +# plot reads them as attention variants, distinguished by shade + marker. +SERIES = { + "hyena": {"label": "HyenaND (nSubQ)", "color": "#C0392B", "marker": "s"}, + "attention": {"label": "Attention (SDPA)", "color": "#3B6FB6", "marker": "o"}, + "flex": {"label": "FlexAttention", "color": "#17A2B8", "marker": "^"}, + "fa4": {"label": "FlashAttention-4", "color": "#9467BD", "marker": "v"}, + "mamba": {"label": "Mamba2", "color": "#7A8B3C", "marker": "D"}, +} +SERIES_ORDER = ["hyena", "attention", "flex", "fa4", "mamba"] + +# Attention-family mixers (all avoid the O(L^2) score matrix via flash-class kernels). +_ATTN_MIXERS = ("attention", "flex", "fa4") + +FAIL_STATUS = {"oom", "error", "timeout"} +FAIL_COLOR = "#2E7D32" # green "OOM x", as in Figure 1 + +# Approx usable HBM per GPU (GB), for the memory-metric ceiling line. +_GPU_MEM_GB = { + "GB300": 279.0, + "B300": 279.0, + "GB200": 186.0, + "B200": 186.0, + "H200": 141.0, + "H100": 80.0, + "A100": 80.0, + "A6000": 48.0, +} + + +def _gpu_mem_gb(device: str | None) -> float | None: + if not device: + return None + for key, gb in _GPU_MEM_GB.items(): + if key in device: + return gb + return None + + +def _format_seq(n: int) -> str: + if n >= 1024 * 1024: + v = n / (1024 * 1024) + return f"{v:.0f}M" if v == int(v) else f"{v:.1f}M" + if n >= 1024: + v = n / 1024 + return f"{v:.0f}K" if v == int(v) else f"{v:.1f}K" + return str(n) + + +def _tick_label(seq_len: int, data_dim: int) -> str: + if data_dim == 1: + return _format_seq(seq_len) # 1D: L == R, single line + r = round(seq_len ** (1.0 / data_dim)) + sup = {2: "²", 3: "³"}.get(data_dim, f"^{data_dim}") + return f"{_format_seq(seq_len)}\n{r}{sup}" + + +def load_rows(path: Path) -> list[dict]: + rows = [] + with path.open() as fh: + for line in fh: + line = line.strip() + if line: + rows.append(json.loads(line)) + return rows + + +def make_plot( + rows: list[dict], + out_path: Path, + show_fail_markers: bool = True, + metric: str = "time", + attn_heads: int | None = None, + gpu_mem_gb: float | None = None, +) -> None: + plt.rcParams.update( + { + "font.family": "serif", + "font.size": 11, + "axes.titlesize": 12, + "axes.titleweight": "bold", + "axes.labelsize": 11, + "xtick.labelsize": 9, + "ytick.labelsize": 9, + "axes.spines.top": False, + "axes.spines.right": False, + "axes.linewidth": 0.8, + "xtick.major.width": 0.8, + "ytick.major.width": 0.8, + } + ) + + mkey = "ms" if metric == "time" else "mem_gb" + + # Drop resolutions where no operator produced a value for this metric, so the + # x-axis ends at the last reachable point. + ok_seq = {int(r["seq_len"]) for r in rows if r.get("status") == "ok" and r.get(mkey) is not None} + rows = [r for r in rows if int(r["seq_len"]) in ok_seq] + + # Dimensionality (L = R^data_dim) and hardware, inferred from the data. + def _dim_of(r: dict) -> int: + if r.get("data_dim"): + return int(r["data_dim"]) + seq_len, res = int(r["seq_len"]), int(r["resolution"]) + return max(1, round(math.log(seq_len) / math.log(res))) if res > 1 else 2 + + data_dim = _dim_of(rows[0]) + + def _common(key): # dominant non-null value (config fields are uniform per run) + vals = [r.get(key) for r in rows if r.get(key) is not None] + return max(set(vals), key=vals.count) if vals else None + + device = _common("device") + # Config subtitle: distinguishes runs that share a title — e.g. the hidden-8 + # (head_dim 4) reach sweep vs the hidden-512 (head_dim 128) flash comparison. + # head_dim is what makes attention flash-eligible, so it explains why the + # HyenaND-vs-attention gap differs between the two sets. + _hidden, _heads, _dt = _common("hidden_dim"), _common("num_heads"), _common("dtype") + _sub = [device] if device else [] + if _hidden: + _cfg = f"hidden {_hidden}" + if _heads: + _cfg += f", head_dim {_hidden // _heads}" + _sub.append(_cfg) + if _dt: + _sub.append(str(_dt)) + subtitle = " · ".join(_sub) + + # Group by mixer -> {seq_len: row} + by_mixer: dict[str, dict[int, dict]] = defaultdict(dict) + for r in rows: + by_mixer[r["mixer"]][int(r["seq_len"])] = r + + all_seq = sorted({int(r["seq_len"]) for r in rows}) + ok_vals = [r[mkey] for r in rows if r.get("status") == "ok" and r.get(mkey) is not None] + if not ok_vals: + raise SystemExit(f"No successful ('ok') {mkey} values in the JSONL — nothing to plot.") + ceiling = max(ok_vals) # where OOM/timeout x-marks with no measured value are parked + + # Widen with the tick count so the two-line "L / R^2" labels don't crowd. + fig_w = min(9.0, max(3.6, 0.6 * len(all_seq) + 1.2)) + fig, ax = plt.subplots(figsize=(fig_w, 3.2), constrained_layout=True) + + fail_x, fail_y = [], [] + for mixer in SERIES_ORDER: + pts = by_mixer.get(mixer) + if not pts: + continue + cfg = SERIES[mixer] + xs = sorted(pts) + + ok_x = [x for x in xs if pts[x].get("status") == "ok" and pts[x].get(mkey) is not None] + + # A series that never produced a single timing never actually ran here + # (missing dep, build/version error, or OOM even at the smallest grid) — + # that is a setup failure, not a scaling wall, so omit it entirely rather + # than paint a wall of x's across every resolution. + if not ok_x: + n_fail = sum(1 for x in xs if pts[x].get("status") in FAIL_STATUS) + if n_fail: + print(f"Omitting '{mixer}': 0/{len(xs)} points ran — {pts[xs[0]].get('detail', '')[:90]}") + continue + + ax.plot( + ok_x, + [pts[x][mkey] for x in ok_x], + color=cfg["color"], + marker=cfg["marker"], + linestyle="-", + linewidth=1.6, + markersize=5.5, + markeredgecolor="black", + markeredgewidth=0.5, + label=cfg["label"], + ) + + # Mark only the FIRST failure per series — the wall where it stops. Later + # failures at larger L are redundant (and clutter/overlap other series). + # A wall-timed timeout carries a measured (over-budget) ms → place the x + # there; a predictive skip has ms=None → park it at the ceiling. + fails = [x for x in xs if pts[x].get("status") in FAIL_STATUS] + if fails: + x = fails[0] # xs is sorted ascending → smallest-L failure + fail_x.append(x) + fail_y.append(pts[x][mkey] if pts[x].get(mkey) is not None else ceiling) + + if fail_x and show_fail_markers: + ax.scatter( + fail_x, + fail_y, + marker="x", + s=70, + linewidths=2.0, + color=FAIL_COLOR, + zorder=5, + label="OOM / timeout", + ) + + # ── Metric-specific overlays ────────────────────────────────────────────── + gaps: dict[str, tuple[int, float] | None] = {} + mem_ceiling: float | None = None + if metric == "time": + # HyenaND vs each slower operator: one double-arrow "N×" at its max gap + # (vs Attention at the largest shared L; vs Mamba at its last point). + def _annotate_gap(slow_key: str, fast_key: str) -> tuple[int, float] | None: + s_pts, f_pts = by_mixer.get(slow_key, {}), by_mixer.get(fast_key, {}) + best = None # (x, ratio, slow, fast) + for x in all_seq: + s, f = s_pts.get(x), f_pts.get(x) + if not s or not f or s.get("status") != "ok" or f.get("status") != "ok": + continue + if s.get("ms") is None or f.get("ms") is None: + continue + ratio = s["ms"] / f["ms"] + if best is None or ratio > best[1]: + best = (x, ratio, s["ms"], f["ms"]) + if best is None or best[1] < 1.5: + return None + x, ratio, s_ms, f_ms = best + ax.annotate( + "", + xy=(x, s_ms), + xytext=(x, f_ms), + arrowprops={"arrowstyle": "<->", "color": "0.4", "lw": 0.8, "mutation_scale": 6}, + ) + ax.text( + x, + math.sqrt(s_ms * f_ms), + f" {ratio:.0f}×", + color="0.25", + fontsize=9, + ha="left", + va="center", + fontweight="bold", + ) + return (x, ratio) + + gaps = {slow: _annotate_gap(slow, "hyena") for slow in ("attention", "mamba")} + title = f"{data_dim}D Forward-Time Scaling" + ylabel = "Forward time (ms)" + else: + # O(L^2) memory a *materialized* (non-flash) attention needs for its + # H×L×L score matrix in the model dtype — the term flash avoids by + # recomputing (and pays for in time). The measured curves are all O(L); + # this dashed reference shows what attention would cost without flash, + # and where it exceeds GPU memory (~256K here). + heads = ( + attn_heads + or next( + (r.get("num_heads") for r in rows if r.get("mixer") in _ATTN_MIXERS and r.get("num_heads")), + None, + ) + or 2 + ) + dbytes = {"bf16": 2, "fp16": 2, "fp32": 4}.get(rows[0].get("dtype", "bf16"), 2) + batch = rows[0].get("batch_size", 1) or 1 + materialized = [batch * heads * (x**2) * dbytes / 1e9 for x in all_seq] + ax.plot( + all_seq, + materialized, + color=SERIES["attention"]["color"], + linestyle="--", + dashes=(4, 2), + linewidth=1.4, + label="Attention (materialized scores)", + ) + mem = gpu_mem_gb if gpu_mem_gb is not None else _gpu_mem_gb(device) + if mem: + ax.axhline(mem, color="0.45", linestyle=":", linewidth=1.0, zorder=1) + ax.text( + all_seq[-1], mem, f" {device or 'GPU'} memory", fontsize=7.5, color="0.4", va="bottom", ha="right" + ) + mem_ceiling = mem # applied as ylim after the log scale is set + title = f"{data_dim}D Peak-Memory Scaling" + ylabel = "Peak memory (GB)" + + fig.suptitle(title, fontsize=12, fontweight="bold") + if subtitle: + ax.set_title(subtitle, fontsize=9, fontweight="normal", color="0.4") + ax.set_xscale("log") + ax.set_yscale("log") + ax.set_xticks(all_seq) + ax.set_xticklabels([_tick_label(s, data_dim) for s in all_seq]) + ax.minorticks_off() + if data_dim == 1: + ax.set_xlabel("Sequence length (tokens)") + else: + ax.set_xlabel(f"{data_dim}D context (tokens $L = R^{data_dim}$)") + ax.set_ylabel(ylabel) + if mem_ceiling: # after set_yscale, else switching to log re-autoscales + ax.set_ylim(top=mem_ceiling * 4) + ax.grid(True, which="major", axis="both", linestyle=":", linewidth=0.5, alpha=0.6) + ax.legend(frameon=False, fontsize=8, loc="upper left", handlelength=1.6, borderaxespad=0.4) + + out_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(out_path, dpi=300, bbox_inches="tight") + pdf_path = out_path.with_suffix(".pdf") + fig.savefig(pdf_path, bbox_inches="tight") + plt.close(fig) + print(f"Saved: {out_path}") + print(f"Saved: {pdf_path}") + for slow, g in gaps.items(): + if g: + print(f"Max {slow}/HyenaND gap: {g[1]:.0f}x at L={g[0]}") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument( + "--input", + type=Path, + default=Path("benchmarks/results/forward_time.jsonl"), + help="JSONL produced by benchmark_forward_time_nd_resolution.py", + ) + parser.add_argument( + "--out", + type=Path, + default=Path("benchmarks/results/forward_time.png"), + ) + parser.add_argument( + "--no-fail-markers", + action="store_true", + help="Hide the green OOM/timeout 'x' markers (each curve just ends where it walls).", + ) + parser.add_argument( + "--metric", + choices=["time", "memory"], + default="time", + help="'time' = forward ms (Figure 1). 'memory' = peak GB, with an O(L^2) " + "materialized-attention reference line and the GPU-memory ceiling.", + ) + parser.add_argument( + "--attn-heads", + type=int, + default=None, + help="Head count for the materialized-scores memory line (default: read from the JSONL, else 2).", + ) + parser.add_argument( + "--gpu-mem-gb", + type=float, + default=None, + help="GPU memory (GB) for the ceiling line (default: inferred from the device name).", + ) + args = parser.parse_args() + + if not args.input.exists(): + raise SystemExit(f"Input JSONL not found: {args.input}") + rows = load_rows(args.input) + if not rows: + raise SystemExit(f"No rows in {args.input}") + make_plot( + rows, + args.out, + show_fail_markers=not args.no_fail_markers, + metric=args.metric, + attn_heads=args.attn_heads, + gpu_mem_gb=args.gpu_mem_gb, + ) + + +if __name__ == "__main__": + main() diff --git a/setup_conda_env.sh b/setup_conda_env.sh index 6663e3e8..58ebf38e 100644 --- a/setup_conda_env.sh +++ b/setup_conda_env.sh @@ -6,15 +6,15 @@ # # Prerequisites: # - conda (Miniforge / Miniconda / Anaconda) -# - NVIDIA GPU with CUDA 12.9 drivers +# - NVIDIA GPU with CUDA 13.0 drivers # - nvcc on PATH (included in the nvcr.io CUDA devel image; on bare metal -# install the CUDA 12.9 toolkit separately) +# install the CUDA 13.0 toolkit separately) set -euo pipefail ENV_NAME=nvsubquadratic PYTHON_VERSION=3.12 -TORCH_INDEX=https://download.pytorch.org/whl/cu129 +TORCH_INDEX=https://download.pytorch.org/whl/cu130 TORCH_VERSION=2.10.0 TORCHVISION_VERSION=0.25.0 @@ -32,7 +32,7 @@ ENV_PREFIX=$(conda env list | grep -E "^${ENV_NAME}[[:space:]]" | awk '{print $N PIP="${ENV_PREFIX}/bin/pip" # ── 2. Install PyTorch ──────────────────────────────────────────────────────── -echo "Installing PyTorch ${TORCH_VERSION} (CUDA 12.9)..." +echo "Installing PyTorch ${TORCH_VERSION} (CUDA 13.0)..." "${PIP}" install --no-cache-dir \ "torch==${TORCH_VERSION}" "torchvision==${TORCHVISION_VERSION}" \ --index-url "${TORCH_INDEX}" diff --git a/tests/conftest.py b/tests/conftest.py index 5279b710..b14ad6b6 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -25,11 +25,11 @@ def _subq_ops_version() -> tuple[int, ...]: - """Return the installed subquadratic-ops-torch-cu12 version as an int tuple.""" + """Return the installed subquadratic-ops-torch-cu13 version as an int tuple.""" try: from importlib.metadata import version - return tuple(int(x) for x in version("subquadratic-ops-torch-cu12").split(".")[:3]) + return tuple(int(x) for x in version("subquadratic-ops-torch-cu13").split(".")[:3]) except Exception: return (0, 0, 0)