Documentation

.github/workflows/build_wheels.yml

@@ -45,7 +45,7 @@ jobs:
             python-builds: 'cp310-* cp311-* cp312-*'
 
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v5
 
     - name: Free Disk Space
       run: |
@@ -62,7 +62,7 @@ jobs:
         df -h
 
     - name: Set up Python
-      uses: actions/setup-python@v5
+      uses: actions/setup-python@v6
       with:
         python-version: '3.11'
 
@@ -132,7 +132,7 @@ jobs:
         done
 
     - name: Upload wheel artifacts
-      uses: actions/upload-artifact@v4
+      uses: actions/upload-artifact@v7
       with:
         name: wheel-torch-${{ matrix.torch-version }}-${{ matrix.torch-cuda }}
         path: wheelhouse/*.whl

.github/workflows/docs.yaml

@@ -0,0 +1,92 @@
+name: Build and deploy documentation
+
+# Publishing happens only on a published release or a manual dispatch with
+# "deploy" checked (see the deploy job's `if:` guard). Pull requests targeting
+# main run the BUILD job only: they produce a downloadable `html-docs` artifact
+# for preview and never touch the published documentation.
+on:
+  release:
+    types: [published]
+  pull_request:
+    branches: [main]
+  workflow_dispatch:
+    inputs:
+      deploy:
+        description: "Publish to GitHub Pages (leave unchecked for a test build you can download and inspect offline)"
+        type: boolean
+        default: false
+
+# Allow the deploy job to publish to GitHub Pages.
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+# Allow one concurrent deployment; do not cancel an in-progress deploy.
+concurrency:
+  group: pages
+  cancel-in-progress: false
+
+jobs:
+  build-documentation:
+    name: Build documentation
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v5
+
+    - name: Set up Python 3.11
+      uses: actions/setup-python@v6
+      with:
+        python-version: "3.11"
+
+    # Same install logic as tests.yml: CPU-only torch, then an editable
+    # build of the package. autodoc imports torch_harmonics, so the compiled
+    # C++ extensions must be built (CPU-only needs no CUDA toolkit).
+    - name: Install dependencies
+      run: |
+        python3 -m pip install --upgrade pip setuptools wheel
+        python3 -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu
+
+    - name: Install package with docs dependencies
+      run: |
+        python3 -m pip install --no-build-isolation -e ".[docs]"
+
+    - name: Build HTML documentation
+      # -W turns Sphinx warnings (broken cross-references, malformed
+      # directives, etc.) into errors; --keep-going reports them all in one
+      # run instead of stopping at the first.
+      run: |
+        make -C docs html SPHINXOPTS="-W --keep-going"
+
+    # Plain downloadable artifact for offline inspection (a zip of the raw
+    # HTML). Produced on every run, including test builds that don't deploy.
+    - name: Upload HTML for download
+      uses: actions/upload-artifact@v7
+      with:
+        name: html-docs
+        path: docs/_build/html
+
+    # Special artifact consumed by the deploy job below (only used when the
+    # deploy job actually runs).
+    - name: Upload Pages artifact
+      uses: actions/upload-pages-artifact@v5
+      with:
+        path: docs/_build/html
+
+  deploy-documentation:
+    name: Deploy documentation
+    needs: build-documentation
+    # Publish on a release, or on a manual run only when "deploy" is checked.
+    # A manual run with deploy unchecked just builds the downloadable artifact.
+    if: github.event_name == 'release' || inputs.deploy
+    runs-on: ubuntu-latest
+
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+
+    steps:
+    - name: Deploy to GitHub Pages
+      id: deployment
+      uses: actions/deploy-pages@v5

.github/workflows/style.yml

@@ -10,12 +10,22 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v5
 
     - name: Set up Python 3.9
-      uses: actions/setup-python@v5
+      uses: actions/setup-python@v6
       with:
         python-version: "3.9"
 
+    # Run pre-commit directly instead of pre-commit/action@v3.0.1, whose
+    # bundled sub-actions still target the deprecated Node.js 20 runtime.
+    - name: Cache pre-commit environments
+      uses: actions/cache@v5
+      with:
+        path: ~/.cache/pre-commit
+        key: pre-commit-${{ hashFiles('.pre-commit-config.yaml') }}
+
     - name: Run pre-commit
-      uses: pre-commit/action@v3.0.1
+      run: |
+        python -m pip install --upgrade pip pre-commit
+        pre-commit run --all-files --show-diff-on-failure --color=always

.github/workflows/tests.yml

@@ -10,10 +10,10 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v5
 
     - name: Set up Python 3.9
-      uses: actions/setup-python@v5
+      uses: actions/setup-python@v6
       with:
         python-version: "3.9"
 

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
       - id: end-of-file-fixer
       - id: check-yaml
       - id: check-added-large-files
-        args: ['--maxkb=500']
+        args: ['--maxkb=2500']
       - id: check-merge-conflict
 
   - repo: https://github.com/psf/black
@@ -32,6 +32,16 @@ repos:
       - id: clang-format
         types_or: [c, c++, cuda]
 
+  # Format the MyST Markdown documentation sources. Scoped to docs/ so it does
+  # not reformat README.md / CONTRIBUTING.md (which have intentional layout).
+  - repo: https://github.com/executablebooks/mdformat
+    rev: 0.7.22
+    hooks:
+      - id: mdformat
+        additional_dependencies:
+          - mdformat-myst
+        files: ^docs/.*\.md$
+
   - repo: local
     hooks:
       - id: license-header-check

.readthedocs.yaml

@@ -0,0 +1,23 @@
+# Read the Docs configuration file
+# https://docs.readthedocs.io/en/stable/config-file/v2.html
+
+version: 2
+
+build:
+  os: ubuntu-22.04
+  tools:
+    python: "3.11"
+
+sphinx:
+  configuration: docs/conf.py
+  fail_on_warning: false
+
+python:
+  install:
+    - method: pip
+      path: .
+      extra_requirements:
+        - docs
+
+formats:
+  - htmlzip

CONTRIBUTING.md

@@ -104,7 +104,7 @@ print('disco:', disco_ok())
 If CUDA is not detected automatically (containers, headless nodes):
 
 ```bash
-export FORCE_CUDA_EXTENSION=1
+export TORCH_HARMONICS_BUILD_CUDA_EXTENSION=1
 export TORCH_CUDA_ARCH_LIST="8.0 8.6 9.0"   # set to GPUs you target; reduces compile time
 pip install -e . --no-build-isolation
 ```
@@ -115,7 +115,8 @@ Custom CUDA extensions require compute capability **≥ 8.0**.
 
 | Variable | Effect |
 |----------|--------|
-| `FORCE_CUDA_EXTENSION=1` | Build CUDA kernels even if CUDA is not detected at configure time |
+| `TORCH_HARMONICS_BUILD_CPP_EXTENSION=1` | Build C++ extensions even if PyTorch is not detected at configure time |
+| `TORCH_HARMONICS_BUILD_CUDA_EXTENSION=1` | Build CUDA kernels even if CUDA is not detected at configure time |
 | `TORCH_CUDA_ARCH_LIST` | Limit NVCC target architectures |
 | `TORCH_HARMONICS_DEBUG=1` | Debug flags (`-O0`, `-g`) for extensions |
 | `TORCH_HARMONICS_PROFILE=1` | NVCC lineinfo / PTXAS verbose (CUDA) |
@@ -166,7 +167,7 @@ docker run --gpus all -it --rm --ipc=host --ulimit memlock=-1 --ulimit stack=671
 | Extensions not found / stale behavior | Rebuild with `pip install -e . --no-build-isolation`; remove old `_C*.so` if needed |
 | No matching wheel on install | Install PyTorch first, then torch-harmonics; or build from source |
 | ABI mismatch | Rebuild from source against the same PyTorch version |
-| CUDA not detected at build time | `export FORCE_CUDA_EXTENSION=1` and set `TORCH_CUDA_ARCH_LIST` |
+| CUDA not detected at build time | `export TORCH_HARMONICS_BUILD_CUDA_EXTENSION=1` and set `TORCH_CUDA_ARCH_LIST` |
 
 ## Running tests
 

Dockerfile

@@ -36,8 +36,8 @@ FROM nvcr.io/nvidia/pytorch:25.12-py3
 RUN pip install parameterized
 
 # The custom CUDA extension does not suppport architerctures < 7.0
-ENV FORCE_CUDA_EXTENSION=1
+ENV TORCH_HARMONICS_BUILD_CUDA_EXTENSION=1
 ENV TORCH_HARMONICS_ENABLE_OPENMP=1
-ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0 10.0+PTX"
+ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0a 10.0a+PTX"
 COPY . /workspace/torch_harmonics
 RUN cd /workspace/torch_harmonics && NVCC_THREADS=4 MAX_JOBS=8 pip install --no-build-isolation .

Dockerfile.examples

@@ -48,7 +48,7 @@ RUN pip install h5py
 RUN cd /opt && git clone https://github.com/SHI-Labs/NATTEN natten  && \
     cd natten && \
     make WITH_CUDA=1 \
-        CUDA_ARCH="8.0;8.6;8.7;9.0;10.0" \
+        CUDA_ARCH="8.0;8.6;8.7;9.0a;10.0a" \
         WORKERS=4
 
 # install torch harmonics
@@ -57,5 +57,5 @@ COPY . /workspace/torch_harmonics
 # The custom CUDA extension does not suppport architerctures < 7.0
 ENV FORCE_CUDA_EXTENSION=1
 ENV TORCH_HARMONICS_ENABLE_OPENMP=1
-ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0 10.0+PTX"
+ENV TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0a 10.0a+PTX"
 RUN cd /workspace/torch_harmonics && NVCC_THREADS=4 MAX_JOBS=8 pip install --no-deps --no-build-isolation .

README.md

@@ -118,14 +118,16 @@ cd torch-harmonics
 pip install --no-build-isolation -e .
 ```
 
-If CUDA devices are not detected automatically (e.g. inside a container), set the `FORCE_CUDA_EXTENSION` flag. Set `TORCH_CUDA_ARCH_LIST` to only the architectures you need to reduce compilation time:
+If CUDA devices are not detected automatically (e.g. inside a container), set the `TORCH_HARMONICS_BUILD_CUDA_EXTENSION` flag. Set `TORCH_CUDA_ARCH_LIST` to only the architectures you need to reduce compilation time:
 
 ```bash
-export FORCE_CUDA_EXTENSION=1
-export TORCH_CUDA_ARCH_LIST="8.0 8.6 9.0 10.0+PTX"
+export TORCH_HARMONICS_BUILD_CUDA_EXTENSION=1
+export TORCH_CUDA_ARCH_LIST="8.0 8.6 9.0a 10.0a+PTX"
 pip install --no-build-isolation -e .
 ```
 
+> **Tip:** Use the `a` suffix (e.g. `9.0a`, `10.0a`) instead of plain `9.0` or `10.0` to enable architecture-specific tensor core instructions (e.g. `wgmma` on Hopper). Some layers such as DISCO benefit significantly from this. The trade-off is that the resulting binary is only compatible with the exact GPU generation it was compiled for.
+
 :warning: Custom CUDA extensions require architectures >= 7.0.
 
 Alternatively, build a Docker container:

benchmarks/run.py

@@ -71,7 +71,7 @@ def load_reference_csv(path: str) -> list[dict]:
     return rows
 
 
-def _lookup_reference(reference: Optional[list[dict]], result: "BenchmarkResult") -> Optional[dict]:
+def _lookup_reference(reference: Optional[list[dict]], result: "BenchmarkResult", reference_arch: Optional[str] = None) -> Optional[dict]:
     """Find the reference row matching result's name and architecture."""
     if reference is None:
         return None
@@ -80,17 +80,20 @@ def _lookup_reference(reference: Optional[list[dict]], result: "BenchmarkResult"
     for row in reference:
         if row["name"] != result.name:
             continue
-        if _arch_match(row, result):
+        if _arch_match(row, result, reference_arch):
             return row
         if not row.get("architecture"):
             fallback = row
     return fallback
 
 
-def _arch_match(ref: dict, result: "BenchmarkResult") -> bool:
+def _arch_match(ref: dict, result: "BenchmarkResult", reference_arch: Optional[str] = None) -> bool:
     """Return False if the reference entry was produced on a different GPU."""
     arch = ref.get("architecture", "")
-    return not arch or _fmt_arch(arch) == _fmt_arch(result.arch)
+    if not arch:
+        return True
+    target = _fmt_arch(reference_arch) if reference_arch else _fmt_arch(result.arch)
+    return _fmt_arch(arch) == target
 
 
 def _speedup(ref_ms: Optional[float], cur_ms: Optional[float]) -> Optional[float]:
@@ -143,6 +146,7 @@ def _render_rows(
     reference: Optional[dict[str, dict]],
     tol: float,
     with_err: bool,
+    reference_arch: Optional[str] = None,
 ) -> tuple[list[str], list[list[str]]]:
     """Return (headers, data_rows) as lists of pre-formatted cell strings."""
     headers = ["name", "architecture", "device", "dtype", "fwd_ms", "bwd_ms"]
@@ -153,7 +157,7 @@ def _render_rows(
 
     rows = []
     for r in results:
-        ref = _lookup_reference(reference, r)
+        ref = _lookup_reference(reference, r, reference_arch)
         cells = [
             r.name,
             _fmt_arch(r.arch),
@@ -216,10 +220,11 @@ def print_table(
     skipped: list[tuple[str, str]],
     reference: Optional[dict[str, dict]] = None,
     tol: float = 0.05,
+    reference_arch: Optional[str] = None,
 ) -> int:
     """Print results table. Returns number of flagged regressions."""
     with_err = any(r.ref_error is not None for r in results)
-    headers, rows = _render_rows(results, reference, tol, with_err)
+    headers, rows = _render_rows(results, reference, tol, with_err, reference_arch)
 
     # insert blank lines between benchmark groups (sht / disco / attention)
     separator_after: set = set()
@@ -236,7 +241,7 @@ def print_table(
     n_regressions = 0
     if reference is not None:
         for r in results:
-            ref = _lookup_reference(reference, r)
+            ref = _lookup_reference(reference, r, reference_arch)
             if ref is not None:
                 fwd_spd = _speedup(ref.get("fwd_ms"), r.fwd_ms)
                 bwd_spd = _speedup(ref.get("bwd_ms"), r.bwd_ms)
@@ -306,6 +311,12 @@ def main() -> int:
     parser.add_argument("--save-json", default=None, metavar="PATH", help="write results to a JSON file")
     parser.add_argument("--save-csv", default=None, metavar="PATH", help="write results to a CSV file")
     parser.add_argument("--reference-csv", default=None, metavar="PATH", help="CSV produced by a previous --save-csv run; " "adds speedup columns and flags regressions")
+    parser.add_argument(
+        "--reference-arch",
+        default=None,
+        metavar="ARCH",
+        help="match reference rows by this architecture instead of the current GPU " "(e.g. 'H100 80GB HBM3' to compare a GB200 run against H100 numbers)",
+    )
     parser.add_argument("--regression-tol", type=float, default=0.05, metavar="F", help="slowdowns within this fraction of 1x are not flagged " "(default 0.05 = 5%%)")
     parser.add_argument("--check-outputs", action="store_true", help="run float64/CPU reference and report ref_l_inf error")
     args = parser.parse_args()
@@ -321,9 +332,12 @@ def main() -> int:
     reference = None
     if args.reference_csv:
         reference = load_reference_csv(args.reference_csv)
-        ref_archs = sorted({r["architecture"] for r in reference if r.get("architecture")})
-        arch_str = ", ".join(_fmt_arch(a) for a in ref_archs) if ref_archs else "unknown"
-        print(f"reference : {args.reference_csv}  ({len(reference)} rows, arch: {arch_str})")
+        ref_archs = sorted({_fmt_arch(r["architecture"]) for r in reference if r.get("architecture")})
+        arch_str = ", ".join(ref_archs) if ref_archs else "unknown"
+        print(f"reference : {args.reference_csv}  ({len(reference)} rows)")
+        print(f"avail arch: {arch_str}")
+        if args.reference_arch:
+            print(f"ref arch  : {args.reference_arch}")
         print(f"tolerance : {args.regression_tol*100:.0f}%")
         print()
 
@@ -354,7 +368,7 @@ def main() -> int:
             results.append(result)
 
     print()
-    n_regressions = print_table(results, skipped, reference=reference, tol=args.regression_tol)
+    n_regressions = print_table(results, skipped, reference=reference, tol=args.regression_tol, reference_arch=args.reference_arch)
 
     if args.save_json:
         save_json(results, args.save_json)

docs/.gitignore

@@ -0,0 +1,2 @@
+_build/
+api/generated/