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/Dockerfile b/Dockerfile index 4297d96a..71816a09 100644 --- a/Dockerfile +++ b/Dockerfile @@ -74,6 +74,83 @@ RUN MAX_JOBS="${MAX_JOBS}" pip install -v --disable-pip-version-check --no-cache --config-settings "--build-option=--cuda_ext" \ git+https://github.com/NVIDIA/apex.git +# ── Mamba baseline: mamba-ssm + causal-conv1d from source ───────────────────── +# Optional comparison baseline (e.g. the 2D forward-time benchmark's Mamba2 +# series); NOT a project dependency, so it is installed explicitly here rather +# than via an extra. Placed before COPY of the full tree so unrelated code +# changes do not trigger a rebuild. The tiny patch script is copied first so we +# can rewrite upstream setup.py before compiling. +# +# Upstream hardcodes -gencode for sm_75..sm_120 and ignores TORCH_CUDA_ARCH_LIST, +# which OOMs / gcc-ICEs under QEMU arm64. We clone, patch gencodes to match +# TORCH_CUDA_ARCH_LIST, and honor NVCC_THREADS (arm64 build sets this to 1). +# +# mamba-ssm depends on an unpinned `torch`, so a plain install lets pip swap +# torch / nvidia-cudnn-cu12 out from under the 2.10 stack — which breaks cuDNN +# (CUDNN_STATUS_NOT_INITIALIZED at runtime). Freeze the current torch/nvidia/ +# triton versions into a constraints file so the mamba install cannot touch them. +# Gated by INSTALL_MAMBA (default false — benchmark-only baseline, not a project dep). +# Build with --build-arg INSTALL_MAMBA=true for the benchmark image; the default lean +# build skips the (slow, QEMU-heavy) source compile and the benchmark's Mamba2 series +# then reports 'unavailable' and is omitted. The patch script is COPYed unconditionally +# (tiny, harmless if unused; COPY cannot be gated). +ARG NVCC_THREADS=4 +ARG INSTALL_MAMBA=false +COPY scripts/docker/patch_mamba_cuda_arches.py /tmp/patch_mamba_cuda_arches.py +RUN if [ "${INSTALL_MAMBA}" != "true" ]; then \ + echo "INSTALL_MAMBA=${INSTALL_MAMBA} — skipping Mamba baseline (mamba-ssm / causal-conv1d)"; \ + else \ + pip freeze | grep -iE '^(torch|torchvision|torchaudio|nvidia-|triton|pytorch-triton)' > /tmp/mamba-constraints.txt && \ + cat /tmp/mamba-constraints.txt && \ + git clone --depth 1 https://github.com/Dao-AILab/causal-conv1d.git /tmp/causal-conv1d && \ + git clone --depth 1 https://github.com/state-spaces/mamba.git /tmp/mamba && \ + python /tmp/patch_mamba_cuda_arches.py /tmp/causal-conv1d/setup.py /tmp/mamba/setup.py && \ + MAX_JOBS="${MAX_JOBS}" NVCC_THREADS="${NVCC_THREADS}" \ + CAUSAL_CONV1D_FORCE_BUILD=TRUE MAMBA_FORCE_BUILD=TRUE \ + pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation \ + -c /tmp/mamba-constraints.txt \ + /tmp/causal-conv1d /tmp/mamba && \ + rm -rf /tmp/causal-conv1d /tmp/mamba ; \ + fi + +# ── FlashAttention-4 baseline (Blackwell / Hopper) ──────────────────────────── +# Optional baseline for the forward-time benchmark's FlashAttention-4 series +# (attn_impl="fa4"); NOT a project dependency. FA4 is a pure-Python CuTe-DSL wheel +# that JIT-compiles its kernel at runtime on the target GPU, so — unlike apex / +# mamba above — there is NO CUDA source build here: this layer is cheap and +# QEMU-safe. It needs a Hopper/Blackwell GPU + CUDA >= 12.3 at *run* time (the +# build host needs no GPU; the JIT fires on the first forward). Alpha release, so +# --pre; on a CUDA 13 base use the `[cu13]` extra on both pins below. +# +# VERSION LOCK: flash-attn-4 pins an EXACT matching CuTe-DSL dev build (b23 -> +# nvidia-cutlass-dsl==4.6.0.dev0). A skew between the two crashes the JIT with +# "fmax() takes 2 positional arguments but 3 given" in flash_attn/cute/softmax.py. +# So (1) install the matched pair explicitly, and (2) EXCLUDE nvidia-cutlass-dsl +# from the torch/CUDA constraints pin — otherwise a DSL already present in the base +# image gets frozen to the wrong version and strands FA4 on a mismatched API. The +# rest of the nvidia/torch/triton stack is still pinned (the cuDNN-clobber guard). +# Gated by INSTALL_FA4 (default false — benchmark-only baseline, not a project dep). +# Build with --build-arg INSTALL_FA4=true for the benchmark image; the default lean +# build skips the alpha flash-attn-4 wheel and the benchmark's FA4 series then reports +# 'unavailable' and is omitted. When enabled, the install IS fatal on failure (so a bad +# version pin fails the build loudly); only the post-install import probe is best-effort +# (braced so its `|| echo` cannot mask an install failure) since importing the CuTe DSL +# may touch the CUDA driver on a GPU-less build host. +ARG FLASH_ATTN4_VERSION=4.0.0b23 +ARG CUTLASS_DSL_VERSION=4.6.0.dev0 +ARG INSTALL_FA4=false +RUN if [ "${INSTALL_FA4}" != "true" ]; then \ + echo "INSTALL_FA4=${INSTALL_FA4} — skipping FlashAttention-4 baseline"; \ + else \ + pip freeze | grep -iE '^(torch|torchvision|torchaudio|nvidia-|triton|pytorch-triton)' \ + | grep -viE '^nvidia-cutlass-dsl' > /tmp/fa4-constraints.txt && \ + pip install --no-cache-dir --pre -c /tmp/fa4-constraints.txt \ + "nvidia-cutlass-dsl==${CUTLASS_DSL_VERSION}" "flash-attn-4==${FLASH_ATTN4_VERSION}" && \ + { python -c "import nvidia_cutlass_dsl as c, flash_attn.cute; from flash_attn.cute import flash_attn_func; \ +print('flash-attn-4 import OK / cutlass-dsl', c.__version__)" \ + || echo "WARNING: flash-attn-4 import check failed at build (JIT may probe CUDA); verify on-GPU." ; } ; \ + fi + # ── Dev deps: cached until requirements-dev.txt changes ────────────────────── COPY requirements-dev.txt . RUN pip install --no-cache-dir -r requirements-dev.txt diff --git a/README.md b/README.md index 4a300031..5796ac5a 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 ``` 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/benchmarks/results/README.md b/benchmarks/results/README.md new file mode 100644 index 00000000..79c83599 --- /dev/null +++ b/benchmarks/results/README.md @@ -0,0 +1,129 @@ +# Forward-time scaling benchmarks (ND analogue of Figure 1) + +Single-layer **forward time vs. context size** for HyenaND (nSubQ) against attention +and Mamba2, in 1D / 2D / 3D. This is the N-dimensional analogue of the paper's +Figure 1 (right): it shows attention's `O(L²)` wall while HyenaND scales as +`O(L log L)`, where the token count is `L = Rᴺ` for a square/cubic resolution `R`. + +Two result sets share one protocol and differ in a **single knob — `hidden_dim`** — +which tells two complementary stories: + +| set | dir | hidden_dim | head_dim | reach | question it answers | +| --------- | -------------------------------------- | ---------- | -------- | -------------- | ------------------------------------------------ | +| **Reach** | [`reach_hidden8/`](reach_hidden8/) | 8 | 4 | **16M** tokens | How *far* does the subquadratic operator scale? | +| **Flash** | [`flash_hidden512/`](flash_hidden512/) | 512 | 128 | 1M (3D: 256K) | Is the win real vs. *optimized flash* attention? | + +Generated by [`benchmarks/benchmark_forward_time_nd_resolution.py`](../benchmark_forward_time_nd_resolution.py) +and plotted by [`scripts/visualization/visualize_forward_time_nd.py`](../../scripts/visualization/visualize_forward_time_nd.py). + +## Shared protocol + +- **Hardware / precision:** one NVIDIA **GB200**, `bf16`, batch 1, under + `torch.inference_mode()` + `torch.autocast("cuda", bf16)`. +- **What is timed:** a single mixer *layer* on a channels-last `[B, R, …, C]` input + (not a full residual net) — faithful to Figure 1's operator-level timing. CUDA-event + timing with warmup and adaptive iteration counts; peak memory via + `torch.cuda.max_memory_allocated`. +- **HyenaND backend:** fused `subq_ops` FFT-conv — causal in 1D, non-causal in 2D; + 3D falls back to `torch_fft` (no 3D `subq_ops` kernel yet). +- **The reach ceiling** in every run is torch's **2³¹ 32-bit index limit**: the qkv + tensor `3 · hidden · Rᴺ` must stay under `2³¹`, so a 64× wider `hidden` (8 → 512) + lowers the max token count ~64× (16M → ~256K–1M). It is *not* a memory or + `subq_ops` limit. +- **Status handling:** each point is `ok` / `oom` / `error` / `timeout` / + `unavailable`; a series that never produced a timing (missing dep, unusable config) + is omitted rather than drawn as a wall of ×'s. + +## The two sets + +### Reach — `reach_hidden8/` (hidden 8, head_dim 4) + +Tiny width chosen so the qkv tensor stays under `2³¹` out to **16M tokens**. +Series: HyenaND, Attention (SDPA), Mamba2. *(Flex/FA4 excluded — they require +head_dim ≥ 16.)* + +- HyenaND reaches 16M in every dim (**44 / 47 / 66 ms** at 16M for 1D/2D/3D). +- Attention is `O(L²)` → **272 s** at 16M (2D), i.e. **5807×** slower than HyenaND + (6167× in 1D, 4102× in 3D). +- Mamba2 walls ~1M (1D/2D) / ~2M (3D) at its `causal_conv1d` grid limit. + +> **Caveat:** at head_dim 4, attention is **not flash-eligible** and runs the fp32 +> *math* backend — a weak baseline. The huge ratios partly reflect crippled +> attention. That is exactly why the Flash set exists. + +### Flash — `flash_hidden512/` (hidden 512, head_dim 128) + +A realistic width at the head_dim **flash kernels are optimized for** (FA3/FA4/cuDNN +are tuned for 128). Series: HyenaND, Attention (SDPA), **FlexAttention**, +**FlashAttention-4**, Mamba2. The three attention kernels are apples-to-apples — +same layer, RoPE, QK-norm, and `bf16` q/k/v; only the kernel differs. + +Forward time at **1M tokens (2D)**: + +| operator | ms | vs. HyenaND | +| -------------------- | ---------- | ----------- | +| **HyenaND (nSubQ)** | 91 | — | +| **FlashAttention-4** | 1481 | **16×** | +| Attention (SDPA) | 5183 | 57× | +| FlexAttention | 7549 | 83× | +| Mamba2 | walls @ 1M | — | + +FA4 is the **fastest attention** (~3.5× over SDPA, ~5× over Flex — Blackwell-optimized), +yet HyenaND still beats it **16×**. Reach caps at 1M (1D/2D) / 256K (3D); in 3D the +`L = R³` axis jumps 256K → 2M, so it stops at 256K (the `128³` step busts the 2³¹ wall). + +## Files + +Per set and per dimension `{1,2,3}d`: + +- `forward_time[_flash]_Nd.png` / `.pdf` — the forward-time plot (Figure-1 style). +- `forward_time[_flash]_Nd_memory.png` / `.pdf` — peak-memory plot, with the `O(L²)` + materialized-attention reference line and the GB200 186 GB ceiling. +- `forward_time[_flash]_Nd.jsonl` — raw per-point records (the plot source). +- `*-.out` — Slurm job logs (untracked). + +Each plot's gray subtitle states its config (`hidden N, head_dim D · bf16`), so the +two sets are never confused. + +## Reproduce + +Build a benchmark image (baselines are opt-in — see +[`scripts/slurm/enroot/build_sqsh.sh`](../../scripts/slurm/enroot/build_sqsh.sh)): + +```bash +INSTALL_BASELINES=true PLATFORM=arm64 scripts/slurm/enroot/build_sqsh.sh # + Mamba2 & FA4 +``` + +Then, on the cluster login node: + +```bash +# Reach set (hidden 8) — the default sweep. Writes forward_time_{1,2,3}d.* +sbatch scripts/slurm/submit_forward_time_nd.sh # 2D +DATA_DIM=1 sbatch scripts/slurm/submit_forward_time_nd.sh +DATA_DIM=3 sbatch scripts/slurm/submit_forward_time_nd.sh + +# Flash set (hidden 512, head_dim 128). Writes forward_time_flash_{1,2,3}d.* +scripts/slurm/submit_forward_time_flash_kernels.sh # 2D +DATA_DIM=1 scripts/slurm/submit_forward_time_flash_kernels.sh +DATA_DIM=3 scripts/slurm/submit_forward_time_flash_kernels.sh +``` + +Jobs write to `benchmarks/results/` (both time and memory plots render in-job); sort +the `forward_time_*` outputs into `reach_hidden8/` and the `forward_time_flash_*` +outputs into `flash_hidden512/`. Re-plot from a JSONL anywhere with matplotlib: + +```bash +python scripts/visualization/visualize_forward_time_nd.py \ + --input benchmarks/results/flash_hidden512/forward_time_flash_2d.jsonl \ + --out /tmp/flash_2d.png # add --metric memory for the memory plot +``` + +## Reading notes + +- **Do not compare absolute ms across the two sets** — different `hidden_dim` (64× + apart). Compare operators *within* a plot; use the subtitle to tell the sets apart. +- **Reach** shows the scaling ceiling and attention's `O(L²)` blow-up; **Flash** shows + the win is real (16×) even against FA4 at a production width. +- **Memory:** every measured operator is `O(L)`; the dashed line is the `O(L²)` memory + a *materialized* (non-flash) attention would need — the cost flash avoids by recompute + — crossing the GB200 ceiling around 256K. diff --git a/benchmarks/results/flash_hidden512/forward_time_flash_1d.jsonl b/benchmarks/results/flash_hidden512/forward_time_flash_1d.jsonl new file mode 100644 index 00000000..15f2735c --- /dev/null +++ b/benchmarks/results/flash_hidden512/forward_time_flash_1d.jsonl @@ -0,0 +1,24 @@ +{"mixer": "hyena", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1967123280400815, "mem_gb": 0.08133316040039062, "iters": 23} +{"mixer": "hyena", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.475222396850586, "mem_gb": 0.2805976867675781, "iters": 30} +{"mixer": "hyena", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.2308064778645833, "mem_gb": 1.8960151672363281, "iters": 30} +{"mixer": "hyena", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 13.773561604817708, "mem_gb": 7.521747589111328, "iters": 30} +{"mixer": "hyena", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 92.91827799479167, "mem_gb": 30.030536651611328, "iters": 30} +{"mixer": "attention", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.7067061106363932, "mem_gb": 0.0664682388305664, "iters": 30} +{"mixer": "attention", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.985272471110026, "mem_gb": 0.2189950942993164, "iters": 30} +{"mixer": "attention", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 22.636588541666665, "mem_gb": 0.8282480239868164, "iters": 30} +{"mixer": "attention", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 327.64879557291664, "mem_gb": 3.2686777114868164, "iters": 15} +{"mixer": "attention", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 5156.254557291667, "mem_gb": 13.030396461486816, "iters": 3} +{"mixer": "flex", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9447999795277914, "mem_gb": 0.05859375, "iters": 3} +{"mixer": "flex", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.6805651982625327, "mem_gb": 0.1875, "iters": 3} +{"mixer": "flex", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 29.455210367838543, "mem_gb": 0.7022705078125, "iters": 30} +{"mixer": "flex", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 473.91836825284093, "mem_gb": 2.7647705078125, "iters": 11} +{"mixer": "flex", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7505.627604166667, "mem_gb": 11.0147705078125, "iters": 3} +{"mixer": "fa4", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.8173226515452067, "mem_gb": 0.05859375, "iters": 3} +{"mixer": "fa4", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9170645395914714, "mem_gb": 0.1875, "iters": 30} +{"mixer": "fa4", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 6.684200541178385, "mem_gb": 0.7022705078125, "iters": 30} +{"mixer": "fa4", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 96.54114583333333, "mem_gb": 2.7647705078125, "iters": 30} +{"mixer": "fa4", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1464.6765950520833, "mem_gb": 11.0147705078125, "iters": 3} +{"mixer": "mamba", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.325354735056559, "mem_gb": 0.3266119956970215, "iters": 3} +{"mixer": "mamba", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.303655497233073, "mem_gb": 0.312161922454834, "iters": 30} +{"mixer": "mamba", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.327561696370443, "mem_gb": 1.163419246673584, "iters": 30} +{"mixer": "mamba", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 8.304791259765626, "mem_gb": 4.568448543548584, "iters": 30} diff --git a/benchmarks/results/flash_hidden512/forward_time_flash_1d.png b/benchmarks/results/flash_hidden512/forward_time_flash_1d.png new file mode 100644 index 00000000..dc2c5eb3 Binary files /dev/null and b/benchmarks/results/flash_hidden512/forward_time_flash_1d.png differ diff --git a/benchmarks/results/flash_hidden512/forward_time_flash_1d_memory.png b/benchmarks/results/flash_hidden512/forward_time_flash_1d_memory.png new file mode 100644 index 00000000..a1ddbb7c Binary files /dev/null and b/benchmarks/results/flash_hidden512/forward_time_flash_1d_memory.png differ diff --git a/benchmarks/results/flash_hidden512/forward_time_flash_2d.jsonl b/benchmarks/results/flash_hidden512/forward_time_flash_2d.jsonl new file mode 100644 index 00000000..25cb43b9 --- /dev/null +++ b/benchmarks/results/flash_hidden512/forward_time_flash_2d.jsonl @@ -0,0 +1,24 @@ +{"mixer": "hyena", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.452396035194397, "mem_gb": 0.15062236785888672, "iters": 8} +{"mixer": "hyena", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1616024604210486, "mem_gb": 0.5623340606689453, "iters": 13} +{"mixer": "hyena", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.640416039360894, "mem_gb": 2.2028350830078125, "iters": 9} +{"mixer": "hyena", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 14.633184432983398, "mem_gb": 8.766796112060547, "iters": 4} +{"mixer": "hyena", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 90.8052978515625, "mem_gb": 29.523633003234863, "iters": 3} +{"mixer": "attention", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.8922431945800782, "mem_gb": 0.0626230239868164, "iters": 30} +{"mixer": "attention", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.169860331217448, "mem_gb": 0.2026376724243164, "iters": 30} +{"mixer": "attention", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 23.09128621419271, "mem_gb": 0.7659921646118164, "iters": 30} +{"mixer": "attention", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 330.49306640625, "mem_gb": 3.0191659927368164, "iters": 15} +{"mixer": "attention", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 5183.457356770833, "mem_gb": 12.031373023986816, "iters": 3} +{"mixer": "flex", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2663359642028809, "mem_gb": 0.0587158203125, "iters": 3} +{"mixer": "flex", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.893301328023275, "mem_gb": 0.18701171875, "iters": 3} +{"mixer": "flex", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 29.809527587890624, "mem_gb": 0.7034912109375, "iters": 30} +{"mixer": "flex", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 477.862744140625, "mem_gb": 2.7691650390625, "iters": 10} +{"mixer": "flex", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7549.233072916667, "mem_gb": 11.0313720703125, "iters": 3} +{"mixer": "fa4", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0554239749908447, "mem_gb": 0.0587158203125, "iters": 3} +{"mixer": "fa4", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.4196212768554688, "mem_gb": 0.18701171875, "iters": 30} +{"mixer": "fa4", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7.08994140625, "mem_gb": 0.7034912109375, "iters": 30} +{"mixer": "fa4", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 98.28483072916667, "mem_gb": 2.7691650390625, "iters": 30} +{"mixer": "fa4", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1481.43603515625, "mem_gb": 11.0313720703125, "iters": 3} +{"mixer": "mamba", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.7725652058919272, "mem_gb": 0.3266119956970215, "iters": 3} +{"mixer": "mamba", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.7265225728352864, "mem_gb": 0.312161922454834, "iters": 30} +{"mixer": "mamba", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.7398602803548178, "mem_gb": 1.163419246673584, "iters": 30} +{"mixer": "mamba", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 8.387642415364583, "mem_gb": 4.568448543548584, "iters": 30} diff --git a/benchmarks/results/flash_hidden512/forward_time_flash_2d.png b/benchmarks/results/flash_hidden512/forward_time_flash_2d.png new file mode 100644 index 00000000..b3418a63 Binary files /dev/null and b/benchmarks/results/flash_hidden512/forward_time_flash_2d.png differ diff --git a/benchmarks/results/flash_hidden512/forward_time_flash_2d_memory.png b/benchmarks/results/flash_hidden512/forward_time_flash_2d_memory.png new file mode 100644 index 00000000..5e0fb434 Binary files /dev/null and b/benchmarks/results/flash_hidden512/forward_time_flash_2d_memory.png differ diff --git a/benchmarks/results/flash_hidden512/forward_time_flash_3d.jsonl b/benchmarks/results/flash_hidden512/forward_time_flash_3d.jsonl new file mode 100644 index 00000000..f8fe9e13 --- /dev/null +++ b/benchmarks/results/flash_hidden512/forward_time_flash_3d.jsonl @@ -0,0 +1,15 @@ +{"mixer": "hyena", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": "torch_fft", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.136843288646025, "mem_gb": 0.1871204376220703, "iters": 17} +{"mixer": "hyena", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": "torch_fft", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7.887974330357143, "mem_gb": 1.3487305641174316, "iters": 21} +{"mixer": "hyena", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": "torch_fft", "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 71.35670166015625, "mem_gb": 10.588771343231201, "iters": 5} +{"mixer": "attention", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.6075071970621745, "mem_gb": 0.0625619888305664, "iters": 30} +{"mixer": "attention", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 5.643666076660156, "mem_gb": 0.3902597427368164, "iters": 30} +{"mixer": "attention", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 323.7279296875, "mem_gb": 3.0186777114868164, "iters": 15} +{"mixer": "flex", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.7944640318552653, "mem_gb": 0.0546875, "iters": 3} +{"mixer": "flex", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 7.430975914001465, "mem_gb": 0.3272705078125, "iters": 4} +{"mixer": "flex", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 469.77587890625, "mem_gb": 2.5147705078125, "iters": 11} +{"mixer": "fa4", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.7151679992675781, "mem_gb": 0.0546875, "iters": 3} +{"mixer": "fa4", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.7867637634277345, "mem_gb": 0.3272705078125, "iters": 30} +{"mixer": "fa4", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 97.94173177083333, "mem_gb": 2.5147705078125, "iters": 30} +{"mixer": "mamba", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.5995093981424966, "mem_gb": 0.3266119956970215, "iters": 3} +{"mixer": "mamba", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.6365119934082033, "mem_gb": 0.595914363861084, "iters": 30} +{"mixer": "mamba", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 512, "num_heads": 4, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 8.542538452148438, "mem_gb": 4.568448543548584, "iters": 30} diff --git a/benchmarks/results/flash_hidden512/forward_time_flash_3d.png b/benchmarks/results/flash_hidden512/forward_time_flash_3d.png new file mode 100644 index 00000000..69d2fcac Binary files /dev/null and b/benchmarks/results/flash_hidden512/forward_time_flash_3d.png differ diff --git a/benchmarks/results/flash_hidden512/forward_time_flash_3d_memory.png b/benchmarks/results/flash_hidden512/forward_time_flash_3d_memory.png new file mode 100644 index 00000000..1f0b933d Binary files /dev/null and b/benchmarks/results/flash_hidden512/forward_time_flash_3d_memory.png differ diff --git a/benchmarks/results/reach_hidden8/forward_time_1d.jsonl b/benchmarks/results/reach_hidden8/forward_time_1d.jsonl new file mode 100644 index 00000000..fbdcd187 --- /dev/null +++ b/benchmarks/results/reach_hidden8/forward_time_1d.jsonl @@ -0,0 +1,21 @@ +{"mixer": "attention", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1094367980957032, "mem_gb": 0.011264324188232422, "iters": 30} +{"mixer": "attention", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.8594603220621744, "mem_gb": 0.013186931610107422, "iters": 30} +{"mixer": "attention", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 5.14368642171224, "mem_gb": 0.029788494110107422, "iters": 30} +{"mixer": "attention", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 69.2323974609375, "mem_gb": 0.09192228317260742, "iters": 30} +{"mixer": "attention", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1068.3568359375, "mem_gb": 0.3438754081726074, "iters": 5} +{"mixer": "attention", "resolution": 4194304, "seq_len": 4194304, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 17067.8984375, "mem_gb": 1.3516879081726074, "iters": 3} +{"mixer": "attention", "resolution": 16777216, "seq_len": 16777216, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 272240.4375, "mem_gb": 5.382937431335449, "iters": 1} +{"mixer": "hyena", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.2253140767415365, "mem_gb": 0.010535717010498047, "iters": 30} +{"mixer": "hyena", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.5129173278808594, "mem_gb": 0.015342235565185547, "iters": 30} +{"mixer": "hyena", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.5983509063720702, "mem_gb": 0.04035139083862305, "iters": 30} +{"mixer": "hyena", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.5913035074869792, "mem_gb": 0.12718915939331055, "iters": 30} +{"mixer": "hyena", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.5504468282063804, "mem_gb": 0.48177289962768555, "iters": 30} +{"mixer": "hyena", "resolution": 4194304, "seq_len": 4194304, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 10.832672119140625, "mem_gb": 1.8999247550964355, "iters": 30} +{"mixer": "hyena", "resolution": 16777216, "seq_len": 16777216, "data_dim": 1, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 44.14231770833333, "mem_gb": 7.5721659660339355, "iters": 30} +{"mixer": "mamba", "resolution": 4096, "seq_len": 4096, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.8152745564778647, "mem_gb": 0.2675185203552246, "iters": 3} +{"mixer": "mamba", "resolution": 16384, "seq_len": 16384, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.8554122924804686, "mem_gb": 0.04393625259399414, "iters": 30} +{"mixer": "mamba", "resolution": 65536, "seq_len": 65536, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.9237940470377604, "mem_gb": 0.14885568618774414, "iters": 30} +{"mixer": "mamba", "resolution": 262144, "seq_len": 262144, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.18571294148763, "mem_gb": 0.5695099830627441, "iters": 30} +{"mixer": "mamba", "resolution": 1048576, "seq_len": 1048576, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 12.136154174804688, "mem_gb": 2.247244358062744, "iters": 30} +{"mixer": "mamba", "resolution": 4194304, "seq_len": 4194304, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid configuration argument\\nSearch for `cudaErrorInvalidConfiguration' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.\\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\\n\")"} +{"mixer": "mamba", "resolution": 16777216, "seq_len": 16777216, "data_dim": 1, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid configuration argument\\nSearch for `cudaErrorInvalidConfiguration' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.\\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\\n\")"} diff --git a/benchmarks/results/reach_hidden8/forward_time_1d.png b/benchmarks/results/reach_hidden8/forward_time_1d.png new file mode 100644 index 00000000..3e3c33f9 Binary files /dev/null and b/benchmarks/results/reach_hidden8/forward_time_1d.png differ diff --git a/benchmarks/results/reach_hidden8/forward_time_1d_memory.png b/benchmarks/results/reach_hidden8/forward_time_1d_memory.png new file mode 100644 index 00000000..b357fe0f Binary files /dev/null and b/benchmarks/results/reach_hidden8/forward_time_1d_memory.png differ diff --git a/benchmarks/results/reach_hidden8/forward_time_2d.jsonl b/benchmarks/results/reach_hidden8/forward_time_2d.jsonl new file mode 100644 index 00000000..be4d5c56 --- /dev/null +++ b/benchmarks/results/reach_hidden8/forward_time_2d.jsonl @@ -0,0 +1,21 @@ +{"mixer": "attention", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.7784586588541667, "mem_gb": 0.011144161224365234, "iters": 30} +{"mixer": "attention", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1518709818522135, "mem_gb": 0.012702465057373047, "iters": 30} +{"mixer": "attention", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 5.1782679239908855, "mem_gb": 0.027842998504638672, "iters": 30} +{"mixer": "attention", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 69.2794189453125, "mem_gb": 0.08412504196166992, "iters": 30} +{"mixer": "attention", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1065.76044921875, "mem_gb": 0.3126559257507324, "iters": 5} +{"mixer": "attention", "resolution": 2048, "seq_len": 4194304, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 17014.041666666668, "mem_gb": 1.2267489433288574, "iters": 3} +{"mixer": "attention", "resolution": 4096, "seq_len": 16777216, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 271898.21875, "mem_gb": 4.883059501647949, "iters": 1} +{"mixer": "hyena", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1649413108825684, "mem_gb": 0.010765552520751953, "iters": 12} +{"mixer": "hyena", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9916480137751653, "mem_gb": 0.017599105834960938, "iters": 26} +{"mixer": "hyena", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9963790453397311, "mem_gb": 0.044452667236328125, "iters": 26} +{"mixer": "hyena", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.0030330160389775, "mem_gb": 0.1475982666015625, "iters": 23} +{"mixer": "hyena", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.6219893561469183, "mem_gb": 0.5636062622070312, "iters": 18} +{"mixer": "hyena", "resolution": 2048, "seq_len": 4194304, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 11.253952026367188, "mem_gb": 2.2276835441589355, "iters": 7} +{"mixer": "hyena", "resolution": 4096, "seq_len": 16777216, "data_dim": 2, "backend": "subq_ops", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 46.82208251953125, "mem_gb": 8.883872985839844, "iters": 3} +{"mixer": "mamba", "resolution": 64, "seq_len": 4096, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.422154744466146, "mem_gb": 0.2675185203552246, "iters": 3} +{"mixer": "mamba", "resolution": 128, "seq_len": 16384, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.358257039388021, "mem_gb": 0.04393625259399414, "iters": 30} +{"mixer": "mamba", "resolution": 256, "seq_len": 65536, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.379887898763021, "mem_gb": 0.14885568618774414, "iters": 30} +{"mixer": "mamba", "resolution": 512, "seq_len": 262144, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.144506581624349, "mem_gb": 0.5695099830627441, "iters": 30} +{"mixer": "mamba", "resolution": 1024, "seq_len": 1048576, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 12.05062255859375, "mem_gb": 2.247244358062744, "iters": 30} +{"mixer": "mamba", "resolution": 2048, "seq_len": 4194304, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid configuration argument\\nSearch for `cudaErrorInvalidConfiguration' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.\\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\\n\")"} +{"mixer": "mamba", "resolution": 4096, "seq_len": 16777216, "data_dim": 2, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid configuration argument\\nSearch for `cudaErrorInvalidConfiguration' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.\\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\\n\")"} diff --git a/benchmarks/results/reach_hidden8/forward_time_2d.png b/benchmarks/results/reach_hidden8/forward_time_2d.png new file mode 100644 index 00000000..bb841f06 Binary files /dev/null and b/benchmarks/results/reach_hidden8/forward_time_2d.png differ diff --git a/benchmarks/results/reach_hidden8/forward_time_2d_memory.png b/benchmarks/results/reach_hidden8/forward_time_2d_memory.png new file mode 100644 index 00000000..174eb5e3 Binary files /dev/null and b/benchmarks/results/reach_hidden8/forward_time_2d_memory.png differ diff --git a/benchmarks/results/reach_hidden8/forward_time_3d.jsonl b/benchmarks/results/reach_hidden8/forward_time_3d.jsonl new file mode 100644 index 00000000..6543bb0d --- /dev/null +++ b/benchmarks/results/reach_hidden8/forward_time_3d.jsonl @@ -0,0 +1,15 @@ +{"mixer": "attention", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.44606186548868815, "mem_gb": 0.011142253875732422, "iters": 30} +{"mixer": "attention", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.3486378987630208, "mem_gb": 0.017459392547607422, "iters": 30} +{"mixer": "attention", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 69.33837890625, "mem_gb": 0.08410978317260742, "iters": 30} +{"mixer": "attention", "resolution": 128, "seq_len": 2097152, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 4270.072916666667, "mem_gb": 0.6173129081726074, "iters": 3} +{"mixer": "attention", "resolution": 256, "seq_len": 16777216, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 272285.96875, "mem_gb": 4.882937431335449, "iters": 1} +{"mixer": "hyena", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": "torch_fft", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9599826335906982, "mem_gb": 0.011609554290771484, "iters": 24} +{"mixer": "hyena", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": "torch_fft", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 0.9508597056070963, "mem_gb": 0.03052997589111328, "iters": 30} +{"mixer": "hyena", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": "torch_fft", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 1.1017120361328125, "mem_gb": 0.17970037460327148, "iters": 30} +{"mixer": "hyena", "resolution": 128, "seq_len": 2097152, "data_dim": 3, "backend": "torch_fft", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 8.126579284667969, "mem_gb": 1.3500127792358398, "iters": 18} +{"mixer": "hyena", "resolution": 256, "seq_len": 16777216, "data_dim": 3, "backend": "torch_fft", "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 66.38068135579427, "mem_gb": 10.715025901794434, "iters": 6} +{"mixer": "mamba", "resolution": 16, "seq_len": 4096, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.251157283782959, "mem_gb": 0.2675185203552246, "iters": 3} +{"mixer": "mamba", "resolution": 32, "seq_len": 32768, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 2.200347646077474, "mem_gb": 0.07988595962524414, "iters": 30} +{"mixer": "mamba", "resolution": 64, "seq_len": 262144, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 3.065058135986328, "mem_gb": 0.5695099830627441, "iters": 30} +{"mixer": "mamba", "resolution": 128, "seq_len": 2097152, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "ok", "ms": 23.239107259114583, "mem_gb": 4.485525608062744, "iters": 30} +{"mixer": "mamba", "resolution": 256, "seq_len": 16777216, "data_dim": 3, "backend": null, "batch_size": 1, "hidden_dim": 8, "dtype": "bf16", "device": "NVIDIA GB200", "status": "error", "ms": null, "mem_gb": null, "iters": null, "detail": "AcceleratorError(\"CUDA error: invalid configuration argument\\nSearch for `cudaErrorInvalidConfiguration' in https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__TYPES.html for more information.\\nCUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.\\nFor debugging consider passing CUDA_LAUNCH_BLOCKING=1\\nCompile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.\\n\")"} diff --git a/benchmarks/results/reach_hidden8/forward_time_3d.png b/benchmarks/results/reach_hidden8/forward_time_3d.png new file mode 100644 index 00000000..3898e6f4 Binary files /dev/null and b/benchmarks/results/reach_hidden8/forward_time_3d.png differ diff --git a/benchmarks/results/reach_hidden8/forward_time_3d_memory.png b/benchmarks/results/reach_hidden8/forward_time_3d_memory.png new file mode 100644 index 00000000..235da1f8 Binary files /dev/null and b/benchmarks/results/reach_hidden8/forward_time_3d_memory.png differ 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/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..6dc412f1 100755 --- a/scripts/slurm/enroot/build_sqsh.sh +++ b/scripts/slurm/enroot/build_sqsh.sh @@ -9,31 +9,69 @@ # 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="9.0"; MAX_JOBS_DEFAULT=""; NVCC_THREADS_DEFAULT="4" ;; + arm64) DOCKER_PLATFORM="linux/arm64"; TARGET_HW="GB200 (ARM64)"; CUDA_ARCHS="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 +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)." +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 +80,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/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()